Merge pull request #459 from koic/bound_client_sse_event_buffering

Bound client-side message buffering in the HTTP transport

README.md

@@ -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,

lib/mcp/client/http.rb

@@ -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
 

test/mcp/client/http_test.rb

@@ -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",