@@ -4373,6 +4373,123 @@ def string
43734373 assert_nil transport . instance_variable_get ( :@reaper_thread )
43744374 end
43754375
4376+ test "close lets a reap in progress finish closing the streams of the sessions it removed" do
4377+ transport = StreamableHTTPTransport . new ( @server , session_idle_timeout : 0.01 )
4378+ init_request = create_rack_request (
4379+ "POST" ,
4380+ "/" ,
4381+ { "CONTENT_TYPE" => "application/json" } ,
4382+ { jsonrpc : "2.0" , method : "initialize" , id : "init" , params : initialize_params } . to_json ,
4383+ )
4384+ session_id = transport . handle_request ( init_request ) [ 1 ] [ "mcp-session-id" ]
4385+
4386+ # A stream whose close blocks until released stands in for a reap caught between removing the session and closing its stream.
4387+ entered = Queue . new
4388+ release = Queue . new
4389+ closed = false
4390+ blocking_stream = Object . new
4391+ blocking_stream . define_singleton_method ( :close ) do
4392+ entered << :closing
4393+ release . pop
4394+ closed = true
4395+ end
4396+ transport . instance_variable_get ( :@sessions ) [ session_id ] [ :get_sse_stream ] = blocking_stream
4397+
4398+ sleep ( 0.02 ) # Wait for the session to expire.
4399+
4400+ # Wake the reaper ahead of its interval; it removes the expired session and blocks inside the close.
4401+ assert wake_reaper_until ( transport ) { !entered . empty? } , "the reaper did not reach the stream's close"
4402+ assert_empty transport . instance_variable_get ( :@sessions )
4403+
4404+ closer = Thread . new { transport . close }
4405+
4406+ refute closer . join ( 0.1 ) , "close returned while the reaper was still closing a stream"
4407+
4408+ release << :go
4409+
4410+ assert closer . join ( 2 ) , "close did not return once the reaper had finished"
4411+ assert closed , "the stream of the reaped session was never closed"
4412+ assert_nil transport . instance_variable_get ( :@reaper_thread )
4413+ ensure
4414+ release << :go
4415+ transport . close
4416+ end
4417+
4418+ test "close tears the transport down even when the reaper died with an exception" do
4419+ transport = StreamableHTTPTransport . new ( @server , session_idle_timeout : 3600 )
4420+ init_request = create_rack_request (
4421+ "POST" ,
4422+ "/" ,
4423+ { "CONTENT_TYPE" => "application/json" } ,
4424+ { jsonrpc : "2.0" , method : "initialize" , id : "init" , params : initialize_params } . to_json ,
4425+ )
4426+ session_id = transport . handle_request ( init_request ) [ 1 ] [ "mcp-session-id" ]
4427+ stream = StringIO . new
4428+ transport . send ( :store_stream_for_session , session_id , stream )
4429+
4430+ # A reap that raises, met by an exception reporter that raises as well, ends the reaper thread with an exception,
4431+ # which `Thread#join` would re-raise.
4432+ original_reporter = MCP . configuration . exception_reporter
4433+ MCP . configuration . exception_reporter = -> ( _error , _context ) { raise "reporter failed" }
4434+ transport . define_singleton_method ( :reap_expired_sessions ) { raise "reap failed" }
4435+ reaper_thread = transport . instance_variable_get ( :@reaper_thread )
4436+ reaper_thread . report_on_exception = false
4437+
4438+ assert wake_reaper_until ( transport ) { !reaper_thread . alive? } , "the reaper did not die"
4439+ assert_raises ( RuntimeError ) { reaper_thread . join }
4440+
4441+ transport . close
4442+
4443+ assert_nil transport . instance_variable_get ( :@reaper_thread )
4444+ assert_empty transport . instance_variable_get ( :@sessions )
4445+ assert stream . closed?
4446+ ensure
4447+ MCP . configuration . exception_reporter = original_reporter if original_reporter
4448+ transport . close
4449+ end
4450+
4451+ test "close kills a reaper still closing a stream once the join timed out" do
4452+ transport = StreamableHTTPTransport . new ( @server , session_idle_timeout : 0.01 )
4453+ init_request = create_rack_request (
4454+ "POST" ,
4455+ "/" ,
4456+ { "CONTENT_TYPE" => "application/json" } ,
4457+ { jsonrpc : "2.0" , method : "initialize" , id : "init" , params : initialize_params } . to_json ,
4458+ )
4459+ session_id = transport . handle_request ( init_request ) [ 1 ] [ "mcp-session-id" ]
4460+
4461+ # A stream whose close never returns while the test runs; the join is made to time out at once instead of after
4462+ # its interval.
4463+ entered = Queue . new
4464+ never = Queue . new
4465+ stuck_stream = Object . new
4466+ stuck_stream . define_singleton_method ( :close ) do
4467+ entered << :closing
4468+ never . pop
4469+ end
4470+ transport . instance_variable_get ( :@sessions ) [ session_id ] [ :get_sse_stream ] = stuck_stream
4471+
4472+ sleep ( 0.02 ) # Wait for the session to expire.
4473+
4474+ assert wake_reaper_until ( transport ) { !entered . empty? } , "the reaper did not reach the stream's close"
4475+
4476+ reaper_thread = transport . instance_variable_get ( :@reaper_thread )
4477+ reaper_thread . stubs ( :join ) . returns ( nil )
4478+ closer = Thread . new { transport . close }
4479+
4480+ assert closer . join ( 2 ) , "close did not return after the join timed out"
4481+ assert wait_until_thread_dead ( reaper_thread ) , "the reaper was not killed"
4482+ assert_nil transport . instance_variable_get ( :@reaper_thread )
4483+ assert_empty transport . instance_variable_get ( :@sessions )
4484+ ensure
4485+ # A failure would otherwise leave the reaper stuck in the close, and the closer with it; releasing the close
4486+ # also keeps the final `transport.close` from blocking on a stream still registered when an early step failed.
4487+ never &.close
4488+ reaper_thread &.kill
4489+ closer &.kill
4490+ transport . close
4491+ end
4492+
43764493 test "reaper thread is not started when session_idle_timeout is nil" do
43774494 transport = StreamableHTTPTransport . new ( @server , session_idle_timeout : nil )
43784495 assert_nil ( transport . instance_variable_get ( :@reaper_thread ) )
@@ -7183,6 +7300,34 @@ def wait_until_asleep(threads)
71837300
71847301 private
71857302
7303+ # Signals the reaper's wakeup under the transport's lock until the block holds or the timeout elapses; the signal
7304+ # is repeated because one sent before the reaper reached its wait would be lost. Returns the final outcome.
7305+ def wake_reaper_until ( transport , timeout : 2 )
7306+ mutex = transport . instance_variable_get ( :@mutex )
7307+ wakeup = transport . instance_variable_get ( :@reaper_wakeup )
7308+ deadline = Process . clock_gettime ( Process ::CLOCK_MONOTONIC ) + timeout
7309+
7310+ loop do
7311+ mutex . synchronize { wakeup . signal }
7312+ return true if yield
7313+ return false if Process . clock_gettime ( Process ::CLOCK_MONOTONIC ) > deadline
7314+
7315+ sleep ( 0.005 )
7316+ end
7317+ end
7318+
7319+ # Polls until `thread` is no longer alive or the timeout elapses; returns the final outcome.
7320+ def wait_until_thread_dead ( thread , timeout : 2 )
7321+ deadline = Process . clock_gettime ( Process ::CLOCK_MONOTONIC ) + timeout
7322+
7323+ loop do
7324+ return true unless thread . alive?
7325+ return false if Process . clock_gettime ( Process ::CLOCK_MONOTONIC ) > deadline
7326+
7327+ sleep ( 0.005 )
7328+ end
7329+ end
7330+
71867331 def initialize_test_session ( id : "init" )
71877332 init_request = create_rack_request (
71887333 "POST" ,
0 commit comments