smith-agents 0.4.0

data/lib/smith/workflow/dsl.rb

@@ -0,0 +1,223 @@
+# frozen_string_literal: true
+
+module Smith
+  class Workflow
+    module DSL
+      def self.included(base)
+        base.extend(ClassMethods)
+      end
+
+      module ClassMethods
+        def inherited(subclass)
+          super
+          subclass.instance_variable_set(:@states, (@states || []).dup)
+          subclass.instance_variable_set(:@transitions, (@transitions || {}).dup)
+          subclass.instance_variable_set(:@initial_state_name, @initial_state_name)
+          subclass.instance_variable_set(:@budget_config, @budget_config&.dup)
+          subclass.instance_variable_set(:@max_transitions_count, @max_transitions_count)
+          subclass.instance_variable_set(:@guardrails_class, @guardrails_class)
+          subclass.instance_variable_set(:@context_manager_class, @context_manager_class)
+          subclass.instance_variable_set(:@seed_messages_builder, @seed_messages_builder)
+          subclass.instance_variable_set(:@persistence_key_builder, @persistence_key_builder)
+          subclass.instance_variable_set(:@persistence_schema_version, @persistence_schema_version)
+          subclass.instance_variable_set(:@migrations, (@migrations || {}).dup)
+          subclass.instance_variable_set(:@seed_validation_mode, @seed_validation_mode)
+          subclass.instance_variable_set(:@idempotency_mode, @idempotency_mode)
+          subclass.instance_variable_set(:@persistence_ttl, @persistence_ttl)
+        end
+
+        def initial_state(name = nil)
+          return @initial_state_name if name.nil?
+
+          @initial_state_name = name
+          state(name)
+        end
+
+        def state(name)
+          @states ||= []
+          @states << name unless @states.include?(name)
+          generate_fail_transition if name == :failed
+        end
+
+        def transition(name, from:, to:, &)
+          @transitions ||= {}
+          @transitions[name] = Transition.new(name, from: from, to: to, &)
+        end
+
+        def budget(**opts)
+          return @budget_config if opts.empty?
+
+          @budget_config = opts
+        end
+
+        def max_transitions(count = nil)
+          return @max_transitions_count if count.nil?
+
+          @max_transitions_count = count
+        end
+
+        def guardrails(klass = nil)
+          return @guardrails_class if klass.nil?
+
+          @guardrails_class = klass
+        end
+
+        def context_manager(klass = nil)
+          return @context_manager_class if klass.nil?
+
+          @context_manager_class = klass
+        end
+
+        def seed_messages(&block)
+          return @seed_messages_builder unless block_given?
+
+          @seed_messages_builder = block
+        end
+
+        def persistence_key(&block)
+          return @persistence_key_builder unless block_given?
+
+          @persistence_key_builder = block
+        end
+
+        # Schema version stamped into every persisted payload's
+        # :schema_version key. Restore compares the stored version with
+        # this value and either passes through (equal), applies
+        # registered migrate_from blocks one step at a time (stored
+        # less than current), or raises Smith::PersistenceSchemaMismatch
+        # (stored greater than current, or unbridged gap).
+        # Pre-versioning payloads (no :schema_version key) are treated
+        # as v1 for backward compatibility.
+        def persistence_schema_version(version = nil)
+          return @persistence_schema_version || 1 if version.nil?
+
+          unless version.is_a?(Integer) && version >= 1
+            raise ArgumentError, "persistence_schema_version must be a positive Integer, got #{version.inspect}"
+          end
+
+          @persistence_schema_version = version
+        end
+
+        # Register a one-step migration from stored version N to N+1.
+        # The block receives the persisted payload Hash (top-level keys
+        # already symbolized) and must return the migrated payload.
+        # Bumping the :schema_version key is the migration's
+        # responsibility but Smith advances defensively if the block
+        # omits it, so migrations stay loop-free.
+        def migrate_from(version, &block)
+          raise ArgumentError, "migrate_from requires a block" unless block
+
+          unless version.is_a?(Integer) && version >= 1
+            raise ArgumentError, "migrate_from version must be a positive Integer, got #{version.inspect}"
+          end
+
+          @migrations ||= {}
+          @migrations[version] = block
+        end
+
+        def migrations
+          @migrations || {}
+        end
+
+        # Controls whether restore validates that the seed_messages
+        # builder still produces the same digest as when this workflow
+        # was originally persisted.
+        #
+        # Modes:
+        #   :off    (default) skip validation entirely. Recommended when
+        #            the seed builder is non-deterministic (timestamps,
+        #            UUIDs, request-scoped data) since drift would
+        #            surface on every restore.
+        #   :warn   log a warning via Smith.config.logger on drift; do
+        #            not raise. Suitable for soft monitoring.
+        #   :strict raise Smith::SeedMismatch on drift. Suitable when
+        #            the seed builder is deterministic (system
+        #            instructions, static templates) and divergence
+        #            indicates a code change that would invalidate the
+        #            persisted conversation context.
+        def seed_validation(mode = nil)
+          return @seed_validation_mode || :off if mode.nil?
+
+          unless %i[strict warn off].include?(mode)
+            raise ArgumentError, "seed_validation must be :strict, :warn, or :off, got #{mode.inspect}"
+          end
+
+          @seed_validation_mode = mode
+        end
+
+        # Controls whether run_persisted! / advance_persisted! stamp a
+        # step_in_progress marker before each advance and clear it
+        # afterward.
+        #
+        # Modes:
+        #   :lax    (default) no marker stamping; restore never raises.
+        #            Safe when agent calls and tools are idempotent, so
+        #            re-running a step on restore is harmless.
+        #   :strict marker is persisted before each advance and cleared
+        #            after. Restore raises
+        #            Smith::StepInProgressOnRestore when the marker is
+        #            set, indicating a previous worker crashed mid-step
+        #            and re-running could double-execute non-idempotent
+        #            agent calls or tools.
+        def idempotency_mode(mode = nil)
+          return @idempotency_mode || :lax if mode.nil?
+
+          unless %i[strict lax].include?(mode)
+            raise ArgumentError, "idempotency_mode must be :strict or :lax, got #{mode.inspect}"
+          end
+
+          @idempotency_mode = mode
+        end
+
+        # Per-workflow TTL override (in seconds). Takes precedence over
+        # Smith.config.persistence_ttl at persist! time. nil (default)
+        # means inherit the global config.
+        #
+        # Hosts typically set this when different workflow classes have
+        # different durability horizons: e.g., short-lived UI sessions
+        # at 1.day.to_i, long-running research workflows at 30.days.to_i.
+        #
+        # Wiring contract: when the resolved TTL is non-nil,
+        # Workflow#persist! forwards it to the adapter as a `ttl:`
+        # kwarg. Shipped adapters (Memory, RedisStore, CacheStore,
+        # ActiveRecordStore) accept this kwarg; external duck-typed
+        # adapters that implement only the bare REQUIRED_METHODS contract
+        # without a `ttl:` kwarg will only break when a host actually
+        # opts into TTL.
+        def persistence_ttl(seconds = nil)
+          return @persistence_ttl if seconds.nil?
+
+          unless seconds.is_a?(Numeric) && seconds.positive?
+            raise ArgumentError,
+                  "persistence_ttl must be a positive Numeric (seconds), got #{seconds.inspect}"
+          end
+
+          @persistence_ttl = seconds
+        end
+
+        def transitions_from(state)
+          (@transitions || {}).values.select { |t| t.from == state }
+        end
+
+        def find_transition(name)
+          (@transitions || {})[name]
+        end
+
+        def from_state(hash)
+          workflow = allocate
+          workflow.send(:restore_state, hash)
+          workflow
+        end
+
+        private
+
+        def generate_fail_transition
+          @transitions ||= {}
+          return if @transitions.key?(:fail)
+
+          @transitions[:fail] = Transition.new(:fail, from: nil, to: :failed)
+        end
+      end
+    end
+  end
+end

data/lib/smith/workflow/durability.rb

@@ -0,0 +1,369 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Smith
+  class Workflow
+    module Durability
+      def self.included(base)
+        base.extend(ClassMethods)
+      end
+
+      module ClassMethods
+        def restore(key, adapter: Smith.persistence_adapter)
+          resolved_key = explicit_persistence_key!(key)
+          payload = fetch_persisted_payload(resolved_key, adapter:)
+          return nil unless payload
+
+          from_state(JSON.parse(payload)).tap do |workflow|
+            workflow.instance_variable_set(:@persistence_key, resolved_key)
+          end
+        end
+
+        def restore_or_initialize(key: nil, context: {}, adapter: Smith.persistence_adapter, **kwargs)
+          restore(resolved_persistence_key(key:, context:), adapter:) || new(context:, **kwargs)
+        end
+
+        # Peek without instantiating: returns true if persisted state
+        # exists for the resolved key, false otherwise. Reuses the
+        # existing private helpers (`resolved_persistence_key`,
+        # `fetch_persisted_payload`) so it doesn't expand the adapter
+        # contract — `Smith::PersistenceAdapter` requires only
+        # `store/fetch/delete`, and this peek piggybacks on `fetch`.
+        # Custom adapters work without changes.
+        #
+        # Hadithi uses this to skip the credits guard at execution time
+        # when persisted state already exists for a workflow key (a
+        # prior attempt's billable work is durable in Redis, OR an
+        # in-flight workflow is being resumed — either way, no NEW
+        # credit authorization is needed).
+        def persisted_state_exists?(key: nil, context: {}, adapter: Smith.persistence_adapter)
+          resolved_key = resolved_persistence_key(key:, context:)
+          !fetch_persisted_payload(resolved_key, adapter:).nil?
+        end
+
+        # Stricter peek: returns true only when persisted state contains
+        # billable work that needs to be preserved (at least one
+        # `usage_entries` entry).
+        #
+        # `persisted_state_exists?` answers "is there any state?" — but
+        # that includes the bare initial-state record Smith writes at
+        # the top of `run_persisted!` BEFORE the first `advance!`. A
+        # worker that dies between that initial `persist!` and the
+        # first model call leaves a Redis key with no billable work.
+        # If the credits guard's bypass keys on `persisted_state_exists?`
+        # alone, a zero-balance user's retry on that abandoned init
+        # state silently runs the first model call (the guard is
+        # skipped because state exists, but the state has nothing to
+        # bill — it's just the workflow's starting state).
+        #
+        # `restorable_billing_state?` returns true only when there's
+        # actual `usage_entries` to bill on idempotent replay. Terminal
+        # state with zero entries is also `false` because there's
+        # nothing to preserve — `run_persisted!` is a no-op on
+        # terminal anyway, so guard outcome doesn't matter for
+        # correctness in that case.
+        #
+        # This calls `restore` (full deserialize) rather than just
+        # `fetch`, so it's heavier than `persisted_state_exists?`. Use
+        # this when you specifically want the billing-aware semantics;
+        # use `persisted_state_exists?` when you only need a key-
+        # presence check.
+        def restorable_billing_state?(key: nil, context: {}, adapter: Smith.persistence_adapter)
+          resolved_key = resolved_persistence_key(key:, context:)
+          payload = fetch_persisted_payload(resolved_key, adapter:)
+          return false if payload.nil?
+
+          workflow = from_state(JSON.parse(payload))
+          entries = workflow.instance_variable_get(:@usage_entries) || []
+          entries.any?
+        end
+
+        def stuck_for?(persistence_key:, threshold:, since: nil, adapter: Smith.persistence_adapter)
+          raise WorkflowError, "persistence_adapter is not configured" if adapter.nil?
+          raise ArgumentError, "persistence_key must not be blank" if persistence_key.nil? || (persistence_key.respond_to?(:strip) && persistence_key.strip.empty?)
+          raise ArgumentError, "threshold must respond to :to_i" unless threshold.respond_to?(:to_i)
+
+          if since && !since.respond_to?(:to_time)
+            raise ArgumentError, "since must respond to :to_time or be nil"
+          end
+
+          threshold_seconds = threshold.to_i
+          now = Time.now.utc
+
+          if Smith::PersistenceAdapters.supports?(adapter, :last_heartbeat)
+            hb = adapter.last_heartbeat(persistence_key)
+            if hb
+              age = (now - hb.to_time.utc).to_f
+              return false if age < threshold_seconds
+            end
+          end
+
+          payload = adapter.fetch(persistence_key)
+          return stuck_for_no_payload?(since, now, threshold_seconds) if payload.nil?
+
+          if !Smith::PersistenceAdapters.supports?(adapter, :last_heartbeat)
+            Smith::PersistenceAdapters.warn_missing_heartbeat(adapter)
+            fallback_age = age_from_payload_updated_at(payload, now)
+            return false if fallback_age.nil? || fallback_age < threshold_seconds
+          end
+
+          !terminal_in_payload?(payload)
+        end
+
+        def heartbeat_age(persistence_key:, adapter: Smith.persistence_adapter)
+          raise WorkflowError, "persistence_adapter is not configured" if adapter.nil?
+
+          if Smith::PersistenceAdapters.supports?(adapter, :last_heartbeat)
+            hb = adapter.last_heartbeat(persistence_key)
+            return [Time.now.utc - hb.to_time.utc, 0.0].max.to_f if hb
+          end
+
+          payload = adapter.fetch(persistence_key)
+          return nil if payload.nil?
+
+          age_from_payload_updated_at(payload, Time.now.utc)
+        end
+
+        def run_persisted!(key: nil, context: {}, adapter: Smith.persistence_adapter, clear: :done, on_step: nil, **kwargs)
+          clear_policy = normalize_clear_policy(clear)
+          resolved_key = resolved_persistence_key(key:, context:)
+          workflow = restore(resolved_key, adapter:) || new(context:, **kwargs)
+          result = workflow.run_persisted!(resolved_key, adapter:, on_step:)
+
+          workflow.clear_persisted!(resolved_key, adapter:) if clear_persisted_after_run?(clear_policy, workflow)
+          result
+        end
+
+        private
+
+        def fetch_persisted_payload(key, adapter:)
+          persistence_adapter!(adapter).fetch(key)
+        end
+
+        def stuck_for_no_payload?(since, now, threshold_seconds)
+          return false if since.nil?
+
+          age = (now - since.to_time.utc).to_f
+          age.clamp(0.0, Float::INFINITY) >= threshold_seconds
+        end
+
+        def age_from_payload_updated_at(payload, now)
+          parsed = JSON.parse(payload)
+          updated_at = parsed["updated_at"] || parsed[:updated_at]
+          return nil if updated_at.nil?
+
+          (now - Time.parse(updated_at.to_s).utc).to_f
+        rescue JSON::ParserError, ArgumentError
+          nil
+        end
+
+        def terminal_in_payload?(payload)
+          parsed = JSON.parse(payload)
+          state_name = parsed["state"] || parsed[:state]
+          class_name = parsed["class"] || parsed[:class]
+          next_transition = parsed["next_transition_name"] || parsed[:next_transition_name]
+
+          return false unless state_name && class_name
+
+          klass = Object.const_get(class_name)
+          klass.transitions_from(state_name.to_sym).empty? && next_transition.nil?
+        rescue JSON::ParserError, NameError, NoMethodError
+          false
+        end
+
+        def persistence_adapter!(adapter)
+          return adapter if adapter
+
+          raise WorkflowError, "persistence_adapter is not configured"
+        end
+
+        def resolve_persistence_key!(key:, context:)
+          return key unless key.nil? || blank_key?(key)
+
+          builder = persistence_key
+          raise WorkflowError, "persistence key is required unless workflow defines persistence_key" unless builder
+
+          derived = if builder.arity == 1
+            builder.call(context)
+          else
+            builder.call
+          end
+
+          return derived unless blank_key?(derived)
+
+          raise WorkflowError, "persistence_key must return a non-blank key"
+        end
+
+        def resolved_persistence_key(key:, context:)
+          resolve_persistence_key!(key:, context:)
+        end
+
+        def explicit_persistence_key!(key)
+          return key unless blank_key?(key)
+
+          raise WorkflowError, "restore requires a non-blank explicit persistence key"
+        end
+
+        def blank_key?(value)
+          return true if value.nil?
+          return true if value.respond_to?(:strip) && value.strip.empty?
+
+          value.respond_to?(:empty?) ? value.empty? : false
+        end
+
+        def normalize_clear_policy(clear)
+          case clear
+          when false, nil
+            false
+          when true, :done
+            :done
+          when :terminal
+            :terminal
+          else
+            raise WorkflowError, "invalid clear policy #{clear.inspect}; expected false, :done, or :terminal"
+          end
+        end
+
+        def clear_persisted_after_run?(clear_policy, workflow)
+          return false if clear_policy == false
+          return workflow.done? if clear_policy == :done
+
+          workflow.terminal?
+        end
+      end
+
+      def run_persisted!(key = nil, adapter: Smith.persistence_adapter, on_step: nil)
+        return build_run_result([]) if terminal?
+
+        resolved_key = resolve_persistence_key!(key)
+        steps = []
+        persist!(resolved_key, adapter:)
+
+        until terminal?
+          if strict_idempotency?
+            mark_step_in_progress!
+            persist!(resolved_key, adapter:)
+          end
+
+          step = advance!
+          steps << step if step
+
+          clear_step_in_progress!
+          persist!(resolved_key, adapter:)
+          invoke_on_step_callback(step, on_step) if step
+        end
+
+        build_run_result(steps)
+      end
+
+      def advance_persisted!(key = nil, adapter: Smith.persistence_adapter, on_step: nil)
+        return if terminal?
+
+        resolved_key = resolve_persistence_key!(key)
+        mark_step_in_progress! if strict_idempotency?
+        persist!(resolved_key, adapter:)
+        step = advance!
+        clear_step_in_progress!
+        persist!(resolved_key, adapter:) if step
+        invoke_on_step_callback(step, on_step) if step
+        step
+      end
+
+      def persist!(key = nil, adapter: Smith.persistence_adapter)
+        resolved_key = resolve_persistence_key!(key)
+        store = persistence_adapter!(adapter)
+        previous_version = @persistence_version || 0
+        next_version = previous_version + 1
+        payload = JSON.generate(to_state.merge(persistence_version: next_version))
+
+        dispatch_store!(store, resolved_key, payload, previous_version: previous_version)
+
+        # Increment ONLY after successful store. On PersistenceVersionConflict
+        # (raised by store_versioned), @persistence_version stays at the
+        # previous value so callers can rescue + restore + retry cleanly.
+        @persistence_version = next_version
+        self
+      end
+
+      def clear_persisted!(key = nil, adapter: Smith.persistence_adapter)
+        resolved_key = resolve_persistence_key!(key)
+        persistence_adapter!(adapter).delete(resolved_key)
+        self
+      end
+
+      def mark_step_in_progress!
+        @step_in_progress = true
+      end
+
+      def clear_step_in_progress!
+        @step_in_progress = false
+      end
+
+      private
+
+      def strict_idempotency?
+        self.class.idempotency_mode == :strict
+      end
+
+      # Forwards the persist payload to the adapter, splatting `ttl:`
+      # only when a TTL is resolved. The empty-Hash splat is a no-op so
+      # external duck-typed adapters that don't accept a `ttl:` kwarg
+      # keep working; they only break if the host opts into TTL.
+      def dispatch_store!(store, key, payload, previous_version:)
+        kwargs = ttl_kwarg(effective_persistence_ttl)
+
+        if Smith::PersistenceAdapters.supports?(store, :store_versioned)
+          store.store_versioned(key, payload, expected_version: previous_version, **kwargs)
+        else
+          Smith::PersistenceAdapters.warn_missing_versioning(store)
+          store.store(key, payload, **kwargs)
+        end
+
+        store.record_heartbeat(key, **kwargs) if Smith::PersistenceAdapters.supports?(store, :record_heartbeat)
+      end
+
+      # Resolves the effective TTL once per persist, with the class-level
+      # DSL override taking precedence over Smith.config.persistence_ttl.
+      # Returns nil when neither is set, which collapses ttl_kwarg to an
+      # empty Hash and preserves the pre-TTL adapter call shape.
+      def effective_persistence_ttl
+        self.class.persistence_ttl || Smith.config.persistence_ttl
+      end
+
+      def ttl_kwarg(ttl)
+        ttl ? { ttl: ttl } : {}
+      end
+
+      def persistence_adapter!(adapter)
+        return adapter if adapter
+
+        raise WorkflowError, "persistence_adapter is not configured"
+      end
+
+      def resolve_persistence_key!(key)
+        unless key.nil? || blank_key?(key)
+          @persistence_key = key
+          return key
+        end
+
+        return @persistence_key unless blank_key?(@persistence_key)
+
+        @persistence_key = self.class.send(:resolve_persistence_key!, key:, context: @context)
+      end
+
+      def blank_key?(value)
+        return true if value.nil?
+        return true if value.respond_to?(:strip) && value.strip.empty?
+
+        value.respond_to?(:empty?) ? value.empty? : false
+      end
+
+      def invoke_on_step_callback(step, callback)
+        callback&.call(step)
+      rescue StandardError => e
+        Smith.config.logger&.error("Smith::Workflow on_step callback error: #{e.message}")
+      end
+    end
+  end
+end