Let the session reaper finish before the transport closes

lib/mcp/server/transports/streamable_http_transport.rb

@@ -347,6 +347,10 @@ def initialize(
 
           @max_stream_lifetime = max_stream_lifetime
 
+          # `close` sets the flag and signals the reaper under `@mutex`, so the reaper leaves its wait at once instead of
+          # being killed in the middle of a reap.
+          @closing = false
+          @reaper_wakeup = ConditionVariable.new
           start_reaper_thread if @session_idle_timeout
         end
 
@@ -355,6 +359,8 @@ def initialize(
         REQUIRED_GET_ACCEPT_TYPES = ["text/event-stream"].freeze
         STREAM_WRITE_ERRORS = [IOError, Errno::EPIPE, Errno::ECONNRESET].freeze
         SESSION_REAP_INTERVAL = 60
+        # Seconds `close` waits for a reap in progress to finish closing its streams before it kills the reaper.
+        REAPER_JOIN_TIMEOUT = 5
 
         # Loopback hosts always accepted by DNS rebinding protection. A locally bound MCP server (the canonical pattern) is
         # protected out of the box; non-loopback deployments widen the list via `allowed_hosts:`.
@@ -453,8 +459,28 @@ def handle_request(request)
           end
         end
 
+        # The reaper is asked to stop and joined before it is killed: `reap_expired_sessions` removes expired sessions
+        # under `@mutex` and closes their streams after releasing it, and a kill landing between the two would leave
+        # those streams open with no session left to find them by. Waking the reaper ends its wait at once, and the join
+        # lets a reap already past its removal finish its closes before the rest of the transport is torn down.
+        # The join is bounded, so a stream whose close never returns cannot hold the shutdown, and a reaper that died
+        # with a `StandardError` does not stop the teardown below: `join` re-raises the exception a thread died with,
+        # which is not this method's to raise. A reaper still alive after the wait is killed as before, leaving
+        # the streams of that reap it had not closed yet as they were.
         def close
-          @reaper_thread&.kill
+          reaper = @mutex.synchronize do
+            @closing = true
+            @reaper_wakeup.signal
+            @reaper_thread
+          end
+
+          begin
+            reaper&.join(REAPER_JOIN_TIMEOUT)
+          rescue StandardError
+            # The reaper is gone either way; the exception that ended it is not this method's to report.
+          end
+
+          reaper.kill if reaper&.alive?
           @reaper_thread = nil
 
           teardown_listen_subscriptions
@@ -772,10 +798,18 @@ def cancel_pending_request(request_id, reason: nil)
 
         private
 
+        # The wait between reaps is a condition variable wait under `@mutex` rather than a sleep, so `close` can end it
+        # by signalling; a spurious wakeup only runs a reap early.
         def start_reaper_thread
           @reaper_thread = Thread.new do
             loop do
-              sleep(SESSION_REAP_INTERVAL)
+              closing = @mutex.synchronize do
+                @reaper_wakeup.wait(@mutex, SESSION_REAP_INTERVAL) unless @closing
+                @closing
+              end
+
+              break if closing
+
               reap_expired_sessions
             rescue StandardError => e
               MCP.configuration.exception_reporter.call(e, error: "Session reaper error")

test/mcp/server/transports/streamable_http_transport_test.rb

@@ -4323,6 +4323,123 @@ def string
           assert_nil transport.instance_variable_get(:@reaper_thread)
         end
 
+        test "close lets a reap in progress finish closing the streams of the sessions it removed" do
+          transport = StreamableHTTPTransport.new(@server, session_idle_timeout: 0.01)
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "init", params: initialize_params }.to_json,
+          )
+          session_id = transport.handle_request(init_request)[1]["mcp-session-id"]
+
+          # A stream whose close blocks until released stands in for a reap caught between removing the session and closing its stream.
+          entered = Queue.new
+          release = Queue.new
+          closed = false
+          blocking_stream = Object.new
+          blocking_stream.define_singleton_method(:close) do
+            entered << :closing
+            release.pop
+            closed = true
+          end
+          transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream] = blocking_stream
+
+          sleep(0.02) # Wait for the session to expire.
+
+          # Wake the reaper ahead of its interval; it removes the expired session and blocks inside the close.
+          assert wake_reaper_until(transport) { !entered.empty? }, "the reaper did not reach the stream's close"
+          assert_empty transport.instance_variable_get(:@sessions)
+
+          closer = Thread.new { transport.close }
+
+          refute closer.join(0.1), "close returned while the reaper was still closing a stream"
+
+          release << :go
+
+          assert closer.join(2), "close did not return once the reaper had finished"
+          assert closed, "the stream of the reaped session was never closed"
+          assert_nil transport.instance_variable_get(:@reaper_thread)
+        ensure
+          release << :go
+          transport.close
+        end
+
+        test "close tears the transport down even when the reaper died with an exception" do
+          transport = StreamableHTTPTransport.new(@server, session_idle_timeout: 3600)
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "init", params: initialize_params }.to_json,
+          )
+          session_id = transport.handle_request(init_request)[1]["mcp-session-id"]
+          stream = StringIO.new
+          transport.send(:store_stream_for_session, session_id, stream)
+
+          # A reap that raises, met by an exception reporter that raises as well, ends the reaper thread with an exception,
+          # which `Thread#join` would re-raise.
+          original_reporter = MCP.configuration.exception_reporter
+          MCP.configuration.exception_reporter = ->(_error, _context) { raise "reporter failed" }
+          transport.define_singleton_method(:reap_expired_sessions) { raise "reap failed" }
+          reaper_thread = transport.instance_variable_get(:@reaper_thread)
+          reaper_thread.report_on_exception = false
+
+          assert wake_reaper_until(transport) { !reaper_thread.alive? }, "the reaper did not die"
+          assert_raises(RuntimeError) { reaper_thread.join }
+
+          transport.close
+
+          assert_nil transport.instance_variable_get(:@reaper_thread)
+          assert_empty transport.instance_variable_get(:@sessions)
+          assert stream.closed?
+        ensure
+          MCP.configuration.exception_reporter = original_reporter if original_reporter
+          transport.close
+        end
+
+        test "close kills a reaper still closing a stream once the join timed out" do
+          transport = StreamableHTTPTransport.new(@server, session_idle_timeout: 0.01)
+          init_request = create_rack_request(
+            "POST",
+            "/",
+            { "CONTENT_TYPE" => "application/json" },
+            { jsonrpc: "2.0", method: "initialize", id: "init", params: initialize_params }.to_json,
+          )
+          session_id = transport.handle_request(init_request)[1]["mcp-session-id"]
+
+          # A stream whose close never returns while the test runs; the join is made to time out at once instead of after
+          # its interval.
+          entered = Queue.new
+          never = Queue.new
+          stuck_stream = Object.new
+          stuck_stream.define_singleton_method(:close) do
+            entered << :closing
+            never.pop
+          end
+          transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream] = stuck_stream
+
+          sleep(0.02) # Wait for the session to expire.
+
+          assert wake_reaper_until(transport) { !entered.empty? }, "the reaper did not reach the stream's close"
+
+          reaper_thread = transport.instance_variable_get(:@reaper_thread)
+          reaper_thread.stubs(:join).returns(nil)
+          closer = Thread.new { transport.close }
+
+          assert closer.join(2), "close did not return after the join timed out"
+          assert wait_until_thread_dead(reaper_thread), "the reaper was not killed"
+          assert_nil transport.instance_variable_get(:@reaper_thread)
+          assert_empty transport.instance_variable_get(:@sessions)
+        ensure
+          # A failure would otherwise leave the reaper stuck in the close, and the closer with it; releasing the close
+          # also keeps the final `transport.close` from blocking on a stream still registered when an early step failed.
+          never&.close
+          reaper_thread&.kill
+          closer&.kill
+          transport.close
+        end
+
         test "reaper thread is not started when session_idle_timeout is nil" do
           transport = StreamableHTTPTransport.new(@server, session_idle_timeout: nil)
           assert_nil(transport.instance_variable_get(:@reaper_thread))
@@ -7133,6 +7250,34 @@ def wait_until_asleep(threads)
 
         private
 
+        # Signals the reaper's wakeup under the transport's lock until the block holds or the timeout elapses; the signal
+        # is repeated because one sent before the reaper reached its wait would be lost. Returns the final outcome.
+        def wake_reaper_until(transport, timeout: 2)
+          mutex = transport.instance_variable_get(:@mutex)
+          wakeup = transport.instance_variable_get(:@reaper_wakeup)
+          deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout
+
+          loop do
+            mutex.synchronize { wakeup.signal }
+            return true if yield
+            return false if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
+
+            sleep(0.005)
+          end
+        end
+
+        # Polls until `thread` is no longer alive or the timeout elapses; returns the final outcome.
+        def wait_until_thread_dead(thread, timeout: 2)
+          deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout
+
+          loop do
+            return true unless thread.alive?
+            return false if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
+
+            sleep(0.005)
+          end
+        end
+
         def initialize_test_session(id: "init")
           init_request = create_rack_request(
             "POST",