Introduce ServerSession for per-connection state
## Motivation and Context The Ruby SDK uses a 1-Server-to-many-sessions model where session-specific state (`@client`, `@logging_message_notification`) is stored on the shared `Server` instance. This causes the last client to connect to overwrite state for all sessions, and requires `session_id` to be threaded through method signatures for session-scoped notifications. The Python SDK avoids this by creating a `ServerSession` per connection that wraps a shared `Server`. Each session holds its own state and naturally scopes notifications through its own stream. ### Changes - Added `MCP::ServerSession` class holding per-session state: client info, logging configuration, and a `transport` reference for session-scoped notification delivery. - `StreamableHTTPTransport#handle_initialization` creates a `ServerSession` per `initialize` request and stores it in the session hash. - `handle_regular_request` routes requests through the session's `ServerSession` instead of the shared `Server` directly. - `Server#handle` and `Server#handle_json` accept an optional `session:` keyword (`ServerSession` instance). - `Server#init` stores client info on the session when available. - `Server#configure_logging_level` stores logging config on the session. - `ServerContext` and `Progress` accept a `notification_target:` which can be either a `ServerSession` (session-scoped) or a `Server` (broadcast). No `session_id:` parameter needed. - `ServerContext#notify_progress` and `#notify_log_message` delegate to the `notification_target` without `session_id` threading. - `StdioTransport` creates a single `ServerSession` on `open`, making the session model transparent across both transports. ### Design Decision This follows the Python SDK's "shared `Server` + per-connection `Session`" pattern. The `Server` holds configuration (tools, prompts, resources, handlers). Each `ServerSession` holds per-connection state (client info, logging level, stream writer). Notifications from `ServerSession` go only to that session's stream. ## How Has This Been Tested? All existing tests pass. All conformance tests pass. Added tests verifying: - Session-scoped log notification is sent only to the originating session. - Session-scoped progress notification is sent only to the originating session. - Each session stores its own client info independently. - Each session stores its own logging level independently. - `StdioTransport#open` creates a `ServerSession` and stores client info on it. ## Breaking Change None for end users. The public API (`Server.new`, `define_tool`,`server_context.report_progress`, `server_context.notify_log_message`, etc.) is unchanged. The following internal API changes affect only SDK internals: - `ServerContext.new` now requires `notification_target:` instead of just`progress:`. - `Progress.new` now takes `notification_target:` instead of `server:`. - `Server#handle` and `Server#handle_json` accept an optional `session:` keyword.
This commit is contained in:
Koichi ITO
parent
a9e45148f6
commit
719e5e7fee
8 changed files with 563 additions and 63 deletions
|
|
@ -15,6 +15,7 @@ module MCP
|
|||
autoload :Resource, "mcp/resource"
|
||||
autoload :ResourceTemplate, "mcp/resource_template"
|
||||
autoload :Server, "mcp/server"
|
||||
autoload :ServerSession, "mcp/server_session"
|
||||
autoload :Tool, "mcp/tool"
|
||||
|
||||
class << self
|
||||
|
|
|
|||
|
|
@ -2,15 +2,15 @@
|
|||
|
||||
module MCP
|
||||
class Progress
|
||||
def initialize(server:, progress_token:)
|
||||
@server = server
|
||||
def initialize(notification_target:, progress_token:)
|
||||
@notification_target = notification_target
|
||||
@progress_token = progress_token
|
||||
end
|
||||
|
||||
def report(progress, total: nil, message: nil)
|
||||
return unless @progress_token
|
||||
|
||||
@server.notify_progress(
|
||||
@notification_target.notify_progress(
|
||||
progress_token: @progress_token,
|
||||
progress: progress,
|
||||
total: total,
|
||||
|
|
|
|||
|
|
@ -111,15 +111,29 @@ module MCP
|
|||
@transport = transport
|
||||
end
|
||||
|
||||
def handle(request)
|
||||
# Processes a parsed JSON-RPC request and returns the response as a Hash.
|
||||
#
|
||||
# @param request [Hash] A parsed JSON-RPC request.
|
||||
# @param session [ServerSession, nil] Per-connection session. Passed by
|
||||
# `ServerSession#handle` for session-scoped notification delivery.
|
||||
# When `nil`, notifications broadcast to all sessions.
|
||||
# @return [Hash, nil] The JSON-RPC response, or `nil` for notifications.
|
||||
def handle(request, session: nil)
|
||||
JsonRpcHandler.handle(request) do |method|
|
||||
handle_request(request, method)
|
||||
handle_request(request, method, session: session)
|
||||
end
|
||||
end
|
||||
|
||||
def handle_json(request)
|
||||
# Processes a JSON-RPC request string and returns the response as a JSON string.
|
||||
#
|
||||
# @param request [String] A JSON-RPC request as a JSON string.
|
||||
# @param session [ServerSession, nil] Per-connection session. Passed by
|
||||
# `ServerSession#handle_json` for session-scoped notification delivery.
|
||||
# When `nil`, notifications broadcast to all sessions.
|
||||
# @return [String, nil] The JSON-RPC response as JSON, or `nil` for notifications.
|
||||
def handle_json(request, session: nil)
|
||||
JsonRpcHandler.handle_json(request) do |method|
|
||||
handle_request(request, method)
|
||||
handle_request(request, method, session: session)
|
||||
end
|
||||
end
|
||||
|
||||
|
|
@ -279,11 +293,12 @@ module MCP
|
|||
end
|
||||
end
|
||||
|
||||
def handle_request(request, method)
|
||||
def handle_request(request, method, session: nil)
|
||||
handler = @handlers[method]
|
||||
unless handler
|
||||
instrument_call("unsupported_method") do
|
||||
add_instrumentation_data(client: @client) if @client
|
||||
client = session&.client || @client
|
||||
add_instrumentation_data(client: client) if client
|
||||
end
|
||||
return
|
||||
end
|
||||
|
|
@ -293,6 +308,8 @@ module MCP
|
|||
->(params) {
|
||||
instrument_call(method) do
|
||||
result = case method
|
||||
when Methods::INITIALIZE
|
||||
init(params, session: session)
|
||||
when Methods::TOOLS_LIST
|
||||
{ tools: @handlers[Methods::TOOLS_LIST].call(params) }
|
||||
when Methods::PROMPTS_LIST
|
||||
|
|
@ -303,10 +320,15 @@ module MCP
|
|||
{ contents: @handlers[Methods::RESOURCES_READ].call(params) }
|
||||
when Methods::RESOURCES_TEMPLATES_LIST
|
||||
{ resourceTemplates: @handlers[Methods::RESOURCES_TEMPLATES_LIST].call(params) }
|
||||
when Methods::TOOLS_CALL
|
||||
call_tool(params, session: session)
|
||||
when Methods::LOGGING_SET_LEVEL
|
||||
configure_logging_level(params, session: session)
|
||||
else
|
||||
@handlers[method].call(params)
|
||||
end
|
||||
add_instrumentation_data(client: @client) if @client
|
||||
client = session&.client || @client
|
||||
add_instrumentation_data(client: client) if client
|
||||
|
||||
result
|
||||
rescue => e
|
||||
|
|
@ -342,8 +364,14 @@ module MCP
|
|||
}.compact
|
||||
end
|
||||
|
||||
def init(params)
|
||||
@client = params[:clientInfo] if params
|
||||
def init(params, session: nil)
|
||||
if params
|
||||
if session
|
||||
session.store_client_info(client: params[:clientInfo], capabilities: params[:capabilities])
|
||||
else
|
||||
@client = params[:clientInfo]
|
||||
end
|
||||
end
|
||||
|
||||
protocol_version = params[:protocolVersion] if params
|
||||
negotiated_version = if Configuration::SUPPORTED_STABLE_PROTOCOL_VERSIONS.include?(protocol_version)
|
||||
|
|
@ -371,7 +399,7 @@ module MCP
|
|||
}.compact
|
||||
end
|
||||
|
||||
def configure_logging_level(request)
|
||||
def configure_logging_level(request, session: nil)
|
||||
if capabilities[:logging].nil?
|
||||
raise RequestHandlerError.new("Server does not support logging", request, error_type: :internal_error)
|
||||
end
|
||||
|
|
@ -381,6 +409,7 @@ module MCP
|
|||
raise RequestHandlerError.new("Invalid log level #{request[:level]}", request, error_type: :invalid_params)
|
||||
end
|
||||
|
||||
session&.configure_logging(logging_message_notification)
|
||||
@logging_message_notification = logging_message_notification
|
||||
|
||||
{}
|
||||
|
|
@ -390,7 +419,7 @@ module MCP
|
|||
@tools.values.map(&:to_h)
|
||||
end
|
||||
|
||||
def call_tool(request)
|
||||
def call_tool(request, session: nil)
|
||||
tool_name = request[:name]
|
||||
|
||||
tool = tools[tool_name]
|
||||
|
|
@ -422,7 +451,7 @@ module MCP
|
|||
|
||||
progress_token = request.dig(:_meta, :progressToken)
|
||||
|
||||
call_tool_with_args(tool, arguments, server_context_with_meta(request), progress_token: progress_token)
|
||||
call_tool_with_args(tool, arguments, server_context_with_meta(request), progress_token: progress_token, session: session)
|
||||
rescue RequestHandlerError
|
||||
raise
|
||||
rescue => e
|
||||
|
|
@ -491,12 +520,13 @@ module MCP
|
|||
parameters.any? { |type, name| type == :keyrest || name == :server_context }
|
||||
end
|
||||
|
||||
def call_tool_with_args(tool, arguments, context, progress_token: nil)
|
||||
def call_tool_with_args(tool, arguments, context, progress_token: nil, session: nil)
|
||||
args = arguments&.transform_keys(&:to_sym) || {}
|
||||
|
||||
if accepts_server_context?(tool.method(:call))
|
||||
progress = Progress.new(server: self, progress_token: progress_token)
|
||||
server_context = ServerContext.new(context, progress: progress)
|
||||
notification_target = session || self
|
||||
progress = Progress.new(notification_target: notification_target, progress_token: progress_token)
|
||||
server_context = ServerContext.new(context, progress: progress, notification_target: notification_target)
|
||||
tool.call(**args, server_context: server_context).to_h
|
||||
else
|
||||
tool.call(**args).to_h
|
||||
|
|
|
|||
|
|
@ -10,17 +10,19 @@ module MCP
|
|||
STATUS_INTERRUPTED = Signal.list["INT"] + 128
|
||||
|
||||
def initialize(server)
|
||||
@server = server
|
||||
super(server)
|
||||
@open = false
|
||||
@session = nil
|
||||
$stdin.set_encoding("UTF-8")
|
||||
$stdout.set_encoding("UTF-8")
|
||||
super
|
||||
end
|
||||
|
||||
def open
|
||||
@open = true
|
||||
@session = ServerSession.new(server: @server, transport: self)
|
||||
while @open && (line = $stdin.gets)
|
||||
handle_json_request(line.strip)
|
||||
response = @session.handle_json(line.strip)
|
||||
send_response(response) if response
|
||||
end
|
||||
rescue Interrupt
|
||||
warn("\nExiting...")
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ module MCP
|
|||
class StreamableHTTPTransport < Transport
|
||||
def initialize(server, stateless: false)
|
||||
super(server)
|
||||
# { session_id => { stream: stream_object }
|
||||
# Maps `session_id` to `{ stream: stream_object, server_session: ServerSession }`.
|
||||
@sessions = {}
|
||||
@mutex = Mutex.new
|
||||
|
||||
|
|
@ -228,18 +228,25 @@ module MCP
|
|||
|
||||
def handle_initialization(body_string, body)
|
||||
session_id = nil
|
||||
server_session = nil
|
||||
|
||||
unless @stateless
|
||||
session_id = SecureRandom.uuid
|
||||
server_session = ServerSession.new(server: @server, transport: self, session_id: session_id)
|
||||
|
||||
@mutex.synchronize do
|
||||
@sessions[session_id] = {
|
||||
stream: nil,
|
||||
server_session: server_session,
|
||||
}
|
||||
end
|
||||
end
|
||||
|
||||
response = @server.handle_json(body_string)
|
||||
response = if server_session
|
||||
server_session.handle_json(body_string)
|
||||
else
|
||||
@server.handle_json(body_string)
|
||||
end
|
||||
|
||||
headers = {
|
||||
"Content-Type" => "application/json",
|
||||
|
|
@ -255,16 +262,24 @@ module MCP
|
|||
end
|
||||
|
||||
def handle_regular_request(body_string, session_id)
|
||||
unless @stateless
|
||||
if session_id && !session_exists?(session_id)
|
||||
return session_not_found_response
|
||||
server_session = nil
|
||||
stream = nil
|
||||
|
||||
if session_id && !@stateless
|
||||
@mutex.synchronize do
|
||||
session = @sessions[session_id]
|
||||
return session_not_found_response unless session
|
||||
|
||||
server_session = session[:server_session]
|
||||
stream = session[:stream]
|
||||
end
|
||||
end
|
||||
|
||||
response = @server.handle_json(body_string)
|
||||
|
||||
# Stream can be nil since stateless mode doesn't retain streams
|
||||
stream = get_session_stream(session_id) if session_id
|
||||
response = if server_session
|
||||
server_session.handle_json(body_string)
|
||||
else
|
||||
@server.handle_json(body_string)
|
||||
end
|
||||
|
||||
if stream
|
||||
send_response_to_stream(stream, response, session_id)
|
||||
|
|
|
|||
|
|
@ -2,15 +2,41 @@
|
|||
|
||||
module MCP
|
||||
class ServerContext
|
||||
def initialize(context, progress:)
|
||||
def initialize(context, progress:, notification_target:)
|
||||
@context = context
|
||||
@progress = progress
|
||||
@notification_target = notification_target
|
||||
end
|
||||
|
||||
# Reports progress for the current tool operation.
|
||||
# The notification is automatically scoped to the originating session.
|
||||
#
|
||||
# @param progress [Numeric] Current progress value.
|
||||
# @param total [Numeric, nil] Total expected value.
|
||||
# @param message [String, nil] Human-readable status message.
|
||||
def report_progress(progress, total: nil, message: nil)
|
||||
@progress.report(progress, total: total, message: message)
|
||||
end
|
||||
|
||||
# Sends a progress notification scoped to the originating session.
|
||||
#
|
||||
# @param progress_token [String, Integer] The token identifying the operation.
|
||||
# @param progress [Numeric] Current progress value.
|
||||
# @param total [Numeric, nil] Total expected value.
|
||||
# @param message [String, nil] Human-readable status message.
|
||||
def notify_progress(progress_token:, progress:, total: nil, message: nil)
|
||||
@notification_target.notify_progress(progress_token: progress_token, progress: progress, total: total, message: message)
|
||||
end
|
||||
|
||||
# Sends a log message notification scoped to the originating session.
|
||||
#
|
||||
# @param data [Object] The log data to send.
|
||||
# @param level [String] Log level (e.g., `"debug"`, `"info"`, `"error"`).
|
||||
# @param logger [String, nil] Logger name.
|
||||
def notify_log_message(data:, level:, logger: nil)
|
||||
@notification_target.notify_log_message(data: data, level: level, logger: logger)
|
||||
end
|
||||
|
||||
def method_missing(name, ...)
|
||||
if @context.respond_to?(name)
|
||||
@context.public_send(name, ...)
|
||||
|
|
|
|||
79
lib/mcp/server_session.rb
Normal file
79
lib/mcp/server_session.rb
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
# frozen_string_literal: true
|
||||
|
||||
require_relative "methods"
|
||||
|
||||
module MCP
|
||||
# Holds per-connection state for a single client session.
|
||||
# Created by the transport layer; delegates request handling to the shared `Server`.
|
||||
class ServerSession
|
||||
attr_reader :session_id, :client, :logging_message_notification
|
||||
|
||||
def initialize(server:, transport:, session_id: nil)
|
||||
@server = server
|
||||
@transport = transport
|
||||
@session_id = session_id
|
||||
@client = nil
|
||||
@client_capabilities = nil # TODO: Use for per-session capability validation.
|
||||
@logging_message_notification = nil
|
||||
end
|
||||
|
||||
def handle(request)
|
||||
@server.handle(request, session: self)
|
||||
end
|
||||
|
||||
def handle_json(request_json)
|
||||
@server.handle_json(request_json, session: self)
|
||||
end
|
||||
|
||||
# Called by `Server#init` during the initialization handshake.
|
||||
def store_client_info(client:, capabilities: nil)
|
||||
@client = client
|
||||
@client_capabilities = capabilities
|
||||
end
|
||||
|
||||
# Called by `Server#configure_logging_level`.
|
||||
def configure_logging(logging_message_notification)
|
||||
@logging_message_notification = logging_message_notification
|
||||
end
|
||||
|
||||
# Sends a progress notification to this session only.
|
||||
def notify_progress(progress_token:, progress:, total: nil, message: nil)
|
||||
params = {
|
||||
"progressToken" => progress_token,
|
||||
"progress" => progress,
|
||||
"total" => total,
|
||||
"message" => message,
|
||||
}.compact
|
||||
|
||||
send_to_transport(Methods::NOTIFICATIONS_PROGRESS, params)
|
||||
rescue => e
|
||||
@server.report_exception(e, notification: "progress")
|
||||
end
|
||||
|
||||
# Sends a log message notification to this session only.
|
||||
def notify_log_message(data:, level:, logger: nil)
|
||||
effective_logging = @logging_message_notification || @server.logging_message_notification
|
||||
return unless effective_logging&.should_notify?(level)
|
||||
|
||||
params = { "data" => data, "level" => level }
|
||||
params["logger"] = logger if logger
|
||||
|
||||
send_to_transport(Methods::NOTIFICATIONS_MESSAGE, params)
|
||||
rescue => e
|
||||
@server.report_exception(e, { notification: "log_message" })
|
||||
end
|
||||
|
||||
private
|
||||
|
||||
# TODO: When Ruby 2.7 support is dropped, replace with a direct call:
|
||||
# `@transport.send_notification(method, params, session_id: @session_id)` and
|
||||
# add `**` to `Transport#send_notification` and `StdioTransport#send_notification`.
|
||||
def send_to_transport(method, params)
|
||||
if @session_id
|
||||
@transport.send_notification(method, params, session_id: @session_id)
|
||||
else
|
||||
@transport.send_notification(method, params)
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
@ -36,8 +36,8 @@ module MCP
|
|||
send_response(response) if response
|
||||
end
|
||||
|
||||
# Send a notification to the client
|
||||
# Returns true if the notification was sent successfully
|
||||
# Send a notification to the client.
|
||||
# Returns true if the notification was sent successfully.
|
||||
def send_notification(method, params = nil)
|
||||
raise NotImplementedError, "Subclasses must implement send_notification"
|
||||
end
|
||||
|
|
|
|||
Loading…
Reference in a new issue