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:
parent
a9e45148f6
commit
719e5e7fee
6 changed files with 563 additions and 63 deletions
|
|
@ -30,14 +30,14 @@ module MCP
|
|||
end
|
||||
|
||||
test "#report is a no-op when progress_token is nil" do
|
||||
progress = Progress.new(server: @server, progress_token: nil)
|
||||
progress = Progress.new(notification_target: @server, progress_token: nil)
|
||||
progress.report(50, total: 100, message: "halfway")
|
||||
|
||||
assert_equal 0, @transport.notifications.size
|
||||
end
|
||||
|
||||
test "#report sends notification when progress_token is present" do
|
||||
progress = Progress.new(server: @server, progress_token: "token-1")
|
||||
progress = Progress.new(notification_target: @server, progress_token: "token-1")
|
||||
progress.report(50, total: 100, message: "halfway")
|
||||
|
||||
assert_equal 1, @transport.notifications.size
|
||||
|
|
@ -50,7 +50,7 @@ module MCP
|
|||
end
|
||||
|
||||
test "#report omits total and message when not provided" do
|
||||
progress = Progress.new(server: @server, progress_token: "token-1")
|
||||
progress = Progress.new(notification_target: @server, progress_token: "token-1")
|
||||
progress.report(50)
|
||||
|
||||
assert_equal 1, @transport.notifications.size
|
||||
|
|
|
|||
|
|
@ -101,6 +101,38 @@ module MCP
|
|||
end
|
||||
end
|
||||
|
||||
test "open creates a ServerSession and processes requests through it" do
|
||||
request = {
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
id: "1",
|
||||
params: {
|
||||
protocolVersion: "2025-11-25",
|
||||
clientInfo: { name: "stdio-client", version: "1.0" },
|
||||
},
|
||||
}
|
||||
input = StringIO.new(JSON.generate(request) + "\n")
|
||||
output = StringIO.new
|
||||
original_stdin = $stdin
|
||||
original_stdout = $stdout
|
||||
|
||||
begin
|
||||
$stdin = input
|
||||
$stdout = output
|
||||
@transport.open
|
||||
|
||||
# Verify a session was created.
|
||||
session = @transport.instance_variable_get(:@session)
|
||||
assert_instance_of(ServerSession, session)
|
||||
|
||||
# Verify client info was stored on the session, not on the server.
|
||||
assert_equal({ name: "stdio-client", version: "1.0" }, session.client)
|
||||
ensure
|
||||
$stdin = original_stdin
|
||||
$stdout = original_stdout
|
||||
end
|
||||
end
|
||||
|
||||
test "handles invalid JSON requests" do
|
||||
invalid_json = "invalid json"
|
||||
output = StringIO.new
|
||||
|
|
|
|||
|
|
@ -586,8 +586,8 @@ module MCP
|
|||
# Monkey-patch handle_json on the server to send a notification when called
|
||||
original_handle_json = @server.method(:handle_json)
|
||||
transport = @transport # Capture the transport in a local variable
|
||||
@server.define_singleton_method(:handle_json) do |request|
|
||||
result = original_handle_json.call(request)
|
||||
@server.define_singleton_method(:handle_json) do |request, **kwargs|
|
||||
result = original_handle_json.call(request, **kwargs)
|
||||
# Send notification while still in request context - broadcast to all sessions
|
||||
transport.send_notification("test_notification", { session: "current" }, **{})
|
||||
result
|
||||
|
|
@ -1357,29 +1357,289 @@ module MCP
|
|||
assert_equal([], response[2])
|
||||
end
|
||||
|
||||
test "handle_regular_request checks session existence under mutex" do
|
||||
init_request = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{ jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
|
||||
)
|
||||
init_response = @transport.handle_request(init_request)
|
||||
session_id = init_response[1]["Mcp-Session-Id"]
|
||||
|
||||
@transport.expects(:session_exists?).with(session_id).returns(true)
|
||||
|
||||
test "handle_regular_request returns 404 for unknown session_id" do
|
||||
request = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{
|
||||
"CONTENT_TYPE" => "application/json",
|
||||
"HTTP_MCP_SESSION_ID" => session_id,
|
||||
"HTTP_MCP_SESSION_ID" => "nonexistent-session",
|
||||
},
|
||||
{ jsonrpc: "2.0", method: "ping", id: "456" }.to_json,
|
||||
)
|
||||
response = @transport.handle_request(request)
|
||||
assert_equal(200, response[0])
|
||||
assert_equal(404, response[0])
|
||||
body = JSON.parse(response[2][0])
|
||||
assert_equal("Session not found", body["error"])
|
||||
end
|
||||
|
||||
test "session-scoped log notification is sent only to the originating session" do
|
||||
server = Server.new(name: "test", tools: [], prompts: [], resources: [])
|
||||
server.logging_message_notification = MCP::LoggingMessageNotification.new(level: "debug")
|
||||
transport = StreamableHTTPTransport.new(server)
|
||||
server.transport = transport
|
||||
|
||||
server.define_tool(name: "log_tool") do |server_context:|
|
||||
server_context.notify_log_message(data: "secret", level: "info")
|
||||
Tool::Response.new([{ type: "text", text: "ok" }])
|
||||
end
|
||||
server.server_context = server
|
||||
|
||||
# Create two sessions.
|
||||
init1 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
id: "1",
|
||||
params: { protocolVersion: "2025-11-25", clientInfo: { name: "a" } },
|
||||
}.to_json,
|
||||
)
|
||||
session1 = transport.handle_request(init1)[1]["Mcp-Session-Id"]
|
||||
|
||||
init2 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
id: "2",
|
||||
params: { protocolVersion: "2025-11-25", clientInfo: { name: "b" } },
|
||||
}.to_json,
|
||||
)
|
||||
session2 = transport.handle_request(init2)[1]["Mcp-Session-Id"]
|
||||
|
||||
# Connect SSE for both sessions.
|
||||
io1 = StringIO.new
|
||||
get1 = create_rack_request("GET", "/", { "HTTP_MCP_SESSION_ID" => session1 })
|
||||
response1 = transport.handle_request(get1)
|
||||
response1[2].call(io1) if response1[2].is_a?(Proc)
|
||||
|
||||
io2 = StringIO.new
|
||||
get2 = create_rack_request("GET", "/", { "HTTP_MCP_SESSION_ID" => session2 })
|
||||
response2 = transport.handle_request(get2)
|
||||
response2[2].call(io2) if response2[2].is_a?(Proc)
|
||||
|
||||
sleep(0.1)
|
||||
|
||||
# Call tool from session 1.
|
||||
tool_request = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json", "HTTP_MCP_SESSION_ID" => session1 },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "tools/call",
|
||||
id: "call-1",
|
||||
params: { name: "log_tool", arguments: {} },
|
||||
}.to_json,
|
||||
)
|
||||
transport.handle_request(tool_request)
|
||||
|
||||
# Session 1 should receive the log notification.
|
||||
io1.rewind
|
||||
assert_includes io1.read, "secret"
|
||||
|
||||
# Session 2 should NOT receive the log notification.
|
||||
io2.rewind
|
||||
refute_includes io2.read, "secret"
|
||||
end
|
||||
|
||||
test "session-scoped progress notification is sent only to the originating session" do
|
||||
server = Server.new(name: "test", tools: [], prompts: [], resources: [])
|
||||
transport = StreamableHTTPTransport.new(server)
|
||||
server.transport = transport
|
||||
|
||||
server.define_tool(name: "progress_tool") do |server_context:|
|
||||
server_context.report_progress(50, total: 100, message: "halfway")
|
||||
Tool::Response.new([{ type: "text", text: "ok" }])
|
||||
end
|
||||
server.server_context = server
|
||||
|
||||
# Create two sessions.
|
||||
init1 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
id: "1",
|
||||
params: { protocolVersion: "2025-11-25", clientInfo: { name: "a" } },
|
||||
}.to_json,
|
||||
)
|
||||
session1 = transport.handle_request(init1)[1]["Mcp-Session-Id"]
|
||||
|
||||
init2 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
id: "2",
|
||||
params: { protocolVersion: "2025-11-25", clientInfo: { name: "b" } },
|
||||
}.to_json,
|
||||
)
|
||||
session2 = transport.handle_request(init2)[1]["Mcp-Session-Id"]
|
||||
|
||||
# Connect SSE for both sessions.
|
||||
io1 = StringIO.new
|
||||
get1 = create_rack_request("GET", "/", { "HTTP_MCP_SESSION_ID" => session1 })
|
||||
response1 = transport.handle_request(get1)
|
||||
response1[2].call(io1) if response1[2].is_a?(Proc)
|
||||
|
||||
io2 = StringIO.new
|
||||
get2 = create_rack_request("GET", "/", { "HTTP_MCP_SESSION_ID" => session2 })
|
||||
response2 = transport.handle_request(get2)
|
||||
response2[2].call(io2) if response2[2].is_a?(Proc)
|
||||
|
||||
sleep(0.1)
|
||||
|
||||
# Call tool from session 1 with a progress token.
|
||||
tool_request = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json", "HTTP_MCP_SESSION_ID" => session1 },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "tools/call",
|
||||
id: "call-1",
|
||||
params: {
|
||||
name: "progress_tool",
|
||||
arguments: {},
|
||||
_meta: { progressToken: "token-1" },
|
||||
},
|
||||
}.to_json,
|
||||
)
|
||||
transport.handle_request(tool_request)
|
||||
|
||||
# Session 1 should receive the progress notification.
|
||||
io1.rewind
|
||||
assert_includes io1.read, "halfway"
|
||||
|
||||
# Session 2 should NOT receive the progress notification.
|
||||
io2.rewind
|
||||
refute_includes io2.read, "halfway"
|
||||
end
|
||||
|
||||
test "each session stores its own client info independently" do
|
||||
server = Server.new(name: "test", tools: [], prompts: [], resources: [])
|
||||
transport = StreamableHTTPTransport.new(server)
|
||||
server.transport = transport
|
||||
|
||||
# Initialize session 1 with client "alpha".
|
||||
init1 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
id: "1",
|
||||
params: {
|
||||
protocolVersion: "2025-11-25",
|
||||
clientInfo: { name: "alpha", version: "1.0" },
|
||||
},
|
||||
}.to_json,
|
||||
)
|
||||
session1 = transport.handle_request(init1)[1]["Mcp-Session-Id"]
|
||||
|
||||
# Initialize session 2 with client "beta".
|
||||
init2 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
id: "2",
|
||||
params: {
|
||||
protocolVersion: "2025-11-25",
|
||||
clientInfo: { name: "beta", version: "2.0" },
|
||||
},
|
||||
}.to_json,
|
||||
)
|
||||
session2 = transport.handle_request(init2)[1]["Mcp-Session-Id"]
|
||||
|
||||
# Each session should have its own client info.
|
||||
sessions = transport.instance_variable_get(:@sessions)
|
||||
assert_equal({ name: "alpha", version: "1.0" }, sessions[session1][:server_session].client)
|
||||
assert_equal({ name: "beta", version: "2.0" }, sessions[session2][:server_session].client)
|
||||
end
|
||||
|
||||
test "each session stores its own logging level independently" do
|
||||
server = Server.new(name: "test", tools: [], prompts: [], resources: [])
|
||||
transport = StreamableHTTPTransport.new(server)
|
||||
server.transport = transport
|
||||
|
||||
# Initialize two sessions.
|
||||
init1 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
id: "1",
|
||||
params: { protocolVersion: "2025-11-25", clientInfo: { name: "a" } },
|
||||
}.to_json,
|
||||
)
|
||||
session1 = transport.handle_request(init1)[1]["Mcp-Session-Id"]
|
||||
|
||||
init2 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
id: "2",
|
||||
params: { protocolVersion: "2025-11-25", clientInfo: { name: "b" } },
|
||||
}.to_json,
|
||||
)
|
||||
session2 = transport.handle_request(init2)[1]["Mcp-Session-Id"]
|
||||
|
||||
# Session 1 sets log level to "error".
|
||||
set_level1 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json", "HTTP_MCP_SESSION_ID" => session1 },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "logging/setLevel",
|
||||
id: "3",
|
||||
params: { level: "error" },
|
||||
}.to_json,
|
||||
)
|
||||
transport.handle_request(set_level1)
|
||||
|
||||
# Session 2 sets log level to "debug".
|
||||
set_level2 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json", "HTTP_MCP_SESSION_ID" => session2 },
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "logging/setLevel",
|
||||
id: "4",
|
||||
params: { level: "debug" },
|
||||
}.to_json,
|
||||
)
|
||||
transport.handle_request(set_level2)
|
||||
|
||||
# Session 1 (error level) should not notify for "info", but should for "error".
|
||||
session1_logging = transport.instance_variable_get(:@sessions)[session1][:server_session].logging_message_notification
|
||||
refute session1_logging.should_notify?("info")
|
||||
assert session1_logging.should_notify?("error")
|
||||
|
||||
# Session 2 (debug level) should notify for both "info" and "debug".
|
||||
session2_logging = transport.instance_variable_get(:@sessions)[session2][:server_session].logging_message_notification
|
||||
assert session2_logging.should_notify?("info")
|
||||
assert session2_logging.should_notify?("debug")
|
||||
end
|
||||
|
||||
private
|
||||
|
|
|
|||
|
|
@ -6,18 +6,18 @@ module MCP
|
|||
class ServerContextTest < ActiveSupport::TestCase
|
||||
test "ServerContext delegates method calls to the underlying context" do
|
||||
context = { user: "test_user" }
|
||||
progress = Progress.new(server: mock, progress_token: nil)
|
||||
progress = Progress.new(notification_target: mock, progress_token: nil)
|
||||
|
||||
server_context = ServerContext.new(context, progress: progress)
|
||||
server_context = ServerContext.new(context, progress: progress, notification_target: mock)
|
||||
|
||||
assert_equal "test_user", server_context[:user]
|
||||
end
|
||||
|
||||
test "ServerContext respond_to? returns true for methods on the underlying context" do
|
||||
context = { user: "test_user" }
|
||||
progress = Progress.new(server: mock, progress_token: nil)
|
||||
progress = Progress.new(notification_target: mock, progress_token: nil)
|
||||
|
||||
server_context = ServerContext.new(context, progress: progress)
|
||||
server_context = ServerContext.new(context, progress: progress, notification_target: mock)
|
||||
|
||||
assert_respond_to server_context, :[]
|
||||
assert_respond_to server_context, :keys
|
||||
|
|
@ -25,18 +25,18 @@ module MCP
|
|||
|
||||
test "ServerContext respond_to? returns false for methods not on the underlying context" do
|
||||
context = { user: "test_user" }
|
||||
progress = Progress.new(server: mock, progress_token: nil)
|
||||
progress = Progress.new(notification_target: mock, progress_token: nil)
|
||||
|
||||
server_context = ServerContext.new(context, progress: progress)
|
||||
server_context = ServerContext.new(context, progress: progress, notification_target: mock)
|
||||
|
||||
refute_respond_to server_context, :nonexistent_method
|
||||
end
|
||||
|
||||
test "ServerContext raises NoMethodError for methods not on the underlying context" do
|
||||
context = { user: "test_user" }
|
||||
progress = Progress.new(server: mock, progress_token: nil)
|
||||
progress = Progress.new(notification_target: mock, progress_token: nil)
|
||||
|
||||
server_context = ServerContext.new(context, progress: progress)
|
||||
server_context = ServerContext.new(context, progress: progress, notification_target: mock)
|
||||
|
||||
assert_raises(NoMethodError) { server_context.nonexistent_method }
|
||||
end
|
||||
|
|
@ -46,9 +46,9 @@ module MCP
|
|||
def context.custom_method
|
||||
"custom_value"
|
||||
end
|
||||
progress = Progress.new(server: mock, progress_token: nil)
|
||||
progress = Progress.new(notification_target: mock, progress_token: nil)
|
||||
|
||||
server_context = ServerContext.new(context, progress: progress)
|
||||
server_context = ServerContext.new(context, progress: progress, notification_target: mock)
|
||||
|
||||
assert_equal "custom_value", server_context.custom_method
|
||||
end
|
||||
|
|
@ -57,7 +57,7 @@ module MCP
|
|||
progress = mock
|
||||
progress.expects(:report).with(50, total: 100, message: nil).once
|
||||
|
||||
server_context = ServerContext.new(nil, progress: progress)
|
||||
server_context = ServerContext.new(nil, progress: progress, notification_target: mock)
|
||||
server_context.report_progress(50, total: 100)
|
||||
end
|
||||
|
||||
|
|
|
|||
|
|
@ -175,5 +175,23 @@ module MCP
|
|||
assert_equal Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED, notifications[2][:method]
|
||||
assert_equal Methods::NOTIFICATIONS_MESSAGE, notifications[3][:method]
|
||||
end
|
||||
|
||||
test "server.notify_log_message works after logging/setLevel via session" do
|
||||
session = ServerSession.new(server: @server, transport: @mock_transport)
|
||||
|
||||
# Client sends logging/setLevel through session.
|
||||
@server.handle(
|
||||
{ jsonrpc: "2.0", id: 1, method: "logging/setLevel", params: { level: "info" } },
|
||||
session: session,
|
||||
)
|
||||
|
||||
# Server-level broadcast should still work because logging level
|
||||
# is stored on both the session and the server.
|
||||
@server.notify_log_message(data: "broadcast log", level: "info")
|
||||
|
||||
log_notifications = @mock_transport.notifications.select { |n| n[:method] == Methods::NOTIFICATIONS_MESSAGE }
|
||||
assert_equal 1, log_notifications.size
|
||||
assert_equal "broadcast log", log_notifications.first[:params]["data"]
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
|
|||
|
|
@ -206,6 +206,20 @@ module MCP
|
|||
assert_instrumentation_data({ method: "ping" })
|
||||
end
|
||||
|
||||
test "unsupported method instrumentation includes client from session" do
|
||||
session = ServerSession.new(server: @server, transport: mock)
|
||||
session.store_client_info(client: { name: "session-client", version: "1.0" })
|
||||
|
||||
request = {
|
||||
jsonrpc: "2.0",
|
||||
method: "does/not/exist",
|
||||
id: 1,
|
||||
}
|
||||
|
||||
@server.handle(request, session: session)
|
||||
assert_instrumentation_data({ method: "unsupported_method", client: { name: "session-client", version: "1.0" } })
|
||||
end
|
||||
|
||||
test "#handle returns nil for notification requests" do
|
||||
request = {
|
||||
jsonrpc: "2.0",
|
||||
|
|
|
|||
Loading…
Reference in a new issue