Bound client-side message buffering in the HTTP transport
## Motivation and Context `MCP::Client::HTTP` buffers server bytes with no upper bound in two places: - SSE responses are parsed incrementally, but the parser accumulates `data:` lines until a blank line dispatches the event, so a server that never terminates an event grows the buffer indefinitely. The abort-on-response signal does not help here, since it only fires after a complete response event has been dispatched. - Non-SSE (JSON) bodies accumulate in `SSEStream#buffer` chunk by chunk with no limit. Both the POST response stream and the standalone GET listening stream are affected. The stdio client already bounds server frames (`max_line_bytes:`, 4 MiB default) and the server transports bound request bodies (`max_request_bytes:`, 4 MiB default); this brings the HTTP client in line. `MCP::Client::HTTP.new` gains `max_message_bytes:` (default 4 MiB), which caps the bytes buffered for a single message - an SSE event or a JSON response body. External byte counting cannot tell consumed-and-discarded bytes (comments, field names, delimiters) from consumed-and-retained ones, so the check measures the parser's own String buffers after each chunk. An over-limit message raises `RequestHandlerError` on the request path and stops the GET listening stream instead of reconnecting to the same over-limit event. ## How Has This Been Tested? New tests in `test/mcp/client/http_test.rb` cover the SSE and JSON rejection paths end-to-end, the per-event reset (events that individually fit must not count against later ones), the env-less `on_data` fallback, and the constructor validation. `bundle exec rake` (tests, RuboCop, and conformance baseline) passes. ## Breaking Changes Responses larger than 4 MiB are now rejected by default; pass a larger g`max_message_bytes:` to accept them.
This commit is contained in:
Koichi ITO
parent
84090a88eb
commit
ced5cd1a95
3 changed files with 170 additions and 3 deletions
|
|
@ -2318,6 +2318,10 @@ response = client.call_tool(
|
|||
|
||||
The server will send `notifications/progress` back to the client during execution.
|
||||
|
||||
`MCP::Client::HTTP.new` accepts an optional `max_message_bytes:` keyword that caps the bytes buffered in memory for a single message from the server -
|
||||
an SSE event or a JSON response body. A message that reaches this limit before completing is rejected as a transport error, preventing unbounded memory growth from
|
||||
a server that never terminates an SSE event. It defaults to `4 * 1024 * 1024` (4 MiB); raise it if your server returns larger responses.
|
||||
|
||||
#### Server-to-Client Requests (Elicitation)
|
||||
|
||||
Servers can send requests back to the client while one of the client's own requests is in flight - for example,
|
||||
|
|
|
|||
|
|
@ -30,6 +30,13 @@ module MCP
|
|||
# (60 seconds for Net::HTTP) would recycle quiet streams too eagerly.
|
||||
SSE_LISTENER_READ_TIMEOUT = 300
|
||||
|
||||
# Upper bound in bytes on a single JSON-RPC message from the server - an SSE event or
|
||||
# a JSON response body - buffered in memory while reading a response. Without a bound,
|
||||
# a server that never terminates an SSE event (or never ends a JSON body) grows
|
||||
# the buffer indefinitely. Matches the 4 MiB default of `MCP::Client::Stdio::MAX_LINE_BYTES`
|
||||
# and the server transports' request cap.
|
||||
MAX_MESSAGE_BYTES = 4 * 1024 * 1024
|
||||
|
||||
# Raised when an `oauth:` provider is paired with an MCP URL that is neither HTTPS nor
|
||||
# a loopback `http://` URL, since a bearer token sent over plain HTTP to a remote host
|
||||
# is trivially observed and stolen.
|
||||
|
|
@ -71,6 +78,12 @@ module MCP
|
|||
class StreamAbort < StandardError; end
|
||||
private_constant :StreamAbort
|
||||
|
||||
# Raised inside the streaming callback when the server sends more bytes than
|
||||
# `max_message_bytes` without completing a message. Translated into
|
||||
# a `RequestHandlerError` (or a listener stop) where the request context is known.
|
||||
class MessageTooLargeError < StandardError; end
|
||||
private_constant :MessageTooLargeError
|
||||
|
||||
# Per-exchange SSE state shared between the initial POST stream and any SEP-1699 reconnection GET streams:
|
||||
# the incrementally parsed JSON-RPC response, the last received SSE event id (for `Last-Event-ID`),
|
||||
# and the server's `retry:` reconnection delay. Non-SSE bodies accumulate in `buffer` for the JSON path.
|
||||
|
|
@ -78,8 +91,9 @@ module MCP
|
|||
attr_reader :buffer, :response, :last_event_id, :retry_ms
|
||||
attr_accessor :abortable
|
||||
|
||||
def initialize(abortable:, on_request: nil)
|
||||
def initialize(abortable:, max_message_bytes: MAX_MESSAGE_BYTES, on_request: nil)
|
||||
@abortable = abortable
|
||||
@max_message_bytes = max_message_bytes
|
||||
@on_request = on_request
|
||||
@buffer = +""
|
||||
@parser = nil
|
||||
|
|
@ -95,14 +109,24 @@ module MCP
|
|||
end
|
||||
|
||||
# Faraday `on_data` streaming callback. SSE chunks are parsed incrementally;
|
||||
# anything else (JSON bodies) accumulates in `buffer`.
|
||||
# anything else (JSON bodies) accumulates in `buffer`. Both paths are bounded
|
||||
# by `max_message_bytes`, checked after the awaited response had its chance to
|
||||
# abort the stream so an over-limit tail cannot mask a response that already arrived.
|
||||
def on_data
|
||||
proc do |chunk, _received_bytes, env|
|
||||
if event_stream?(env)
|
||||
feed(chunk)
|
||||
raise StreamAbort if @abortable && @response
|
||||
|
||||
if parser_buffered_bytes > @max_message_bytes
|
||||
raise MessageTooLargeError, "Server SSE event exceeds #{@max_message_bytes} bytes without completing"
|
||||
end
|
||||
else
|
||||
@buffer << chunk
|
||||
|
||||
if @buffer.bytesize > @max_message_bytes
|
||||
raise MessageTooLargeError, "Server response body exceeds #{@max_message_bytes} bytes"
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
@ -154,6 +178,20 @@ module MCP
|
|||
end
|
||||
end
|
||||
|
||||
# Bytes the parser has buffered for a not-yet-dispatched event: the partial line
|
||||
# and the accumulated `data:` field values (plus the small event type and last event id buffers).
|
||||
# External byte counting cannot tell consumed-and-discarded bytes (comments, field names, delimiters)
|
||||
# from consumed-and-retained ones, so the parser's own String buffers are measured instead;
|
||||
# summing every String instance variable keeps the measurement valid if the parser renames or adds buffers.
|
||||
# A parser rewrite that buffered elsewhere would make this return 0 and quietly lift the cap -
|
||||
# the regression tests feeding an over-limit event guard against that.
|
||||
def parser_buffered_bytes
|
||||
parser.instance_variables.sum do |name|
|
||||
value = parser.instance_variable_get(name)
|
||||
value.is_a?(String) ? value.bytesize : 0
|
||||
end
|
||||
end
|
||||
|
||||
def parser
|
||||
@parser ||= begin
|
||||
require "event_stream_parser"
|
||||
|
|
@ -169,7 +207,13 @@ module MCP
|
|||
|
||||
attr_reader :url, :session_id, :protocol_version, :server_info, :oauth
|
||||
|
||||
def initialize(url:, headers: {}, oauth: nil, &block)
|
||||
def initialize(url:, headers: {}, oauth: nil, max_message_bytes: MAX_MESSAGE_BYTES, &block)
|
||||
# `nil` or a non-positive value would make the buffering unbounded and silently
|
||||
# disable the protection, so reject it up front.
|
||||
unless max_message_bytes.is_a?(Integer) && max_message_bytes > 0
|
||||
raise ArgumentError, "max_message_bytes must be a positive Integer"
|
||||
end
|
||||
|
||||
if oauth && !MCP::Client::OAuth::Discovery.secure_url?(url)
|
||||
# Mask credentials (userinfo) and query parameters before quoting the URL in the error message
|
||||
# so they cannot leak into logs.
|
||||
|
|
@ -183,6 +227,7 @@ module MCP
|
|||
@headers = headers
|
||||
@faraday_customizer = block
|
||||
@oauth = oauth
|
||||
@max_message_bytes = max_message_bytes
|
||||
# Snapshot the canonical URL at construction time. This single value
|
||||
# serves two related roles, both of which need to see the query string:
|
||||
#
|
||||
|
|
@ -336,6 +381,7 @@ module MCP
|
|||
# so the response object (and its `Mcp-Session-Id` header) is always available for capture.
|
||||
stream = SSEStream.new(
|
||||
abortable: method.to_s != MCP::Methods::INITIALIZE,
|
||||
max_message_bytes: @max_message_bytes,
|
||||
on_request: ->(message) { dispatch_server_request(message) },
|
||||
)
|
||||
|
||||
|
|
@ -354,6 +400,12 @@ module MCP
|
|||
capture_session_info(method, response, body) if response
|
||||
|
||||
body
|
||||
rescue MessageTooLargeError => e
|
||||
raise RequestHandlerError.new(
|
||||
e.message,
|
||||
{ method: method, params: params },
|
||||
error_type: :internal_error,
|
||||
)
|
||||
rescue Faraday::BadRequestError => e
|
||||
raise RequestHandlerError.new(
|
||||
"The #{method} request is invalid",
|
||||
|
|
@ -699,6 +751,7 @@ module MCP
|
|||
def listen_for_server_requests
|
||||
stream = SSEStream.new(
|
||||
abortable: false,
|
||||
max_message_bytes: @max_message_bytes,
|
||||
on_request: ->(message) { dispatch_server_request(message) },
|
||||
)
|
||||
consecutive_failures = 0
|
||||
|
|
@ -714,6 +767,9 @@ module MCP
|
|||
end
|
||||
|
||||
consecutive_failures = 0
|
||||
rescue MessageTooLargeError
|
||||
# Reconnecting would let the server stream the same over-limit event again, so stop listening instead of retrying.
|
||||
break
|
||||
rescue Faraday::Error => e
|
||||
break if e.response&.dig(:status) == 405
|
||||
|
||||
|
|
|
|||
|
|
@ -853,6 +853,113 @@ module MCP
|
|||
assert_empty(stream.buffer)
|
||||
end
|
||||
|
||||
def test_send_request_rejects_sse_event_exceeding_max_message_bytes
|
||||
request = {
|
||||
jsonrpc: "2.0",
|
||||
id: "test_id",
|
||||
method: "tools/list",
|
||||
}
|
||||
client = HTTP.new(url: url, max_message_bytes: 64)
|
||||
|
||||
stub_request(:post, url).with(
|
||||
body: request.to_json,
|
||||
).to_return(
|
||||
status: 200,
|
||||
headers: { "Content-Type" => "text/event-stream" },
|
||||
body: "data: #{"a" * 128}\n",
|
||||
)
|
||||
|
||||
error = assert_raises(RequestHandlerError) do
|
||||
client.send_request(request: request)
|
||||
end
|
||||
|
||||
assert_includes(error.message, "Server SSE event exceeds 64 bytes")
|
||||
assert_equal(:internal_error, error.error_type)
|
||||
end
|
||||
|
||||
def test_send_request_rejects_json_body_exceeding_max_message_bytes
|
||||
request = {
|
||||
jsonrpc: "2.0",
|
||||
id: "test_id",
|
||||
method: "tools/list",
|
||||
}
|
||||
client = HTTP.new(url: url, max_message_bytes: 64)
|
||||
|
||||
stub_request(:post, url).with(
|
||||
body: request.to_json,
|
||||
).to_return(
|
||||
status: 200,
|
||||
headers: { "Content-Type" => "application/json" },
|
||||
body: { result: { tools: ["a" * 128] } }.to_json,
|
||||
)
|
||||
|
||||
error = assert_raises(RequestHandlerError) do
|
||||
client.send_request(request: request)
|
||||
end
|
||||
|
||||
assert_includes(error.message, "Server response body exceeds 64 bytes")
|
||||
assert_equal(:internal_error, error.error_type)
|
||||
end
|
||||
|
||||
def test_send_request_accepts_sse_events_that_each_stay_within_max_message_bytes
|
||||
request = {
|
||||
jsonrpc: "2.0",
|
||||
id: "test_id",
|
||||
method: "tools/list",
|
||||
}
|
||||
client = HTTP.new(url: url, max_message_bytes: 256)
|
||||
|
||||
# The notifications and the response together exceed the limit; each event stays
|
||||
# under it, so dispatched events must not count against the ones that follow.
|
||||
notification = { jsonrpc: "2.0", method: "notifications/progress", params: { token: "a" * 100 } }.to_json
|
||||
sse_body = "data: #{notification}\n\n" * 3 +
|
||||
"data: {\"jsonrpc\":\"2.0\",\"id\":\"test_id\",\"result\":{\"tools\":[]}}\n\n"
|
||||
|
||||
stub_request(:post, url).with(
|
||||
body: request.to_json,
|
||||
).to_return(
|
||||
status: 200,
|
||||
headers: { "Content-Type" => "text/event-stream" },
|
||||
body: sse_body,
|
||||
)
|
||||
|
||||
response = client.send_request(request: request)
|
||||
|
||||
assert_equal({ "tools" => [] }, response["result"])
|
||||
end
|
||||
|
||||
def test_sse_stream_rejects_unterminated_event_exceeding_max_message_bytes
|
||||
stream = HTTP.const_get(:SSEStream).new(abortable: false, max_message_bytes: 64)
|
||||
env = stub(response_headers: { "content-type" => "text/event-stream" })
|
||||
chunk = "data: #{"a" * 128}"
|
||||
|
||||
error = assert_raises(HTTP.const_get(:MessageTooLargeError)) do
|
||||
stream.on_data.call(chunk, chunk.bytesize, env)
|
||||
end
|
||||
|
||||
assert_includes(error.message, "Server SSE event exceeds 64 bytes")
|
||||
end
|
||||
|
||||
def test_sse_stream_bounds_buffered_chunks_when_env_is_unavailable
|
||||
stream = HTTP.const_get(:SSEStream).new(abortable: true, max_message_bytes: 64)
|
||||
|
||||
error = assert_raises(HTTP.const_get(:MessageTooLargeError)) do
|
||||
stream.on_data.call("a" * 65, 65)
|
||||
end
|
||||
|
||||
assert_includes(error.message, "Server response body exceeds 64 bytes")
|
||||
end
|
||||
|
||||
def test_initialize_rejects_max_message_bytes_that_is_not_a_positive_integer
|
||||
[0, -1, nil, "4"].each do |value|
|
||||
error = assert_raises(ArgumentError) do
|
||||
HTTP.new(url: url, max_message_bytes: value)
|
||||
end
|
||||
|
||||
assert_equal("max_message_bytes must be a positive Integer", error.message)
|
||||
end
|
||||
end
|
||||
|
||||
def test_send_request_dispatches_server_request_to_registered_handler
|
||||
request = {
|
||||
jsonrpc: "2.0",
|
||||
|
|
|
|||
Loading…
Reference in a new issue