mistri 0.4.1 → 0.6.0

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

@@ -3,20 +3,22 @@
 module Mistri
   # What a tool handler may know about the run it executes inside: the
   # caller's session, the abort signal, the event stream, and the host's
-  # own context object. Handlers take it as an optional second argument — a
+  # own context object. Handlers take it as an optional second argument: a
   # proc ignores it invisibly, a lambda opts in by accepting two
   # parameters. Sub-agents are built on it; any tool that spawns work,
   # links records to the session, or streams progress can use it the same
   # way.
   #
-  # app carries whatever the host passes as Agent.new(context:) — the
-  # acting user, a tenant, a request — untouched. The gem provides the
+  # app carries whatever the host passes as Agent.new(context:) (the
+  # acting user, a tenant, a request), untouched. The gem provides the
   # slot, never the vocabulary:
   #
   #   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)

data/lib/mistri/tool_executor.rb

@@ -8,77 +8,168 @@ module Mistri
   # concurrently up to max_concurrency; each runs inside the Rails executor
   # when Rails is present, so ActiveRecord connections return to the pool.
   #
-  # A tool that raises becomes an in-band error string the model can read. An
-  # abort never starts a not-yet-started call: it gets an interrupted result
-  # instead, so the turn always pairs and the session replays cleanly.
+  # A tool that fails becomes an in-band ToolResult with an explicit error
+  # fact. An abort never starts a not-yet-started call: it gets an interrupted
+  # result instead, so the turn always pairs and the session replays cleanly.
   module ToolExecutor
+    # Separates Mistri's deadline from a handler's own Timeout::Error.
+    class InvocationTimeout < StandardError
+    end
+    private_constant :InvocationTimeout
+
     INTERRUPTED = "[interrupted: this tool call never ran]"
+    OUTCOME_UNKNOWN = "[interrupted: this tool call's outcome is unavailable; the tool may have " \
+                      "executed, so verify its effects before retrying]"
+    VERIFY_BEFORE_RETRY = "The tool may have completed partially; verify its effects before " \
+                          "retrying."
+    COMMITTED = Object.new.freeze
+    private_constant :COMMITTED
+
+    PreparedContext = Class.new(ToolContext) do
+      def arguments_prepared? = true
+    end
+    private_constant :PreparedContext
 
     module_function
 
     def call(calls, tools_by_name, signal: nil, max_concurrency: 4, session: nil, emit: nil,
              app: nil)
+      call_with_outcomes(
+        calls,
+        tools_by_name,
+        signal:,
+        max_concurrency:,
+        session:,
+        emit:,
+        app:,
+        prepared_arguments: false
+      ).map { |call, result, seconds, _committed| [call, result, seconds] }
+    end
+
+    # Agent prepares model arguments before approval and execution. This
+    # lower-level form preserves that boundary and exposes commitment so hooks
+    # cannot run for queued calls that an abort prevented from starting.
+    def call_with_outcomes(calls, tools_by_name, signal: nil, max_concurrency: 4, session: nil,
+                           emit: nil, app: nil, prepared_arguments: false)
       return [] if calls.empty?
 
-      context = ToolContext.new(session: session, signal: signal, emit: thread_safe(emit),
-                                app: app)
+      delivery = EventDelivery.wrap(emit, passthrough: [InvocationTimeout])
+      context_class = prepared_arguments ? PreparedContext : ToolContext
+      context = context_class.new(session: session, signal: signal,
+                                  emit: thread_safe(delivery, signal), app: app)
       results = Array.new(calls.length)
       queue = Queue.new
+      errors = Queue.new
       calls.each_with_index { |call, index| queue << [call, index] }
       workers = max_concurrency.clamp(1, calls.length)
-      Array.new(workers) { worker(queue, results, tools_by_name, context) }.each(&:join)
-      calls.zip(results).map do |call, entry|
-        value, seconds = entry || [INTERRUPTED, nil]
-        [call, value, seconds]
+      Array.new(workers) { worker(queue, results, tools_by_name, context, errors) }.each(&:join)
+      unless errors.empty?
+        error = errors.pop
+        raise EventDelivery.unwrap(error, delivery)
+      end
+
+      calls.each_with_index.map do |call, index|
+        entry = results[index]
+        entry = [failure(OUTCOME_UNKNOWN), nil, true] if entry.equal?(COMMITTED)
+        value, seconds, committed = entry || [failure(INTERRUPTED), nil, false]
+        [call, value, seconds, committed]
       end
     end
 
-    def worker(queue, results, tools_by_name, context)
+    def worker(queue, results, tools_by_name, context, errors)
       Thread.new do
         loop do
+          break unless errors.empty?
+
           call, index = begin
             queue.pop(true)
           rescue ThreadError
             break
           end
           if context.signal&.aborted?
-            results[index] = [INTERRUPTED, nil]
+            results[index] = [failure(INTERRUPTED), nil, false]
             next
           end
 
-          started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-          value = run_one(call, tools_by_name, context)
-          results[index] = [value, Process.clock_gettime(Process::CLOCK_MONOTONIC) - started]
+          results[index] = run_one(call, index, results, tools_by_name, context)
         end
+      rescue StandardError => e
+        errors << e
       end
     end
 
-    def run_one(call, tools_by_name, context)
+    def run_one(call, index, results, tools_by_name, context)
       tool = tools_by_name[call.name]
-      return "Error: unknown tool #{call.name.inspect}" unless tool
+      return [failure("Error: unknown tool #{call.name.inspect}"), nil, false] unless tool
+
+      return [failure(INTERRUPTED), nil, false] unless commit(call, context)
+
+      results[index] = COMMITTED
+      [*invoke_one(tool, call, context), true]
+    end
 
-      with_rails_executor { invoke(tool, call, context) }
+    def invoke_one(tool, call, context)
+      started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      value = with_rails_executor { invoke(tool, call, context) }
+      [value, elapsed(started)]
+    rescue EventDelivery::Failure
+      raise
+    rescue InvocationTimeout
+      content = "Error running tool #{call.name.inspect}: timed out after #{tool.timeout}s. " \
+                "#{VERIFY_BEFORE_RETRY}"
+      [failure(content), elapsed(started)]
     rescue StandardError => e
-      "Error running tool #{call.name.inspect}: #{e.class}: #{e.message}"
+      [failure("Error running tool #{call.name.inspect}: #{e.class}: #{e.message}. " \
+               "#{VERIFY_BEFORE_RETRY}"),
+       elapsed(started)]
     end
 
-    # A tool with a timeout answers in band when it stalls, so one hung
-    # handler cannot stall the whole run.
     def invoke(tool, call, context)
       return tool.call(call.arguments, context) unless tool.timeout
 
-      Timeout.timeout(tool.timeout) { tool.call(call.arguments, context) }
-    rescue Timeout::Error
-      "Error running tool #{call.name.inspect}: timed out after #{tool.timeout}s"
+      Timeout.timeout(tool.timeout, InvocationTimeout) do
+        tool.call(call.arguments, context)
+      end
+    end
+
+    def commit(call, context)
+      return false if context.signal&.aborted?
+      return true unless context.emit
+
+      context.emit.call(Event.new(type: :tool_started, tool_call: call))
+    end
+
+    def failure(content) = ToolResult.new(content:, error: true)
+
+    def elapsed(started)
+      started && (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started)
     end
 
     # Concurrent tools share the caller's sink; sinks are not required to be
     # thread-safe, so forwarded events serialize here.
-    def thread_safe(emit)
-      return nil unless emit
+    def thread_safe(delivery, signal)
+      return nil unless delivery
 
       mutex = Mutex.new
-      ->(event) { mutex.synchronize { emit.call(event) } }
+      # Already-committed workers may report after a sibling fails delivery;
+      # give all of them the first failure without calling the broken sink.
+      failure = nil
+      lambda do |event|
+        mutex.synchronize do
+          next false if event.type == :tool_started && (failure || signal&.aborted?)
+          raise failure if failure
+
+          result = begin
+            delivery.call(event)
+          rescue InvocationTimeout
+            raise
+          rescue EventDelivery::Failure => e
+            failure = e
+            raise
+          end
+          event.type == :tool_started ? true : result
+        end
+      end
     end
 
     def with_rails_executor(&)