mistri 0.4.1 → 0.6.0

data/lib/mistri/mcp/wires.rb

@@ -13,22 +13,34 @@ module Mistri
       # Streamable HTTP: requests POST to one endpoint, responses arrive as
       # JSON or an SSE stream. Sessions and bearer auth live here.
       class Http
-        def initialize(url:, token:, headers:, open_timeout:, read_timeout:)
-          uri = URI(url)
+        HEADER_NAME = /\A[!#$%&'*+\-.^_`|~0-9A-Za-z]+\z/
+        INVALID_HEADER_VALUE = /[\x00-\x1F\x7F]/n
+        RESERVED_HEADERS = %w[host content-length transfer-encoding connection keep-alive
+                              proxy-authenticate proxy-authorization proxy-connection te trailer
+                              upgrade accept content-type accept-encoding mcp-session-id
+                              mcp-protocol-version].freeze
+        private_constant :HEADER_NAME, :INVALID_HEADER_VALUE, :RESERVED_HEADERS
+
+        def initialize(url:, token:, headers:, open_timeout:, read_timeout:, allow_non_public:,
+                       max_record_bytes: DEFAULT_MAX_RECORD_BYTES)
+          uri, resolver = Egress.resolver(url, allow_non_public:, label: "MCP URL",
+                                               timeout: open_timeout)
           @path = uri.path.empty? ? "/" : uri.path
           @path = "#{@path}?#{uri.query}" if uri.query
-          @transport = Transport.new(origin: "#{uri.scheme}://#{uri.host}:#{uri.port}",
-                                     open_timeout: open_timeout, read_timeout: read_timeout)
+          @transport = Transport.new(origin: Egress.origin(uri), address_resolver: resolver,
+                                     open_timeout: open_timeout, read_timeout: read_timeout,
+                                     max_record_bytes: max_record_bytes)
           @token = token
-          @headers = headers
+          @headers = normalized_headers(headers)
           @session_id = nil
           @protocol_version = nil
         end
 
         attr_writer :protocol_version
 
-        def call(payload, &)
-          meta = @transport.post_either(@path, body: payload, headers: request_headers, &)
+        def call(payload, replayable:, &)
+          meta = @transport.post_either(@path, body: payload, headers: request_headers,
+                                               replayable: replayable, &)
           capture_session(meta)
           nil
         end
@@ -43,12 +55,33 @@ module Mistri
 
         def refreshable? = @token.respond_to?(:call)
 
-        def reset_session = @session_id = nil
+        def reset_session
+          @session_id = nil
+          @protocol_version = nil
+        end
 
         def close = @transport.close
 
         private
 
+        def normalized_headers(headers)
+          headers.each_with_object({}) do |(key, value), copy|
+            name = key.to_s
+            content = value.to_s
+            unless name.ascii_only? && HEADER_NAME.match?(name)
+              raise ConfigurationError, "#{name.inspect} is not a valid HTTP header name"
+            end
+            if RESERVED_HEADERS.include?(name.downcase)
+              raise ConfigurationError, "#{name} is managed by the MCP transport"
+            end
+            if INVALID_HEADER_VALUE.match?(content.b)
+              raise ConfigurationError, "#{name} contains invalid control characters"
+            end
+
+            copy[name.dup.freeze] = content.dup.freeze
+          end.freeze
+        end
+
         def request_headers
           headers = { "Accept" => "application/json, text/event-stream" }
           headers.merge!(@headers)
@@ -72,14 +105,24 @@ module Mistri
       # credentials in its environment, as the spec prescribes for local
       # servers. Its stderr stays attached for honest local debugging.
       class Stdio
-        def initialize(command:, env: {}, read_timeout: 120)
+        READ_SIZE = 16 * 1024
+        BLANK_RECORD = /\A[[:space:]]*\z/
+        private_constant :READ_SIZE, :BLANK_RECORD
+
+        def initialize(command:, env: {}, read_timeout: 120,
+                       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
+
           @command = Array(command).map(&:to_s)
           @env = env.transform_keys(&:to_s).transform_values(&:to_s)
           @read_timeout = read_timeout
+          @max_record_bytes = max_record_bytes
           @pid = nil
         end
 
-        def call(payload)
+        def call(payload, **)
           spawn_server unless @pid
           write(payload)
           loop do
@@ -105,11 +148,13 @@ module Mistri
         def protocol_version=(_version); end
 
         def close
-          return unless @pid
-
-          [@stdin, @stdout].each { |io| io.close unless io.closed? }
-          terminate
-          @pid = nil
+          if @pid
+            [@stdin, @stdout].each { |io| io.close unless io.closed? }
+            terminate
+            @pid = nil
+          end
+          @read_buffer&.clear
+          nil
         end
 
         private
@@ -118,6 +163,8 @@ module Mistri
           child_in, @stdin = IO.pipe
           @stdout, child_out = IO.pipe
           @pid = Process.spawn(@env, *@command, in: child_in, out: child_out)
+          @read_buffer = +"".b
+          @newline_search_offset = 0
           child_in.close
           child_out.close
         end
@@ -126,24 +173,69 @@ module Mistri
           @stdin.write("#{JSON.generate(payload)}\n")
           @stdin.flush
         rescue Errno::EPIPE
-          raise Error, "the MCP server closed its input"
+          raise WireError, "the MCP server closed its input"
         end
 
         # The spec requires stdout to carry only protocol messages, so a
         # line that is not one is corruption worth failing loudly on.
         def read_record
+          deadline = if @read_timeout
+                       Process.clock_gettime(Process::CLOCK_MONOTONIC) + @read_timeout
+                     end
           loop do
-            ready = @stdout.wait_readable(@read_timeout)
-            raise Error, "timed out waiting for the MCP server" unless ready
+            line = complete_record
+            return JSON.parse(line) if line
 
-            line = @stdout.gets
-            raise Error, "the MCP server exited" if line.nil?
-            next if line.strip.empty?
+            oversized_record! if @read_buffer.bytesize > @max_record_bytes
+            chunk = read_chunk(deadline)
+            return JSON.parse(eof_record) unless chunk
 
-            return JSON.parse(line)
+            @read_buffer << chunk
           end
         rescue JSON::ParserError
-          raise Error, "the MCP server wrote non-protocol output on stdout"
+          raise WireError, "the MCP server wrote non-protocol output on stdout"
+        end
+
+        def complete_record
+          while (newline = @read_buffer.index("\n", @newline_search_offset))
+            oversized_record! if newline > @max_record_bytes
+            line = @read_buffer.slice!(0, newline + 1)
+            @newline_search_offset = 0
+            return line unless BLANK_RECORD.match?(line)
+          end
+          @newline_search_offset = @read_buffer.bytesize
+          nil
+        end
+
+        def read_chunk(deadline)
+          loop do
+            remaining = deadline && (deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC))
+            if remaining && !remaining.positive?
+              raise WireError, "timed out waiting for the MCP server"
+            end
+
+            ready = @stdout.wait_readable(remaining)
+            raise WireError, "timed out waiting for the MCP server" unless ready
+
+            room = @max_record_bytes - @read_buffer.bytesize + 1
+            chunk = @stdout.read_nonblock([READ_SIZE, room].min, exception: false)
+            return chunk unless chunk == :wait_readable
+          end
+        end
+
+        def eof_record
+          raise WireError, "the MCP server exited" if @read_buffer.empty?
+
+          line = @read_buffer.slice!(0, @read_buffer.bytesize)
+          raise WireError, "the MCP server exited" if BLANK_RECORD.match?(line)
+
+          line
+        end
+
+        def oversized_record!
+          error = ResponseTooLargeError.new(kind: :stdio_record, limit: @max_record_bytes)
+          close
+          raise error
         end
 
         def terminate

data/lib/mistri/mcp.rb

@@ -5,7 +5,7 @@ require "json"
 module Mistri
   # Bridge Model Context Protocol servers into Mistri tools: list a server's
   # tools, hand them to an agent, and everything the harness already does
-  # composes — approval gates on third-party write tools, retries, sub-agent
+  # composes: approval gates on third-party write tools, retries, sub-agent
   # pools, the ui channel.
   #
   #   client = Mistri::MCP::Client.new(url: "https://mcp.linear.app/mcp",
@@ -17,6 +17,10 @@ module Mistri
   # {"name", "description", "inputSchema"} hashes) and call_tool(name, args)
   # bridges the same way, so the official mcp gem's client plugs in too.
   module MCP
+    EMPTY_INPUT_SCHEMA = {
+      "type" => "object", "properties" => {}.freeze, "additionalProperties" => false
+    }.freeze
+
     # A protocol-level failure: a JSON-RPC error, a missing response, an
     # unsupported negotiation.
     class Error < Mistri::Error
@@ -28,42 +32,69 @@ module Mistri
       end
     end
 
+    # The stdio peer failed while a protocol message was in flight.
+    class WireError < Error; end
+
     # The server expired this client's session (a 404 with a session
     # attached); the spec says start a fresh one, and Client does.
     class SessionExpired < Error; end
 
+    # A remote URL cannot be used without weakening the MCP network boundary.
+    class UnsafeURLError < ConfigurationError; end
+
     module_function
 
     # The server's tools as Mistri tools. allow/deny filter by remote name,
     # prefix namespaces local names ("linear__create_issue") against
     # collisions, and gates marks tools needing human approval
     # (gates: { "create_issue" => true }, or needs_approval: for all).
-    def tools(client, allow: nil, deny: [], prefix: nil, needs_approval: false, gates: {})
+    def tools(client, allow: nil, deny: [], prefix: nil, needs_approval: false, gates: {},
+              complete_argument_validator: nil)
       listed = client.tools
       listed = listed.select { |tool| allow.include?(tool["name"]) } if allow
       listed = listed.reject { |tool| deny.include?(tool["name"]) }
       listed.map do |tool|
-        bridge(client, tool, prefix: prefix, gate: gates.fetch(tool["name"], needs_approval))
+        bridge(client, tool, prefix: prefix, gate: gates.fetch(tool["name"], needs_approval),
+                             complete_argument_validator: complete_argument_validator)
       end
     end
 
-    def bridge(client, spec, prefix: nil, gate: false)
+    def bridge(client, spec, prefix: nil, gate: false, complete_argument_validator: nil)
       remote = spec.fetch("name")
       local = prefix ? "#{prefix}__#{remote}" : remote
+      if spec.key?("inputSchema") && spec["inputSchema"].nil?
+        raise ConfigurationError, "inputSchema must be an object, not null"
+      end
+
+      raw_schema = spec.fetch("inputSchema", EMPTY_INPUT_SCHEMA)
+      input_schema = Schema.validate_mcp!(
+        raw_schema, complete: !complete_argument_validator.nil?
+      )
       Tool.define(local, spec["description"].to_s,
-                  input_schema: spec["inputSchema"] || Tool::EMPTY_SCHEMA,
-                  needs_approval: gate) do |args|
+                  input_schema: input_schema,
+                  needs_approval: gate,
+                  complete_argument_validator: complete_argument_validator) do |args|
         answer(client.call_tool(remote, args || {}))
       end
+    rescue ConfigurationError => e
+      raise ConfigurationError, "MCP tool #{remote.inspect}: #{e.message}"
     end
 
     # An MCP result becomes model-readable content: text joins, images ride
-    # as image blocks, and isError answers in band so the model can react.
+    # as image blocks, and isError remains structured failure data end to end.
     def answer(result)
       blocks = Array(result["content"]).map { |block| convert(block) }
       if result["isError"]
-        text = blocks.grep(String).join("\n")
-        return "MCP tool error: #{text.empty? ? "unknown error" : text}"
+        text_blocks = blocks.grep(String)
+        if result["structuredContent"]
+          structured = JSON.generate(result["structuredContent"])
+          text_blocks << structured unless text_blocks.include?(structured)
+        end
+        text = text_blocks.join("\n")
+        content = "MCP tool error: #{text.empty? ? "unknown error" : text}"
+        non_text = blocks.grep_v(String)
+        content = [content, *non_text] unless non_text.empty?
+        return ToolResult.new(content:, error: true)
       end
       if blocks.empty? && result["structuredContent"]
         return JSON.generate(result["structuredContent"])
@@ -91,6 +122,9 @@ module Mistri
   end
 end
 
+require_relative "mcp/egress"
 require_relative "mcp/wires"
 require_relative "mcp/client"
 require_relative "mcp/oauth"
+
+Mistri::MCP.private_constant :Egress, :WireError, :EMPTY_INPUT_SCHEMA

data/lib/mistri/message.rb

@@ -10,26 +10,33 @@ module Mistri
   # them, which is what lets a later turn replay history across models, plus
   # usage and the stop reason.
   #
-  # ui is a tool result's host-only channel: it persists with the message and
-  # rides its :tool_result event, but no serializer ever sends it to a model.
-  # error is the machine-readable failure on an errored turn (ErrorData
-  # shape); error_message stays the human story.
+  # ui is a tool result's host-only channel. tool_error is true or false on
+  # newly produced tool results and nil on legacy records that predate the
+  # field. error is the machine-readable failure on an errored provider turn
+  # (ErrorData shape); error_message stays the human story.
   class Message < Data.define(:role, :content, :tool_call_id, :tool_name,
                               :model, :provider, :usage, :stop_reason, :error_message, :ui,
-                              :error)
+                              :error, :tool_error)
     ROLES = %i[system user assistant tool].freeze
 
     def initialize(role:, content: nil, tool_call_id: nil, tool_name: nil, model: nil,
                    provider: nil, usage: nil, stop_reason: nil, error_message: nil, ui: nil,
-                   error: nil)
+                   error: nil, tool_error: nil)
       role = role.to_sym
       raise ArgumentError, "unknown role #{role.inspect}" unless ROLES.include?(role)
       if stop_reason && !StopReason.valid?(stop_reason)
         raise ArgumentError, "unknown stop reason #{stop_reason.inspect}"
       end
 
+      unless tool_error.nil?
+        unless [true, false].include?(tool_error)
+          raise ArgumentError, "tool_error must be true, false, or nil"
+        end
+        raise ArgumentError, "tool_error is only valid on tool messages" if role != :tool
+      end
+
       super(role:, content: Content.wrap(content).freeze, tool_call_id:, tool_name:,
-            model:, provider:, usage:, stop_reason:, error_message:, ui:, error:)
+            model:, provider:, usage:, stop_reason:, error_message:, ui:, error:, tool_error:)
     end
 
     def self.system(content) = new(role: :system, content:)
@@ -50,8 +57,8 @@ module Mistri
       new(role: :assistant, content: [*Content.wrap(content), *tool_calls], **meta)
     end
 
-    def self.tool(content:, tool_call_id:, tool_name: nil, ui: nil)
-      new(role: :tool, content:, tool_call_id:, tool_name:, ui:)
+    def self.tool(content:, tool_call_id:, tool_name: nil, ui: nil, tool_error: false)
+      new(role: :tool, content:, tool_call_id:, tool_name:, ui:, tool_error:)
     end
 
     def self.from_h(hash)
@@ -62,13 +69,15 @@ module Mistri
           model: h["model"], provider: h["provider"]&.to_sym,
           usage: h["usage"] && Usage.from_h(h["usage"]),
           stop_reason: h["stop_reason"]&.to_sym, error_message: h["error_message"],
-          ui: h["ui"], error: h["error"])
+          ui: h["ui"], error: h["error"], tool_error: h["tool_error"])
     end
 
     def system? = role == :system
     def user? = role == :user
     def assistant? = role == :assistant
     def tool? = role == :tool
+    def tool_error? = tool_error == true
+    def tool_error_known? = !tool_error.nil?
 
     # Every Text block joined, or nil when the turn carried no text.
     def text
@@ -84,7 +93,8 @@ module Mistri
     # never with new(**to_h).
     def to_h
       { role:, content: content.map(&:to_h), tool_call_id:, tool_name:, model:,
-        provider:, usage: usage&.to_h, stop_reason:, error_message:, ui:, error: }.compact
+        provider:, usage: usage&.to_h, stop_reason:, error_message:, ui:, error:,
+        tool_error: }.compact
     end
   end
 end

data/lib/mistri/models.rb

@@ -1,43 +1,181 @@
 # frozen_string_literal: true
 
+require "json"
+
 module Mistri
   # The model catalog: capability data for known models, with graceful
   # passthrough for unknown ones. An id missing here still works everywhere;
-  # the catalog only improves defaults (output ceilings, provider inference),
-  # so a brand-new model is usable the day it ships.
+  # the catalog supplies capability defaults and paid direct-API pricing, so a
+  # brand-new model remains usable unless the host requests a cost ceiling.
   module Models
+    EMPTY_PRICING = [].freeze
+
+    # One request's standard rates and optional prompt-size tier.
+    RateCard = Data.define(:rates, :above, :higher) do
+      def initialize(rates:, above: nil, higher: nil)
+        if above.nil? != higher.nil?
+          raise ArgumentError, "a pricing threshold needs both above and higher"
+        end
+
+        super(rates: rates.freeze, above:, higher: higher&.freeze)
+      end
+
+      def rates_for(usage)
+        usage && above && usage.prompt_tokens > above ? higher : rates
+      end
+    end
+
+    # A rate card that becomes active at one absolute instant.
+    Schedule = Data.define(:effective_at, :card)
+
     # thinking is how the model accepts a reasoning request: :adaptive (the
     # model decides), :budget (a token budget), or :effort. It is what keeps
     # a provider from sending an unsupported thinking shape that 400s.
-    Model = Data.define(:id, :provider, :max_output, :context_window, :thinking)
+    #
+    # context_window is the provider's published context or input limit.
+    # Anthropic and OpenAI share it with the next output; Gemini publishes a
+    # separate inputTokenLimit and outputTokenLimit.
+    #
+    # Pricing is selected per request from its prompt size. Rates are paid
+    # standard direct-API list prices. cache_write is the provider's baseline
+    # write rate; Usage prices a reported longer-retention slice separately.
+    Model = Data.define(:id, :provider, :max_output, :context_window, :thinking) do
+      def initialize(id:, provider:, max_output:, context_window:, thinking:, pricing: [])
+        @pricing = pricing.sort_by(&:effective_at).freeze
+        super(id:, provider:, max_output:, context_window:, thinking:)
+      end
+
+      def rates(usage: nil, at: Time.now)
+        pricing.reverse_each.find { |entry| entry.effective_at <= at }&.card&.rates_for(usage)
+      end
+
+      def priced? = !pricing.empty?
+
+      def with(**changes)
+        updated_pricing = changes.delete(:pricing) { pricing }
+        self.class.new(**to_h, **changes, pricing: updated_pricing)
+      end
+
+      def ==(other) = super && pricing == other.send(:pricing)
+
+      alias_method :eql?, :==
+
+      def hash = [super, pricing].hash
+
+      def _dump(_level)
+        serialized = pricing.map do |entry|
+          { effective_at: entry.effective_at.to_i, rates: entry.card.rates,
+            above: entry.card.above, higher: entry.card.higher }
+        end
+        JSON.generate(to_h.merge(pricing: serialized))
+      end
+
+      def self._load(payload)
+        attributes = JSON.parse(payload, symbolize_names: true)
+        serialized = attributes.delete(:pricing)
+        pricing = serialized.map do |entry|
+          card = RateCard.new(rates: entry[:rates], above: entry[:above],
+                              higher: entry[:higher])
+          Schedule.new(effective_at: Time.at(entry[:effective_at]).utc, card:)
+        end
+        attributes[:provider] = attributes[:provider].to_sym
+        attributes[:thinking] = attributes[:thinking].to_sym
+        new(**attributes, pricing:)
+      end
+
+      private
+
+      def pricing = @pricing || EMPTY_PRICING
+    end
+
+    EPOCH = Time.at(0).utc.freeze
+
+    def self.price(rates = nil, from: EPOCH, above: nil, higher: nil, **flat_rates)
+      rates ||= flat_rates
+      Schedule.new(effective_at: from,
+                   card: RateCard.new(rates:, above:, higher:))
+    end
+    private_class_method :price
 
     CATALOG = [
-      ["claude-fable-5", :anthropic, 128_000, 200_000, :adaptive],
-      ["claude-opus-4-8", :anthropic, 128_000, 200_000, :adaptive],
-      ["claude-opus-4-7", :anthropic, 128_000, 200_000, :adaptive],
-      ["claude-opus-4-6", :anthropic, 128_000, 200_000, :adaptive],
-      ["claude-sonnet-5", :anthropic, 128_000, 200_000, :adaptive],
-      ["claude-sonnet-4-6", :anthropic, 128_000, 200_000, :adaptive],
-      ["claude-haiku-4-5", :anthropic, 64_000, 200_000, :budget],
-      ["gpt-5.5", :openai, 128_000, 400_000, :effort],
-      ["gpt-5.4", :openai, 128_000, 400_000, :effort],
-      ["gpt-5-nano", :openai, 128_000, 400_000, :effort],
-      ["gemini-3.5-flash", :gemini, 65_536, 1_048_576, :level],
-      ["gemini-3.1-pro-preview", :gemini, 65_536, 1_048_576, :level],
-      ["gemini-2.5-pro", :gemini, 65_536, 1_048_576, :budget],
-      ["gemini-2.5-flash", :gemini, 65_536, 1_048_576, :budget]
-    ].to_h do |id, provider, max_output, context_window, thinking|
-      [id, Model.new(id:, provider:, max_output:, context_window:, thinking:)]
+      ["claude-fable-5", :anthropic, 128_000, 1_000_000, :adaptive,
+       [price(input: 10.0, output: 50.0, cache_read: 1.0, cache_write: 12.5)]],
+      ["claude-opus-4-8", :anthropic, 128_000, 1_000_000, :adaptive,
+       [price(input: 5.0, output: 25.0, cache_read: 0.5, cache_write: 6.25)]],
+      ["claude-opus-4-7", :anthropic, 128_000, 1_000_000, :adaptive,
+       [price(input: 5.0, output: 25.0, cache_read: 0.5, cache_write: 6.25)]],
+      ["claude-opus-4-6", :anthropic, 128_000, 1_000_000, :adaptive,
+       [price(input: 5.0, output: 25.0, cache_read: 0.5, cache_write: 6.25)]],
+      ["claude-sonnet-5", :anthropic, 128_000, 1_000_000, :adaptive,
+       [price(input: 2.0, output: 10.0, cache_read: 0.2, cache_write: 2.5),
+        price({ input: 3.0, output: 15.0, cache_read: 0.3, cache_write: 3.75 },
+              from: Time.utc(2026, 9, 1))]],
+      ["claude-sonnet-4-6", :anthropic, 128_000, 1_000_000, :adaptive,
+       [price(input: 3.0, output: 15.0, cache_read: 0.3, cache_write: 3.75)]],
+      ["claude-haiku-4-5", :anthropic, 64_000, 200_000, :budget,
+       [price(input: 1.0, output: 5.0, cache_read: 0.1, cache_write: 1.25)]],
+      ["gpt-5.6-sol", :openai, 128_000, 1_050_000, :effort,
+       [price({ input: 5.0, output: 30.0, cache_read: 0.5, cache_write: 6.25 },
+              above: 272_000,
+              higher: { input: 10.0, output: 45.0, cache_read: 1.0,
+                        cache_write: 12.5 })]],
+      ["gpt-5.6-terra", :openai, 128_000, 1_050_000, :effort,
+       [price({ input: 2.5, output: 15.0, cache_read: 0.25, cache_write: 3.125 },
+              above: 272_000,
+              higher: { input: 5.0, output: 22.5, cache_read: 0.5,
+                        cache_write: 6.25 })]],
+      ["gpt-5.6-luna", :openai, 128_000, 1_050_000, :effort,
+       [price({ input: 1.0, output: 6.0, cache_read: 0.1, cache_write: 1.25 },
+              above: 272_000,
+              higher: { input: 2.0, output: 9.0, cache_read: 0.2,
+                        cache_write: 2.5 })]],
+      ["gpt-5.5", :openai, 128_000, 1_050_000, :effort,
+       [price({ input: 5.0, output: 30.0, cache_read: 0.5 },
+              above: 272_000, higher: { input: 10.0, output: 45.0, cache_read: 1.0 })]],
+      ["gpt-5.4", :openai, 128_000, 1_050_000, :effort,
+       [price({ input: 2.5, output: 15.0, cache_read: 0.25 },
+              above: 272_000, higher: { input: 5.0, output: 22.5, cache_read: 0.5 })]],
+      ["gpt-5-nano", :openai, 128_000, 400_000, :effort,
+       [price(input: 0.05, output: 0.4, cache_read: 0.005)]],
+      ["gemini-3.5-flash", :gemini, 65_536, 1_048_576, :level,
+       [price(input: 1.5, output: 9.0, cache_read: 0.15)]],
+      ["gemini-3.1-pro-preview", :gemini, 65_536, 1_048_576, :level,
+       [price({ input: 2.0, output: 12.0, cache_read: 0.2 },
+              above: 200_000, higher: { input: 4.0, output: 18.0, cache_read: 0.4 })]],
+      ["gemini-2.5-pro", :gemini, 65_536, 1_048_576, :budget,
+       [price({ input: 1.25, output: 10.0, cache_read: 0.125 },
+              above: 200_000, higher: { input: 2.5, output: 15.0, cache_read: 0.25 })]],
+      ["gemini-2.5-flash", :gemini, 65_536, 1_048_576, :budget,
+       [price(input: 0.3, output: 2.5, cache_read: 0.03)]]
+    ].to_h do |row|
+      id, provider, max_output, context_window, thinking, pricing = row
+      [id, Model.new(id:, provider:, max_output:, context_window:, thinking:, pricing:)]
     end.freeze
 
-    # Dated aliases resolve to their base entry: claude-opus-4-8-20260115 and
-    # gpt-5.4-2025-04-14 both match their base ids.
+    ALIASES = { "gpt-5.6" => "gpt-5.6-sol" }.freeze
+
+    # Published family aliases and dated ids resolve to their canonical entry.
     def self.find(id)
-      CATALOG[id] || CATALOG[id.to_s.sub(/-\d{8}\z/, "").sub(/-\d{4}-\d{2}-\d{2}\z/, "")]
+      exact = CATALOG[id]
+      return exact if exact
+
+      base = id.to_s.sub(/-\d{8}\z/, "").sub(/-\d{4}-\d{2}-\d{2}\z/, "")
+      CATALOG[ALIASES.fetch(base, base)]
     end
 
     def self.max_output(id) = find(id)&.max_output
 
+    # Output capacity that occupies the published context limit. Gemini's
+    # output limit is independent, so its compaction headroom is input-only.
+    def self.shared_output(id)
+      model = find(id)
+      model&.max_output unless model&.provider == :gemini
+    end
+
     def self.thinking(id) = find(id)&.thinking
+
+    def self.rates(id, usage: nil, at: Time.now) = find(id)&.rates(usage:, at:)
+
+    def self.priced?(id) = find(id)&.priced? || false
   end
 end