diff --git a/README.md b/README.md index b4a82b1..1bf1670 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,7 @@ of the MCP specification. - [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) @@ -1184,6 +1185,35 @@ try await server.start(transport: transport) { clientInfo, clientCapabilities in } ``` +### 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 diff --git a/Sources/MCP/Base/Authorization/OAuthAuthorizer.swift b/Sources/MCP/Base/Authorization/OAuthAuthorizer.swift index 2cac349..b2e5be3 100644 --- a/Sources/MCP/Base/Authorization/OAuthAuthorizer.swift +++ b/Sources/MCP/Base/Authorization/OAuthAuthorizer.swift @@ -217,7 +217,6 @@ public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable { return nil } if accessToken.isExpired() { - tokenStorage.clear() return nil } return "\(OAuthTokenType.bearer) \(accessToken.value)" @@ -367,7 +366,6 @@ public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable { 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() else { return } guard token.isExpired(skewSeconds: configuration.proactiveRefreshWindowSeconds) else { return } @@ -441,8 +439,10 @@ public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable { candidates.append(fallback) } + let fallbackIssuer = try? discoveryClient.metadataDiscovery + .authorizationServerFallbackIssuer(from: endpoint) let metadata = try await discoveryClient.fetchProtectedResourceMetadata( - candidates: candidates, session: session) + candidates: candidates, fallbackIssuer: fallbackIssuer, session: session) try validateProtectedResource(metadata: metadata, endpoint: endpoint) self.protectedResourceMetadata = metadata @@ -716,10 +716,17 @@ public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable { expiresAt: expiresAt, scopes: scopeSet, authorizationServer: selectedAuthorizationServer, - refreshToken: decoded.refreshToken + 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 { @@ -823,7 +830,8 @@ public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable { expiresAt: nil, scopes: requestedScopes ?? [], authorizationServer: authorizationServer, - refreshToken: nil + refreshToken: nil, + clientID: nonEmptyClientID() )) } diff --git a/Sources/MCP/Base/Authorization/OAuthClientRegistrar.swift b/Sources/MCP/Base/Authorization/OAuthClientRegistrar.swift index 2bbdf13..6267421 100644 --- a/Sources/MCP/Base/Authorization/OAuthClientRegistrar.swift +++ b/Sources/MCP/Base/Authorization/OAuthClientRegistrar.swift @@ -79,7 +79,7 @@ struct OAuthClientRegistrar: Sendable { let responseTypes: [String] switch configuration.grantType { case .authorizationCode: - grantTypes = [OAuthGrantTypeValue.authorizationCode] + grantTypes = [OAuthGrantTypeValue.authorizationCode, OAuthGrantTypeValue.refreshToken] responseTypes = [OAuthParameterName.code] case .clientCredentials: grantTypes = [OAuthGrantTypeValue.clientCredentials] diff --git a/Sources/MCP/Base/Authorization/OAuthDiscoveryClient.swift b/Sources/MCP/Base/Authorization/OAuthDiscoveryClient.swift index 55bdb09..8795ed7 100644 --- a/Sources/MCP/Base/Authorization/OAuthDiscoveryClient.swift +++ b/Sources/MCP/Base/Authorization/OAuthDiscoveryClient.swift @@ -7,7 +7,7 @@ import Foundation /// Internal protocol for fetching OAuth discovery metadata. protocol OAuthDiscoveryFetching: Sendable { var metadataDiscovery: any OAuthMetadataDiscovering { get } - func fetchProtectedResourceMetadata(candidates: [URL], session: URLSession) async throws -> OAuthProtectedResourceMetadata + func fetchProtectedResourceMetadata(candidates: [URL], fallbackIssuer: URL?, session: URLSession) async throws -> OAuthProtectedResourceMetadata func fetchAuthorizationServerMetadata(candidates: [URL], session: URLSession) async throws -> (server: URL, metadata: OAuthAuthorizationServerMetadata) } @@ -29,8 +29,13 @@ struct OAuthDiscoveryClient: Sendable { } /// 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() @@ -56,15 +61,28 @@ struct OAuthDiscoveryClient: Sendable { continue } } + if let fallbackIssuer { + return OAuthProtectedResourceMetadata( + resource: nil, + authorizationServers: [fallbackIssuer], + scopesSupported: nil + ) + } throw OAuthAuthorizationError.metadataDiscoveryFailed } - /// Fetches Authorization Server Metadata from the first candidate that returns a valid response. + /// 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 @@ -95,22 +113,28 @@ struct OAuthDiscoveryClient: Sendable { let asMetadata = try decoder.decode( OAuthAuthorizationServerMetadata.self, from: data) - // RFC 8414 §3: issuer field must match the candidate server URL. - // Absent issuer is tolerated (some servers omit it). - if let metadataIssuer = asMetadata.issuer { - guard metadataIssuer.absoluteString.lowercased() + // Prefer metadata whose issuer matches the candidate URL (RFC 8414 §3). + let issuerMatches = + asMetadata.issuer == nil + || asMetadata.issuer?.absoluteString.lowercased() == candidateServer.absoluteString.lowercased() - else { - continue - } + 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) } - - return (server: candidateServer, metadata: asMetadata) } catch { continue } } } + + if let firstValid { + return firstValid + } throw OAuthAuthorizationError.authorizationServerMetadataDiscoveryFailed } } diff --git a/Sources/MCP/Base/Authorization/OAuthModels.swift b/Sources/MCP/Base/Authorization/OAuthModels.swift index d5451c3..28043b9 100644 --- a/Sources/MCP/Base/Authorization/OAuthModels.swift +++ b/Sources/MCP/Base/Authorization/OAuthModels.swift @@ -113,7 +113,7 @@ struct OAuthTokenErrorResponse: Decodable { /// /// 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 { +public struct OAuthAccessToken: Sendable, Codable { /// The raw bearer token string for use in the `Authorization` header. public let value: String @@ -135,6 +135,14 @@ public struct OAuthAccessToken: Sendable { /// 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, @@ -142,7 +150,8 @@ public struct OAuthAccessToken: Sendable { expiresAt: Date?, scopes: Set, authorizationServer: URL?, - refreshToken: String? + refreshToken: String?, + clientID: String? = nil ) { self.value = value self.tokenType = tokenType @@ -150,6 +159,7 @@ public struct OAuthAccessToken: Sendable { 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. diff --git a/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift b/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift index 36c517e..bd56191 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift @@ -207,6 +207,24 @@ struct SSEEvent: Sendable { } } +// 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. diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift index 3f1ec89..c7a4d48 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift @@ -29,7 +29,7 @@ import Logging /// 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 { +public actor StatefulHTTPServerTransport: Transport, HTTPContextProviding { public nonisolated let logger: Logger // MARK: - Dependencies @@ -54,6 +54,11 @@ public actor StatefulHTTPServerTransport: Transport { /// 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. @@ -280,6 +285,7 @@ public actor StatefulHTTPServerTransport: Transport { // 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) @@ -406,6 +412,7 @@ public actor StatefulHTTPServerTransport: Transport { continuation.finish() requestSSEContinuations.removeValue(forKey: requestID) } + httpRequestContexts.removeValue(forKey: requestID) } private func routeServerInitiatedMessage(_ data: Data) { @@ -535,6 +542,13 @@ public actor StatefulHTTPServerTransport: Transport { 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 @@ -552,6 +566,7 @@ public actor StatefulHTTPServerTransport: Transport { continuation.finish() } requestSSEContinuations.removeAll() + httpRequestContexts.removeAll() // Close standalone GET stream standaloneSSEContinuation?.finish() diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift index 654cf70..abe9576 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift @@ -29,7 +29,7 @@ import Logging /// - Session management is handled externally or not needed /// /// For full streaming and session support, use ``StatefulHTTPServerTransport`` instead. -public actor StatelessHTTPServerTransport: Transport { +public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { public nonisolated let logger: Logger // MARK: - Dependencies @@ -52,6 +52,11 @@ public actor StatelessHTTPServerTransport: Transport { /// 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. @@ -207,11 +212,16 @@ public actor StatelessHTTPServerTransport: Transport { return .accepted() case .request(let id, _): - return await handleJSONRPCRequest(body, requestID: id) + return await handleJSONRPCRequest(body, requestID: id, request: request) } } - private func handleJSONRPCRequest(_ body: Data, requestID: String) async -> HTTPResponse { + private func handleJSONRPCRequest( + _ body: Data, + requestID: String, + request: HTTPRequest + ) async -> HTTPResponse { + httpRequestContexts[requestID] = request // Yield the incoming message to the server incomingContinuation.yield(body) @@ -222,15 +232,23 @@ public actor StatelessHTTPServerTransport: Transport { 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 { @@ -245,6 +263,7 @@ public actor StatelessHTTPServerTransport: Transport { logger.debug("Cancelled waiter for request", metadata: ["requestID": "\(id)"]) } responseWaiters.removeAll() + httpRequestContexts.removeAll() // Close incoming stream incomingContinuation.finish() diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 99fbe5c..060b51f 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -298,13 +298,41 @@ public actor Server { // MARK: - Request Context - /// The JSON-RPC request ID of the currently executing method handler. + /// 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. Accessible package-wide for - /// transports that need to identify the active request (e.g. closing an - /// SSE stream mid-call for reconnection testing per SEP-1699). - @TaskLocal package static var currentRequestID: ID? = nil + /// 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 @@ -750,10 +778,17 @@ public actor Server { 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 currentRequestID as a task local so handlers can identify the active request. + // Set currentHandlerContext as a task local so handlers see it. var handlerTask: Task, Error>! - Server.$currentRequestID.withValue(request.id) { + Server.$currentHandlerContext.withValue(handlerContext) { handlerTask = Task, Error> { do { // Check if task was cancelled before starting diff --git a/Sources/MCPConformance/Server/main.swift b/Sources/MCPConformance/Server/main.swift index d15e1b7..6a7ccea 100644 --- a/Sources/MCPConformance/Server/main.swift +++ b/Sources/MCPConformance/Server/main.swift @@ -180,8 +180,7 @@ func createConformanceServer(state: ServerState, transport: StatefulHTTPServerTr // 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.currentRequestID { - + if let requestID = Server.currentHandlerContext?.id { await transport.closeSSEStream(forRequestID: requestID.description) } // Wait briefly for the client to reconnect before sending the response. diff --git a/Tests/MCPTests/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift index 3127534..8e1c89b 100644 --- a/Tests/MCPTests/HTTPServerTransportTests.swift +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -1052,3 +1052,218 @@ struct StatelessHTTPServerTransportTests { #expect(response.statusCode == 500) } } + +// MARK: - HTTPContextProviding / Server.currentHandlerContext + +@Suite("Server.currentHandlerContext Tests") +struct ServerHandlerContextTests { + + // MARK: Stateful transport — HTTPContextProviding + + @Test("Stateful: httpRequestContext returns in-flight request, cleared after response") + func testStatefulContextLifecycle() async throws { + let sessionID = "test-session" + let transport = makeStatefulTransport( + sessionIDGenerator: FixedSessionIDGenerator(sessionID: sessionID) + ) + _ = try await initializeSession(transport: transport, sessionID: sessionID) + + // POST a tools/list request with a distinctive Authorization header and path. + let body = makeRequestBody(id: "42", method: "tools/list") + let httpRequest = HTTPRequest( + method: "POST", + headers: [ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Mcp-Session-Id": sessionID, + "Authorization": "Bearer token-xyz", + ], + body: body, + path: "/mcp" + ) + let response = await transport.handleRequest(httpRequest) + #expect(response.statusCode == 200) + + // In-flight: context surfaces through the protocol lookup. + let inFlight = await transport.httpRequestContext(for: .string("42")) + #expect(inFlight != nil) + #expect(inFlight?.header("Authorization") == "Bearer token-xyz") + #expect(inFlight?.path == "/mcp") + + // Drain the SSE stream in the background so the response routes cleanly. + if case .stream(let stream, _) = response { + Task { for try await _ in stream {} } + } + + // Send the response — transport clears the context alongside the SSE map. + try await transport.send(makeResponseBody(id: "42")) + try await Task.sleep(for: .milliseconds(50)) + + let afterResponse = await transport.httpRequestContext(for: .string("42")) + #expect(afterResponse == nil) + } + + @Test("Stateful: context cleared on session termination") + func testStatefulContextClearedOnTerminate() async throws { + let sessionID = "term-session" + let transport = makeStatefulTransport( + sessionIDGenerator: FixedSessionIDGenerator(sessionID: sessionID) + ) + _ = try await initializeSession(transport: transport, sessionID: sessionID) + + // Register an in-flight request to populate the context map. + let body = makeRequestBody(id: "7", method: "tools/list") + _ = await transport.handleRequest( + makeStatefulPOSTRequest( + body: body, + sessionID: sessionID, + authorization: "Bearer keep" + ) + ) + #expect(await transport.httpRequestContext(for: .string("7")) != nil) + + // DELETE terminates the session. + _ = await transport.handleRequest(makeDELETERequest(sessionID: sessionID)) + + #expect(await transport.httpRequestContext(for: .string("7")) == nil) + } + + // MARK: Stateless transport — HTTPContextProviding + + @Test("Stateless: context available during handling, cleared after response") + func testStatelessContextLifecycle() async throws { + let transport = makeStatelessTransport() + try await transport.connect() + + let httpRequest = HTTPRequest( + method: "POST", + headers: [ + "Content-Type": "application/json", + "Accept": "application/json", + "Authorization": "Bearer stateless-token", + ], + body: makeRequestBody(id: "99", method: "tools/list"), + path: "/mcp" + ) + + // handleRequest blocks until a response is sent — run it concurrently. + let handleTask = Task { await transport.handleRequest(httpRequest) } + + // Give the handler time to register its waiter (and store the context). + try await Task.sleep(for: .milliseconds(50)) + + let inFlight = await transport.httpRequestContext(for: .string("99")) + #expect(inFlight != nil) + #expect(inFlight?.header("Authorization") == "Bearer stateless-token") + #expect(inFlight?.path == "/mcp") + + // Send the response to unblock the waiter. + try await transport.send(makeResponseBody(id: "99")) + _ = await handleTask.value + + let afterResponse = await transport.httpRequestContext(for: .string("99")) + #expect(afterResponse == nil) + + await transport.disconnect() + } + + @Test("Stateless: unknown request id returns nil") + func testStatelessContextUnknownID() async throws { + let transport = makeStatelessTransport() + try await transport.connect() + let none = await transport.httpRequestContext(for: .string("does-not-exist")) + #expect(none == nil) + await transport.disconnect() + } + + // MARK: Server integration — end-to-end + + @Test("Server handler reads HTTP context via Server.currentHandlerContext") + func testServerHandlerReadsHTTPContext() async throws { + actor Captured { + var auth: String? + var path: String? + var id: ID? + func record(auth: String?, path: String?, id: ID?) { + self.auth = auth + self.path = path + self.id = id + } + } + let captured = Captured() + + let transport = makeStatelessTransport() + let server = Server(name: "TestServer", version: "1.0") + + await server.withMethodHandler(CallTool.self) { _ in + let ctx = Server.currentHandlerContext + await captured.record( + auth: ctx?.httpContext?.header("Authorization"), + path: ctx?.httpContext?.path, + id: ctx?.id + ) + return CallTool.Result(content: [.text(text: "ok", annotations: nil, _meta: nil)]) + } + + try await server.start(transport: transport) + defer { Task { await server.stop() } } + + let body = try JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", + "id": "call-1", + "method": "tools/call", + "params": ["name": "whoami"] as [String: Any], + ]) + let httpRequest = HTTPRequest( + method: "POST", + headers: [ + "Content-Type": "application/json", + "Accept": "application/json", + "Authorization": "Bearer server-integration-token", + ], + body: body, + path: "/mcp" + ) + + let httpResponse = await transport.handleRequest(httpRequest) + #expect(httpResponse.statusCode == 200) + + #expect(await captured.auth == "Bearer server-integration-token") + #expect(await captured.path == "/mcp") + #expect(await captured.id == .string("call-1")) + } + + @Test("Non-HTTP transport yields nil httpContext") + func testNonHTTPTransportNilContext() async throws { + actor Captured { + var httpContext: HTTPRequest? + var sawDispatch = false + func set(_ v: HTTPRequest?) { httpContext = v; sawDispatch = true } + } + let captured = Captured() + + let transport = MockTransport() + let server = Server(name: "TestServer", version: "1.0") + + // Start first — `start` calls registerDefaultHandlers which would otherwise + // clobber our override. + try await server.start(transport: transport) + + // Override the Ping handler so we can observe the task-local during dispatch. + await server.withMethodHandler(Ping.self) { _ in + await captured.set(Server.currentHandlerContext?.httpContext) + return Empty() + } + + try await transport.queue(request: Ping.request(.init())) + + // Wait for dispatch. + try await Task.sleep(for: .milliseconds(100)) + + #expect(await captured.sawDispatch == true) + #expect(await captured.httpContext == nil) + + await server.stop() + await transport.disconnect() + } +} diff --git a/Tests/MCPTests/OAuthAuthorizerTests.swift b/Tests/MCPTests/OAuthAuthorizerTests.swift index b079652..0384307 100644 --- a/Tests/MCPTests/OAuthAuthorizerTests.swift +++ b/Tests/MCPTests/OAuthAuthorizerTests.swift @@ -64,7 +64,7 @@ final class MockDiscoveryClient: OAuthDiscoveryFetching, @unchecked Sendable { ) } - func fetchProtectedResourceMetadata(candidates: [URL], session: URLSession) async throws -> OAuthProtectedResourceMetadata { + func fetchProtectedResourceMetadata(candidates: [URL], fallbackIssuer: URL?, session: URLSession) async throws -> OAuthProtectedResourceMetadata { fetchProtectedResourceMetadataCallCount += 1 return protectedResourceMetadataResult } @@ -100,6 +100,10 @@ final class MockTokenClient: OAuthTokenRequesting, @unchecked Sendable { final class MockClientRegistrar: OAuthClientRegistering, @unchecked Sendable { var registerCallCount = 0 + var registrationResult: ( + response: OAuthClientRegistrationResponse, + updatedAuthentication: OAuthConfiguration.TokenEndpointAuthentication + )? func register( configuration: OAuthConfiguration, @@ -110,7 +114,7 @@ final class MockClientRegistrar: OAuthClientRegistering, @unchecked Sendable { updatedAuthentication: OAuthConfiguration.TokenEndpointAuthentication )? { registerCallCount += 1 - return nil + return registrationResult } } @@ -381,4 +385,45 @@ struct OAuthAuthorizerTests { #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/OAuthDiscoveryClientTests.swift b/Tests/MCPTests/OAuthDiscoveryClientTests.swift index 31b74b8..b962bb6 100644 --- a/Tests/MCPTests/OAuthDiscoveryClientTests.swift +++ b/Tests/MCPTests/OAuthDiscoveryClientTests.swift @@ -49,6 +49,7 @@ import Testing let metadata = try await makeClient().fetchProtectedResourceMetadata( candidates: [URL(string: "https://example.com/.well-known/oauth-protected-resource")!], + fallbackIssuer: nil, session: session ) let expected = OAuthProtectedResourceMetadata( @@ -76,6 +77,7 @@ import Testing 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( @@ -103,6 +105,7 @@ import Testing 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( @@ -112,7 +115,7 @@ import Testing #expect(metadata == expected) } - @Test("Throws metadataDiscoveryFailed when all candidates fail") + @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 @@ -126,11 +129,34 @@ import Testing 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") @@ -162,23 +188,23 @@ import Testing #expect(metadata == expectedMetadata) } - @Test("Skips candidate when issuer field does not match") - func testFetchAuthorizationServerMetadataSkipsIssuerMismatch() async throws { - let wrongIssuerBody = try makeASMetadataBody(issuer: "https://other.example.com") + @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, wrongIssuerBody) + return (response, body) } - await #expect(throws: OAuthAuthorizationError.self) { - try await makeClient().fetchAuthorizationServerMetadata( - candidates: [URL(string: "https://auth.example.com")!], - session: session - ) - } + 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")