smith-agents 0.7.0 → 0.8.0

data/lib/smith/tool/argument_snapshot_traversal.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class ArgumentSnapshotTraversal
+      MAX_DEPTH = 100
+      MAX_NODES = 100_000
+      MAX_BYTES = 1 * 1024 * 1024
+
+      extend Dry::Initializer
+
+      param :value
+
+      def call
+        initialize_traversal
+        result = copy
+        ArgumentSnapshotResult.new(
+          value: result,
+          node_count: @accounting.node_count,
+          byte_count: @accounting.byte_count
+        )
+      end
+
+      private
+
+      def initialize_traversal
+        @active = {}.compare_by_identity
+        @copies = {}.compare_by_identity
+        @metrics = {}.compare_by_identity
+        @accounting = ArgumentSnapshotAccounting.new(
+          max_depth: MAX_DEPTH,
+          max_nodes: MAX_NODES,
+          max_bytes: MAX_BYTES
+        )
+        @scalar_snapshot = ArgumentScalarSnapshot.new(
+          byte_counter: @accounting.method(:add_bytes!),
+          byte_validator: @accounting.method(:validate_bytes!)
+        )
+        @container_reader = ArgumentContainerReader.new(
+          scalar_snapshot: @scalar_snapshot,
+          byte_count: -> { @accounting.byte_count }
+        )
+      end
+
+      def copy
+        result = nil
+        pending = [[:visit, value, nil, nil, 0, nil]]
+
+        until pending.empty?
+          entry = pending.pop
+          if entry.first == :finish
+            finish_container!(entry)
+            next
+          end
+
+          _, source, parent, key, depth, parent_metrics = entry
+          @accounting.visit!(depth)
+          type = @container_reader.type(source)
+          copied = if type
+                     copy_container(source, pending, depth, type, parent_metrics)
+                   else
+                     copy_scalar(source, parent_metrics)
+                   end
+          parent ? parent[key] = copied : result = copied
+        end
+
+        result
+      end
+
+      def copy_container(source, pending, depth, type, parent_metrics)
+        raise Error, "tool arguments contain a cyclic value" if @active.key?(source)
+        return copy_reused_container(source, depth, parent_metrics) if @copies.key?(source)
+
+        size = @container_reader.size(source, type)
+        @accounting.reserve_children!(size)
+        captured = @container_reader.snapshot(source, type)
+        unless @container_reader.size(captured, type) == size
+          raise Error, "tool argument container changed while it was being captured"
+        end
+
+        target = @container_reader.target(type, size)
+        metrics = { node_count: 1, byte_count: 0, max_depth: 0 }
+        @copies[source] = target
+        @active[source] = true
+        pending << [:finish, source, target, parent_metrics, metrics]
+        @container_reader.append_children(pending:, source: captured, target:, depth:, metrics:)
+        target
+      end
+
+      def finish_container!(entry)
+        _, source, target, parent_metrics, metrics = entry
+        @active.delete(source)
+        target.freeze
+        @metrics[source] = metrics.freeze
+        @accounting.propagate!(parent_metrics, metrics)
+      end
+
+      def copy_reused_container(source, depth, parent_metrics)
+        metrics = @metrics.fetch(source)
+        @accounting.account_reused!(metrics, depth)
+        @accounting.propagate!(parent_metrics, metrics)
+        @copies.fetch(source)
+      end
+
+      def copy_scalar(source, parent_metrics)
+        previous_bytes = @accounting.byte_count
+        copied = @scalar_snapshot.copy(source)
+        if parent_metrics
+          parent_metrics[:node_count] += 1
+          parent_metrics[:byte_count] += @accounting.byte_count - previous_bytes
+          parent_metrics[:max_depth] = [parent_metrics[:max_depth], 1].max
+        end
+        copied
+      end
+    end
+  end
+end

data/lib/smith/tool/bounded_completion_context.rb

@@ -0,0 +1,131 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module BoundedCompletionContext
+      include BoundedCompletionControls
+      include FailFastCompletion
+      include GracefulCompletion
+
+      SUPPORTED_RUBY_LLM = Gem::Requirement.new("= 1.16.0")
+      REQUIRED_HOOKS = %i[
+        add_tool_result_message forced_tool_choice? handle_concurrent_tool_calls handle_sequential_tool_calls
+        params reset_tool_choice run_callbacks
+      ].freeze
+
+      private_constant :SUPPORTED_RUBY_LLM, :REQUIRED_HOOKS
+
+      def self.install(chat)
+        version = Gem::Version.new(RubyLLM::VERSION)
+        raise Error, "bounded completion requires RubyLLM 1.16.0" unless SUPPORTED_RUBY_LLM.satisfied_by?(version)
+
+        missing = REQUIRED_HOOKS.reject { |hook| chat.respond_to?(hook, true) }
+        valid = missing.empty? && chat.respond_to?(:tool_prefs) && chat.tool_prefs.respond_to?(:replace)
+        raise Error, "unsupported RubyLLM tool-loop interface for bounded completion" unless valid
+
+        chat.extend(self) unless chat.singleton_class < self
+        chat.__send__(:install_bounded_completion_context)
+        chat
+      end
+
+      def complete(&stream)
+        allowance = Tool.current_tool_call_allowance
+        owner = [Thread.current, Fiber.current]
+        iterative = allowance.is_a?(CallAllowance)
+
+        bounded_completion_guard.around_completion(owner, reentrant: !iterative) do
+          return super(&stream) unless iterative
+
+          validate_reserved_params!
+          validate_supported_tools! if allowance.complete_on_exhaustion?
+          state = bounded_completion_guard.state_for(allowance)
+          loop do
+            result = complete_with_budget_policy(state) { super(&stream) }
+            return result unless state.consume_continuation!
+          end
+        end
+      end
+
+      private
+
+      def install_bounded_completion_context
+        return if @smith_bounded_completion_guard
+
+        @smith_bounded_completion_guard = BoundedCompletionGuard.new
+      end
+
+      def complete_with_budget_policy(state, &completion)
+        return complete_with_fail_fast_policy(state, &completion) unless state.allowance.complete_on_exhaustion?
+
+        if state.finalization_required? || state.allowance.remaining.zero?
+          complete_without_tools(state, &completion)
+        else
+          with_bounded_tool_call_preference(state.allowance.remaining) do
+            with_sequential_tool_execution(&completion)
+          end
+        end
+      end
+
+      def handle_tool_calls(response, &)
+        state = current_bounded_completion_state
+        return super unless state
+
+        tool_calls = validated_tool_calls(response)
+        return dispatch_fail_fast_calls(response, tool_calls, state) unless state.allowance.complete_on_exhaustion?
+        return reject_finalization_violation(tool_calls, state) if state.finalization_started?
+
+        admitted, result = dispatch_admitted_calls(response, tool_calls, state)
+        return result if admitted
+
+        complete_rejected_batch(response, tool_calls, state)
+      end
+
+      def execute_tool(tool_call)
+        admission = bounded_completion_guard.admission_for(tool_call)
+        return super unless admission
+
+        CallAdmission.around(admission) { super }
+      end
+
+      def current_bounded_completion_state
+        allowance = Tool.current_tool_call_allowance
+        return unless allowance.is_a?(CallAllowance)
+
+        bounded_completion_guard.state_for(allowance)
+      end
+
+      def dispatch_admitted_calls(response, tool_calls, state)
+        result = nil
+        admitted = bounded_completion_guard.with_admitted_calls(
+          tool_calls,
+          tools:,
+          allowance: state.allowance,
+          ledger: Tool.current_ledger
+        ) do
+          result = handle_sequential_tool_calls(tool_calls)
+        end
+        return [false, nil] unless admitted
+
+        reset_tool_choice if forced_tool_choice?
+        state.request_continuation! unless result
+        [true, result || response]
+      end
+
+      def complete_rejected_batch(response, tool_calls, state)
+        state.request_finalization!
+        append_budget_rejections(tool_calls, state.allowance)
+        reset_tool_choice if forced_tool_choice?
+        state.request_continuation!
+        response
+      end
+
+      def bounded_completion_guard
+        @smith_bounded_completion_guard || raise(Error, "bounded completion context is not installed")
+      end
+
+      def smith_tool_call_admissions(tool_calls)
+        bounded_completion_guard.admissions_for(tool_calls)
+      end
+    end
+  end
+end

data/lib/smith/tool/bounded_completion_controls.rb

@@ -0,0 +1,113 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module BoundedCompletionControls
+      RESERVED_TOOL_PARAM_KEYS = %w[
+        parallel_tool_calls tool_choice tool_config toolConfig tools
+      ].freeze
+
+      private_constant :RESERVED_TOOL_PARAM_KEYS
+
+      private
+
+      def append_budget_rejections(tool_calls, allowance)
+        error = nil
+        Thread.handle_interrupt(Object => :never) do
+          tool_calls.each_value do |tool_call|
+            begin
+              run_callbacks(:before_message, :new_message)
+            rescue Exception => e # rubocop:disable Lint/RescueException
+              error ||= e
+            end
+
+            begin
+              add_tool_result_message(tool_call, budget_rejection(allowance))
+            rescue Exception => e # rubocop:disable Lint/RescueException
+              error ||= e
+            end
+          end
+        end
+        raise error if error
+      end
+
+      def budget_rejection(allowance)
+        {
+          error: {
+            code: "tool_call_budget_exhausted",
+            message: "Tool call was not executed. Complete from available tool results.",
+            remaining: allowance.remaining
+          }
+        }
+      end
+
+      def validated_tool_calls(response)
+        ExecutionBatchCollection.capture(response.tool_calls)
+      end
+
+      def reject_finalization_violation(tool_calls, state)
+        append_budget_rejections(tool_calls, state.allowance)
+        raise BoundedCompletionError, "provider requested a tool during tool-disabled budget finalization"
+      end
+
+      def validate_supported_tools!
+        unsupported = tools.values.grep_v(Smith::Tool)
+        return if unsupported.empty?
+
+        names = unsupported.map(&:name).sort.join(", ")
+        raise BoundedCompletionError,
+              "tool_budget_exhaustion :complete requires Smith::Tool bindings; unsupported: #{names}"
+      end
+
+      def validate_reserved_params!
+        keys = params.respond_to?(:keys) ? params.keys.map(&:to_s) : []
+        reserved = keys & RESERVED_TOOL_PARAM_KEYS
+        return if reserved.empty?
+
+        raise BoundedCompletionError,
+              "bounded completion reserves provider tool params: #{reserved.sort.join(", ")}"
+      end
+
+      def with_bounded_tool_call_preference(remaining)
+        return yield unless tool_call_cardinality_supported?
+
+        original_calls = tool_prefs[:calls]
+        begin
+          tool_prefs[:calls] = remaining > 1 ? :many : :one
+          yield
+        ensure
+          tool_prefs[:calls] = original_calls
+        end
+      end
+
+      def tool_call_cardinality_supported?
+        current_params = params || {}
+        mode = current_params[:openai_api_mode] || current_params["openai_api_mode"]
+        return true if mode.to_s == "responses"
+
+        selected_model = respond_to?(:model) ? model : nil
+        selected_model&.supports?(:parallel_tool_calls) || false
+      end
+
+      def with_tools_disabled
+        original_tools = tools.dup
+        original_preferences = tool_prefs.dup
+        tools.clear
+        tool_prefs[:choice] = :none
+        tool_prefs[:calls] = tool_call_cardinality_supported? ? :one : nil
+        yield
+      ensure
+        tools.replace(original_tools) if original_tools
+        tool_prefs.replace(original_preferences) if original_preferences
+      end
+
+      def with_sequential_tool_execution
+        original = @concurrency
+        @concurrency = nil
+        yield
+      ensure
+        @concurrency = original
+      end
+    end
+  end
+end

data/lib/smith/tool/bounded_completion_guard.rb

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class BoundedCompletionGuard
+      def initialize
+        @state = nil
+        @state_mutex = Mutex.new
+        @owner = nil
+        @depth = 0
+        @owner_mutex = Mutex.new
+        @admissions = nil
+      end
+
+      def around_completion(owner, reentrant:, &block)
+        Thread.handle_interrupt(Object => :never) do
+          enter(owner, reentrant:)
+          begin
+            Thread.handle_interrupt(Object => :immediate, &block)
+          ensure
+            leave(owner)
+          end
+        end
+      end
+
+      def state_for(allowance)
+        @state_mutex.synchronize do
+          return @state if @state&.allowance.equal?(allowance)
+
+          @state = BoundedCompletionState.new(allowance: allowance)
+        end
+      end
+
+      def with_admitted_calls(tool_calls, tools:, allowance:, ledger:, &block)
+        Thread.handle_interrupt(Object => :never) do
+          reservation = allowance.reserve_batch(tool_calls.each_value.map(&:name), ledger:)
+          return false unless reservation
+
+          previous = @admissions
+          @admissions = build_admissions(tool_calls, tools, reservation)
+          begin
+            Thread.handle_interrupt(Object => :immediate, &block)
+          ensure
+            begin
+              reservation.settle!
+            ensure
+              @admissions = previous
+            end
+          end
+          true
+        end
+      end
+
+      def admission_for(tool_call)
+        @admissions&.fetch(tool_call, nil)
+      end
+
+      def admissions_for(tool_calls)
+        tool_calls.each_with_object({}.compare_by_identity) do |tool_call, admissions|
+          admission = admission_for(tool_call)
+          admissions[tool_call] = admission if admission
+        end.freeze
+      end
+
+      private
+
+      def build_admissions(tool_calls, tools, reservation)
+        tool_calls.each_value.with_object({}.compare_by_identity) do |tool_call, admissions|
+          tool = tools[tool_call.name.to_sym]
+          next unless tool
+
+          admissions[tool_call] = CallAdmission.new(tool:, reservation:)
+        end
+      end
+
+      def enter(owner, reentrant:)
+        @owner_mutex.synchronize do
+          if @owner && (@owner != owner || !reentrant)
+            raise Error, "concurrent or reentrant completion on one bounded Smith chat is unsupported"
+          end
+
+          # Each outer completion derives fresh budget policy state, so a
+          # finalized earlier ask can never leak its terminal state into a
+          # later independently bounded completion on the same chat.
+          @state_mutex.synchronize { @state = nil } if @depth.zero?
+          @owner = owner
+          @depth += 1
+        end
+      end
+
+      def leave(owner)
+        @owner_mutex.synchronize do
+          return unless @owner == owner
+
+          @depth -= 1
+          @owner = nil if @depth.zero?
+        end
+      end
+    end
+  end
+end

data/lib/smith/tool/bounded_completion_installation.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module BoundedCompletionInstallation
+      def complete(&)
+        allowance = Tool.current_tool_call_allowance
+        return super unless allowance.is_a?(CallAllowance)
+        return super if singleton_class < BoundedCompletionContext
+
+        BoundedCompletionContext.install(self)
+        complete(&)
+      end
+
+      private
+
+      def smith_tool_call_admissions(_tool_calls) = nil
+    end
+  end
+end

data/lib/smith/tool/bounded_completion_state.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class BoundedCompletionState
+      extend Dry::Initializer
+
+      option :allowance
+
+      def initialize(...)
+        super
+        @state = :open
+        @continuation_required = false
+        @mutex = Mutex.new
+      end
+
+      def request_continuation!
+        @mutex.synchronize { @continuation_required = true }
+      end
+
+      def consume_continuation!
+        @mutex.synchronize do
+          required = @continuation_required
+          @continuation_required = false
+          required
+        end
+      end
+
+      def request_finalization!
+        @mutex.synchronize { @state = :required if @state == :open }
+      end
+
+      def begin_finalization
+        @mutex.synchronize do
+          return :active if @state == :finalizing
+
+          @state = :finalizing
+          :started
+        end
+      end
+
+      def abort_finalization!
+        @mutex.synchronize { @state = :required if @state == :finalizing }
+      end
+
+      def finalization_required?
+        @mutex.synchronize { @state == :required }
+      end
+
+      def finalization_started?
+        @mutex.synchronize { @state == :finalizing }
+      end
+    end
+  end
+end

data/lib/smith/tool/budget_enforcement.rb

@@ -9,15 +9,18 @@ module Smith
         allowance = self.class.current_tool_call_allowance
         ledger = self.class.current_ledger
         workflow_active = ledger&.limits&.key?(:tool_calls)
+        admission = CallAdmission.current
 
-        if allowance.is_a?(CallAllowance)
-          return allowance.charge! { commit_workflow_tool_call!(ledger, workflow_active) }
-        end
-        if allowance.is_a?(Hash)
-          return CallAllowance.charge_legacy!(allowance) { commit_workflow_tool_call!(ledger, workflow_active) }
-        end
+        return if admission&.claim(self)
 
-        commit_workflow_tool_call!(ledger, workflow_active)
+        charge_agent_allowance!(allowance) { commit_workflow_tool_call!(ledger, workflow_active) }
+      end
+
+      def charge_agent_allowance!(allowance, &)
+        return allowance.charge!(name, &) if allowance.is_a?(CallAllowance)
+        return CallAllowance.charge_legacy!(allowance, &) if allowance.is_a?(Hash)
+
+        yield
       end
 
       def commit_workflow_tool_call!(ledger, workflow_active)
@@ -26,6 +29,11 @@ module Smith
         reservation = ledger.reserve!(:tool_calls, 1)
         ledger.reconcile!(reservation, 1)
       end
+
+      def mark_tool_execution_started!
+        self.class.current_tool_execution_tracker&.mark_started!
+        self.class.__send__(:current_tool_dispatch_start_handler)&.call(self)
+      end
     end
   end
 end

data/lib/smith/tool/call_admission.rb

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