smith-agents 0.5.0 → 0.6.1

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)

data/lib/smith/workflow/fanout_execution.rb

@@ -35,29 +35,32 @@ module Smith
         )
 
         branch_calls = branches.map do |branch_key, agent_name|
-          proc do |signal|
-            run_fanout_branch(branch_key, agent_name, branch_agent_classes.fetch(branch_key), env, signal)
-          end
+          PreparedBranchExecution.instance_method(:prepared_branch).bind_call(
+            self,
+            FanoutExecution.instance_method(:run_fanout_branch),
+            branch_key,
+            agent_name,
+            branch_agent_classes.fetch(branch_key),
+            env
+          )
         end
 
         Parallel.execute(branches: branch_calls)
       end
 
       def run_fanout_branch(branch_key, agent_name, agent_class, env, signal)
-        setup_fanout_branch_context(env, @ledger, agent_class)
-
-        with_agent_context(agent_class) do
-          branch_ledger = effective_call_ledger
-          reserved = reserve_fanout_branch_call(branch_ledger, env.branch_estimates[branch_key], agent_class)
-          begin
-            result = guarded_fanout_branch_call(agent_class, env, signal)
-            finalize_named_branch(branch_key, agent_name, result, branch_ledger, reserved).tap { reserved = nil }
-          ensure
-            settle_budget_on_failure(branch_ledger, reserved, Thread.current[:smith_last_agent_result]) if reserved
+        with_branch_context(env, @ledger, agent_class:) do
+          with_agent_context(agent_class) do
+            branch_ledger = effective_call_ledger
+            reserved = reserve_fanout_branch_call(branch_ledger, env.branch_estimates[branch_key], agent_class)
+            begin
+              result = guarded_fanout_branch_call(agent_class, env, signal)
+              finalize_named_branch(branch_key, agent_name, result, branch_ledger, reserved).tap { reserved = nil }
+            ensure
+              settle_budget_on_failure(branch_ledger, reserved, Thread.current[:smith_last_agent_result]) if reserved
+            end
           end
         end
-      ensure
-        teardown_branch_context(env)
       end
 
       def guarded_fanout_branch_call(agent_class, env, signal)
@@ -89,9 +92,7 @@ module Smith
       end
 
       def run_fanout_agent_input_guardrails(branch_agent_classes)
-        branch_agent_classes.each_value do |agent_class|
-          run_agent_input_guardrails(agent_class)
-        end
+        branch_agent_classes.each_value { |agent_class| run_agent_input_guardrails(agent_class) }
       end
 
       def fanout_branch_estimates(branches, branch_agent_classes)

data/lib/smith/workflow/message_value_normalizer.rb

@@ -19,6 +19,7 @@ module Smith
       extend Dry::Initializer
 
       param :value
+      option :label, default: proc { "session message" }
 
       def call
         @active = {}.compare_by_identity
@@ -50,12 +51,12 @@ module Smith
         when Float then copy_float(item)
         when true, false, nil then item
         else
-          reject!("session message contains unsupported value #{item.class}")
+          reject!("#{label} contains unsupported value #{item.class}")
         end
       end
 
       def copy_container(source)
-        reject!("session message contains a cyclic value") if @active.key?(source)
+        reject!("#{label} contains a cyclic value") if @active.key?(source)
 
         @active[source] = true
         yield
@@ -77,8 +78,8 @@ module Smith
         pairs = []
         remaining_nodes = MAX_NODES - @nodes
         HASH_EACH_PAIR.bind_call(source) do |key, nested|
-          reject!("session message batch exceeds maximum size #{MAX_NODES}") if pairs.length == remaining_nodes
-          reject!("session message Hash keys must be strings or symbols") unless key.is_a?(String) || key.is_a?(Symbol)
+          reject!("#{label} exceeds maximum size #{MAX_NODES}") if pairs.length == remaining_nodes
+          reject!("#{label} Hash keys must be strings or symbols") unless key.is_a?(String) || key.is_a?(Symbol)
 
           pairs << [copy_string(key.to_s), nested]
         end
@@ -88,7 +89,7 @@ module Smith
       def validate_unique_keys!(pairs)
         return unless pairs.each_cons(2).any? { |left, right| left.first == right.first }
 
-        reject!("session message contains duplicate canonical Hash keys")
+        reject!("#{label} contains duplicate canonical Hash keys")
       end
 
       def copy_array(source, depth:)
@@ -99,7 +100,7 @@ module Smith
 
       def copy_string(string)
         @bytes += StringSnapshot.bytesize(string)
-        reject!("session message batch exceeds maximum bytes #{MAX_BYTES}") if @bytes > MAX_BYTES
+        reject!("#{label} exceeds maximum bytes #{MAX_BYTES}") if @bytes > MAX_BYTES
 
         StringSnapshot.copy(string, freeze: true)
       end
@@ -107,22 +108,22 @@ module Smith
       def copy_integer(integer)
         return integer if INTEGER_RANGE.cover?(integer)
 
-        reject!("session message integers must fit a signed 64-bit value")
+        reject!("#{label} integers must fit a signed 64-bit value")
       end
 
       def copy_float(float)
         return float if float.finite?
 
-        reject!("session message contains a non-finite Float")
+        reject!("#{label} contains a non-finite Float")
       end
 
       def visit!(depth)
-        reject!("session message batch exceeds maximum depth #{MAX_DEPTH}") if depth > MAX_DEPTH
+        reject!("#{label} exceeds maximum depth #{MAX_DEPTH}") if depth > MAX_DEPTH
 
         @nodes += 1
         return if @nodes <= MAX_NODES
 
-        reject!("session message batch exceeds maximum size #{MAX_NODES}")
+        reject!("#{label} exceeds maximum size #{MAX_NODES}")
       end
 
       def reject!(message)

data/lib/smith/workflow/parallel_execution.rb

@@ -2,10 +2,14 @@
 
 require_relative "branch_env"
 require_relative "parallel_agent_binding"
+require_relative "thread_context_snapshot"
 
 module Smith
   class Workflow
     module ParallelExecution
+      NO_PARALLEL_BINDING = Object.new.freeze
+      private_constant :NO_PARALLEL_BINDING
+
       private
 
       def execute_parallel_step(transition, prepared_input: nil)
@@ -22,26 +26,27 @@ module Smith
         )
         ledger = @ledger
         branches = Array.new(count) do |i|
-          proc { |signal| run_branch(transition, i, env, ledger, signal) }
+          PreparedBranchExecution
+            .instance_method(:prepared_branch)
+            .bind_call(self, ParallelExecution.instance_method(:run_branch), transition, i, env, ledger)
         end
         Parallel.execute(branches: branches)
       end
 
       def run_branch(transition, index, env, ledger, signal)
-        setup_branch_context(env, ledger)
-        Thread.current[:smith_parallel_agent_binding] = ParallelAgentBinding.new(self, transition, env.agent_class)
-        with_agent_context(env.agent_class) do
-          branch_ledger = effective_call_ledger
-          reserved = reserve_branch_call(branch_ledger, env, ledger)
-          begin
-            result = guarded_branch_call(transition, env, signal)
-            finalize_branch(transition, index, result, branch_ledger, reserved).tap { reserved = nil }
-          ensure
-            settle_budget_on_failure(branch_ledger, reserved, Thread.current[:smith_last_agent_result]) if reserved
+        binding = ParallelAgentBinding.new(self, transition, env.agent_class)
+        with_branch_context(env, ledger, parallel_agent_binding: binding) do
+          with_agent_context(env.agent_class) do
+            branch_ledger = effective_call_ledger
+            reserved = reserve_branch_call(branch_ledger, env, ledger)
+            begin
+              result = guarded_branch_call(transition, env, signal)
+              finalize_branch(transition, index, result, branch_ledger, reserved).tap { reserved = nil }
+            ensure
+              settle_budget_on_failure(branch_ledger, reserved, Thread.current[:smith_last_agent_result]) if reserved
+            end
           end
         end
-      ensure
-        teardown_branch_context(env)
       end
 
       def reserve_branch_call(branch_ledger, env, workflow_ledger)
@@ -50,6 +55,29 @@ module Smith
         reserve_serial_budget(branch_ledger) if branch_ledger
       end
 
+      def with_branch_context(
+        env,
+        ledger,
+        parallel_agent_binding: NO_PARALLEL_BINDING,
+        agent_class: nil,
+        &block
+      )
+        snapshot = ThreadContextSnapshot.new
+        snapshot.around do
+          if agent_class
+            setup_fanout_branch_context(env, ledger, agent_class)
+          else
+            setup_branch_context(env, ledger)
+          end
+          unless parallel_agent_binding.equal?(NO_PARALLEL_BINDING)
+            Thread.current[:smith_parallel_agent_binding] = parallel_agent_binding
+          end
+          Thread.handle_interrupt(Object => :immediate, &block)
+        ensure
+          teardown_branch_context(env)
+        end
+      end
+
       def setup_branch_context(env, ledger)
         env.setup_thread
         Tool.current_ledger = ledger

data/lib/smith/workflow/prepared_branch_execution.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Smith
+  class Workflow
+    module PreparedBranchExecution
+      private
+
+      def prepared_branch(implementation, *arguments)
+        unless @split_step_active_execution_authorization
+          return proc { |signal| __send__(implementation.name, *arguments, signal) }
+        end
+
+        proc do |signal|
+          run = proc { implementation.bind_call(self, *arguments, signal) }
+          PreparedBranchExecution.instance_method(:within_prepared_branch_execution).bind_call(self, &run)
+        end
+      end
+
+      def within_prepared_branch_execution(&)
+        authorization = @split_step_active_execution_authorization
+        return yield unless authorization
+
+        PreparedStepExecutionAuthorization
+          .instance_method(:within_branch_execution!)
+          .bind_call(authorization, &)
+      end
+    end
+
+    private_constant :PreparedBranchExecution
+  end
+end

data/lib/smith/workflow/prepared_step_execution_authorization.rb

@@ -3,12 +3,15 @@
 require_relative "prepared_step"
 require_relative "prepared_step_dispatch"
 require_relative "prepared_step_execution_scope"
+require_relative "process_local"
 require_relative "split_step_persistence/execution_binding_snapshot"
 require_relative "../errors"
 
 module Smith
   class Workflow
     class PreparedStepExecutionAuthorization
+      include ProcessLocal
+
       attr_reader :prepared_step, :dispatch_claim
 
       def initialize(prepared_step:, dispatch_claim:, execution_bindings:)
@@ -55,42 +58,23 @@ module Smith
       end
 
       def active_in_current_execution?
-        issued_in_current_process? && @execution_scope.active_for?(Thread.current)
-      end
-
-      def initialize_copy(_source)
-        raise TypeError, "prepared-step execution authorizations cannot be copied"
-      end
-
-      def _dump(_depth)
-        raise TypeError, "prepared-step execution authorizations cannot be serialized"
-      end
-
-      def encode_with(_coder)
-        raise TypeError, "prepared-step execution authorizations cannot be serialized"
+        issued_in_current_process? && @execution_scope.active_for?(Thread.current, Fiber.current)
       end
 
-      def init_with(_coder)
-        raise TypeError, "prepared-step execution authorizations cannot be deserialized"
-      end
-
-      def as_json(*)
-        raise TypeError, "prepared-step execution authorizations cannot be serialized"
-      end
+      private
 
-      def to_json(*)
-        raise TypeError, "prepared-step execution authorizations cannot be serialized"
+      def within_branch_execution!(&)
+        ensure_current_process!
+        @execution_scope.within_branch(&)
       end
 
-      private
-
-      def activate_execution!(thread)
+      def activate_execution!
         ensure_current_process!
-        @execution_scope.activate!(thread)
+        @execution_scope.activate!(Thread.current, Fiber.current)
       end
 
-      def close_execution!(thread = nil)
-        @execution_scope.close!(thread)
+      def close_execution!
+        @execution_scope.close!(Thread.current, Fiber.current)
       end
 
       def ensure_active_execution!
@@ -100,7 +84,8 @@ module Smith
       end
 
       def ensure_binding_access!
-        accessible = issued_in_current_process? && @execution_scope.binding_accessible_for?(Thread.current)
+        accessible = issued_in_current_process? &&
+                     @execution_scope.binding_accessible_for?(Thread.current, Fiber.current)
         return if accessible
 
         raise WorkflowError, "prepared-step execution authorization is outside its binding access scope"

data/lib/smith/workflow/prepared_step_execution_scope.rb

@@ -1,45 +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)
-        @mutex.synchronize do
-          raise WorkflowError, "prepared-step execution scope is no longer available" unless @phase == :issued
+      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
+            @phase = :active
+            @thread = thread
+            @fiber = fiber
+          end
         end
       end
 
-      def close!(thread = nil)
-        @mutex.synchronize do
-          if @phase == :active && thread && !@thread.equal?(thread)
-            raise WorkflowError, "prepared-step execution scope belongs to another thread"
+      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
 
-          @phase = :closed
-          @thread = nil
+      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)
-        @mutex.synchronize { @phase == :active && @thread.equal?(thread) }
+      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)
+      def binding_accessible_for?(thread, fiber = Fiber.current)
         @mutex.synchronize do
-          @phase == :issued || (@phase == :active && @thread.equal?(thread))
+          @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/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