squishling 0.3.0 → 0.4.0

data/lib/squishling/observer.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+module Squishling
+  # The around-style hook for tools that must see an LLM request while it is made (a tracing span, a Coolhand
+  # declared workload), which squawk's after-the-fact report can't do. An observer is any object responding to
+  # around(call), which yields once to make the request. See docs/observers.md.
+  #
+  # Observers get call metadata (Call) and, back from the yield, the request's token Usage: never the model's
+  # output or the method's inputs.
+  module Observer
+    # What an observer is told about one provider request. handed_off is true when squish! sent the call to the
+    # LLM (typically from a rescue) and false when squish_when or a missing implementation did. role is nil for an
+    # escalation attempt, :sample (with sample "a", "b", ...) or :judge, like squawk's metadata. attempt and
+    # attempts count the (expanded) escalation path the request belongs to. request is :chat, or :judgment for a
+    # System One judge.
+    Call = Data.define(:label, :purpose, :harness, :handed_off, :role, :sample, :attempt, :attempts, :model, :provider,
+      :request)
+
+    # What an observer's block returns: the request's token usage and the model that answered, never the content.
+    # A count the provider didn't report is nil.
+    Usage = Data.define(:input_tokens, :output_tokens, :model)
+
+    BUILT_IN = %i[coolhand].freeze
+
+    module_function
+
+    # The observer itself, nil (auto-detect), or false (none). A built-in's name (:coolhand) or anything that
+    # responds to around is accepted, checked here so a mistake fails at setup, not inside the first request.
+    def validate(observer, where)
+      return observer if observer.nil? || observer == false
+
+      if observer.is_a?(Symbol)
+        return observer if BUILT_IN.include?(observer)
+
+        raise ConfigurationError, "#{where} observer must be one of #{BUILT_IN.map(&:inspect).join(', ')}, " \
+                                  "an object that responds to around, or false, got #{observer.inspect}"
+      end
+      return observer if observer.respond_to?(:around)
+
+      raise ConfigurationError, "#{where} observer must respond to around(call) { ... } (or be false), " \
+                                "got #{observer.inspect}"
+    end
+
+    # The observer that wraps a request: the configured one, the Coolhand observer when nothing is configured and
+    # Coolhand is loaded and set up, or nil.
+    def resolve(configured)
+      case configured
+      when false then nil
+      when nil then Observers::Coolhand.new if Observers::Coolhand.available?
+      when Symbol then Observers::Coolhand.new
+      else configured
+      end
+    end
+
+    # Runs the request inside the resolved observer, if there is one. The block runs on the calling thread, so
+    # for samples it is the worker thread that makes the request. The observer gets the request's Usage back from
+    # its yield, not the response, so the model's output can't leave the process through it, and what it returns
+    # is ignored, and so is any rescue of the request's error: the error is raised again. An observer that never
+    # yields would make the request vanish, and one that yields twice would send it twice: both are setup mistakes.
+    def around(call)
+      observer = resolve(Squishling.config.observer)
+      return yield unless observer
+
+      yielded = false
+      response = nil
+      failure = nil
+      observer.around(call) do
+        if yielded
+          failure = ConfigurationError.new("#{call.label}: #{observer.class}#around must yield only once")
+          raise failure
+        end
+
+        yielded = true
+        begin
+          response = yield
+        rescue StandardError => e
+          failure = e
+          raise
+        end
+        usage(response)
+      end
+      unless yielded
+        raise ConfigurationError,
+          "#{call.label}: the #{observer.class} observer's around must yield to make the request"
+      end
+
+      raise failure if failure # an observer can't swallow the request's error: Squishling's taxonomy holds
+
+      response
+    end
+
+    def usage(response)
+      tokens = response.respond_to?(:tokens) ? response.tokens : nil
+      model = response.respond_to?(:model) ? response.model : nil
+      Usage.new(input_tokens: tokens&.input, output_tokens: tokens&.output,
+        model: model.respond_to?(:id) ? model.id : model)
+    end
+    private_class_method :usage
+  end
+end

data/lib/squishling/observers/coolhand.rb

@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+module Squishling
+  module Observers
+    # Declares each request's workload and template to Coolhand (Coolhand.with_context), so Coolhand files the
+    # logged call under them instead of matching it to a template. The workload is the squished method
+    # ("Class#method", described as that, plus "(rescue)" when squish! handed the call over); the template is the
+    # call's place in it: "Class#method/squish:L2" (an escalation attempt), "/sample:a", or "/judge". The model
+    # isn't in the key: it's on the log, and the place is the stable identity.
+    # Used automatically when the coolhand gem is loaded and configured (see Observer.resolve). See docs/coolhand.md.
+    class Coolhand
+      REPORTER = "squishling"
+
+      # Coolhand rejects keys longer than this.
+      MAX_KEY_LENGTH = 255
+
+      # Whether Coolhand is loaded, new enough to declare workloads, and set up to send logs.
+      def self.available?
+        return false unless defined?(::Coolhand) && ::Coolhand.respond_to?(:with_context)
+
+        configuration = ::Coolhand.respond_to?(:configuration) && ::Coolhand.configuration
+        return false unless configuration
+
+        configuration.enabled != false && !configuration.api_key.to_s.strip.empty?
+      rescue StandardError
+        false # nobody asked for this observer, so a Coolhand that can't say whether it's set up isn't ours to break
+      end
+
+      def self.load_library
+        require "coolhand"
+      end
+      private_class_method :load_library
+
+      # Raises ConfigurationError when the coolhand gem is missing or too old to declare workloads.
+      def initialize
+        self.class.send(:load_library) unless defined?(::Coolhand)
+        return if ::Coolhand.respond_to?(:with_context)
+
+        raise ConfigurationError, "observer :coolhand needs a coolhand gem version with Coolhand.with_context " \
+                                  "(declared workloads); upgrade the coolhand gem"
+      rescue LoadError
+        raise ConfigurationError, "observer :coolhand needs the coolhand gem; add `gem \"coolhand\"` to your Gemfile"
+      end
+
+      def around(call, &)
+        workload = call.label
+        template = "#{workload}/#{role_key(call)}"
+        return yield unless declarable?(workload) && declarable?(template)
+
+        ::Coolhand.with_context(
+          reporter: REPORTER,
+          workload: { key: workload, name: workload, description: description(call) },
+          template: { key: template, name: "#{workload} · #{role_name(call)}" },
+          &
+        )
+      end
+
+      private
+
+      # The squished method, marked "(rescue)" when squish! handed it to the LLM. Not the purpose text, which can be
+      # a proc that differs per call.
+      def description(call)
+        call.handed_off ? "#{call.label} (rescue)" : call.label
+      end
+
+      # Coolhand rejects, and drops the whole log for, a key over the limit, with surrounding whitespace, or with
+      # a leading, trailing or doubled "/".
+      def declarable?(key)
+        key.length <= MAX_KEY_LENGTH && key == key.strip && !key.start_with?("/") && !key.end_with?("/") &&
+          !key.include?("//")
+      end
+
+      def role_key(call)
+        case call.role
+        when :sample then "sample:#{call.sample}"
+        when :judge then "judge"
+        else "squish:L#{call.attempt}"
+        end
+      end
+
+      def role_name(call)
+        role_key(call).tr(":", " ")
+      end
+    end
+  end
+end

data/lib/squishling/params.rb

@@ -37,16 +37,21 @@ module Squishling
       raise ConfigurationError, "#{label} must be a Hash, got #{params.class}" unless params.is_a?(Hash)
 
       params = params.transform_keys(&:to_sym)
-      reserved = params.keys & RESERVED_KEYS
-      if reserved.any?
-        raise ConfigurationError, "#{label}: #{reserved.join(', ')} can't be set through params " \
-                                  "(controlled by Squishling/RubyLLM; use model:/provider:/output_schema)"
-      end
+      reject_reserved!(params, RESERVED_KEYS, label, "; use model:/provider:/output_schema")
       reject_nested_reserved!(params, label)
       params[:thinking] = normalize_thinking(params[:thinking], label) unless params[:thinking].nil?
       params.freeze
     end
 
+    # Raises when params sets any of the keys Squishling owns. hint is appended inside the parenthetical.
+    def reject_reserved!(params, keys, label, hint = "")
+      reserved = params.keys & keys
+      return if reserved.empty?
+
+      raise ConfigurationError, "#{label}: #{reserved.join(', ')} can't be set through params " \
+                                "(controlled by Squishling/RubyLLM#{hint})"
+    end
+
     # Nested keys are compared as symbols so a string-keyed "responseMimeType" can't slip past the check.
     # A non-Hash container would replace (not merge into) the one RubyLLM builds, wiping the strict format.
     # Validated containers are copied and frozen so they can't be mutated after the check.

data/lib/squishling/rotation.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Squishling
+  # How many times each squished method has been called with models to rotate between (see ModelPath: steps tied on
+  # priority). The count is the call's turn: each tier of tied models starts that many steps further along, so
+  # successive calls give every model an even share of the first attempts. Counted per class and method, across
+  # threads.
+  module Rotation
+    # class => { method name => calls so far }. Weak keys, so classes replaced by code reloading can be collected.
+    COUNTS = ObjectSpace::WeakKeyMap.new
+    COUNTS_LOCK = Mutex.new
+    private_constant :COUNTS, :COUNTS_LOCK
+
+    class << self
+      # This call's turn, counting from 0, and counts it.
+      def advance(klass, name)
+        COUNTS_LOCK.synchronize do
+          counts = (COUNTS[klass] ||= {})
+          turn = counts.fetch(name, 0)
+          counts[name] = turn + 1
+          turn
+        end
+      end
+    end
+  end
+end

data/lib/squishling/run/messages.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Squishling
+  class Run
+    # The follow-up messages a Run sends a model after an output was rejected.
+    module Messages
+      private
+
+      def retry_message(errors)
+        "Your previous response was rejected:\n- #{errors.join("\n- ")}\nRespond again with corrected JSON only."
+      end
+
+      # The first message to a fresh chat after an earlier model's output was rejected. The rejected output is
+      # model-generated and may cross providers, so it is capped at Invoker::MAX_FORWARDED_CHARS.
+      def escalation_message(input, raw, errors)
+        previous = rejected_text(raw)
+        previous = "(an empty response)" if previous.strip.empty? || raw.nil?
+        if previous.length > Invoker::MAX_FORWARDED_CHARS
+          omitted = previous.length - Invoker::MAX_FORWARDED_CHARS
+          previous = "#{previous[0, Invoker::MAX_FORWARDED_CHARS]}... [truncated, #{omitted} more characters]"
+        end
+        "#{input}\n\nA previous attempt at this request returned:\n#{previous}\n" \
+          "It was rejected:\n- #{errors.join("\n- ")}\nRespond with corrected JSON only."
+      end
+
+      # The rejected output as text. Output that was rejected for being unrepresentable in JSON (Infinity, invalid
+      # UTF-8) can't be serialized or measured as is, so it is scrubbed or replaced rather than raising.
+      def rejected_text(raw)
+        return utf8(raw).scrub("?") if raw.is_a?(String)
+
+        JSON.generate(raw)
+      rescue JSON::JSONError
+        "(a response JSON can't represent)"
+      end
+
+      # A binary-tagged string counts every byte as valid, so it is read as UTF-8 before scrubbing.
+      def utf8(text)
+        text.encoding == Encoding::BINARY ? text.dup.force_encoding(Encoding::UTF_8) : text
+      end
+    end
+  end
+end

data/lib/squishling/run.rb

@@ -0,0 +1,296 @@
+# frozen_string_literal: true
+
+module Squishling
+  # One squished call, as a harness sees it (see Harness::Base): the primitives that ask the model for outputs,
+  # and everything they do for free: model, provider, and params resolution, schema and squish_validate checks,
+  # retries in the same chat, escalation, squawk reports, and Squishling's own errors when attempts run out.
+  # Each returns typed results, the same type as the deterministic path.
+  #
+  # A stage is the attempts of one escalation step (a model with its attempts:), as an Array of
+  # ModelPath::Step. Primitives take stages, or anything that declares one (a model name or a step Hash).
+  #
+  # Primitives must be called from the thread that runs the harness. Only the provider requests inside them run
+  # on other threads; parsing, validation, and every callback (squish_validate, compare:, procs) stay on this one.
+  class Run
+    include Messages
+
+    private_constant :Messages
+
+    # One of #samples' outputs, between rounds.
+    Sample = Struct.new(:name, :stage, :chat, :rejected, :result)
+
+    SAMPLE_NAMES = ("a".."z").to_a.freeze
+    private_constant :Sample, :SAMPLE_NAMES
+
+    attr_reader :receiver, :inputs, :steps
+
+    def initialize(definition:, receiver:, inputs:, prompt:, payload:, purpose:, harness:)
+      @definition = definition
+      @receiver = receiver
+      @inputs = inputs
+      @prompt = prompt
+      @payload = payload
+      @purpose = purpose
+      @harness = harness
+      @steps = definition.escalation_path.freeze
+      @client = LLMClient.new(definition.label, purpose:, harness: harness.type, handed_off: definition.handed_off)
+      @output_check = OutputCheck.new(definition, receiver, inputs)
+      @owner = Thread.current
+      @produced = {}.compare_by_identity
+    end
+
+    # Whether value is a result one of this run's primitives returned (not one built by hand).
+    def produced?(value)
+      @produced.key?(value)
+    end
+
+    # The squished method, as "Class#method".
+    def label
+      @definition.label
+    end
+
+    # The escalation's steps grouped into stages: runs of identical consecutive steps, the attempts of one model.
+    def stages
+      @stages ||= steps.slice_when { |before, after| before != after }.map(&:freeze).freeze
+    end
+
+    # A stage from a model name or a step Hash (model:, provider:, params:, attempts:, forward_rejected:), the
+    # way a judge step is resolved: it keeps its own provider, and its params merge over the method's.
+    # A stage (an Array of steps) is returned as is, and a single step is a stage of one attempt.
+    def stage(declaration)
+      return checked_stage([declaration]) if declaration.is_a?(ModelPath::Step)
+
+      if declaration.is_a?(Array)
+        return checked_stage(declaration) if declaration.any? && declaration.all?(ModelPath::Step)
+
+        raise ConfigurationError, "#{label}: a stage is the steps of one model, got #{declaration.inspect}"
+      end
+
+      step = Harness::OptionTypes.check(:step, declaration, label, :"harness step")
+      ModelPath.steps([step], provider: nil, params: @definition.params).freeze
+    end
+
+    # Makes each attempt of the stages in order (the escalation's steps when none are given) until an output
+    # passes the schema and squish_validate. Consecutive attempts on the same step continue the same chat, so the
+    # model sees what it got wrong; a new step starts a fresh chat told about the last rejected output (unless
+    # the step sets forward_rejected: false). Raises InvalidOutputError, or LLMError, when the last attempt fails.
+    def escalate(*stages)
+      owner!
+      path = stages.empty? ? steps : stages.flat_map { |declared| stage(declared) }
+      escalate_with(path, @prompt)
+    end
+
+    # One sample per stage (the attempts of one model), each asked on its own chat, each retried per its stage's
+    # attempts like #escalate, with the requests of every round sent concurrently. Returns the outputs in order.
+    # Samples are named "a", "b", ... in squawk reports.
+    def samples(stages)
+      owner!
+      raise ConfigurationError, "#{label}: samples takes an Array of stages" unless stages.is_a?(Array)
+
+      stages = stages.map { |declared| stage(declared) }
+      raise ConfigurationError, "#{label}: samples needs at least one stage" if stages.empty?
+      if stages.size > SAMPLE_NAMES.size
+        raise ConfigurationError, "#{label}: samples takes at most #{SAMPLE_NAMES.size} stages"
+      end
+
+      if (mixed = stages.find { |sample_stage| sample_stage.uniq.size > 1 })
+        raise ConfigurationError, "#{label}: a sample is asked on one model, so its stage can't mix steps " \
+                                  "(got #{mixed.map(&:display_name).join(', ')})"
+      end
+
+      samples = stages.zip(SAMPLE_NAMES).map { |sample_stage, name| Sample.new(name, sample_stage) }
+      (1..stages.map(&:size).max).each do |attempt|
+        pending = samples.reject(&:result)
+        break if pending.empty?
+
+        round(pending, attempt)
+      end
+      samples.map(&:result)
+    end
+
+    # A judge that picks between two outputs (see Judge#verdict, which returns a Judge::Verdict), or rejects
+    # both. judge: is a model name or a Hash with model:, provider:, params:, and attempts: (as an escalation
+    # step, without priority:), plus type: :chat (default) or :judgment and, for :judgment, min_confidence:.
+    # Without one, the judge is the escalation's next stage after the first. instructions: is a String, or a
+    # Proc evaluated against the receiver, replacing Harness::DEFAULT_JUDGE_INSTRUCTIONS. The judge is resolved
+    # here, so a missing one raises ConfigurationError before any request is made.
+    def judge(judge: nil, instructions: nil)
+      owner!
+      judge = judge_step(judge) unless judge.nil?
+      Judge.new(definition: @definition, receiver: @receiver, client: @client, judge:, instructions:, path: steps,
+        purpose: @purpose, payload: @payload, escalate: method(:escalate_with), log_failure: method(:log_failure),
+        guard: method(:owner!), harness_type: @harness.type)
+    end
+
+    # Whether two outputs agree: compare (a Proc, evaluated against the receiver with the two results and the
+    # method's inputs, truthy for agreement) or, without one, whether they are exactly equal.
+    # compare is the developer's own code, so its exceptions propagate unwrapped.
+    def agree?(first, second, compare = nil)
+      owner!
+      return Schema.jsonify(first.to_h) == Schema.jsonify(second.to_h) unless compare
+
+      @receiver.instance_exec(first, second, **@inputs, &compare) ? true : false
+    end
+
+    # Logs a warning, "[Squishling] Class#method message", to Squishling.config.logger, if there is one.
+    def log(message)
+      Squishling.config.logger&.warn("[Squishling] #{label} #{message}")
+    end
+
+    private
+
+    # Steps made or changed by hand get the params checks a declared step does (no reserved keys).
+    def checked_stage(stage)
+      stage.each { |step| Params.normalize(step.params, "#{label} harness step params") }
+      stage.frozen? ? stage : stage.dup.freeze
+    end
+
+    # A judge declaration as the escalation step Judge takes: one step (no priority:), with its type.
+    def judge_step(declaration)
+      step = ModelPath.normalize_step(declaration, "#{label} judge", extras: Judge::STEP_KEYS)
+      ModelPath.single_step!(step, "#{label} judge: a judge is one step, so it")
+
+      type = step.fetch(:type, :chat)
+      type = :chat if type.nil?
+      step.merge(type: type.is_a?(String) ? type.to_sym : type)
+    end
+
+    def owner!
+      return if Thread.current.equal?(@owner)
+
+      raise ConfigurationError, "#{label}: a harness must call its run's primitives from the thread it runs on"
+    end
+
+    # Makes each attempt of the path until an output passes (see #escalate). The judge's prompt (role :judge)
+    # skips squish_validate and names the judge in errors and logs.
+    def escalate_with(path, prompt)
+      chat = nil
+      rejected = nil # [raw content, errors] of the last invalid output
+      models = []
+
+      path.each.with_index(1) do |step, attempt|
+        models << step.model
+        last = attempt == path.size
+        if chat && step == path[attempt - 2]
+          message = retry_message(rejected.last)
+        else
+          chat = start_chat(step, prompt)
+          message = rejected && step.forward_rejected ? escalation_message(prompt.input, *rejected) : prompt.input
+        end
+
+        begin
+          response = @client.ask(chat, message, step, attempt:, attempts: path.size, role: prompt.role)
+        rescue LLMError => e
+          squawk(prompt, path, attempt, nil, e)
+          raise if last
+
+          # The failed request may be half-recorded in this chat, so the next attempt starts a fresh one.
+          chat = nil
+          log_failure(prompt.role, attempt, path, e.message)
+          next
+        end
+
+        result, errors = @output_check.call(response.content, prompt)
+        if errors.empty?
+          squawk(prompt, path, attempt, response, nil)
+          return produce(result, prompt)
+        end
+
+        error = InvalidOutputError.new(errors:, raw: response.content, attempts: attempt, models: models.dup,
+          source: prompt.role == :judge ? "Judge" : "LLM")
+        squawk(prompt, path, attempt, response, error)
+        raise error if last
+
+        rejected = [response.content, errors]
+        log_failure(prompt.role, attempt, path, errors.join("; "))
+      end
+    end
+
+    # One round of #samples: the pending samples' requests, sent concurrently, then each outcome recorded.
+    def round(pending, attempt)
+      jobs = pending.map do |sample|
+        step = sample.stage[attempt - 1]
+        message = sample_message(sample, step)
+        chat = sample.chat
+        attempts = sample.stage.size
+        -> { @client.ask(chat, message, step, attempt:, attempts:, role: :sample, sample: sample.name) }
+      end
+      outcomes = LLMClient.concurrently(jobs)
+      # A setup mistake (or anything that isn't a failed request) ends the call, whichever sample hit it.
+      fatal = outcomes.map(&:last).find { |error| error && !error.is_a?(LLMError) }
+      raise fatal if fatal
+
+      pending.zip(outcomes).each { |sample, outcome| record(sample, outcome, attempt) }
+    end
+
+    # The next message for a sample: a retry in its chat after invalid output, otherwise a fresh chat.
+    def sample_message(sample, step)
+      return retry_message(sample.rejected.last) if sample.chat
+
+      sample.chat = start_chat(step, @prompt)
+      sample.rejected && step.forward_rejected ? escalation_message(@prompt.input, *sample.rejected) : @prompt.input
+    end
+
+    def record(sample, (response, error), attempt)
+      stage = sample.stage
+      last = attempt == stage.size
+      role = "sample #{sample.name}"
+      if error
+        squawk(@prompt, stage, attempt, nil, error, sample:)
+        raise error if last
+
+        sample.chat = nil
+        return log_failure(role, attempt, stage, error.message)
+      end
+
+      result, errors = @output_check.call(response.content, @prompt)
+      if errors.empty?
+        squawk(@prompt, stage, attempt, response, nil, sample:)
+        return sample.result = produce(result, @prompt)
+      end
+
+      invalid = InvalidOutputError.new(errors:, raw: response.content, attempts: attempt,
+        models: stage.first(attempt).map(&:model))
+      squawk(@prompt, stage, attempt, response, invalid, sample:)
+      raise invalid if last
+
+      sample.rejected = [response.content, errors]
+      log_failure(role, attempt, stage, errors.join("; "))
+    end
+
+    # Remembers a result of the method's own output (a judge's verdict isn't one), so a harness can only return
+    # what a primitive produced.
+    def produce(result, prompt)
+      @produced[result] = true unless prompt.role == :judge
+      result
+    end
+
+    def start_chat(step, prompt)
+      @client.chat(step, instructions: prompt.instructions, schema: prompt.schema.llm_schema)
+    end
+
+    # Runs the observability hook, if any, with this attempt's raw output (nil when the call itself failed),
+    # the error that ended it (nil when it was accepted), and what was asked. role is nil for escalation
+    # attempts, :sample (with the sample's name) or :judge.
+    def squawk(prompt, path, attempt, response, error, sample: nil)
+      hook = @definition.squawk
+      return unless hook
+
+      step = path[attempt - 1]
+      metadata = {
+        label:, attempt:, attempts: path.size, final: attempt == path.size,
+        model: response&.model || step.model, provider: step.provider, params: step.params, input: prompt.input,
+        usage: response&.tokens&.to_h, role: sample ? :sample : prompt.role, sample: sample&.name
+      }
+      Squawk.call(hook, output: response&.content, metadata:, error:)
+    end
+
+    def log_failure(role, attempt, path, reason)
+      step = path[attempt - 1]
+      next_step = path[attempt]
+      action = next_step == step ? "retrying" : "escalating to #{next_step.display_name}"
+      who = role ? "#{role} attempt" : "attempt"
+      log("#{who} #{attempt} of #{path.size} (#{step.display_name}) failed, #{action}: #{reason}")
+    end
+  end
+end

data/lib/squishling/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Squishling
-  VERSION = "0.3.0"
+  VERSION = "0.4.0"
 end

data/lib/squishling.rb

@@ -10,7 +10,16 @@ require_relative "squishling/errors"
 require_relative "squishling/configuration"
 require_relative "squishling/params"
 require_relative "squishling/squawk"
+require_relative "squishling/observer"
+require_relative "squishling/observers/coolhand"
 require_relative "squishling/model_path"
+require_relative "squishling/rotation"
+require_relative "squishling/harness/option_types"
+require_relative "squishling/harness/registry"
+require_relative "squishling/harness/base"
+require_relative "squishling/harness/escalation"
+require_relative "squishling/harness/squishsum"
+require_relative "squishling/harness/judged_squishsum"
 require_relative "squishling/harness"
 require_relative "squishling/result"
 require_relative "squishling/schema"
@@ -19,6 +28,8 @@ require_relative "squishling/appendices"
 require_relative "squishling/definition"
 require_relative "squishling/llm_client"
 require_relative "squishling/output_check"
+require_relative "squishling/run/messages"
+require_relative "squishling/run"
 require_relative "squishling/invoker"
 require_relative "squishling/judge"
 require_relative "squishling/router"
@@ -42,6 +53,27 @@ module Squishling
       @config = Configuration.new
     end
 
+    # Forgets the registered harnesses (the built-ins stay), e.g. between specs. Kept apart from reset_config!
+    # because harnesses are registered once at boot, and a suite that resets its config between examples
+    # shouldn't lose them.
+    def reset_harnesses!
+      Harness::Registry.reset!
+    end
+
+    # Registers a custom harness under a name, usable anywhere a built-in type is (harness: :name, or
+    # harness: { type: :name, option: value }, on a class, a method, a squish! call, or config.default_harness).
+    # Pass a Harness::Base subclass, or a block that becomes the body of one (see Harness::Base). Registered
+    # globally, so do it before the classes that use it load (an initializer). Raises ConfigurationError for a
+    # name that is invalid or already taken, including the built-in :escalation, :squishsum, :judged_squishsum.
+    def register_harness(name, klass = nil, &)
+      Harness::Registry.register(name, klass, &)
+    end
+
+    # The names of the harnesses that can be declared: the built-ins, then those registered.
+    def harnesses
+      Harness.types
+    end
+
     def included(base)
       raise ConfigurationError, "Squishling can only be included in a class" unless base.is_a?(Class)