smith-agents 0.4.4 → 0.5.0

data/lib/smith/persistence_adapters/memory.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "monitor"
+require "securerandom"
 
 module Smith
   module PersistenceAdapters
@@ -15,7 +16,10 @@ module Smith
     # Smith.config.persistence_adapter is nil AND Smith.config.test_mode
     # is true (typically set in spec_helper.rb).
     class Memory
-      def initialize
+      attr_reader :persistence_identity
+
+      def initialize(identity: "memory:#{SecureRandom.uuid}")
+        @persistence_identity = identity.to_s.dup.freeze
         @store = {}
         @heartbeats = {}
         @monitor = Monitor.new
@@ -23,21 +27,16 @@ module Smith
 
       def store(key, payload, ttl: Smith.config.persistence_ttl)
         @monitor.synchronize do
-          @store[key] = { payload: payload, expires_at: ttl ? Time.now.utc + ttl : nil }
+          @store[key] = entry_for(payload, ttl)
         end
       end
 
       def fetch(key)
         @monitor.synchronize do
-          entry = @store[key]
+          entry = live_entry(key)
           next nil if entry.nil?
 
-          if entry[:expires_at] && entry[:expires_at] < Time.now.utc
-            @store.delete(key)
-            next nil
-          end
-
-          entry[:payload]
+          copy_payload(entry[:payload])
         end
       end
 
@@ -48,6 +47,9 @@ module Smith
         end
       end
 
+      def transaction_open? = false
+      def transaction_identity = nil
+
       def record_heartbeat(key, ttl: Smith.config.persistence_ttl)
         @monitor.synchronize do
           @heartbeats[key] = { at: Time.now.utc, expires_at: ttl ? Time.now.utc + ttl : nil }
@@ -76,7 +78,7 @@ module Smith
       # across all versioned adapters).
       def store_versioned(key, payload, expected_version:, ttl: Smith.config.persistence_ttl)
         @monitor.synchronize do
-          entry = @store[key]
+          entry = live_entry(key)
           if entry
             current_version = parse_version(entry[:payload])
             if current_version != expected_version
@@ -84,9 +86,21 @@ module Smith
                 key: key, expected: expected_version, actual: current_version
               )
             end
+          else
+            VersionExpectation.validate_missing!(key, expected_version)
           end
-          @store[key] = { payload: payload, expires_at: ttl ? Time.now.utc + ttl : nil }
+          @store[key] = entry_for(payload, ttl)
+        end
+      end
+
+      def replace_exact(key, payload, expected_payload:, ttl: Smith.config.persistence_ttl)
+        @monitor.synchronize do
+          entry = live_entry(key)
+          raise PersistencePayloadConflict.new(key:) unless entry && entry[:payload] == expected_payload
+
+          @store[key] = entry_for(payload, ttl)
         end
+        payload
       end
 
       def clear!
@@ -95,10 +109,27 @@ module Smith
 
       private
 
+      def live_entry(key)
+        entry = @store[key]
+        return unless entry
+        return entry unless entry[:expires_at] && entry[:expires_at] < Time.now.utc
+
+        @store.delete(key)
+        nil
+      end
+
+      def entry_for(payload, ttl)
+        { payload: copy_payload(payload), expires_at: ttl ? Time.now.utc + ttl : nil }
+      end
+
+      def copy_payload(payload)
+        payload.dup
+      rescue TypeError
+        payload
+      end
+
       def parse_version(payload)
-        JSON.parse(payload).fetch("persistence_version", 0)
-      rescue JSON::ParserError
-        0
+        PayloadVersion.call(payload)
       end
     end
   end

data/lib/smith/persistence_adapters/payload_version.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Smith
+  module PersistenceAdapters
+    module PayloadVersion
+      module_function
+
+      def call(payload)
+        document = payload.is_a?(String) ? JSON.parse(payload) : payload
+        raise TypeError, "persisted workflow payload must be a JSON object" unless document.is_a?(Hash)
+
+        version = document.fetch("persistence_version", 0)
+        unless version.is_a?(Integer) && version >= 0
+          raise TypeError, "persisted workflow persistence_version must be a non-negative integer"
+        end
+
+        version
+      rescue JSON::ParserError
+        0
+      end
+    end
+  end
+end

data/lib/smith/persistence_adapters/redis_client_access.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Smith
+  module PersistenceAdapters
+    module RedisClientAccess
+      private
+
+      def client
+        return @client if instance_variable_defined?(:@client)
+
+        @client_resolution_mutex.synchronize do
+          return @client if instance_variable_defined?(:@client)
+
+          resolved = redis_client?(@redis_source) ? @redis_source : resolve_redis_source
+          raise ArgumentError, "Redis client is required" unless resolved
+
+          @client = resolved
+        end
+      end
+
+      def resolve_redis_source
+        @redis_source.respond_to?(:call) ? @redis_source.call : @redis_source
+      end
+
+      def redis_client?(candidate)
+        candidate.respond_to?(:get) && candidate.respond_to?(:set) && candidate.respond_to?(:del)
+      end
+
+      def without_reconnection
+        redis = client
+        return redis.without_reconnect { yield redis } if redis.respond_to?(:without_reconnect)
+        return redis.disable_reconnection { yield redis } if redis.respond_to?(:disable_reconnection)
+
+        raise ArgumentError, "Redis client must expose without_reconnect or disable_reconnection for CAS writes"
+      end
+    end
+  end
+end

data/lib/smith/persistence_adapters/redis_exact_write.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  module PersistenceAdapters
+    class RedisExactWrite
+      extend Dry::Initializer
+
+      option :client
+      option :key
+      option :storage_key
+      option :payload
+      option :expected_payload
+      option :ttl
+
+      def call
+        result = client.watch(storage_key) do
+          validate_current!
+          enqueue_write
+        end
+        return payload unless result.nil?
+
+        raise PersistencePayloadConflict.new(key:)
+      end
+
+      private
+
+      def validate_current!
+        return if client.get(storage_key) == expected_payload
+
+        client.unwatch
+        raise PersistencePayloadConflict.new(key:)
+      end
+
+      def enqueue_write
+        client.multi do |transaction|
+          options = ttl ? { ex: ttl } : {}
+          transaction.set(storage_key, payload, **options)
+        end
+      end
+    end
+  end
+end

data/lib/smith/persistence_adapters/redis_store.rb

@@ -2,9 +2,13 @@
 
 require "time"
 
+require_relative "redis_client_access"
+
 module Smith
   module PersistenceAdapters
     class RedisStore
+      include RedisClientAccess
+
       # Redis transient errors — narrow list; non-transient errors
       # (CommandError, etc.) propagate up immediately. Pattern matches
       # Redis::BaseConnectionError if loaded (covers Connection/Timeout)
@@ -24,9 +28,13 @@ module Smith
         end + [Errno::ECONNREFUSED, Errno::ETIMEDOUT, Errno::EPIPE]
       end
 
-      def initialize(redis:, namespace: "smith")
+      attr_reader :persistence_identity
+
+      def initialize(redis:, namespace: "smith", identity: nil)
         @redis_source = redis
-        @namespace = namespace
+        @namespace = namespace.nil? ? nil : namespace.to_s.dup.freeze
+        @persistence_identity = identity.to_s.dup.freeze if identity
+        @client_resolution_mutex = Mutex.new
       end
 
       def store(key, payload, ttl: Smith.config.persistence_ttl)
@@ -51,6 +59,9 @@ module Smith
         end
       end
 
+      def transaction_open? = false
+      def transaction_identity = nil
+
       def record_heartbeat(key, ttl: Smith.config.persistence_ttl)
         Retry.with_retries(operation: :record_heartbeat, transient: self.class.transient_errors) do
           iso = Time.now.utc.iso8601
@@ -77,47 +88,37 @@ module Smith
       # Smith::PersistenceVersionConflict on a stale expected_version
       # OR on EXEC failure (WATCH detected concurrent write).
       def store_versioned(key, payload, expected_version:, ttl: Smith.config.persistence_ttl)
-        Retry.with_retries(operation: :store_versioned, transient: self.class.transient_errors) do
-          namespaced_key = namespaced(key)
-          result = client.watch(namespaced_key) do
-            current = client.get(namespaced_key)
-            if current && (current_version = parse_version(current)) != expected_version
-              client.unwatch
-              raise Smith::PersistenceVersionConflict.new(
-                key: key, expected: expected_version, actual: current_version
-              )
-            end
-
-            client.multi do |tx|
-              if ttl
-                tx.set(namespaced_key, payload, ex: ttl)
-              else
-                tx.set(namespaced_key, payload)
-              end
-            end
-          end
-
-          if result.nil?
-            raise Smith::PersistenceVersionConflict.new(
-              key: key, expected: expected_version, actual: :concurrent
-            )
-          end
-
-          result
+        without_reconnection do |redis|
+          RedisVersionedWrite.new(
+            client: redis,
+            key: key,
+            storage_key: namespaced(key),
+            payload: payload,
+            expected_version: expected_version,
+            ttl: ttl
+          ).call
         end
+      rescue *self.class.transient_errors => e
+        raise Smith::PersistenceIOError.new(operation: :store_versioned, cause: e)
       end
 
-      private
-
-      def client
-        @client ||= begin
-          resolved = @redis_source.respond_to?(:call) ? @redis_source.call : @redis_source
-          raise ArgumentError, "Redis client is required" unless resolved
-
-          resolved
+      def replace_exact(key, payload, expected_payload:, ttl: Smith.config.persistence_ttl)
+        without_reconnection do |redis|
+          RedisExactWrite.new(
+            client: redis,
+            key: key,
+            storage_key: namespaced(key),
+            payload: payload,
+            expected_payload: expected_payload,
+            ttl: ttl
+          ).call
         end
+      rescue *self.class.transient_errors => e
+        raise Smith::PersistenceIOError.new(operation: :replace_exact, cause: e)
       end
 
+      private
+
       def namespaced(key)
         [@namespace, key].compact.join(":")
       end
@@ -125,12 +126,6 @@ module Smith
       def namespaced_heartbeat(key)
         [@namespace, "heartbeat", key].compact.join(":")
       end
-
-      def parse_version(payload)
-        JSON.parse(payload).fetch("persistence_version", 0)
-      rescue JSON::ParserError
-        0
-      end
     end
   end
 end

data/lib/smith/persistence_adapters/redis_versioned_write.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  module PersistenceAdapters
+    class RedisVersionedWrite
+      extend Dry::Initializer
+
+      option :client
+      option :key
+      option :storage_key
+      option :payload
+      option :expected_version
+      option :ttl
+
+      def call
+        result = client.watch(storage_key) do
+          current = client.get(storage_key)
+          validate_current!(current)
+          enqueue_write
+        end
+        return result unless result.nil?
+
+        raise PersistenceVersionConflict.new(
+          key: key,
+          expected: expected_version,
+          actual: :concurrent
+        )
+      end
+
+      private
+
+      def validate_current!(current)
+        return VersionExpectation.validate_missing!(key, expected_version) unless current
+
+        current_version = PayloadVersion.call(current)
+        return if current_version == expected_version
+
+        client.unwatch
+        raise PersistenceVersionConflict.new(
+          key: key,
+          expected: expected_version,
+          actual: current_version
+        )
+      end
+
+      def enqueue_write
+        client.multi do |transaction|
+          options = ttl ? { ex: ttl } : {}
+          transaction.set(storage_key, payload, **options)
+        end
+      end
+    end
+  end
+end

data/lib/smith/persistence_adapters/retry.rb

@@ -16,27 +16,39 @@ module Smith
 
       def with_retries(operation:, transient:, policy: Smith.config.persistence_retry_policy,
                        logger: Smith.config.logger)
-        attempts = policy.fetch(:attempts, 3)
-        base = policy.fetch(:base_delay, 0.1)
-        max_delay = policy.fetch(:max_delay, 1.0)
+        schedule = schedule_for(policy)
         last_error = nil
 
-        attempts.times do |i|
+        schedule.attempts.times do |i|
           return yield
         rescue *transient => e
           last_error = e
-          break if i == attempts - 1
+          break if i == schedule.attempts - 1
 
-          delay = [base * (2**i), max_delay].min
-          logger&.warn(
-            "Smith::PersistenceAdapters::Retry #{operation} attempt #{i + 1}/#{attempts} failed: " \
-            "#{e.class}: #{e.message}; sleeping #{delay}s"
-          )
+          delay = schedule.delay(i + 1)
+          log_retry(logger, operation, e, { attempt: i + 1, attempts: schedule.attempts, delay: })
           sleep(delay)
         end
 
         raise Smith::PersistenceIOError.new(operation: operation, cause: last_error)
       end
+
+      def schedule_for(policy)
+        ExponentialBackoff.new(
+          attempts: policy.fetch(:attempts, 3),
+          base_delay: policy.fetch(:base_delay, 0.1),
+          max_delay: policy.fetch(:max_delay, 1.0),
+          jitter: 0
+        )
+      end
+
+      def log_retry(logger, operation, error, retry_context)
+        logger&.warn(
+          "Smith::PersistenceAdapters::Retry #{operation} " \
+          "attempt #{retry_context.fetch(:attempt)}/#{retry_context.fetch(:attempts)} failed: " \
+          "#{error.class}: #{error.message}; sleeping #{retry_context.fetch(:delay)}s"
+        )
+      end
     end
   end
 end

data/lib/smith/persistence_adapters/version_expectation.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Smith
+  module PersistenceAdapters
+    module VersionExpectation
+      module_function
+
+      def validate_missing!(key, expected_version)
+        return if expected_version.is_a?(Integer) && expected_version.zero?
+
+        raise Smith::PersistenceVersionConflict.new(
+          key: key,
+          expected: expected_version,
+          actual: :missing
+        )
+      end
+    end
+  end
+end

data/lib/smith/persistence_adapters.rb

@@ -1,8 +1,16 @@
 # frozen_string_literal: true
 
 require "monitor"
+require_relative "persistence_adapters/active_record_connection_errors"
+require_relative "persistence_adapters/active_record_initial_write"
+require_relative "persistence_adapters/active_record_exact_write"
+require_relative "persistence_adapters/active_record_exact_store"
+require_relative "persistence_adapters/payload_version"
+require_relative "persistence_adapters/version_expectation"
 require_relative "persistence_adapters/cache_store"
 require_relative "persistence_adapters/rails_cache"
+require_relative "persistence_adapters/redis_versioned_write"
+require_relative "persistence_adapters/redis_exact_write"
 require_relative "persistence_adapters/redis_store"
 require_relative "persistence_adapters/active_record_store"
 require_relative "persistence_adapters/memory"
@@ -23,7 +31,10 @@ module Smith
     # check support via `supports?(adapter, capability)` and fall back
     # gracefully (e.g., Workflow#persist! warns once and uses plain
     # `store` when `store_versioned` is missing).
-    OPTIONAL_METHODS = %i[store_versioned record_heartbeat last_heartbeat].freeze
+    OPTIONAL_METHODS = %i[
+      store_versioned replace_exact persistence_identity record_heartbeat last_heartbeat
+      transaction_open? transaction_identity
+    ].freeze
 
     def self.resolve(adapter, **options)
       return nil if adapter.nil?
@@ -34,22 +45,7 @@ module Smith
         return validate!(instance)
       end
 
-      built_in = case adapter.to_sym
-                 when :cache_store
-                   CacheStore.new(**options)
-                 when :rails_cache, :solid_cache
-                   RailsCache.new(**options)
-                 when :redis
-                   RedisStore.new(**options)
-                 when :active_record
-                   ActiveRecordStore.new(**options)
-                 when :memory
-                   Memory.new
-                 else
-                   raise ArgumentError, "Unknown persistence adapter #{adapter.inspect}"
-                 end
-
-      validate!(built_in)
+      validate!(build(adapter, options))
     end
 
     def self.adapter_like?(adapter)
@@ -60,9 +56,21 @@ module Smith
       return adapter if adapter_like?(adapter)
 
       missing = REQUIRED_METHODS.reject { |method_name| adapter.respond_to?(method_name) }
-      raise ArgumentError, "Persistence adapter must implement #{missing.join(', ')}"
+      raise ArgumentError, "Persistence adapter must implement #{missing.join(", ")}"
     end
 
+    def self.build(adapter, options)
+      case adapter.to_sym
+      when :cache_store then CacheStore.new(**options)
+      when :rails_cache, :solid_cache then RailsCache.new(**options)
+      when :redis then RedisStore.new(**options)
+      when :active_record then ActiveRecordStore.new(**options)
+      when :memory then Memory.new(**options)
+      else raise ArgumentError, "Unknown persistence adapter #{adapter.inspect}"
+      end
+    end
+    private_class_method :build
+
     # Capability introspection used by Workflow#persist! to decide
     # whether the adapter supports optimistic locking via store_versioned.
     def self.supports?(adapter, capability)

data/lib/smith/persistence_payload_conflict.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module Smith
+  class PersistencePayloadConflict < Error
+    attr_reader :key
+
+    def initialize(key:)
+      @key = key
+      super("persisted payload conflict for #{key.inspect}: the exact expected payload is no longer current")
+    end
+  end
+end

data/lib/smith/tool/budget_enforcement.rb

@@ -22,8 +22,8 @@ module Smith
 
       def commit_tool_call_charges!(ledger, allowance, workflow_active)
         if workflow_active
-          ledger.reserve!(:tool_calls, 1)
-          ledger.reconcile!(:tool_calls, 1, 1)
+          reservation = ledger.reserve!(:tool_calls, 1)
+          ledger.reconcile!(reservation, 1)
         end
 
         allowance[:remaining] -= 1 if allowance

data/lib/smith/types.rb

@@ -12,5 +12,6 @@ module Smith
     # the constants we need as direct module constants.
     const_set(:String, const_get(:String))
     const_set(:Integer, const_get(:Integer))
+    const_set(:Sha256Hex, const_get(:String).constrained(format: /\A[0-9a-f]{64}\z/))
   end
 end

data/lib/smith/version.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
 module Smith
-  VERSION = "0.4.4"
+  VERSION = "0.5.0"
+  EXECUTION_SEMANTICS_VERSION = "1"
 end

data/lib/smith/workflow/branch_env.rb

@@ -4,7 +4,7 @@ module Smith
   class Workflow
     # rubocop:disable Style/RedundantStructKeywordInit
     BranchEnv = Struct.new(
-      :prepared_input, :guardrail_sources, :scoped_store, :branch_estimates, :deadline,
+      :prepared_input, :guardrail_sources, :scoped_store, :branch_estimates, :deadline, :agent_class,
       keyword_init: true
     ) do
       def setup_thread

data/lib/smith/workflow/budget_integration.rb

@@ -15,8 +15,7 @@ module Smith
       def reserve_branch_budget(ledger, branch_estimates:)
         return nil unless ledger && branch_estimates
 
-        branch_estimates.each { |dim, amount| ledger.reserve!(dim, amount) }
-        branch_estimates
+        reserve_estimates!(ledger, branch_estimates)
       end
 
       def compute_branch_estimates(ledger, branch_count:, agent_budget: nil)
@@ -29,13 +28,16 @@ module Smith
         end
       end
 
-      def reconcile_branch_budget(ledger, estimates, agent_result: nil)
-        return unless ledger && estimates
+      def reconcile_branch_budget(ledger, reservation, agent_result: nil)
+        return unless ledger && reservation
 
         actuals = extract_actuals(agent_results_for_settlement(agent_result))
-        estimates.each do |dim, amt|
-          ledger.reconcile!(dim, amt, actual_for_dimension(dim, actuals[:tokens], actuals[:cost]))
+        settlement = reservation.amounts.to_h do |dimension, _amount|
+          [dimension, actual_for_dimension(dimension, actuals[:tokens], actuals[:cost])]
         end
+        return ledger.reconcile_many!(reservation, actual: settlement) if settlement.length > 1
+
+        ledger.reconcile!(reservation, settlement.values.first)
       end
 
       def extract_actuals(agent_results)
@@ -66,10 +68,12 @@ module Smith
         0
       end
 
-      def release_branch_budget(ledger, estimates)
-        return unless ledger && estimates
+      def release_branch_budget(ledger, reservation)
+        return unless ledger && reservation
+
+        return ledger.release_many!(reservation) if reservation.amounts.length > 1
 
-        estimates.each { |dim, amount| ledger.release!(dim, amount) }
+        ledger.release!(reservation)
       end
 
       def settle_budget_on_failure(ledger, estimates, agent_result)
@@ -91,8 +95,14 @@ module Smith
           est[dim] = cap ? [remaining, cap].min : remaining
         end
 
-        estimates.each { |dim, amount| ledger.reserve!(dim, amount) }
-        estimates
+        reserve_estimates!(ledger, estimates)
+      end
+
+      def reserve_estimates!(ledger, estimates)
+        return ledger.reserve_many!(estimates) if estimates.length > 1
+
+        dimension, amount = estimates.first
+        ledger.reserve!(dimension, amount)
       end
 
       def finalize_branch(transition, index, result, ledger, reserved)