smith-agents 0.4.5 → 0.5.0

data/lib/smith/workflow/message_value_normalizer.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+require_relative "../errors"
+require_relative "string_snapshot"
+
+module Smith
+  class Workflow
+    class MessageValueNormalizer
+      MAX_DEPTH = 64
+      MAX_NODES = 100_000
+      MAX_BYTES = 1 * 1024 * 1024
+      INTEGER_RANGE = (-(2**63)..((2**63) - 1))
+      ARRAY_EACH = Array.instance_method(:each)
+      HASH_EACH_PAIR = Hash.instance_method(:each_pair)
+      private_constant :ARRAY_EACH, :HASH_EACH_PAIR
+
+      extend Dry::Initializer
+
+      param :value
+
+      def call
+        @active = {}.compare_by_identity
+        @nodes = 0
+        @bytes = 0
+        copy_value(value, depth: 0)
+      end
+
+      private
+
+      def copy_value(item, depth:)
+        visit!(depth)
+        return copy_container_value(item, depth:) if item.is_a?(Hash) || item.is_a?(Array)
+
+        copy_scalar(item)
+      end
+
+      def copy_container_value(item, depth:)
+        copy_container(item) do
+          item.is_a?(Hash) ? copy_hash(item, depth:) : copy_array(item, depth:)
+        end
+      end
+
+      def copy_scalar(item)
+        case item
+        when String then copy_string(item)
+        when Symbol then copy_string(item.to_s)
+        when Integer then copy_integer(item)
+        when Float then copy_float(item)
+        when true, false, nil then item
+        else
+          reject!("session message contains unsupported value #{item.class}")
+        end
+      end
+
+      def copy_container(source)
+        reject!("session message contains a cyclic value") if @active.key?(source)
+
+        @active[source] = true
+        yield
+      ensure
+        @active.delete(source)
+      end
+
+      def copy_hash(source, depth:)
+        pairs = canonical_pairs(source)
+        pairs.sort_by!(&:first)
+        validate_unique_keys!(pairs)
+
+        pairs.each_with_object({}) do |(key, nested), copy|
+          copy[key] = copy_value(nested, depth: depth + 1)
+        end.freeze
+      end
+
+      def canonical_pairs(source)
+        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)
+
+          pairs << [copy_string(key.to_s), nested]
+        end
+        pairs
+      end
+
+      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")
+      end
+
+      def copy_array(source, depth:)
+        copy = []
+        ARRAY_EACH.bind_call(source) { |nested| copy << copy_value(nested, depth: depth + 1) }
+        copy.freeze
+      end
+
+      def copy_string(string)
+        @bytes += StringSnapshot.bytesize(string)
+        reject!("session message batch exceeds maximum bytes #{MAX_BYTES}") if @bytes > MAX_BYTES
+
+        StringSnapshot.copy(string, freeze: true)
+      end
+
+      def copy_integer(integer)
+        return integer if INTEGER_RANGE.cover?(integer)
+
+        reject!("session message 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")
+      end
+
+      def visit!(depth)
+        reject!("session message batch 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}")
+      end
+
+      def reject!(message)
+        raise WorkflowError, message
+      end
+    end
+  end
+end

data/lib/smith/workflow/nested_execution.rb

@@ -21,10 +21,20 @@ module Smith
       end
 
       def build_child_workflow(child_class)
+        @split_step_active_execution_authorization&.verify_workflow!(child_class)
         child = child_class.new(context: @context.dup, ledger: @ledger, created_at: @created_at)
+        unless child.instance_of?(child_class)
+          raise WorkflowError, "nested workflow constructor returned #{child.class} instead of #{child_class}"
+        end
+
+        @split_step_active_execution_authorization&.verify_workflow!(child_class)
         child.instance_variable_set(:@execution_namespace, @execution_namespace)
         child.instance_variable_set(:@inherited_deadline, wall_clock_deadline)
         child.instance_variable_set(:@inherited_scoped_artifacts, Smith.scoped_artifacts)
+        child.instance_variable_set(
+          :@split_step_active_execution_authorization,
+          @split_step_active_execution_authorization
+        )
         child
       end
 

data/lib/smith/workflow/orchestrator_worker.rb

@@ -23,13 +23,13 @@ module Smith
 
       def execute_orchestration_step(transition, prepared_input: nil)
         state = OrchestrationState.new(transition.orchestrator_config, prepared_input)
-        state.orchestrator_class = Agent::Registry.fetch!(
+        state.orchestrator_class = resolve_registered_agent!(
           state.config[:orchestrator],
           workflow_class: self.class,
           transition_name: transition.name,
           role: :orchestrator
         )
-        state.worker_class = Agent::Registry.fetch!(
+        state.worker_class = resolve_registered_agent!(
           state.config[:worker],
           workflow_class: self.class,
           transition_name: transition.name,
@@ -77,11 +77,7 @@ module Smith
 
       def run_worker_with_schema(worker_class, task, worker_output_schema)
         input = [{ role: :user, content: task.to_json }]
-        original_schema = worker_class.output_schema
-        worker_class.output_schema(worker_output_schema)
-        invoke_agent_with_budget(worker_class, input)
-      ensure
-        worker_class.output_schema(original_schema)
+        invoke_agent_with_budget(worker_class, input, output_schema: worker_output_schema)
       end
 
       def prepare_orchestrator_input(prepared_input, round, worker_results)

data/lib/smith/workflow/parallel/cancellation_signal.rb

@@ -3,17 +3,26 @@
 module Smith
   class Workflow
     class Parallel
-      CancellationSignal = Struct.new(:cancelled, :mutex) do
+      class CancellationSignal
         def initialize
-          super(false, Mutex.new)
+          @cancelled = false
+          @reason = nil
+          @mutex = Mutex.new
         end
 
-        def cancel!
-          mutex.synchronize { self.cancelled = true }
+        def cancel!(error = nil)
+          @mutex.synchronize do
+            @reason ||= error
+            @cancelled = true
+          end
         end
 
         def cancelled?
-          mutex.synchronize { cancelled }
+          @mutex.synchronize { @cancelled }
+        end
+
+        def reason
+          @mutex.synchronize { @reason }
         end
       end
     end

data/lib/smith/workflow/parallel/execution_context.rb

@@ -0,0 +1,113 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  class Workflow
+    class Parallel
+      class ExecutionContext
+        extend Dry::Initializer
+
+        THREAD_KEY = :smith_parallel_execution_context
+        DEPTH_KEY = :smith_parallel_execution_depth
+
+        option :executor
+        option :signal
+        option :concurrency
+        option :nesting_limit
+        option :top_level_branch_count
+
+        def initialize(...)
+          super
+          @pending_top_level = top_level_branch_count
+          @active_top_level = 0
+          @reserved_nested_workers = 0
+          @mutex = Mutex.new
+        end
+
+        def self.current
+          Fiber[THREAD_KEY]
+        end
+
+        def self.current_depth
+          Fiber[DEPTH_KEY] || 0
+        end
+
+        def within(depth: self.class.current_depth)
+          previous = self.class.current
+          previous_depth = Fiber[DEPTH_KEY]
+          Fiber[THREAD_KEY] = self
+          Fiber[DEPTH_KEY] = depth
+          yield
+        ensure
+          Fiber[THREAD_KEY] = previous
+          Fiber[DEPTH_KEY] = previous_depth
+        end
+
+        def raise_if_cancelled!
+          raise(signal.reason || Cancellation.new("cancelled")) if signal.cancelled?
+        end
+
+        def next_nesting_depth!
+          depth = self.class.current_depth + 1
+          return depth if depth <= nesting_limit
+
+          raise WorkflowError, "parallel nesting exceeds configured limit #{nesting_limit}"
+        end
+
+        def top_level_started!
+          @mutex.synchronize do
+            @pending_top_level -= 1
+            @active_top_level += 1
+            validate_worker_count!
+          end
+        end
+
+        def top_level_finished!
+          @mutex.synchronize do
+            @active_top_level -= 1
+            validate_worker_count!
+          end
+        end
+
+        def reserve_workers(requested)
+          @mutex.synchronize do
+            reserved = [requested, available_workers].min
+            @reserved_nested_workers += reserved
+            validate_worker_count!
+            reserved
+          end
+        end
+
+        def release_workers(count)
+          @mutex.synchronize do
+            @reserved_nested_workers -= count
+            validate_worker_count!
+          end
+        end
+
+        private
+
+        def validate_worker_count!
+          valid = @pending_top_level >= 0 && @active_top_level >= 0 && @reserved_nested_workers >= 0
+          return if valid && occupied_workers <= concurrency
+
+          raise WorkflowError, "parallel worker accounting exceeded its concurrency bound"
+        end
+
+        def available_workers
+          concurrency - occupied_workers
+        end
+
+        def occupied_workers
+          @active_top_level + reserved_top_level_workers + @reserved_nested_workers
+        end
+
+        def reserved_top_level_workers
+          capacity = [concurrency - @active_top_level - @reserved_nested_workers, 0].max
+          [@pending_top_level, capacity].min
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/parallel/nested_execution.rb

@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+require "concurrent"
+require "dry-initializer"
+
+module Smith
+  class Workflow
+    class Parallel
+      class NestedExecution
+        extend Dry::Initializer
+
+        option :branches
+        option :context
+
+        def initialize(...)
+          super
+          @depth = context.next_nesting_depth!
+        end
+
+        def call
+          Thread.handle_interrupt(Object => :never) do
+            prepare
+            Thread.handle_interrupt(Object => :immediate) { execute }
+          rescue Exception => e # rubocop:disable Lint/RescueException
+            context.signal.cancel!(e)
+            raise
+          ensure
+            drain_futures
+            release_workers
+          end
+        end
+
+        private
+
+        def prepare
+          context.raise_if_cancelled!
+          @reserved_workers = context.reserve_workers(branches.length - 1)
+          return if @reserved_workers.zero?
+
+          @values = Array.new(branches.length)
+          @cursor = Concurrent::AtomicFixnum.new(0)
+          @futures = []
+          @reserved_workers.times { @futures << build_future }
+        end
+
+        def execute
+          return execute_inline if @reserved_workers.zero?
+
+          execute_concurrently
+        end
+
+        def execute_concurrently
+          fulfilled, reasons = execute_workers
+          raise_worker_failure!(fulfilled, reasons)
+          @values
+        end
+
+        def execute_workers
+          current_error = capture_error { work }
+          fulfilled, _values, reasons = Concurrent::Promises.zip(*@futures).result
+          [fulfilled, [context.signal.reason, current_error, *reasons]]
+        end
+
+        def raise_worker_failure!(fulfilled, reasons)
+          error = Parallel.preferred_error(reasons)
+          raise(error || Cancellation.new("parallel execution failed")) unless fulfilled && !error
+        end
+
+        def build_future
+          Concurrent::Promises.future_on(context.executor, context, @depth) do |execution_context, depth|
+            execution_context.within(depth:) { work }
+          end
+        end
+
+        def work
+          while (index = next_index)
+            @values[index] = execute_branch(index)
+          end
+        rescue StandardError => e
+          context.signal.cancel!(e)
+          raise
+        end
+
+        def next_index
+          context.raise_if_cancelled!
+          index = @cursor.increment - 1
+          return if index >= branches.length
+
+          context.raise_if_cancelled!
+          index
+        end
+
+        def execute_branch(index)
+          context.within(depth: @depth) { branches.fetch(index).call(context.signal) }
+        end
+
+        def execute_inline
+          branches.map do |branch|
+            context.raise_if_cancelled!
+            context.within(depth: @depth) { branch.call(context.signal) }
+          end
+        rescue StandardError => e
+          context.signal.cancel!(e)
+          raise
+        end
+
+        def capture_error
+          yield
+          nil
+        rescue StandardError => e
+          e
+        end
+
+        def drain_futures
+          Concurrent::Promises.zip(*@futures).result if @futures&.any?
+        end
+
+        def release_workers
+          context.release_workers(@reserved_workers) if @reserved_workers&.positive?
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/parallel/root_execution.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+require "concurrent"
+require "dry-initializer"
+
+module Smith
+  class Workflow
+    class Parallel
+      class RootExecution
+        extend Dry::Initializer
+
+        option :branches
+
+        def call
+          Thread.handle_interrupt(Object => :never) do
+            prepare
+            values = Thread.handle_interrupt(Object => :immediate) { resolve(@futures) }
+            @completed = true
+            values
+          rescue Exception => e # rubocop:disable Lint/RescueException
+            @signal&.cancel!(e)
+            raise
+          ensure
+            cleanup
+          end
+        end
+
+        private
+
+        def prepare
+          @completed = false
+          @signal = CancellationSignal.new
+          @concurrency = Smith.config.parallel_concurrency
+          @nesting_limit = Smith.config.parallel_nesting_limit
+          @executor = build_executor
+          @futures = build_futures(build_context)
+        end
+
+        def build_executor
+          Concurrent::FixedThreadPool.new(
+            @concurrency,
+            max_queue: [@concurrency, branches.length].max,
+            fallback_policy: :abort
+          )
+        end
+
+        def build_context
+          ExecutionContext.new(
+            executor: @executor,
+            signal: @signal,
+            concurrency: @concurrency,
+            nesting_limit: @nesting_limit,
+            top_level_branch_count: branches.length
+          )
+        end
+
+        def build_futures(context)
+          branches.map { future_for(_1, context) }
+        end
+
+        def future_for(branch, context)
+          Concurrent::Promises.future_on(context.executor, branch, context) do |callable, execution_context|
+            execution_context.top_level_started!
+            execution_context.raise_if_cancelled!
+            execution_context.within(depth: 0) { callable.call(execution_context.signal) }
+          rescue StandardError => e
+            execution_context.signal.cancel!(e)
+            raise
+          ensure
+            execution_context.top_level_finished!
+          end
+        end
+
+        def resolve(futures)
+          fulfilled, values, reasons = Concurrent::Promises.zip(*futures).result
+          raise(@signal.reason || Parallel.preferred_error(reasons)) unless fulfilled
+
+          values
+        end
+
+        def cleanup
+          @signal&.cancel!(Cancellation.new("parallel execution interrupted")) unless @completed
+          return unless @executor
+
+          @executor.shutdown
+          @executor.wait_for_termination
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/parallel.rb

@@ -1,8 +1,10 @@
 # frozen_string_literal: true
 
-require "concurrent"
-
+require_relative "parallel/cancellation"
 require_relative "parallel/cancellation_signal"
+require_relative "parallel/execution_context"
+require_relative "parallel/nested_execution"
+require_relative "parallel/root_execution"
 
 module Smith
   class Workflow
@@ -10,35 +12,30 @@ module Smith
       def self.resolve_branch_count(transition, context)
         count = transition.agent_opts[:count]
         resolved = count.respond_to?(:call) ? count.call(context) : (count || 1)
-        return resolved if resolved.is_a?(Integer) && resolved.positive?
-
-        raise WorkflowError, "parallel branch count must be a positive integer"
+        validate_branch_count!(resolved)
       end
 
-      def self.execute(branches:)
-        signal = CancellationSignal.new
-
-        futures = branches.map do |branch|
-          Concurrent::Promises.future(branch, signal) do |b, s|
-            b.call(s)
-          rescue StandardError
-            s.cancel!
-            raise
-          end
+      def self.validate_branch_count!(count)
+        unless count.is_a?(Integer) && count.positive?
+          raise WorkflowError, "parallel branch count must be a positive integer"
         end
 
-        fulfilled, values, reasons = Concurrent::Promises.zip(*futures).result
+        limit = Smith.config.parallel_branch_limit
+        raise WorkflowError, "parallel branch count exceeds configured limit #{limit}" if count > limit
 
-        unless fulfilled
-          error = preferred_error(reasons)
-          raise error
-        end
+        count
+      end
+
+      def self.execute(branches:)
+        validate_branch_count!(branches.length)
+        context = ExecutionContext.current
+        return NestedExecution.new(branches:, context:).call if context
 
-        values
+        RootExecution.new(branches:).call
       end
 
       def self.preferred_error(reasons)
-        errors = reasons.compact
+        errors = Array(reasons).compact
         errors.find { |error| !error.is_a?(Cancellation) } || errors.first
       end
     end

data/lib/smith/workflow/parallel_agent_binding.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  class Workflow
+    class ParallelAgentBinding
+      extend Dry::Initializer
+
+      param :workflow
+      param :transition
+      param :agent_class
+
+      def initialize(...)
+        super
+        freeze
+      end
+
+      def resolve(workflow:, transition:)
+        agent_class if @workflow.equal?(workflow) && @transition.equal?(transition)
+      end
+    end
+  end
+end

data/lib/smith/workflow/parallel_execution.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "branch_env"
+require_relative "parallel_agent_binding"
 
 module Smith
   class Workflow
@@ -8,7 +9,7 @@ module Smith
       private
 
       def execute_parallel_step(transition, prepared_input: nil)
-        count = Parallel.resolve_branch_count(transition, @context)
+        count = @resolved_parallel_branch_count || Parallel.resolve_branch_count(transition, @context)
         agent_class = resolve_agent_class(transition)
         estimates = compute_branch_estimates(@ledger, branch_count: count, agent_budget: agent_class&.budget)
         env = BranchEnv.new(
@@ -16,7 +17,8 @@ module Smith
           guardrail_sources: Tool.current_guardrails,
           scoped_store: propagate_scoped_artifacts,
           branch_estimates: estimates,
-          deadline: wall_clock_deadline
+          deadline: wall_clock_deadline,
+          agent_class:
         )
         ledger = @ledger
         branches = Array.new(count) do |i|
@@ -27,7 +29,8 @@ module Smith
 
       def run_branch(transition, index, env, ledger, signal)
         setup_branch_context(env, ledger)
-        with_agent_context(resolve_agent_class(transition)) do
+        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
@@ -60,6 +63,7 @@ module Smith
         clear_failed_billable_attempts
         Tool.current_ledger = nil
         Tool.current_tool_result_collector = nil
+        Thread.current[:smith_parallel_agent_binding] = nil
         env.teardown_thread
       end