squishling 0.4.0 → 0.5.1

data/lib/squishling/adapters/openai_options.rb

@@ -0,0 +1,114 @@
+# frozen_string_literal: true
+
+require "ipaddr"
+require "uri"
+
+module Squishling
+  module Adapters
+    # Checks the options of Adapters::OpenAICompatible, raising ConfigurationError (without echoing a value that may
+    # be a secret) for a bad one.
+    module OpenAIOptions
+      # The URL itself is never in the message: it may carry credentials.
+      INVALID_BASE_URL = "adapter :openai_compatible base_url must be an http(s) URL without credentials (use " \
+                         "api_key:), a query, or a fragment"
+      # Headers that carry a secret as a matter of course (OpenAI-style, Azure OpenAI and API Management, Google, and
+      # proxies), by name: an API or subscription key, a token, auth, a secret, a credential, or a cookie.
+      CREDENTIAL_HEADER = /(api|subscription)[-_]?key|token|auth|secret|credential|cookie/i
+      CLEARTEXT_CREDENTIAL = "adapter :openai_compatible won't send credentials (the api_key or a header such as " \
+                             "Authorization or api-key) over plain http to a host that isn't this machine; use " \
+                             "https, or pass allow_insecure_http: true to accept the risk on a network you trust"
+      INVALID_HEADER = "adapter :openai_compatible headers (and the api_key) can't contain line breaks or NUL"
+
+      module_function
+
+      def endpoint(base_url)
+        uri = URI.parse(base_url.to_s.strip.sub(%r{/+\z}, ""))
+        unless uri.is_a?(URI::HTTP) && uri.host && uri.userinfo.nil? && uri.query.nil? && uri.fragment.nil?
+          raise ConfigurationError, INVALID_BASE_URL
+        end
+
+        URI.parse("#{uri}/chat/completions")
+      rescue URI::InvalidURIError
+        raise ConfigurationError, INVALID_BASE_URL
+      end
+
+      def headers(headers)
+        unless headers.is_a?(Hash) && headers.all? { |name, value| name.is_a?(String) && value.is_a?(String) }
+          raise ConfigurationError, "adapter :openai_compatible headers must be a Hash of String names to String values"
+        end
+        raise ConfigurationError, INVALID_HEADER if headers.flatten.any? { |text| unsafe_header?(text) }
+
+        # A copy, so the caller can't change a header after the check above.
+        headers.to_h { |name, value| [name.dup.freeze, value.dup.freeze] }.freeze
+      end
+
+      # Net::HTTP raises, with the value in the message, for a header that carries a line break.
+      def unsafe_header?(text)
+        text.match?(/[\r\n\0]/)
+      end
+
+      # Whether a credential would travel in the clear: over http, to a host that isn't loopback (a name is judged by
+      # its spelling, never resolved, so a lookup can't change the answer).
+      def cleartext?(endpoint)
+        endpoint.scheme == "http" && !loopback?(endpoint.hostname)
+      end
+
+      def loopback?(host)
+        host = host.downcase.delete_suffix(".")
+        host == "localhost" || host.end_with?(".localhost") || IPAddr.new(host).loopback?
+      rescue IPAddr::Error
+        false
+      end
+
+      # Whether a request to endpoint, with this key (a String, possibly blank) and these headers, would send a
+      # credential in the clear.
+      def cleartext_credential?(endpoint, key, headers)
+        return false unless cleartext?(endpoint)
+
+        !key.to_s.strip.empty? || headers.each_key.any? { |name| name.match?(CREDENTIAL_HEADER) }
+      end
+
+      def flag(name, value)
+        return value if [true, false].include?(value)
+
+        raise ConfigurationError, "adapter :openai_compatible #{name} must be true or false, got #{value.inspect}"
+      end
+
+      # Defaults to the time the three timeouts allow together: connecting, writing the request, one read.
+      def total_timeout(value, open_timeout, read_timeout)
+        return open_timeout + (2 * read_timeout) if value.nil?
+
+        number(:total_timeout, value)
+      end
+
+      def default_model(model)
+        return nil if model.nil?
+        return model.to_s if (model.is_a?(String) || model.is_a?(Symbol)) && !model.to_s.strip.empty?
+
+        raise ConfigurationError, "adapter :openai_compatible default_model must be a model name, got #{model.inspect}"
+      end
+
+      def number(name, value, type: Numeric)
+        return value if value.is_a?(type) && value.positive?
+
+        kind = type == Integer ? "Integer" : "number"
+        raise ConfigurationError, "adapter :openai_compatible #{name} must be a positive #{kind}, got #{value.inspect}"
+      end
+
+      def retries(value)
+        return value if value.is_a?(Integer) && !value.negative?
+
+        raise ConfigurationError,
+          "adapter :openai_compatible max_retries must be an Integer of 0 or more, got #{value.inspect}"
+      end
+
+      def choice(name, value, choices)
+        symbol = value.is_a?(String) ? value.to_sym : value
+        return symbol if choices.include?(symbol)
+
+        raise ConfigurationError,
+          "adapter :openai_compatible #{name} must be one of #{choices.join(', ')}, got #{value.inspect}"
+      end
+    end
+  end
+end

data/lib/squishling/adapters/openai_transport.rb

@@ -0,0 +1,123 @@
+# frozen_string_literal: true
+
+require "net/http"
+require "timeout"
+
+module Squishling
+  module Adapters
+    # The HTTP exchange of Adapters::OpenAICompatible, kept apart so the limits sit in one place: one POST whose
+    # response is read in chunks, with a cap on its size (the declared length, then the bytes actually received, for an
+    # error body too) and on the total time it may take, which read_timeout alone doesn't bound (it restarts at every
+    # read, so a slow drip of headers or body can hold a worker thread for as long as the server likes). Compression
+    # is refused, not inflated: the request asks for an identity encoding, and a response that comes back encoded
+    # anyway is dropped, so a small body can't expand past the cap. Redirects aren't followed: the next hop could be
+    # another host, and the Authorization header would go with it. A limit that is hit raises a Refused, whose
+    # message names the limit and never carries the URL, a header, or any of the response.
+    class OpenAITransport
+      class Refused < StandardError; end
+      class ResponseTooLarge < Refused; end
+      class DeadlineExceeded < Refused; end
+      class UnexpectedEncoding < Refused; end
+
+      # What is left of a response once it has been read: its status (a String, as Net::HTTPResponse has it), whole
+      # body, and Retry-After header.
+      Reply = Struct.new(:code, :body, :retry_after) do
+        def success?
+          code.start_with?("2")
+        end
+      end
+
+      MONOTONIC = -> { Process.clock_gettime(Process::CLOCK_MONOTONIC) }
+      private_constant :MONOTONIC
+
+      def initialize(endpoint:, open_timeout:, read_timeout:, total_timeout:, max_response_bytes:, clock: MONOTONIC)
+        @endpoint = endpoint
+        @open_timeout = open_timeout
+        @read_timeout = read_timeout
+        @total_timeout = total_timeout
+        @max_response_bytes = max_response_bytes
+        @clock = clock
+      end
+
+      # total_timeout and the byte cap apply per call, so each retry gets its own.
+      def post(body, headers)
+        deadline = @clock.call + @total_timeout
+        perform(body, headers, deadline)
+      rescue DeadlineExceeded
+        # Net::HTTP re-raises what interrupts a connect with the host and port in the message; ours is built afresh.
+        expired!
+      rescue Timeout::Error
+        expired! if @clock.call >= deadline
+
+        raise
+      end
+
+      private
+
+      # Net::HTTP reads the status line and headers before it yields the response, with read_timeout restarting at every
+      # byte, so the deadline can't be checked there: a Timeout cuts the whole exchange off at total_timeout instead.
+      # The connection is dropped whichever way the block ends, so nothing is left half-used. The deadline is checked
+      # again at each chunk, which also stops a body that arrives in time but whose clock says otherwise.
+      def perform(body, headers, deadline)
+        request = Net::HTTP::Post.new(@endpoint.request_uri, headers)
+        request.body = body
+        http = connection
+        Timeout.timeout(@total_timeout, DeadlineExceeded, expired_message) do
+          http.start { http.request(request) { |response| return read(response, deadline) } }
+        end
+      end
+
+      def connection
+        Net::HTTP.new(@endpoint.hostname, @endpoint.port).tap do |http|
+          http.use_ssl = @endpoint.scheme == "https"
+          http.open_timeout = @open_timeout
+          http.read_timeout = http.write_timeout = @read_timeout
+        end
+      end
+
+      def read(response, deadline)
+        Reply.new(response.code, body_of(response, deadline), response["retry-after"])
+      end
+
+      # The whole body, within the limits. On a failure status a body that is compressed or over the cap is dropped
+      # rather than refused: the status alone says what kind of failure it is (a proxy's gzipped 401 is still a
+      # credentials problem, not a transport one), and nothing in the body is read for it anyway.
+      def body_of(response, deadline)
+        check_headers(response)
+        body = String.new(encoding: Encoding::BINARY)
+        response.read_body do |chunk|
+          expired! unless @clock.call < deadline
+          too_large! if body.bytesize + chunk.bytesize > @max_response_bytes
+          body << chunk
+        end
+        body
+      rescue ResponseTooLarge, UnexpectedEncoding
+        raise if response.code.start_with?("2")
+
+        String.new(encoding: Encoding::BINARY)
+      end
+
+      def check_headers(response)
+        encoding = response["content-encoding"].to_s.strip.downcase
+        unless encoding.empty? || encoding == "identity"
+          raise UnexpectedEncoding, "the response was compressed although the request asked for none"
+        end
+
+        length = response["content-length"].to_s.strip
+        too_large! if length.match?(/\A\d+\z/) && length.to_i > @max_response_bytes
+      end
+
+      def too_large!
+        raise ResponseTooLarge, "the response was larger than max_response_bytes (#{@max_response_bytes})"
+      end
+
+      def expired!
+        raise DeadlineExceeded, expired_message, cause: nil
+      end
+
+      def expired_message
+        "the request ran past total_timeout (#{@total_timeout}s)"
+      end
+    end
+  end
+end

data/lib/squishling/adapters/ruby_llm.rb

@@ -0,0 +1,229 @@
+# frozen_string_literal: true
+
+module Squishling
+  module Adapters
+    # Sends requests through RubyLLM (https://rubyllm.com): OpenAI, Anthropic Claude, Google Gemini, AWS Bedrock,
+    # OpenRouter, and any other provider it supports. Configure provider API keys in RubyLLM as usual
+    # (RubyLLM.configure); transient HTTP failures are retried by RubyLLM itself (RubyLLM.config.max_retries).
+    # The ruby_llm gem is loaded when the adapter is built, so Squishling doesn't depend on it. See docs/adapters.md.
+    class RubyLLM < Adapter::Base
+      MINIMUM_VERSION = Gem::Version.new("2.1")
+
+      # The ruby_llm gem isn't installed (Adapter.detect tells that apart from a gem that is, but too old).
+      class MissingGemError < ConfigurationError; end
+
+      # Raises ConfigurationError when the ruby_llm gem is missing. keep_cause: true keeps RubyLLM's original exception
+      # as the #cause of the errors raised for a failed request (its message and response hold the provider's response
+      # body); by default #cause is a sanitized stand-in (Adapter::ProviderFailure) with only the class and HTTP status.
+      def initialize(keep_cause: false)
+        super()
+        unless [true, false].include?(keep_cause)
+          raise ConfigurationError, "adapter :ruby_llm keep_cause: must be true or false, got #{keep_cause.inspect}"
+        end
+
+        @keep_cause = keep_cause
+        require "ruby_llm"
+        check_version!
+        @lib = gem_lib_dir
+      rescue LoadError => e
+        raise unless e.path == "ruby_llm"
+
+        raise MissingGemError, "adapter :ruby_llm needs the ruby_llm gem; add `gem \"ruby_llm\"` to your Gemfile"
+      end
+
+      # One chat on the request's model: the system prompt, the strict output schema, and the params, then the
+      # conversation so far replayed, answering its last message. Models missing from RubyLLM's registry (e.g. newly
+      # released ones) are only usable when a provider is named, so RubyLLM is told to assume they exist.
+      def complete(request)
+        chat = build_chat(request)
+        *history, last = request.messages
+        history.each { |message| chat.add_message(role: message[:role], content: message[:content]) }
+        response = guarded(request) { chat.ask(last[:content]) }
+        Adapter::Response.new(content: response.content, model: response.model, tokens: tokens_of(response),
+          finish_reason: response.finish_reason)
+      end
+
+      def judgments?
+        true
+      end
+
+      # A System One judgment (RubyLLM.judge) on one step. The step's params are sent as provider options.
+      # RubyLLM raises ArgumentError for a judgment it can't build (e.g. a model on a provider that takes none, or
+      # a provider option its protocol reserves), and a bare RubyLLM::Error without an HTTP response for a provider
+      # that doesn't support judgments at all. An ArgumentError from anywhere else (an instrumentation subscriber
+      # runs inside the call) is the developer's own and propagates as-is.
+      def judge(input, questions:, step:, label:)
+        request = judgment_request(label, step)
+        guarded(request) do
+          ::RubyLLM.judge(input, questions:, provider_options: step.params, **model_options(step.model, step.provider))
+        rescue ::RubyLLM::ModelNotFoundError, ::RubyLLM::Error, ArgumentError => e
+          raise unless setup_failure?(e)
+
+          raise ConfigurationError, "#{label}: #{e.message}"
+        end
+      end
+
+      # RubyLLM's instrumentation context (the workflow and its captured context) is carried to the worker threads
+      # the way RubyLLM's own concurrent tool calls do it. Unlike those, the threads don't enter the Rails executor:
+      # the caller already holds it, and a second share can deadlock against a pending code reload.
+      def capture_context
+        instrumentation = ::RubyLLM::Support::Instrumentation if defined?(::RubyLLM::Support::Instrumentation)
+        return unless instrumentation.respond_to?(:current_workflow) && instrumentation.respond_to?(:capture_context)
+
+        [instrumentation.current_workflow, instrumentation.capture_context]
+      end
+
+      def with_context(context, &)
+        return yield unless context
+
+        instrumentation = ::RubyLLM::Support::Instrumentation
+        instrumentation.with_workflow(context.first) { instrumentation.with_context(context.last, &) }
+      end
+
+      private
+
+      def tokens_of(response)
+        response.tokens ? response.tokens.to_h : {}
+      end
+
+      # Whether a failure from RubyLLM.judge is a judgment RubyLLM can't make, rather than a failed request or the
+      # developer's own error.
+      def setup_failure?(error)
+        case error
+        when ::RubyLLM::ModelNotFoundError then true
+        when ArgumentError then raised_by_ruby_llm?(error)
+        else error.instance_of?(::RubyLLM::Error) && error.response.nil?
+        end
+      end
+
+      def judgment_request(label, step)
+        Adapter::Request.new(label:, model: step.model, provider: step.provider, instructions: nil, messages: [],
+          schema: nil, params: step.params)
+      end
+
+      def build_chat(request)
+        chat = ::RubyLLM.chat(**model_options(request.model, request.provider))
+        chat.with_instructions(request.instructions)
+        chat.with_schema(request.schema)
+        apply_params(chat, request)
+        chat
+      rescue ::RubyLLM::ModelNotFoundError, ::RubyLLM::ConfigurationError => e
+        raise configuration_error(request, e.message)
+      end
+
+      # RubyLLM validates some settings locally (e.g. an impossible thinking budget for the model)
+      # and raises ArgumentError before any request is sent.
+      def apply_params(chat, request)
+        params = request.params
+        rest = params.except(:temperature, :thinking, :max_output_tokens)
+        chat.with_temperature(params[:temperature]) if params.key?(:temperature)
+        chat.with_max_output_tokens(params[:max_output_tokens]) if params.key?(:max_output_tokens)
+        apply_thinking(chat, params[:thinking]) if params.key?(:thinking)
+        chat.with_provider_options(rest) if rest.any?
+      rescue ArgumentError => e
+        raise configuration_error(request, "invalid params #{params.keys.inspect} (#{e.message})")
+      end
+
+      # true: the model's default thinking; false: off; a Hash: { effort:, budget:, display: }.
+      def apply_thinking(chat, thinking)
+        thinking.is_a?(Hash) ? chat.with_thinking(**thinking) : chat.with_thinking(thinking)
+      end
+
+      # Transient HTTP failures are already retried by RubyLLM; anything that still fails is surfaced as an
+      # LLMError, which moves on to the next attempt of the escalation (or propagates from the last one), or as a
+      # ConfigurationError for a setup mistake (see #http_failure for which statuses are which): escalating or a
+      # fallback would otherwise hide it on every call.
+      # A provider's or transport's own message never goes into ours: a provider can echo request content or part of
+      # an API key in it, and ours reach logs and error trackers.
+      # Error trackers report the whole cause chain, and RubyLLM's error message and response hold the provider's
+      # response body, so #cause is a sanitized stand-in (class, HTTP status, backtrace); the original is there only
+      # with keep_cause: true. RubyLLM's own ConfigurationError, and the errors #build_chat, #apply_params, and #judge
+      # convert, are local (they name a missing setting or a bad param and hold no provider response), so they are
+      # raised as they are and Ruby sets the original as their #cause.
+      def guarded(request)
+        yield
+      rescue ::RubyLLM::ConfigurationError => e
+        raise configuration_error(request, "#{e.class}: #{e.message}")
+      rescue ::RubyLLM::Error, Faraday::Error => e
+        raise failure_for(request, e), cause: cause_of(e)
+      end
+
+      # Sorted by the HTTP status, not by RubyLLM's error class: RubyLLM's middleware picks a class from the provider's
+      # message text on a 400, so a bad parameter whose message happens to say "per day" arrives as a RateLimitError.
+      # An error with no response is sorted by its class. The one class honored over the status is RubyLLM's
+      # ContextLengthExceededError, since RubyLLM is where its provider's wording for an overlong prompt is known.
+      def failure_for(request, error)
+        via = describe_failure(error)
+        return llm_error(request, via) if error.is_a?(::RubyLLM::ContextLengthExceededError)
+
+        status = response_field(error.response, :status) || class_status(error)
+        http_failure(request, status:, body: response_field(error.response, :body), via:)
+      end
+
+      # The status RubyLLM's classes stand for, for an error raised without a response.
+      def class_status(error)
+        case error
+        when ::RubyLLM::UnauthorizedError then 401
+        when ::RubyLLM::PaymentRequiredError then 402
+        when ::RubyLLM::ForbiddenError then 403
+        when ::RubyLLM::BadRequestError then 400
+        end
+      end
+
+      # The error's class and HTTP status, when it has one, which is all of it that is ours to repeat.
+      def describe_failure(error)
+        status = status_of(error)
+        status ? "#{error.class}, HTTP #{status}" : error.class.to_s
+      end
+
+      def status_of(error)
+        response_field(error.response, :status)
+      end
+
+      def cause_of(error)
+        @keep_cause ? error : sanitized_cause(error, status: status_of(error))
+      end
+
+      # A field of an error's response, which is a Faraday::Response, or a Hash for a Faraday error.
+      def response_field(response, field)
+        response.is_a?(Hash) ? response[field] : (response.public_send(field) if response.respond_to?(field))
+      end
+
+      def model_options(model, provider)
+        options = { model:, provider: }.compact
+        options[:assume_model_exists] = true if model && provider && !known_model?(model, provider)
+        options
+      end
+
+      def known_model?(model, provider)
+        ::RubyLLM.models.find(model, provider:)
+        true
+      rescue ::RubyLLM::ModelNotFoundError
+        false
+      end
+
+      def raised_by_ruby_llm?(error)
+        error.backtrace&.first&.start_with?("#{@lib}/") || false
+      end
+
+      # The adapter uses RubyLLM 2.1 APIs (judgments, with_max_output_tokens, Tokens#to_h), and no gemspec pins it.
+      def check_version!
+        version = Gem.loaded_specs["ruby_llm"]&.version
+        return if version.nil? || version >= MINIMUM_VERSION
+
+        raise ConfigurationError, "adapter :ruby_llm needs ruby_llm #{MINIMUM_VERSION} or newer, got #{version}"
+      end
+
+      # Where RubyLLM's own code lives, to tell its ArgumentErrors from the developer's (see #judge).
+      # It is read from the file that was loaded, which is the path RubyLLM's backtrace frames carry.
+      def gem_lib_dir
+        entry = $LOADED_FEATURES.find { |feature| feature.end_with?("/lib/ruby_llm.rb") }
+        return File.dirname(entry) if entry
+
+        File.join(Gem.loaded_specs.fetch("ruby_llm").full_gem_path, "lib")
+      end
+    end
+
+    Adapter::Registry.register_built_in(:ruby_llm, RubyLLM)
+  end
+end

data/lib/squishling/budget.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module Squishling
+  # A cap on what one squished call may spend across every request it makes: max_total_output_tokens: and
+  # max_total_input_tokens:, each a positive Integer. A budget is a Hash of the limits one level declares
+  # ({ output:, input: }), layered config -> class -> method -> squish! like generation params. `false` lifts a limit an
+  # outer level set. See docs/failures.md.
+  module Budget
+    # The limit, as declared, for each kind of token, and the counter in Tally#used it is checked against.
+    KEYS = { output: :max_total_output_tokens, input: :max_total_input_tokens }.freeze
+
+    module_function
+
+    # One level's limits, validated. A limit left nil is not set at this level.
+    def build(label, output: nil, input: nil)
+      { output:, input: }.compact.to_h { |kind, value| [kind, check(value, KEYS.fetch(kind), label)] }.freeze
+    end
+
+    # value as is when it is a positive Integer, or false (no limit). Raises ConfigurationError otherwise.
+    def check(value, name, label)
+      return value if value == false || (value.is_a?(Integer) && value.positive?)
+
+      raise ConfigurationError,
+        "#{label}: #{name} must be a positive Integer (or false for no limit), got #{value.inspect}"
+    end
+
+    # The limits that apply: later levels override earlier ones limit by limit, and false removes a limit.
+    def resolve(*layers)
+      layers.compact.reduce({}) { |merged, layer| merged.merge(layer) }.select { |_kind, max| max }
+    end
+
+    # What one call has spent so far, counted from every response it receives. Requests of one round run on worker
+    # threads, so the totals are taken under a lock. A response without token counts adds nothing, so a provider that
+    # reports no usage never exhausts a budget.
+    class Tally
+      def initialize(budget)
+        @budget = budget
+        @used = { input_tokens: 0, output_tokens: 0 }
+        @lock = Mutex.new
+      end
+
+      # Counts a response's tokens: an Adapter::Response, or a judgment (see Observer.usage).
+      def add(response)
+        usage = Observer.usage(response)
+        @lock.synchronize do
+          @used[:input_tokens] += count(usage.input_tokens)
+          @used[:output_tokens] += count(usage.output_tokens)
+        end
+        response
+      end
+
+      # The first limit that has been reached ({ limit:, max:, used: }, used being both totals), or nil.
+      def exhausted
+        @lock.synchronize do
+          @budget.each do |kind, max|
+            next if @used.fetch(:"#{kind}_tokens") < max
+
+            return { limit: KEYS.fetch(kind), max:, used: @used.dup }
+          end
+          nil
+        end
+      end
+
+      private
+
+      def count(tokens)
+        tokens.is_a?(Integer) && tokens.positive? ? tokens : 0
+      end
+    end
+  end
+end

data/lib/squishling/class_methods.rb

@@ -9,7 +9,7 @@ module Squishling
     #   squishling model: "claude-sonnet-5-5", purpose: "...", output_schema: MySchema
     # model: is a single attempt; escalation: is a list of models tried in order (see ModelPath):
     #   squishling escalation: [{ model: "claude-haiku-4-5", attempts: 2 }, "claude-sonnet-5-5", "claude-opus-5-5"]
-    # Pass provider: alongside model: for models missing from RubyLLM's registry, e.g.
+    # Pass provider: alongside model: for models missing from RubyLLM's registry (the RubyLLM adapter), e.g.
     #   squishling model: "gpt-7-preview", provider: :openai
     # provider: needs a model: or escalation: (declared on the class), and a subclass's model never inherits its
     # parent's provider.
@@ -19,11 +19,19 @@ module Squishling
     #   squishling append_to_purpose: ["The Ruby that handles well-formed input:", self]
     # harness: chooses how the escalation is used (see Harness):
     #   squishling harness: :judged_squishsum
-    # squawk: is called after every LLM attempt with the raw output (see Configuration#squawk); false silences
+    # max_total_output_tokens:/max_total_input_tokens: cap the tokens one call may use across all its requests
+    # (see Budget); false lifts a cap inherited from the parent class or the configuration:
+    #   squishling max_total_output_tokens: 25_000
+    # squawk: is called after each LLM attempt (narrowed by squawk_on:), and once more when a token budget or a
+    # disagreement stops the call, with the raw output (see Configuration#squawk); false silences
     # an inherited one:
     #   squishling squawk: ->(output:, metadata:, error:) { Tracer.record(output, metadata, error) }
+    # squawk_on: limits which attempts reach the hook: :all (the default), :rejections, :retries, :final, or a
+    # list of them (any one holding runs the hook):
+    #   squishling squawk: hook, squawk_on: %i[rejections final]
     def squishling(model: nil, escalation: nil, provider: nil, params: nil, harness: nil, purpose: nil,
-      append_to_purpose: nil, output_schema: nil, squawk: nil)
+      append_to_purpose: nil, output_schema: nil, squawk: nil, squawk_on: nil, max_total_output_tokens: nil,
+      max_total_input_tokens: nil)
       path = ModelPath.declare(model, escalation, to_s)
       # A provider belongs to the model declared at the same level. Without a model here it would apply to
       # nothing (the model comes from the config or a parent class), so it is an error unless this class already
@@ -41,14 +49,24 @@ module Squishling
       self.purpose(purpose) if purpose
       self.append_to_purpose(append_to_purpose) unless append_to_purpose.nil?
       self.output_schema(output_schema) if output_schema
-      @squishling_squawk = Squawk.validate(squawk, to_s) unless squawk.nil?
+      declare_squawk(squawk, squawk_on)
+      declare_budget(max_total_output_tokens, max_total_input_tokens)
       self
     end
 
+    # The token limits declared on this class and its ancestors, limit by limit (see Budget).
+    def squishling_budget
+      squishling_inherited(:squishling_budget, {}).merge(@squishling_budget || {})
+    end
+
     def squishling_squawk
       squishling_lookup(:@squishling_squawk)
     end
 
+    def squishling_squawk_on
+      squishling_lookup(:@squishling_squawk_on)
+    end
+
     # The class's (or nearest ancestor's) model or escalation, as normalized steps.
     def squishling_model_path
       squishling_lookup(:@squishling_model_path)
@@ -147,11 +165,14 @@ module Squishling
     #   squish :triage, purpose: "...", escalation: %w[claude-haiku-4-5 claude-sonnet-5-5], when: ->(**) { true },
     #                   fallback: ->(error, **) { { priority: "medium" } }, append_to_purpose: [...],
     #                   validate: ->(result, **) { "team is required" if result.team.empty? },
-    #                   harness: :squishsum, squawk: ->(output:, error:, **) { Tracer.record(output, error) } do
+    #                   harness: :squishsum, max_total_output_tokens: 25_000,
+    #                   squawk: ->(output:, error:, **) { Tracer.record(output, error) },
+    #                   squawk_on: :retries do
     #     string :priority
     #   end
     def squish(*names, purpose: nil, append_to_purpose: nil, output_schema: nil, model: nil, escalation: nil,
-      provider: nil, params: nil, harness: nil, when: nil, fallback: nil, validate: nil, squawk: nil, &schema_block)
+      provider: nil, params: nil, harness: nil, when: nil, fallback: nil, validate: nil, squawk: nil,
+      squawk_on: nil, max_total_output_tokens: nil, max_total_input_tokens: nil, &schema_block)
       schema = schema_block ? Schematist::Schema.create(&schema_block) : output_schema
       model = ModelPath.declare(model, escalation, "#{self} squish")
       raise ConfigurationError, "#{self} squish: provider: needs a model: or escalation:" if provider && !model
@@ -162,8 +183,11 @@ module Squishling
         append_to_purpose = Appendices.normalize(append_to_purpose, "#{self} squish append_to_purpose")
       end
       squawk = Squawk.validate(squawk, "#{self} squish")
+      squawk_on = Squawk.validate_on(squawk_on, "#{self} squish")
+      budget = Budget.build("#{self} squish", output: max_total_output_tokens, input: max_total_input_tokens)
       options = { purpose:, append_to_purpose:, output_schema: schema, model:, provider:, params:, harness:,
-                  predicate: binding.local_variable_get(:when), fallback:, validator: validate, squawk: }.compact
+                  predicate: binding.local_variable_get(:when), fallback:, validator: validate, squawk:,
+                  squawk_on:, budget: }.compact
 
       names.map(&:to_sym).each do |name|
         (@squishling_methods ||= {})[name] = options
@@ -192,8 +216,20 @@ module Squishling
 
     private
 
+    # nil leaves a level unset (so the parent's applies); false for squawk: is stored, to silence it.
+    def declare_squawk(squawk, squawk_on)
+      @squishling_squawk = Squawk.validate(squawk, to_s) unless squawk.nil?
+      @squishling_squawk_on = Squawk.validate_on(squawk_on, "#{self} squishling") unless squawk_on.nil?
+    end
+
+    # Adds the class's own token limits to those it declared before.
+    def declare_budget(output, input)
+      @squishling_budget = (@squishling_budget || {}).merge(Budget.build(to_s, output:, input:)).freeze
+    end
+
     # Each class gets its own prepended wrapper so overrides in subclasses are routed too.
     def squishling_install_wrapper
+      ClassRegistry.register(self)
       @squishling_wrapper = Wrapper.new
       prepend(@squishling_wrapper)