smith-agents 0.7.0 → 0.8.0

data/lib/smith/tool/execution_batch_sources.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class ExecutionBatchSources
+      extend Dry::Initializer
+
+      param :entries
+
+      def call
+        bytes = 0
+        metadata = entries.each_with_object([]) do |(key, tool_call), captured|
+          source = ExecutionBatchSourceMetadata.capture(key:, tool_call:)
+          bytes += source.metadata_byte_count
+          validate_bytes!(bytes)
+          captured << source
+        end.freeze
+        metadata.map(&:to_source_call).freeze
+      end
+
+      private
+
+      def validate_bytes!(bytes)
+        return if bytes <= ExecutionBatchSourceCall::MAX_METADATA_BYTES
+
+        raise Error, "provider tool batch metadata exceeds #{ExecutionBatchSourceCall::MAX_METADATA_BYTES} bytes"
+      end
+    end
+  end
+end

data/lib/smith/tool/execution_batch_state.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class ExecutionBatchState
+      SETTLED_STATES = %i[executed failure_notified failure_notifying].freeze
+      private_constant :SETTLED_STATES
+
+      extend Dry::Initializer
+
+      option :requests
+
+      def initialize(...)
+        super
+        @states = requests.each_key.to_h { [_1, :pending] }.compare_by_identity
+        @dispatch_claims = {}.compare_by_identity
+        @notification_attempted = {}.compare_by_identity
+        @unsettled_tool_calls = requests.each_key.to_a.freeze
+        @unsettled_index = 0
+        @mutex = Mutex.new
+      end
+
+      def claim_dispatch!(tool_call)
+        @mutex.synchronize do
+          unless @states[tool_call] == :pending
+            raise ToolDispatchRejected, "admitted tool invocation was already claimed"
+          end
+
+          claim = Object.new.freeze
+          @dispatch_claims[tool_call] = claim
+          @states[tool_call] = :claimed
+          claim
+        end
+      end
+
+      def dispatch_claimed?(tool_call, claim)
+        @mutex.synchronize do
+          @dispatch_claims[tool_call].equal?(claim) && %i[claimed started].include?(@states[tool_call])
+        end
+      end
+
+      def mark_started!(tool_call, claim:) = transition_dispatch!(tool_call, claim, from: :claimed, to: :started)
+
+      def mark_executed!(tool_call, claim:) = transition_dispatch!(tool_call, claim, from: :started, to: :executed)
+
+      def started?(tool_call) = @mutex.synchronize { %i[started executed].include?(@states[tool_call]) }
+
+      def claim_failure_request(tool_call)
+        @mutex.synchronize do
+          state = @states[tool_call]
+          return if state.nil? || SETTLED_STATES.include?(state) || @notification_attempted.key?(tool_call)
+
+          claim_notification!(tool_call)
+          [requests.fetch(tool_call), state]
+        end
+      end
+
+      def claim_unsettled_request
+        @mutex.synchronize do
+          while @unsettled_index < @unsettled_tool_calls.length
+            tool_call = next_unsettled_tool_call
+            state = @states[tool_call]
+            next if SETTLED_STATES.include?(state) || @notification_attempted.key?(tool_call)
+
+            claim_notification!(tool_call)
+            return [tool_call, requests.fetch(tool_call), state]
+          end
+          nil
+        end
+      end
+
+      def complete_failure_notification!(tool_call)
+        @mutex.synchronize do
+          @states[tool_call] = :failure_notified if @states[tool_call] == :failure_notifying
+        end
+      end
+
+      def release_failure_notification!(tool_call, state)
+        @mutex.synchronize do
+          @states[tool_call] = state if @states[tool_call] == :failure_notifying
+        end
+      end
+
+      private
+
+      def next_unsettled_tool_call
+        tool_call = @unsettled_tool_calls[@unsettled_index]
+        @unsettled_index += 1
+        tool_call
+      end
+
+      def claim_notification!(tool_call)
+        @notification_attempted[tool_call] = true
+        @states[tool_call] = :failure_notifying
+      end
+
+      def transition_dispatch!(tool_call, claim, from:, to:)
+        @mutex.synchronize do
+          unless @dispatch_claims[tool_call].equal?(claim) && @states[tool_call] == from
+            raise ToolDispatchRejected, "admitted tool invocation has an invalid dispatch state"
+          end
+
+          @states[tool_call] = to
+        end
+      end
+    end
+  end
+end

data/lib/smith/tool/execution_dispatch.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module ExecutionDispatch
+      private
+
+      def execute_tool(tool_call)
+        batch = execution_batch(tool_call)
+        return super unless batch
+
+        dispatch_call = batch.dispatch_for(tool_call)
+        return execute_unmanaged_tool(batch, tool_call) { super } unless batch.request_for(dispatch_call)
+
+        execute_managed_tool(batch, dispatch_call) { super(dispatch_call) }
+      rescue Exception => e # rubocop:disable Lint/RescueException
+        record_batch_failure(batch, e) if batch
+        raise
+      end
+
+      def execute_managed_tool(batch, dispatch_call, &)
+        current_tool = tools[dispatch_call.name.to_sym]
+        claim = dispatch_claim(batch, dispatch_call)
+        Tool::ScopedContext.around(dispatch_context(batch, dispatch_call, current_tool, claim)) do
+          execute_admitted_tool(batch, dispatch_call, current_tool:, claim:, &)
+        end
+      end
+
+      def dispatch_claim(batch, dispatch_call)
+        Tool::ScopedContext.capture.fetch(:current_tool_dispatch_claim) || batch.claim_dispatch!(dispatch_call)
+      end
+
+      def dispatch_context(batch, dispatch_call, current_tool, claim)
+        request = batch.request_for(dispatch_call)
+        batch.context.merge(
+          current_invocation: request.invocation,
+          current_tool_dispatch_claim: claim,
+          current_tool_dispatch_start_handler: dispatch_start_handler(batch, dispatch_call, current_tool, claim)
+        ).freeze
+      end
+
+      def dispatch_start_handler(batch, dispatch_call, current_tool, claim)
+        lambda do |executing_tool|
+          batch.mark_started!(dispatch_call, claim:) if executing_tool.equal?(current_tool)
+        end
+      end
+
+      def execute_unmanaged_tool(batch, tool_call)
+        Tool::ScopedContext.around(batch.context) do
+          mark_unmanaged_tool_dispatch_started!(tool_call)
+          yield
+        end
+      end
+
+      def execute_admitted_tool(batch, tool_call, current_tool:, claim:, &)
+        batch.verify_dispatch!(tool_call, current_tool:, claim:)
+        return yield if current_tool.__send__(:invocation_argument_error, tool_call.arguments)
+
+        result = dispatch_with_admission(current_tool, claim, batch.admission_for(tool_call), &)
+        batch.mark_executed!(tool_call, claim:)
+        result
+      end
+
+      def dispatch_with_admission(tool, claim, admission, &block)
+        ExecutionAuthority.around(tool:, dispatch_claim: claim) do
+          admission ? CallAdmission.around(admission, &block) : block.call
+        end
+      end
+
+      def mark_unmanaged_tool_dispatch_started!(tool_call)
+        tool = tools[tool_call.name.to_sym]
+        return unless tool.is_a?(RubyLLM::Tool) && !tool.is_a?(Smith::Tool)
+
+        Tool.current_tool_execution_tracker&.mark_started!
+      end
+    end
+  end
+end

data/lib/smith/tool/execution_failure_handling.rb

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module ExecutionFailureHandling
+      private
+
+      def notify_invocation_failure(batch, tool_call, error)
+        return unless batch&.host_admitted?
+
+        entry = batch&.claim_failure_request(tool_call)
+        return unless entry
+
+        request, state = entry
+        batch.context.fetch(:current_invocation_failure_handler)&.call(request:, error:)
+        batch.complete_failure_notification!(tool_call)
+      rescue StandardError => e
+        batch.release_failure_notification!(tool_call, state)
+        raise ToolFailureNotificationFailed.new(
+          dispatch_error: error,
+          notification_error: e
+        ), cause: e
+      end
+
+      def notify_unsettled_batch_failures(batch, error)
+        return unless batch&.host_admitted?
+
+        handler = batch.context.fetch(:current_invocation_failure_handler)
+        handler_error = notify_unsettled_requests(batch, error, handler)
+        raise handler_error if handler_error
+      end
+
+      def notify_unsettled_requests(batch, error, handler)
+        first_error = nil
+        while (entry = batch.claim_unsettled_request)
+          tool_call, request, state = entry
+          failure = batch_dispatch_failure(error, state)
+          begin
+            handler&.call(request:, error: failure)
+            batch.complete_failure_notification!(tool_call)
+          rescue StandardError => e
+            batch.release_failure_notification!(tool_call, state)
+            first_error ||= ToolFailureNotificationFailed.new(
+              dispatch_error: failure,
+              notification_error: e
+            )
+          end
+        end
+        first_error
+      end
+
+      def batch_dispatch_failure(error, state)
+        return error.dispatch_error if error.is_a?(ToolFailureNotificationFailed)
+        return error if state == :started || exact_dispatch_failure?(error)
+
+        ToolDispatchRejected.new("tool batch dispatch failed before execution")
+      end
+
+      def invocation_dispatch_failure(batch, tool_call, error)
+        return error if exact_dispatch_failure?(error) || batch&.started?(tool_call)
+        return error unless batch&.host_admitted?
+
+        ToolDispatchRejected.new("tool invocation failed before execution")
+      end
+
+      def exact_dispatch_failure?(error)
+        error.is_a?(ToolDispatchRejected) || error.is_a?(ToolExecutionNotAdmitted)
+      end
+
+      def record_batch_failure(batch, error)
+        if !error.is_a?(StandardError)
+          batch.fatal_failures.push(error)
+        elsif error.is_a?(ToolFailureNotificationFailed)
+          batch.notification_failures.push(error)
+        elsif error.is_a?(ToolCaptureFailed)
+          batch.capture_failures.push(error)
+        elsif Smith::Errors.retry_forbidden?(error)
+          batch.terminal_failures.push(error)
+        end
+      end
+
+      def prioritized_batch_failure(batch)
+        return unless batch
+
+        %i[fatal_failures notification_failures capture_failures terminal_failures].each do |queue|
+          failure = first_failure(batch.public_send(queue))
+          return failure if failure
+        end
+        nil
+      end
+
+      def first_failure(queue)
+        return unless queue
+
+        queue.pop(true)
+      rescue ThreadError
+        nil
+      end
+    end
+  end
+end

data/lib/smith/tool/execution_lifecycle.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module ExecutionLifecycle
+      private
+
+      def prepare_tool_execution!(kwargs)
+        ensure_capture_ready!
+        run_before_execute_hook!(kwargs)
+        check_dispatch_deadline!
+        check_privilege!(kwargs)
+        check_authorization!(kwargs)
+        run_tool_guardrails!(kwargs)
+        check_dispatch_deadline!
+        charge_tool_call!
+        mark_tool_execution_started!
+      end
+
+      def perform_with_duration(kwargs)
+        start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        result = perform(**kwargs)
+        duration = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start
+        [result, duration]
+      end
+    end
+  end
+end

data/lib/smith/tool/execution_tracker.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class ExecutionTracker
+      def initialize
+        @started = false
+        @mutex = Mutex.new
+      end
+
+      def mark_started!
+        @mutex.synchronize { @started = true }
+      end
+
+      def started?
+        @mutex.synchronize { @started }
+      end
+    end
+  end
+end

data/lib/smith/tool/fail_fast_completion.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module FailFastCompletion
+      private
+
+      def complete_with_fail_fast_policy(state, &)
+        with_bounded_tool_call_preference(state.allowance.remaining, &)
+      end
+
+      def dispatch_fail_fast_calls(response, tool_calls, state)
+        result = if concurrency
+                   handle_concurrent_tool_calls(tool_calls)
+                 else
+                   handle_sequential_tool_calls(tool_calls)
+                 end
+        reset_tool_choice if forced_tool_choice?
+        state.request_continuation! unless result
+        result || response
+      end
+    end
+  end
+end

data/lib/smith/tool/graceful_completion.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module GracefulCompletion
+      private
+
+      def complete_without_tools(state, &)
+        unless state.begin_finalization == :started
+          raise BoundedCompletionError, "tool-disabled budget finalization may run only once"
+        end
+
+        with_tools_disabled(&)
+      rescue Exception # rubocop:disable Lint/RescueException
+        state.abort_finalization!
+        raise
+      end
+    end
+  end
+end

data/lib/smith/tool/invocation.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require "dry-struct"
+
+require_relative "../types"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class Invocation < Dry::Struct
+      OwnedString = Types::String.constructor { |value| value.is_a?(String) ? value.dup.freeze : value }
+      private_constant :OwnedString
+
+      attribute :tool_call_id, OwnedString.optional
+      attribute :tool_name, OwnedString.constrained(min_size: 1)
+      attribute :ordinal, Types::Integer.constrained(gt: 0)
+      attribute :batch_ordinal, Types::Integer.constrained(gt: 0)
+      attribute :batch_size, Types::Integer.constrained(gt: 0)
+
+      def initialize(...)
+        super
+        raise ArgumentError, "tool invocation batch ordinal exceeds its batch size" if batch_ordinal > batch_size
+
+        freeze
+      end
+    end
+  end
+end

data/lib/smith/tool/invocation_request.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require "dry-struct"
+
+require_relative "../types"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class InvocationRequest < Dry::Struct
+      MAX_ARGUMENT_DEPTH = ArgumentSnapshot::MAX_DEPTH
+      MAX_ARGUMENT_NODES = ArgumentSnapshot::MAX_NODES
+      MAX_ARGUMENT_BYTES = ArgumentSnapshot::MAX_BYTES
+
+      ToolClass = Types::Any.constructor do |value|
+        raise TypeError, "expected a Smith::Tool subclass" unless value.is_a?(Class) && value <= Smith::Tool
+
+        value
+      end
+      private_constant :ToolClass
+
+      attribute :invocation, Types.Instance(Invocation)
+      attribute :tool_class, ToolClass
+      attribute :arguments, Types::Hash
+
+      attr_reader :argument_node_count, :argument_byte_count
+
+      def initialize(...)
+        super
+        snapshot = ArgumentSnapshot.new(arguments).call
+        @argument_node_count = snapshot.node_count
+        @argument_byte_count = snapshot.byte_count
+        @attributes = @attributes.merge(arguments: snapshot.value).freeze
+        freeze
+      end
+    end
+  end
+end

data/lib/smith/tool/invocation_sequence.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class InvocationSequence
+      def initialize(next_ordinal: 1)
+        unless next_ordinal.is_a?(Integer) && next_ordinal.positive?
+          raise ArgumentError, "next tool invocation ordinal must be positive"
+        end
+
+        @next_ordinal = next_ordinal
+        @mutex = Mutex.new
+      end
+
+      def reserve(size)
+        raise ArgumentError, "tool invocation batch size must be positive" unless size.is_a?(Integer) && size.positive?
+
+        @mutex.synchronize do
+          first = @next_ordinal
+          @next_ordinal += size
+          first
+        end
+      end
+    end
+  end
+end

data/lib/smith/tool/legacy_call_allowance.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class LegacyCallAllowance
+      MUTEX = Mutex.new
+      ALLOWANCE_MUTEXES = ObjectSpace::WeakMap.new
+      private_constant :MUTEX, :ALLOWANCE_MUTEXES
+
+      def self.charge!(allowance)
+        allowance_mutex(allowance).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
+
+            yield if block_given?
+            allowance[:remaining] = remaining - 1
+          end
+        end
+      end
+
+      def self.reserve_batch(allowance, size, ledger: nil)
+        allowance_mutex(allowance).synchronize do
+          Thread.handle_interrupt(Object => :never) do
+            reserve_synchronized(allowance, size, ledger)
+          end
+        end
+      rescue BudgetExceeded
+        nil
+      end
+
+      def self.allowance_mutex(allowance)
+        MUTEX.synchronize { ALLOWANCE_MUTEXES[allowance] ||= Mutex.new }
+      end
+      private_class_method :allowance_mutex
+
+      def self.reserve_synchronized(allowance, size, ledger)
+        remaining = allowance[:remaining]
+        return unless remaining.is_a?(Integer) && remaining >= size
+
+        workflow_ledger = ledger if ledger&.limits&.key?(:tool_calls)
+        ledger_reservation = workflow_ledger&.reserve!(:tool_calls, size)
+        allowance[:remaining] = remaining - size
+        CallReservation.new(
+          limit: size,
+          ledger: workflow_ledger,
+          ledger_reservation:
+        )
+      end
+      private_class_method :reserve_synchronized
+
+      private_class_method :new
+    end
+  end
+end

data/lib/smith/tool/scoped_context.rb

@@ -8,8 +8,15 @@ module Smith
         current_deadline: :smith_tool_deadline,
         current_ledger: :smith_tool_ledger,
         current_tool_call_allowance: :smith_tool_call_allowance,
+        current_tool_execution_tracker: :smith_tool_execution_tracker,
+        current_tool_dispatch_claim: :smith_tool_dispatch_claim,
+        current_tool_dispatch_start_handler: :smith_tool_dispatch_start_handler,
         current_tool_result_collector: :smith_tool_result_collector,
-        current_invocation_context: :smith_tool_invocation_context
+        current_invocation_context: :smith_tool_invocation_context,
+        current_invocation_batch_admitter: :smith_tool_invocation_batch_admitter,
+        current_invocation_failure_handler: :smith_tool_invocation_failure_handler,
+        current_invocation_sequence: :smith_tool_invocation_sequence,
+        current_invocation: :smith_tool_invocation
       }.freeze
 
       CONTEXT_KEYS.each do |reader, key|
@@ -17,6 +24,18 @@ module Smith
         define_method("#{reader}=") { |value| Thread.current[key] = value }
       end
 
+      private :current_invocation_batch_admitter,
+              :current_invocation_batch_admitter=,
+              :current_invocation_failure_handler,
+              :current_invocation_failure_handler=,
+              :current_invocation_sequence,
+              :current_invocation_sequence=,
+              :current_tool_dispatch_claim,
+              :current_tool_dispatch_claim=,
+              :current_tool_dispatch_start_handler,
+              :current_tool_dispatch_start_handler=,
+              :current_invocation=
+
       def self.capture
         CONTEXT_KEYS.to_h { |reader, key| [reader, Thread.current[key]] }.freeze
       end
@@ -36,27 +55,60 @@ module Smith
         end
       end
 
-      def with_invocation_context(value, &block)
+      def with_invocation_context(value, invocation_sequence: InvocationSequence.new, batch_admitter: nil,
+                                  failure_handler: nil, &block)
+        raise ArgumentError, "block required" unless block
+        unless invocation_sequence.is_a?(InvocationSequence)
+          raise ArgumentError, "invocation sequence must be a Smith::Tool::InvocationSequence"
+        end
+
+        validate_callback!(batch_admitter, "batch admitter")
+        validate_callback!(failure_handler, "failure handler")
+        if batch_admitter && !failure_handler
+          raise ArgumentError, "invocation failure handler is required when a batch admitter is configured"
+        end
+
+        values = ScopedContext.capture.merge(
+          current_invocation_context: value,
+          current_invocation_batch_admitter: batch_admitter,
+          current_invocation_failure_handler: failure_handler,
+          current_invocation_sequence: invocation_sequence,
+          current_invocation: nil
+        ).freeze
+        ScopedContext.around(values, &block)
+      end
+
+      def with_call_budget(budget, on_exhaustion: :raise, &block)
         raise ArgumentError, "block required" unless block
 
-        previous = current_invocation_context
+        previous = current_tool_call_allowance
+        if previous && !previous.is_a?(CallAllowance)
+          raise ArgumentError, "a scoped tool call budget cannot nest under a legacy Hash tool call allowance"
+        end
+
+        allowance = previous ? previous.scope(budget, on_exhaustion:) : CallAllowance.new(budget, on_exhaustion:)
         Thread.handle_interrupt(Object => :never) do
-          self.current_invocation_context = value
+          self.current_tool_call_allowance = allowance
           begin
             Thread.handle_interrupt(Object => :immediate, &block)
           ensure
-            self.current_invocation_context = previous
+            self.current_tool_call_allowance = previous
           end
         end
       end
 
       def self.install(values)
-        CONTEXT_KEYS.each do |reader, key|
-          Thread.current[key] = values.fetch(reader)
-        end
+        CONTEXT_KEYS.each { |reader, key| Thread.current[key] = values.fetch(reader) }
       end
       private_class_method :install
 
+      def validate_callback!(callback, name)
+        return if callback.nil? || callback.respond_to?(:call)
+
+        raise ArgumentError, "invocation #{name} must respond to #call"
+      end
+      private :validate_callback!
+
       def self.validate!(values)
         complete = values.is_a?(Hash) && values.length == CONTEXT_KEYS.length && CONTEXT_KEYS.each_key.all? do |key|
           values.key?(key)