mistri 0.5.0 → 0.6.1

data/lib/mistri/tool.rb

@@ -4,50 +4,81 @@ require "json"
 
 module Mistri
   # A tool the agent can call: a name, a description, a JSON Schema for its
-  # arguments, and a handler. The handler receives the parsed arguments hash
-  # (string keys, exactly as the model sent them) and returns a String, a Hash
-  # (serialized as JSON), or content blocks, so a tool can hand back images as
-  # naturally as text.
+  # arguments, and a handler. Model calls reach the handler as canonical
+  # argument hashes, optionally normalized by the tool before policy sees
+  # them. Trusted direct calls keep ordinary Ruby Hash semantics.
   class Tool
     # A no-argument tool still needs a valid object schema; providers reject a
     # bare empty hash.
-    EMPTY_SCHEMA = { type: "object", properties: {} }.freeze
+    EMPTY_SCHEMA = {
+      type: "object", properties: {}.freeze, required: [].freeze,
+      additionalProperties: false
+    }.freeze
 
     attr_reader :name, :description, :input_schema, :timeout
 
-    # Define a tool. Give the argument shape as a raw JSON Schema hash via
-    # input_schema:, or build it in Ruby with a schema: block.
+    # Define a tool. Give the argument shape as a JSON Schema value via
+    # input_schema:, or build it in Ruby with a schema: block. The result is
+    # canonicalized and owned once so provider and local semantics cannot drift.
     #
     #   Tool.define("get_weather", "Weather for a city",
     #               schema: -> { string :city, "City name", required: true }) do |args|
     #     Weather.for(args["city"])
     #   end
     def self.define(name, description, input_schema: nil, schema: nil, **, &handler)
-      input_schema ||= schema ? Schema.build(&schema) : EMPTY_SCHEMA
+      raise ArgumentError, "choose input_schema or schema, not both" if schema && !input_schema.nil?
+
+      input_schema = schema ? Schema.build(&schema) : EMPTY_SCHEMA if input_schema.nil?
       new(name: name, description: description, input_schema: input_schema, **, &handler)
     end
 
     def initialize(name:, description:, input_schema: EMPTY_SCHEMA, eager_input_streaming: false,
-                   needs_approval: false, ends_turn: false, timeout: nil, &handler)
+                   needs_approval: false, ends_turn: false, timeout: nil,
+                   argument_normalizer: nil, argument_validator: nil,
+                   complete_argument_validator: nil, &handler)
       raise ArgumentError, "tool #{name.inspect} needs a handler block" unless handler
+      unless argument_normalizer.nil? || argument_normalizer.respond_to?(:call)
+        raise ArgumentError, "argument_normalizer must be callable"
+      end
+      unless argument_validator.nil? || argument_validator.respond_to?(:call)
+        raise ArgumentError, "argument_validator must be callable"
+      end
+      unless complete_argument_validator.nil? || complete_argument_validator.respond_to?(:call)
+        raise ArgumentError, "complete_argument_validator must be callable"
+      end
+      if argument_validator && complete_argument_validator
+        raise ArgumentError,
+              "choose argument_validator or complete_argument_validator, not both"
+      end
 
       @name = name.to_s
       @description = description
-      @input_schema = input_schema
+      @schema_validator = Schema.tool_validator(
+        input_schema, complete: !complete_argument_validator.nil?
+      )
+      @input_schema = @schema_validator.schema
       @eager_input_streaming = eager_input_streaming
       @needs_approval = needs_approval
       @ends_turn = ends_turn
       @timeout = timeout
+      @argument_normalizer = argument_normalizer
+      @argument_validator = argument_validator
+      @complete_argument_validator = complete_argument_validator
       @handler = handler
     end
 
-    # A handler may return a ToolResult to speak on two channels; its ui
-    # payload is canonicalized through JSON here so the live event and a
-    # reloaded session read the identical shape.
+    # A handler may return a ToolResult to add host-only UI or declare a
+    # model-readable failure. Its ui payload is canonicalized through JSON
+    # here so the live event and a reloaded session read the identical shape.
     #
     # Handlers receive (arguments, context). A proc that declares one
     # parameter ignores the context invisibly; a lambda opts in by arity.
+    # Direct calls are trusted host invocations: they apply the compatibility
+    # normalizer, but do not canonicalize or validate as model calls do. The
+    # executor marks arguments the Agent already prepared so subclasses keep
+    # the historical #call extension point without running normalization twice.
     def call(arguments, context = ToolContext.new)
+      arguments = normalize_arguments(arguments || {}) unless prepared_context?(context)
       result = invoke(arguments || {}, context)
       return serialize_result(result) unless result.is_a?(ToolResult)
 
@@ -55,6 +86,31 @@ module Mistri
                   ui: result.ui && JSON.parse(JSON.generate(result.ui)))
     end
 
+    # Normalization is an explicit per-tool migration boundary, never a
+    # global coercion policy. Agent calls this once before policy; direct
+    # trusted invocations get the same compatibility behavior through #call.
+    def normalize_arguments(arguments)
+      return arguments unless @argument_normalizer
+
+      normalized = @argument_normalizer.call(arguments)
+      raise ArgumentError, "argument_normalizer must return a Hash" unless normalized.is_a?(Hash)
+
+      normalized
+    end
+
+    # Core validation is always authoritative for its portable subset. A
+    # supplemental validator adds domain rules; an explicitly complete one
+    # additionally owns schema interactions core cannot represent.
+    def argument_violations(arguments)
+      validate_arguments(arguments, owned: false)
+    end
+
+    # Agent has already moved the value through ToolCall's ownership boundary,
+    # so its hot path can avoid copying the same immutable JSON twice.
+    def prepared_argument_violations(arguments)
+      validate_arguments(arguments, owned: true)
+    end
+
     # Whether this call should pause for a human. true/false, or a callable
     # given the parsed arguments so a tool can gate only the risky calls
     # (needs_approval: ->(args) { args["amount"].to_i > 100 }).
@@ -79,6 +135,30 @@ module Mistri
 
     private
 
+    def prepared_context?(context)
+      context.respond_to?(:arguments_prepared?) && context.arguments_prepared?
+    end
+
+    def validate_arguments(arguments, owned:)
+      unless owned
+        arguments, error = ToolArguments.canonicalize(arguments)
+        return ["$ validation limit exceeded"] if ToolArguments.resource_error?(error)
+        return ["$ must be valid JSON"] if error
+      end
+
+      errors = @schema_validator.violations(arguments, owned: true)
+      custom_validator = @complete_argument_validator || @argument_validator
+      return errors if errors.any? || !custom_validator
+
+      custom = custom_validator.call(arguments, @input_schema)
+      unless custom.is_a?(Array) && custom.all?(String)
+        name = @complete_argument_validator ? "complete_argument_validator" : "argument_validator"
+        raise TypeError, "#{name} must return an Array of Strings"
+      end
+
+      custom
+    end
+
     def invoke(arguments, context)
       if @handler.lambda? && @handler.arity.between?(0, 1)
         @handler.arity.zero? ? @handler.call : @handler.call(arguments)

data/lib/mistri/tool_arguments.rb

@@ -0,0 +1,377 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Mistri
+  # Owns model-supplied JSON before it crosses a durable or policy boundary.
+  module ToolArguments # rubocop:disable Metrics/ModuleLength -- one JSON ownership boundary
+    MAX_DEPTH = 64
+    MAX_NODES = 10_000
+    MAX_LEXICAL_TOKENS = (MAX_NODES * 2) + 1
+    MAX_BYTES = 8 * 1024 * 1024
+    MAX_NUMBER_BYTES = 64 * 1024
+    EMPTY_OBJECT = {}.freeze
+    ERROR_CODES = %w[
+      invalid_arguments invalid_encoding invalid_json invalid_key duplicate_key
+      cyclic_value too_deep too_large too_many_nodes number_too_large
+      non_finite_number non_json_value incomplete
+    ].freeze
+    RESOURCE_ERRORS = %w[too_deep too_large too_many_nodes number_too_large].freeze
+
+    Failure = Class.new(StandardError) do
+      attr_reader :code
+
+      def initialize(code)
+        @code = code
+        super
+      end
+    end
+    private_constant :Failure
+
+    module_function
+
+    def canonicalize(value)
+      state = { nodes: 0, bytes: 0, active: {} }
+      value = copy(value, state, 0)
+      raise Failure, "too_large" unless serialized_size_within_limit?(value)
+
+      [value, nil]
+    rescue Failure => e
+      [nil, e.code]
+    end
+
+    # Counts the generated JSON envelope without materializing it. Integer
+    # digits are conservatively bounded so a hostile Bignum cannot allocate an
+    # oversized decimal string before the limit is enforced.
+    def serialized_size_within_limit?(value, limit: MAX_BYTES)
+      return false unless limit.is_a?(Integer) && limit >= 0
+
+      state = { limit:, remaining: limit, active: {} }
+      count_json(value, state, 0)
+      true
+    rescue Failure
+      false
+    end
+
+    # Provider assemblers parse completed argument payloads strictly. Partial
+    # JSON belongs only to transient streaming snapshots, never executable calls.
+    def parse_json(raw)
+      return [nil, "invalid_json"] unless raw.is_a?(String)
+      return [nil, "too_large"] if raw.bytesize > MAX_BYTES
+      if (resource_error = raw_json_resource_error(raw))
+        return [nil, resource_error]
+      end
+
+      [JSON.parse(raw, max_nesting: MAX_DEPTH + 1), nil]
+    rescue JSON::NestingError
+      [nil, "too_deep"]
+    rescue JSON::ParserError, TypeError
+      [nil, "invalid_json"]
+    end
+
+    def normalize_error(code)
+      candidate = code.to_s
+      ERROR_CODES.include?(candidate) ? candidate.freeze : "invalid_arguments"
+    end
+
+    def resource_error?(code) = RESOURCE_ERRORS.include?(code)
+
+    # A lexical pass bounds parser allocation without duplicating JSON's
+    # grammar. Counting keys as tokens leaves room for every valid document
+    # under MAX_NODES while stopping wide hostile inputs before JSON.parse.
+    # This byte-state machine must branch on JSON token starts without building a tree.
+    # rubocop:disable Metrics/CyclomaticComplexity, Metrics/MethodLength, Metrics/PerceivedComplexity
+    def raw_json_resource_error(raw)
+      tokens = 0
+      depth = 0
+      index = 0
+      while index < raw.bytesize
+        byte = raw.getbyte(index)
+        case byte
+        when 34
+          tokens += 1
+          return "too_many_nodes" if tokens > MAX_LEXICAL_TOKENS
+
+          index = string_end(raw, index + 1)
+        when 91, 123
+          tokens += 1
+          return "too_many_nodes" if tokens > MAX_LEXICAL_TOKENS
+
+          depth += 1
+          return "too_deep" if depth > MAX_DEPTH + 1
+
+          index += 1
+        when 93, 125
+          depth -= 1 if depth.positive?
+          index += 1
+        when 45, 48..57
+          tokens += 1
+          return "too_many_nodes" if tokens > MAX_LEXICAL_TOKENS
+
+          number_start = index
+          index = number_end(raw, index + 1)
+          return "number_too_large" if index - number_start > MAX_NUMBER_BYTES
+        when 116
+          tokens += 1
+          return "too_many_nodes" if tokens > MAX_LEXICAL_TOKENS
+
+          index += raw.byteslice(index, 4) == "true" ? 4 : 1
+        when 102
+          tokens += 1
+          return "too_many_nodes" if tokens > MAX_LEXICAL_TOKENS
+
+          index += raw.byteslice(index, 5) == "false" ? 5 : 1
+        when 110
+          tokens += 1
+          return "too_many_nodes" if tokens > MAX_LEXICAL_TOKENS
+
+          index += raw.byteslice(index, 4) == "null" ? 4 : 1
+        else
+          index += 1
+        end
+      end
+      nil
+    end
+    # rubocop:enable Metrics/CyclomaticComplexity, Metrics/MethodLength, Metrics/PerceivedComplexity
+
+    def string_end(raw, index)
+      while index < raw.bytesize
+        case raw.getbyte(index)
+        when 34 then return index + 1
+        when 92 then index += 2
+        else index += 1
+        end
+      end
+      index
+    end
+    private_class_method :string_end
+
+    def number_end(raw, index)
+      while index < raw.bytesize
+        byte = raw.getbyte(index)
+        break unless byte == 43 || byte == 45 || byte == 46 || byte == 69 || byte == 101 ||
+                     byte.between?(48, 57)
+
+        index += 1
+      end
+      index
+    end
+    private_class_method :number_end
+
+    # Anthropic and Gemini require an object when an earlier invalid call is
+    # replayed. The paired error result preserves the truth; this shape only
+    # keeps the provider transcript structurally valid.
+    def replay_object(call)
+      return EMPTY_OBJECT if call.arguments_error
+      return EMPTY_OBJECT unless call.arguments.is_a?(Hash)
+
+      call.arguments
+    end
+
+    # OpenAI carries function arguments as a JSON string, so every completed
+    # JSON value can replay exactly; only an unparseable value needs a placeholder.
+    def replay_value(call)
+      call.arguments_error ? EMPTY_OBJECT : call.arguments
+    end
+
+    # Assemblers own freshly parsed previews, so they can freeze the tree in
+    # place without the copy that completed values require.
+    def freeze_partial(value)
+      return value if value.frozen?
+
+      case value
+      when Hash
+        value.each do |key, nested|
+          freeze_partial(key)
+          freeze_partial(nested)
+        end
+      when Array
+        value.each { |nested| freeze_partial(nested) }
+      end
+      value.freeze if value.respond_to?(:freeze)
+      value
+    end
+
+    def copy(value, state, depth)
+      raise Failure, "too_deep" if depth > MAX_DEPTH
+
+      state[:nodes] += 1
+      raise Failure, "too_large" if state[:nodes] > MAX_NODES
+
+      case value
+      when nil, true, false then value
+      when Integer then copy_integer(value)
+      when Float then copy_float(value)
+      when String then copy_string(value, state)
+      when Array then copy_array(value, state, depth)
+      when Hash then copy_hash(value, state, depth)
+      else raise Failure, "non_json_value"
+      end
+    end
+    private_class_method :copy
+
+    def copy_float(value)
+      raise Failure, "non_finite_number" unless value.finite?
+
+      value
+    end
+    private_class_method :copy_float
+
+    def copy_integer(value)
+      upper_bound = integer_decimal_upper_bound(value)
+      raise Failure, "number_too_large" if upper_bound > MAX_NUMBER_BYTES + 1
+      if upper_bound > MAX_NUMBER_BYTES && value.to_s.bytesize > MAX_NUMBER_BYTES
+        raise Failure, "number_too_large"
+      end
+
+      value
+    end
+    private_class_method :copy_integer
+
+    def copy_string(value, state)
+      string = if value.encoding == Encoding::UTF_8
+                 raise Failure, "too_large" if value.bytesize > MAX_BYTES - state[:bytes]
+                 raise Failure, "invalid_encoding" unless value.valid_encoding?
+
+                 value.dup
+               else
+                 value.encode(Encoding::UTF_8)
+               end
+      state[:bytes] += string.bytesize
+      raise Failure, "too_large" if state[:bytes] > MAX_BYTES
+
+      string.freeze
+    rescue EncodingError
+      raise Failure, "invalid_encoding"
+    end
+    private_class_method :copy_string
+
+    def copy_array(value, state, depth)
+      within(value, state) do
+        copied = []
+        # Array#map preallocates from an untrusted length before the node ceiling fires.
+        # rubocop:disable Style/MapIntoArray
+        value.each { |item| copied << copy(item, state, depth + 1) }
+        # rubocop:enable Style/MapIntoArray
+        copied.freeze
+      end
+    end
+    private_class_method :copy_array
+
+    def copy_hash(value, state, depth)
+      within(value, state) do
+        value.each_with_object({}) do |(key, item), copy|
+          key = copy_key(key, state)
+          raise Failure, "duplicate_key" if copy.key?(key)
+
+          copy[key] = copy(item, state, depth + 1)
+        end.freeze
+      end
+    end
+    private_class_method :copy_hash
+
+    def copy_key(key, state)
+      raise Failure, "invalid_key" unless key.is_a?(String) || key.is_a?(Symbol)
+
+      copy_string(key.to_s, state)
+    end
+    private_class_method :copy_key
+
+    def within(value, state)
+      identity = value.object_id
+      raise Failure, "cyclic_value" if state[:active].key?(identity)
+
+      state[:active][identity] = true
+      inserted = true
+      yield
+    ensure
+      state[:active].delete(identity) if inserted
+    end
+    private_class_method :within
+
+    def count_json(value, state, depth)
+      raise Failure, "too_deep" if depth > MAX_DEPTH
+
+      case value
+      when nil, true then consume(state, 4)
+      when false then consume(state, 5)
+      when Integer then count_integer(value, state)
+      when Float then count_float(value, state)
+      when String then count_string(value, state)
+      when Array then count_array(value, state, depth)
+      when Hash then count_hash(value, state, depth)
+      else raise Failure, "non_json_value"
+      end
+    end
+    private_class_method :count_json
+
+    def count_float(value, state)
+      raise Failure, "non_finite_number" unless value.finite?
+
+      consume(state, JSON.generate(value).bytesize)
+    end
+    private_class_method :count_float
+
+    def count_integer(value, state)
+      copy_integer(value)
+      upper_bound = integer_decimal_upper_bound(value)
+      # A bit length can straddle one decimal boundary. One byte of bounded
+      # slack lets the exact count decide without exposing an oversized allocation.
+      raise Failure, "too_large" if upper_bound > state[:limit] + 1
+
+      consume(state, value.to_s.bytesize)
+    end
+    private_class_method :count_integer
+
+    def integer_decimal_upper_bound(value)
+      digits = ((value.bit_length * 30_103) / 100_000) + 1
+      digits + (value.negative? ? 1 : 0)
+    end
+    private_class_method :integer_decimal_upper_bound
+
+    def count_string(value, state)
+      valid_utf8 = value.encoding == Encoding::UTF_8 && value.valid_encoding?
+      raise Failure, "invalid_encoding" unless valid_utf8
+
+      short_controls = value.count("\b\t\n\f\r")
+      controls = value.count("\u0000-\u001f")
+      escaped = value.count('"') + value.count("\\")
+      size = value.bytesize + 2 + escaped + short_controls + ((controls - short_controls) * 5)
+      consume(state, size)
+    end
+    private_class_method :count_string
+
+    def count_array(value, state, depth)
+      within(value, state) do
+        consume(state, 2)
+        value.each_with_index do |item, index|
+          consume(state, 1) unless index.zero?
+          count_json(item, state, depth + 1)
+        end
+      end
+    end
+    private_class_method :count_array
+
+    def count_hash(value, state, depth)
+      within(value, state) do
+        consume(state, 2)
+        value.each_with_index do |(key, item), index|
+          raise Failure, "invalid_key" unless key.is_a?(String)
+
+          consume(state, 1) unless index.zero?
+          count_string(key, state)
+          consume(state, 1)
+          count_json(item, state, depth + 1)
+        end
+      end
+    end
+    private_class_method :count_hash
+
+    def consume(state, bytes)
+      raise Failure, "too_large" if bytes > state[:remaining]
+
+      state[:remaining] -= bytes
+    end
+    private_class_method :consume
+  end # rubocop:enable Metrics/ModuleLength
+end

data/lib/mistri/tool_call.rb

@@ -1,18 +1,52 @@
 # frozen_string_literal: true
 
+require_relative "tool_arguments"
+
 module Mistri
-  # One tool invocation requested by the model. `arguments` is the parsed Hash
-  # exactly as the model emitted it, string keys and all; the top level is owned
-  # and frozen, nested values stay as parsed. `signature` is the opaque
-  # reasoning payload some providers attach to the call and require echoed back,
-  # such as Gemini's thoughtSignature.
-  ToolCall = Data.define(:id, :name, :arguments, :signature) do
-    def initialize(id:, name:, arguments: {}, signature: nil)
-      super(id:, name:, arguments: arguments.dup.freeze, signature:)
+  # One immutable tool invocation requested by the model. Completed arguments
+  # are deeply owned JSON; arguments_error records why an unsafe value became
+  # nil instead of letting it cross the execution boundary. Signature carries
+  # opaque replay state; provider_call_id distinguishes an optional wire ID
+  # from Mistri's always-present session correlation ID.
+  ToolCall = Data.define(:id, :name, :arguments, :signature, :arguments_error,
+                         :provider_call_id) do
+    def initialize(id:, name:, arguments: {}, signature: nil, arguments_error: nil,
+                   provider_call_id: nil, canonicalize: true)
+      id = id.dup.freeze if id.is_a?(String)
+      name = name.dup.freeze if name.is_a?(String)
+      signature = signature.dup.freeze if signature.is_a?(String)
+      provider_call_id = provider_call_id.dup.freeze if provider_call_id.is_a?(String)
+      @arguments_owned = canonicalize || !arguments_error.nil?
+      if arguments_error.nil?
+        arguments, arguments_error = ToolArguments.canonicalize(arguments) if canonicalize
+      else
+        arguments = nil
+        arguments_error = ToolArguments.normalize_error(arguments_error)
+      end
+      super(id:, name:, arguments:, signature:, arguments_error:, provider_call_id:)
     end
 
     def type = :tool_call
 
-    def to_h = { type: :tool_call, id:, name:, arguments:, signature: }.compact
+    def arguments_error? = !arguments_error.nil?
+
+    def arguments_owned? = @arguments_owned
+
+    # Data#with bypasses custom initializers; this value must re-enter the
+    # ownership boundary whenever a normalizer replaces its arguments.
+    def with(id: self.id, name: self.name, arguments: self.arguments,
+             signature: self.signature, arguments_error: self.arguments_error,
+             provider_call_id: self.provider_call_id)
+      self.class.new(id:, name:, arguments:, signature:, arguments_error:, provider_call_id:)
+    end
+
+    def to_h
+      hash = { type: :tool_call, id:, name: }.compact
+      hash[:arguments] = arguments
+      hash[:signature] = signature unless signature.nil?
+      hash[:arguments_error] = arguments_error unless arguments_error.nil?
+      hash[:provider_call_id] = provider_call_id unless provider_call_id.nil?
+      hash
+    end
   end
 end

data/lib/mistri/tool_context.rb

@@ -15,8 +15,10 @@ module Mistri
   #
   #   agent = Mistri.agent("claude-opus-4-8", tools: tools,
   #                        context: { traveler: current_traveler })
-  #   Mistri::Tool.define("book_hotel", "Books the chosen hotel.") do |args, context|
-  #     Bookings.create(args, traveler: context.app[:traveler])
+  #   Mistri::Tool.define("book_hotel", "Books the chosen hotel.", schema: -> {
+  #     string :hotel_id, "Hotel ID", required: true
+  #   }) do |args, context|
+  #     Bookings.create(args.fetch("hotel_id"), traveler: context.app[:traveler])
   #   end
   ToolContext = Data.define(:session, :signal, :emit, :app) do
     def initialize(session: nil, signal: nil, emit: nil, app: nil)