smith-agents 0.5.0 → 0.6.0

data/lib/smith/workflow/composite/plan.rb

@@ -0,0 +1,130 @@
+# frozen_string_literal: true
+
+require_relative "../../types"
+require_relative "../../version"
+require_relative "../prepared_step_dispatch"
+require_relative "branch"
+require_relative "enums"
+require_relative "payload"
+require_relative "payload_digest"
+require_relative "plan_integrity"
+
+module Smith
+  class Workflow
+    module Composite
+      class Plan < Payload
+        VERSION = 1
+        MAX_BRANCHES = 10_000
+        RESUME_POLICY = :incomplete_only
+        FAILURE_POLICY = :host_committed_primary
+        REDUCTION_POLICY = :ordered_all_success
+        RETRY_POLICY = :none
+        ENUM_ATTRIBUTES = %i[kind resume_policy failure_policy reduction_policy retry_policy].freeze
+        private_constant :ENUM_ATTRIBUTES
+        OwnedString = Types::String.constructor { |value| value.is_a?(String) ? value.dup.freeze : value }
+        private_constant :OwnedString
+
+        attribute :version, Types::Integer.enum(VERSION)
+        attribute :execution_semantics_version, OwnedString.constrained(min_size: 1, max_size: 32)
+        attribute :dispatch, Types.Instance(PreparedStepDispatch)
+        attribute :kind, Types::Symbol.enum(:parallel, :fanout)
+        attribute :transition, OwnedString.constrained(min_size: 1, max_size: 256)
+        attribute :from, OwnedString.constrained(min_size: 1, max_size: 256)
+        attribute :execution_namespace, OwnedString.constrained(format: PreparedStep::UUID_PATTERN)
+        attribute :branches, Types::Array.of(Types.Instance(Branch))
+        attribute :input_digest, OwnedString.constrained(format: PreparedStep::DIGEST_PATTERN)
+        attribute :budget_state_digest, OwnedString.constrained(format: PreparedStep::DIGEST_PATTERN)
+        attribute :resume_policy, Types::Symbol.enum(RESUME_POLICY)
+        attribute :failure_policy, Types::Symbol.enum(FAILURE_POLICY)
+        attribute :reduction_policy, Types::Symbol.enum(REDUCTION_POLICY)
+        attribute :retry_policy, Types::Symbol.enum(RETRY_POLICY)
+        attribute :plan_digest, OwnedString.constrained(format: PreparedStep::DIGEST_PATTERN)
+
+        class << self
+          def build(**attributes)
+            values = {
+              version: VERSION,
+              execution_semantics_version: Smith::EXECUTION_SEMANTICS_VERSION,
+              dispatch: attributes.fetch(:dispatch),
+              kind: attributes.fetch(:kind),
+              transition: attributes.fetch(:transition).to_s,
+              from: attributes.fetch(:from).to_s,
+              execution_namespace: attributes.fetch(:execution_namespace),
+              branches: attributes.fetch(:branches),
+              input_digest: attributes.fetch(:input_digest),
+              budget_state_digest: attributes.fetch(:budget_state_digest),
+              resume_policy: RESUME_POLICY,
+              failure_policy: FAILURE_POLICY,
+              reduction_policy: REDUCTION_POLICY,
+              retry_policy: RETRY_POLICY
+            }
+            new(values.merge(plan_digest: PayloadDigest.call(serializable(values))))
+          end
+
+          def normalize_attributes(attributes)
+            normalized = super
+            normalize_dispatch!(normalized)
+            normalize_branches!(normalized)
+            normalize_policies!(normalized)
+            normalized
+          end
+
+          def preflight_attributes!(attributes)
+            branches = raw_attribute(attributes, :branches)
+            validate_branch_count!(branches) if branches.is_a?(Array)
+            attributes
+          end
+
+          def serializable(values)
+            values.merge(
+              dispatch: values.fetch(:dispatch).to_h,
+              branches: values.fetch(:branches).map(&:to_h)
+            )
+          end
+
+          def validate_branch_count!(branches_or_count)
+            count = if branches_or_count.is_a?(Array)
+                      Array.instance_method(:length).bind_call(branches_or_count)
+                    else
+                      branches_or_count
+                    end
+            return if count&.between?(1, MAX_BRANCHES)
+
+            raise ArgumentError, "composite plan branch count is outside the transport limit"
+          end
+
+          private
+
+          def normalize_dispatch!(attributes)
+            dispatch = attributes[:dispatch]
+            attributes[:dispatch] = PreparedStepDispatch.deserialize(dispatch) unless
+              dispatch.is_a?(PreparedStepDispatch)
+          end
+
+          def normalize_branches!(attributes)
+            branches = attributes[:branches]
+            validate_branch_count!(branches)
+            normalized = []
+            Array.instance_method(:each).bind_call(branches) do |branch|
+              normalized << (branch.is_a?(Branch) ? branch : Branch.deserialize(branch))
+            end
+            attributes[:branches] = normalized.freeze
+          end
+
+          def normalize_policies!(attributes)
+            ENUM_ATTRIBUTES.each { |key| attributes[key] = Enums.normalize(key, attributes[key]) }
+          end
+        end
+
+        def initialize(attributes)
+          self.class.preflight_attributes!(attributes)
+          owned = self.class.normalize_attributes(attributes)
+          super(owned)
+          PlanIntegrity.new(self).call
+        end
+
+        def execution_for(branch) = BranchExecution.build(plan: self, branch:)
+      end
+    end
+  end
+end

data/lib/smith/workflow/composite/plan_integrity.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require "dry/initializer"
+
+require_relative "payload_digest"
+
+module Smith
+  class Workflow
+    module Composite
+      class PlanIntegrity
+        extend Dry::Initializer
+
+        param :plan
+
+        def call
+          validate_execution_semantics!
+          validate_branches!
+          validate_digest!
+          plan
+        end
+
+        private
+
+        def validate_execution_semantics!
+          return if plan.execution_semantics_version == Smith::EXECUTION_SEMANTICS_VERSION
+
+          raise ArgumentError, "composite plan execution semantics do not match"
+        end
+
+        def validate_branches!
+          keys = {}
+          plan.branches.each_with_index do |branch, ordinal|
+            raise ArgumentError, "composite plan ordinals must be contiguous" unless branch.ordinal == ordinal
+            raise ArgumentError, "composite plan branch keys must be unique" if keys.key?(branch.key)
+
+            keys[branch.key] = true
+          end
+        end
+
+        def validate_digest!
+          attributes = plan.to_h.except(:plan_digest).merge(
+            dispatch: plan.dispatch.to_h,
+            branches: plan.branches.map(&:to_h)
+          )
+          return if plan.plan_digest == PayloadDigest.call(attributes)
+
+          raise ArgumentError, "composite plan digest does not match"
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/composite/planner.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+require_relative "branch"
+require_relative "plan"
+
+module Smith
+  class Workflow
+    module Composite
+      class Planner
+        extend Dry::Initializer
+
+        option :dispatch
+        option :kind
+        option :transition
+        option :from
+        option :execution_namespace
+        option :branch_specs
+        option :input_digest
+        option :budget_state_digest
+
+        def call
+          Plan.validate_branch_count!(branch_specs)
+          Plan.build(
+            dispatch:,
+            kind:,
+            transition:,
+            from:,
+            execution_namespace:,
+            branches: build_branches,
+            input_digest:,
+            budget_state_digest:
+          )
+        end
+
+        private
+
+        def build_branches
+          branch_specs.each_with_index.map do |spec, ordinal|
+            Branch.build(
+              ordinal:,
+              key: spec.fetch(:key),
+              agent: spec.fetch(:agent),
+              binding_identity: spec.fetch(:binding_identity),
+              budget: spec.fetch(:budget, {})
+            )
+          end.freeze
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/composite/preparation.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require_relative "../../types"
+require_relative "input"
+require_relative "plan"
+
+module Smith
+  class Workflow
+    module Composite
+      class Preparation < Dry::Struct
+        attribute :plan, Types.Instance(Plan)
+        attribute :input, Types.Instance(Input)
+
+        def initialize(attributes)
+          super
+          unless plan.input_digest == input.digest
+            raise ArgumentError,
+                  "composite preparation input does not match plan"
+          end
+
+          self.attributes.freeze
+          freeze
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/composite/reducer.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+require_relative "branch_outcome"
+require_relative "budget_math"
+require_relative "effects"
+require_relative "outcome_accumulator"
+require_relative "plan"
+require_relative "reduction"
+
+module Smith
+  class Workflow
+    module Composite
+      class Reducer
+        extend Dry::Initializer
+
+        option :plan
+        option :outcomes
+        option :primary_failure, optional: true
+
+        def call
+          outcome_set = OutcomeAccumulator.new(plan:, outcomes:).call
+          ordered = outcome_set.ordered
+          validate_effects!(ordered)
+          effects = merged_effects(ordered)
+          failure_seen, selected_failure = failure_state(ordered)
+          return successful_reduction(outcome_set.output, effects) unless failure_seen
+
+          failed_reduction(selected_failure, effects)
+        end
+
+        private
+
+        def validate_effects!(ordered)
+          usage_ids = {}
+          ordered.each { validate_outcome_effects!(_1, usage_ids) }
+        end
+
+        def validate_outcome_effects!(outcome, usage_ids)
+          branch = plan.branches.fetch(outcome.ordinal)
+          validate_budget!(outcome.effects, branch.budget)
+          validate_usage!(outcome.effects.usage_entries, branch, usage_ids)
+        end
+
+        def validate_usage!(entries, branch, usage_ids)
+          entries.each { validate_usage_entry!(_1, branch, usage_ids) }
+        end
+
+        def validate_usage_entry!(entry, branch, usage_ids)
+          unless entry.fetch("agent_name") == branch.agent
+            raise ArgumentError, "composite usage entry does not match its branch agent"
+          end
+
+          usage_id = entry.fetch("usage_id")
+          raise ArgumentError, "composite usage entry is duplicated" if usage_ids.key?(usage_id)
+
+          usage_ids[usage_id] = true
+        end
+
+        def validate_budget!(effects, envelope)
+          consumed = effects.budget_consumed
+          validate_budget_dimensions!(consumed, envelope)
+          envelope.each { |dimension, limit| validate_budget_dimension!(dimension, limit, consumed, effects) }
+        end
+
+        def validate_budget_dimensions!(consumed, envelope)
+          return if (consumed.keys - envelope.keys).empty?
+
+          raise ArgumentError, "composite branch consumed an undeclared budget dimension"
+        end
+
+        def validate_budget_dimension!(dimension, limit, consumed, effects)
+          amount = consumed.fetch(dimension, 0)
+          unless amount.is_a?(Numeric) && amount.finite? && amount >= 0 && amount <= limit
+            raise ArgumentError, "composite branch budget consumption exceeds its envelope"
+          end
+
+          expected = expected_consumption(dimension, effects)
+          return if decimal_equal?(amount, expected)
+
+          raise ArgumentError, "composite branch budget consumption does not match recorded usage"
+        end
+
+        def expected_consumption(dimension, effects)
+          return effects.total_tokens if %w[total_tokens token_limit].include?(dimension)
+          return effects.total_cost if dimension == "total_cost"
+
+          0
+        end
+
+        def decimal_equal?(left, right)
+          Budget::DecimalContext.call { BigDecimal(left.to_s) == BigDecimal(right.to_s) }
+        end
+
+        def merged_effects(ordered)
+          Effects.new(
+            usage_entries: ordered.flat_map { _1.effects.usage_entries },
+            tool_results: ordered.flat_map { _1.effects.tool_results },
+            budget_consumed: BudgetMath.sum(ordered.map { _1.effects.budget_consumed })
+          )
+        end
+
+        def failure_state(ordered)
+          expected_key = primary_failure&.to_s
+          failure_seen = false
+          selected_failure = nil
+          ordered.each do |outcome|
+            next unless outcome.failed?
+
+            failure_seen = true
+            selected_failure = outcome if outcome.branch_key == expected_key
+          end
+          [failure_seen, selected_failure]
+        end
+
+        def successful_reduction(output, effects)
+          raise ArgumentError, "primary failure must be absent for a successful composite" if primary_failure
+
+          Reduction.new(status: :succeeded, output:, error: nil,
+                        failed_branch_key: nil, effects:)
+        end
+
+        def failed_reduction(failure, effects)
+          raise ArgumentError, "primary failure must identify a failed branch" unless failure
+
+          Reduction.new(status: :failed, output: nil, error: failure.error,
+                        failed_branch_key: failure.branch_key, effects:)
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/composite/reduction.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require_relative "../../types"
+require_relative "effects"
+require_relative "error"
+
+module Smith
+  class Workflow
+    module Composite
+      class Reduction < Dry::Struct
+        attribute :status, Types::Symbol.enum(:succeeded, :failed)
+        attribute? :output, Types::Any.optional
+        attribute? :error, Types.Instance(Error).optional
+        attribute? :failed_branch_key, Types::String.optional
+        attribute :effects, Types.Instance(Effects)
+
+        def initialize(attributes)
+          super
+          valid = status == :succeeded ? error.nil? && failed_branch_key.nil? : error && failed_branch_key
+          raise ArgumentError, "composite reduction fields do not match status" unless valid
+
+          self.attributes.freeze
+          freeze
+        end
+
+        def succeeded? = status == :succeeded
+        def failed? = status == :failed
+      end
+    end
+  end
+end

data/lib/smith/workflow/composite/value_budget.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+require_relative "../../errors"
+require_relative "../string_snapshot"
+
+module Smith
+  class Workflow
+    module Composite
+      class ValueBudget
+        extend Dry::Initializer
+
+        ARRAY_EACH = Array.instance_method(:each)
+        HASH_EACH_PAIR = Hash.instance_method(:each_pair)
+        private_constant :ARRAY_EACH, :HASH_EACH_PAIR
+
+        option :max_bytes
+        option :max_nodes
+        option :max_depth
+        option :label
+
+        def initialize(...)
+          super
+          @bytes = 0
+          @nodes = 0
+        end
+
+        def add(value, depth: 0)
+          pending = [[value, depth]]
+          until pending.empty?
+            item, item_depth = pending.pop
+            visit!(item, item_depth, pending)
+          end
+          self
+        end
+
+        private
+
+        def visit!(item, depth, pending)
+          validate_visit!(depth)
+          dispatch_value(item, depth, pending)
+        end
+
+        def dispatch_value(item, depth, pending)
+          case item
+          when Hash then enqueue_hash(item, depth, pending)
+          when Array then enqueue_array(item, depth, pending)
+          when String, Symbol then add_bytes!(item.to_s)
+          when Float then validate_float!(item)
+          when Integer, true, false, nil then nil
+          else raise WorkflowError, "#{label} contains unsupported value #{item.class}"
+          end
+        end
+
+        def validate_visit!(depth)
+          raise WorkflowError, "#{label} exceeds maximum depth #{max_depth}" if depth > max_depth
+
+          @nodes += 1
+          raise WorkflowError, "#{label} exceeds maximum size #{max_nodes}" if @nodes > max_nodes
+        end
+
+        def enqueue_hash(hash, depth, pending)
+          HASH_EACH_PAIR.bind_call(hash) do |key, value|
+            unless key.is_a?(String) || key.is_a?(Symbol)
+              raise WorkflowError, "#{label} contains unsupported Hash key #{key.class}"
+            end
+
+            pending << [value, depth + 1] << [key, depth + 1]
+          end
+        end
+
+        def enqueue_array(array, depth, pending)
+          ARRAY_EACH.bind_call(array) { |value| pending << [value, depth + 1] }
+        end
+
+        def add_bytes!(string)
+          @bytes += StringSnapshot.bytesize(string)
+          raise WorkflowError, "#{label} exceeds maximum bytes #{max_bytes}" if @bytes > max_bytes
+        end
+
+        def validate_float!(float)
+          return if float.finite?
+
+          raise WorkflowError, "#{label} contains a non-finite Float"
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/composite_branch_execution_authorization.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require_relative "prepared_step_execution_authorization"
+require_relative "composite/branch_execution"
+require_relative "composite/input"
+
+module Smith
+  class Workflow
+    class CompositeBranchExecutionAuthorization < PreparedStepExecutionAuthorization
+      attr_reader :execution_digest, :plan_digest, :input_digest, :branch_digest,
+                  :branch_ordinal, :execution_namespace
+
+      def initialize(execution:, input:, **attributes)
+        validate_payloads!(execution, input)
+        @execution_digest = execution.digest.dup.freeze
+        @plan_digest = execution.plan_digest.dup.freeze
+        @input_digest = input.digest.dup.freeze
+        assign_branch_identity(execution.branch)
+        @execution_namespace = execution.execution_namespace.dup.freeze
+        super(**attributes)
+      end
+
+      private
+
+      def validate_payloads!(execution, input)
+        return if execution.is_a?(Composite::BranchExecution) && input.is_a?(Composite::Input)
+
+        raise ArgumentError, "composite branch authorization requires typed execution and input"
+      end
+
+      def assign_branch_identity(branch)
+        @branch_digest = branch.digest.dup.freeze
+        @branch_ordinal = branch.ordinal
+      end
+
+      def verify_composite_branch!(execution, input)
+        validate_payloads!(execution, input)
+        expected = [
+          execution_digest, plan_digest, branch_digest, branch_ordinal,
+          execution_namespace, input_digest
+        ]
+        actual = [
+          execution.digest, execution.plan_digest, execution.branch.digest,
+          execution.branch.ordinal, execution.execution_namespace, input.digest
+        ]
+        return if actual == expected
+
+        raise WorkflowError, "composite branch authorization does not match execution"
+      end
+    end
+  end
+end

data/lib/smith/workflow/deadline_enforcement.rb

@@ -2,6 +2,8 @@
 
 require "time"
 
+require_relative "thread_context_snapshot"
+
 module Smith
   class Workflow
     module DeadlineEnforcement
@@ -19,17 +21,32 @@ module Smith
         [wall_clock_deadline, call_dl].compact.min
       end
 
-      def with_agent_context(agent_class)
+      def with_agent_context(agent_class, &block)
         saved_deadline = Tool.current_deadline
         saved_call_ledger = Thread.current[:smith_call_ledger]
+        snapshot = ThreadContextSnapshot.new(
+          tool_attributes: %i[current_deadline current_tool_call_allowance],
+          thread_keys: %i[smith_call_deadline smith_call_ledger],
+          scoped_artifacts: false
+        )
+        snapshot.around do
+          apply_agent_context(agent_class)
+          Thread.handle_interrupt(Object => :immediate, &block)
+        ensure
+          restore_agent_context(saved_deadline, saved_call_ledger)
+        end
+      end
+
+      def apply_agent_context(agent_class)
         apply_agent_deadline(agent_class)
         narrow_tool_deadline!
         apply_agent_tool_calls(agent_class)
         apply_agent_call_ledger(agent_class)
-        yield
-      ensure
-        Tool.current_deadline = saved_deadline
-        Thread.current[:smith_call_ledger] = saved_call_ledger
+      end
+
+      def restore_agent_context(deadline, call_ledger)
+        Tool.current_deadline = deadline
+        Thread.current[:smith_call_ledger] = call_ledger
         clear_agent_deadline
         clear_agent_tool_calls
       end

data/lib/smith/workflow/execution.rb

@@ -2,13 +2,17 @@
 
 require_relative "agent_result"
 require_relative "execution_binding_resolution"
+require_relative "prepared_branch_execution"
 require_relative "step_completion"
+require_relative "step_context"
 
 module Smith
   class Workflow
     module Execution
       include ExecutionBindingResolution
+      include PreparedBranchExecution
       include StepCompletion
+      include StepContext
       include Agent::Lifecycle
       include NestedExecution
       include EvaluatorOptimizer
@@ -21,14 +25,12 @@ module Smith
       private
 
       def execute_step(transition)
-        setup_step_context
+        with_step_context(transition) { execute_step_body(transition) }
+      end
+
+      def execute_step_body(transition)
         output = with_scoped_artifacts { run_with_retry_policy(transition) }
         StepCompletion.instance_method(:complete_step).bind_call(self, transition, output)
-      rescue StandardError => e
-        @outcome = nil
-        GuardrailIntegration.instance_method(:handle_step_failure).bind_call(self, transition, e)
-      ensure
-        teardown_step_context
       end
 
       def setup_step_context
@@ -37,16 +39,16 @@ module Smith
         Tool.current_tool_result_collector = tool_result_collector
       end
 
-      def teardown_step_context
-        Tool.current_guardrails = nil
-        Tool.current_deadline = nil
-        Tool.current_ledger = nil
-        Tool.current_tool_result_collector = nil
-        Smith.scoped_artifacts = nil
-      end
-
       def run_guarded_step(transition)
+        return apply_composite_reduction!(transition) if @composite_reduction
+
         @resolved_parallel_branch_count = preflight_branch_count(transition)
+        run_standard_guarded_step(transition)
+      ensure
+        @resolved_parallel_branch_count = nil
+      end
+
+      def run_standard_guarded_step(transition)
         return dispatch_step(transition) if transition.deterministic?
         return run_guarded_fanout_step(transition) if transition.fanout?
 
@@ -62,8 +64,6 @@ module Smith
         validate_data_volume!(output, agent_class)
         run_output_guardrails(output, agent_class)
         resolve_router_output(transition, output)
-      ensure
-        @resolved_parallel_branch_count = nil
       end
 
       def preflight_branch_count(transition)