ruby-utcp 1.1.5 → 1.1.6

data/lib/utcp/protocols/mcp.rb

@@ -7,12 +7,14 @@ module UTCP
     MAX_MESSAGE_BYTES = 16 * 1024 * 1024
     MAX_STDERR_BYTES = 64 * 1024
 
-    def initialize(config, timeout, max_message_bytes: MAX_MESSAGE_BYTES)
+    def initialize(config, timeout, max_message_bytes: MAX_MESSAGE_BYTES, max_response_bytes: nil)
       @timeout = Float(timeout)
       raise ValidationError, "MCP timeout must be finite and greater than zero" unless @timeout.finite? && @timeout.positive?
 
       @max_message_bytes = Integer(max_message_bytes)
       raise ValidationError, "MCP message limit must be greater than zero" unless @max_message_bytes.positive?
+      @max_response_bytes = max_response_bytes.nil? ? @max_message_bytes : Integer(max_response_bytes)
+      raise ValidationError, "MCP response limit must be greater than zero" unless @max_response_bytes.positive?
 
       command = config["command"]
       command = command.first if command.is_a?(Array)
@@ -41,8 +43,9 @@ module UTCP
         message = { "jsonrpc" => "2.0", "id" => identifier, "method" => method }
         message["params"] = params unless params.nil?
         write_message(message, deadline)
+        budget = ResponseByteBudget.new(@max_response_bytes, "MCP stdio")
         loop do
-          response = read_message(deadline)
+          response = read_message(deadline, budget)
           next unless response["id"] == identifier
           raise ToolCallError, "MCP error #{response["error"].inspect}" if response["error"]
 
@@ -99,20 +102,22 @@ module UTCP
       raise ToolCallError, "MCP stdio write failed: #{error.message}"
     end
 
-    def read_message(deadline)
+    def read_message(deadline, budget = ResponseByteBudget.new(@max_response_bytes, "MCP stdio"))
       loop do
         remaining_time(deadline)
         if (index = @read_buffer.index("\n"))
-          response = JSON.parse(@read_buffer.slice!(0, index + 1))
+          bytes = @read_buffer.slice!(0, index + 1)
+          budget.consume(bytes)
+          response = JSON.parse(bytes)
           raise SerializerValidationError, "MCP stdio response must be an object" unless response.is_a?(Hash)
 
           return response
         end
-        if @read_buffer.bytesize >= @max_message_bytes
-          raise SerializerValidationError, "MCP stdio message exceeds #{@max_message_bytes} bytes"
+        if @read_buffer.bytesize >= budget.remaining
+          raise SerializerValidationError, "MCP stdio message exceeds #{@max_response_bytes} bytes (max_response_bytes)"
         end
 
-        chunk = @stdout.read_nonblock([4096, @max_message_bytes - @read_buffer.bytesize].min, exception: false)
+        chunk = @stdout.read_nonblock([4096, budget.remaining - @read_buffer.bytesize].min, exception: false)
         case chunk
         when :wait_readable then wait_for_io(deadline)
         when nil then raise ToolCallError, "MCP stdio server closed the stream"
@@ -211,7 +216,7 @@ module UTCP
       values = []
       parser = SSEParser.new
       parser.feed(body) { |value| values << value }
-      parser.finish { |value| values << value }
+      parser.finish
       values
     end
   end
@@ -229,10 +234,11 @@ module UTCP
       assert_mcp_template!(template)
       tools = []
       errors = []
+      budget = ResponseByteBudget.new(template.max_response_bytes, "MCP discovery")
       template.servers.each do |server_name, config|
         begin
           session = session_for(client, template, server_name, config)
-          each_list_item(session, "tools/list", "tools") do |tool|
+          each_list_item(session, "tools/list", "tools", budget) do |tool|
             tools << Tool.new(
               name: "#{server_name}.#{tool.fetch("name")}",
               description: tool["description"].to_s,
@@ -241,7 +247,7 @@ module UTCP
               tool_call_template: template
             )
           end
-          add_resource_tools(client, template, server_name, session, tools) if template.register_resources_as_tools
+          add_resource_tools(client, template, server_name, session, tools, budget) if template.register_resources_as_tools
         rescue StandardError => error
           errors << "#{server_name}: #{error.message}"
         end
@@ -280,6 +286,7 @@ module UTCP
                else
                  session.request("tools/call", "name" => local_name, "arguments" => Utils.stringify_keys(tool_args || {}))
                end
+      ResponseByteBudget.new(template.max_response_bytes, "MCP").consume_value(result)
       process_mcp_result(result, tool_name: tool_name)
     rescue Error
       raise
@@ -308,7 +315,7 @@ module UTCP
         body: message,
         content_type: "application/json",
         timeout: template.timeout,
-        sensitive_headers: sensitive.uniq
+        sensitive_headers: sensitive.uniq, max_response_bytes: template.max_response_bytes
       )
       {
         body: response.body.to_s,
@@ -330,7 +337,7 @@ module UTCP
       assert_no_auth!(template, context: "MCP stdio") unless http_transport || @session_factory
       # A session belongs to the credentials and endpoint used to initialize it.
       fingerprint = Digest::SHA256.hexdigest(JSON.generate([
-        config, template.auth&.to_h, template.protocol_version, template.timeout
+        config, template.auth&.to_h, template.protocol_version, template.timeout, template.max_response_bytes
       ]))
       key = [client, template.name, server_name, fingerprint]
       @sessions_mutex.synchronize do
@@ -343,7 +350,7 @@ module UTCP
                   elsif http_transport
                     MCPHTTPSession.new(server_config, snapshot, self)
                   else
-                    MCPStdioSession.new(server_config, snapshot.timeout)
+                    MCPStdioSession.new(server_config, snapshot.timeout, max_response_bytes: snapshot.max_response_bytes)
                   end
         begin
           initialize_session(session, snapshot)
@@ -356,20 +363,22 @@ module UTCP
     end
 
     def initialize_session(session, template)
-      session.request("initialize", {
+      response = session.request("initialize", {
         "protocolVersion" => template.protocol_version,
         "capabilities" => {},
         "clientInfo" => { "name" => "ruby-utcp", "version" => VERSION }
       })
+      ResponseByteBudget.new(template.max_response_bytes, "MCP initialization").consume_value(response)
       session.notify("notifications/initialized", {})
     end
 
-    def each_list_item(session, method, collection)
+    def each_list_item(session, method, collection, budget)
       cursor = nil
       seen = {}
       loop do
         params = cursor.nil? ? {} : { "cursor" => cursor }
         result = Utils.hash!(session.request(method, params) || {}, "MCP #{method} result")
+        budget.consume_value(result)
         Utils.array!(result.fetch(collection, []), "MCP #{collection}").each { |item| yield item }
         cursor = result["nextCursor"]
         break if cursor.nil?
@@ -381,8 +390,8 @@ module UTCP
       end
     end
 
-    def add_resource_tools(client, template, server_name, session, tools)
-      each_list_item(session, "resources/list", "resources") do |resource|
+    def add_resource_tools(client, template, server_name, session, tools, budget)
+      each_list_item(session, "resources/list", "resources", budget) do |resource|
         safe_name = resource.fetch("name", resource.fetch("uri")).to_s.gsub(/[^[:alnum:]_]/, "_")
         local_name = "resource_#{safe_name}"
         @sessions_mutex.synchronize do

data/lib/utcp/protocols/sse.rb

@@ -16,7 +16,9 @@ module UTCP
 
     def call_tool(client, tool_name, tool_args, template)
       values = []
-      call_tool_streaming(client, tool_name, tool_args, template) { |value| values << value }
+      with_collection_timeout(template) do
+        call_tool_streaming(client, tool_name, tool_args, template) { |value| values << value }
+      end
       values
     end
 
@@ -25,11 +27,16 @@ module UTCP
 
       assert_sse_template!(template)
       with_stream_response(template, tool_args || {}, accept: "text/event-stream") do |response|
-        parser = SSEParser.new(event_type: template.event_type)
+        parser = SSEParser.new(event_type: template.event_type, max_event_bytes: template.max_event_bytes)
+        count = 0
         response.read_body do |chunk|
-          parser.feed(chunk) { |event| yield event }
+          parser.feed(chunk) do |event|
+            count += 1
+            raise ToolCallError, "SSE response exceeds max_response_items" if count > template.max_response_items
+            yield event
+          end
         end
-        parser.finish { |event| yield event }
+        parser.finish
       end
     rescue Error
       raise
@@ -47,30 +54,57 @@ module UTCP
   end
 
   class SSEParser
-    def initialize(event_type: nil)
+    def initialize(event_type: nil, max_event_bytes: ResponseLimits::DEFAULT_MAX_EVENT_BYTES)
       @event_type = event_type
-      @buffer = +""
+      @maximum = Integer(max_event_bytes)
+      raise ValidationError, "max_event_bytes must be greater than zero" unless @maximum.positive?
+      @buffer = +"".b
+      @skip_lf = false
+      @first_line = true
+      @event_bytes = 0
       @fields = reset_fields
     end
 
     def feed(chunk)
-      @buffer << chunk.to_s.gsub("\r\n", "\n").gsub("\r", "\n")
-      while (index = @buffer.index("\n"))
-        line = @buffer.slice!(0..index).chomp
-        process_line(line) { |event| yield event }
+      bytes = chunk.to_s.b
+      offset = 0
+      while offset < bytes.bytesize
+        if @skip_lf
+          @skip_lf = false
+          offset += 1 if bytes.getbyte(offset) == 10
+          next if offset == bytes.bytesize
+        end
+        index = bytes.index(/[\r\n]/, offset)
+        length = (index || bytes.bytesize) - offset
+        @event_bytes += length
+        raise ToolCallError, "SSE event exceeds max_event_bytes (#{@maximum})" if @event_bytes > @maximum
+        @buffer << bytes.byteslice(offset, length)
+        break unless index
+
+        line = @buffer
+        @buffer = +"".b
+        if @first_line
+          line = line.delete_prefix("\xEF\xBB\xBF".b)
+          @first_line = false
+        end
+        @skip_lf = bytes.getbyte(index) == 13
+        offset = index + 1
+        process_line(line.force_encoding(Encoding::UTF_8)) { |event| yield event }
       end
     end
 
     def finish
-      process_line(@buffer) { |event| yield event } unless @buffer.empty?
-      dispatch { |event| yield event } unless @fields[:data].empty?
+      # EOF is not an event delimiter; discard an unfinished event.
       @buffer.clear
+      @fields = reset_fields
+      @event_bytes = 0
     end
 
     private
 
     def process_line(line)
       if line.empty?
+        @event_bytes = 0
         dispatch { |event| yield event }
         return
       end
@@ -79,7 +113,9 @@ module UTCP
       field, value = line.split(":", 2)
       value = value.to_s.sub(/\A /, "")
       case field
-      when "data" then @fields[:data] << value
+      when "data"
+        @fields[:has_data] = true
+        @fields[:data] << value << "\n"
       when "event" then @fields[:event] = value
       when "id" then @fields[:id] = value unless value.include?("\0")
       when "retry" then @fields[:retry] = Integer(value) rescue nil
@@ -89,17 +125,20 @@ module UTCP
     def dispatch
       fields = @fields
       @fields = reset_fields
-      return if fields[:data].empty?
+      return unless fields[:has_data]
       return if @event_type && fields[:event] != @event_type
 
-      payload = fields[:data].join("\n")
-      yield JSON.parse(payload)
-    rescue JSON::ParserError
-      yield payload
+      payload = fields[:data].delete_suffix("\n")
+      value = begin
+        JSON.parse(payload)
+      rescue JSON::ParserError
+        payload
+      end
+      yield value
     end
 
     def reset_fields
-      { data: [], event: nil, id: nil, retry: nil }
+      { data: +"", has_data: false, event: nil, id: nil, retry: nil }
     end
   end
   SseCommunicationProtocol = SSEProtocol

data/lib/utcp/protocols/streamable_http.rb

@@ -16,9 +16,11 @@ module UTCP
     def call_tool(client, tool_name, tool_args, template)
       chunks = []
       binary = false
-      call_tool_streaming(client, tool_name, tool_args, template) do |chunk|
-        binary ||= chunk.is_a?(String) && chunk.encoding == Encoding::BINARY
-        chunks << chunk
+      with_collection_timeout(template) do
+        call_tool_streaming(client, tool_name, tool_args, template) do |chunk|
+          binary ||= chunk.is_a?(String) && chunk.encoding == Encoding::BINARY
+          chunks << chunk
+        end
       end
       binary && chunks.all? { |chunk| chunk.is_a?(String) } ? chunks.join.b : chunks
     end
@@ -27,21 +29,32 @@ module UTCP
       return enum_for(__method__, _client, tool_name, tool_args, template) unless block_given?
 
       assert_stream_template!(template)
+      count = 0
+      emit = lambda do |item|
+        count += 1
+        raise ToolCallError, "Streamable HTTP response exceeds max_response_items" if count > template.max_response_items
+        yield item
+      end
       with_stream_response(template, tool_args || {}) do |response|
         content_type = response["content-type"].to_s.downcase
         if content_type.include?("application/x-ndjson") || content_type.include?("application/json-seq")
-          stream_json_lines(response) { |item| yield item }
+          separator = content_type.include?("application/json-seq") ? "\x1E" : "\n"
+          stream_json_lines(response, template.max_event_bytes, separator, &emit)
         elsif content_type.include?("application/json")
           body = +""
-          response.read_body { |chunk| body << chunk }
-          yield decode_json_or_text(body) unless body.empty?
+          response.read_body do |chunk|
+            check_event_size!(body.bytesize + chunk.bytesize, template.max_event_bytes)
+            body << chunk
+          end
+          emit.call(decode_json_or_text(body)) unless body.empty?
         else
           response.read_body do |chunk|
             bytes = chunk.to_s.b
             offset = 0
             while offset < bytes.bytesize
-              yield bytes.byteslice(offset, template.chunk_size)
-              offset += template.chunk_size
+              size = [template.chunk_size, template.max_event_bytes].min
+              emit.call(bytes.byteslice(offset, size))
+              offset += size
             end
           end
         end
@@ -60,17 +73,30 @@ module UTCP
       raise ValidationError, "streamable HTTP protocol requires a StreamableHttpCallTemplate"
     end
 
-    def stream_json_lines(response)
-      buffer = +""
+    def stream_json_lines(response, maximum, separator = "\n")
+      buffer = +"".b
       response.read_body do |chunk|
-        buffer << chunk.to_s
-        while (index = buffer.index("\n"))
-          line = buffer.slice!(0..index).strip.sub(/\A\x1E/, "")
-          yield decode_json_or_text(line) unless line.empty?
+        bytes = chunk.to_s.b
+        offset = 0
+        while offset < bytes.bytesize
+          index = bytes.index(separator, offset)
+          length = (index || bytes.bytesize) - offset
+          check_event_size!(buffer.bytesize + length, maximum)
+          buffer << bytes.byteslice(offset, length)
+          break unless index
+
+          line = buffer.strip
+          buffer = +"".b
+          yield decode_json_or_text(line.force_encoding(Encoding::UTF_8)) unless line.empty?
+          offset = index + 1
         end
       end
-      tail = buffer.strip.sub(/\A\x1E/, "")
-      yield decode_json_or_text(tail) unless tail.empty?
+      tail = buffer.strip
+      yield decode_json_or_text(tail.force_encoding(Encoding::UTF_8)) unless tail.empty?
+    end
+
+    def check_event_size!(size, maximum)
+      raise ToolCallError, "Streamable HTTP event exceeds max_event_bytes (#{maximum})" if size > maximum
     end
   end
   StreamableHttpCommunicationProtocol = StreamableHTTPProtocol

data/lib/utcp/protocols/tcp.rb

@@ -71,6 +71,7 @@ module UTCP
 
     def read_framed(socket, template, timeout)
       deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout
+      maximum = [template.max_response_size, template.max_response_bytes].min
       case template.framing_strategy
       when "length_prefix"
         prefix = read_exact(socket, template.length_prefix_bytes, deadline)
@@ -81,16 +82,19 @@ module UTCP
           [8, "big"] => "Q>", [8, "little"] => "Q<"
         }.fetch([template.length_prefix_bytes, template.length_prefix_endian])
         length = prefix.unpack1(unpack)
-        raise ToolCallError, "TCP response exceeds max_response_size" if length > template.max_response_size
+        raise ToolCallError, "TCP response exceeds max_response_bytes or max_response_size" if length > maximum
 
         read_exact(socket, length, deadline)
       when "delimiter"
         read_until(socket, escaped_delimiter(template.message_delimiter, template.interpret_escape_sequences),
-                   template.max_response_size, deadline)
+                   maximum, deadline)
       when "fixed_length"
+        if template.fixed_message_length > maximum
+          raise ToolCallError, "TCP response exceeds max_response_bytes or max_response_size"
+        end
         read_exact(socket, template.fixed_message_length, deadline)
       when "stream"
-        read_stream(socket, template.max_response_size, deadline)
+        read_stream(socket, maximum, deadline)
       end
     end
 
@@ -112,25 +116,33 @@ module UTCP
       raise ValidationError, "message_delimiter cannot be empty" if delimiter.empty?
 
       result = +"".b
-      until result.end_with?(delimiter)
-        raise ToolCallError, "TCP response exceeds max_response_size" if result.bytesize >= maximum
+      loop do
+        if (index = result.index(delimiter))
+          raise ToolCallError, "TCP response exceeds max_response_bytes or max_response_size" if index > maximum
+          return result.byteslice(0, index)
+        end
+        if result.bytesize >= maximum + delimiter.bytesize
+          raise ToolCallError, "TCP response exceeds max_response_bytes or max_response_size"
+        end
 
         wait_readable!(socket, deadline, "TCP read")
-        result << socket.readpartial([4096, maximum - result.bytesize].min)
+        result << socket.readpartial([4096, maximum + delimiter.bytesize - result.bytesize].min)
       end
-      result.byteslice(0, result.bytesize - delimiter.bytesize)
     rescue EOFError
       raise ToolCallError, "TCP connection closed before the message delimiter"
     end
 
     def read_stream(socket, maximum, deadline)
       result = +"".b
-      while result.bytesize < maximum
+      loop do
         remaining = deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC)
         break unless remaining.positive? && IO.select([socket], nil, nil, remaining)
 
         begin
-          result << socket.readpartial([4096, maximum - result.bytesize].min)
+          result << socket.readpartial([4096, maximum + 1 - result.bytesize].min)
+          if result.bytesize > maximum
+            raise ToolCallError, "TCP response exceeds max_response_bytes or max_response_size"
+          end
         rescue EOFError
           break
         end

data/lib/utcp/protocols/udp.rb

@@ -42,9 +42,11 @@ module UTCP
       return nil if response_count.zero?
 
       deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + socket_timeout_seconds(template)
+      budget = ResponseByteBudget.new(template.max_response_bytes, "UDP")
       values = response_count.times.map do
         wait_readable!(socket, deadline, "UDP read")
         payload = socket.recv(65_535)
+        budget.consume(payload)
         decode_socket_payload(payload, template.response_byte_format)
       end
       values.length == 1 ? values.first : values