mistri 0.5.0 → 0.6.1

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

@@ -6,29 +6,45 @@ module Mistri
       # Folds the Messages API stream into the event union, building the
       # assistant message block by block. Every emitted event carries an
       # immutable snapshot of the message so far; in-flight tool arguments
-      # parse via PartialJson so consumers can read them mid-stream.
+      # expose a bounded, cached PartialJson preview so hostile fragmentation
+      # cannot make snapshot construction quadratic.
       #
       # Unknown event and block types are skipped by contract: the API adds
       # types over time and a live stream must survive them.
-      class Assembler
-        def initialize(model:)
+      class Assembler # rubocop:disable Metrics/ClassLength -- one stream owns block order
+        def initialize(model:, catalog_pricing: true)
           @model = model
+          @catalog_pricing = catalog_pricing
+          @pricing_at = Time.now
           @blocks = []
           @current = nil
-          @usage = Usage.zero
+          @usage = Usage.new
           @stop_reason = nil
           @done = false
+          @next_wire_index = 0
+          @open_wire_index = nil
+          @message_phase = false
         end
 
+        STREAM_EVENTS = %w[
+          message_start content_block_start content_block_delta content_block_stop
+          message_delta message_stop error
+        ].freeze
+
         def feed(record, &)
-          case record["type"]
-          when "message_start" then @usage = parse_usage(record.dig("message", "usage"))
+          type = record["type"]
+          if @done && STREAM_EVENTS.include?(type)
+            return protocol_error("stream continued after message_stop", &)
+          end
+
+          case type
+          when "message_start" then message_start(record)
           when "content_block_start" then start_block(record, &)
           when "content_block_delta" then delta_block(record, &)
           when "content_block_stop" then stop_block(record, &)
-          when "message_delta" then message_delta(record)
-          when "message_stop" then @done = true
-          when "error" then @error = wire_error(record["error"])
+          when "message_delta" then message_delta(record, &)
+          when "message_stop" then stop_message(&)
+          when "error" then @error ||= wire_error(record["error"])
           end
         end
 
@@ -38,15 +54,24 @@ module Mistri
         # reading as a cancellation.
         def finish(&emit)
           return fail_stream(@error, &emit) if @error
+          if @refused
+            return fail_stream(refusal_error,
+                               usage_known: @done && @usage_authoritative, &emit)
+          end
           return fail_stream("stream ended without message_stop", &emit) unless @done
+          if @open_wire_index
+            return fail_stream("content block ended without content_block_stop", &emit)
+          end
 
+          invalidate_cost unless @usage_authoritative
           @message = assemble(stop_reason: @stop_reason || StopReason::STOP)
           emit&.call(Event.new(type: :done, reason: @message.stop_reason, message: @message))
           @message
         end
 
         def abort(&emit)
-          finalize_current
+          close_current(interrupted: true, &emit)
+          invalidate_cost
           @message = assemble(stop_reason: StopReason::ABORTED, error_message: "aborted")
           emit&.call(Event.new(type: :error, reason: StopReason::ABORTED, message: @message,
                                error_message: "aborted"))
@@ -57,12 +82,22 @@ module Mistri
         # as retryable, not fold into prose.
         def wire_error(payload)
           message = payload&.dig("message") || "provider error"
-          klass = payload&.dig("type").to_s.include?("overloaded") ? OverloadedError : ProviderError
-          klass.new(message)
+          type = payload&.dig("type").to_s
+          klass = case type
+                  when "authentication_error" then AuthenticationError
+                  when "rate_limit_error" then RateLimitError
+                  when "overloaded_error" then OverloadedError
+                  when "api_error" then ServerError
+                  when "timeout_error" then ProviderError
+                  else InvalidRequestError
+                  end
+          detail = type.empty? ? message : "#{type}: #{message}"
+          klass.new(detail)
         end
 
-        def fail_stream(reason, &emit)
-          finalize_current
+        def fail_stream(reason, usage_known: false, &emit)
+          close_current(interrupted: true, &emit)
+          invalidate_cost unless usage_known
           text = case reason
                  when ProviderError then "#{reason.class}: #{reason.describe}"
                  when Exception then "#{reason.class}: #{reason.message}"
@@ -77,35 +112,110 @@ module Mistri
 
         def message = @message ||= finish
 
-        Builder = Struct.new(:kind, :index, :text, :json, :signature, :id, :name, :redacted)
+        Builder = Struct.new(:kind, :index, :text, :json, :signature, :id, :name, :redacted,
+                             :argument_bytes, :argument_error, :argument_preview, :preview_bytes)
+        DELTA_KINDS = {
+          "text_delta" => :text,
+          "thinking_delta" => :thinking,
+          "signature_delta" => :thinking,
+          "input_json_delta" => :toolcall
+        }.freeze
+
+        PREVIEW_EAGER_BYTES = 4 * 1024
+        PREVIEW_STEP_BYTES = 4 * 1024
+        PREVIEW_MAX_BYTES = 64 * 1024
+        MAX_SIGNATURE_BYTES = ToolArguments::MAX_BYTES
+        MAX_REFUSAL_BYTES = 2048
+
+        def start(&emit)
+          emit&.call(Event.new(type: :start, partial: assemble))
+        end
 
         private
 
+        def message_start(record)
+          raw = record.dig("message", "usage")
+          @service_tier = raw&.fetch("service_tier", nil)
+          @usage = priced(parse_usage(raw))
+        end
+
         def start_block(record, &)
+          return if @error
+          return protocol_error("stream opened content after the message delta phase", &) \
+            if @message_phase
+          if @open_wire_index
+            return protocol_error("stream opened a content block before closing the prior block", &)
+          end
+
+          index = record["index"]
+          unless index == @next_wire_index
+            return protocol_error("stream opened a content block at an unexpected index", &)
+          end
+
+          @open_wire_index = index
+          @next_wire_index += 1
+
           block = record["content_block"] || {}
+          if block["type"] == "tool_use" && block.key?("input") &&
+             (!block["input"].is_a?(Hash) || !block["input"].empty?)
+            return protocol_error("tool_use started with nonempty input before its deltas", &)
+          end
+
           kind = { "text" => :text, "thinking" => :thinking, "redacted_thinking" => :thinking,
                    "tool_use" => :toolcall }[block["type"]]
           return unless kind
 
           @current = Builder.new(kind, @blocks.size, +"", +"", nil,
-                                 block["id"], block["name"], block["type"] == "redacted_thinking")
+                                 block["id"], block["name"], block["type"] == "redacted_thinking",
+                                 0, nil, ToolArguments::EMPTY_OBJECT, 0)
           @current.signature = block["data"] if @current.redacted
           emit_event(:"#{kind}_start", content_index: @current.index, &)
         end
 
         def delta_block(record, &)
+          return if @error
+          return protocol_error("stream sent content after the message delta phase", &) \
+            if @message_phase
+          unless @open_wire_index && record["index"] == @open_wire_index
+            return protocol_error("stream sent a delta for a different content block", &)
+          end
           return unless @current
 
           delta = record["delta"] || {}
+          expected = DELTA_KINDS[delta["type"]]
+          return unless expected
+          unless @current.kind == expected
+            return protocol_error("stream sent a delta outside its matching content block", &)
+          end
+
+          route_delta(delta, &)
+        end
+
+        def route_delta(delta, &)
           case delta["type"]
-          when "text_delta" then text_delta(delta["text"], &)
-          when "thinking_delta" then thinking_delta(delta["thinking"], &)
+          when "text_delta"
+            text_delta(delta["text"], &)
+          when "thinking_delta"
+            thinking_delta(delta["thinking"], &)
           when "signature_delta"
-            @current.signature = "#{@current.signature}#{delta["signature"]}"
-          when "input_json_delta" then input_delta(delta["partial_json"], &)
+            signature_delta(delta["signature"])
+          when "input_json_delta"
+            input_delta(delta["partial_json"], &)
           end
         end
 
+        def signature_delta(fragment)
+          fragment = fragment.to_s
+          current = @current.signature || +""
+          if fragment.bytesize > MAX_SIGNATURE_BYTES - current.bytesize
+            raise ResponseTooLargeError.new(kind: :thinking_signature,
+                                            limit: MAX_SIGNATURE_BYTES)
+          end
+
+          @current.signature = current
+          current << fragment
+        end
+
         def text_delta(text, &)
           @current.text << text.to_s
           emit_event(:text_delta, content_index: @current.index, delta: text, &)
@@ -117,54 +227,122 @@ module Mistri
         end
 
         def input_delta(fragment, &)
-          @current.json << fragment.to_s
+          append_arguments(@current, fragment.to_s)
           emit_event(:toolcall_delta, content_index: @current.index, delta: fragment, &)
         end
 
-        def stop_block(_record, &)
-          return unless @current
+        def append_arguments(builder, fragment)
+          return if builder.argument_error
 
-          block = finalize_current
-          kind = block.is_a?(ToolCall) ? :toolcall : block.type
-          fields = { content_index: @blocks.size - 1 }
-          fields[:tool_call] = block if block.is_a?(ToolCall)
-          fields[:content] = @blocks.last.is_a?(ToolCall) ? nil : builder_text(block)
-          emit_event(:"#{kind}_end", **fields.compact, &)
+          if fragment.bytesize > ToolArguments::MAX_BYTES - builder.argument_bytes
+            builder.argument_error = "too_large"
+            # String#clear may retain its backing allocation.
+            builder.json = +""
+            return
+          end
+
+          builder.json << fragment
+          builder.argument_bytes += fragment.bytesize
+          refresh_argument_preview(builder)
+        end
+
+        def refresh_argument_preview(builder)
+          target = [builder.argument_bytes, PREVIEW_MAX_BYTES].min
+          eager = target <= PREVIEW_EAGER_BYTES
+          milestone = target - builder.preview_bytes >= PREVIEW_STEP_BYTES
+          return unless eager || milestone || target == PREVIEW_MAX_BYTES
+          return if target == builder.preview_bytes
+
+          source = if target == builder.json.bytesize
+                     builder.json
+                   else
+                     builder.json.byteslice(0, target)
+                   end
+          builder.argument_preview = ToolArguments.freeze_partial(PartialJson.parse(source))
+          builder.preview_bytes = target
+        end
+
+        def stop_block(record, &)
+          return if @error
+          return protocol_error("stream stopped content after the message delta phase", &) \
+            if @message_phase
+          unless @open_wire_index && record["index"] == @open_wire_index
+            return protocol_error("stream stopped a different content block", &)
+          end
+
+          close_current(&) if @current
+          @open_wire_index = nil
         end
 
-        def message_delta(record)
+        def message_delta(record, &)
+          @message_phase = true
+          if @open_wire_index
+            return protocol_error("message delta arrived before the content block stopped", &)
+          end
+
           reason = record.dig("delta", "stop_reason")
           @stop_reason = map_stop_reason(reason) if reason
+          if reason == "refusal"
+            @refused = true
+            @refusal = record.dig("delta", "stop_details")
+          end
           # message_delta usage is cumulative; merge output counts over the
           # opening snapshot rather than summing.
           output = record.dig("usage", "output_tokens")
-          @usage = @usage.with(output: output.to_i) if output
+          return unless output
+
+          @usage = priced(@usage.with(output: output.to_i))
+          @usage_authoritative = true
+        end
+
+        def stop_message(&)
+          protocol_error("message stopped before the content block stopped", &) if @open_wire_index
+          @done = true
         end
 
-        def finalize_current
+        def close_current(interrupted: false, &)
           return unless @current
 
-          built = build_block(@current)
+          built = build_block(@current, interrupted:)
           @blocks << built
           @current = nil
+          kind = built.is_a?(ToolCall) ? :toolcall : built.type
+          fields = { content_index: @blocks.size - 1 }
+          fields[:tool_call] = built if built.is_a?(ToolCall)
+          fields[:content] = builder_text(built) unless built.is_a?(ToolCall)
+          emit_event(:"#{kind}_end", **fields, &)
           built
         end
 
-        def build_block(builder)
+        def build_block(builder, final: true, interrupted: false)
           case builder.kind
           when :text then Content::Text.new(text: builder.text)
           when :thinking
-            Content::Thinking.new(thinking: builder.text, signature: builder.signature,
-                                  redacted: builder.redacted)
+            signature = interrupted ? nil : builder.signature
+            redacted = interrupted ? false : builder.redacted
+            Content::Thinking.new(thinking: builder.text, signature:, redacted:)
           when :toolcall
+            arguments, error = if interrupted
+                                 [nil, "incomplete"]
+                               elsif final
+                                 if builder.argument_error
+                                   [nil, builder.argument_error]
+                                 else
+                                   parsed_arguments(builder.json)
+                                 end
+                               else
+                                 [builder.argument_preview, nil]
+                               end
             ToolCall.new(id: builder.id, name: builder.name,
-                         arguments: parsed_arguments(builder.json), signature: nil)
+                         arguments:, signature: nil, arguments_error: error,
+                         canonicalize: final)
           end
         end
 
         def parsed_arguments(json)
-          parsed = json.strip.empty? ? {} : PartialJson.parse(json)
-          parsed.is_a?(Hash) ? parsed : {}
+          return [{}, nil] if json.empty?
+
+          ToolArguments.parse_json(json)
         end
 
         def builder_text(block)
@@ -175,23 +353,71 @@ module Mistri
           emit&.call(Event.new(type:, partial: assemble, **fields))
         end
 
+        def protocol_error(message, &)
+          @error ||= ProviderError.new(message)
+          close_current(interrupted: true, &)
+          @blocks.map! do |block|
+            if block.is_a?(ToolCall)
+              block.with(arguments: nil,
+                         arguments_error: "incomplete")
+            else
+              block
+            end
+          end
+          @open_wire_index = nil
+          nil
+        end
+
         def assemble(**meta)
           blocks = @blocks.dup
-          blocks << build_block(@current) if @current
+          blocks << build_block(@current, final: false) if @current
           Message.assistant(content: blocks, model: @model, provider: :anthropic,
                             usage: @usage, **meta)
         end
 
-        # pause_turn (a server tool paused a long turn) maps to tool_use so the
-        # loop continues the turn rather than ending it.
+        # pause_turn (a server tool paused a long turn) maps to tool_use so
+        # the loop continues the turn rather than ending it; a filled context
+        # window is a truncation, per the API's own guidance.
         def map_stop_reason(reason)
           { "end_turn" => StopReason::STOP, "stop_sequence" => StopReason::STOP,
             "max_tokens" => StopReason::LENGTH, "tool_use" => StopReason::TOOL_USE,
+            "model_context_window_exceeded" => StopReason::LENGTH,
             "pause_turn" => StopReason::TOOL_USE }.fetch(reason, StopReason::STOP)
         end
 
+        # The API's guidance for a refusal is a different model, never a
+        # retry of this one, so it fails fast; stop_details names the policy
+        # category when the API provides one.
+        def refusal_error
+          details = bounded_refusal(@refusal&.dig("category"),
+                                    @refusal&.dig("explanation"))
+          text = +"the model refused to respond"
+          text << " (#{details})" unless details.empty?
+          InvalidRequestError.new(text)
+        end
+
+        def bounded_refusal(*values)
+          output = +""
+          values.each do |value|
+            next unless value.is_a?(String)
+
+            separator = output.empty? ? "" : ": "
+            remaining = MAX_REFUSAL_BYTES - output.bytesize - separator.bytesize
+            break unless remaining.positive?
+
+            output << separator << utf8_prefix(value, remaining)
+          end
+          output
+        end
+
+        def utf8_prefix(value, limit)
+          prefix = value.byteslice(0, limit).dup.force_encoding(value.encoding)
+          prefix.encode(Encoding::UTF_8, invalid: :replace, undef: :replace, replace: "?")
+                .scrub.byteslice(0, limit).to_s.force_encoding(Encoding::UTF_8).scrub
+        end
+
         def parse_usage(raw)
-          return Usage.zero unless raw
+          return Usage.new unless raw
 
           cache_creation = raw["cache_creation"] || {}
           Usage.new(input: raw["input_tokens"].to_i, output: raw["output_tokens"].to_i,
@@ -199,7 +425,19 @@ module Mistri
                     cache_write: raw["cache_creation_input_tokens"].to_i,
                     cache_write_1h: cache_creation["ephemeral_1h_input_tokens"].to_i)
         end
-      end
+
+        def priced(usage)
+          return usage unless @catalog_pricing
+          return usage unless @service_tier == "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 # rubocop:enable Metrics/ClassLength
     end
   end
 end

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

@@ -42,30 +42,35 @@ module Mistri
         end
 
         def message(msg)
-          { role: msg.role.to_s, content: msg.content.filter_map { |block| block(block) } }
+          own = msg.provider == :anthropic
+          { role: msg.role.to_s,
+            content: msg.content.filter_map { |content| block(content, own:) } }
         end
 
         def tool_results(group)
           { role: "user", content: group.map do |msg|
-            blocks = msg.content.filter_map { |block| block(block) }
+            blocks = msg.content.filter_map { |content| block(content, own: false) }
             # The API rejects an empty tool_result; a space stands in for a
             # tool that returned nothing.
             blocks = [{ type: "text", text: " " }] if blocks.empty?
-            { type: "tool_result", tool_use_id: msg.tool_call_id, content: blocks }
+            result = { type: "tool_result", tool_use_id: msg.tool_call_id, content: blocks }
+            result[:is_error] = true if msg.tool_error?
+            result
           end }
         end
 
         # Returns nil for a block the API would reject (empty text, unusable
         # thinking), so callers filter_map it out.
-        def block(block)
+        def block(block, own: true)
           case block
           when Content::Text then text_block(block)
-          when Content::Thinking then thinking_block(block)
+          when Content::Thinking then thinking_block(block, own:)
           when Content::Image
             { type: "image",
               source: { type: "base64", media_type: block.mime_type, data: block.data } }
           when ToolCall
-            { type: "tool_use", id: block.id, name: block.name, input: block.arguments }
+            { type: "tool_use", id: block.id, name: block.name,
+              input: ToolArguments.replay_object(block) }
           else
             raise SchemaError, "cannot serialize #{block.class} for Anthropic"
           end
@@ -80,9 +85,9 @@ module Mistri
         # its opaque payload; a normal thinking block missing its signature
         # (an aborted turn cut before signature_delta) cannot replay, so it
         # degrades to its text, or drops when even that is empty.
-        def thinking_block(block)
-          return { type: "redacted_thinking", data: block.signature } if block.redacted?
-          if block.signature
+        def thinking_block(block, own: true)
+          return { type: "redacted_thinking", data: block.signature } if own && block.redacted?
+          if own && block.signature
             return { type: "thinking", thinking: block.thinking,
                      signature: block.signature }
           end

data/lib/mistri/providers/anthropic.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "schema_capabilities"
+
 module Mistri
   module Providers
     # The Anthropic Messages API, streamed. Defaults target the current model
@@ -10,41 +12,58 @@ module Mistri
     # raising: the loop decides whether to retry, and the host always gets a
     # message back.
     class Anthropic
+      DEFAULT_ORIGIN = "https://api.anthropic.com"
       VERSION_HEADER = "2023-06-01"
       DEFAULT_THINKING = { type: "adaptive", display: "summarized" }.freeze
 
       # Messages API parameters passed through verbatim from a stream override.
-      PASSTHROUGH = %i[temperature top_p top_k stop_sequences metadata
-                       tool_choice service_tier].freeze
+      PASSTHROUGH = %i[temperature top_p top_k stop_sequences metadata tool_choice].freeze
       # The ceiling for an uncatalogued model: high enough for headroom, low
       # enough that every current model accepts it. Catalog a model to unlock
       # its real output limit.
       UNKNOWN_MODEL_MAX_TOKENS = 64_000
 
-      def initialize(api_key:, model: "claude-opus-4-8", origin: "https://api.anthropic.com",
+      def initialize(api_key:, model: "claude-opus-4-8", origin: DEFAULT_ORIGIN,
                      max_tokens: nil, thinking: DEFAULT_THINKING, cache: true,
+                     service_tier: nil, catalog_pricing: nil,
                      **transport_options)
         @api_key = api_key
         @model = model
         @max_tokens = max_tokens
         @thinking = thinking
         @cache = cache
+        @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?
+        @catalog_pricing && Models.priced?(model) && @service_tier.to_s == "standard_only"
+      end
+
+      def native_output_schema(schema)
+        return unless Models.find(model)&.provider == :anthropic
+
+        SchemaCapabilities.derive(schema, :anthropic)
+      end
+
       def stream(messages:, system: nil, tools: [], signal: nil, **overrides, &emit)
+        delivery = EventDelivery.wrap(emit)
         model = overrides.fetch(:model, @model)
-        assembler = Anthropic::Assembler.new(model: model)
+        assembler = Anthropic::Assembler.new(model: model, catalog_pricing: @catalog_pricing)
+        assembler.start(&delivery)
         body = build_body(model, messages, system, tools, overrides)
         outcome = @transport.stream_post("/v1/messages", body: body, headers: headers,
                                                          signal: signal) do |record|
           assembler.feed(
-            record, &emit
+            record, &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
@@ -53,6 +72,8 @@ module Mistri
 
       private
 
+      def official_origin?(origin) = origin.to_s.delete_suffix("/") == DEFAULT_ORIGIN
+
       def build_body(model, messages, system, tools, overrides)
         body = {
           model: model,
@@ -63,6 +84,8 @@ module Mistri
         system_blocks = Serializer.system_blocks(system, cache: @cache)
         body[:system] = system_blocks if system_blocks
         body[:tools] = Serializer.tools(tools) if tools.any?
+        service_tier = overrides.fetch(:service_tier, @service_tier)
+        body[:service_tier] = service_tier if service_tier
         thinking = thinking_for(model, overrides)
         body[:thinking] = thinking if thinking
         if (schema = overrides[:output_schema])