activeagent 1.2.0 → 1.3.0

data/lib/active_agent/delegation/runner.rb

@@ -0,0 +1,283 @@
+# frozen_string_literal: true
+
+require "timeout"
+
+module ActiveAgent
+  module Delegation
+    # Executes one delegated call: budget check, sub-agent generation, result
+    # validation, accounting.
+    #
+    # The runner is what makes a sub-agent a *tool* rather than a method call.
+    # It answers to the calling model in the model's own terms — a refused call
+    # comes back as a structured result the model can reason about, not as an
+    # exception that ends the conversation.
+    #
+    # @api private
+    class Runner
+      # @return [Definition]
+      attr_reader :definition
+      # @return [ActiveAgent::Base] the delegating agent instance
+      attr_reader :owner
+
+      # @param definition [Definition]
+      # @param owner [ActiveAgent::Base]
+      def initialize(definition, owner:)
+        @definition = definition
+        @owner      = owner
+      end
+
+      # Runs the delegated call.
+      #
+      # @param arguments [Hash] arguments supplied by the calling model
+      # @return [Object] the sub-agent's result, or a structured error the
+      #   calling model can act on
+      # @raise [BudgetExceededError] when the budget policy is +:raise+
+      # @raise [InvalidResultError] when the returns policy is +:raise+
+      def call(**arguments)
+        if (violation = budget_violation)
+          return refuse(violation)
+        end
+
+        arguments = coerce_arguments(arguments)
+
+        instrument(arguments) do |payload|
+          started  = clock
+          response = nil
+
+          begin
+            response = with_timeout { generate(arguments) }
+          rescue Timeout::Error
+            duration = clock - started
+            record(duration: duration)
+            payload[:status]      = :timed_out
+            payload[:duration_ms] = (duration * 1_000).round
+
+            next timed_out(duration)
+          end
+
+          duration = clock - started
+          usage    = response.usage
+          model    = response.model.presence || generation_model(response)
+          cost     = Pricing.cost_for(usage: usage, model: model, rates: budget.rates)
+
+          record(tokens: usage&.total_tokens, cost: cost, duration: duration)
+
+          payload[:status]      = :ok
+          payload[:model]       = model
+          payload[:duration_ms] = (duration * 1_000).round
+          payload[:usage]       = usage
+          payload[:cost]        = cost
+          payload[:ledger]      = ledgers.last.to_h
+
+          result(response, payload)
+        end
+      end
+
+      private
+
+      # @return [Budget]
+      def budget = definition.budget
+
+      # Agent-wide ledger first, then this delegation's own — a call has to
+      # clear both.
+      #
+      # @return [Array<Ledger>]
+      def ledgers
+        @ledgers ||= [ owner.delegation_ledger, owner.delegation_ledger_for(definition.tool_name) ]
+      end
+
+      # @return [Budget::Violation, nil]
+      def budget_violation
+        owner.class.delegation_budget.violation_for(ledgers.first) || budget.violation_for(ledgers.last)
+      end
+
+      # @param tokens [Integer, nil]
+      # @param cost [Float, nil]
+      # @param duration [Float]
+      # @return [void]
+      def record(tokens: 0, cost: nil, duration: 0.0)
+        ledgers.each { |ledger| ledger.record(tokens: tokens, cost: cost, duration: duration) }
+      end
+
+      # Drops arguments the contract never declared, so a model that
+      # hallucinates an extra key gets a working call instead of an
+      # +ArgumentError+ from the sub-agent's method signature.
+      #
+      # @param arguments [Hash]
+      # @return [Hash]
+      def coerce_arguments(arguments)
+        arguments = arguments.symbolize_keys
+        return arguments if definition.schema.empty?
+
+        arguments.slice(*definition.schema.keys)
+      end
+
+      # @param arguments [Hash]
+      # @return [ActiveAgent::Providers::Common::PromptResponse]
+      def generate(arguments)
+        agent = definition.resolved_agent_class.new
+        agent.params = resolved_params(arguments)
+        agent.process(definition.action, **arguments)
+
+        definition.backend.apply(agent)
+        apply_returns_format(agent)
+        inherit_trace_id(agent)
+
+        agent.process_prompt
+      end
+
+      # A delegated generation is part of its parent's work, so it carries the
+      # parent's trace id — otherwise the sub-agent's tokens and latency land
+      # in a separate trace and the budget you set has nothing to show for it.
+      #
+      # @param agent [ActiveAgent::Base]
+      # @return [void]
+      def inherit_trace_id(agent)
+        trace_id = owner.prompt_options[:trace_id]
+        agent.prompt_options[:trace_id] ||= trace_id if trace_id
+      end
+
+      # A declared +returns+ schema *is* the response format for the delegated
+      # call; the sub-agent does not have to restate it.
+      #
+      # @param agent [ActiveAgent::Base]
+      # @return [void]
+      def apply_returns_format(agent)
+        return unless definition.structured?
+
+        agent.prompt_options[:response_format] = {
+          type: :json_schema,
+          json_schema: definition.returns.to_response_format(name: "#{definition.tool_name}_result")
+        }
+      end
+
+      # @param arguments [Hash]
+      # @return [Hash]
+      def resolved_params(arguments)
+        case (params = definition.params)
+        when nil    then {}
+        when Hash   then params
+        when Symbol then owner.send(params)
+        when Proc   then params.arity == 1 ? owner.instance_exec(arguments, &params) : owner.instance_exec(&params)
+        else
+          raise ArgumentError, "Delegation params must be a Hash, Symbol or Proc, got #{params.inspect}"
+        end.to_h.symbolize_keys
+      end
+
+      # @yield the delegated generation
+      # @return [Object]
+      def with_timeout
+        return yield unless budget.timeout
+
+        Timeout.timeout(budget.timeout) { yield }
+      end
+
+      # @param response [ActiveAgent::Providers::Common::PromptResponse]
+      # @param payload [Hash] instrumentation payload
+      # @return [Object]
+      def result(response, payload)
+        message = response.message
+        return nil if message.nil?
+
+        return text_of(message) unless definition.structured?
+
+        parsed  = message.respond_to?(:parsed_json) ? message.parsed_json : nil
+        missing = definition.returns.missing_keys(parsed)
+        return parsed if missing.empty?
+
+        payload[:status]  = :invalid_result
+        payload[:missing] = missing
+
+        invalid(missing, text_of(message))
+      end
+
+      # @param message [Object]
+      # @return [String]
+      def text_of(message)
+        message.respond_to?(:text) ? message.text : message.content.to_s
+      end
+
+      # @param response [ActiveAgent::Providers::Common::PromptResponse]
+      # @return [String, nil]
+      def generation_model(response)
+        response.context.is_a?(Hash) ? response.context[:model] : nil
+      end
+
+      # @param violation [Budget::Violation]
+      # @return [Hash]
+      def refuse(violation)
+        ActiveSupport::Notifications.instrument("delegation_refused.active_agent", instrument_payload.merge(
+          status: :budget_exceeded, limit: violation.limit, allowed: violation.allowed, used: violation.used
+        ))
+
+        if budget.policy == :raise
+          raise BudgetExceededError.new(violation, definition: definition)
+        end
+
+        { error: "budget_exceeded", limit: violation.limit.to_s, allowed: violation.allowed, used: violation.used, message: violation.message }
+      end
+
+      # @param duration [Float]
+      # @return [Hash]
+      def timed_out(duration)
+        if budget.policy == :raise
+          raise TimeoutError, "#{definition.agent_class}##{definition.action} exceeded its #{budget.timeout}s delegation timeout"
+        end
+
+        {
+          error: "timeout",
+          limit: "timeout",
+          allowed: budget.timeout,
+          used: duration.round(3),
+          message: "The #{definition.tool_name} delegation timed out after #{budget.timeout}s. " \
+                   "Answer with the information you already have, or call it with a smaller request."
+        }
+      end
+
+      # @param missing [Array<Symbol>]
+      # @param content [String]
+      # @return [Hash]
+      def invalid(missing, content)
+        if definition.on_invalid == :raise
+          raise InvalidResultError, "#{definition.agent_class}##{definition.action} returned a result missing " \
+            "required #{"key".pluralize(missing.size)}: #{missing.join(", ")}"
+        end
+
+        {
+          error: "invalid_result",
+          missing: missing.map(&:to_s),
+          message: "The #{definition.tool_name} delegation returned a result missing required " \
+                   "#{"key".pluralize(missing.size)}: #{missing.join(", ")}.",
+          content: content
+        }
+      end
+
+      # @param arguments [Hash]
+      # @yield [Hash] instrumentation payload
+      def instrument(arguments, &block)
+        ActiveSupport::Notifications.instrument(
+          "delegate.active_agent", instrument_payload.merge(arguments: arguments), &block
+        )
+      end
+
+      # @return [Hash]
+      def instrument_payload
+        {
+          agent: owner.class.name,
+          delegate: definition.agent_class.name,
+          action: definition.action,
+          tool: definition.tool_name.to_s,
+          provider: definition.backend.provider,
+          budget: budget.to_h.except(:rates)
+        }.compact
+      end
+
+      # Monotonic so a clock adjustment mid-generation can't corrupt a latency budget.
+      #
+      # @return [Float]
+      def clock
+        Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      end
+    end
+  end
+end

data/lib/active_agent/delegation/schema.rb

@@ -0,0 +1,220 @@
+# frozen_string_literal: true
+
+module ActiveAgent
+  module Delegation
+    # Declarative JSON Schema for a delegated agent's inputs and outputs.
+    #
+    # A delegation is only as good as its contract: the calling model needs to
+    # know exactly what a sub-agent accepts, and the calling *agent* needs to
+    # know what shape comes back. Schema builds both from a small DSL, from a
+    # plain JSON Schema hash, or from any class that responds to
+    # +to_json_schema+ (see {ActiveAgent::SchemaGenerator}).
+    #
+    # @example DSL
+    #   Schema.build do
+    #     string  :text, required: true, description: "Document to summarize"
+    #     integer :max_points, description: "How many bullets to return"
+    #     array   :tags, of: :string, description: "Topic tags"
+    #   end
+    #
+    # @example Plain JSON Schema
+    #   Schema.build(type: "object", properties: { text: { type: "string" } }, required: [ "text" ])
+    #
+    # @example ActiveModel / ActiveRecord
+    #   Schema.build(ContactForm) # ContactForm includes ActiveAgent::SchemaGenerator
+    class Schema
+      # Scalar types that get a one-line DSL helper.
+      SCALAR_TYPES = %i[string integer number boolean].freeze
+
+      class << self
+        # Coerces any supported schema source into a Schema.
+        #
+        # @param source [Schema, Hash, Class, nil] existing schema, raw JSON Schema,
+        #   or a class responding to +to_json_schema+
+        # @yield DSL block (evaluated against a new Schema, or yielded when arity is 1)
+        # @return [Schema]
+        # @raise [ArgumentError] when the source cannot be interpreted
+        def build(source = nil, &block)
+          schema =
+            case source
+            when Schema then source
+            when Hash   then from_hash(source)
+            when nil    then new
+            else
+              if source.respond_to?(:to_json_schema)
+                from_hash(source.to_json_schema)
+              else
+                raise ArgumentError, "Cannot build a delegation schema from #{source.inspect}. " \
+                  "Pass a Hash, a class responding to #to_json_schema, or use the block DSL."
+              end
+            end
+
+          if block
+            block.arity == 1 ? block.call(schema) : schema.instance_eval(&block)
+          end
+
+          schema
+        end
+
+        # @param hash [Hash] JSON Schema object
+        # @return [Schema]
+        def from_hash(hash)
+          hash = hash.deep_symbolize_keys
+
+          new.tap do |schema|
+            (hash[:properties] || {}).each { |name, definition| schema.property(name, definition) }
+            schema.required(*Array(hash[:required]))
+            schema.additional_properties(hash.fetch(:additionalProperties, hash.fetch(:additional_properties, false)))
+          end
+        end
+      end
+
+      def initialize
+        @properties           = {}
+        @required             = []
+        @additional_properties = false
+      end
+
+      # @return [Boolean] true when nothing has been declared
+      def empty?
+        @properties.empty?
+      end
+
+      # Declares a property from an already-built JSON Schema fragment.
+      #
+      # @param name [Symbol, String]
+      # @param definition [Hash] JSON Schema fragment (e.g. +{ type: "string" }+)
+      # @return [void]
+      def property(name, definition)
+        @properties[name.to_sym] = definition.deep_symbolize_keys
+      end
+
+      # Declares a property.
+      #
+      # @param name [Symbol, String]
+      # @param type [Symbol, String] JSON Schema type
+      # @param required [Boolean] whether the calling model must supply it
+      # @param description [String, nil] shown to the calling model — write it for a reader who
+      #   has never seen your code
+      # @param options [Hash] any other JSON Schema keyword (+enum+, +format+, +minimum+, ...)
+      # @yield nested DSL for object properties
+      # @return [void]
+      def param(name, type = :string, required: false, description: nil, **options, &block)
+        definition = { type: type.to_s, description: description }.compact.merge(options)
+
+        if block
+          nested = self.class.build(&block)
+          definition = definition.merge(nested.to_json_schema.except(:type))
+          definition[:type] = "object"
+        end
+
+        property(name, definition)
+        self.required(name) if required
+      end
+
+      SCALAR_TYPES.each do |type|
+        define_method(type) do |name, **options|
+          param(name, type, **options)
+        end
+      end
+
+      # Declares an object property.
+      #
+      # @param name [Symbol, String]
+      # @yield nested DSL
+      def object(name, **options, &block)
+        param(name, :object, **options, &block)
+      end
+
+      # Declares an array property.
+      #
+      # @param name [Symbol, String]
+      # @param of [Symbol, String, Hash, nil] item type, or a JSON Schema fragment for items
+      # @yield nested DSL describing an object item type
+      def array(name, of: nil, **options, &block)
+        items =
+          if block
+            self.class.build(&block).to_json_schema
+          elsif of.is_a?(Hash)
+            of.deep_symbolize_keys
+          elsif of
+            { type: of.to_s }
+          end
+
+        param(name, :array, **options.merge({ items: items }.compact))
+      end
+
+      # Marks properties as required.
+      #
+      # @param names [Array<Symbol, String>]
+      # @return [Array<Symbol>]
+      def required(*names)
+        names.flatten.each { |name| @required |= [ name.to_sym ] }
+        @required
+      end
+
+      # @param value [Boolean]
+      def additional_properties(value = true)
+        @additional_properties = value
+      end
+
+      # @return [Array<Symbol>] declared property names
+      def keys
+        @properties.keys
+      end
+
+      # @return [Array<Symbol>] required property names
+      def required_keys
+        @required.dup
+      end
+
+      # @return [Hash] JSON Schema object
+      #
+      # @note Deep-duplicated: providers normalize tool definitions in place,
+      #   and a schema is shared by every generation that exposes it.
+      def to_json_schema
+        {
+          type: "object",
+          properties: @properties.deep_dup,
+          required: @required.map(&:to_s),
+          additionalProperties: @additional_properties
+        }
+      end
+      alias_method :to_h, :to_json_schema
+
+      # JSON Schema for a +response_format+ payload.
+      #
+      # ActiveAgent camelizes response-format schema *keys* on the way to the
+      # provider, but the +required+ array holds string *values* — so they are
+      # camelized here to keep the emitted schema internally consistent.
+      # {Assistant#parsed_json} underscores the keys again on the way back, so
+      # agent code only ever sees snake_case.
+      #
+      # @param name [String] schema name reported to the provider
+      # @param strict [Boolean]
+      # @return [Hash]
+      def to_response_format(name:, strict: true)
+        schema = to_json_schema
+        schema[:required] = schema[:required].map { |key| key.camelize(:lower) }
+
+        { name: name, schema: schema, strict: strict }
+      end
+
+      # Validates a parsed payload against the declared required properties.
+      #
+      # This is a deliberately shallow check: it catches the failure that
+      # actually happens in practice (a model omitting a required key) without
+      # pulling a full JSON Schema validator into the gem's dependencies.
+      #
+      # @param payload [Hash, nil]
+      # @return [Array<Symbol>] missing required keys
+      def missing_keys(payload)
+        return required_keys if payload.nil?
+        return [] unless payload.is_a?(Hash)
+
+        keys = payload.keys.map { |key| key.to_s.underscore.to_sym }
+        required_keys.reject { |key| keys.include?(key) }
+      end
+    end
+  end
+end

data/lib/active_agent/providers/_base_provider.rb

@@ -54,8 +54,11 @@ module ActiveAgent
       attr_internal :options, :context, :trace_id,   # Setup
                     :request, :message_stack,        # Runtime
                     :stream_broadcaster, :streaming, # Callback (Streams)
+                    :stream_completion_pending,      # Callback (Streams)
+                    :stream_completion_result,       # Callback (Streams)
                     :tools_function,                 # Callback (Tools)
                     :usage_stack,                    # Usage Tracking
+                    :stream_usage_index,             # Usage Tracking (Streams)
                     :max_tool_turns, :tool_turns     # Tool-loop safety
 
       # Upper bound on tool-calling round-trips within one generation. A
@@ -118,6 +121,9 @@ module ActiveAgent
         self.context            = kwargs
         self.message_stack      = []
         self.usage_stack        = []
+        self.stream_completion_pending = false
+        self.stream_completion_result  = nil
+        self.stream_usage_index        = nil
       end
 
       # Generates prompt preview without executing the API call.
@@ -179,6 +185,9 @@ module ActiveAgent
       #
       # @return [ActiveAgent::Providers::Common::PromptResponse]
       def resolve_prompt
+        # Each turn streams its own usage; see record_stream_usage.
+        self.stream_usage_index = nil
+
         api_parameters = api_request_build(prepare_prompt_request, prompt_request_type)
         api_response = instrument("prompt.provider.active_agent") do |payload|
           raw_response = with_exception_handling { api_prompt_execute(api_parameters) }
@@ -193,6 +202,12 @@ module ActiveAgent
           raw_response
         end
 
+        # A stream that deferred its completion has already run
+        # process_prompt_finished (including any tool-call recursion) from
+        # stream_finished!, and holds the response it produced. Calling it a
+        # second time here would re-run that work.
+        return stream_completion_result if stream_completion_result
+
         process_prompt_finished(api_response)
       end
 
@@ -231,7 +246,11 @@ module ActiveAgent
       # @return [Hash] API request parameters
       def api_request_build(request, request_type)
         parameters          = request_type.serialize(request)
-        parameters[:stream] = process_stream if request.try(:stream)
+
+        if request.try(:stream)
+          parameters[:stream] = process_stream
+          parameters.deep_merge!(api_stream_usage_parameters)
+        end
 
         if options.extra_headers.present?
           parameters[:request_options] = { extra_headers: options.extra_headers }.deep_merge(parameters[:request_options] || {})
@@ -240,6 +259,20 @@ module ActiveAgent
         parameters
       end
 
+      # Extra request parameters needed to make the provider report token usage
+      # while streaming.
+      #
+      # A streaming request returns its usage on a final chunk rather than in a
+      # response body, and several providers only send that chunk when the
+      # request asks for it. Providers that need such a flag override this;
+      # the default asks for nothing, so a provider that reports usage
+      # unconditionally — or not at all — is unaffected.
+      #
+      # @return [Hash]
+      def api_stream_usage_parameters
+        {}
+      end
+
       # @return [Proc] for each response chunk
       def process_stream
         proc do |api_response_chunk|
@@ -258,10 +291,27 @@ module ActiveAgent
           api_prompt_executer.create(**parameters)
         else
           api_prompt_executer.stream(**parameters.except(:stream)).each(&parameters[:stream])
+          stream_finished!
           nil
         end
       end
 
+      # Runs the deferred end-of-generation work once a stream has drained.
+      #
+      # A chunk handler that completed the message marks the generation
+      # pending rather than finishing inline, because providers emit their
+      # usage chunk after the content is done — finishing inline would build
+      # the response before usage is recorded. Handlers that finish inline
+      # leave the flag unset and this is a no-op.
+      #
+      # @return [Object, nil] result of process_prompt_finished, if deferred
+      def stream_finished!
+        return unless stream_completion_pending
+
+        self.stream_completion_pending = false
+        self.stream_completion_result  = process_prompt_finished
+      end
+
       # Returns provider-specific API executer for prompt requests.
       #
       # Since all currently implemented providers use stainless gems, subclasses
@@ -305,6 +355,52 @@ module ActiveAgent
         fail NotImplementedError, "Subclass expected to implement"
       end
 
+      # Records token usage carried on a streaming chunk.
+      #
+      # The non-streaming path pushes onto usage_stack in resolve_prompt, from
+      # the response body. A streaming request has no response body — the
+      # provider streams chunks and api_prompt_execute returns nil — so usage
+      # would otherwise be lost, and every streamed generation reports zero
+      # tokens and zero cost. Chunk handlers call this when they see usage so
+      # the two paths converge on the same usage_stack.
+      #
+      # Ignores blank and all-zero payloads: providers send `usage: null` on
+      # ordinary content chunks, and pushing those would add empty entries to
+      # a stack that is summed with reduce(:+).
+      #
+      # A streamed usage payload is a running total for the turn, not a delta,
+      # so the turn keeps a single entry that later payloads replace. Chat
+      # Completions sends exactly one, on a final chunk, but Gemini's
+      # OpenAI-compatible endpoint repeats a cumulative usage on every chunk —
+      # summing those would report a turn's tokens many times over. Tool
+      # calling still accumulates across turns: each turn enters resolve_prompt
+      # and starts a fresh entry.
+      #
+      # @param raw_usage [Hash, Object, nil] provider-shaped usage payload
+      # @return [void]
+      def record_stream_usage(raw_usage)
+        return if raw_usage.blank?
+
+        # from_provider_usage only reads hashes, and the stainless gems hand
+        # back model objects (e.g. OpenAI::Models::CompletionUsage), so an
+        # unconverted object is silently dropped.
+        raw_usage = raw_usage.deep_to_h if raw_usage.respond_to?(:deep_to_h)
+        raw_usage = raw_usage.to_h if !raw_usage.is_a?(Hash) && raw_usage.respond_to?(:to_h)
+
+        usage = Common::Usage.from_provider_usage(raw_usage)
+        return if usage.blank?
+        return if usage.total_tokens.to_i.zero? &&
+                  usage.input_tokens.to_i.zero? &&
+                  usage.output_tokens.to_i.zero?
+
+        if stream_usage_index
+          usage_stack[stream_usage_index] = usage
+        else
+          self.stream_usage_index = usage_stack.length
+          usage_stack.push(usage)
+        end
+      end
+
       # Broadcasts stream open event.
       #
       # Fires once per request cycle, even during multi-turn tool calling.