ruby-pi 0.1.8 → 0.1.10

data/lib/ruby_pi/llm/anthropic.rb

@@ -340,7 +340,7 @@ module RubyPi
         end
 
         handle_error_response(response) unless response.success?
-        parse_response(JSON.parse(response.body))
+        parse_response(parse_json_response(response.body))
       end
 
       # Executes a streaming request to the Anthropic API, yielding events.
@@ -366,26 +366,28 @@ 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 — not after the entire
-        # response has been buffered.
-        #
-        # The buffer is BINARY because 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 that already holds non-ASCII text raises
-        # Encoding::CompatibilityError. Each complete line is re-encoded to
-        # UTF-8 (and scrubbed) before parsing, so deltas reach the caller as
-        # valid UTF-8 strings.
-        sse_buffer = (+"").force_encoding(Encoding::BINARY)
+        sse_parser = SSEParser.new(provider: provider_name)
+        stream_completed = false
         response_status = nil
 
         # Accumulate error response body separately so ApiError gets the
         # full body even though on_data consumed the chunks.
         error_body = (+"").force_encoding(Encoding::BINARY)
 
+        process_payload = proc do |data_str|
+          next if data_str == "[DONE]"
+
+          data = sse_parser.parse_json(data_str)
+          stream_completed = true if data["type"] == "message_stop"
+          stream_state = process_anthropic_stream_event(
+            data, accumulated_text, accumulated_tool_calls,
+            current_tool_call, current_tool_json, usage_data, finish_reason, block
+          )
+          current_tool_call = stream_state[:current_tool_call]
+          current_tool_json = stream_state[:current_tool_json]
+          finish_reason = stream_state[:finish_reason]
+        end
+
         response = with_transport_errors do
           conn.post("/v1/messages") do |req|
             req.headers["Content-Type"] = "application/json"
@@ -395,7 +397,7 @@ module RubyPi
           # Without this, Faraday buffers the entire response body before
           # returning, which means no deltas reach the caller until the model
           # finishes generating (fake streaming).
-          req.options.on_data = proc do |chunk, overall_received_bytes, env|
+          req.options.on_data = proc do |chunk, _overall_received_bytes, env|
             response_status ||= env&.status
 
             # If the HTTP status indicates an error, accumulate the body for
@@ -403,41 +405,11 @@ module RubyPi
             # calls on_data for error responses too, which would otherwise
             # consume the body and leave response.body empty.
             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 each SSE event exactly as before ---
-              # Process the SSE event and update mutable locals from the
-              # returned hash. This keeps all streaming state method-local,
-              # avoiding thread-unsafe instance variables.
-              stream_state = process_anthropic_stream_event(
-                data, accumulated_text, accumulated_tool_calls,
-                current_tool_call, current_tool_json, usage_data, finish_reason, block
-              )
-              current_tool_call = stream_state[:current_tool_call]
-              current_tool_json = stream_state[:current_tool_json]
-              finish_reason = stream_state[:finish_reason]
-            end
+            sse_parser.feed(chunk, &process_payload)
           end
           end # conn.post
         end # with_transport_errors
@@ -452,25 +424,13 @@ module RubyPi
           handle_error_response(error_response, override_body: error_body_str)
         end
 
-        # Process any remaining data in the buffer after the connection closes
-        sse_buffer.force_encoding(Encoding::UTF_8).scrub.each_line do |line|
-          line = line.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
-          stream_state = process_anthropic_stream_event(
-            data, accumulated_text, accumulated_tool_calls,
-            current_tool_call, current_tool_json, usage_data, finish_reason, block
+        sse_parser.finish(&process_payload)
+
+        unless stream_completed
+          raise RubyPi::StreamingProtocolError.new(
+            "anthropic stream ended before message_stop",
+            status_code: nil
           )
-          current_tool_call = stream_state[:current_tool_call]
-          current_tool_json = stream_state[:current_tool_json]
-          finish_reason = stream_state[:finish_reason]
         end
 
         # (Event processing is now handled incrementally by the on_data callback
@@ -532,11 +492,14 @@ module RubyPi
           delta = data["delta"] || {}
           if delta["type"] == "text_delta"
             text = delta["text"] || ""
-            accumulated_text << text
+            append_stream_data!(accumulated_text, text, limit: MAX_STREAM_OUTPUT_BYTES, label: "stream output")
             block.call(StreamEvent.new(type: :text_delta, data: text))
           elsif delta["type"] == "input_json_delta"
             json_chunk = delta["partial_json"] || ""
-            current_tool_json << json_chunk
+            append_stream_data!(
+              current_tool_json, json_chunk,
+              limit: MAX_STREAM_TOOL_ARGUMENT_BYTES, label: "tool arguments"
+            )
             block.call(StreamEvent.new(type: :tool_call_delta, data: {
               id: current_tool_call&.dig(:id),
               partial_json: json_chunk
@@ -554,11 +517,10 @@ module RubyPi
                         else
                           begin
                             JSON.parse(current_tool_json)
-                          rescue JSON::ParserError => e
+                          rescue JSON::ParserError
                             raise RubyPi::ProviderError.new(
                               "Failed to parse streaming tool call arguments for " \
-                              "'#{current_tool_call[:name]}': #{e.message} " \
-                              "(accumulated JSON: #{current_tool_json.inspect})",
+                              "'#{current_tool_call[:name]}'",
                               provider: :anthropic
                             )
                           end
@@ -568,6 +530,7 @@ module RubyPi
               name: current_tool_call[:name],
               arguments: arguments
             )
+            enforce_stream_tool_call_limit!(accumulated_tool_calls.size)
             current_tool_call = nil
             current_tool_json = +""
           end

data/lib/ruby_pi/llm/base_provider.rb

@@ -7,6 +7,9 @@
 # providers (Gemini, Anthropic, OpenAI) must subclass this and implement the
 # abstract methods.
 
+require "json"
+require "time"
+
 module RubyPi
   module LLM
     # Abstract base class that defines the contract every LLM provider must
@@ -30,6 +33,11 @@ module RubyPi
     #     end
     #   end
     class BaseProvider
+      MAX_STREAM_OUTPUT_BYTES = 16 * 1024 * 1024
+      MAX_STREAM_TOOL_ARGUMENT_BYTES = 1024 * 1024
+      MAX_STREAM_TOOL_CALLS = 128
+      MAX_ERROR_BODY_BYTES = 64 * 1024
+      RETRY_AFTER_CEILING = 60.0
       # @return [Integer] maximum number of retry attempts
       attr_reader :max_retries
 
@@ -53,6 +61,8 @@ module RubyPi
         @max_retries = max_retries || @config.max_retries
         @retry_base_delay = retry_base_delay || @config.retry_base_delay
         @retry_max_delay = retry_max_delay || @config.retry_max_delay
+
+        validate_retry_configuration!
       end
 
       # Sends a completion request to the LLM provider with automatic retry
@@ -71,10 +81,22 @@ module RubyPi
       # @raise [RubyPi::TimeoutError] on request timeouts
       def complete(messages:, tools: [], stream: false, &block)
         attempt = 0
+        partial_chars = 0
+        partial_tool_calls = false
+
+        attempt_block = if stream && block
+                          proc do |event|
+                            partial_chars += event.data.to_s.length if event.text_delta?
+                            partial_tool_calls = true if event.tool_call_delta?
+                            block.call(event)
+                          end
+                        else
+                          block
+                        end
 
         begin
           attempt += 1
-          perform_complete(messages: messages, tools: tools, stream: stream, &block)
+          perform_complete(messages: messages, tools: tools, stream: stream, &attempt_block)
         rescue RubyPi::AuthenticationError
           # Authentication errors are not retryable — raise immediately
           raise
@@ -93,9 +115,23 @@ module RubyPi
           # `attempt <= @max_retries` allows retries on attempts 1..3, so we get
           # 3 retries + 1 initial = 4 total attempts. Previously used `< @max_retries`
           # which was off-by-one (only 2 retries with max_retries: 3).
-          if attempt <= @max_retries
+          if retryable_error?(e) && attempt <= @max_retries
             delay = retry_delay_for(e, attempt)
             log_retry(attempt, delay, e)
+
+            if stream && block
+              block.call(StreamEvent.new(type: :retry_start, data: {
+                provider: provider_name,
+                error: e.message,
+                attempt: attempt + 1,
+                partial_output: partial_chars.positive? || partial_tool_calls,
+                partial_chars: partial_chars,
+                partial_tool_calls: partial_tool_calls
+              }))
+              partial_chars = 0
+              partial_tool_calls = false
+            end
+
             sleep(delay)
             retry
           else
@@ -124,6 +160,31 @@ module RubyPi
 
       private
 
+      def append_stream_data!(buffer, chunk, limit:, label:)
+        if buffer.bytesize + chunk.to_s.bytesize > limit
+          raise RubyPi::StreamingProtocolError.new(
+            "#{provider_name} #{label} exceeded #{limit} bytes",
+            status_code: nil
+          )
+        end
+
+        buffer << chunk
+      end
+
+      def append_error_body(error_body, chunk)
+        remaining = MAX_ERROR_BODY_BYTES - error_body.bytesize
+        error_body << chunk.to_s.b.byteslice(0, remaining) if remaining.positive?
+      end
+
+      def enforce_stream_tool_call_limit!(count)
+        return if count <= MAX_STREAM_TOOL_CALLS
+
+        raise RubyPi::StreamingProtocolError.new(
+          "#{provider_name} stream exceeded #{MAX_STREAM_TOOL_CALLS} tool calls",
+          status_code: nil
+        )
+      end
+
       # Performs the actual completion request. Subclasses MUST implement this
       # method with provider-specific HTTP logic.
       #
@@ -132,15 +193,10 @@ module RubyPi
       # @param stream [Boolean] streaming mode flag
       # @yield [event] optional block for streaming events
       # @return [RubyPi::LLM::Response]
-      def perform_complete(messages:, tools:, stream:, &block)
+      def perform_complete(...)
         raise RubyPi::AbstractMethodError, :perform_complete
       end
 
-      # Maximum delay (seconds) honored from a server-provided Retry-After
-      # header. Caps pathological or misconfigured server values so a single
-      # 429 cannot stall the client indefinitely.
-      RETRY_AFTER_CEILING = 60.0
-
       # Picks the delay before the next retry. A server-provided Retry-After
       # on a 429 takes precedence over the local exponential backoff: the
       # server knows its own cooldown window, and retrying earlier just burns
@@ -159,6 +215,17 @@ module RubyPi
         end
       end
 
+      # Only transient failures should be retried. Deterministic client errors
+      # (400, 404, 422, etc.) cannot improve on an identical retry and merely
+      # waste quota and add the full backoff delay.
+      def retryable_error?(error)
+        return true if error.is_a?(RubyPi::RateLimitError) || error.is_a?(RubyPi::TimeoutError)
+        return false unless error.is_a?(RubyPi::ApiError)
+
+        status = error.status_code
+        status.nil? || status == 408 || status == 409 || status == 425 || status >= 500
+      end
+
       # Calculates the backoff delay for a given retry attempt using
       # exponential backoff with jitter.
       #
@@ -170,6 +237,27 @@ module RubyPi
         [base + jitter, @retry_max_delay].min
       end
 
+      # Validates constructor-level overrides. Configuration's writers perform
+      # the same checks, but provider keyword overrides bypass those writers.
+      # In particular, an infinite max_retries value makes the retry loop run
+      # forever because every finite attempt is <= Infinity.
+      def validate_retry_configuration!
+        unless @max_retries.is_a?(Integer) && @max_retries >= 0
+          raise ArgumentError,
+                "max_retries must be a non-negative integer, got #{@max_retries.inspect}"
+        end
+
+        {
+          retry_base_delay: @retry_base_delay,
+          retry_max_delay: @retry_max_delay
+        }.each do |name, value|
+          next if value.is_a?(Numeric) && !value.is_a?(Complex) && value.finite? && value >= 0
+
+          raise ArgumentError,
+                "#{name} must be a finite non-negative number, got #{value.inspect}"
+        end
+      end
+
       # Logs a retry attempt if a logger is configured.
       #
       # @param attempt [Integer] current attempt number
@@ -269,7 +357,7 @@ module RubyPi
             response_body: body
           )
         when 429
-          retry_after = response.headers["retry-after"]&.to_f
+          retry_after = parse_retry_after(response.headers["retry-after"])
           raise RubyPi::RateLimitError.new(
             "#{provider_name} rate limit exceeded (HTTP 429)",
             retry_after: retry_after,
@@ -284,6 +372,36 @@ module RubyPi
         end
       end
 
+      # Parses a successful JSON response into a Hash while preserving the
+      # provider error contract. A proxy can truncate a nominal HTTP 200 body;
+      # treating that as a transient ApiError allows retry/fallback instead of
+      # leaking JSON::ParserError or accepting an empty response.
+      def parse_json_response(body)
+        JSON.parse(body)
+      rescue JSON::ParserError => e
+        raise RubyPi::ApiError.new(
+          "#{provider_name} returned malformed JSON: #{e.message}",
+          status_code: nil,
+          response_body: body
+        )
+      end
+
+      # Parses both Retry-After forms allowed by HTTP: delta-seconds and an
+      # HTTP-date. Fractional seconds are retained for compatibility with API
+      # providers and fast test environments. Invalid or past dates return nil
+      # so the normal exponential backoff policy applies.
+      def parse_retry_after(value)
+        return nil if value.nil? || value.to_s.strip.empty?
+
+        seconds = Float(value, exception: false)
+        return seconds if seconds&.finite? && seconds.positive?
+
+        delay = Time.httpdate(value.to_s) - Time.now
+        delay.positive? ? delay : nil
+      rescue ArgumentError
+        nil
+      end
+
     end
   end
 end

data/lib/ruby_pi/llm/fallback.rb

@@ -157,8 +157,18 @@ module RubyPi
         # carries partial_output/partial_chars so consumers can deterministically
         # truncate what they already rendered.
         partial_chars = 0
+        partial_tool_calls = false
         counting_block = proc do |event|
-          partial_chars += event.data.to_s.length if event.text_delta?
+          if event.retry_start?
+            # The provider tells consumers to discard its previous attempt,
+            # so only count text from the current attempt for a later fallback.
+            partial_chars = 0
+            partial_tool_calls = false
+          elsif event.text_delta?
+            partial_chars += event.data.to_s.length
+          elsif event.tool_call_delta?
+            partial_tool_calls = true
+          end
           block.call(event)
         end
 
@@ -187,10 +197,11 @@ module RubyPi
           #                    this amount if appending to a shared buffer)
           block.call(StreamEvent.new(type: :fallback_start, data: {
             failed_provider: @primary.provider_name,
-            error: e.message,
+            error: safe_error_message(e),
             fallback_provider: @fallback.provider_name,
-            partial_output: partial_chars.positive?,
-            partial_chars: partial_chars
+            partial_output: partial_chars.positive? || partial_tool_calls,
+            partial_chars: partial_chars,
+            partial_tool_calls: partial_tool_calls
           }))
 
           # Stream directly from the fallback to the consumer's block.
@@ -213,10 +224,14 @@ module RubyPi
 
         logger.warn(
           "[RubyPi::Fallback] Primary provider (#{@primary.provider_name}/#{@primary.model_name}) " \
-          "failed with #{error.class}: #{error.message}. " \
+          "failed with #{error.class}: #{safe_error_message(error)}. " \
           "Falling back to #{@fallback.provider_name}/#{@fallback.model_name}."
         )
       end
+
+      def safe_error_message(error)
+        error.message.to_s.gsub(/[[:cntrl:]]+/, " ").slice(0, 512)
+      end
     end
   end
 end

data/lib/ruby_pi/llm/gemini.rb

@@ -281,7 +281,7 @@ module RubyPi
         end
 
         handle_error_response(response) unless response.success?
-        parse_response(JSON.parse(response.body))
+        parse_response(parse_json_response(response.body))
       end
 
       # Executes a streaming request to the Gemini API, yielding events.
@@ -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