smith-agents 0.11.0 → 0.12.1

data/lib/smith/agent/provider_completion.rb

@@ -1,25 +1,43 @@
 # frozen_string_literal: true
 
 require_relative "provider_failure_handling"
+require_relative "content_bytes"
+require_relative "content_parts"
 require_relative "invocation_preparation"
 require_relative "provider_attempt"
 require_relative "provider_call_timing"
+require_relative "provider_request"
 require_relative "provider_candidate_sequence"
+require_relative "provider_streaming"
+require_relative "observed_model"
 
 module Smith
   class Agent
     module ProviderCompletion
       include ProviderFailureHandling
       include InvocationPreparation
+      include ContentParts
       include ProviderCallTiming
+      include ProviderRequest
+      include ProviderStreaming
+      include ObservedModel
 
       private
 
-      def complete_with_provider(agent_class, prepared_input, output_schema:)
-        candidates = ProviderCandidateSequence.new(build_model_chain(agent_class))
+      # The input's documents and images are read once (ContentBytes), and
+      # its cached parts checked against every Anthropic model in the chain,
+      # before the first attempt.
+      def complete_with_provider(agent_class, prepared_input, shape:)
+        chain = build_model_chain(agent_class).tap { check_chain_cache_breakpoints!(_1, prepared_input) }
+        ContentBytes.with(prepared_input) { attempt_candidates(agent_class, prepared_input, chain, shape) }
+      end
+
+      def attempt_candidates(agent_class, prepared_input, chain, shape)
+        candidates = ProviderCandidateSequence.new(chain)
         candidates.each do |model_reference, index|
           check_deadline! if index.positive?
-          attempt = attempt_model(agent_class, prepared_input, model_reference, output_schema:, attempt_index: index)
+          withdraw_streamed(index.positive? ? :fallback : :retry)
+          attempt = attempt_model(agent_class, prepared_input, model_reference, shape:, attempt_index: index)
           return attempt if attempt.success?
 
           candidates.suppress(account_failed_provider(attempt, model_reference))
@@ -33,7 +51,7 @@ module Smith
         # Fail closed: an empty candidate chain (agent without a model)
         # or any exhausted sequence must surface a typed diagnostic, not
         # the sequence itself destructured into nils.
-        raise Smith::AgentError,
+        raise Smith::WorkflowError,
               "no executable model candidate for #{agent_class}; declare a model or fallback_models"
       end
 
@@ -88,52 +106,37 @@ module Smith
       # cohesive lifecycle (identity, prepared chat, observed model, timed
       # completion, prefix accounting on failure); splitting it would scatter
       # the rescue-path accounting away from what it accounts for.
-      def attempt_model(agent_class, prepared_input, model_reference, output_schema:, attempt_index:)
+      def attempt_model(agent_class, prepared_input, model_reference, shape:, attempt_index:)
         facts = { attempt_id: SecureRandom.uuid, attempt_index:, agent_name: agent_class.register_as }
-        chat = prepared_attempt_chat(agent_class, prepared_input, model_reference, output_schema:)
+        chat = prepared_attempt_chat(agent_class, prepared_input, model_reference, shape)
         message_count = chat_message_count(chat)
         observed_reference = observed_model_reference(chat, fallback: model_reference)
         timer = ProviderCallTiming::Timer.start
-        response = chat.complete
+        response, stream, mode, call_facts = complete_chat(chat, observed_reference, collected = [], facts)
         timer.stop
-        completion = Completion.from_messages(response: response, messages: new_chat_messages(chat, message_count))
+        completion = Completion.from_messages(response:, messages: new_chat_messages(chat, message_count), stream:,
+                                              mode:)
 
-        completed_provider_attempt(completion, observed_reference, timer, facts)
+        completed_provider_attempt(completion, observed_reference, timer, facts.merge(call_facts.to_h))
       rescue StandardError => e
         timer&.stop
         observed_reference ||= observed_model_reference(chat, fallback: model_reference)
         prefix_usages = account_completed_prefix(
-          agent_class, observed_reference, new_chat_messages(chat, message_count), attempt_id: facts[:attempt_id]
+          agent_class, observed_reference, new_chat_messages(chat, message_count),
+          attempt_id: facts[:attempt_id], stream_rounds: streamed_rounds(collected || [])
         )
         attempt = failed_provider_attempt(e, observed_reference, timer, facts, prefix_usages)
-        raise unless provider_failure?(e)
+        raise unless provider_failure?(e) || unsupported_content?(e)
 
         attempt
       end
       # rubocop:enable Metrics/AbcSize
 
-      def observed_model_reference(chat, fallback:)
-        model = observable_model(chat)
-        return fallback unless model.respond_to?(:id) && !model.id.to_s.empty?
-
-        ModelReference.coerce(model.id, provider: observed_provider(model, fallback))
-      rescue StandardError
-        fallback
-      end
-
-      def observable_model(chat)
-        resolved_chat = chat.respond_to?(:to_llm) ? chat.to_llm : chat
-        resolved_chat.model if resolved_chat.respond_to?(:model)
-      end
-
-      def observed_provider(model, fallback)
-        model.respond_to?(:provider) && !model.provider.to_s.empty? ? model.provider : fallback.provider
-      end
-
-      def prepared_attempt_chat(agent_class, prepared_input, model_reference, output_schema:)
+      def prepared_attempt_chat(agent_class, prepared_input, model_reference, shape)
         chat = agent_class.chat(**model_reference.chat_options, **bridge_workflow_inputs(agent_class))
-        add_prepared_input(chat, prepared_input)
-        output_schema ? chat.with_schema(output_schema) : chat
+        add_prepared_input(chat, prepared_input, observed_model_reference(chat, fallback: model_reference).provider)
+        continue_paused_loop(chat)
+        shaped_chat(chat, model_reference, shape)
       end
 
       def chat_message_count(chat)

data/lib/smith/agent/provider_failure_handling.rb

@@ -49,7 +49,8 @@ module Smith
       end
 
       def fallback_eligible?(error)
-        transient_failure?(error) ||
+        unsupported_content?(error) ||
+          transient_failure?(error) ||
           transport_failure?(error) ||
           provider_account_failure?(error) ||
           model_permission_failure?(error) ||
@@ -80,6 +81,9 @@ module Smith
 
       def terminal_provider_error(error, model_reference)
         return Smith::AgentError.new(error.message) if transient_failure?(error) || transport_failure?(error)
+        if unsupported_content?(error)
+          return Smith::ContentUnavailable.new("#{model_reference} takes none of this content: #{error.message}")
+        end
 
         Smith::ProviderPermanentFailure.new(
           error.message,

data/lib/smith/agent/provider_request.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Smith
+  class Agent
+    # What one provider call asks for beyond its messages: the output schema
+    # as the provider takes it, and a larger output limit when a cut-off
+    # response is retried.
+    module ProviderRequest
+      Shape = Data.define(:output_schema, :output_limit)
+
+      private
+
+      def shaped_chat(chat, model_reference, shape)
+        provider = observed_model_reference(chat, fallback: model_reference).provider
+        chat = with_output_limit(chat, provider, shape.output_limit) if shape.output_limit
+        shape.output_schema ? chat.with_schema(provider_output_schema(shape.output_schema, provider)) : chat
+      end
+
+      def with_output_limit(chat, provider, limit)
+        return chat unless chat.respond_to?(:with_params)
+
+        params = (chat.params if chat.respond_to?(:params)) || {}
+        chat.with_params(**::RubyLLM::Utils.deep_merge(params, Providers::OutputLimit.params(provider, params, limit)))
+      end
+
+      # For Anthropic without the constraints it rejects (SchemaConstraints);
+      # unchanged otherwise.
+      def provider_output_schema(output_schema, provider)
+        return output_schema unless provider.to_s == "anthropic"
+
+        payload = Providers::SchemaConstraints.json_schema(output_schema)
+        payload ? Providers::SchemaConstraints.for_provider(payload, provider) : output_schema
+      end
+    end
+  end
+end

data/lib/smith/agent/provider_streaming.rb

@@ -0,0 +1,155 @@
+# frozen_string_literal: true
+
+module Smith
+  class Agent
+    # The provider calls of a streamed step (Smith::Streaming): before a
+    # call that replaces an earlier one, the earlier call's text is
+    # withdrawn; each call is announced and its text chunks sent. The stop
+    # reason is read from the stream (Providers::StreamFacts), as a
+    # streamed reply keeps no raw body.
+    module ProviderStreaming
+      # Room for what a pause adds besides its messages (its retry, the
+      # attempt's usage).
+      PAUSE_MARGIN = 4096
+
+      private
+
+      # Withdraws the text a streamed step sent, for `reason`, when a call
+      # is about to replace it.
+      def withdraw_streamed(reason) = @step_stream&.withdraw(reason)
+
+      # The chat's reply and, for a streamed step, what Completion needs to
+      # read its facts. The stream's collection goes into `collected` too,
+      # for the facts of a call that fails part way. Requests over the wire
+      # are held on the step's budget ledger, if it has one, keyed by the
+      # attempt's `attempt_id` (`facts`).
+      def complete_chat(chat, model_reference, collected, facts = {})
+        return capture_batch_call(chat, model_reference) if @batch_capturing
+
+        entry = take_batch_result
+        delivered = deliver_batch_result(chat, entry) if entry
+        return delivered if delivered
+
+        remembered_completion(chat, model_reference, facts) do |resumed|
+          live_completion(chat, model_reference, collected, facts, resumed)
+        end
+      end
+
+      # The call over the wire; `resumed` (RememberedCalls) says where a
+      # call a response memo kept part of goes on from.
+      def live_completion(chat, model_reference, collected, facts, resumed)
+        ledger = run_budget_ledger
+        return wire_completion(chat, model_reference, collected) unless ledger
+
+        start = resumed ? resumed.start : chat_message_count(chat)
+        holding = { pausable: budget_pausable?, rounds: resumed&.rounds || 0 }
+        Providers::RequestHolds.within(chat_provider(chat), ledger, facts[:attempt_id], **holding) do |holder|
+          held_completion(holder, chat, start) { wire_completion(chat, model_reference, collected) }
+        end
+      end
+
+      # A refusal of a tool loop's later round pauses the loop (its step
+      # retries for budget, BudgetPause): the error carries the messages to
+      # continue from (from `start`, a memo's kept rounds included). A loop
+      # whose messages cannot be kept goes on, unheld, from its tools'
+      # results (the chat keeps them), as when its step cannot pause.
+      def held_completion(holder, chat, start)
+        yield
+      rescue BudgetUnavailable => e
+        raise unless holder.pausable && holder.rounds > 1 && e.continuation.nil?
+
+        paused = paused_loop(e, chat, start)
+        raise paused if paused
+
+        holder.pausable = false
+        retry
+      end
+
+      def paused_loop(error, chat, start)
+        continuation = Budget::Continuation.dump(@budget_continuation, chat.messages.drop(start || 0))
+        BudgetUnavailable.new(error.message, continuation:) if continuation && storable_pause?(continuation)
+      end
+
+      # Whether the run's state, with the paused loop kept, stays within
+      # persistence_max_bytes: a pause it could not store would lose the
+      # loop's paid rounds, so the loop goes on instead.
+      def storable_pause?(continuation)
+        limit = Smith.config.persistence_max_bytes
+        return true unless limit
+
+        JSON.generate(to_state).bytesize + JSON.generate(continuation).bytesize + PAUSE_MARGIN <= limit
+      rescue StandardError
+        false
+      end
+
+      def wire_completion(chat, model_reference, collected)
+        stream = @step_stream
+        return [chat.complete, nil] unless stream
+
+        stream.start(attempt: @step_attempt, model: model_reference.model_id, provider: model_reference.provider&.to_s)
+        response, rounds = streamed_completion(chat, stream, collected)
+        [response, { rounds:, output_limit: request_output_limit(chat, model_reference) }]
+      end
+
+      # The chat's reply, its text sent as it comes (without <think> text),
+      # and the facts of each of its requests.
+      def streamed_completion(chat, stream, collected)
+        response = nil
+        thinking = Streaming::ThinkTags.new
+        rounds = Providers::StreamFacts.collect(chat_provider(chat)) do |collection|
+          collected << collection
+          response = chat.complete { |chunk| stream.delta(thinking.strip(chunk.content, collection.rounds.length)) }
+        end
+        [response, rounds]
+      end
+
+      # The batch result as the chat's reply, and whether the request the
+      # call rendered is other than the one sent (a fact of its attempt). A
+      # result its call will not use (`on_request_change` :resend or :fail)
+      # is set aside, to be charged; :resend returns nil, and the call goes
+      # over the wire.
+      def deliver_batch_result(chat, entry)
+        delivery = @batch_delivery
+        response, changed = Providers::BatchExchange.deliver(chat_provider(chat), entry.result, entry.digest,
+                                                             on_change: delivery.on_change) { chat.complete }
+        [response, nil, :batch, { request_changed: changed || nil }]
+      rescue BatchRequestChanged
+        delivery.set_aside(Attribution.current&.branch_key, entry)
+        raise unless delivery.on_change == :resend
+      end
+
+      # A batch step's call: the provider must have a batch interface (only
+      # Anthropic's, so far) and the agent no tools (a batch request is one
+      # turn). The request is thrown out of the call before it is sent.
+      def capture_batch_call(chat, model_reference)
+        provider = chat_provider(chat)
+        unless provider.respond_to?(:slug) && provider.slug.to_s == "anthropic"
+          raise BatchModeUnavailable, "#{model_reference.provider || "this provider"} has no batch interface in Smith"
+        end
+        raise BatchModeUnavailable, "a batch request is one turn: its agent cannot have tools" if chat_tools?(chat)
+
+        @batch_capture = Providers::BatchExchange::Exchange.new(provider:, holding: !run_budget_ledger.nil?)
+        Providers::BatchExchange.within(@batch_capture) { chat.complete }
+        raise BatchModeUnavailable, "the call did not reach its provider's request"
+      end
+
+      def chat_tools?(chat)
+        tools = chat.tools if chat.respond_to?(:tools)
+        tools.respond_to?(:empty?) && !tools.empty?
+      end
+
+      # The facts of a streamed call that failed, if it streamed.
+      def streamed_rounds(collected) = collected.first&.facts
+
+      def chat_provider(chat)
+        llm_chat = chat.respond_to?(:to_llm) ? chat.to_llm : chat
+        llm_chat.instance_variable_get(:@provider)
+      end
+
+      def request_output_limit(chat, model_reference)
+        params = chat.params if chat.respond_to?(:params)
+        Providers::OutputLimit.current(model_reference.provider, params.is_a?(Hash) ? params : {}, model_reference)
+      end
+    end
+  end
+end

data/lib/smith/agent/provider_usage.rb

@@ -3,12 +3,33 @@
 require "dry-struct"
 
 require_relative "../types"
+require_relative "../providers/response_facts"
 
 module Smith
   class Agent
+    # The usage one provider response reported. `input_tokens` counts only
+    # uncached input, which is how RubyLLM reports every provider: tokens read
+    # from or written to a prompt cache are `cache_read_tokens` and
+    # `cache_write_tokens`. `thinking_tokens` are already inside
+    # `output_tokens` and are reported, never added. The optional counts are
+    # nil when the provider did not report them.
     class ProviderUsage < Dry::Struct
+      COUNT_READERS = {
+        cache_read_tokens: %i[cache_read_tokens cached_tokens],
+        cache_write_tokens: %i[cache_write_tokens cache_creation_tokens],
+        thinking_tokens: %i[thinking_tokens]
+      }.freeze
+      OPTIONAL_COUNTS = COUNT_READERS.keys.freeze
+
       attribute :input_tokens, Types::Integer
       attribute :output_tokens, Types::Integer
+      attribute :cache_read_tokens, Types::Integer.optional.default(nil)
+      attribute :cache_write_tokens, Types::Integer.optional.default(nil)
+      attribute :thinking_tokens, Types::Integer.optional.default(nil)
+      attribute :stop_reason, Types::String.optional.default(nil)
+      attribute :provider_message_id, Types::String.optional.default(nil)
+      # How the response was made: :stream, :batch, or nil (synchronous).
+      attribute :mode, Types::Symbol.optional.default(nil)
 
       def self.from_message(message)
         return unless message.respond_to?(:input_tokens) && message.respond_to?(:output_tokens)
@@ -17,17 +38,63 @@ module Smith
         output_tokens = message.output_tokens
         return unless input_tokens.is_a?(Integer) && output_tokens.is_a?(Integer)
 
-        new(input_tokens:, output_tokens:)
+        new(input_tokens:, output_tokens:, **optional_counts(message), **Providers::ResponseFacts.from_message(message))
+      end
+
+      def self.optional_counts(message)
+        COUNT_READERS.transform_values do |readers|
+          readers.lazy.filter_map { |reader| count_from(message, reader) }.first
+        end
       end
 
+      def self.count_from(message, reader)
+        return unless message.respond_to?(reader)
+
+        value = message.public_send(reader)
+        value if value.is_a?(Integer) && value >= 0
+      end
+
+      # The stop reason and message id are those of the last response, the
+      # one that ended the attempt.
       def self.sum(usages)
         return if usages.empty?
 
-        new(input_tokens: usages.sum(&:input_tokens), output_tokens: usages.sum(&:output_tokens))
+        counts = OPTIONAL_COUNTS.to_h do |name|
+          values = usages.filter_map(&name)
+          [name, values.empty? ? nil : values.sum]
+        end
+        new(input_tokens: usages.sum(&:input_tokens), output_tokens: usages.sum(&:output_tokens), **counts,
+            stop_reason: usages.last.stop_reason, provider_message_id: usages.last.provider_message_id,
+            mode: usages.last.mode)
+      end
+      private_class_method :optional_counts, :count_from
+
+      # This usage with the stop reason and message id of `facts`, a
+      # ResponseFacts result, in place of its own: a tool loop's summed
+      # usage carries its final response's facts even when that response
+      # reported no usage.
+      def with_response_facts(facts)
+        self.class.new(**to_h.except(:stop_reason, :provider_message_id), **facts)
+      end
+
+      # This usage as a streamed response's, with the facts its stream
+      # carried (Providers::StreamFacts).
+      def streamed(facts) = with_response_facts(facts).new(mode: :stream)
+
+      # Optional members are omitted when nil, so usage without cache or
+      # response details keeps its earlier shape.
+      def to_h
+        super.compact
+      end
+      alias to_hash to_h
+
+      # Every input token the provider processed, cached or not.
+      def total_input_tokens
+        input_tokens + (cache_read_tokens || 0) + (cache_write_tokens || 0)
       end
 
       def total_tokens
-        input_tokens + output_tokens
+        total_input_tokens + output_tokens
       end
     end
   end

data/lib/smith/agent/usage_entry_recording.rb

@@ -32,7 +32,9 @@ module Smith
           attempt_kind: attempt_kind,
           recorded_at: Time.now.utc.iso8601(6),
           attempt_id: attempt_id,
-          **ambient_attribution_fields
+          **ambient_attribution_fields,
+          **Workflow::AgentResult::USAGE_DETAIL_MEMBERS.to_h { |member| [member, agent_result.public_send(member)] },
+          mode: agent_result.mode
         )
       end
 
@@ -52,7 +54,7 @@ module Smith
 
       def accumulate_usage(agent_result, entry)
         @usage_mutex.synchronize do
-          @total_tokens = (@total_tokens || 0) + agent_result.input_tokens + agent_result.output_tokens
+          @total_tokens = (@total_tokens || 0) + agent_result.total_tokens
           @total_cost = (@total_cost || 0.0) + (agent_result.cost || 0.0)
           @usage_entries << entry
         end

data/lib/smith/agent/usage_traces.rb

@@ -16,6 +16,8 @@ module Smith
           data: {
             input_tokens: agent_result.input_tokens,
             output_tokens: agent_result.output_tokens,
+            cache_read_tokens: agent_result.cache_read_tokens,
+            cache_write_tokens: agent_result.cache_write_tokens,
             model: agent_result.model_used,
             provider: agent_result.provider_used
           }.compact
@@ -30,7 +32,9 @@ module Smith
       # traces is not a spend total. Unpriced usage emits nothing, and the
       # caller gates out partially metered or partially priced invocations
       # so an incomplete figure is never presented as the invocation cost.
-      def emit_cost_trace(agent_result, invocation_cost)
+      # `cache_unpriced: true` says cache tokens were priced at 0 because the
+      # catalog has no rate for them.
+      def emit_cost_trace(agent_result, invocation_cost, cache_unpriced: false)
         return if invocation_cost.nil?
 
         Smith::Trace.record(
@@ -40,7 +44,10 @@ module Smith
             model: agent_result.model_used,
             provider: agent_result.provider_used,
             input_tokens: agent_result.input_tokens,
-            output_tokens: agent_result.output_tokens
+            output_tokens: agent_result.output_tokens,
+            cache_read_tokens: agent_result.cache_read_tokens,
+            cache_write_tokens: agent_result.cache_write_tokens,
+            cache_unpriced: (true if cache_unpriced)
           }.compact
         )
       end

data/lib/smith/agent/usage_tracking.rb

@@ -13,64 +13,70 @@ module Smith
 
       private
 
+      # A failed attempt's own reported usage, cache counts included, is
+      # recorded, priced and settled into the budget like any billed
+      # response, so its usage entry agrees with its :provider_call trace.
       def account_failed_attempt(error, model_reference, agent_class, attempt_id: nil)
-        return unless error.respond_to?(:input_tokens) && error.respond_to?(:output_tokens)
+        usage = ProviderUsage.from_message(error)
+        return unless usage
 
-        input = error.input_tokens
-        output = error.output_tokens
-        return unless input.is_a?(Integer) && output.is_a?(Integer)
-
-        record_failed_usage(agent_class, model_reference, input, output, attempt_id:)
+        record_billed_usages(agent_class, [usage], :failed_attempt, model_reference, attempt_id:)
       end
 
-      # Returns the provider usages it recorded, so the attempt's trace can
-      # carry the same totals its usage entries hold.
-      def account_completed_prefix(agent_class, model_reference, messages, attempt_id: nil)
-        completion = Completion.from_messages(response: nil, messages: messages)
-        record_completion_usage(agent_class, completion, :partial_attempt, model_reference, attempt_id:)
+      # The completed rounds of a tool loop whose attempt then failed were
+      # billed: they are recorded and settled into the budget like a failed
+      # attempt's usage. Returns the provider usages recorded, so the
+      # attempt's trace carries the same totals its usage entries hold.
+      def account_completed_prefix(agent_class, model_reference, messages, attempt_id: nil, stream_rounds: nil)
+        stream = stream_rounds && { rounds: stream_rounds, output_limit: nil }
+        completion = Completion.from_messages(response: nil, messages:, stream:)
+        record_billed_usages(agent_class, completion.provider_usages, :partial_attempt, model_reference, attempt_id:)
         completion.provider_usages
       end
 
-      def record_failed_usage(agent_class, model_reference, input_tokens, output_tokens, attempt_id: nil)
-        model_reference = coerce_model_reference(model_reference)
-        cost = Smith::Pricing.compute_cost(
-          model: model_reference.model_id,
-          provider: model_reference.provider,
-          input_tokens:,
-          output_tokens:
-        )
-        agent_result = Workflow::AgentResult.new(
-          content: nil,
-          input_tokens:,
-          output_tokens:,
-          cost: cost,
-          model_used: model_reference.model_id,
-          provider_used: model_reference.provider
-        )
-        Thread.current[:smith_failed_agent_results] ||= []
-        Thread.current[:smith_failed_agent_results] << agent_result
-        record_usage(agent_class, agent_result, :failed_attempt, model_reference, attempt_id:)
+      # Usage that is billed but is not the invocation's result joins the
+      # failed billable attempts, which budget settlement reads beside it.
+      def record_billed_usages(agent_class, usages, attempt_kind, model_reference, attempt_id:)
+        results = record_response_usages(agent_class, usages, attempt_kind, model_reference, attempt_id:)
+        add_failed_agent_results(results)
+      end
+
+      # A new Array, never the one in place: a step's thread context
+      # (ThreadContextSnapshot) puts back the list it started with, which
+      # must not carry this step's failed attempts into the next.
+      def add_failed_agent_results(results)
+        Thread.current[:smith_failed_agent_results] = [*Array(Thread.current[:smith_failed_agent_results]), *results]
       end
 
-      def snapshot_and_finalize(agent_class, completion, model_reference, attempt_id: nil)
+      # Records and settles one completed attempt's usage and returns its
+      # result, whether or not the result becomes the output.
+      def account_attempt(agent_class, completion, model_reference, attempt_id: nil)
         model_reference = coerce_model_reference(model_reference)
-        agent_result = Workflow::AgentResult.new(
-          content: completion.content,
-          input_tokens: completion.input_tokens,
-          output_tokens: completion.output_tokens,
-          cost: nil,
-          model_used: model_reference.model_id,
-          provider_used: model_reference.provider
-        )
+        agent_result = completion_agent_result(completion, model_reference)
         Thread.current[:smith_last_agent_result] = agent_result
         emit_token_usage(agent_result)
         account_completion!(agent_class, completion, model_reference, agent_result, attempt_id)
+        agent_result
+      end
 
-        agent_result.content = run_after_completion(agent_class, agent_result.content, @context)
+      # The output has been through `after_completion` (OutputIntegrity).
+      def finalize_output(agent_class, agent_result)
         raise_blank_output!(agent_class, agent_result)
         agent_result
       end
 
+      def completion_agent_result(completion, model_reference)
+        Workflow::AgentResult.new(
+          content: completion.content,
+          input_tokens: completion.input_tokens,
+          output_tokens: completion.output_tokens,
+          cost: nil,
+          model_used: model_reference.model_id,
+          provider_used: model_reference.provider,
+          **Workflow::AgentResult.usage_details(completion.usage)
+        )
+      end
+
       def raise_blank_output!(agent_class, agent_result)
         return unless blank_agent_output?(agent_result.content)
 
@@ -93,26 +99,33 @@ module Smith
       # budget settlement and result surfaces read) and the recorded entries,
       # the budget ledger, and the :cost trace all agree. The trace emits
       # only for fully metered, fully priced invocations, so a partial
-      # figure is never presented as the invocation cost; the partial sum
-      # still settles the budget because it is what was verifiably billed.
+      # figure is never presented as the invocation cost; the partial sums
+      # of cost and tokens still settle the budget because they are what
+      # was verifiably billed.
       def account_completion!(agent_class, completion, model_reference, agent_result, attempt_id)
-        invocation_cost, fully_priced = record_completion_usage(
+        invocation_cost, fully_priced, results = record_completion_usage(
           agent_class, completion, :completed_attempt, model_reference, attempt_id:
         )
         agent_result.cost = invocation_cost
-        emit_cost_trace(agent_result, invocation_cost) if fully_priced
+        agent_result.metered_tokens = results.sum(&:total_tokens) unless completion.usage_complete
+        return unless fully_priced
+
+        emit_cost_trace(agent_result, invocation_cost, cache_unpriced: results.any?(&:cache_unpriced))
       end
 
+      # Prices the result and notes whether cache tokens were priced at 0
+      # for lack of a rate, which the :cost trace reports.
       def compute_agent_cost(agent_result)
         return unless agent_result.usage_known?
 
-        model = agent_result.model_used
-        agent_result.cost = Smith::Pricing.compute_cost(
-          model: model,
-          provider: agent_result.provider_used,
-          input_tokens: agent_result.input_tokens,
-          output_tokens: agent_result.output_tokens
-        )
+        usage = {
+          model: agent_result.model_used, provider: agent_result.provider_used,
+          input_tokens: agent_result.input_tokens, output_tokens: agent_result.output_tokens,
+          cache_read_tokens: agent_result.cache_read_tokens || 0,
+          cache_write_tokens: agent_result.cache_write_tokens || 0
+        }
+        agent_result.cost = Smith::Pricing.compute_cost(**usage, mode: agent_result.mode || :sync)
+        agent_result.cache_unpriced = Smith::Pricing.cache_unpriced?(**usage.except(:output_tokens))
       end
     end
   end

data/lib/smith/agent.rb

@@ -34,6 +34,9 @@ module Smith
         subclass.instance_variable_set(:@fallback_models_block, @fallback_models_block)
         subclass.instance_variable_set(:@model_block, @model_block)
         subclass.instance_variable_set(:@tool_budget_exhaustion, @tool_budget_exhaustion)
+        %i[@on_truncation @on_refusal @output_contract @output_repairs].each do |name|
+          subclass.instance_variable_set(name, instance_variable_get(name))
+        end
         subclass.instance_variable_set(:@execution_identity, nil)
         subclass.instance_variable_set(:@registered_name, nil)
       end
@@ -148,8 +151,10 @@ require_relative "agent/dynamic_configuration"
 require_relative "agent/fallback_configuration"
 require_relative "agent/reserved_input_bridge"
 require_relative "agent/chat_construction"
+require_relative "agent/output_policy"
 
 Smith::Agent.extend(Smith::Agent::DynamicConfiguration)
 Smith::Agent.extend(Smith::Agent::FallbackConfiguration)
 Smith::Agent.extend(Smith::Agent::ReservedInputBridge)
 Smith::Agent.extend(Smith::Agent::ChatConstruction)
+Smith::Agent.extend(Smith::Agent::OutputPolicy)