actionagent 1.8.1 → 1.9.1

data/app/models/action_agent/scenario_catalog.rb

@@ -0,0 +1,172 @@
+# frozen_string_literal: true
+
+module ActionAgent
+  # A catalog of evaluation scenarios kept in the dashboard: the products an
+  # owner evaluates, each with named sets of scenarios, as
+  # ActiveAgent::Evals::Catalog reads and writes them.
+  #
+  # The catalog is imported from a YAML document (pasted, uploaded, read from
+  # a connected repository at a ref, or built through the API), keeps the
+  # canonical document and its digest, and writes the document to Active
+  # Storage when the host has it (see ActionAgent.active_storage), so the
+  # scenarios live in the database, in a file a person can version, and in
+  # the host's storage, never only in code. A set runs by becoming an
+  # evaluation of the agent under test (ScenarioSet#materialize!), so every
+  # run's results, traces and recordings are the evaluation's.
+  class ScenarioCatalog < ApplicationRecord
+    include Ownable
+    owned_by :account, :user
+
+    SOURCE_KINDS = %w[upload repository api].freeze
+    KEY_FORMAT = /\A[A-Za-z0-9][A-Za-z0-9_.\/-]*\z/
+    # Bounds an import: the document's bytes and the scenarios it holds.
+    MAX_BYTES = 2.megabytes
+    MAX_SCENARIOS = 2000
+
+    has_many :products, -> { order(:position, :id) }, class_name: "ActionAgent::ScenarioProduct",
+      foreign_key: :scenario_catalog_id, inverse_of: :catalog, dependent: :destroy
+
+    # The document as last imported or exported, per ActionAgent.active_storage.
+    has_one_attached :document_file, **ActionAgent.attachment_options if ActionAgent.active_storage_macros?
+
+    validates :key, presence: true, format: { with: KEY_FORMAT }, length: { maximum: 120 },
+      uniqueness: { scope: [ :account_id, :user_id ] }
+    validates :name, presence: true, length: { maximum: 200 }
+    validates :source_kind, inclusion: { in: SOURCE_KINDS }
+
+    scope :ordered, -> { order(:name, :id) }
+
+    # Whether the document can be written to Active Storage here.
+    def self.attachments_available?
+      ActionAgent.active_storage_available? && method_defined?(:document_file)
+    end
+
+    def metadata
+      value = self[:metadata]
+      value.is_a?(Hash) ? value : {}
+    end
+
+    def sets
+      ScenarioSet.where(scenario_product_id: products.select(:id))
+    end
+
+    def scenarios
+      CatalogScenario.where(scenario_set_id: sets.select(:id))
+    end
+
+    def scenario_count
+      scenarios.count
+    end
+
+    # The catalog as the framework reads it, built from the records.
+    # @return [ActiveAgent::Evals::Catalog]
+    def to_catalog
+      ActiveAgent::Evals::Catalog.new([ to_document ], name: key)
+    end
+
+    # The canonical YAML of the records: the whole catalog, or only the
+    # products, sets or scenarios named by id (with the products and sets
+    # that hold them), as a document of this catalog's key that a merge
+    # import (ScenarioCatalogImport mode "merge") reads back into it.
+    def export_yaml(product_ids: nil, set_ids: nil, scenario_ids: nil)
+      selection = { product_ids: ids(product_ids), set_ids: ids(set_ids), scenario_ids: ids(scenario_ids) }
+      return to_catalog.to_yaml if selection.values.all?(&:nil?)
+
+      ActiveAgent::Evals::Catalog.new([ to_document(**selection) ], name: key).to_yaml
+    end
+
+    # A key for a new product, set or scenario: +text+ as a key, made unique
+    # among +taken+ with a numeric suffix.
+    def self.key_for(text, taken)
+      base = text.to_s.parameterize(separator: "_").first(60).presence || "item"
+      base = "k_#{base}" unless base.match?(KEY_FORMAT)
+      candidate = base
+      suffix = 1
+      candidate = "#{base}_#{suffix += 1}" while taken.include?(candidate)
+      candidate
+    end
+
+    # Rewrites the stored document and digest from the records, after a
+    # change made through the API rather than an import.
+    def refresh_document!
+      catalog = to_catalog
+      update!(document: catalog.to_yaml, digest: catalog.digest)
+    end
+
+    # Writes the document to Active Storage when the host has it, unless
+    # the attached file already carries this digest.
+    # @return [Boolean] whether a file is attached now
+    def sync_to_storage!
+      return false unless self.class.attachments_available?
+      return true if document_file.attached? && synced_digest == digest
+
+      document_file.attach(io: StringIO.new(document.to_s), filename: "#{key.tr('/', '-')}.yml", content_type: "application/x-yaml")
+      update!(synced_at: Time.current, synced_digest: digest)
+      true
+    end
+
+    # Replaces the records with the attached document, for a database
+    # restored behind its storage.
+    # @return [ScenarioCatalog]
+    def restore_from_storage!
+      raise ActiveRecord::RecordNotFound, "no document is attached to catalog #{key}" unless self.class.attachments_available? && document_file.attached?
+
+      ScenarioCatalogImport.new(owner: owner, document: document_file.download, catalog: self, source_kind: source_kind, source_path: source_path).call
+    end
+
+    def synced?
+      synced_digest.present? && synced_digest == digest
+    end
+
+    def summary
+      {
+        id: id,
+        key: key,
+        name: name,
+        description: description,
+        source_kind: source_kind,
+        source_path: source_path,
+        digest: digest,
+        synced: synced?,
+        synced_at: synced_at,
+        storage_available: self.class.attachments_available?,
+        product_count: products.size,
+        scenario_count: scenario_count,
+        updated_at: updated_at
+      }
+    end
+
+    private
+
+    # The catalog's document, narrowed to the products, sets and scenarios
+    # named by id when any are (see #export_yaml). A product or set is kept
+    # whole when it is named, and otherwise only for what it holds that is.
+    def to_document(product_ids: nil, set_ids: nil, scenario_ids: nil)
+      narrowed = [ product_ids, set_ids, scenario_ids ].any?
+      documents = products.includes(sets: :scenarios).filter_map do |product|
+        next product.to_document unless narrowed
+        next product.to_document if product_ids&.include?(product.id)
+
+        sets = product.sets.filter_map do |set|
+          next set.to_document if set_ids&.include?(set.id)
+
+          chosen = set.scenarios.select { |scenario| scenario_ids&.include?(scenario.id) }
+          set.to_document.merge("scenarios" => chosen.map(&:to_document)) if chosen.any?
+        end
+        product.to_document.merge("sets" => sets) if sets.any?
+      end
+      {
+        "catalog" => key,
+        "name" => name,
+        "description" => description,
+        "products" => documents
+      }.merge(metadata).compact
+    end
+
+    def ids(value)
+      return nil if value.blank?
+
+      Array(value).flat_map { |entry| entry.to_s.split(",") }.filter_map { |entry| Integer(entry.strip, exception: false) }
+    end
+  end
+end

data/app/models/action_agent/scenario_product.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module ActionAgent
+  # A product in a scenario catalog: what its sets of scenarios run against.
+  # That is a dashboard agent, a project (whose booted app the replays reach,
+  # with its sandbox's browser), or, until one is chosen, the `agent:` name
+  # the document gave.
+  class ScenarioProduct < ApplicationRecord
+    belongs_to :catalog, class_name: "ActionAgent::ScenarioCatalog", foreign_key: :scenario_catalog_id, inverse_of: :products
+    belongs_to :agent, class_name: "ActionAgent::Agent", optional: true
+    belongs_to :project, class_name: "ActionAgent::Project", optional: true
+    has_many :sets, -> { order(:position, :id) }, class_name: "ActionAgent::ScenarioSet",
+      foreign_key: :scenario_product_id, inverse_of: :product, dependent: :destroy
+
+    validates :key, presence: true, format: { with: ScenarioCatalog::KEY_FORMAT }, length: { maximum: 120 },
+      uniqueness: { scope: :scenario_catalog_id }
+    validates :name, presence: true, length: { maximum: 200 }
+
+    def metadata
+      value = self[:metadata]
+      value.is_a?(Hash) ? value : {}
+    end
+
+    # The agent a set of this product runs against when the caller names
+    # none: the product's agent, else the project's agent under test, else
+    # the owner's agent named by the document.
+    # @param candidates [ActiveRecord::Relation, nil] the agents the caller may run
+    # @return [Agent, nil]
+    def target_agent(candidates = nil)
+      agent || project&.target_agent || (agent_name.present? && candidates ? candidates.find_by(name: agent_name) : nil)
+    end
+
+    def to_document
+      {
+        "key" => key,
+        "name" => name,
+        "description" => description,
+        "agent" => agent_name.presence || agent&.name,
+        "sets" => sets.map(&:to_document)
+      }.merge(metadata).compact
+    end
+
+    def summary
+      {
+        id: id,
+        key: key,
+        name: name,
+        description: description,
+        agent_name: agent_name,
+        agent: agent && { id: agent.id, name: agent.name },
+        project: project && { id: project.id, name: project.name, repository: project.repository, ref: project.checkout_ref },
+        set_count: sets.size,
+        scenario_count: sets.sum { |set| set.scenarios.size }
+      }
+    end
+  end
+end

data/app/models/action_agent/scenario_set.rb

@@ -0,0 +1,138 @@
+# frozen_string_literal: true
+
+module ActionAgent
+  # A named set of scenarios in a catalog product: a smoke set, the UAT of
+  # one pull request, a regression list. A set runs by becoming an
+  # evaluation of the agent under test (#materialize!), named
+  # "<catalog>/<product>/<set>", whose scenarios are replaced from the set
+  # each time, so the evaluation's runs, results, traces and recordings are
+  # the set's history, and the catalog document can change without changing
+  # what an earlier run asked (EvaluationScenarioResult keeps a snapshot).
+  class ScenarioSet < ApplicationRecord
+    # Raised by #materialize! and #run! when no agent is known to run against.
+    class NoAgent < StandardError; end
+
+    belongs_to :product, class_name: "ActionAgent::ScenarioProduct", foreign_key: :scenario_product_id, inverse_of: :sets
+    belongs_to :evaluation, class_name: "ActionAgent::Evaluation", optional: true
+    has_many :scenarios, -> { order(:position, :id) }, class_name: "ActionAgent::CatalogScenario",
+      foreign_key: :scenario_set_id, inverse_of: :set, dependent: :destroy
+
+    validates :key, presence: true, format: { with: ScenarioCatalog::KEY_FORMAT }, length: { maximum: 120 },
+      uniqueness: { scope: :scenario_product_id }
+    validates :name, presence: true, length: { maximum: 200 }
+
+    delegate :catalog, to: :product
+
+    def judge
+      value = self[:judge]
+      value.is_a?(Hash) ? value : {}
+    end
+
+    def criteria
+      Array(self[:criteria]).select { |criterion| criterion.is_a?(Hash) }
+    end
+
+    def metadata
+      value = self[:metadata]
+      value.is_a?(Hash) ? value : {}
+    end
+
+    # The evaluation's name: the catalog, product and set keys.
+    def evaluation_name
+      "#{catalog.key}/#{product.key}/#{key}".first(EvaluationReportImport::MAX_EVALUATION_NAME)
+    end
+
+    # The reference an evaluation and its runs keep to this set.
+    def catalog_reference
+      {
+        "catalog_id" => catalog.id,
+        "catalog_key" => catalog.key,
+        "product_key" => product.key,
+        "set_key" => key,
+        "set_id" => id,
+        "digest" => catalog.digest
+      }
+    end
+
+    # Creates or refreshes the evaluation of +agent+ this set runs as, with
+    # the set's judge and criteria, and its scenarios replaced from the
+    # set's (a scenario the set dropped is disabled, not destroyed, so runs
+    # that scored it still resolve their results). Remembers the evaluation.
+    #
+    # @param agent [Agent] the agent under test; the product's by default
+    # @return [Evaluation]
+    def materialize!(agent: nil)
+      agent ||= product.target_agent
+      raise NoAgent, "Choose the agent to run #{evaluation_name} against: the product names none" if agent.nil?
+
+      attributes = scenarios.each_with_index.map { |scenario, index| scenario.evaluation_attributes(index) }
+      transaction do
+        target = evaluation if evaluation&.agent_id == agent.id
+        target ||= agent.evaluations.find_or_initialize_by(name: evaluation_name)
+        target.judge_kind = judge["kind"].to_s.presence || target.judge_kind.presence || "rules"
+        target.judge_model = judge["model"].to_s.presence if judge["model"].present?
+        target.criteria = criteria if criteria.present?
+        target.config = target.config.merge("catalog" => catalog_reference)
+        # An evaluation with neither criteria nor scenarios is invalid, so a
+        # new one gets its scenarios before the save.
+        attributes.each { |attrs| target.scenarios.build(attrs.except("enabled")) } if target.new_record?
+        target.save!
+        target.replace_scenarios!(attributes, on_removed: :disable) unless target.previously_new_record?
+        update!(evaluation: target) if evaluation_id != target.id
+        target
+      end
+    end
+
+    # Materializes and runs the set. Against a project, the run waits for the
+    # project's sandbox and reaches it with its browser, as the project's own
+    # evaluation does (ProjectEvaluationJob); otherwise the evaluation runs
+    # in the background with the given selection (models, keys).
+    #
+    # @return [EvaluationRun]
+    def run!(agent: nil, project: nil, mount_url: nil, **selection)
+      agent ||= project&.target_agent || product.target_agent
+      if project && agent.nil?
+        raise NoAgent, "Choose the agent to evaluate in project #{project.name} first"
+      end
+
+      target = materialize!(agent: agent)
+      reference = { "catalog" => catalog_reference }
+      if project
+        sandbox = project.current_sandbox_session
+        raise NoAgent, "Boot project #{project.name} before running #{evaluation_name} against it" if sandbox.nil?
+
+        run = target.evaluation_runs.create!(status: :pending,
+          selection: reference.merge("project_id" => project.id, "sandbox_id" => sandbox.session_id).merge(selection.deep_stringify_keys))
+        ProjectEvaluationJob.perform_later(project.id, run.id, mount_url)
+        run
+      else
+        target.run_later!(**selection.deep_stringify_keys.merge(reference).deep_symbolize_keys)
+      end
+    end
+
+    def to_document
+      {
+        "key" => key,
+        "name" => name,
+        "description" => description,
+        "judge" => judge.presence,
+        "criteria" => criteria.presence,
+        "scenarios" => scenarios.map(&:to_document)
+      }.merge(metadata).compact
+    end
+
+    def summary
+      {
+        id: id,
+        key: key,
+        name: name,
+        description: description,
+        judge: judge,
+        criteria_count: criteria.size,
+        scenario_count: scenarios.size,
+        evaluation: evaluation && { id: evaluation.id, name: evaluation.name, agent_id: evaluation.agent_id },
+        latest_run: evaluation&.latest_run&.then { |run| { id: run.id, status: run.status, created_at: run.created_at } }
+      }
+    end
+  end
+end

data/app/models/action_agent/session_recording.rb

@@ -7,13 +7,36 @@ module ActionAgent
 
     belongs_to :agent_run, optional: true
     belongs_to :sandbox_session, optional: true
+    # The conversation a recording belongs to. A recording of a conversation
+    # outlives any one of its runs.
+    belongs_to :agent_context, optional: true
 
     has_many :recording_actions, dependent: :destroy
     has_many :recording_snapshots, dependent: :destroy
+    # Deleted rather than destroyed: destroying an event looks up its
+    # attachment, and a host that loads Active Storage without migrating it
+    # has no table to look in. Attached payloads are purged first.
+    has_many :recording_events, dependent: :delete_all
+
+    before_destroy :purge_attached_event_payloads, prepend: true
 
     enum :status, { recording: 0, completed: 1, failed: 2 }
 
+    # Where a recording came from: an agent's browser, or a person using the
+    # dashboard. Recordings made before this was recorded have none.
+    SOURCES = %w[agent dashboard].freeze
+
+    # Browser state that must never leave the server in a read response: the
+    # handoff state a recording carries is a copy of the visitor's cookies
+    # and web storage. Only the handoff endpoint returns it, to the owner,
+    # when they continue the session.
+    SENSITIVE_STATE_KEYS = %w[cookies session_storage local_storage].freeze
+
+    INGEST_TOKEN_PREFIX = "aarec_"
+    INGEST_TOKEN_TTL = 2.hours
+
     validates :status, presence: true
+    validates :source, inclusion: { in: SOURCES }, allow_nil: true
     validate :must_have_parent, unless: -> { demo_recording? || user_session? }
 
     scope :recent, -> { order(created_at: :desc) }
@@ -36,18 +59,20 @@ module ActionAgent
     # The owner column is written here, at creation: the index and recent
     # endpoints scope through it, and a recording nothing ever stamped was
     # invisible in the list to the very person who made it. An explicit
-    # +owner+ wins; otherwise the recording inherits the owner of the sandbox
-    # or agent it records. Ownable#owner= is a no-op in a single-user
-    # install, where nothing is owned.
-    def self.start!(agent_run: nil, sandbox_session: nil, name: nil, owner: nil)
+    # +owner+ wins; otherwise the recording inherits the owner of the sandbox,
+    # agent or conversation it records. Ownable#owner= is a no-op in a
+    # single-user install, where nothing is owned.
+    def self.start!(agent_run: nil, sandbox_session: nil, agent_context: nil, source: nil, name: nil, owner: nil)
       recording = new(
         agent_run: agent_run,
         sandbox_session: sandbox_session,
-        name: name || generate_name(agent_run, sandbox_session),
+        agent_context: agent_context,
+        source: source,
+        name: name || generate_name(agent_run, sandbox_session, agent_context),
         status: :recording,
         metadata: { started_at: Time.current.iso8601 }
       )
-      recording.owner = owner || inherited_owner(agent_run, sandbox_session)
+      recording.owner = owner || inherited_owner(agent_run, sandbox_session, agent_context)
       recording.save!
       recording
     end
@@ -71,11 +96,73 @@ module ActionAgent
       recording
     end
 
-    # Whoever owns the sandbox or agent a recording is made against. Both
-    # models declare the same owner candidates as this one, so the record
-    # they hand back is of the class this install owns things through.
-    def self.inherited_owner(agent_run, sandbox_session)
-      sandbox_session&.owner || agent_run&.agent&.owner
+    # Whoever owns the sandbox, agent or conversation's agent a recording is
+    # made against. Both models declare the same owner candidates as this
+    # one, so the record they hand back is of the class this install owns
+    # things through.
+    def self.inherited_owner(agent_run, sandbox_session, agent_context = nil)
+      conversation_agent = agent_context&.contextable
+      sandbox_session&.owner || agent_run&.agent&.owner || (conversation_agent.owner if conversation_agent.is_a?(Agent))
+    end
+
+    # Returns +value+ without SENSITIVE_STATE_KEYS in any Hash it holds, at
+    # any depth.
+    def self.without_sensitive_state(value)
+      case value
+      when Hash
+        value.each_with_object({}) do |(key, item), clean|
+          next if SENSITIVE_STATE_KEYS.include?(key.to_s)
+
+          clean[key] = without_sensitive_state(item)
+        end
+      when Array then value.map { |item| without_sensitive_state(item) }
+      else value
+      end
+    end
+
+    # The digest an ingest token is stored and looked up as.
+    def self.ingest_token_digest(token)
+      Digest::SHA256.hexdigest(token.to_s)
+    end
+
+    # Issues the token a browser posts this recording's events with, and
+    # returns it. Only its digest is stored, and issuing another replaces it.
+    # The token expires after +expires_in+, or when the recording's sandbox
+    # does if that is sooner. See #ingest_token_valid? for when it is
+    # accepted.
+    # @return [String]
+    def issue_ingest_token!(expires_in: INGEST_TOKEN_TTL)
+      token = "#{INGEST_TOKEN_PREFIX}#{SecureRandom.base58(32)}"
+      expires_at = [ Time.current + expires_in, sandbox_session&.expires_at ].compact.min
+      update!(ingest_token_digest: self.class.ingest_token_digest(token), ingest_token_expires_at: expires_at)
+      token
+    end
+
+    # Whether +token+ may post events to this recording now: it is the token
+    # last issued for this recording, it has not expired, the recording is
+    # still recording, and the recording's sandbox, if it has one, is active.
+    def ingest_token_valid?(token)
+      return false if token.blank? || ingest_token_digest.blank?
+      return false unless recording? && ingest_token_expires_at&.future?
+      return false if sandbox_session_id && !sandbox_live?
+
+      ActiveSupport::SecurityUtils.secure_compare(ingest_token_digest, self.class.ingest_token_digest(token))
+    end
+
+    # Stores one event the server observed, such as an agent's browser
+    # action, and counts it. Unlike a browser's batches, server events are
+    # not capped.
+    # @param data [Hash] the event
+    # @param started_at [Time] when it began
+    # @param finished_at [Time] when it ended
+    # @return [RecordingEvent]
+    def record_server_event!(kind:, data:, started_at:, finished_at: started_at)
+      event = recording_events.new(kind: kind, clock_offset_ms: 0)
+      event.events = [ { "at" => RecordingEvent.milliseconds(started_at), "data" => data } ]
+      event.occurred_to = [ finished_at, event.occurred_from ].max
+      event.save!
+      self.class.update_counters(id, event_count: 1, event_bytes: event.byte_size)
+      event
     end
 
     # Record a browser action
@@ -136,10 +223,14 @@ module ActionAgent
     # same redaction the /actions endpoint applies: this is what #show
     # renders, and it used to hand back the cleartext password that
     # /actions had just redacted for the same action.
+    #
+    # Each entry names its action `action_type`, as /actions entries do.
+    # `type` carries the same value for readers of the older key.
     def timeline
       recording_actions.order(:sequence).map do |action|
         {
           id: action.id,
+          action_type: action.action_type,
           type: action.action_type,
           sequence: action.sequence,
           timestamp_ms: action.timestamp_ms,
@@ -153,17 +244,33 @@ module ActionAgent
 
     private
 
+    def purge_attached_event_payloads
+      return unless RecordingEvent.attachments_available?
+
+      ActiveStorage::Attachment
+        .where(record_type: RecordingEvent.polymorphic_name, name: "payload_file", record_id: recording_events.select(:id))
+        .includes(:blob)
+        .find_each(&:purge_later)
+    end
+
+    def sandbox_live?
+      live = SandboxSession.active.where(id: sandbox_session_id)
+      live.where(expires_at: nil).or(live.where(expires_at: Time.current..)).exists?
+    end
+
     def must_have_parent
-      return if agent_run.present? || sandbox_session.present?
+      return if agent_run.present? || sandbox_session.present? || agent_context.present?
 
-      errors.add(:base, "must belong to an agent_run or sandbox_session")
+      errors.add(:base, "must belong to an agent_run, sandbox_session or agent_context")
     end
 
-    def self.generate_name(agent_run, sandbox_session)
+    def self.generate_name(agent_run, sandbox_session, agent_context = nil)
       prefix = if agent_run&.agent
         agent_run.agent.name.parameterize
       elsif sandbox_session&.agent_template
         sandbox_session.agent_template.name.parameterize
+      elsif agent_context&.agent_name.present?
+        agent_context.agent_name.parameterize
       else
         "session"
       end

data/app/models/action_agent/telemetry_trace.rb

@@ -47,7 +47,28 @@ module ActionAgent
     # The release of the agent this trace came from (see #attach_agent_version!).
     belongs_to :agent_version, class_name: "ActionAgent::AgentVersion", optional: true
 
-    scope :for_account, ->(account) { where(account: account) if ActionAgent.multi_tenant? }
+    # The tenant's traces in multi-tenant mode; every trace otherwise. The
+    # account is matched by class, as Ownable.for_owner matches an owner: a
+    # user maps to its tenant through ActionAgent.tenant_for, and nil or any
+    # other class scopes to nothing.
+    scope :for_account, lambda { |account|
+      next unless ActionAgent.multi_tenant?
+
+      account = tenant_account(account)
+      account ? where(account_id: account.id) : none
+    }
+
+    # +account+ as an instance of the configured account class, or nil.
+    def self.tenant_account(account)
+      return nil if account.nil?
+
+      account_class = ActionAgent.account_class.presence&.safe_constantize
+      return nil if account_class.nil?
+      return account if account.is_a?(account_class)
+
+      tenant = ActionAgent.tenant_for(account)
+      tenant if tenant.is_a?(account_class)
+    end
 
     # Creates a TelemetryTrace from an ingested trace payload.
     #
@@ -192,24 +213,7 @@ module ActionAgent
       root_span = spans.find { |s| s["parent_span_id"].nil? } || spans.first || {}
 
       total_duration = root_span["duration_ms"]
-
-      # Instrumentation mirrors LLM token usage onto the root span for
-      # display, so summing every span double-counts. When child spans
-      # carry token data, they are the source of truth; the root span only
-      # counts for single-span traces.
-      counted_spans = spans.reject { |s| s["parent_span_id"].nil? }
-      counted_spans = spans if counted_spans.none? { |s| span_token_sum(s).positive? }
-
-      total_input = 0
-      total_output = 0
-      total_thinking = 0
-
-      counted_spans.each do |span|
-        tokens = span["tokens"] || {}
-        total_input += (tokens["input"] || 0)
-        total_output += (tokens["output"] || 0)
-        total_thinking += (tokens["thinking"] || 0)
-      end
+      total_input, total_output, total_thinking = token_totals(spans)
 
       # Extract agent info from root span attributes
       attributes = root_span["attributes"] || {}
@@ -262,6 +266,57 @@ module ActionAgent
     end
 
 
+    # The input, output and thinking tokens +spans+ add up to.
+    #
+    # Instrumentation mirrors LLM token usage onto the root span for display,
+    # so summing every span double-counts. When child spans carry token data,
+    # they are the source of truth; the root span only counts for
+    # single-span traces.
+    #
+    # @return [Array(Integer, Integer, Integer)]
+    # @api private
+    def self.token_totals(spans)
+      counted = spans.reject { |s| s["parent_span_id"].nil? }
+      counted = spans if counted.none? { |s| span_token_sum(s).positive? }
+
+      counted.each_with_object([ 0, 0, 0 ]) do |span, totals|
+        tokens = span["tokens"] || {}
+        totals[0] += (tokens["input"] || 0)
+        totals[1] += (tokens["output"] || 0)
+        totals[2] += (tokens["thinking"] || 0)
+      end
+    end
+
+    # Adds a later execution segment of the same run to this trace: a run that
+    # paused for input records its resumed segment under the same trace id.
+    # The segment's root span becomes a child of this trace's root, and the
+    # totals, status and error cover both.
+    #
+    # @param segment [Array<Hash>] the segment's flattened spans
+    # @return [self]
+    def append_segment!(segment)
+      parent_id = root_span&.dig("span_id")
+      segment = Array(segment).map do |span|
+        span = span.to_h.stringify_keys
+        span["parent_span_id"].nil? && parent_id ? span.merge("parent_span_id" => parent_id) : span
+      end
+      combined = Array(spans) + segment
+      input, output, thinking = self.class.token_totals(combined)
+      segment_root = segment.find { |span| span["parent_span_id"] == parent_id }
+      error_span = segment.find { |span| span["status"] == STATUS_ERROR }
+
+      update!(
+        spans: combined,
+        total_input_tokens: input,
+        total_output_tokens: output,
+        total_thinking_tokens: thinking,
+        total_duration_ms: total_duration_ms.to_f + segment_root&.dig("duration_ms").to_f,
+        status: error_span ? STATUS_ERROR : status,
+        error_message: error_message.presence || error_span&.dig("attributes", "error.message")
+      )
+      self
+    end
+
     # Sums a span's token counts (used to decide which spans carry the
     # authoritative token data during ingestion).
     #