Support session expiry controls for StreamableHTTPTransport
## Motivation and Context The MCP specification recommends expiring session IDs to reduce session hijacking risks: https://modelcontextprotocol.io/specification/latest/basic/security_best_practices#session-hijacking Several other SDKs (Go, C#, PHP, Python) already implement session expiry controls, but the Ruby SDK did not. Sessions persisted indefinitely unless explicitly deleted or a stream error occurred, leaving abandoned sessions to accumulate in memory. This adds a `session_idle_timeout:` option to `StreamableHTTPTransport#initialize`. When set, sessions that receive no HTTP requests for the specified duration (in seconds) are automatically expired. Expired sessions return 404 on subsequent requests (GET and POST), matching the MCP specification's behavior for terminated sessions. Each request resets the idle timer, so actively used sessions are not interrupted. A background reaper thread periodically cleans up expired sessions to handle orphaned sessions that receive no further requests. The reaper only starts when `session_idle_timeout` is configured. The default is `nil` (no expiry) for backward compatibility, consistent with the Python SDK's approach. The Python SDK recommends 1800 seconds (30 minutes) for most deployments: https://github.com/modelcontextprotocol/python-sdk/pull/2022 Resolves #265. ## How Has This Been Tested? Added tests for session expiry, idle timeout reset, reaper thread lifecycle, input validation, and default behavior in `streamable_http_transport_test.rb`. All existing tests continue to pass. ## Breaking Change None. The default value of `session_idle_timeout` is `nil`, which preserves the existing behavior of sessions never expiring. The new `last_active_at` field in the internal session hash is not part of the public API. Existing code that instantiates `StreamableHTTPTransport.new(server)` or `StreamableHTTPTransport.new(server, stateless: true)` continues to work without changes.
This commit is contained in:
Koichi ITO
parent
8969edf4b9
commit
49fc501abc
1 changed files with 445 additions and 9 deletions
|
|
@ -8,18 +8,30 @@ module MCP
|
|||
class Server
|
||||
module Transports
|
||||
class StreamableHTTPTransport < Transport
|
||||
def initialize(server, stateless: false)
|
||||
def initialize(server, stateless: false, session_idle_timeout: nil)
|
||||
super(server)
|
||||
# Maps `session_id` to `{ stream: stream_object, server_session: ServerSession }`.
|
||||
# Maps `session_id` to `{ stream: stream_object, server_session: ServerSession, last_active_at: float_from_monotonic_clock }`.
|
||||
@sessions = {}
|
||||
@mutex = Mutex.new
|
||||
|
||||
@stateless = stateless
|
||||
@session_idle_timeout = session_idle_timeout
|
||||
|
||||
if @session_idle_timeout
|
||||
if @stateless
|
||||
raise ArgumentError, "session_idle_timeout is not supported in stateless mode."
|
||||
elsif @session_idle_timeout <= 0
|
||||
raise ArgumentError, "session_idle_timeout must be a positive number."
|
||||
end
|
||||
end
|
||||
|
||||
start_reaper_thread if @session_idle_timeout
|
||||
end
|
||||
|
||||
REQUIRED_POST_ACCEPT_TYPES = ["application/json", "text/event-stream"].freeze
|
||||
REQUIRED_GET_ACCEPT_TYPES = ["text/event-stream"].freeze
|
||||
STREAM_WRITE_ERRORS = [IOError, Errno::EPIPE, Errno::ECONNRESET].freeze
|
||||
SESSION_REAP_INTERVAL = 60
|
||||
|
||||
def handle_request(request)
|
||||
case request.env["REQUEST_METHOD"]
|
||||
|
|
@ -35,6 +47,9 @@ module MCP
|
|||
end
|
||||
|
||||
def close
|
||||
@reaper_thread&.kill
|
||||
@reaper_thread = nil
|
||||
|
||||
@mutex.synchronize do
|
||||
@sessions.each_key { |session_id| cleanup_session_unsafe(session_id) }
|
||||
end
|
||||
|
|
@ -56,6 +71,11 @@ module MCP
|
|||
session = @sessions[session_id]
|
||||
return false unless session && session[:stream]
|
||||
|
||||
if session_expired?(session)
|
||||
cleanup_session_unsafe(session_id)
|
||||
return false
|
||||
end
|
||||
|
||||
begin
|
||||
send_to_stream(session[:stream], notification)
|
||||
true
|
||||
|
|
@ -75,6 +95,11 @@ module MCP
|
|||
@sessions.each do |sid, session|
|
||||
next unless session[:stream]
|
||||
|
||||
if session_expired?(session)
|
||||
failed_sessions << sid
|
||||
next
|
||||
end
|
||||
|
||||
begin
|
||||
send_to_stream(session[:stream], notification)
|
||||
sent_count += 1
|
||||
|
|
@ -97,6 +122,39 @@ module MCP
|
|||
|
||||
private
|
||||
|
||||
def start_reaper_thread
|
||||
@reaper_thread = Thread.new do
|
||||
loop do
|
||||
sleep(SESSION_REAP_INTERVAL)
|
||||
reap_expired_sessions
|
||||
rescue StandardError => e
|
||||
MCP.configuration.exception_reporter.call(e, error: "Session reaper error")
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
def reap_expired_sessions
|
||||
return unless @session_idle_timeout
|
||||
|
||||
expired_streams = @mutex.synchronize do
|
||||
@sessions.each_with_object([]) do |(session_id, session), streams|
|
||||
next unless session_expired?(session)
|
||||
|
||||
streams << session[:stream] if session[:stream]
|
||||
@sessions.delete(session_id)
|
||||
end
|
||||
end
|
||||
|
||||
expired_streams.each do |stream|
|
||||
# Closing outside the mutex is safe because expired sessions are already
|
||||
# removed from `@sessions` above, so other threads will not find them
|
||||
# and will not attempt to close the same stream.
|
||||
stream.close
|
||||
rescue
|
||||
nil
|
||||
end
|
||||
end
|
||||
|
||||
def send_to_stream(stream, data)
|
||||
message = data.is_a?(String) ? data : data.to_json
|
||||
stream.write("data: #{message}\n\n")
|
||||
|
|
@ -145,7 +203,9 @@ module MCP
|
|||
session_id = extract_session_id(request)
|
||||
|
||||
return missing_session_id_response unless session_id
|
||||
return session_not_found_response unless session_exists?(session_id)
|
||||
|
||||
error_response = validate_and_touch_session(session_id)
|
||||
return error_response if error_response
|
||||
return session_already_connected_response if get_session_stream(session_id)
|
||||
|
||||
setup_sse_stream(session_id)
|
||||
|
|
@ -242,6 +302,7 @@ module MCP
|
|||
@sessions[session_id] = {
|
||||
stream: nil,
|
||||
server_session: server_session,
|
||||
last_active_at: Process.clock_gettime(Process::CLOCK_MONOTONIC),
|
||||
}
|
||||
end
|
||||
end
|
||||
|
|
@ -269,13 +330,16 @@ module MCP
|
|||
server_session = nil
|
||||
stream = nil
|
||||
|
||||
if session_id && !@stateless
|
||||
@mutex.synchronize do
|
||||
session = @sessions[session_id]
|
||||
return session_not_found_response unless session
|
||||
unless @stateless
|
||||
if session_id
|
||||
error_response = validate_and_touch_session(session_id)
|
||||
return error_response if error_response
|
||||
|
||||
server_session = session[:server_session]
|
||||
stream = session[:stream]
|
||||
@mutex.synchronize do
|
||||
session = @sessions[session_id]
|
||||
server_session = session[:server_session] if session
|
||||
stream = session[:stream] if session
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
|
|
@ -292,6 +356,22 @@ module MCP
|
|||
end
|
||||
end
|
||||
|
||||
def validate_and_touch_session(session_id)
|
||||
@mutex.synchronize do
|
||||
return session_not_found_response unless (session = @sessions[session_id])
|
||||
return unless @session_idle_timeout
|
||||
|
||||
if session_expired?(session)
|
||||
cleanup_session_unsafe(session_id)
|
||||
return session_not_found_response
|
||||
end
|
||||
|
||||
session[:last_active_at] = Process.clock_gettime(Process::CLOCK_MONOTONIC)
|
||||
end
|
||||
|
||||
nil
|
||||
end
|
||||
|
||||
def get_session_stream(session_id)
|
||||
@mutex.synchronize { @sessions[session_id]&.fetch(:stream, nil) }
|
||||
end
|
||||
|
|
@ -397,6 +477,12 @@ module MCP
|
|||
)
|
||||
raise # Re-raise to exit the keepalive loop
|
||||
end
|
||||
|
||||
def session_expired?(session)
|
||||
return false unless @session_idle_timeout
|
||||
|
||||
Process.clock_gettime(Process::CLOCK_MONOTONIC) - session[:last_active_at] > @session_idle_timeout
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
|
|||
Loading…
Reference in a new issue