smith-agents 0.7.0 → 0.8.0

data/lib/smith/tool/call_allowance.rb

@@ -3,59 +3,132 @@
 module Smith
   class Tool < RubyLLM::Tool
     class CallAllowance
-      LEGACY_MUTEX = Mutex.new
-      LEGACY_ALLOWANCE_MUTEXES = ObjectSpace::WeakMap.new
-      private_constant :LEGACY_MUTEX
-      private_constant :LEGACY_ALLOWANCE_MUTEXES
+      EXHAUSTION_POLICIES = %i[raise complete].freeze
 
       def self.charge_legacy!(allowance)
-        legacy_allowance_mutex(allowance).synchronize do
+        LegacyCallAllowance.charge!(allowance) { yield if block_given? }
+      end
+
+      def initialize(remaining, on_exhaustion: :raise)
+        invalid = (remaining.is_a?(Integer) && remaining.negative?) ||
+                  (!remaining.is_a?(Integer) && !remaining.is_a?(CallBudget))
+        raise ArgumentError, "tool call allowance must be a non-negative integer" if invalid
+
+        initialize_scope(CallBudget.coerce(remaining), on_exhaustion:, parent: nil)
+      end
+
+      def scope(budget, on_exhaustion: @on_exhaustion)
+        self.class.allocate.tap do |allowance|
+          allowance.__send__(:initialize_scope, CallBudget.coerce(budget), on_exhaustion:, parent: self)
+        end
+      end
+
+      def charge!(tool_name = nil)
+        batch = CallBatch.coerce(exact? ? [tool_name] : 1, exact: exact?)
+
+        synchronize do
           Thread.handle_interrupt(Object => :never) do
-            remaining = allowance[:remaining]
-            unless remaining.is_a?(Integer) && remaining.positive?
-              raise BudgetExceeded, "agent tool_calls budget exceeded"
-            end
+            raise BudgetExceeded, "agent tool_calls budget exceeded" unless reservable?(batch)
 
             yield if block_given?
-            allowance[:remaining] = remaining - 1
+            consume!(batch)
           end
         end
       end
 
-      def self.legacy_allowance_mutex(allowance)
-        LEGACY_MUTEX.synchronize do
-          LEGACY_ALLOWANCE_MUTEXES[allowance] ||= Mutex.new
-        end
+      def remaining
+        synchronize { lineage.map { _1.__send__(:counter).remaining }.min }
       end
-      private_class_method :legacy_allowance_mutex
 
-      def initialize(remaining)
-        unless remaining.is_a?(Integer) && remaining >= 0
-          raise ArgumentError, "tool call allowance must be a non-negative integer"
-        end
+      def remaining_for(tool_name)
+        name = canonical_tool_name(tool_name)
+        synchronize do
+          exact_scopes = lineage.select { _1.__send__(:budget).exact? }
+          return nil if exact_scopes.empty?
 
-        @remaining = remaining
-        @mutex = Mutex.new
+          exact_scopes.map { _1.__send__(:counter).remaining_for(name) }.min
+        end
       end
 
-      def charge!
-        @mutex.synchronize do
+      def reserve_batch(tool_names_or_size, ledger: nil)
+        batch = CallBatch.coerce(tool_names_or_size, exact: exact?)
+
+        synchronize do
           Thread.handle_interrupt(Object => :never) do
-            raise BudgetExceeded, "agent tool_calls budget exceeded" unless @remaining.positive?
+            return unless reservable?(batch)
 
-            yield if block_given?
-            @remaining -= 1
+            workflow_ledger = ledger if ledger&.limits&.key?(:tool_calls)
+            ledger_reservation = workflow_ledger&.reserve!(:tool_calls, batch.size)
+            reservation = CallReservation.new(
+              limit: batch.size,
+              ledger: workflow_ledger,
+              ledger_reservation: ledger_reservation
+            )
+
+            consume!(batch)
+            reservation
           end
         end
+      rescue BudgetExceeded
+        nil
       end
 
-      def remaining
-        @mutex.synchronize { @remaining }
+      def complete_on_exhaustion?
+        @on_exhaustion == :complete
+      end
+
+      def exact?
+        lineage.any? { _1.__send__(:budget).exact? }
+      end
+
+      def used?
+        synchronize { counter.used? }
       end
 
       def [](key)
         remaining if key == :remaining
       end
+
+      protected
+
+      attr_reader :budget
+
+      private
+
+      def initialize_scope(budget, on_exhaustion:, parent:)
+        unless EXHAUSTION_POLICIES.include?(on_exhaustion)
+          raise ArgumentError, "tool call exhaustion policy must be :raise or :complete"
+        end
+
+        @budget = budget
+        @parent = parent
+        @lineage = [self, *Array(parent&.__send__(:lineage))].freeze
+        @mutex = parent ? parent.__send__(:shared_mutex) : Mutex.new
+        @counter = CallAllowanceCounter.new(budget)
+        @on_exhaustion = on_exhaustion
+        self
+      end
+
+      attr_reader :lineage, :counter
+
+      def shared_mutex = @mutex
+
+      def synchronize(&) = @mutex.synchronize(&)
+
+      def canonical_tool_name(value)
+        name = value.to_s
+        raise ArgumentError, "tool name must not be empty" if name.empty?
+
+        name
+      end
+
+      def reservable?(batch)
+        lineage.all? { _1.__send__(:counter).available?(batch) }
+      end
+
+      def consume!(batch)
+        lineage.each { _1.__send__(:counter).consume!(batch) }
+      end
     end
   end
 end

data/lib/smith/tool/call_allowance_counter.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class CallAllowanceCounter
+      def initialize(budget)
+        @budget = budget
+        @remaining = budget.total
+        @remaining_by_tool = budget.tool_limits&.dup
+      end
+
+      attr_reader :remaining
+
+      def remaining_for(tool_name)
+        @remaining_by_tool.fetch(tool_name, 0)
+      end
+
+      def used?
+        remaining < @budget.total
+      end
+
+      def available?(batch)
+        remaining >= batch.size && tool_counts_available?(batch.counts)
+      end
+
+      def consume!(batch)
+        @remaining -= batch.size
+        return unless @budget.exact?
+
+        batch.counts.each { |name, count| @remaining_by_tool[name] -= count }
+      end
+
+      private
+
+      def tool_counts_available?(counts)
+        return true unless @budget.exact?
+        return false unless counts
+
+        counts.all? { |name, count| @remaining_by_tool.fetch(name, 0) >= count }
+      end
+    end
+  end
+end

data/lib/smith/tool/call_batch.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class CallBatch
+      attr_reader :size, :counts
+
+      def self.coerce(value, exact:)
+        return from_size(value, exact:) if value.is_a?(Integer)
+        raise ArgumentError, "tool call batch size must be a positive integer" unless value.is_a?(Array)
+
+        names = value.map { canonical_tool_name(_1) }
+        validate_size!(names.length)
+        new(size: names.length, counts: names.tally.freeze)
+      end
+
+      def self.from_size(size, exact:)
+        validate_size!(size)
+        raise ArgumentError, "exact tool call allowance requires tool names for batch reservation" if exact
+
+        new(size:, counts: nil)
+      end
+      private_class_method :from_size
+
+      def self.canonical_tool_name(value)
+        name = value.to_s
+        raise ArgumentError, "tool name must not be empty" if name.empty?
+
+        name
+      end
+      private_class_method :canonical_tool_name
+
+      def self.validate_size!(size)
+        return if size.is_a?(Integer) && size.positive?
+
+        raise ArgumentError, "tool call batch size must be a positive integer"
+      end
+      private_class_method :validate_size!
+
+      def initialize(size:, counts:)
+        @size = size
+        @counts = counts
+        freeze
+      end
+    end
+  end
+end

data/lib/smith/tool/call_budget.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require "dry-struct"
+
+require_relative "../types"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class CallBudget < Dry::Struct
+      OwnedString = Types::String.constructor { |value| value.is_a?(String) ? value.dup.freeze : value }
+      ToolName = OwnedString.constrained(min_size: 1)
+      Limit = Types::Integer.constrained(gteq: 0)
+      PositiveLimit = Types::Integer.constrained(gt: 0)
+      ToolLimits = Types::Hash.map(ToolName, PositiveLimit)
+
+      private_constant :OwnedString, :ToolName, :Limit, :PositiveLimit, :ToolLimits
+
+      attribute :total, Limit
+      attribute :tool_limits, ToolLimits.optional.default(nil)
+
+      def self.coerce(value)
+        return value if value.is_a?(self)
+        return new(total: value) if value.is_a?(Integer)
+
+        raise ArgumentError, "tool call budget must be an integer or Smith::Tool::CallBudget"
+      end
+
+      def initialize(...)
+        super
+        validate_tool_limits!
+        tool_limits&.freeze
+        freeze
+      end
+
+      def exact? = !tool_limits.nil?
+
+      def limit_for(tool_name)
+        tool_limits&.fetch(tool_name.to_s, nil)
+      end
+
+      private
+
+      def validate_tool_limits!
+        return unless tool_limits
+        return if total <= tool_limits.values.sum
+
+        raise ArgumentError, "aggregate tool call allowance exceeds the sum of its per-tool limits"
+      end
+    end
+  end
+end

data/lib/smith/tool/call_reservation.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class CallReservation
+      extend Dry::Initializer
+
+      option :limit
+      option :ledger
+      option :ledger_reservation
+
+      def initialize(...)
+        super
+        @claimed = 0
+        @settled = false
+        @mutex = Mutex.new
+      end
+
+      def claim
+        @mutex.synchronize do
+          return false if @settled || @claimed >= limit
+
+          @claimed += 1
+          true
+        end
+      end
+
+      def settle!
+        Thread.handle_interrupt(Exception => :never) do
+          @mutex.synchronize do
+            return if @settled
+
+            ledger&.reconcile!(ledger_reservation, @claimed)
+            @settled = true
+          end
+        end
+      end
+    end
+  end
+end

data/lib/smith/tool/chat_execution_callbacks.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module ChatExecutionCallbacks
+      private
+
+      def execute_tool_with_callbacks(tool_call)
+        batch = execution_batch(tool_call)
+        dispatch_call = batch&.dispatch_for(tool_call)
+        return super unless dispatch_call && batch.request_for(dispatch_call)
+
+        execute_managed_tool_callbacks(batch, dispatch_call) do
+          super(batch.source_for(dispatch_call))
+        end
+      end
+
+      def execute_managed_tool_callbacks(batch, dispatch_call, &block)
+        claim = batch.claim_dispatch!(dispatch_call)
+        with_dispatch_claim(batch, claim) { complete_managed_callbacks(batch, dispatch_call, &block) }
+      rescue StandardError => e
+        handle_managed_callback_failure(batch, dispatch_call, claim, e)
+      rescue Exception => e # rubocop:disable Lint/RescueException
+        record_batch_failure(batch, e)
+        raise
+      end
+
+      def with_dispatch_claim(batch, claim, &)
+        context = batch.context.merge(current_tool_dispatch_claim: claim).freeze
+        Tool::ScopedContext.around(context, &)
+      end
+
+      def complete_managed_callbacks(batch, dispatch_call)
+        result = yield
+        notify_rejected_callback_dispatch(batch, dispatch_call) unless batch.started?(dispatch_call)
+        result
+      end
+
+      def handle_managed_callback_failure(batch, dispatch_call, claim, error)
+        record_batch_failure(batch, error)
+        raise error unless claim
+        raise error if error.is_a?(ToolFailureNotificationFailed)
+
+        failure = invocation_dispatch_failure(batch, dispatch_call, error)
+        notify_invocation_failure(batch, dispatch_call, failure)
+        raise failure
+      end
+
+      def notify_rejected_callback_dispatch(batch, dispatch_call)
+        notify_invocation_failure(
+          batch,
+          dispatch_call,
+          ToolDispatchRejected.new("tool invocation arguments were rejected before execution")
+        )
+      end
+    end
+  end
+end

data/lib/smith/tool/chat_execution_context.rb

@@ -3,6 +3,12 @@
 module Smith
   class Tool < RubyLLM::Tool
     module ChatExecutionContext
+      include BoundedCompletionInstallation
+      include ExecutionFailureHandling
+      include ExecutionDispatch
+      include ChatExecutionCallbacks
+      include ExecutionBatchLifecycle
+
       def self.install(chat)
         return chat unless chat.respond_to?(:tools) && chat.tools.respond_to?(:values)
         unless chat.respond_to?(:execute_tool, true)
@@ -10,96 +16,9 @@ module Smith
         end
 
         chat.extend(self) unless chat.singleton_class < self
-        chat.__send__(:install_smith_tool_execution_context)
+        chat.__send__(:smith_tool_execution_batches)
         chat
       end
-
-      private
-
-      def install_smith_tool_execution_context
-        return if defined?(@smith_tool_execution_batches_mutex)
-
-        @smith_tool_execution_batches = {}.compare_by_identity
-        @smith_tool_execution_batches_mutex = Mutex.new
-      end
-
-      def execute_tool(tool_call)
-        batch = execution_batch(tool_call)
-        return super unless batch
-
-        Tool::ScopedContext.around(batch.fetch(:context)) { super }
-      rescue Exception => e # rubocop:disable Lint/RescueException
-        record_batch_failure(batch, e) if batch
-        raise
-      end
-
-      def execute_tools_concurrently(tool_calls, ...)
-        batch = execution_batch_context
-        register_execution_batch(tool_calls, batch)
-        super
-      rescue Exception => e # rubocop:disable Lint/RescueException
-        raise e unless e.is_a?(StandardError)
-
-        fatal_failure = first_failure(batch&.fetch(:fatal_failures))
-        raise fatal_failure if fatal_failure
-
-        capture_failure = first_failure(batch&.fetch(:capture_failures))
-        raise capture_failure if capture_failure
-
-        raise
-      ensure
-        unregister_execution_batch(tool_calls, batch)
-      end
-
-      def execution_batch_context
-        {
-          context: Tool::ScopedContext.capture,
-          capture_failures: Queue.new,
-          fatal_failures: Queue.new
-        }.freeze
-      end
-
-      def register_execution_batch(tool_calls, batch)
-        @smith_tool_execution_batches_mutex.synchronize do
-          tool_calls.each_value do |tool_call|
-            raise Error, "tool call is already active on this chat" if @smith_tool_execution_batches.key?(tool_call)
-
-            @smith_tool_execution_batches[tool_call] = batch
-          end
-        end
-      end
-
-      def unregister_execution_batch(tool_calls, batch)
-        return unless tool_calls && batch
-
-        @smith_tool_execution_batches_mutex.synchronize do
-          tool_calls.each_value do |tool_call|
-            @smith_tool_execution_batches.delete(tool_call) if @smith_tool_execution_batches[tool_call].equal?(batch)
-          end
-        end
-      end
-
-      def execution_batch(tool_call)
-        @smith_tool_execution_batches_mutex.synchronize do
-          @smith_tool_execution_batches[tool_call]
-        end
-      end
-
-      def record_batch_failure(batch, error)
-        if error.is_a?(ToolCaptureFailed)
-          batch.fetch(:capture_failures).push(error)
-        elsif !error.is_a?(StandardError)
-          batch.fetch(:fatal_failures).push(error)
-        end
-      end
-
-      def first_failure(queue)
-        return unless queue
-
-        queue.pop(true)
-      rescue ThreadError
-        nil
-      end
     end
   end
 end

data/lib/smith/tool/execution_authority.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class ExecutionAuthority
+      THREAD_KEY = :smith_tool_execution_authority
+      private_constant :THREAD_KEY
+
+      def self.current = Thread.current[THREAD_KEY]
+
+      def self.around(tool:, dispatch_claim:, &block)
+        raise ArgumentError, "block required" unless block
+
+        previous = current
+        Thread.handle_interrupt(Object => :never) do
+          Thread.current[THREAD_KEY] = new(tool:, dispatch_claim:)
+          begin
+            Thread.handle_interrupt(Object => :immediate, &block)
+          ensure
+            Thread.current[THREAD_KEY] = previous
+          end
+        end
+      end
+
+      def initialize(tool:, dispatch_claim:)
+        @tool = tool
+        @dispatch_claim = dispatch_claim
+        @available = true
+        @mutex = Mutex.new
+      end
+
+      def claim(tool, dispatch_claim)
+        @mutex.synchronize do
+          return false unless @available && @tool.equal?(tool) && @dispatch_claim.equal?(dispatch_claim)
+
+          @available = false
+          true
+        end
+      end
+    end
+  end
+end

data/lib/smith/tool/execution_authorization.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module ExecutionAuthorization
+      private
+
+      def authorize_tool_execution!
+        claim = self.class.__send__(:current_tool_dispatch_claim)
+        return unless claim
+        return if ExecutionAuthority.current&.claim(self, claim)
+
+        raise ToolExecutionNotAdmitted, "managed Smith tool execution requires its exact admitted dispatch authority"
+      end
+    end
+  end
+end

data/lib/smith/tool/execution_batch.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class ExecutionBatch
+      extend Dry::Initializer
+
+      option :context
+      option :tool_calls
+      option :source_calls
+      option :dispatch_calls
+      option :source_calls_by_dispatch
+      option :requests
+      option :targets
+      option :call_admissions, default: proc { {}.compare_by_identity.freeze }
+      option :call_reservation, optional: true
+
+      attr_reader :capture_failures, :fatal_failures, :notification_failures, :terminal_failures
+
+      def initialize(...)
+        super
+        @state = ExecutionBatchState.new(requests:)
+        @mutex = Mutex.new
+        initialize_failure_queues
+        @host_admitted = false
+        @host_admission_required = context.fetch(:current_invocation_batch_admitter) && requests.any?
+      end
+
+      def request_for(tool_call) = requests[tool_call]
+
+      def target_for(tool_call) = targets[tool_call]
+
+      def admission_for(tool_call) = call_admissions[tool_call]
+
+      def dispatch_for(tool_call) = requests.key?(tool_call) ? tool_call : dispatch_calls[tool_call]
+
+      def source_for(tool_call) = source_calls_by_dispatch[tool_call]
+
+      def claim_dispatch!(tool_call) = @state.claim_dispatch!(tool_call)
+
+      def verify_dispatch!(tool_call, current_tool:, claim:)
+        request = requests.fetch(tool_call)
+        target = targets.fetch(tool_call)
+        return if dispatch_admitted? &&
+                  @state.dispatch_claimed?(tool_call, claim) &&
+                  target_unchanged?(current_tool, target, request) &&
+                  call_unchanged?(tool_call, request)
+
+        raise ToolDispatchRejected, "admitted tool invocation changed before dispatch"
+      end
+
+      def mark_started!(tool_call, claim:) = @state.mark_started!(tool_call, claim:)
+
+      def mark_executed!(tool_call, claim:) = @state.mark_executed!(tool_call, claim:)
+
+      def started?(tool_call) = @state.started?(tool_call)
+
+      def mark_host_admitted!
+        @mutex.synchronize { @host_admitted = true }
+      end
+
+      def host_admitted?
+        @mutex.synchronize { @host_admitted }
+      end
+
+      def claim_failure_request(tool_call) = @state.claim_failure_request(tool_call)
+
+      def claim_unsettled_request = @state.claim_unsettled_request
+
+      def complete_failure_notification!(tool_call) = @state.complete_failure_notification!(tool_call)
+
+      def release_failure_notification!(tool_call, state) = @state.release_failure_notification!(tool_call, state)
+
+      def settle! = call_reservation&.settle!
+
+      private
+
+      def initialize_failure_queues
+        @capture_failures = Queue.new
+        @fatal_failures = Queue.new
+        @notification_failures = Queue.new
+        @terminal_failures = Queue.new
+      end
+
+      def dispatch_admitted? = !@host_admission_required || host_admitted?
+
+      def target_unchanged?(current_tool, target, request)
+        current_tool.equal?(target) && current_tool.instance_of?(request.tool_class)
+      end
+
+      def call_unchanged?(tool_call, request)
+        tool_call.id.equal?(request.invocation.tool_call_id) &&
+          tool_call.name.to_s == request.invocation.tool_name &&
+          tool_call.arguments.equal?(request.arguments)
+      end
+    end
+  end
+end