Rename session[:stream] to session[:get_sse_stream]
Follow-up to #294. The generic `:stream` key does not convey that it refers to the GET SSE stream, making it harder to distinguish from `session[:post_request_streams]` introduced in #294.
This commit is contained in:
parent
40b048b8f4
commit
65069ed1cc
2 changed files with 22 additions and 22 deletions
|
|
@ -15,7 +15,7 @@ module MCP
|
|||
|
||||
def initialize(server, stateless: false, session_idle_timeout: nil)
|
||||
super(server)
|
||||
# Maps `session_id` to `{ stream: stream_object, server_session: ServerSession, last_active_at: float_from_monotonic_clock }`.
|
||||
# Maps `session_id` to `{ get_sse_stream: stream_object, server_session: ServerSession, last_active_at: float_from_monotonic_clock }`.
|
||||
@sessions = {}
|
||||
@mutex = Mutex.new
|
||||
|
||||
|
|
@ -61,7 +61,7 @@ module MCP
|
|||
end
|
||||
|
||||
removed_sessions.each do |session|
|
||||
close_stream_safely(session[:stream])
|
||||
close_stream_safely(session[:get_sse_stream])
|
||||
close_post_request_streams(session)
|
||||
end
|
||||
end
|
||||
|
|
@ -113,7 +113,7 @@ module MCP
|
|||
failed_sessions = []
|
||||
|
||||
@sessions.each do |sid, session|
|
||||
next unless (stream = session[:stream])
|
||||
next unless (stream = session[:get_sse_stream])
|
||||
|
||||
if session_expired?(session)
|
||||
failed_sessions << sid
|
||||
|
|
@ -247,7 +247,7 @@ module MCP
|
|||
end
|
||||
|
||||
removed_sessions.each do |session|
|
||||
close_stream_safely(session[:stream])
|
||||
close_stream_safely(session[:get_sse_stream])
|
||||
close_post_request_streams(session)
|
||||
end
|
||||
end
|
||||
|
|
@ -334,7 +334,7 @@ module MCP
|
|||
end
|
||||
|
||||
if session
|
||||
close_stream_safely(session[:stream])
|
||||
close_stream_safely(session[:get_sse_stream])
|
||||
close_post_request_streams(session)
|
||||
end
|
||||
end
|
||||
|
|
@ -358,7 +358,7 @@ module MCP
|
|||
def cleanup_and_collect_stream(session_id, streams_to_close)
|
||||
return unless (removed = cleanup_session_unsafe(session_id))
|
||||
|
||||
streams_to_close << removed[:stream]
|
||||
streams_to_close << removed[:get_sse_stream]
|
||||
removed[:post_request_streams]&.each_value { |stream| streams_to_close << stream }
|
||||
end
|
||||
|
||||
|
|
@ -449,7 +449,7 @@ module MCP
|
|||
|
||||
@mutex.synchronize do
|
||||
@sessions[session_id] = {
|
||||
stream: nil,
|
||||
get_sse_stream: nil,
|
||||
server_session: server_session,
|
||||
last_active_at: Process.clock_gettime(Process::CLOCK_MONOTONIC),
|
||||
}
|
||||
|
|
@ -543,7 +543,7 @@ module MCP
|
|||
if related_request_id
|
||||
session.dig(:post_request_streams, related_request_id)
|
||||
else
|
||||
session[:stream]
|
||||
session[:get_sse_stream]
|
||||
end
|
||||
end
|
||||
|
||||
|
|
@ -572,7 +572,7 @@ module MCP
|
|||
end
|
||||
|
||||
if removed
|
||||
close_stream_safely(removed[:stream])
|
||||
close_stream_safely(removed[:get_sse_stream])
|
||||
|
||||
removed[:post_request_streams]&.each_value do |stream|
|
||||
close_stream_safely(stream)
|
||||
|
|
@ -583,7 +583,7 @@ module MCP
|
|||
end
|
||||
|
||||
def get_session_stream(session_id)
|
||||
@mutex.synchronize { @sessions[session_id]&.fetch(:stream, nil) }
|
||||
@mutex.synchronize { @sessions[session_id]&.fetch(:get_sse_stream, nil) }
|
||||
end
|
||||
|
||||
def session_exists?(session_id)
|
||||
|
|
@ -626,8 +626,8 @@ module MCP
|
|||
def store_stream_for_session(session_id, stream)
|
||||
@mutex.synchronize do
|
||||
session = @sessions[session_id]
|
||||
if session && !session[:stream]
|
||||
session[:stream] = stream
|
||||
if session && !session[:get_sse_stream]
|
||||
session[:get_sse_stream] = stream
|
||||
else
|
||||
# Either session was removed, or another request already established a stream.
|
||||
stream.close
|
||||
|
|
@ -652,13 +652,13 @@ module MCP
|
|||
end
|
||||
|
||||
def session_active_with_stream?(session_id)
|
||||
@mutex.synchronize { @sessions.key?(session_id) && @sessions[session_id][:stream] }
|
||||
@mutex.synchronize { @sessions.key?(session_id) && @sessions[session_id][:get_sse_stream] }
|
||||
end
|
||||
|
||||
def send_keepalive_ping(session_id)
|
||||
@mutex.synchronize do
|
||||
if @sessions[session_id] && @sessions[session_id][:stream]
|
||||
send_ping_to_stream(@sessions[session_id][:stream])
|
||||
if @sessions[session_id] && @sessions[session_id][:get_sse_stream]
|
||||
send_ping_to_stream(@sessions[session_id][:get_sse_stream])
|
||||
end
|
||||
end
|
||||
rescue *STREAM_WRITE_ERRORS => e
|
||||
|
|
|
|||
|
|
@ -346,7 +346,7 @@ module MCP
|
|||
|
||||
# Simulate an active SSE stream by storing a stream object in the session
|
||||
mock_stream = StringIO.new
|
||||
@transport.instance_variable_get(:@sessions)[session_id][:stream] = mock_stream
|
||||
@transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream] = mock_stream
|
||||
|
||||
# Attempt a second GET request for the same session
|
||||
get_request = create_rack_request(
|
||||
|
|
@ -377,14 +377,14 @@ module MCP
|
|||
# Establish stream A
|
||||
stream_a = StringIO.new
|
||||
@transport.send(:store_stream_for_session, session_id, stream_a)
|
||||
assert_equal stream_a, @transport.instance_variable_get(:@sessions)[session_id][:stream]
|
||||
assert_equal stream_a, @transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream]
|
||||
|
||||
# Attempt to store stream B (simulating a racing request)
|
||||
stream_b = StringIO.new
|
||||
@transport.send(:store_stream_for_session, session_id, stream_b)
|
||||
|
||||
# Stream A should still be the active stream
|
||||
assert_equal stream_a, @transport.instance_variable_get(:@sessions)[session_id][:stream]
|
||||
assert_equal stream_a, @transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream]
|
||||
|
||||
# Stream B should have been closed
|
||||
assert stream_b.closed?
|
||||
|
|
@ -928,7 +928,7 @@ module MCP
|
|||
end
|
||||
end
|
||||
|
||||
@transport.instance_variable_get(:@sessions)[session_id][:stream] = mock_stream
|
||||
@transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream] = mock_stream
|
||||
|
||||
result = @transport.send_notification("test", { message: "test" }, session_id: session_id)
|
||||
|
||||
|
|
@ -973,7 +973,7 @@ module MCP
|
|||
# The broken request_stream should be removed.
|
||||
refute @transport.instance_variable_get(:@sessions)[session_id][:post_request_streams].key?(related_id)
|
||||
# GET SSE stream should still be intact.
|
||||
assert @transport.instance_variable_get(:@sessions)[session_id][:stream]
|
||||
assert @transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream]
|
||||
end
|
||||
|
||||
test "active_stream does not fall back to GET SSE when related_request_id is given but request_stream is missing" do
|
||||
|
|
@ -2378,7 +2378,7 @@ module MCP
|
|||
mutex.unlock
|
||||
end
|
||||
end
|
||||
transport.instance_variable_get(:@sessions)[session_id][:stream] = mock_stream
|
||||
transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream] = mock_stream
|
||||
|
||||
sleep(0.02) # Wait for session to expire.
|
||||
|
||||
|
|
@ -2495,7 +2495,7 @@ module MCP
|
|||
|
||||
# Attach a mock stream to the session
|
||||
stream = StringIO.new
|
||||
transport.instance_variable_get(:@sessions)[session_id][:stream] = stream
|
||||
transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream] = stream
|
||||
|
||||
# Wait for the session to exceed the idle timeout (0.01s)
|
||||
sleep(0.02)
|
||||
|
|
|
|||
Loading…
Reference in a new issue