mistri 0.4.1 → 0.6.0

data/lib/mistri/providers/fake.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "json"
+require "securerandom"
 
 module Mistri
   module Providers
@@ -26,10 +27,17 @@ module Mistri
 
       def initialize(turns: [], chunk_size: 12)
         @turns = turns.map { |turn| turn.transform_keys(&:to_sym) }
+        @prices_usage = @turns.all? do |turn|
+          !turn[:usage] || turn[:usage].cost.known?
+        end
         @chunk_size = [chunk_size, 1].max
         @requests = []
+        @tool_call_prefix = SecureRandom.uuid
+        @tool_call_sequence = 0
       end
 
+      def prices_usage? = @prices_usage
+
       def stream(messages: [], **options, &emit)
         # Snapshot the array: the loop appends replies to it in place.
         @requests << { messages: messages.dup, options: }
@@ -67,18 +75,40 @@ module Mistri
         kind == :text ? Content::Text.new(text:) : Content::Thinking.new(thinking: text)
       end
 
-      def stream_tool_call(spec, position, blocks, emit)
+      # Small scripted arguments parse eagerly on every chunk for precise UI
+      # tests. Real assemblers cap their refresh schedule under hostile fragmentation.
+      def stream_tool_call(spec, _position, blocks, emit)
         spec = spec.transform_keys(&:to_sym)
-        call = ToolCall.new(id: spec[:id] || "call_#{position + 1}", name: spec[:name],
-                            arguments: (spec[:arguments] || {}).transform_keys(&:to_s))
+        arguments = spec.key?(:arguments) ? spec[:arguments] : {}
+        id = spec.key?(:id) ? spec[:id] : next_tool_call_id
+        call = ToolCall.new(id:, name: spec[:name], arguments:, signature: spec[:signature],
+                            arguments_error: spec[:arguments_error],
+                            provider_call_id: spec[:provider_call_id])
         index = blocks.size
         emit_event(emit, :toolcall_start, blocks, content_index: index)
-        emit_event(emit, :toolcall_delta, blocks, content_index: index,
-                                                  delta: JSON.generate(call.arguments))
+        built = +""
+        JSON.generate(call.arguments).scan(/.{1,#{@chunk_size}}/m) do |chunk|
+          built << chunk
+          emit_event(emit, :toolcall_delta, blocks + [in_progress(call, built)],
+                     content_index: index, delta: chunk)
+        end
         blocks << call
         emit_event(emit, :toolcall_end, blocks, content_index: index, tool_call: call)
       end
 
+      def in_progress(call, json)
+        parsed = ToolArguments.freeze_partial(PartialJson.parse(json))
+        ToolCall.new(id: call.id, name: call.name,
+                     arguments: parsed, signature: call.signature,
+                     provider_call_id: call.provider_call_id,
+                     canonicalize: false)
+      end
+
+      def next_tool_call_id
+        @tool_call_sequence += 1
+        "fake_#{@tool_call_prefix}_#{@tool_call_sequence}"
+      end
+
       def finish(turn, blocks, emit)
         reason = turn[:stop_reason] ||
                  (blocks.any?(ToolCall) ? StopReason::TOOL_USE : StopReason::STOP)
@@ -90,7 +120,7 @@ module Mistri
       def finish_error(turn, blocks, emit)
         error = { "type" => turn.fetch(:error_type, "Error") }
         error["status"] = turn[:status] if turn[:status]
-        message = assemble(blocks, usage: Usage.zero, stop_reason: StopReason::ERROR,
+        message = assemble(blocks, usage: turn[:usage] || Usage.zero, stop_reason: StopReason::ERROR,
                                    error_message: turn[:error], error: error)
         emit&.call(Event.new(type: :error, reason: StopReason::ERROR, message:,
                              error_message: turn[:error]))

data/lib/mistri/providers/gemini/assembler.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "json"
+require "securerandom"
 
 module Mistri
   module Providers
@@ -13,46 +14,104 @@ module Mistri
       # Thought signatures ride on individual parts and are captured onto the
       # block they arrived with, verbatim, for replay.
       class Assembler
-        def initialize(model:)
+        def initialize(model:, catalog_pricing: true, service_tier: nil)
           @model = model
+          @catalog_pricing = catalog_pricing
+          @service_tier = service_tier
+          @pricing_at = Time.now
           @blocks = []
           @current = nil
-          @usage = Usage.zero
+          @usage = Usage.new
           @finish_reason = nil
+          @terminal = false
+        end
+
+        # Finish reasons that are the provider's verdict on the content
+        # itself: a retry of the same input meets the same filter, and a
+        # harness that re-rolls against a policy verdict is machinery for
+        # evading it, so these fail fast.
+        VERDICTS = %w[SAFETY RECITATION LANGUAGE BLOCKLIST PROHIBITED_CONTENT SPII
+                      IMAGE_SAFETY IMAGE_PROHIBITED_CONTENT IMAGE_RECITATION].freeze
+        # Missing replay state accuses the request rather than the generated
+        # content, but retrying the same history is equally deterministic.
+        INPUT_ERRORS = %w[MISSING_THOUGHT_SIGNATURE].freeze
+        # The model fumbled its own output (an invalid or runaway tool call,
+        # a missing image), or the API stopped for a reason it cannot name
+        # (OTHER is documented as "Unknown reason"). The input stands accused
+        # of nothing, so these error retryably; a regeneration usually lands.
+        FUMBLES = %w[MALFORMED_FUNCTION_CALL UNEXPECTED_TOOL_CALL TOO_MANY_TOOL_CALLS
+                     NO_IMAGE OTHER IMAGE_OTHER MALFORMED_RESPONSE
+                     FINISH_REASON_UNSPECIFIED].freeze
+
+        def start(&emit)
+          emit&.call(Event.new(type: :start, partial: assemble))
         end
 
         def feed(record, &)
-          if (error = record["error"])
-            @error = ProviderError.new(error["message"] || "provider error",
-                                       status: error["code"])
-            return
-          end
+          return after_terminal(record, &) if @terminal
+          return record_error(record["error"]) if record["error"]
+
+          fold_record(record, &)
+        end
+
+        def after_terminal(record, &)
+          return ingest_usage(record) unless terminal_content?(record)
+
+          protocol_error("stream continued after its terminal record", &)
+        end
+
+        def record_error(error)
+          @error = ProviderError.new(error["message"] || "provider error",
+                                     status: error["code"])
+          @terminal = true
+        end
 
+        def fold_record(record, &)
+          block = record.dig("promptFeedback", "blockReason").to_s
+          @block_reason = block unless block.empty? || block == "BLOCK_REASON_UNSPECIFIED"
           candidate = record.dig("candidates", 0) || {}
-          Array(candidate.dig("content", "parts")).each { |part| fold_part(part, &) }
+          Array(candidate.dig("content", "parts")).each_with_index do |part, index|
+            fold_part(part, part_boundary: index.positive?, &)
+          end
           @finish_reason = candidate["finishReason"] if candidate["finishReason"]
-          @usage = parse_usage(record["usageMetadata"]) if record["usageMetadata"]
+          @terminal = true if @finish_reason || @block_reason
+          close_current(&) if @terminal
+          ingest_usage(record)
         end
+        private :after_terminal, :record_error, :fold_record
 
         # A stream that ended without a finishReason was truncated, not
         # cancelled: fail it so the loop can treat it as retryable.
         def finish(&emit)
           return fail_stream(@error, &emit) if @error
+          if (refused = blocked)
+            return fail_stream(refused, usage_known: @usage_authoritative, &emit)
+          end
           return fail_stream("stream ended without a finish reason", &emit) unless @finish_reason
 
+          if @blocks.any?(ToolCall) && !%w[STOP MAX_TOKENS].include?(@finish_reason)
+            invalidate_tool_calls
+            reason = ProviderError.new("generation ended before its tool calls were confirmed: " \
+                                       "#{@finish_reason}")
+            return fail_stream(reason, usage_known: @usage_authoritative, &emit)
+          end
+
           close_current(&emit)
+          invalidate_cost unless @usage_authoritative
           @message = assemble(stop_reason: stop_reason)
           emit&.call(Event.new(type: :done, reason: @message.stop_reason, message: @message))
           @message
         end
 
         def abort(&)
-          close_current
+          close_current(&)
+          invalidate_cost
           terminal(StopReason::ABORTED, "aborted", &)
         end
 
-        def fail_stream(reason, &)
-          close_current
+        def fail_stream(reason, usage_known: false, &)
+          close_current(&)
+          invalidate_cost unless usage_known
           text = case reason
                  when ProviderError then "#{reason.class}: #{reason.describe}"
                  when Exception then "#{reason.class}: #{reason.message}"
@@ -67,16 +126,52 @@ module Mistri
 
         private
 
-        def fold_part(part, &)
+        def fold_part(part, part_boundary: false, &)
           if part.key?("functionCall")
             fold_function_call(part, &)
           elsif part.key?("text")
-            fold_text(part, part["thought"] ? :thinking : :text, &)
+            fold_text(part, part["thought"] ? :thinking : :text, part_boundary:, &)
+          end
+        end
+
+        def terminal_content?(record)
+          return true if record["error"]
+          return true unless record.dig("promptFeedback", "blockReason").to_s.empty?
+
+          candidate = record.dig("candidates", 0)
+          candidate && (candidate["finishReason"] ||
+            !Array(candidate.dig("content", "parts")).empty?)
+        end
+
+        def ingest_usage(record)
+          raw = record["usageMetadata"]
+          return unless raw
+
+          @service_tier = raw["serviceTier"] if raw.key?("serviceTier")
+          @usage = priced(parse_usage(raw))
+          @usage_authoritative = true if @terminal
+        end
+
+        def protocol_error(message, &)
+          @error ||= ProviderError.new(message)
+          close_current(&)
+          invalidate_tool_calls
+          nil
+        end
+
+        def invalidate_tool_calls
+          @blocks.map! do |block|
+            if block.is_a?(ToolCall)
+              block.with(arguments: nil, arguments_error: "incomplete")
+            else
+              block
+            end
           end
         end
 
-        def fold_text(part, kind, &)
-          close_current(&) if @current && @current.kind != kind
+        def fold_text(part, kind, part_boundary:, &)
+          close_current(&) if @current && (@current.kind != kind || part_boundary ||
+                                           @current.signature)
           unless @current
             @current = Builder.new(kind, @blocks.size, +"", nil)
             emit_event(:"#{kind}_start", content_index: @current.index, &)
@@ -92,14 +187,15 @@ module Mistri
         def fold_function_call(part, &)
           close_current(&)
           call_spec = part["functionCall"] || {}
-          arguments = call_spec["args"].is_a?(Hash) ? call_spec["args"] : {}
-          call = ToolCall.new(id: call_spec["id"] || "call_#{@blocks.size + 1}",
+          arguments = call_spec.key?("args") ? call_spec["args"] : {}
+          wire_id = call_spec["id"]
+          call = ToolCall.new(id: wire_id || "gemini_#{SecureRandom.uuid}",
                               name: call_spec["name"], arguments: arguments,
-                              signature: part["thoughtSignature"])
+                              signature: part["thoughtSignature"], provider_call_id: wire_id)
           index = @blocks.size
           emit_event(:toolcall_start, content_index: index, &)
           emit_event(:toolcall_delta, content_index: index,
-                                      delta: JSON.generate(arguments), &)
+                                      delta: JSON.generate(call.arguments), &)
           @blocks << call
           emit_event(:toolcall_end, content_index: index, tool_call: call, &)
         end
@@ -124,9 +220,26 @@ module Mistri
           end
         end
 
+        # A blocked prompt arrives as promptFeedback with no candidates; a
+        # blocked candidate arrives as a verdict finishReason. Either way the
+        # error carries the wire word, so hosts tell SAFETY from RECITATION
+        # without parsing prose.
+        def blocked
+          if @block_reason
+            return InvalidRequestError.new("the prompt was blocked: #{@block_reason}")
+          end
+          return unless @finish_reason
+          if VERDICTS.include?(@finish_reason) || INPUT_ERRORS.include?(@finish_reason)
+            return InvalidRequestError.new("generation stopped: #{@finish_reason}")
+          end
+          return unless FUMBLES.include?(@finish_reason)
+
+          ProviderError.new("generation stopped: #{@finish_reason}")
+        end
+
         def stop_reason
-          return StopReason::TOOL_USE if @blocks.any?(ToolCall)
           return StopReason::LENGTH if @finish_reason == "MAX_TOKENS"
+          return StopReason::TOOL_USE if @blocks.any?(ToolCall)
 
           StopReason::STOP
         end
@@ -157,6 +270,20 @@ module Mistri
                     output: raw["candidatesTokenCount"].to_i + reasoning,
                     cache_read: cache_read, reasoning: reasoning)
         end
+
+        def priced(usage)
+          return usage unless @catalog_pricing
+
+          standard = @service_tier.nil? || %w[unspecified standard].include?(@service_tier.to_s)
+          return usage unless standard
+
+          rates = Models.rates(@model, usage:, at: @pricing_at)
+          rates ? usage.with_cost(rates) : usage
+        end
+
+        def invalidate_cost
+          @usage = @usage.with(cost: @usage.cost.with(known: false))
+        end
       end
     end
   end

data/lib/mistri/providers/gemini/serializer.rb

@@ -12,14 +12,28 @@ module Mistri
       # signature would be rejected. Thinking summaries are output-only and
       # never replay. Consecutive user turns stay separate: Gemini accepts
       # them, and mixing a text part into a functionResponse turn makes it
-      # answer an empty candidate.
+      # answer an empty candidate. Completed tool exchanges from another
+      # provider become ordinary text: Gemini rejects foreign functionCall
+      # history because it has no Gemini thought signature.
       module Serializer
         module_function
 
         def contents(history)
-          groups = history.reject(&:system?).chunk_while { |a, b| a.tool? && b.tool? }
+          result_metadata = tool_result_metadata(history)
+          groups = history.reject(&:system?).chunk_while do |a, b|
+            a.tool? && b.tool? && native_tool_result?(a, result_metadata) ==
+              native_tool_result?(b, result_metadata)
+          end
           groups.filter_map do |group|
-            group.first.tool? ? tool_turn(group) : turn(group.first)
+            if group.first.tool?
+              if native_tool_result?(group.first, result_metadata)
+                tool_turn(group, result_metadata)
+              else
+                foreign_tool_turn(group)
+              end
+            else
+              turn(group.first)
+            end
           end
         end
 
@@ -33,7 +47,7 @@ module Mistri
           declarations = definitions.map do |tool|
             spec = tool.transform_keys(&:to_sym)
             { name: spec[:name], description: spec[:description],
-              parameters: spec[:input_schema] }
+              parametersJsonSchema: spec[:input_schema] }
           end
           [{ functionDeclarations: declarations }]
         end
@@ -47,14 +61,28 @@ module Mistri
 
         # Gemini pairs a functionResponse to its call by NAME; a wrong name
         # silently mismatches, so a missing one fails loudly instead.
-        def tool_turn(group)
+        def tool_turn(group, result_metadata = {}.compare_by_identity)
           { role: "user", parts: group.map do |msg|
             unless msg.tool_name
               raise SchemaError, "Gemini tool results need tool_name to pair with their call"
             end
 
-            { functionResponse: { name: msg.tool_name,
-                                  response: { "result" => result_text(msg) } } }
+            key = msg.tool_error? ? "error" : "result"
+            response = { name: msg.tool_name, response: { key => result_text(msg) } }
+            wire_id = result_metadata.dig(msg, :provider_call_id)
+            response[:id] = wire_id if wire_id
+            { functionResponse: response }
+          end }
+        end
+
+        def foreign_tool_turn(group)
+          { role: "user", parts: group.map do |msg|
+            raise SchemaError, "Gemini tool results need tool_name" unless msg.tool_name
+
+            label = msg.tool_error? ? "error" : "result"
+            text = "Tool #{JSON.generate(msg.tool_name)} #{label} " \
+                   "(call #{JSON.generate(msg.tool_call_id)}): #{result_text(msg)}"
+            { text: text }
           end }
         end
 
@@ -84,10 +112,51 @@ module Mistri
             case block
             when Content::Text then signed({ text: block.text }, block.signature, own)
             when ToolCall
-              signed({ functionCall: { name: block.name, args: block.arguments } },
-                     block.signature, own)
+              own ? native_function_call(block) : foreign_function_call(block)
+            end
+          end
+        end
+
+        def native_function_call(block)
+          arguments = ToolArguments.replay_object(block)
+          signature = block.signature if arguments.equal?(block.arguments)
+          call = { name: block.name, args: arguments }
+          call[:id] = block.provider_call_id if block.provider_call_id
+          signed({ functionCall: call }, signature, true)
+        end
+
+        def foreign_function_call(block)
+          arguments = if block.arguments_error
+                        "unavailable (#{block.arguments_error})"
+                      else
+                        JSON.generate(block.arguments)
+                      end
+          { text: "Tool call #{JSON.generate(block.name)} " \
+                  "(call #{JSON.generate(block.id)}) with arguments: #{arguments}" }
+        end
+
+        def native_tool_result?(message, result_metadata)
+          result_metadata.dig(message, :provider) == :gemini
+        end
+
+        # Legacy histories may reuse call_N across turns. Correlating each
+        # result to the nearest preceding occurrence keeps a later Gemini call
+        # from rewriting an earlier foreign result's wire semantics.
+        def tool_result_metadata(history)
+          pending = Hash.new { |calls, id| calls[id] = [] }
+          metadata = {}.compare_by_identity
+          history.each do |message|
+            if message.assistant?
+              message.tool_calls.each do |call|
+                pending[call.id] << { provider: message.provider,
+                                      provider_call_id: call.provider_call_id }.freeze
+              end
+            elsif message.tool?
+              matched = pending[message.tool_call_id].pop
+              metadata[message] = matched if matched
             end
           end
+          metadata
         end
 
         def signed(part, signature, own)

data/lib/mistri/providers/gemini.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "schema_capabilities"
+
 module Mistri
   module Providers
     # The Gemini API (v1beta generateContent), streamed over SSE and
@@ -11,30 +13,50 @@ module Mistri
     # maxOutputTokens is omitted for the same reason: the API defaults to the
     # model's ceiling.
     class Gemini
+      DEFAULT_ORIGIN = "https://generativelanguage.googleapis.com"
       DEFAULT_THINKING = { includeThoughts: true }.freeze
 
       def initialize(api_key:, model: "gemini-2.5-flash",
-                     origin: "https://generativelanguage.googleapis.com",
-                     thinking: DEFAULT_THINKING, **transport_options)
+                     origin: DEFAULT_ORIGIN, thinking: DEFAULT_THINKING,
+                     service_tier: nil, catalog_pricing: nil, **transport_options)
         @api_key = api_key
         @model = model
         @thinking = thinking
+        @service_tier = service_tier
+        @catalog_pricing = catalog_pricing.nil? ? official_origin?(origin) : catalog_pricing
         @transport = Transport.new(origin: origin, **transport_options)
       end
 
       attr_reader :model
 
+      def prices_usage?
+        tier_known = @service_tier.nil? || %w[unspecified standard].include?(@service_tier.to_s)
+        @catalog_pricing && Models.priced?(model) && tier_known
+      end
+
+      def native_output_schema(schema)
+        return unless Models.find(model)&.provider == :gemini
+
+        SchemaCapabilities.derive(schema, :gemini)
+      end
+
       def stream(messages:, system: nil, tools: [], signal: nil, **overrides, &emit)
+        delivery = EventDelivery.wrap(emit)
         model = overrides.fetch(:model, @model)
-        assembler = Gemini::Assembler.new(model: model)
+        service_tier = overrides.fetch(:service_tier, @service_tier)
+        assembler = Gemini::Assembler.new(model: model, catalog_pricing: @catalog_pricing,
+                                          service_tier:)
+        assembler.start(&delivery)
         body = build_body(messages, system, tools, overrides)
         path = "/v1beta/models/#{model}:streamGenerateContent?alt=sse"
         outcome = @transport.stream_post(path, body: body, headers: headers,
                                                signal: signal) do |record|
           assembler.feed(record,
-                         &emit)
+                         &delivery)
         end
-        outcome == :aborted ? assembler.abort(&emit) : assembler.finish(&emit)
+        outcome == :aborted ? assembler.abort(&delivery) : assembler.finish(&delivery)
+      rescue EventDelivery::Failure => e
+        raise EventDelivery.unwrap(e, delivery)
       rescue Error => e
         assembler.fail_stream(e, &emit)
       end
@@ -43,11 +65,15 @@ module Mistri
 
       private
 
+      def official_origin?(origin) = origin.to_s.delete_suffix("/") == DEFAULT_ORIGIN
+
       def build_body(messages, system, tools, overrides)
         body = { contents: Serializer.contents(messages) }
         instruction = Serializer.system_instruction(system)
         body[:systemInstruction] = instruction if instruction
         body[:tools] = Serializer.tools(tools) if tools.any?
+        service_tier = overrides.fetch(:service_tier, @service_tier)
+        body[:serviceTier] = service_tier if service_tier
         config = {}
         thinking = overrides.fetch(:thinking, @thinking)
         config[:thinkingConfig] = thinking if thinking