Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 22 additions & 3 deletions lib/mcp/server.rb
Original file line number Diff line number Diff line change
Expand Up @@ -619,8 +619,24 @@ def handle_request(request, method, session: nil, related_request_id: nil)

# `initialize` MUST NOT be cancelled (MCP spec 2025-11-25, cancellation item 2),
# so do not track it in the in-flight registry.
cancellation = if related_request_id && method != Methods::INITIALIZE
session&.register_in_flight(related_request_id)
cancellation = nil
if related_request_id && method != Methods::INITIALIZE && session
cancellation = session.register_in_flight(related_request_id)

# The spec puts the uniqueness obligation on the sender - "The request ID MUST NOT have been previously used by
# the requestor within the same session" - and says nothing about what a receiver does with a duplicate.
# Answering one is the only option that stays correct: the id routes request-scoped messages back to
# the request that caused them, and the transport's rule is that those messages "SHOULD relate to
# the originating client request", which a second live request under the same id makes impossible to honor
# for either of them. Refused the same way a duplicate `initialize` is, and for the same reason:
# so that a repeated id cannot silently displace state negotiated by the first one.
if cancellation.nil?
raise RequestHandlerError.new(
"Invalid Request: request id #{related_request_id.inspect} is already in flight",
request,
error_type: :invalid_request,
)
end
end

->(params) {
Expand Down Expand Up @@ -727,7 +743,10 @@ def handle_request(request, method, session: nil, related_request_id: nil)
reported_exception = wrapped
raise wrapped
ensure
session&.unregister_in_flight(related_request_id) if related_request_id
# `cancellation` is non-nil exactly when this request claimed the id above, so this also keeps `initialize`
# (which never registers) from evicting an in-flight registration under a reused id when the duplicate-`initialize`
# refusal raises out of the handler.
session&.unregister_in_flight(related_request_id, cancellation: cancellation) if related_request_id && cancellation
end
}
end
Expand Down
39 changes: 35 additions & 4 deletions lib/mcp/server/transports/streamable_http_transport.rb
Original file line number Diff line number Diff line change
Expand Up @@ -450,8 +450,13 @@ def drop_broken_stream(session_id, stream, related_request_id)

@mutex.synchronize do
session = @sessions[session_id]
if related_request_id && session&.dig(:post_request_streams, related_request_id)
session[:post_request_streams].delete(related_request_id)
if related_request_id
# Unregister only our own stream: removing on the id alone would drop whichever stream currently holds it,
# which is not necessarily the one that failed. The failed stream is closed either way, and a request-scoped
# failure never reaches the session teardown below.
registered = session&.dig(:post_request_streams, related_request_id)
session[:post_request_streams].delete(related_request_id) if registered.equal?(stream)

streams_to_close << stream
else
cleanup_and_collect_stream(session_id, streams_to_close)
Expand Down Expand Up @@ -1595,6 +1600,14 @@ def handle_regular_request(body_string, session_id, related_request_id: nil)
end
end

# `Server` refuses a duplicate id as well, but only once the request reaches it. The SSE branch below
# registers this request's stream under that id first, so without this check the colliding request
# would take over the routing entry for the moment it takes to be rejected, and its `ensure` would then
# clear the entry the original request still needs.
if related_request_id && server_session&.in_flight?(related_request_id)
return request_id_conflict_response
end

if session_id && !@stateless && !@enable_json_response
handle_request_with_sse_response(body_string, session_id, server_session, related_request_id: related_request_id)
else
Expand All @@ -1618,7 +1631,11 @@ def handle_request_with_sse_response(body_string, session_id, server_session, re
session = @sessions[session_id]
if session && related_request_id
session[:post_request_streams] ||= {}
session[:post_request_streams][related_request_id] = stream

# Claim the id only while it is free. `handle_regular_request` already refused the colliding request,
# so reaching an occupied slot means a race got past that check; leaving the first stream in place keeps
# its messages going where they belong.
session[:post_request_streams][related_request_id] ||= stream
end
end

Expand All @@ -1630,7 +1647,11 @@ def handle_request_with_sse_response(body_string, session_id, server_session, re
if related_request_id
@mutex.synchronize do
session = @sessions[session_id]
session[:post_request_streams]&.delete(related_request_id) if session
# Only retire our own registration: a request that never claimed the id, or one whose claim has
# already been replaced, must not unregister the stream that owns it.
registered = session&.dig(:post_request_streams, related_request_id)

session[:post_request_streams].delete(related_request_id) if registered.equal?(stream)
end
end

Expand Down Expand Up @@ -1849,6 +1870,16 @@ def session_already_connected_response
)
end

# The POST counterpart of the GET conflict above. A request id already in flight cannot be given
# a stream of its own, because the id is what routes request-scoped messages back.
def request_id_conflict_response
json_rpc_error_response(
status: 409,
code: JsonRpcHandler::ErrorCode::INVALID_REQUEST,
message: "Conflict: Request id is already in flight for this session",
)
end

def setup_sse_stream(session_id)
body = create_sse_body(session_id)

Expand Down
32 changes: 27 additions & 5 deletions lib/mcp/server_session.rb
Original file line number Diff line number Diff line change
Expand Up @@ -58,19 +58,41 @@ def lock_era!(era)
@era = era
end

# Registers a `Cancellation` token for an in-flight request.
# Registers a `Cancellation` token for an in-flight request, or returns `nil` when `request_id` is already in flight.
# The request id is the only key that routes request-scoped notifications, server-to-client requests,
# and `notifications/cancelled` back to the request that caused them, so a second live request under the same id
# has no destination of its own. Rather than let the newcomer displace the registration, report the collision
# and leave the first request intact; the caller turns that into an Invalid Request.
def register_in_flight(request_id)
return if request_id.nil?

cancellation = Cancellation.new(request_id: request_id)
@in_flight_mutex.synchronize { @in_flight[request_id] = cancellation }
cancellation
registered = @in_flight_mutex.synchronize do
next false if @in_flight.key?(request_id)

@in_flight[request_id] = cancellation
true
end

registered ? cancellation : nil
end

def unregister_in_flight(request_id)
# Removes an in-flight registration. Passing the `Cancellation` that `register_in_flight` returned removes
# the entry only while it is still that one, so a request can never evict a registration it does not own.
def unregister_in_flight(request_id, cancellation: nil)
return if request_id.nil?

@in_flight_mutex.synchronize { @in_flight.delete(request_id) }
@in_flight_mutex.synchronize do
next if cancellation && !@in_flight[request_id].equal?(cancellation)

@in_flight.delete(request_id)
end
end

# Whether `request_id` is currently in flight, so a transport can refuse a colliding request
# before registering any state of its own for it.
def in_flight?(request_id)
!lookup_in_flight(request_id).nil?
end

def lookup_in_flight(request_id)
Expand Down
134 changes: 134 additions & 0 deletions test/mcp/server/transports/streamable_http_transport_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -3431,6 +3431,110 @@ def string
assert_equal "ok", result[:content][:text]
end

test "refuses a POST reusing an in-flight request id and keeps the original stream" do
server = Server.new(name: "test", tools: [], prompts: [], resources: [])
transport = StreamableHTTPTransport.new(server)
server.transport = transport

gate = Queue.new
server.define_tool(name: "victim_tool") do |server_context:|
server_context.report_progress(1, message: "first-frame")
gate.pop
server_context.report_progress(2, message: "second-frame")
Tool::Response.new([{ type: "text", text: "done" }])
end

session_id, server_session = start_session(transport)

victim = transport.handle_request(colliding_tool_call(session_id, "req-1"))
victim_stream = TestStream.new
victim_thread = Thread.new { victim[2].call(victim_stream) }
sleep(0.01) until server_session.lookup_in_flight("req-1")

attacker = transport.handle_request(colliding_tool_call(session_id, "req-1"))

assert_equal 409, attacker[0]
assert_equal(
JsonRpcHandler::ErrorCode::INVALID_REQUEST,
JSON.parse(attacker[2][0]).dig("error", "code"),
)

gate << :go
victim_thread.join

# The registration survived the refusal, so the frame emitted after it still reaches
# the request that asked for it.
assert_includes victim_stream.string, "first-frame"
assert_includes victim_stream.string, "second-frame"
end

test "a finishing request leaves a stream another request registered under the same id" do
server = Server.new(name: "test", tools: [], prompts: [], resources: [])
transport = StreamableHTTPTransport.new(server)
server.transport = transport

gate = Queue.new
server.define_tool(name: "victim_tool") do |server_context:|
gate.pop
Tool::Response.new([{ type: "text", text: "done" }])
end

session_id, server_session = start_session(transport)

victim = transport.handle_request(colliding_tool_call(session_id, "req-1"))
victim_stream = TestStream.new
victim_thread = Thread.new { victim[2].call(victim_stream) }
sleep(0.01) until server_session.lookup_in_flight("req-1")

# Stand in for a stream that won the registration in a race the pre-dispatch refusal normally prevents.
# Finishing the other request must not unregister it.
foreign_stream = TestStream.new
sessions = transport.instance_variable_get(:@sessions)
sessions[session_id][:post_request_streams]["req-1"] = foreign_stream

gate << :go
victim_thread.join

assert_same foreign_stream, sessions[session_id][:post_request_streams]["req-1"]
end

test "a broken request-scoped stream drops only itself, even when it is not the registered one" do
server = Server.new(name: "test", tools: [], prompts: [], resources: [])
transport = StreamableHTTPTransport.new(server)
server.transport = transport

session_id, = start_session(transport)
sessions = transport.instance_variable_get(:@sessions)
registered = TestStream.new
sessions[session_id][:post_request_streams] = { "req-1" => registered }

transport.send(:drop_broken_stream, session_id, TestStream.new, "req-1")

assert sessions.key?(session_id), "a request-scoped failure must not tear down the session"
assert_same registered, sessions[session_id][:post_request_streams]["req-1"]
end

test "allows a request id to be reused once the earlier request has finished" do
server = Server.new(name: "test", tools: [], prompts: [], resources: [])
transport = StreamableHTTPTransport.new(server)
server.transport = transport

server.define_tool(name: "victim_tool") do |server_context:|
Tool::Response.new([{ type: "text", text: "done" }])
end

session_id, = start_session(transport)

2.times do
response = transport.handle_request(colliding_tool_call(session_id, "req-1"))

assert_equal 200, response[0]
stream = TestStream.new
response[2].call(stream)
assert_includes stream.string, "done"
end
end

test "JSON response mode returns accepted when cancellation suppresses response" do
server = Server.new(name: "test", tools: [], prompts: [], resources: [])
transport = StreamableHTTPTransport.new(server, enable_json_response: true)
Expand Down Expand Up @@ -6199,6 +6303,36 @@ def install_mutex_probe_stream(session_id, related_request_id: nil)
writes
end

# Initializes `transport` and returns its session id together with the `ServerSession`,
# which the in-flight request id tests poll to know when a request has really started.
def start_session(transport)
request = create_rack_request(
"POST",
"/",
{ "CONTENT_TYPE" => "application/json" },
{ jsonrpc: "2.0", method: "initialize", id: "init", params: initialize_params }.to_json,
)
session_id = transport.handle_request(request)[1]["mcp-session-id"]

[session_id, transport.instance_variable_get(:@sessions)[session_id][:server_session]]
end

# A `tools/call` for `victim_tool` under a caller-chosen request id, with a progress token so
# the tool's `report_progress` has somewhere to go.
def colliding_tool_call(session_id, request_id)
create_rack_request(
"POST",
"/",
{ "CONTENT_TYPE" => "application/json", "HTTP_MCP_SESSION_ID" => session_id },
{
jsonrpc: "2.0",
id: request_id,
method: "tools/call",
params: { name: "victim_tool", arguments: {}, _meta: { progressToken: "tok" } },
}.to_json,
)
end

def create_rack_request(method, path, headers, body = nil)
default_accept = case method
when "POST"
Expand Down
69 changes: 69 additions & 0 deletions test/mcp/server_cancellation_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
module MCP
class ServerCancellationTest < ActiveSupport::TestCase
include InstrumentationTestHelper
include InitializeParamsTestHelper

class MockTransport < Transport
attr_reader :requests, :notifications, :cancelled_request_ids
Expand Down Expand Up @@ -333,6 +334,74 @@ def handle_request(request); end
assert_includes @mock_transport.cancelled_request_ids, "req-9"
end

test "register_in_flight refuses an id that is already in flight" do
first = @session.register_in_flight("req-1")

assert first
assert_nil @session.register_in_flight("req-1"), "a second live request cannot share the id"
assert_same first, @session.lookup_in_flight("req-1"), "the first registration must survive"
end

test "unregister_in_flight leaves a registration it does not own" do
owner = @session.register_in_flight("req-1")
@session.unregister_in_flight("req-1", cancellation: Cancellation.new(request_id: "req-1"))

assert_same owner, @session.lookup_in_flight("req-1")

@session.unregister_in_flight("req-1", cancellation: owner)

assert_nil @session.lookup_in_flight("req-1")
end

test "in_flight? reports whether an id is registered" do
refute @session.in_flight?("req-1")

@session.register_in_flight("req-1")

assert @session.in_flight?("req-1")
end

test "a request reusing an in-flight id is answered with Invalid Request" do
@server.define_tool(name: "slow") do |server_context:|
sleep(0.2)
Tool::Response.new([{ type: "text", text: "ok" }])
end

request = {
jsonrpc: "2.0",
id: "req-1",
method: Methods::TOOLS_CALL,
params: { name: "slow", arguments: {} },
}

in_flight = Thread.new { @session.handle(request) }
sleep(0.01) until @session.lookup_in_flight("req-1")

duplicate = @session.handle(request)

assert_equal(
JsonRpcHandler::ErrorCode::INVALID_REQUEST,
duplicate.dig(:error, :code) || duplicate.dig("error", "code"),
)
in_flight.join
end

test "a refused duplicate initialize leaves an in-flight registration under its reused id" do
@session.handle(jsonrpc: "2.0", id: "init", method: Methods::INITIALIZE, params: initialize_params)

owner = @session.register_in_flight("req-1")

# `initialize` bypasses the duplicate-id refusal (it is never in flight itself), so its
# rejection path is the one place a reused id reaches a handler while the id is still live.
duplicate = @session.handle(jsonrpc: "2.0", id: "req-1", method: Methods::INITIALIZE, params: initialize_params)

assert_equal(
JsonRpcHandler::ErrorCode::INVALID_REQUEST,
duplicate.dig(:error, :code) || duplicate.dig("error", "code"),
)
assert_same owner, @session.lookup_in_flight("req-1"), "the registration must survive the refused initialize"
end

test "parent cancellation propagates to nested server-to-client requests" do
@session.instance_variable_set(:@client_capabilities, { elicitation: {} })

Expand Down