smith-agents 0.4.5 → 0.6.0

data/lib/smith/workflow/prepared_step_execution_scope.rb

@@ -0,0 +1,111 @@
+# frozen_string_literal: true
+
+require_relative "../errors"
+require_relative "process_local"
+
+module Smith
+  class Workflow
+    class PreparedStepExecutionScope
+      include ProcessLocal
+
+      def initialize
+        @mutex = Mutex.new
+        @phase = :issued
+        @thread = nil
+        @fiber = nil
+        @branch_fibers = {}.compare_by_identity
+      end
+
+      def activate!(thread, fiber = Fiber.current)
+        Thread.handle_interrupt(Object => :never) do
+          @mutex.synchronize do
+            raise WorkflowError, "prepared-step execution scope is no longer available" unless @phase == :issued
+
+            @phase = :active
+            @thread = thread
+            @fiber = fiber
+          end
+        end
+      end
+
+      def close!(thread = nil, fiber = Fiber.current)
+        Thread.handle_interrupt(Object => :never) do
+          @mutex.synchronize do
+            if @phase == :active && thread && !owner?(@thread, @fiber, thread, fiber)
+              raise WorkflowError, "prepared-step execution scope belongs to another thread or fiber"
+            end
+            raise WorkflowError, "prepared-step execution still has active branch fibers" unless @branch_fibers.empty?
+
+            @phase = :closed
+            @thread = nil
+            @fiber = nil
+          end
+        end
+      end
+
+      def within_branch(&block)
+        thread = Thread.current
+        fiber = Fiber.current
+        Thread.handle_interrupt(Object => :never) do
+          enter_branch!(thread, fiber)
+          begin
+            Thread.handle_interrupt(Object => :immediate, &block)
+          ensure
+            leave_branch!(thread, fiber)
+          end
+        end
+      end
+
+      def active_for?(thread, fiber = Fiber.current)
+        @mutex.synchronize do
+          @phase == :active && (owner?(@thread, @fiber, thread, fiber) || branch_owner?(thread, fiber))
+        end
+      end
+
+      def binding_accessible_for?(thread, fiber = Fiber.current)
+        @mutex.synchronize do
+          @phase == :issued ||
+            (@phase == :active && (owner?(@thread, @fiber, thread, fiber) || branch_owner?(thread, fiber)))
+        end
+      end
+
+      private
+
+      def enter_branch!(thread, fiber)
+        @mutex.synchronize do
+          raise WorkflowError, "prepared-step execution scope is not active" unless @phase == :active
+
+          entry = @branch_fibers[fiber]
+          if entry && !entry.fetch(:thread).equal?(thread)
+            raise WorkflowError, "prepared-step branch fiber belongs to another thread"
+          end
+
+          @branch_fibers[fiber] = { thread:, count: entry ? entry.fetch(:count) + 1 : 1 }
+        end
+      end
+
+      def leave_branch!(thread, fiber)
+        @mutex.synchronize do
+          entry = @branch_fibers.fetch(fiber) do
+            raise WorkflowError, "prepared-step branch execution fiber is not active"
+          end
+          unless entry.fetch(:thread).equal?(thread)
+            raise WorkflowError, "prepared-step branch fiber belongs to another thread"
+          end
+
+          count = entry.fetch(:count)
+          count == 1 ? @branch_fibers.delete(fiber) : entry[:count] = count - 1
+        end
+      end
+
+      def branch_owner?(thread, fiber)
+        entry = @branch_fibers[fiber]
+        !!(entry && entry.fetch(:thread).equal?(thread))
+      end
+
+      def owner?(owner_thread, owner_fiber, thread, fiber)
+        owner_thread.equal?(thread) && owner_fiber.equal?(fiber)
+      end
+    end
+  end
+end

data/lib/smith/workflow/prepared_step_recovery.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+require "dry-struct"
+
+require_relative "../types"
+require_relative "prepared_step"
+require_relative "prepared_step_dispatch"
+
+module Smith
+  class Workflow
+    class PreparedStepRecovery < Dry::Struct
+      attribute :prepared_step, Types.Instance(PreparedStep)
+      attribute? :dispatch_claim, Types.Instance(PreparedStepDispatch).optional
+      attribute :execution_status, Types::Symbol.enum(:not_started)
+
+      def self.not_started(witness)
+        case witness
+        when PreparedStep
+          new(prepared_step: witness, dispatch_claim: nil, execution_status: :not_started)
+        when PreparedStepDispatch
+          new(prepared_step: witness.prepared_step, dispatch_claim: witness, execution_status: :not_started)
+        else
+          raise ArgumentError, "recovery witness must be a PreparedStep or PreparedStepDispatch"
+        end
+      end
+
+      def initialize(attributes)
+        super
+        if dispatch_claim && dispatch_claim.prepared_step.to_h != prepared_step.to_h
+          raise ArgumentError, "dispatch claim must belong to the prepared step"
+        end
+
+        self.attributes.freeze
+        freeze
+      end
+    end
+  end
+end

data/lib/smith/workflow/process_local.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Smith
+  class Workflow
+    module ProcessLocal
+      def initialize_copy(_source)
+        raise TypeError, "#{self.class} cannot be copied"
+      end
+
+      def _dump(_depth)
+        raise TypeError, "#{self.class} cannot be serialized"
+      end
+
+      def encode_with(_coder)
+        raise TypeError, "#{self.class} cannot be serialized"
+      end
+
+      def init_with(_coder)
+        raise TypeError, "#{self.class} cannot be deserialized"
+      end
+
+      def as_json(*)
+        raise TypeError, "#{self.class} cannot be serialized"
+      end
+
+      def to_json(*)
+        raise TypeError, "#{self.class} cannot be serialized"
+      end
+    end
+
+    private_constant :ProcessLocal
+  end
+end

data/lib/smith/workflow/retry_execution.rb

@@ -9,6 +9,7 @@ module Smith
         config = transition.retry_config
         return run_guarded_step(transition) unless config
 
+        schedule = retry_schedule(config)
         attempt = 0
         begin
           attempt += 1
@@ -16,7 +17,7 @@ module Smith
         rescue StandardError => e
           raise unless retry_transition_error?(config, e, attempt)
 
-          sleep_for_retry(config, attempt)
+          sleep_for_retry(schedule, attempt)
           retry
         end
       end
@@ -32,20 +33,28 @@ module Smith
         end
       end
 
-      def sleep_for_retry(config, failed_attempt)
-        delay = retry_delay(config, failed_attempt)
+      def sleep_for_retry(schedule, failed_attempt)
+        delay = retry_delay(schedule, failed_attempt)
         sleep(delay) if delay.positive?
       end
 
-      def retry_delay(config, failed_attempt)
-        delay = config.fetch(:backoff) * (2**[failed_attempt - 1, 0].max)
-        max_delay = config[:max_delay]
-        delay = [delay, max_delay].min if max_delay
+      def retry_delay(config_or_schedule, failed_attempt)
+        schedule = if config_or_schedule.is_a?(ExponentialBackoff)
+                     config_or_schedule
+                   else
+                     retry_schedule(config_or_schedule)
+                   end
+        schedule.delay(failed_attempt, random: method(:rand))
+      end
 
-        jitter = config.fetch(:jitter)
-        delay += rand * jitter if jitter.positive?
-        delay = [delay, max_delay].min if max_delay
-        delay
+      def retry_schedule(config)
+        ExponentialBackoff.new(
+          attempts: config.fetch(:attempts),
+          base_delay: config.fetch(:backoff),
+          max_delay: config[:max_delay],
+          jitter: config.fetch(:jitter),
+          delay_label: "backoff"
+        )
       end
     end
   end

data/lib/smith/workflow/split_step_persistence/boundary.rb

@@ -21,19 +21,50 @@ module Smith
 
         def initialize_copy(source)
           super
-          @split_step_mutex = Mutex.new
-          return unless source.instance_variable_get(:@split_step_phase)
+          copied_split_step_phase = @split_step_phase
+          detach_copied_runtime_state!(source)
+          return unless copied_split_step_phase
 
           @split_step_phase = :copied_boundary
+          reset_copied_execution_state!
+          reset_copied_dispatch_state!
+        end
+        private :initialize_copy
+
+        private
+
+        def reset_copied_execution_state!
           @split_step_token = nil
+          @split_step_prepared_descriptor = nil
+          @split_step_transaction_identity = nil
           @split_step_execution_thread = nil
+          @split_step_active_execution_authorization = nil
+          @split_step_execution_result = nil
           @split_step_advance_permit = false
+          clear_split_step_execution_verification!
+          clear_split_step_execution_authorization!
+        end
+
+        def reset_copied_dispatch_state!
+          @split_step_dispatch_token = nil
+          @split_step_dispatch_descriptor = nil
+          @split_step_attempted_dispatch_descriptor = nil
+          @split_step_dispatch_thread = nil
+          @split_step_pre_dispatch_payload = nil
+          @split_step_attempted_dispatch_payload = nil
           @split_step_preparation_thread = nil
           @split_step_persist_permit = false
         end
-        private :initialize_copy
 
-        private
+        def detach_copied_runtime_state!(source)
+          source_messages = source.instance_variable_get(:@session_messages)
+          @split_step_mutex = Mutex.new
+          @persisted_keys_mutex = Mutex.new
+          @tool_results_mutex = Mutex.new
+          @usage_mutex = Mutex.new
+          detach_split_step_execution_state!
+          @session_messages = snapshot_value(source_messages || [])
+        end
 
         def ensure_split_step_execution_allowed!
           @split_step_mutex.synchronize do
@@ -91,24 +122,6 @@ module Smith
             @split_step_execution_thread.equal?(Thread.current) &&
             @split_step_advance_permit
         end
-
-        def clear_split_step_boundary!
-          @split_step_phase = nil
-          @split_step_transition_name = nil
-          @split_step_transition = nil
-          @split_step_transition_signature = nil
-          @split_step_origin_state = nil
-          @split_step_token = nil
-          @split_step_persistence_key = nil
-          @split_step_adapter = nil
-          remove_instance_variable(:@split_step_persistence_ttl) if
-            instance_variable_defined?(:@split_step_persistence_ttl)
-          @split_step_preparation_payload = nil
-          @split_step_checkpoint_digest = nil
-          @split_step_checkpoint_version = nil
-          @split_step_preparation_thread = nil
-          @split_step_persist_permit = false
-        end
       end
     end
   end

data/lib/smith/workflow/split_step_persistence/boundary_reset.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module Smith
+  class Workflow
+    module SplitStepPersistence
+      module BoundaryReset
+        private
+
+        def clear_split_step_boundary!
+          @split_step_phase = nil
+          @split_step_transition_name = nil
+          @split_step_transition = nil
+          @split_step_transition_signature = nil
+          @split_step_origin_state = nil
+          @split_step_token = nil
+          @split_step_active_execution_authorization = nil
+          @split_step_execution_result = nil
+          clear_split_step_descriptor_state!
+          @split_step_persistence_key = nil
+          @split_step_adapter = nil
+          remove_instance_variable(:@split_step_persistence_ttl) if
+            instance_variable_defined?(:@split_step_persistence_ttl)
+          @split_step_preparation_payload = nil
+          @split_step_checkpoint_digest = nil
+          @split_step_checkpoint_version = nil
+          @split_step_preparation_thread = nil
+          @split_step_persist_permit = false
+        end
+
+        def clear_split_step_descriptor_state!
+          @split_step_prepared_descriptor = nil
+          @split_step_transaction_identity = nil
+          @split_step_preparation_digest = nil
+          @split_step_pending_descriptor = nil
+          @split_step_dispatch_started = false
+          @split_step_dispatch_token = nil
+          @split_step_dispatch_descriptor = nil
+          @split_step_attempted_dispatch_descriptor = nil
+          @split_step_dispatch_thread = nil
+          @split_step_pre_dispatch_payload = nil
+          @split_step_attempted_dispatch_payload = nil
+          clear_split_step_execution_verification!
+          clear_split_step_execution_authorization!
+          remove_instance_variable(:@split_step_definition_digest) if
+            instance_variable_defined?(:@split_step_definition_digest)
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/split_step_persistence/canonical_payload_digest.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+require "digest"
+require "dry-initializer"
+require "json"
+
+module Smith
+  class Workflow
+    module SplitStepPersistence
+      class CanonicalPayloadDigest
+        extend Dry::Initializer
+
+        MAX_BYTES = 4 * 1024 * 1024
+        MAX_NODES = 100_000
+
+        param :payload
+
+        def self.call(payload) = new(payload).call
+
+        def call
+          if payload.bytesize > MAX_BYTES
+            raise WorkflowError, "split-step preparation payload exceeds maximum bytes #{MAX_BYTES}"
+          end
+
+          @nodes = 0
+          canonical = canonicalize(JSON.parse(payload))
+          Digest::SHA256.hexdigest(JSON.generate(canonical))
+        end
+
+        private
+
+        def canonicalize(value)
+          visit!
+          case value
+          when Hash then canonical_hash(value)
+          when Array then value.map { |item| canonicalize(item) }
+          when Float then value == value.to_i ? value.to_i : value
+          else value
+          end
+        end
+
+        def canonical_hash(value)
+          value.keys.sort.to_h { |key| [key, canonicalize(value.fetch(key))] }
+        end
+
+        def visit!
+          @nodes += 1
+          return if @nodes <= MAX_NODES
+
+          raise WorkflowError, "split-step preparation payload exceeds maximum size #{MAX_NODES}"
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/split_step_persistence/checkpoint.rb

@@ -91,11 +91,9 @@ module Smith
         end
 
         def capture_split_step_payload!(payload, next_version, checkpoint_claimed:)
-          validate_split_step_marker!(payload, expected: true) if @split_step_phase == :preparing || checkpoint_claimed
-          if @split_step_phase == :preparing
-            payload = split_step_preparation_payload(payload)
-            @split_step_preparation_payload = payload
-          end
+          return capture_split_step_preparation_payload!(payload, next_version) if @split_step_phase == :preparing
+
+          validate_split_step_marker!(payload, expected: true) if checkpoint_claimed
           return payload unless checkpoint_claimed
 
           checkpoint_payload = split_step_checkpoint_payload(payload)

data/lib/smith/workflow/split_step_persistence/composite_branch_authorization.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+require_relative "../composite/contract"
+require_relative "../composite_branch_execution_authorization"
+require_relative "execution_binding_snapshot"
+
+module Smith
+  class Workflow
+    module SplitStepPersistence
+      module CompositeBranchAuthorization
+        include Composite::Contract
+
+        private
+
+        def authorize_prepared_composite_branch_execution!(execution:, input:)
+          authorization = nil
+          handed_off = false
+          authorization = CompositeBranchAuthorization
+                          .instance_method(:interrupt_safe_composite_branch_authorization)
+                          .bind_call(self, execution, input)
+          handed_off = true
+          authorization
+        ensure
+          unless handed_off
+            ExecutionLifecycle
+              .instance_method(:release_failed_execution_authorization!)
+              .bind_call(self, authorization)
+          end
+        end
+
+        def with_composite_branch_authorization(execution, input)
+          Thread.handle_interrupt(Object => :never) do
+            authorization = interrupt_safe_composite_branch_authorization(execution, input)
+            begin
+              Thread.handle_interrupt(Object => :immediate) { yield(authorization) }
+            ensure
+              ExecutionLifecycle
+                .instance_method(:release_failed_execution_authorization!)
+                .bind_call(self, authorization)
+            end
+          end
+        end
+
+        def interrupt_safe_composite_branch_authorization(execution, input)
+          Thread.handle_interrupt(Object => :never) do
+            Composite::Contract
+              .instance_method(:validate_composite_branch_payload_types!)
+              .bind_call(self, execution, input)
+            verification_token = ExecutionAuthorization
+                                 .instance_method(:claim_framework_execution_verification!)
+                                 .bind_call(self)
+            begin
+              CompositeBranchAuthorization
+                .instance_method(:authorize_claimed_composite_branch_execution!)
+                .bind_call(self, verification_token, execution, input)
+            ensure
+              ExecutionAuthorization
+                .instance_method(:restore_framework_execution_verification!)
+                .bind_call(self, verification_token)
+            end
+          end
+        end
+
+        def authorize_claimed_composite_branch_execution!(verification_token, execution, input)
+          ExecutionVerification
+            .instance_method(:verify_claimed_split_step_execution!)
+            .bind_call(self, verification_token)
+          authorization = CompositeBranchAuthorization
+                          .instance_method(:build_composite_branch_execution_authorization)
+                          .bind_call(self, execution, input)
+          ExecutionAuthorization
+            .instance_method(:activate_split_step_execution_authorization!)
+            .bind_call(self, authorization, verification_token)
+          Composite::Contract
+            .instance_method(:validate_composite_branch_execution!)
+            .bind_call(self, authorization, execution, input)
+          authorization
+        end
+
+        def build_composite_branch_execution_authorization(execution, input)
+          role = execution.kind == :parallel ? :agent : :fanout_agent
+          bindings = ExecutionBindingSnapshot.capture_agent(
+            @split_step_transition,
+            workflow_class: self.class,
+            name: execution.branch.agent,
+            role:
+          )
+          CompositeBranchExecutionAuthorization.new(
+            execution:,
+            input:,
+            prepared_step: @split_step_prepared_descriptor,
+            dispatch_claim: @split_step_dispatch_descriptor,
+            execution_bindings: bindings
+          )
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/split_step_persistence/composite_branch_effects.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require_relative "../composite/effects"
+
+module Smith
+  class Workflow
+    module SplitStepPersistence
+      module CompositeBranchEffects
+        private
+
+        def composite_effect_offsets
+          usage = @usage_mutex.synchronize { @usage_entries.length }
+          tools = @tool_results_mutex.synchronize { @tool_results.length }
+          [usage, tools]
+        end
+
+        def capture_composite_branch_effects(usage_offset, tool_offset, branch)
+          usage = @usage_mutex.synchronize do
+            @usage_entries.drop(usage_offset).map { composite_usage_entry(_1, branch) }
+          end
+          tools = @tool_results_mutex.synchronize { @tool_results.drop(tool_offset) }
+          consumed = @ledger ? @ledger.consumed : {}
+          [Composite::Effects.new(usage_entries: usage, tool_results: tools, budget_consumed: consumed), nil]
+        rescue WorkflowError, ArgumentError => e
+          safe = Composite::Effects.new(usage_entries: usage || [], tool_results: [], budget_consumed: consumed || {})
+          [safe, e]
+        end
+
+        def composite_usage_entry(entry, branch)
+          attributes = entry.to_h
+          recorded_agent = attributes[:agent_name]
+          if recorded_agent && recorded_agent.to_s != branch.agent
+            raise WorkflowError, "composite usage entry does not match the executed branch agent"
+          end
+
+          attributes.merge(agent_name: branch.agent)
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/split_step_persistence/composite_branch_execution.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+require_relative "../composite/contract"
+require_relative "composite_branch_effects"
+require_relative "composite_branch_outcome"
+
+module Smith
+  class Workflow
+    module SplitStepPersistence
+      module CompositeBranchExecution
+        include Composite::Contract
+        include CompositeBranchEffects
+        include CompositeBranchOutcome
+
+        private
+
+        def execute_authorized_composite_branch!(authorization, execution:, input:)
+          authorization = validate_composite_authorization!(authorization)
+          perform_authorized_prepared_step_execution!(authorization) do
+            validate_composite_branch_execution!(authorization, execution, input)
+            [nil, execute_composite_branch(authorization, execution, input)]
+          end
+        end
+
+        def execute_composite_branch(authorization, execution, input)
+          branch = execution.branch
+          previous_ledger = @ledger
+          previous_execution_namespace = @execution_namespace
+          @ledger = composite_branch_ledger(branch)
+          @execution_namespace = execution.execution_namespace
+          offsets = composite_effect_offsets
+          output, error = capture_composite_branch do
+            dispatch_composite_branch(authorization, execution, input)
+          end
+          effects, effect_error = capture_composite_branch_effects(*offsets, branch)
+          composite_branch_outcome(execution, output, error || effect_error, effects)
+        ensure
+          @ledger = previous_ledger
+          @execution_namespace = previous_execution_namespace
+        end
+
+        def capture_composite_branch(&block)
+          output = within_raw_step_context { with_scoped_artifacts(&block) }
+          [output, nil]
+        rescue StandardError => e
+          [nil, e]
+        end
+
+        def dispatch_composite_branch(authorization, execution, input)
+          branch = execution.branch
+          transition = @split_step_transition
+          agent_role = execution.kind == :parallel ? :agent : :fanout_agent
+          agent_class = captured_agent(authorization, transition, branch.agent, agent_role)
+          apply_tool_guardrails(agent_class) if execution.kind == :parallel
+          environment = composite_branch_environment(execution, input, agent_class, transition)
+          execute_composite_branch_kind(execution, transition, agent_class, environment)
+        end
+
+        def execute_composite_branch_kind(execution, transition, agent_class, environment)
+          branch = execution.branch
+          signal = Parallel::CancellationSignal.new
+          return run_branch(transition, branch.ordinal, environment, @ledger, signal) if execution.kind == :parallel
+
+          branch_key, agent = fetch_composite_fanout_branch(transition, branch.key)
+          run_fanout_branch(branch_key, agent, agent_class, environment, signal)
+        end
+
+        def composite_branch_environment(execution, input, agent_class, transition)
+          branch = execution.branch
+          budget = branch.budget.transform_keys(&:to_sym)
+          branch_key = fetch_composite_fanout_branch(transition, branch.key).first unless execution.kind == :parallel
+          estimates = execution.kind == :parallel ? budget : { branch_key => budget }
+          BranchEnv.new(
+            prepared_input: input.agent_messages,
+            guardrail_sources: Tool.current_guardrails,
+            scoped_store: propagate_scoped_artifacts,
+            branch_estimates: estimates,
+            deadline: wall_clock_deadline,
+            agent_class: execution.kind == :parallel ? agent_class : nil
+          )
+        end
+
+        def composite_branch_ledger(branch)
+          return if branch.budget.empty?
+
+          Budget::Ledger.new(limits: branch.budget.transform_keys(&:to_sym))
+        end
+
+        def fetch_composite_fanout_branch(transition, branch_key)
+          Transition.instance_method(:fetch_fanout_branch!).bind_call(transition, branch_key)
+        end
+      end
+    end
+  end
+end