Delay SSE GET connection until after session ID is established (#97)

* Delay SSE GET connection until after session ID is established

- Add a signaling mechanism to ensure the SSE streaming task waits for the initial session ID to be set before attempting the GET connection.
- Use a TaskGroup to race between session ID signal and a 10-second timeout, logging the outcome.
- Trigger the signal when a session ID is received for the first time from any response.
- Add detailed logging for all code paths to aid debugging and clarify connection timing.
- Clean up signal resources on disconnect to prevent leaks.

This addresses the initialization timing issue where the SSE GET could be attempted before the session was established, as discussed in MCP PR #91.

* Update tests for SSE timing

* Use configurable timeout for acquiring session ID

* Remove Actor protocol conformance for HTTPClientTransport

* Configure shorter session ID timeout for tests

* Add documentation comments to HTTPClientTransport

* Rename sessionIDWaitTimeout to sseInitializationTimeout

Author: mattt

Sources/MCP/Base/Transports/HTTPClientTransport.swift

@@ -9,28 +9,66 @@ import Logging
     import FoundationNetworking
 #endif
 
-public actor HTTPClientTransport: Actor, Transport {
+/// An implementation of the MCP Streamable HTTP transport protocol for clients.
+///
+/// This transport implements the [Streamable HTTP transport](https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#streamable-http)
+/// specification from the Model Context Protocol.
+///
+/// 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
+/// - Automatic reconnection for dropped SSE streams
+/// - 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.
+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?
     private let streaming: Bool
     private var streamingTask: Task<Void, Never>?
+
+    /// 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
+
     private var isConnected = false
     private let messageStream: AsyncThrowingStream<Data, Swift.Error>
     private let messageContinuation: AsyncThrowingStream<Data, Swift.Error>.Continuation
 
+    private var initialSessionIDSignalTask: Task<Void, Never>?
+    private var initialSessionIDContinuation: CheckedContinuation<Void, Never>?
+
+    /// 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)
+    ///   - logger: Optional logger instance for transport events
     public init(
         endpoint: URL,
         configuration: URLSessionConfiguration = .default,
         streaming: Bool = true,
+        sseInitializationTimeout: TimeInterval = 10,
         logger: Logger? = nil
     ) {
         self.init(
             endpoint: endpoint,
             session: URLSession(configuration: configuration),
             streaming: streaming,
+            sseInitializationTimeout: sseInitializationTimeout,
             logger: logger
         )
     }
@@ -39,11 +77,13 @@ public actor HTTPClientTransport: Actor, Transport {
         endpoint: URL,
         session: URLSession,
         streaming: Bool = false,
+        sseInitializationTimeout: TimeInterval = 10,
         logger: Logger? = nil
     ) {
         self.endpoint = endpoint
         self.session = session
         self.streaming = streaming
+        self.sseInitializationTimeout = sseInitializationTimeout
 
         // Create message stream
         var continuation: AsyncThrowingStream<Data, Swift.Error>.Continuation!
@@ -58,11 +98,37 @@ public actor HTTPClientTransport: Actor, Transport {
             )
     }
 
+    // Setup the initial session ID signal
+    private func setupInitialSessionIDSignal() {
+        self.initialSessionIDSignalTask = Task {
+            await withCheckedContinuation { continuation in
+                self.initialSessionIDContinuation = continuation
+                // This task will suspend here until continuation.resume() is called
+            }
+        }
+    }
+
+    // Trigger the initial session ID signal when a session ID is established
+    private func triggerInitialSessionIDSignal() {
+        if let continuation = self.initialSessionIDContinuation {
+            continuation.resume()
+            self.initialSessionIDContinuation = nil  // Consume the continuation
+            logger.debug("Initial session ID signal triggered for SSE task.")
+        }
+    }
+
     /// Establishes connection with the transport
+    ///
+    /// This prepares the transport for communication and sets up SSE streaming
+    /// if streaming mode is enabled. The actual HTTP connection happens with the
+    /// first message sent.
     public func connect() async throws {
         guard !isConnected else { return }
         isConnected = true
 
+        // Setup initial session ID signal
+        setupInitialSessionIDSignal()
+
         if streaming {
             // Start listening to server events
             streamingTask = Task { await startListeningForServerEvents() }
@@ -72,6 +138,9 @@ public actor HTTPClientTransport: Actor, Transport {
     }
 
     /// Disconnects from the transport
+    ///
+    /// This terminates any active connections, cancels the streaming task,
+    /// and releases any resources being used by the transport.
     public func disconnect() async {
         guard isConnected else { return }
         isConnected = false
@@ -86,10 +155,28 @@ public actor HTTPClientTransport: Actor, Transport {
         // Clean up message stream
         messageContinuation.finish()
 
+        // Cancel the initial session ID signal task if active
+        initialSessionIDSignalTask?.cancel()
+        initialSessionIDSignalTask = nil
+        // Resume the continuation if it's still pending to avoid leaks
+        initialSessionIDContinuation?.resume()
+        initialSessionIDContinuation = nil
+
         logger.info("HTTP clienttransport disconnected")
     }
 
     /// Sends data through an HTTP POST request
+    ///
+    /// This sends a JSON-RPC message to the server via HTTP POST and processes
+    /// the response according to the MCP Streamable HTTP specification. It handles:
+    ///
+    /// - Adding appropriate Accept headers for both JSON and SSE
+    /// - Including the session ID in requests if one has been established
+    /// - Processing different response types (JSON vs SSE)
+    /// - Handling HTTP error codes according to the specification
+    ///
+    /// - Parameter data: The JSON-RPC message to send
+    /// - Throws: MCPError for transport failures or server errors
     public func send(_ data: Data) async throws {
         guard isConnected else {
             throw MCPError.internalError("Transport not connected")
@@ -129,7 +216,12 @@ public actor HTTPClientTransport: Actor, Transport {
 
             // Extract session ID if present
             if let newSessionID = httpResponse.value(forHTTPHeaderField: "Mcp-Session-Id") {
+                let wasSessionIDNil = (self.sessionID == nil)
                 self.sessionID = newSessionID
+                if wasSessionIDNil {
+                    // Trigger signal on first session ID
+                    triggerInitialSessionIDSignal()
+                }
                 logger.debug("Session ID received", metadata: ["sessionID": "\(newSessionID)"])
             }
 
@@ -161,7 +253,12 @@ public actor HTTPClientTransport: Actor, Transport {
 
             // Extract session ID if present
             if let newSessionID = httpResponse.value(forHTTPHeaderField: "Mcp-Session-Id") {
+                let wasSessionIDNil = (self.sessionID == nil)
                 self.sessionID = newSessionID
+                if wasSessionIDNil {
+                    // Trigger signal on first session ID
+                    triggerInitialSessionIDSignal()
+                }
                 logger.debug("Session ID received", metadata: ["sessionID": "\(newSessionID)"])
             }
 
@@ -238,13 +335,29 @@ public actor HTTPClientTransport: Actor, Transport {
     }
 
     /// Receives data in an async sequence
+    ///
+    /// This returns an AsyncThrowingStream that emits Data objects representing
+    /// each JSON-RPC message received from the server. This includes:
+    ///
+    /// - Direct responses to client requests
+    /// - Server-initiated messages delivered via SSE streams
+    ///
+    /// - Returns: An AsyncThrowingStream of Data objects
     public func receive() -> AsyncThrowingStream<Data, Swift.Error> {
         return messageStream
     }
 
     // MARK: - SSE
 
     /// Starts listening for server events using SSE
+    ///
+    /// This establishes a long-lived HTTP connection using Server-Sent Events (SSE)
+    /// to enable server-to-client push messaging. It handles:
+    ///
+    /// - Waiting for session ID if needed
+    /// - Opening the SSE connection
+    /// - Automatic reconnection on connection drops
+    /// - Processing received events
     private func startListeningForServerEvents() async {
         #if os(Linux)
             // SSE is not fully supported on Linux
@@ -257,6 +370,63 @@ public actor HTTPClientTransport: Actor, Transport {
             // This is the original code for platforms that support SSE
             guard isConnected else { return }
 
+            // Wait for the initial session ID signal, but only if sessionID isn't already set
+            if self.sessionID == nil, let signalTask = self.initialSessionIDSignalTask {
+                logger.debug("SSE streaming task waiting for initial sessionID signal...")
+
+                // Race the signalTask against a timeout
+                let timeoutTask = Task {
+                    try? await Task.sleep(for: .seconds(self.sseInitializationTimeout))
+                    return false
+                }
+
+                let signalCompletionTask = Task {
+                    await signalTask.value
+                    return true  // Indicates signal received
+                }
+
+                // Use TaskGroup to race the two tasks
+                var signalReceived = false
+                do {
+                    signalReceived = try await withThrowingTaskGroup(of: Bool.self) { group in
+                        group.addTask {
+                            await signalCompletionTask.value
+                        }
+                        group.addTask {
+                            await timeoutTask.value
+                        }
+
+                        // Take the first result and cancel the other task
+                        if let firstResult = try await group.next() {
+                            group.cancelAll()
+                            return firstResult
+                        }
+                        return false
+                    }
+                } catch {
+                    logger.error("Error while waiting for session ID signal: \(error)")
+                }
+
+                // Clean up tasks
+                timeoutTask.cancel()
+
+                if signalReceived {
+                    logger.debug("SSE streaming task proceeding after initial sessionID signal.")
+                } else {
+                    logger.warning(
+                        "Timeout waiting for initial sessionID signal. SSE stream will proceed (sessionID might be nil)."
+                    )
+                }
+            } else if self.sessionID != nil {
+                logger.debug(
+                    "Initial sessionID already available. Proceeding with SSE streaming task immediately."
+                )
+            } else {
+                logger.info(
+                    "Proceeding with SSE connection attempt; sessionID is nil. This might be expected for stateless servers or if initialize hasn't provided one yet."
+                )
+            }
+
             // Retry loop for connection drops
             while isConnected && !Task.isCancelled {
                 do {
@@ -274,6 +444,11 @@ public actor HTTPClientTransport: Actor, Transport {
 
     #if !os(Linux)
         /// Establishes an SSE connection to the server
+        ///
+        /// This initiates a GET request to the server endpoint with appropriate
+        /// headers to establish an SSE stream according to the MCP specification.
+        ///
+        /// - Throws: MCPError for connection failures or server errors
         private func connectToEventStream() async throws {
             guard isConnected else { return }
 
@@ -309,13 +484,23 @@ public actor HTTPClientTransport: Actor, Transport {
 
             // Extract session ID if present
             if let newSessionID = httpResponse.value(forHTTPHeaderField: "Mcp-Session-Id") {
+                let wasSessionIDNil = (self.sessionID == nil)
                 self.sessionID = newSessionID
+                if wasSessionIDNil {
+                    // Trigger signal on first session ID, though this is unlikely to happen here
+                    // as GET usually follows a POST that would have already set the session ID
+                    triggerInitialSessionIDSignal()
+                }
                 logger.debug("Session ID received", metadata: ["sessionID": "\(newSessionID)"])
             }
 
             try await self.processSSE(stream)
         }
 
+        /// Processes an SSE byte stream, extracting events and delivering them
+        ///
+        /// - Parameter stream: The URLSession.AsyncBytes stream to process
+        /// - Throws: Error for stream processing failures
         private func processSSE(_ stream: URLSession.AsyncBytes) async throws {
             do {
                 for try await event in stream.events {