From f5401e25b371fd1b4f1180af09149c43879f9605 Mon Sep 17 00:00:00 2001 From: EhsanAzish80 Date: Thu, 22 Oct 2026 20:34:20 +0300 Subject: [PATCH 1/3] Allow stateful HTTP sessions to recover after DELETE --- README.md | 3 +- .../StatefulHTTPServerSessionManager.swift | 164 +++++++++++++++ .../StatefulHTTPServerTransport.swift | 8 +- Sources/MCP/Server/Server.swift | 5 + Sources/MCPConformance/Server/HTTPApp.swift | 135 ++----------- Tests/MCPTests/HTTPServerTransportTests.swift | 186 ++++++++++++++++++ 6 files changed, 373 insertions(+), 128 deletions(-) create mode 100644 Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift diff --git a/README.md b/README.md index 5bf1aaaa..1c938f10 100644 --- a/README.md +++ b/README.md @@ -1343,7 +1343,8 @@ The Swift SDK provides multiple built-in transports: | [`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 | +| [`StatefulHTTPServerTransport`](/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift) | HTTP server transport for one session, with SSE streaming for server-initiated messages | All platforms with Foundation | One client session | +| [`StatefulHTTPServerSessionManager`](/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift) | Routes an HTTP endpoint to a separate server and transport per session, including after DELETE | All platforms with Foundation | Stateful HTTP endpoints serving multiple sessions | | [`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 | diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift new file mode 100644 index 00000000..6356b716 --- /dev/null +++ b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift @@ -0,0 +1,164 @@ +import Foundation +import Logging + +/// Routes Streamable HTTP requests to a separate server and transport per session. +/// Use one manager for the HTTP endpoint so a client can initialize a new session +/// after its previous session ends. +public actor StatefulHTTPServerSessionManager { + public typealias ServerFactory = @Sendable (String, StatefulHTTPServerTransport) async throws -> Server + + private struct Session { + let server: Server + let transport: StatefulHTTPServerTransport + var lastAccessedAt: Date + } + + private struct FixedSessionIDGenerator: SessionIDGenerator { + let sessionID: String + func generateSessionID() -> String { sessionID } + } + + private let serverFactory: ServerFactory + private let validationPipeline: (any HTTPRequestValidationPipeline)? + private let retryInterval: Int? + private let sessionTimeout: TimeInterval + private let logger: Logger + private var sessions: [String: Session] = [:] + private var cleanupTask: Task? + private var stopped = false + + /// `serverFactory` must return a new `Server` for each session. + public init( + sessionTimeout: TimeInterval = 3600, + validationPipeline: (any HTTPRequestValidationPipeline)? = nil, + retryInterval: Int? = nil, + logger: Logger? = nil, + serverFactory: @escaping ServerFactory + ) { + self.sessionTimeout = sessionTimeout + self.validationPipeline = validationPipeline + self.retryInterval = retryInterval + self.logger = logger ?? Logger( + label: "mcp.transport.http.server.sessions", + factory: { _ in SwiftLogNoOpLogHandler() } + ) + self.serverFactory = serverFactory + } + + public func handleRequest(_ request: HTTPRequest) async -> HTTPResponse { + guard !stopped else { + return .error(statusCode: 503, .internalError("Server has stopped")) + } + + if cleanupTask == nil { + cleanupTask = Task { [weak self] in + while !Task.isCancelled { + try? await Task.sleep(for: .seconds(60)) + guard !Task.isCancelled, let self else { return } + await self.removeExpiredSessions() + } + } + } + + await removeExpiredSessions() + guard !stopped else { + return .error(statusCode: 503, .internalError("Server has stopped")) + } + + if let sessionID = request.header(HTTPHeaderName.sessionID) { + guard var session = sessions[sessionID] else { + return .error(statusCode: 404, .invalidRequest("Session not found")) + } + + if request.method.uppercased() == "DELETE" { + sessions.removeValue(forKey: sessionID) + let response = await session.transport.handleRequest(request) + if response.statusCode == 200 { + await session.server.stop() + } else if !stopped { + sessions[sessionID] = session + } else { + await session.server.stop() + } + return response + } + + session.lastAccessedAt = Date() + sessions[sessionID] = session + return await session.transport.handleRequest(request) + } + + guard request.method.uppercased() == "POST", + let body = request.body, + let kind = JSONRPCMessageKind(data: body), + kind.isInitializeRequest + else { + return .error(statusCode: 400, .invalidRequest("Missing session ID")) + } + + let sessionID = UUID().uuidString + let transport = StatefulHTTPServerTransport( + sessionIDGenerator: FixedSessionIDGenerator(sessionID: sessionID), + validationPipeline: validationPipeline, + retryInterval: retryInterval, + logger: logger + ) + + var server: Server? + do { + let createdServer = try await serverFactory(sessionID, transport) + server = createdServer + try await createdServer.start(transport: transport) + guard !stopped else { + await createdServer.stop() + return .error(statusCode: 503, .internalError("Server has stopped")) + } + + sessions[sessionID] = Session( + server: createdServer, transport: transport, lastAccessedAt: Date() + ) + let response = await transport.handleRequest(request) + if case .error = response { + sessions.removeValue(forKey: sessionID) + await createdServer.stop() + return response + } + return response + } catch { + if let server { + await server.stop() + } else { + await transport.disconnect() + } + return .error( + statusCode: 500, + .internalError("Failed to create session: \(error.localizedDescription)") + ) + } + } + + public func stop() async { + stopped = true + cleanupTask?.cancel() + cleanupTask = nil + let activeSessions = sessions + sessions.removeAll() + for (_, session) in activeSessions { + await session.server.stop() + } + } + + private func removeExpiredSessions() async { + let now = Date() + let expiredIDs = sessions.compactMap { sessionID, session in + now.timeIntervalSince(session.lastAccessedAt) > sessionTimeout ? sessionID : nil + } + for sessionID in expiredIDs { + guard let session = sessions[sessionID], + now.timeIntervalSince(session.lastAccessedAt) > sessionTimeout + else { continue } + sessions.removeValue(forKey: sessionID) + await session.server.stop() + } + } +} diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift index c7a4d485..85020842 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift @@ -1,7 +1,7 @@ import Foundation import Logging -/// A stateful HTTP server transport that manages sessions and uses SSE for streaming responses. +/// A stateful HTTP server transport for one session, using SSE for responses. /// /// This transport implements the MCP Streamable HTTP specification with full session management: /// - Assigns a session ID during initialization (via `Mcp-Session-Id` header) @@ -13,15 +13,17 @@ import Logging /// ## Usage /// /// ```swift -/// let transport = StatefulHTTPServerTransport() // Uses UUID by default +/// let transport = StatefulHTTPServerTransport() // One transport per session /// /// // Start the MCP server with this transport /// try await server.start(transport: transport) /// -/// // In your HTTP framework handler: +/// // In your HTTP framework handler for this session: /// let response = await transport.handleRequest(httpRequest) /// // Convert response to your framework's response type and return it /// ``` +/// Use ``StatefulHTTPServerSessionManager`` for an endpoint that creates new +/// sessions after DELETE. A terminated transport cannot be reused. /// /// ## Framework Integration /// diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 55a8e567..5ff81cd5 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -285,6 +285,11 @@ public actor Server { task?.cancel() task = nil + for (_, requestTask) in pendingRequestTasks { + requestTask.cancel() + } + pendingRequestTasks.removeAll() + // Clear pending requests with errors let pendingRequestsToCancel = self.pendingRequests self.pendingRequests = [:] diff --git a/Sources/MCPConformance/Server/HTTPApp.swift b/Sources/MCPConformance/Server/HTTPApp.swift index 6edba56f..8854376c 100644 --- a/Sources/MCPConformance/Server/HTTPApp.swift +++ b/Sources/MCPConformance/Server/HTTPApp.swift @@ -43,23 +43,14 @@ actor HTTPApp { } /// Factory function to create MCP Server instances for each session. - typealias ServerFactory = @Sendable (String, StatefulHTTPServerTransport) async throws -> Server + typealias ServerFactory = StatefulHTTPServerSessionManager.ServerFactory private let configuration: Configuration - private let serverFactory: ServerFactory - private let validationPipeline: (any HTTPRequestValidationPipeline)? + private let sessionManager: StatefulHTTPServerSessionManager 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. @@ -77,12 +68,17 @@ actor HTTPApp { logger: Logger? = nil ) { self.configuration = configuration - self.serverFactory = serverFactory - self.validationPipeline = validationPipeline self.logger = logger ?? Logger( label: "mcp.http.app", factory: { _ in SwiftLogNoOpLogHandler() } ) + self.sessionManager = StatefulHTTPServerSessionManager( + sessionTimeout: configuration.sessionTimeout, + validationPipeline: validationPipeline, + retryInterval: configuration.retryInterval, + logger: self.logger, + serverFactory: serverFactory + ) } /// Convenience initializer with individual parameters. @@ -132,14 +128,12 @@ actor HTTPApp { 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() + await sessionManager.stop() try? await channel?.close() channel = nil logger.info("MCP HTTP application stopped") @@ -155,114 +149,7 @@ actor HTTPApp { /// - 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) - } - } + await sessionManager.handleRequest(request) } } diff --git a/Tests/MCPTests/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift index 8e1c89b4..589976a4 100644 --- a/Tests/MCPTests/HTTPServerTransportTests.swift +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -158,6 +158,25 @@ private func drainSSEStream( return await collector.getChunks() } +private func nextValue( + from stream: AsyncStream, + timeout: Duration = .seconds(1) +) async -> T? { + await withTaskGroup(of: T?.self) { group in + group.addTask { + var iterator = stream.makeAsyncIterator() + return await iterator.next() + } + group.addTask { + try? await Task.sleep(for: timeout) + return nil + } + let value = await group.next() ?? nil + group.cancelAll() + return value + } +} + /// 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( @@ -592,6 +611,173 @@ struct StatefulHTTPServerTransportTests { #expect(response.statusCode == 404) } + @Test("Session manager creates a new server after DELETE") + func testServerInitializeAfterDelete() async throws { + let manager = StatefulHTTPServerSessionManager( + validationPipeline: StandardValidationPipeline(validators: [SessionValidator()]) + ) { _, _ in Server(name: "test", version: "1.0") } + + func responseText(_ response: HTTPResponse) async throws -> String { + guard case .stream(let stream, _) = response else { + Issue.record("Expected an SSE response") + return "" + } + var text = "" + for try await chunk in stream { + text += String(decoding: chunk, as: UTF8.self) + } + return text + } + + let first = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody(id: "first")) + ) + let firstText = try await responseText(first) + #expect(firstText.contains("serverInfo")) + guard let firstSessionID = first.headers[HTTPHeaderName.sessionID] else { + Issue.record("Expected first session ID") + return + } + + let deleted = await manager.handleRequest(makeDELETERequest(sessionID: firstSessionID)) + #expect(deleted.statusCode == 200) + + let second = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody(id: "second")) + ) + let secondText = try await responseText(second) + #expect(secondText.contains("serverInfo")) + #expect(!secondText.contains("\"error\"")) + #expect(second.headers[HTTPHeaderName.sessionID] != firstSessionID) + + let stale = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeRequestBody(), sessionID: firstSessionID) + ) + #expect(stale.statusCode == 404) + await manager.stop() + } + + @Test("Concurrent sessions with the same request ID stay isolated") + func testConcurrentSessionsAreIsolated() async throws { + let manager = StatefulHTTPServerSessionManager( + validationPipeline: StandardValidationPipeline(validators: [SessionValidator()]) + ) { sessionID, _ in Server(name: sessionID, version: "1.0") } + + async let first = manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody(id: "shared")) + ) + async let second = manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody(id: "shared")) + ) + let responses = await [first, second] + #expect(responses[0].statusCode == 200) + #expect(responses[1].statusCode == 200) + + guard let firstID = responses[0].headers[HTTPHeaderName.sessionID], + let secondID = responses[1].headers[HTTPHeaderName.sessionID] + else { + Issue.record("Expected session IDs") + await manager.stop() + return + } + #expect(firstID != secondID) + + for (response, sessionID) in zip(responses, [firstID, secondID]) { + guard case .stream(let stream, _) = response else { + Issue.record("Expected an SSE response") + continue + } + var body = "" + for try await chunk in stream { + body += String(decoding: chunk, as: UTF8.self) + } + #expect(body.contains(sessionID)) + } + await manager.stop() + } + + @Test("Deleting a slow session cannot send its response to a new session") + func testDeletedSessionCannotLeakResponse() async throws { + actor SessionOrder { + private var count = 0 + func next() -> Int { + count += 1 + return count + } + } + let order = SessionOrder() + let (entered, enteredContinuation) = AsyncStream.makeStream() + let (cancelled, cancelledContinuation) = AsyncStream.makeStream() + + let manager = StatefulHTTPServerSessionManager( + validationPipeline: StandardValidationPipeline(validators: [SessionValidator()]) + ) { _, _ in + let number = await order.next() + let server = Server(name: "server-\(number)", version: "1.0") + await server.withMethodHandler(CallTool.self) { _ in + if number == 1 { + enteredContinuation.yield(()) + do { + try await Task.sleep(for: .seconds(5)) + } catch is CancellationError { + cancelledContinuation.yield(()) + throw CancellationError() + } + } + return CallTool.Result(content: [ + .text(text: "server-\(number)", annotations: nil, _meta: nil) + ]) + } + return server + } + + func textOf(_ response: HTTPResponse) async throws -> String { + guard case .stream(let stream, _) = response else { + Issue.record("Expected SSE response") + return "" + } + var text = "" + for try await chunk in stream { + text += String(decoding: chunk, as: UTF8.self) + } + return text + } + let callBody = try JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", "id": "shared", "method": "tools/call", + "params": ["name": "test"] as [String: Any], + ]) + + let first = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody()) + ) + _ = try await textOf(first) + let oldID = try #require(first.headers[HTTPHeaderName.sessionID]) + let oldCall = await manager.handleRequest( + makeStatefulPOSTRequest(body: callBody, sessionID: oldID) + ) + let oldText = Task { try await textOf(oldCall) } + _ = try #require(await nextValue(from: entered)) + + let deleted = await manager.handleRequest(makeDELETERequest(sessionID: oldID)) + #expect(deleted.statusCode == 200) + _ = try #require(await nextValue(from: cancelled)) + + let second = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody()) + ) + _ = try await textOf(second) + let newID = try #require(second.headers[HTTPHeaderName.sessionID]) + #expect(newID != oldID) + let newCall = await manager.handleRequest( + makeStatefulPOSTRequest(body: callBody, sessionID: newID) + ) + let newText = try await textOf(newCall) + #expect(newText.contains("server-2")) + #expect(!newText.contains("server-1")) + #expect(!(try await oldText.value).contains("server-1")) + await manager.stop() + } + // MARK: - Terminated State @Test("All methods return 404 when terminated") From ec9d42312730a4e687daa4a374d730b51fc3a0e2 Mon Sep 17 00:00:00 2001 From: EhsanAzish80 Date: Sat, 24 Oct 2026 21:03:56 +0300 Subject: [PATCH 2/3] Validate and limit stateful HTTP sessions --- .../StatefulHTTPServerSessionManager.swift | 36 ++++- Tests/MCPTests/HTTPServerTransportTests.swift | 134 ++++++++++++++++++ 2 files changed, 164 insertions(+), 6 deletions(-) diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift index 6356b716..2f5f9ab8 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift @@ -19,24 +19,35 @@ public actor StatefulHTTPServerSessionManager { } private let serverFactory: ServerFactory - private let validationPipeline: (any HTTPRequestValidationPipeline)? + private let validationPipeline: any HTTPRequestValidationPipeline private let retryInterval: Int? private let sessionTimeout: TimeInterval + private let maxSessions: Int private let logger: Logger private var sessions: [String: Session] = [:] + private var creatingSessions = 0 private var cleanupTask: Task? private var stopped = false /// `serverFactory` must return a new `Server` for each session. + /// `maxSessions` limits active sessions, including sessions still being created. public init( sessionTimeout: TimeInterval = 3600, + maxSessions: Int = 100, validationPipeline: (any HTTPRequestValidationPipeline)? = nil, retryInterval: Int? = nil, logger: Logger? = nil, serverFactory: @escaping ServerFactory ) { self.sessionTimeout = sessionTimeout - self.validationPipeline = validationPipeline + self.maxSessions = max(0, maxSessions) + 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.sessions", @@ -96,6 +107,21 @@ public actor StatefulHTTPServerSessionManager { return .error(statusCode: 400, .invalidRequest("Missing session ID")) } + let context = HTTPValidationContext( + httpMethod: "POST", + isInitializationRequest: true, + supportedProtocolVersions: Version.supported + ) + if let errorResponse = validationPipeline.validate(request, context: context) { + return errorResponse + } + + guard sessions.count + creatingSessions < maxSessions else { + return .error(statusCode: 503, .internalError("Session limit reached")) + } + creatingSessions += 1 + defer { creatingSessions -= 1 } + let sessionID = UUID().uuidString let transport = StatefulHTTPServerTransport( sessionIDGenerator: FixedSessionIDGenerator(sessionID: sessionID), @@ -130,10 +156,8 @@ public actor StatefulHTTPServerSessionManager { } else { await transport.disconnect() } - return .error( - statusCode: 500, - .internalError("Failed to create session: \(error.localizedDescription)") - ) + logger.error("Failed to create session: \(error)") + return .error(statusCode: 500, .internalError("Failed to create session")) } } diff --git a/Tests/MCPTests/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift index 589976a4..a5cee297 100644 --- a/Tests/MCPTests/HTTPServerTransportTests.swift +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -657,6 +657,140 @@ struct StatefulHTTPServerTransportTests { await manager.stop() } + @Test("Session manager validates initialize before creating a server") + func testSessionManagerValidatesBeforeFactory() async throws { + actor FactoryCalls { + var count = 0 + func record() { count += 1 } + } + let calls = FactoryCalls() + let bearer = BearerTokenValidator( + resourceMetadataURL: authResourceMetadataURL, + resourceIdentifier: authResourceIdentifier, + tokenValidator: { _, _, _ in .valid(BearerTokenInfo()) } + ) + let manager = StatefulHTTPServerSessionManager( + validationPipeline: StandardValidationPipeline(validators: [bearer, SessionValidator()]) + ) { _, _ in + await calls.record() + return Server(name: "test", version: "1.0") + } + + let unauthorized = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody()) + ) + #expect(unauthorized.statusCode == 401) + #expect(await calls.count == 0) + + let authorized = await manager.handleRequest( + makeStatefulPOSTRequest( + body: makeInitializeBody(), authorization: "Bearer valid-token" + ) + ) + #expect(authorized.statusCode == 200) + #expect(await calls.count == 1) + await manager.stop() + } + + @Test("Session manager limits active sessions and frees a slot after DELETE") + func testSessionManagerLimitsSessions() async throws { + actor FactoryCalls { + var count = 0 + func record() { count += 1 } + } + let calls = FactoryCalls() + let manager = StatefulHTTPServerSessionManager( + maxSessions: 1, + validationPipeline: StandardValidationPipeline(validators: [SessionValidator()]) + ) { _, _ in + await calls.record() + return Server(name: "test", version: "1.0") + } + + let first = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody(id: "first")) + ) + #expect(first.statusCode == 200) + let firstID = try #require(first.headers[HTTPHeaderName.sessionID]) + + let rejected = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody(id: "second")) + ) + #expect(rejected.statusCode == 503) + #expect(await calls.count == 1) + + let deleted = await manager.handleRequest(makeDELETERequest(sessionID: firstID)) + #expect(deleted.statusCode == 200) + let next = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody(id: "third")) + ) + #expect(next.statusCode == 200) + #expect(next.headers[HTTPHeaderName.sessionID] != firstID) + #expect(await calls.count == 2) + await manager.stop() + } + + @Test("Session manager counts sessions while their servers are being created") + func testSessionManagerLimitsPendingSessions() async throws { + actor FactoryGate { + var didStart = false + var started: CheckedContinuation? + var release: CheckedContinuation? + + func waitForFactory() async { + if didStart { return } + await withCheckedContinuation { started = $0 } + } + + func pauseFactory() async { + didStart = true + started?.resume() + await withCheckedContinuation { release = $0 } + } + + func resumeFactory() { release?.resume() } + } + let gate = FactoryGate() + let manager = StatefulHTTPServerSessionManager( + maxSessions: 1, + validationPipeline: StandardValidationPipeline(validators: [SessionValidator()]) + ) { _, _ in + await gate.pauseFactory() + return Server(name: "test", version: "1.0") + } + + let first = Task { + await manager.handleRequest(makeStatefulPOSTRequest(body: makeInitializeBody())) + } + await gate.waitForFactory() + let rejected = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody(id: "second")) + ) + #expect(rejected.statusCode == 503) + await gate.resumeFactory() + #expect(await first.value.statusCode == 200) + await manager.stop() + } + + @Test("Session manager does not expose a server factory error") + func testSessionManagerFactoryError() async { + struct FactoryError: Error, LocalizedError { + var errorDescription: String? { "private factory details" } + } + let manager = StatefulHTTPServerSessionManager( + validationPipeline: StandardValidationPipeline(validators: [SessionValidator()]) + ) { _, _ in + throw FactoryError() + } + let response = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody()) + ) + #expect(response.statusCode == 500) + let body = String(decoding: response.bodyData ?? Data(), as: UTF8.self) + #expect(!body.contains("private factory details")) + await manager.stop() + } + @Test("Concurrent sessions with the same request ID stay isolated") func testConcurrentSessionsAreIsolated() async throws { let manager = StatefulHTTPServerSessionManager( From bf10e4d38d1c3255471697798427bcc516c69e30 Mon Sep 17 00:00:00 2001 From: EhsanAzish80 Date: Sun, 25 Oct 2026 16:46:37 +0300 Subject: [PATCH 3/3] Validate managed HTTP initialization once --- .../StatefulHTTPServerSessionManager.swift | 13 ++++--- .../StatefulHTTPServerTransport.swift | 17 ++++++++- Tests/MCPTests/HTTPServerTransportTests.swift | 38 ++++++++++++++++++- 3 files changed, 59 insertions(+), 9 deletions(-) diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift index 2f5f9ab8..907476f3 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerSessionManager.swift @@ -22,7 +22,7 @@ public actor StatefulHTTPServerSessionManager { private let validationPipeline: any HTTPRequestValidationPipeline private let retryInterval: Int? private let sessionTimeout: TimeInterval - private let maxSessions: Int + private let maxSessions: Int? private let logger: Logger private var sessions: [String: Session] = [:] private var creatingSessions = 0 @@ -30,17 +30,18 @@ public actor StatefulHTTPServerSessionManager { private var stopped = false /// `serverFactory` must return a new `Server` for each session. - /// `maxSessions` limits active sessions, including sessions still being created. + /// Set `maxSessions` to limit active sessions, including sessions still being created. + /// The default is `nil`, so deployments can choose a limit for their capacity. public init( sessionTimeout: TimeInterval = 3600, - maxSessions: Int = 100, + maxSessions: Int? = nil, validationPipeline: (any HTTPRequestValidationPipeline)? = nil, retryInterval: Int? = nil, logger: Logger? = nil, serverFactory: @escaping ServerFactory ) { self.sessionTimeout = sessionTimeout - self.maxSessions = max(0, maxSessions) + self.maxSessions = maxSessions.map { max(0, $0) } self.validationPipeline = validationPipeline ?? StandardValidationPipeline(validators: [ OriginValidator.localhost(), AcceptHeaderValidator(mode: .sseRequired), @@ -116,7 +117,7 @@ public actor StatefulHTTPServerSessionManager { return errorResponse } - guard sessions.count + creatingSessions < maxSessions else { + if let maxSessions, sessions.count + creatingSessions >= maxSessions { return .error(statusCode: 503, .internalError("Session limit reached")) } creatingSessions += 1 @@ -143,7 +144,7 @@ public actor StatefulHTTPServerSessionManager { sessions[sessionID] = Session( server: createdServer, transport: transport, lastAccessedAt: Date() ) - let response = await transport.handleRequest(request) + let response = await transport.handleValidatedInitializationRequest(request) if case .error = response { sessions.removeValue(forKey: sessionID) await createdServer.stop() diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift index 85020842..f4f4b3a0 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift @@ -195,9 +195,20 @@ public actor StatefulHTTPServerTransport: Transport, HTTPContextProviding { } } + func handleValidatedInitializationRequest(_ request: HTTPRequest) -> HTTPResponse { + if terminated { + return .error( + statusCode: 404, + .invalidRequest("Not Found: Session has been terminated"), + sessionID: sessionID + ) + } + return handlePost(request, skipValidation: true) + } + // MARK: - POST Handler - private func handlePost(_ request: HTTPRequest) -> HTTPResponse { + private func handlePost(_ request: HTTPRequest, skipValidation: Bool = false) -> HTTPResponse { // Parse body first so we can determine if it's an initialization request guard let body = request.body, !body.isEmpty else { return .error( @@ -224,7 +235,9 @@ public actor StatefulHTTPServerTransport: Transport, HTTPContextProviding { ) // Run validation pipeline - if let errorResponse = validationPipeline.validate(request, context: context) { + if !skipValidation, + let errorResponse = validationPipeline.validate(request, context: context) + { return errorResponse } diff --git a/Tests/MCPTests/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift index a5cee297..b4c48a32 100644 --- a/Tests/MCPTests/HTTPServerTransportTests.swift +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -659,15 +659,35 @@ struct StatefulHTTPServerTransportTests { @Test("Session manager validates initialize before creating a server") func testSessionManagerValidatesBeforeFactory() async throws { + final class ValidationCalls: @unchecked Sendable { + private let lock = NSLock() + private var value = 0 + + func record() { + lock.lock() + value += 1 + lock.unlock() + } + + var count: Int { + lock.lock() + defer { lock.unlock() } + return value + } + } actor FactoryCalls { var count = 0 func record() { count += 1 } } let calls = FactoryCalls() + let validationCalls = ValidationCalls() let bearer = BearerTokenValidator( resourceMetadataURL: authResourceMetadataURL, resourceIdentifier: authResourceIdentifier, - tokenValidator: { _, _, _ in .valid(BearerTokenInfo()) } + tokenValidator: { _, _, _ in + validationCalls.record() + return .valid(BearerTokenInfo()) + } ) let manager = StatefulHTTPServerSessionManager( validationPipeline: StandardValidationPipeline(validators: [bearer, SessionValidator()]) @@ -689,6 +709,22 @@ struct StatefulHTTPServerTransportTests { ) #expect(authorized.statusCode == 200) #expect(await calls.count == 1) + #expect(validationCalls.count == 1) + await manager.stop() + } + + @Test("Session manager has no default session cap") + func testSessionManagerDefaultLimit() async { + let manager = StatefulHTTPServerSessionManager( + validationPipeline: StandardValidationPipeline(validators: [SessionValidator()]) + ) { _, _ in Server(name: "test", version: "1.0") } + + for index in 0..<101 { + let response = await manager.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody(id: "\(index)")) + ) + #expect(response.statusCode == 200) + } await manager.stop() }