ruby-pi 0.1.9 → 0.1.10

data/lib/ruby_pi/llm/gemini.rb

@@ -301,20 +301,59 @@ module RubyPi
         usage_data = {}
         finish_reason = nil
 
-        # Buffer for incomplete SSE lines across on_data chunks. Faraday's
-        # on_data callback delivers raw bytes as they arrive from the network,
-        # which may split SSE events mid-line. We accumulate a line buffer and
-        # process complete lines incrementally so that deltas reach the caller
-        # as soon as each SSE event is fully received.
-        # BINARY buffer: chunks arrive as ASCII-8BIT and may end mid-way
-        # through a multi-byte UTF-8 character; appending such a chunk to a
-        # UTF-8 buffer holding non-ASCII text raises
-        # Encoding::CompatibilityError. Complete lines are re-encoded to
-        # UTF-8 (and scrubbed) before parsing.
-        sse_buffer = (+"").force_encoding(Encoding::BINARY)
+        sse_parser = SSEParser.new(provider: provider_name)
+        stream_completed = false
         response_status = nil
         error_body = (+"").force_encoding(Encoding::BINARY)
 
+        process_payload = proc do |data_str|
+          next if data_str == "[DONE]"
+
+          data = sse_parser.parse_json(data_str)
+          block_reason = data.dig("promptFeedback", "blockReason")
+          if block_reason && !block_reason.to_s.empty?
+            stream_completed = true
+            finish_reason = block_reason.to_s.downcase
+          end
+
+          candidates = data["candidates"] || []
+          candidate = candidates.first
+          next unless candidate
+
+          parts = candidate.dig("content", "parts") || []
+          parts.each do |part|
+            if part.key?("text")
+              text_chunk = part["text"]
+              append_stream_data!(accumulated_text, text_chunk, limit: MAX_STREAM_OUTPUT_BYTES, label: "stream output")
+              block.call(StreamEvent.new(type: :text_delta, data: text_chunk))
+            elsif part.key?("functionCall")
+              fc = part["functionCall"]
+              tool_call = ToolCall.new(
+                id: "gemini_#{SecureRandom.hex(8)}",
+                name: fc["name"],
+                arguments: fc["args"] || {}
+              )
+              accumulated_tool_calls << tool_call
+              enforce_stream_tool_call_limit!(accumulated_tool_calls.size)
+              block.call(StreamEvent.new(type: :tool_call_delta, data: tool_call.to_h))
+            end
+          end
+
+          if candidate["finishReason"] && !candidate["finishReason"].to_s.empty?
+            finish_reason = candidate["finishReason"].to_s.downcase
+            stream_completed = true
+          end
+
+          if data.key?("usageMetadata")
+            meta = data["usageMetadata"]
+            usage_data = {
+              prompt_tokens: meta["promptTokenCount"],
+              completion_tokens: meta["candidatesTokenCount"],
+              total_tokens: meta["totalTokenCount"]
+            }
+          end
+        end
+
         response = with_transport_errors do
           conn.post(url) do |req|
             req.headers["Content-Type"] = "application/json"
@@ -330,77 +369,11 @@ module RubyPi
             # If the HTTP status indicates an error, accumulate the body for
             # the error handler instead of parsing it as SSE events.
             if response_status && response_status >= 400
-              error_body << chunk.b
+              append_error_body(error_body, chunk)
               next
             end
 
-            sse_buffer << chunk.b
-            # Process all complete lines in the buffer. A complete line holds
-            # complete UTF-8 sequences (multi-byte characters split across
-            # chunks are repaired by the buffering), so re-encode it to UTF-8
-            # here; scrub guards against a server sending invalid bytes.
-            while (line_end = sse_buffer.index("\n"))
-              line = sse_buffer.slice!(0, line_end + 1).force_encoding(Encoding::UTF_8).scrub.strip
-              next if line.empty?
-              next unless line.start_with?("data: ")
-
-              data_str = line.sub(/\Adata: /, "")
-              next if data_str == "[DONE]"
-
-              begin
-                data = JSON.parse(data_str)
-              rescue JSON::ParserError
-                next
-              end
-
-              # Process this SSE event
-              candidates = data.dig("candidates") || []
-              candidate = candidates.first
-              next unless candidate
-
-              parts = candidate.dig("content", "parts") || []
-              parts.each do |part|
-                if part.key?("text")
-                  text_chunk = part["text"]
-                  accumulated_text << text_chunk
-                  block.call(StreamEvent.new(type: :text_delta, data: text_chunk))
-                elsif part.key?("functionCall")
-                  fc = part["functionCall"]
-                  tool_call = ToolCall.new(
-                    # Generate a globally-unique ID per tool call. A simple
-                    # length-based counter ("gemini_0", "gemini_1") collides
-                    # across turns since each response restarts numbering at
-                    # 0, breaking any caller that uses ID as a hash key for
-                    # observability or result correlation.
-                    id: "gemini_#{SecureRandom.hex(8)}",
-                    name: fc["name"],
-                    arguments: fc["args"] || {}
-                  )
-                  accumulated_tool_calls << tool_call
-                  block.call(StreamEvent.new(type: :tool_call_delta, data: tool_call.to_h))
-                end
-              end
-
-              # Parse the actual finish reason from the streaming response
-              # instead of hardcoding "stop". Gemini sends finishReason in
-              # the candidate object (e.g., "STOP", "MAX_TOKENS", "SAFETY").
-              # Coerce via to_s before downcase so a non-String payload can
-              # never raise NoMethodError mid-stream (mirrors the &.to_s in
-              # the non-streaming parse path).
-              if candidate["finishReason"]
-                finish_reason = candidate["finishReason"].to_s.downcase
-              end
-
-              # Capture usage metadata if present
-              if data.key?("usageMetadata")
-                meta = data["usageMetadata"]
-                usage_data = {
-                  prompt_tokens: meta["promptTokenCount"],
-                  completion_tokens: meta["candidatesTokenCount"],
-                  total_tokens: meta["totalTokenCount"]
-                }
-              end
-            end
+            sse_parser.feed(chunk, &process_payload)
           end
           end # conn.post
         end # with_transport_errors
@@ -413,6 +386,15 @@ module RubyPi
           handle_error_response(response, override_body: error_body_str)
         end
 
+        sse_parser.finish(&process_payload)
+
+        unless stream_completed
+          raise RubyPi::StreamingProtocolError.new(
+            "gemini stream ended before a finish reason",
+            status_code: nil
+          )
+        end
+
         # Signal completion
         block.call(StreamEvent.new(type: :done))
 
@@ -420,7 +402,7 @@ module RubyPi
           content: accumulated_text.empty? ? nil : accumulated_text,
           tool_calls: accumulated_tool_calls,
           usage: usage_data,
-          finish_reason: finish_reason || "stop"
+          finish_reason: finish_reason
         )
       end
 

data/lib/ruby_pi/llm/openai.rb

@@ -201,11 +201,10 @@ module RubyPi
                               begin
                                 JSON.parse(tc_args)
                                 tc_args
-                              rescue JSON::ParserError => e
+                              rescue JSON::ParserError
                                 raise RubyPi::ProviderError.new(
                                   "Invalid JSON in assistant tool_call.arguments " \
-                                  "for tool '#{tc_name || "unknown"}': #{e.message} " \
-                                  "(raw: #{tc_args.inspect})",
+                                  "for tool '#{tc_name || "unknown"}'",
                                   provider: :openai
                                 )
                               end
@@ -315,20 +314,63 @@ module RubyPi
         # Issue #21: Accumulate usage data from the final SSE chunk
         streaming_usage = {}
 
-        # Buffer for incomplete SSE lines across on_data chunks. Faraday's
-        # on_data callback delivers raw bytes as they arrive from the network,
-        # which may split SSE events mid-line. We accumulate a line buffer and
-        # process complete lines incrementally so that deltas reach the caller
-        # as soon as each SSE event is fully received.
-        # BINARY buffer: chunks arrive as ASCII-8BIT and may end mid-way
-        # through a multi-byte UTF-8 character; appending such a chunk to a
-        # UTF-8 buffer holding non-ASCII text raises
-        # Encoding::CompatibilityError. Complete lines are re-encoded to
-        # UTF-8 (and scrubbed) before parsing.
-        sse_buffer = (+"").force_encoding(Encoding::BINARY)
+        sse_parser = SSEParser.new(provider: provider_name)
+        stream_completed = false
         response_status = nil
         error_body = (+"").force_encoding(Encoding::BINARY)
 
+        process_payload = proc do |data_str|
+          if data_str == "[DONE]"
+            stream_completed = true
+            next
+          end
+
+          data = sse_parser.parse_json(data_str)
+          if data.key?("usage") && data["usage"]
+            usage_info = data["usage"]
+            streaming_usage = {
+              prompt_tokens: usage_info["prompt_tokens"],
+              completion_tokens: usage_info["completion_tokens"],
+              total_tokens: usage_info["total_tokens"]
+            }
+          end
+
+          choice = (data["choices"] || []).first
+          next unless choice
+
+          delta = choice["delta"] || {}
+          finish_reason = choice["finish_reason"] if choice["finish_reason"]
+
+          if delta.key?("content") && delta["content"]
+            text = delta["content"]
+            append_stream_data!(accumulated_text, text, limit: MAX_STREAM_OUTPUT_BYTES, label: "stream output")
+            block.call(StreamEvent.new(type: :text_delta, data: text))
+          end
+
+          (delta["tool_calls"] || []).each do |tc_delta|
+            index = tc_delta["index"] || 0
+            tool_call_accumulators[index] ||= { id: nil, name: +"", arguments: +"" }
+            enforce_stream_tool_call_limit!(tool_call_accumulators.size)
+            acc = tool_call_accumulators[index]
+            acc[:id] = tc_delta["id"] if tc_delta["id"]
+            acc[:name] << tc_delta["function"]["name"] if tc_delta.dig("function", "name")
+
+            if tc_delta.dig("function", "arguments")
+              append_stream_data!(
+                acc[:arguments], tc_delta["function"]["arguments"],
+                limit: MAX_STREAM_TOOL_ARGUMENT_BYTES, label: "tool arguments"
+              )
+            end
+
+            block.call(StreamEvent.new(type: :tool_call_delta, data: {
+              index: index,
+              id: acc[:id],
+              name: acc[:name].dup.freeze,
+              arguments_fragment: tc_delta.dig("function", "arguments") || ""
+            }))
+          end
+        end
+
         response = with_transport_errors do
           conn.post("/v1/chat/completions") do |req|
             req.headers["Content-Type"] = "application/json"
@@ -344,83 +386,11 @@ module RubyPi
             # If the HTTP status indicates an error, accumulate the body for
             # the error handler instead of parsing it as SSE events.
             if response_status && response_status >= 400
-              error_body << chunk.b
+              append_error_body(error_body, chunk)
               next
             end
 
-            sse_buffer << chunk.b
-            # Process all complete lines in the buffer. A complete line holds
-            # complete UTF-8 sequences (multi-byte characters split across
-            # chunks are repaired by the buffering), so re-encode it to UTF-8
-            # here; scrub guards against a server sending invalid bytes.
-            while (line_end = sse_buffer.index("\n"))
-              line = sse_buffer.slice!(0, line_end + 1).force_encoding(Encoding::UTF_8).scrub.strip
-              next if line.empty?
-              next unless line.start_with?("data: ")
-
-              data_str = line.sub(/\Adata: /, "")
-              next if data_str == "[DONE]"
-
-              begin
-                data = JSON.parse(data_str)
-              rescue JSON::ParserError
-                next
-              end
-
-              # Issue #21: Capture usage data from the final SSE chunk.
-              # OpenAI sends usage in a dedicated chunk when include_usage is true.
-              if data.key?("usage") && data["usage"]
-                usage_info = data["usage"]
-                streaming_usage = {
-                  prompt_tokens: usage_info["prompt_tokens"],
-                  completion_tokens: usage_info["completion_tokens"],
-                  total_tokens: usage_info["total_tokens"]
-                }
-              end
-
-              # Process this SSE event
-              choices = data["choices"] || []
-              choice = choices.first
-              next unless choice
-
-              delta = choice["delta"] || {}
-              finish_reason = choice["finish_reason"] if choice["finish_reason"]
-
-              # Handle text content deltas
-              if delta.key?("content") && delta["content"]
-                text = delta["content"]
-                accumulated_text << text
-                block.call(StreamEvent.new(type: :text_delta, data: text))
-              end
-
-              # Handle tool call deltas
-              if delta.key?("tool_calls")
-                delta["tool_calls"].each do |tc_delta|
-                  index = tc_delta["index"] || 0
-
-                  # Initialize accumulator for this tool call
-                  tool_call_accumulators[index] ||= { id: nil, name: +"", arguments: +"" }
-                  acc = tool_call_accumulators[index]
-
-                  acc[:id] = tc_delta["id"] if tc_delta["id"]
-
-                  if tc_delta.dig("function", "name")
-                    acc[:name] << tc_delta["function"]["name"]
-                  end
-
-                  if tc_delta.dig("function", "arguments")
-                    acc[:arguments] << tc_delta["function"]["arguments"]
-                  end
-
-                  block.call(StreamEvent.new(type: :tool_call_delta, data: {
-                    index: index,
-                    id: acc[:id],
-                    name: acc[:name],
-                    arguments_fragment: tc_delta.dig("function", "arguments") || ""
-                  }))
-                end
-              end
-            end
+            sse_parser.feed(chunk, &process_payload)
           end
           end # conn.post
         end # with_transport_errors
@@ -433,6 +403,15 @@ module RubyPi
           handle_error_response(response, override_body: error_body_str)
         end
 
+        sse_parser.finish(&process_payload)
+
+        unless stream_completed
+          raise RubyPi::StreamingProtocolError.new(
+            "openai stream ended before data: [DONE]",
+            status_code: nil
+          )
+        end
+
         # Build final tool calls from accumulators
         # Issue #12: Guard JSON.parse against empty strings. An empty string
         # is truthy in Ruby, so the previous `empty? ? {} : JSON.parse(...)` check
@@ -543,9 +522,9 @@ module RubyPi
 
         # Attempt to parse the JSON string
         JSON.parse(raw_args)
-      rescue JSON::ParserError => e
+      rescue JSON::ParserError
         raise RubyPi::ProviderError.new(
-          "Failed to parse tool call arguments from OpenAI: #{e.message} (raw: #{raw_args.inspect})",
+          "Failed to parse tool call arguments from OpenAI",
           provider: :openai
         )
       end

data/lib/ruby_pi/llm/sse_parser.rb

@@ -0,0 +1,125 @@
+# frozen_string_literal: true
+
+require "json"
+
+module RubyPi
+  module LLM
+    # Incremental, bounded Server-Sent Events decoder.
+    #
+    # Faraday delivers arbitrary byte chunks, including chunks split inside a
+    # UTF-8 character or line ending. This decoder keeps bytes in BINARY form
+    # until a complete SSE event is available, then validates UTF-8 and yields
+    # the joined value of its data fields.
+    class SSEParser
+      DEFAULT_MAX_EVENT_BYTES = 1_048_576
+
+      def initialize(provider:, max_event_bytes: DEFAULT_MAX_EVENT_BYTES)
+        @provider = provider
+        @max_event_bytes = max_event_bytes
+        @buffer = (+"").force_encoding(Encoding::BINARY)
+        @data_lines = []
+        @event_bytes = 0
+      end
+
+      def feed(chunk, &block)
+        @buffer << chunk.to_s.b
+        extract_lines(finishing: false, &block)
+        enforce_size!(@buffer.bytesize + @event_bytes)
+      end
+
+      def finish(&block)
+        extract_lines(finishing: true, &block)
+
+        unless @buffer.empty?
+          line = @buffer.slice!(0, @buffer.bytesize)
+          process_line(line, &block)
+        end
+
+        dispatch(&block)
+      end
+
+      def parse_json(payload)
+        JSON.parse(payload)
+      rescue JSON::ParserError => e
+        raise RubyPi::StreamingProtocolError.new(
+          "#{@provider} stream contained malformed JSON: #{e.message}",
+          status_code: nil
+        )
+      end
+
+      private
+
+      def extract_lines(finishing:, &block)
+        loop do
+          line_end = next_line_ending
+          break unless line_end
+
+          if @buffer.getbyte(line_end) == 13 && line_end == @buffer.bytesize - 1 && !finishing
+            break
+          end
+
+          terminator_size = if @buffer.getbyte(line_end) == 13 && @buffer.getbyte(line_end + 1) == 10
+                              2
+                            else
+                              1
+                            end
+          line = @buffer.slice!(0, line_end + terminator_size)
+          line = line.byteslice(0, line.bytesize - terminator_size)
+          process_line(line, &block)
+        end
+      end
+
+      def next_line_ending
+        cr = @buffer.index("\r")
+        lf = @buffer.index("\n")
+        [cr, lf].compact.min
+      end
+
+      def process_line(raw_line, &block)
+        line = raw_line.dup.force_encoding(Encoding::UTF_8)
+        unless line.valid_encoding?
+          raise RubyPi::StreamingProtocolError.new(
+            "#{@provider} stream contained invalid UTF-8",
+            status_code: nil
+          )
+        end
+
+        if line.empty?
+          dispatch(&block)
+          return
+        end
+
+        return if line.start_with?(":")
+
+        field, separator, value = line.partition(":")
+        return unless field == "data"
+
+        value = "" unless separator == ":"
+        value = value.byteslice(1..) if value.start_with?(" ")
+        separator_bytes = @data_lines.empty? ? 0 : 1
+        prospective_bytes = @event_bytes + separator_bytes + value.bytesize
+        enforce_size!(prospective_bytes)
+        @data_lines << value
+        @event_bytes = prospective_bytes
+      end
+
+      def dispatch
+        return if @data_lines.empty?
+
+        payload = @data_lines.join("\n")
+        @data_lines.clear
+        @event_bytes = 0
+        yield payload
+      end
+
+      def enforce_size!(bytes)
+        return if bytes <= @max_event_bytes
+
+        raise RubyPi::StreamingProtocolError.new(
+          "#{@provider} stream event exceeded #{@max_event_bytes} bytes",
+          status_code: nil
+        )
+      end
+    end
+  end
+end