Merge pull request #521 from koic/reject_duplicate_in_flight_request_ids
Reject duplicate in-flight JSON-RPC request IDs
This commit is contained in:
Koichi ITO
GitHub
commit
4eb66155e2
2 changed files with 287 additions and 12 deletions
|
|
@ -3431,6 +3431,110 @@ module MCP
|
|||
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)
|
||||
|
|
@ -6199,6 +6303,36 @@ module MCP
|
|||
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"
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ require "test_helper"
|
|||
module MCP
|
||||
class ServerCancellationTest < ActiveSupport::TestCase
|
||||
include InstrumentationTestHelper
|
||||
include InitializeParamsTestHelper
|
||||
|
||||
class MockTransport < Transport
|
||||
attr_reader :requests, :notifications, :cancelled_request_ids
|
||||
|
|
@ -333,6 +334,74 @@ module MCP
|
|||
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: {} })
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue