Skip to content

Commit 6601fd4

Browse files
committed
Close a refused SSE stream outside the registry lock
## Motivation and Context `StreamableHTTPTransport` closes every stream outside its registry lock, so a close that blocks on the peer cannot hold every other session on `@mutex`. One path broke that rule: when a `GET` stream could not be registered, because the session had gone or another stream had taken its place, `store_stream_for_session` closed the refused stream inside the lock, and an error raised by that close escaped the Rack body. The refused stream is now closed after the lock is released and through `close_stream_safely`, like the streams the transport drops everywhere else. ## How Has This Been Tested? Three new tests in `test/mcp/server/transports/streamable_http_transport_test.rb`: two refuse a second stream for a session, one checking from inside the stream's `close` that the lock is free, the other that a `close` raising `IOError` is swallowed and the session keeps serving, and the third refuses a stream for a session that is gone, checking the lock the same way. Against the previous library, the first and the third find the lock held and the second raises. ## Breaking Changes None.
1 parent c090dd6 commit 6601fd4

2 files changed

Lines changed: 66 additions & 14 deletions

File tree

‎lib/mcp/server/transports/streamable_http_transport.rb‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2444,22 +2444,24 @@ def create_sse_body(session_id, auth_info = nil)
24442444
# The expiry of the token that authenticated the GET is kept with the stream so the keepalive loop and
24452445
# the delivery paths can close the stream once that token expires.
24462446
def store_stream_for_session(session_id, stream, auth_info = nil)
2447-
@mutex.synchronize do
2447+
stored = @mutex.synchronize do
24482448
session = @sessions[session_id]
2449-
if session && !session[:get_sse_stream]
2450-
session[:get_sse_stream] = stream
2451-
session[:get_sse_stream_expires_at] = stream_token_expiry(auth_info)
2452-
# The expiry is nil for a stream opened without a token, and callers read a falsy return as "not stored",
2453-
# so the stream itself is the return value.
2454-
stream
2455-
else
2456-
# Either session was removed, or another request already established a stream.
2457-
stream.close
2458-
# `stream.close` may return a truthy value depending on the stream class.
2459-
# Explicitly return nil to guarantee a falsy return for callers.
2460-
nil
2461-
end
2449+
next false unless session && !session[:get_sse_stream]
2450+
2451+
session[:get_sse_stream] = stream
2452+
session[:get_sse_stream_expires_at] = stream_token_expiry(auth_info)
2453+
true
24622454
end
2455+
2456+
# The expiry is nil for a stream opened without a token, and callers read a falsy return as "not stored",
2457+
# so the stream itself is the return value.
2458+
return stream if stored
2459+
2460+
# Either the session was removed, or another request already established a stream. The refused stream is
2461+
# closed outside the lock like every other stream this transport closes, so a close that blocks on the peer
2462+
# cannot hold every other session on `@mutex`.
2463+
close_stream_safely(stream)
2464+
nil
24632465
end
24642466

24652467
# The thread acts on the stream it was started for and on no other: the session may have detached that

‎test/mcp/server/transports/streamable_http_transport_test.rb‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -743,6 +743,56 @@ def string
743743
assert stream_b.closed?
744744
end
745745

746+
test "store_stream_for_session closes a refused stream outside the mutex" do
747+
session_id = initialize_test_session
748+
@transport.send(:store_stream_for_session, session_id, StringIO.new)
749+
750+
mutex = @transport.instance_variable_get(:@mutex)
751+
closed_outside_mutex = false
752+
refused = Object.new
753+
refused.define_singleton_method(:close) do
754+
if mutex.try_lock
755+
closed_outside_mutex = true
756+
mutex.unlock
757+
end
758+
end
759+
760+
assert_nil @transport.send(:store_stream_for_session, session_id, refused)
761+
assert closed_outside_mutex, "the refused stream was closed while the mutex was held"
762+
end
763+
764+
test "store_stream_for_session swallows an error from closing a refused stream" do
765+
session_id = initialize_test_session
766+
@transport.send(:store_stream_for_session, session_id, StringIO.new)
767+
768+
failing = Object.new
769+
failing.define_singleton_method(:close) { raise IOError, "already closed by the peer" }
770+
771+
assert_nil @transport.send(:store_stream_for_session, session_id, failing)
772+
assert_equal 200, @transport.handle_request(create_rack_request(
773+
"POST",
774+
"/",
775+
{ "CONTENT_TYPE" => "application/json", "HTTP_MCP_SESSION_ID" => session_id },
776+
{ jsonrpc: "2.0", method: "ping", id: "ping-1" }.to_json,
777+
))[0]
778+
end
779+
780+
test "store_stream_for_session closes the stream outside the mutex when the session is gone" do
781+
# The session was removed between the GET's checks and its body running, so there is nothing to attach to.
782+
mutex = @transport.instance_variable_get(:@mutex)
783+
closed_outside_mutex = false
784+
orphan = Object.new
785+
orphan.define_singleton_method(:close) do
786+
if mutex.try_lock
787+
closed_outside_mutex = true
788+
mutex.unlock
789+
end
790+
end
791+
792+
assert_nil @transport.send(:store_stream_for_session, "no-such-session", orphan)
793+
assert closed_outside_mutex, "the orphaned stream was closed while the mutex was held"
794+
end
795+
746796
test "handles GET request with invalid session ID" do
747797
request = create_rack_request(
748798
"GET",

0 commit comments

Comments
 (0)