From e4e56122498fbd6bec1393e52bfa330513bca913 Mon Sep 17 00:00:00 2001 From: Koichi ITO Date: Fri, 9 Oct 2026 12:27:20 +0900 Subject: [PATCH] Let the session reaper finish before the transport closes ## Motivation and Context `StreamableHTTPTransport#close` stopped the session reaper with `Thread#kill`. The reaper removes expired sessions under the transport's lock and closes their streams after releasing it, so a kill landing between the two left those streams open with no session left to find them by: nothing closed them afterwards, since `close` tears down only the sessions still registered. The reaper now waits between reaps on a condition variable under the lock instead of sleeping, and `close` sets a flag, signals it, and joins the thread for up to `REAPER_JOIN_TIMEOUT` seconds before killing it. A reaper waiting leaves its wait at once; one past its removal finishes closing the streams it collected before `close` goes on to the rest of the transport. The join is bounded so a stream whose close never returns cannot hold the shutdown, and its outcome never decides whether the teardown runs: `Thread#join` re-raises the exception a thread died with, which `close` swallows, since the reaper reported what it could when it happened and the sessions still have to go. A spurious wakeup only runs a reap early. ## How Has This Been Tested? Three new tests in `test/mcp/server/transports/streamable_http_transport_test.rb`: one wakes the reaper while a reaped session's stream blocks inside its close, and checks that `close` returns only once that close completed and that the stream was closed; one ends the reaper with an exception, through a reap that raises and an exception reporter that raises as well, and checks that `close` still tears the sessions and their streams down; one makes the join time out at once while the reaper is inside a close that never returns, and checks that `close` returns and kills it. Against the previous library, with the reaper's sleep stubbed out so that it reaps at once, `close` returned while the stream was still open and the stream was never closed; the second and the third are what a plain join would have failed, with `close` re-raising the reaper's exception before any teardown, and never returning. ## Breaking Changes None. `close` now waits up to five seconds for a reap in progress to finish closing its streams. --- .../transports/streamable_http_transport.rb | 38 ++++- .../streamable_http_transport_test.rb | 145 ++++++++++++++++++ 2 files changed, 181 insertions(+), 2 deletions(-) diff --git a/lib/mcp/server/transports/streamable_http_transport.rb b/lib/mcp/server/transports/streamable_http_transport.rb index 7d02841f..a50591c9 100644 --- a/lib/mcp/server/transports/streamable_http_transport.rb +++ b/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") diff --git a/test/mcp/server/transports/streamable_http_transport_test.rb b/test/mcp/server/transports/streamable_http_transport_test.rb index 30be27fe..c8265760 100644 --- a/test/mcp/server/transports/streamable_http_transport_test.rb +++ b/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",