mistri 0.4.1 → 0.6.0

data/lib/mistri/tool_result.rb

@@ -1,23 +1,28 @@
 # frozen_string_literal: true
 
 module Mistri
-  # A two-channel tool result: content goes to the model, ui goes only to the
-  # host. The ui payload rides the tool message and its :tool_result event,
-  # persists with the session for transcript re-renders, and never reaches a
-  # provider. Return one from a handler when the UI needs more than the model
-  # should read or pay for: full query rows behind a compact answer, the
-  # updated document behind "saved".
+  # A tool result with model content, host-only UI, and an explicit failure
+  # fact. Both channels and the error bit persist with the tool message; only
+  # content and provider-supported error signaling reach the model.
   #
-  #   Tool.define("edit_page", "Edits the page.") do |args|
-  #     page = apply(args)
+  #   Tool.define("edit_page", "Edits the page.", schema: -> {
+  #     object :changes, "Page changes", required: true
+  #   }) do |args|
+  #     page = apply(args.fetch("changes"))
   #     Mistri::ToolResult.new(content: "Updated.", ui: { "html" => page })
   #   end
   #
   # ui must be JSON-serializable; it is stored and delivered in canonical
   # JSON form (string keys), the same shape a reloaded session reads.
-  ToolResult = Data.define(:content, :ui) do
-    def initialize(content:, ui: nil)
+  ToolResult = Data.define(:content, :ui, :error) do
+    def initialize(content:, ui: nil, error: false)
+      unless [true, false].include?(error)
+        raise ArgumentError, "tool result error must be true or false"
+      end
+
       super
     end
+
+    def error? = error
   end
 end

data/lib/mistri/tools/edit_file.rb

@@ -2,6 +2,9 @@
 
 module Mistri
   module Tools
+    MAX_ATOMIC_EDIT_ATTEMPTS = 3
+    private_constant :MAX_ATOMIC_EDIT_ATTEMPTS
+
     module_function
 
     # The model-facing shape is flat {path, old_string, new_string,
@@ -10,28 +13,79 @@ module Mistri
     # back in band with the closest region and its exact difference, so the
     # model's retry is one shot.
     def edit_file(workspace)
+      atomic = atomic_workspace?(workspace)
       Tool.define("edit_file",
                   "Replace an exact snippet of a document. Copy old_string verbatim from " \
                   "read_file output including whitespace, without line-number prefixes. " \
                   "It must match exactly one place; add surrounding lines to make it " \
                   "unique, or set replace_all to change every occurrence.",
                   eager_input_streaming: true,
+                  argument_normalizer: Tools.method(:tolerate),
                   schema: lambda {
                     string :path, "Document path", required: true
                     string :old_string, "Exact text to replace (whitespace matters)", required: true
                     string :new_string, "Replacement text", required: true
                     boolean :replace_all, "Replace every occurrence instead of exactly one"
                   }) do |args|
-        args = Tools.tolerate(args)
-        with_document(workspace, args) do |content|
-          result = Edit.replace(content, args["old_string"], args["new_string"],
-                                replace_all: args["replace_all"] == true)
-          workspace.write(args["path"], result.content)
-          "Replaced #{result.count} occurrence(s) in #{args["path"]}"
+        result = if atomic
+                   replace_atomically(workspace, args)
+                 else
+                   replace_legacy(workspace, args)
+                 end
+        next result if result.is_a?(ToolResult)
+
+        "Replaced #{result.count} occurrence(s) in #{args["path"]}"
+      rescue EditError, WorkspaceConflictError => e
+        ToolResult.new(content: "edit_file failed: #{e.message}", error: true)
+      end
+    end
+
+    def replace_legacy(workspace, args)
+      with_document(workspace, args) do |content|
+        result = replacement(content, args)
+        workspace.write(args["path"], result.content)
+        result
+      end
+    end
+
+    def replace_atomically(workspace, args)
+      MAX_ATOMIC_EDIT_ATTEMPTS.times do |attempt|
+        snapshot = workspace.snapshot(args["path"])
+        return missing_document(args["path"]) unless snapshot
+        unless snapshot.is_a?(Workspace::Snapshot)
+          raise TypeError, "workspace snapshot must be a Mistri::Workspace::Snapshot"
+        end
+
+        result = replacement(snapshot.content, args)
+        begin
+          committed = workspace.compare_and_write(
+            args["path"], result.content, expected_revision: snapshot.revision
+          )
+          unless committed.is_a?(Workspace::Snapshot)
+            raise TypeError,
+                  "workspace compare_and_write must return a Mistri::Workspace::Snapshot"
+          end
+          unless same_content_bytes?(committed.content, result.content)
+            return ToolResult.new(
+              content: "The write to #{args["path"].inspect} committed, but storage " \
+                       "transformed the resulting document. Use read_file before continuing.",
+              error: true
+            )
+          end
+          return result
+        rescue WorkspaceConflictError
+          raise if attempt == MAX_ATOMIC_EDIT_ATTEMPTS - 1
         end
-      rescue EditError => e
-        "edit_file failed: #{e.message}"
       end
     end
+
+    def replacement(content, args)
+      Edit.replace(content, args["old_string"], args["new_string"],
+                   replace_all: args["replace_all"] == true)
+    end
+
+    def same_content_bytes?(left, right)
+      left == right || (left.bytesize == right.bytesize && left.b == right.b)
+    end
   end
 end

data/lib/mistri/tools.rb

@@ -6,6 +6,8 @@ module Mistri
   # row in a documents table, or an actual file. The names stay read_file and
   # edit_file because those are the tool names models are trained on.
   module Tools
+    ATOMIC_WORKSPACE_METHODS = %i[snapshot compare_and_write].freeze
+    private_constant :ATOMIC_WORKSPACE_METHODS
     ALIASES = { "oldText" => "old_string", "old" => "old_string", "search" => "old_string",
                 "newText" => "new_string", "new" => "new_string", "replace" => "new_string",
                 "replaceAll" => "replace_all", "file" => "path", "filename" => "path" }.freeze
@@ -23,15 +25,50 @@ module Mistri
 
     def with_document(workspace, args)
       content = workspace.read(args["path"])
-      return "No document at #{args["path"].inspect}. Use list_files to see paths." if content.nil?
+      return missing_document(args["path"]) if content.nil?
 
       yield content
     end
 
-    # Absorb the drift real models produce: alias keys, stringly booleans,
-    # unknown keys dropped by simply never being read.
+    def missing_document(path)
+      ToolResult.new(
+        content: "No document at #{path.inspect}. Use list_files to see paths.",
+        error: true
+      )
+    end
+
+    # Atomic writes are an explicit backend claim. A false or absent claim
+    # preserves the legacy four-method port; a true but incomplete claim fails
+    # before the model can rely on safety the backend does not implement.
+    def atomic_workspace?(workspace)
+      return false unless workspace.respond_to?(:atomic_writes?)
+
+      supported = workspace.atomic_writes?
+      unless [true, false].include?(supported)
+        raise ConfigurationError, "workspace atomic_writes? must return true or false"
+      end
+      return false unless supported
+
+      missing = ATOMIC_WORKSPACE_METHODS.reject { |method| workspace.respond_to?(method) }
+      unless missing.empty?
+        raise ConfigurationError,
+              "atomic workspace is missing #{missing.map(&:inspect).join(" and ")}"
+      end
+
+      true
+    end
+
+    # Absorb the drift real models produce for edit_file only. Ambiguous
+    # aliases fail instead of letting hash insertion order choose an edit.
     def tolerate(args)
-      normalized = args.to_h { |key, value| [ALIASES.fetch(key.to_s, key.to_s), value] }
+      normalized = args.each_with_object({}) do |(key, value), copy|
+        canonical = ALIASES.fetch(key.to_s, key.to_s)
+        if copy.key?(canonical)
+          raise ArgumentError, "multiple arguments map to #{canonical.inspect}"
+        end
+
+        copy[canonical] = value
+      end
       case normalized["replace_all"]
       when "true", "1", 1 then normalized["replace_all"] = true
       when "false", "0", 0, nil then normalized["replace_all"] = false

data/lib/mistri/transport.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "net/http"
+require "openssl"
 require "json"
 require "uri"
 
@@ -14,15 +15,66 @@ module Mistri
   # Streaming reads abort two ways: cooperatively between fragments, and hard,
   # by closing the socket from the abort signal's callback, so a stalled read
   # stops immediately instead of waiting out the read timeout.
+  # An address_resolver supplies a fresh validated set for each connection
+  # cycle; candidates are tried before the request is sent while the original
+  # hostname still owns Host, SNI, and TLS validation.
+  # JSON bodies and individual SSE lines share one configurable byte ceiling;
+  # a stream may contain any number of individually safe lines.
   class Transport
     KEEP_ALIVE_SECONDS = 30
+    ERROR_PREVIEW_BYTES = 500
+    BLANK_BODY = /\A[[:space:]]*\z/
+    CONNECT_ERRORS = [IOError, SocketError, SystemCallError, Timeout::Error,
+                      Net::HTTPBadResponse, OpenSSL::SSL::SSLError].freeze
+
+    # Net::HTTP reconnects expired keep-alives internally. Resolving inside
+    # connect makes every MCP connection cycle cross the egress boundary.
+    class ResolvedHTTP < Net::HTTP
+      attr_writer :address_resolver
+
+      private
+
+      def connect
+        addresses = Array(@address_resolver.call)
+        raise ConfigurationError, "address_resolver returned no addresses" if addresses.empty?
+
+        original_timeout = @open_timeout
+        deadline = if original_timeout
+                     Process.clock_gettime(Process::CLOCK_MONOTONIC) + original_timeout
+                   end
+        failure = nil
+        addresses.each do |address|
+          remaining = deadline && (deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC))
+          break if remaining && !remaining.positive?
+
+          @ipaddr = address
+          @open_timeout = remaining if remaining
+          begin
+            return super
+          rescue *CONNECT_ERRORS => e
+            failure = e
+          end
+        end
+        raise failure || Net::OpenTimeout.new("all approved addresses timed out")
+      ensure
+        @open_timeout = original_timeout if original_timeout
+      end
+    end
+    private_constant :ResolvedHTTP, :CONNECT_ERRORS, :BLANK_BODY
+
+    def initialize(origin:, open_timeout: 15, read_timeout: 300, write_timeout: 60,
+                   address_resolver: nil, max_record_bytes: DEFAULT_MAX_RECORD_BYTES)
+      unless max_record_bytes.is_a?(Integer) && max_record_bytes.positive?
+        raise ConfigurationError, "max_record_bytes: must be a positive integer"
+      end
 
-    def initialize(origin:, open_timeout: 15, read_timeout: 300, write_timeout: 60)
       @origin = origin.to_s.chomp("/")
       @uri = URI(@origin)
       @open_timeout = open_timeout
       @read_timeout = read_timeout
       @write_timeout = write_timeout
+      @address_resolver = address_resolver
+      @max_record_bytes = max_record_bytes
       @mutex = Mutex.new
       @connection = nil
     end
@@ -30,11 +82,19 @@ module Mistri
     # POST and decode a JSON response body. Retries once on a dead idle
     # socket, so it suits idempotent endpoints.
     def post(path, body:, headers: {})
-      response = @mutex.synchronize do
-        with_retry { connection.request(build_request(path, body, headers)) }
+      @mutex.synchronize do
+        with_retry do
+          parsed = nil
+          connection.request(build_request(path, body, headers)) do |response|
+            raise_for_status(response)
+            parsed = JSON.parse(read_json_body(response))
+          end
+          parsed
+        end
+      rescue ResponseTooLargeError, ResponseTooComplexError, ProviderError, JSON::ParserError
+        teardown
+        raise
       end
-      raise_for_status(response)
-      JSON.parse(response.body)
     end
 
     # POST and stream the SSE response, yielding each decoded data record.
@@ -47,32 +107,10 @@ module Mistri
 
     # POST for Streamable-HTTP endpoints (the MCP shape) that answer either
     # a JSON body or an SSE stream: yields each JSON record either way and
-    # returns the response headers, downcased. Retries only a dead idle
-    # socket that failed before any response started, so a side-effecting
-    # call can never run twice.
-    def post_either(path, body:, headers: {}, &block)
-      @mutex.synchronize do
-        retried = false
-        begin
-          started = false
-          response_headers = nil
-          connection.request(build_request(path, body, headers, streaming: true)) do |response|
-            started = true
-            raise_for_status(response)
-            response_headers = response.to_hash.transform_values(&:first)
-            read_either(response, &block)
-          end
-          response_headers
-        rescue IOError, SocketError, SystemCallError, Timeout::Error => e
-          teardown
-          if started || retried || e.is_a?(Timeout::Error)
-            raise ProviderError, "connection failed: #{e.message}"
-          end
-
-          retried = true
-          retry
-        end
-      end
+    # returns the response headers, downcased. A replayable request retries
+    # once when a dead idle socket fails before any response starts.
+    def post_either(path, body:, headers: {}, replayable: true, &block)
+      @mutex.synchronize { post_either_locked(path, body, headers, replayable, &block) }
     end
 
     def close
@@ -81,14 +119,44 @@ module Mistri
 
     private
 
+    def post_either_locked(path, body, headers, replayable, &block)
+      retried = false
+      begin
+        started = false
+        response_headers = nil
+        connection.request(build_request(path, body, headers, streaming: true)) do |response|
+          started = true
+          raise_for_status(response)
+          response_headers = response.to_hash.transform_values(&:first)
+          read_either(response, &block)
+        end
+        response_headers
+      rescue ResponseTooLargeError, ResponseTooComplexError, ProviderError
+        teardown
+        raise
+      rescue IOError, SocketError, SystemCallError, Timeout::Error, JSON::ParserError,
+             Net::HTTPBadResponse, OpenSSL::SSL::SSLError => e
+        teardown
+        unless replayable
+          raise AmbiguousDeliveryError, "#{AmbiguousDeliveryError.default_message}: #{e.message}"
+        end
+        if started || retried || e.is_a?(Timeout::Error)
+          raise ProviderError, "connection failed: #{e.message}"
+        end
+
+        retried = true
+        retry
+      end
+    end
+
     def read_either(response, &block)
       if response["content-type"].to_s.include?("text/event-stream")
-        sse = SSE.new
+        sse = SSE.new(max_record_bytes: @max_record_bytes)
         response.read_body { |chunk| sse.feed(chunk, &block) }
         sse.finish(&block)
       else
-        raw = response.read_body
-        block.call(JSON.parse(raw)) unless raw.to_s.strip.empty?
+        raw = read_json_body(response)
+        block.call(JSON.parse(raw)) unless BLANK_BODY.match?(raw)
       end
     end
 
@@ -115,6 +183,9 @@ module Mistri
         # than let the next request read stale frames.
         teardown if aborted
         aborted ? :aborted : nil
+      rescue EventDelivery::Failure, ResponseTooLargeError, ResponseTooComplexError, ProviderError
+        teardown
+        raise
       rescue IOError, SocketError, SystemCallError, Timeout::Error => e
         teardown
         return :aborted if signal&.aborted?
@@ -130,7 +201,7 @@ module Mistri
     end
 
     def read_stream(response, signal, &block)
-      sse = SSE.new
+      sse = SSE.new(max_record_bytes: @max_record_bytes)
       aborted = false
       response.read_body do |fragment|
         if signal&.aborted?
@@ -146,17 +217,32 @@ module Mistri
     def build_request(path, body, headers, streaming: false)
       request = Net::HTTP::Post.new(URI("#{@origin}#{path}"))
       request["Content-Type"] = "application/json"
-      if streaming
-        request["Accept"] = "text/event-stream"
-        # Net::HTTP silently negotiates gzip, and its inflater buffers the
-        # whole stream, delivering "live" events in one burst at the end.
-        request["Accept-Encoding"] = "identity"
-      end
+      request["Accept"] = "text/event-stream" if streaming
       headers.each { |key, value| request[key] = value }
+      # Identity keeps the byte ceiling meaningful before an untrusted
+      # compressed expansion and preserves immediate SSE delivery.
+      request["Accept-Encoding"] = "identity"
       request.body = JSON.generate(body)
       request
     end
 
+    def read_json_body(response)
+      declared = response["content-length"]
+      if declared&.match?(/\A\d+\z/) && declared.to_i > @max_record_bytes
+        raise ResponseTooLargeError.new(kind: :json_body, limit: @max_record_bytes)
+      end
+
+      body = +""
+      response.read_body do |fragment|
+        if body.bytesize + fragment.bytesize > @max_record_bytes
+          raise ResponseTooLargeError.new(kind: :json_body, limit: @max_record_bytes)
+        end
+
+        body << fragment
+      end
+      body
+    end
+
     def with_retry
       attempted = false
       begin
@@ -176,13 +262,21 @@ module Mistri
     end
 
     def connection
-      @connection ||= Net::HTTP.new(@uri.host, @uri.port).tap do |http|
+      @connection ||= begin
+        http = if @address_resolver
+                 ResolvedHTTP.new(@uri.hostname, @uri.port, nil).tap do |resolved|
+                   resolved.address_resolver = @address_resolver
+                 end
+               else
+                 Net::HTTP.new(@uri.hostname, @uri.port)
+               end
         http.use_ssl = @uri.scheme == "https"
         http.open_timeout = @open_timeout
         http.read_timeout = @read_timeout
         http.write_timeout = @write_timeout
         http.keep_alive_timeout = KEEP_ALIVE_SECONDS
         http.start
+        http
       end
     end
 
@@ -204,10 +298,21 @@ module Mistri
       status = response.code.to_i
       return if (200..299).cover?(status)
 
+      preview = +""
+      response.read_body do |fragment|
+        remaining = ERROR_PREVIEW_BYTES - preview.bytesize
+        preview << fragment.byteslice(0, remaining) if remaining.positive?
+        raise status_error(response, status, preview) if preview.bytesize >= ERROR_PREVIEW_BYTES
+      end
+      raise status_error(response, status, preview)
+    end
+
+    def status_error(response, status, preview)
       klass = error_class(status)
-      options = { status: status, body: response.read_body.to_s[0, 500] }
+      body = preview.dup.force_encoding(Encoding::UTF_8).scrub("?")
+      options = { status: status, body: body }
       options[:retry_after] = retry_after(response) if klass == RateLimitError
-      raise klass.new(**options)
+      klass.new(**options)
     end
 
     def error_class(status)

data/lib/mistri/usage.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "json"
+
 module Mistri
   # Token accounting for an assistant turn, and the dollar cost of those tokens.
   #
@@ -9,37 +11,77 @@ module Mistri
   # of `cache_write` held for an hour, which bills at twice the input rate.
   class Usage < Data.define(:input, :output, :cache_read, :cache_write,
                             :cache_write_1h, :reasoning, :cost)
+    # Dollar amounts plus whether every contributing token had known pricing.
+    # Known stays outside the released Data shape so pattern matching and
+    # positional construction remain compatible.
     Cost = Data.define(:input, :output, :cache_read, :cache_write, :total) do
-      def self.zero = new(input: 0.0, output: 0.0, cache_read: 0.0, cache_write: 0.0, total: 0.0)
+      def initialize(input:, output:, cache_read:, cache_write:, total:, known: true)
+        @known = known == true
+        super(input:, output:, cache_read:, cache_write:, total:)
+      end
+
+      def self.zero
+        new(input: 0.0, output: 0.0, cache_read: 0.0, cache_write: 0.0,
+            total: 0.0, known: true)
+      end
+
+      def self.unknown
+        new(input: 0.0, output: 0.0, cache_read: 0.0, cache_write: 0.0,
+            total: 0.0, known: false)
+      end
+
+      def known? = defined?(@known) ? @known : true
+
+      def with(**changes)
+        self.class.new(**to_h, **changes)
+      end
 
       def +(other)
         self.class.new(input: input + other.input, output: output + other.output,
                        cache_read: cache_read + other.cache_read,
                        cache_write: cache_write + other.cache_write,
-                       total: total + other.total)
+                       total: total + other.total, known: known? && other.known?)
       end
+
+      def ==(other) = super && known? == other.known?
+
+      alias_method :eql?, :==
+
+      def hash = [super, known?].hash
+
+      def to_h = super.merge(known: known?)
+
+      def _dump(_level) = JSON.generate(to_h)
+
+      def self._load(payload) = new(**JSON.parse(payload, symbolize_names: true))
     end
 
     def initialize(input: 0, output: 0, cache_read: 0, cache_write: 0,
-                   cache_write_1h: 0, reasoning: 0, cost: Cost.zero)
+                   cache_write_1h: 0, reasoning: 0, cost: Cost.unknown)
       super
     end
 
-    def self.zero = new
+    def self.zero = new(cost: Cost.zero)
 
     def self.from_h(hash)
       h = (hash || {}).transform_keys(&:to_s)
       c = (h["cost"] || {}).transform_keys(&:to_s)
-      new(input: h.fetch("input", 0).to_i, output: h.fetch("output", 0).to_i,
-          cache_read: h.fetch("cache_read", 0).to_i, cache_write: h.fetch("cache_write", 0).to_i,
-          cache_write_1h: h.fetch("cache_write_1h", 0).to_i,
-          reasoning: h.fetch("reasoning", 0).to_i,
-          cost: Cost.new(input: c.fetch("input", 0).to_f, output: c.fetch("output", 0).to_f,
-                         cache_read: c.fetch("cache_read", 0).to_f,
-                         cache_write: c.fetch("cache_write", 0).to_f,
-                         total: c.fetch("total", 0).to_f))
+      counts = { input: h.fetch("input", 0).to_i, output: h.fetch("output", 0).to_i,
+                 cache_read: h.fetch("cache_read", 0).to_i,
+                 cache_write: h.fetch("cache_write", 0).to_i,
+                 cache_write_1h: h.fetch("cache_write_1h", 0).to_i,
+                 reasoning: h.fetch("reasoning", 0).to_i }
+      amounts = { input: c.fetch("input", 0).to_f, output: c.fetch("output", 0).to_f,
+                  cache_read: c.fetch("cache_read", 0).to_f,
+                  cache_write: c.fetch("cache_write", 0).to_f,
+                  total: c.fetch("total", 0).to_f }
+      known = c.fetch("known", false) == true
+      known ||= !c.key?("known") && counts.values.all?(&:zero?) && amounts.values.all?(&:zero?)
+      new(**counts, cost: Cost.new(**amounts, known:))
     end
 
+    def prompt_tokens = input + cache_read + cache_write
+
     def total_tokens = input + output + cache_read + cache_write
 
     # A copy with cost computed from per-million-token rates. The 1h cache-write
@@ -54,7 +96,9 @@ module Mistri
         cache_read: rate(rates, :cache_read) * cache_read,
         cache_write: (rate(rates, :cache_write) * short) + (long_write_rate(rates) * long)
       }
-      with(cost: Cost.new(**computed, total: computed.values.sum))
+      with(cost: Cost.new(**computed,
+                          total: computed.values.sum,
+                          known: rates_cover?(rates, short, long)))
     end
 
     def +(other)
@@ -75,5 +119,13 @@ module Mistri
     def long_write_rate(rates)
       rates.key?(:cache_write_1h) ? rate(rates, :cache_write_1h) : rate(rates, :input) * 2
     end
+
+    def rates_cover?(rates, short, long)
+      (input.zero? || rates.key?(:input)) &&
+        (output.zero? || rates.key?(:output)) &&
+        (cache_read.zero? || rates.key?(:cache_read)) &&
+        (short.zero? || rates.key?(:cache_write)) &&
+        (long.zero? || rates.key?(:cache_write_1h) || rates.key?(:input))
+    end
   end
 end

data/lib/mistri/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Mistri
-  VERSION = "0.4.1"
+  VERSION = "0.6.0"
 end