smith-agents 0.11.0 → 0.12.0

data/lib/smith/workflow/wait_execution.rb

@@ -0,0 +1,180 @@
+# frozen_string_literal: true
+
+require "time"
+
+require_relative "await_step"
+require_relative "message_value_normalizer"
+
+module Smith
+  class Workflow
+    # The wait a run is in: an `await` for `event` (delivered under `token`
+    # by `resume!`), a `delay`, a map's or call's wait for its `children`, a
+    # batch step's wait for the results of its `children` requests
+    # (`:batch`, event `:batch_result`, with its `attempt`), or a durable
+    # retry's backoff (`:retry`, with the `attempt` it waits to make and the
+    # class name of the `error` that failed the last one), entered at
+    # `entered_at` and due to expire at `deadline_at`. Hosts index it to deliver events and to schedule the
+    # `expire_wait!` call; Smith never sleeps, polls or wakes a run itself.
+    Wait = Data.define(:kind, :event, :token, :transition, :into, :entered_at, :deadline_at, :children, :attempt,
+                       :error)
+
+    # What a step that entered a wait returns instead of its output.
+    WAIT_PENDING = Object.new.freeze
+    private_constant :WAIT_PENDING
+
+    # Entering, resuming and expiring `await` and `delay` steps. Entering a
+    # wait is not a step: `step_count` stays, so the step's position and its
+    # token (a key under a position reserved for the wait) are the same when
+    # the wait is entered, when its block runs again after a crash, and when
+    # it completes. The completion (`resume!`, or `expire_wait!` past the
+    # deadline) is the step. One completion at a time claims the wait, so a
+    # delivery and an expiry racing on one instance cannot both win.
+    module WaitExecution
+      def waiting? = !@awaiting.nil?
+
+      # The run's wait, or nil.
+      def awaiting
+        record = @awaiting
+        WaitRecord.to_wait(record, awaited_transition(record), @pending_retry, @pending_batch) if record
+      end
+
+      # Delivers the event this run awaits: writes the normalized payload to
+      # the step's `into:` key and completes the step. A delivery can arrive
+      # before the wait is stored (the block's request answered at once);
+      # when the run stands at that await and the token is the one it will
+      # use, the step completes at once, without running the block.
+      # Returns `:resumed`, `:failed` when the completion took the failure
+      # path, or `:stale` (changing nothing) for any other event or token.
+      def resume!(event:, token:, payload:)
+        transition = claim_wait { resumable_transition(event.to_s, token.to_s) }
+        return :stale unless transition
+        return complete_batch_wait(transition, batch_results!(transition, payload), :resumed) if batch_waiting?
+
+        value = claimed_payload(transition, payload)
+        complete_claimed_wait(transition, :resumed) do
+          transition.child_runs? ? deliver_child_results(transition, value) : write_wait_payload(transition, value)
+          StepCompletion.instance_method(:complete_step).bind_call(self, transition, nil)
+        end
+      end
+
+      # Ends a wait whose deadline has passed at `now` (a Time, or an
+      # ISO 8601 String): a `delay` continues to its `to:`, an `await` takes
+      # `on_timeout`, or else its `on_failure` with Smith::WaitTimedOut, and
+      # a durable retry's wait ends with its step next (RetryWait). Returns
+      # `:expired`, `:not_due`, or `:stale` unless the run waits under this
+      # token.
+      def expire_wait!(token:, now: Time.now.utc)
+        now = wait_time(now)
+        record = @awaiting
+        not_due = false
+        transition = claim_wait do
+          next if record.nil? || !record.equal?(@awaiting) || record["token"] != token.to_s
+
+          not_due = now < wait_deadline(record)
+          awaited_transition(record) unless not_due
+        end
+        return not_due ? :not_due : :stale unless transition
+
+        expire_claimed_wait(transition, record)
+      end
+
+      private
+
+      # Runs the step's block, if any, then either finishes inline (the
+      # block called `resolve!`) or records the wait and returns
+      # WAIT_PENDING, which tells `advance!` the step has not completed. The
+      # payload key and the store's TTL are checked before the block sends
+      # anything.
+      def execute_await_step(transition)
+        check_deadline!
+        seconds = WaitRecord.seconds(transition.await_config.fetch(:timeout), snapshot_value(@context))
+        check_wait_admissible!(transition, seconds)
+        token = await_token_for(transition)
+        settled = settle_on_entry(transition, token)
+        return settled unless wait_pending?(settled)
+
+        enter_wait(transition, token, seconds)
+        WAIT_PENDING
+      end
+
+      # Runs the step's entry: an await's block, or a map's or call's
+      # launch. The step settles at once when the block resolves the wait or
+      # the launcher ran the children itself; otherwise WAIT_PENDING.
+      def settle_on_entry(transition, token)
+        if transition.child_runs?
+          results = launch_child_runs(transition, token)
+          return WAIT_PENDING unless results
+
+          payload = child_results_payload(collectable_results(transition.children_config, results))
+          return deliver_child_results(transition, payload, @launched_children)
+        end
+
+        step = run_await_block(transition, token) if transition.await_config[:block]
+        step&.resolved? ? write_wait_payload(transition, wait_payload(step.payload)) : WAIT_PENDING
+      end
+
+      def wait_pending?(value) = value.equal?(WAIT_PENDING)
+
+      # The delivered payload, normalized for the claimed step: a map's
+      # results checked against its children, anything else bounded like a
+      # session message. A payload that does not fit raises to the host and
+      # frees the wait.
+      def claimed_payload(transition, payload)
+        return wait_payload(payload) unless transition.child_runs?
+
+        entries = child_results_payload(collectable_results(transition.children_config, payload))
+        checked_child_results(entries, expected_children(transition), transition.children_config[:kind])
+      rescue StandardError
+        @wait_claimed = false
+        raise
+      end
+
+      # The wait's token: a key of the step's position under a segment only
+      # waits use, so no `idempotency_key` scope can produce it, and
+      # computable before the wait is stored. A wait is never nested (an
+      # in-process child cannot wait), so the position is the run's own.
+      def await_token_for(transition)
+        config = transition.await_config
+        wait_token_for(transition, [config.fetch(:kind), config[:event]].compact)
+      end
+
+      def wait_token_for(transition, scope)
+        IdempotencyKey.__send__(:root_scope, ArtifactIntegration.instance_method(:execution_namespace).bind_call(self))
+                      .within(:step, @step_count + 1, transition.name.to_s).within(:wait).key(scope)
+      end
+
+      def check_wait_admissible!(transition, seconds)
+        writable_wait_key!(transition) if transition.await_config[:into]
+        check_wait_ttl!(seconds)
+      end
+
+      def check_wait_ttl!(seconds)
+        ttl = self.class.persistence_ttl || Smith.config.persistence_ttl
+        return unless @persistence_key && ttl && seconds >= ttl
+
+        raise WorkflowError, "a wait of #{seconds} seconds does not end within the persistence TTL of #{ttl} " \
+                             "seconds; the stored run would expire before the wait ends"
+      end
+
+      def run_await_block(transition, token)
+        step = AwaitStep.new(context: snapshot_value(@context), await_token: token,
+                             idempotency_scope: IdempotencyKey.__send__(:current_scope!),
+                             attempt: @step_attempt || 1, transition_name: transition.name)
+        transition.await_config.fetch(:block).call(step)
+        step
+      end
+
+      def enter_wait(transition, token, seconds)
+        children = @launched_children if transition.child_runs?
+        @awaiting = WaitRecord.build(transition, token, seconds, Time.now.utc, children:)
+        record_wait_trace(transition, :waiting)
+      end
+
+      def record_wait_trace(transition, outcome)
+        Smith::Trace.record(
+          type: :transition, data: { transition: transition.name, from: transition.from, to: transition.to, outcome: }
+        )
+      end
+    end
+  end
+end

data/lib/smith/workflow/wait_record.rb

@@ -0,0 +1,157 @@
+# frozen_string_literal: true
+
+require "time"
+
+module Smith
+  class Workflow
+    # The persisted form of a run's wait: a frozen Hash of frozen Strings, so
+    # it round-trips through JSON unchanged and no caller can change it.
+    # Built on entry, checked on restore against the workflow class and the
+    # run's position, and read as a Wait.
+    module WaitRecord
+      FIELDS = %w[kind event token transition into entered_at deadline_at children].freeze
+      TOKEN = /\A[0-9a-f]{64}\z/
+      ZONED = /(?:Z|[+-]\d\d:\d\d)\z/
+
+      module_function
+
+      # `children` is how many runs a map or call launched; nil otherwise.
+      def build(transition, token, seconds, now, children: nil)
+        config = transition.await_config
+        frozen(timing(transition, token, seconds, now).merge(
+                 "kind" => config.fetch(:kind).to_s, "event" => config[:event]&.to_s, "into" => config[:into]&.to_s,
+                 "children" => children
+               ))
+      end
+
+      # A batch step's wait for the results of its `count` requests.
+      def build_batch(transition, token, seconds, now, count)
+        frozen(timing(transition, token, seconds, now)
+          .merge("kind" => "batch", "event" => Batch::EVENT, "into" => nil, "children" => count))
+      end
+
+      # A durable retry's wait: no event or context key.
+      def build_retry(transition, token, seconds, now)
+        frozen(timing(transition, token, seconds, now)
+          .merge("kind" => "retry", "event" => nil, "into" => nil, "children" => nil))
+      end
+
+      def timing(transition, token, seconds, now)
+        { "token" => token, "transition" => transition.name.to_s, "entered_at" => now.utc.iso8601(6),
+          "deadline_at" => (now + seconds).utc.iso8601(6) }
+      end
+
+      # A wait's timeout in seconds: the declared number, or what its
+      # callable returns for the context, which must be positive.
+      def seconds(timeout, context)
+        seconds = timeout.respond_to?(:call) ? timeout.call(context) : timeout
+        return seconds if (seconds.is_a?(Integer) || (seconds.is_a?(Float) && seconds.finite?)) && seconds.positive?
+
+        raise WorkflowError, "a wait's timeout must be a positive number of seconds"
+      end
+
+      # The event and context key read as the transition declares them; a
+      # retry's attempt and error read from the run's pending retry, a
+      # batch's attempt from its pending batch.
+      def to_wait(record, transition, pending_retry, pending_batch = nil)
+        config = wait_config(record, transition)
+        pending = { "retry" => pending_retry, "batch" => pending_batch }[record["kind"]]
+        Wait.new(
+          kind: config.fetch(:kind), event: config[:event], token: record.fetch("token"), transition: transition.name,
+          into: config[:into], entered_at: Time.iso8601(record.fetch("entered_at")),
+          deadline_at: Time.iso8601(record.fetch("deadline_at")), children: record["children"],
+          attempt: pending&.fetch("attempt"), error: pending&.fetch("error", nil)
+        )
+      end
+
+      def wait_config(record, transition)
+        { "retry" => { kind: :retry }, "batch" => { kind: :batch, event: Batch::EVENT.to_sym } }
+          .fetch(record["kind"]) { transition.await_config }
+      end
+
+      # A persisted wait, checked: a malformed one raises SerializationError,
+      # and one the class no longer declares exactly (the transition, its
+      # kind, event and context key) or that does not start from the run's
+      # state is refused like any unknown position.
+      def restore(value, workflow_class:, persistence_key:, state:, next_transition:)
+        return if value.nil?
+
+        record = stringified(value)
+        unless well_formed?(record) && next_transition.nil?
+          raise Smith::SerializationError, "persisted workflow awaiting is malformed"
+        end
+
+        transition = workflow_class.transition_named(record["transition"])
+        return frozen(record.slice(*FIELDS)) if declared_wait?(transition, record, state)
+
+        raise Smith::UnknownTransitionOnRestore.new(workflow: workflow_class.name, persistence_key:,
+                                                    transition: record["transition"].to_sym)
+      end
+
+      def stringified(value)
+        raise Smith::SerializationError, "persisted workflow awaiting is malformed" unless value.is_a?(Hash)
+
+        value.to_h { |key, field| [key.to_s, field.is_a?(Symbol) ? field.to_s : field] }
+      end
+
+      def well_formed?(record)
+        %w[await delay children retry batch].include?(record["kind"]) && TOKEN.match?(record["token"].to_s) &&
+          nonempty_string?(record["transition"]) && await_fields?(record) && times?(record) && children_count?(record)
+      end
+
+      # A map's or call's wait holds how many children it launched, a
+      # batch's how many requests it sent.
+      def children_count?(record)
+        count = record["children"]
+        case record["kind"]
+        when "batch" then count.is_a?(Integer) && count.between?(1, 10_000)
+        when "children" then count.is_a?(Integer) && count.between?(0, 10_000)
+        else count.nil?
+        end
+      end
+
+      # An await or a map's wait for its children names its event and
+      # context key; a batch its event alone; a delay or a retry neither.
+      def await_fields?(record)
+        event, into = record.values_at("event", "into")
+        case record["kind"]
+        when "batch" then event == Batch::EVENT && into.nil?
+        when "delay", "retry" then event.nil? && into.nil?
+        else nonempty_string?(event) && nonempty_string?(into)
+        end
+      end
+
+      # Both times carry their zone, and the deadline is not before entry.
+      def times?(record)
+        entered, deadline = record.values_at("entered_at", "deadline_at")
+        [entered, deadline].all? { |value| value.is_a?(String) && ZONED.match?(value) } &&
+          Time.iso8601(entered) <= Time.iso8601(deadline)
+      rescue ArgumentError
+        false
+      end
+
+      def declared_wait?(transition, record, state)
+        return false unless transition && [nil, state.to_s].include?(transition.from&.to_s)
+
+        declared_kind?(transition, record)
+      end
+
+      def declared_kind?(transition, record)
+        case record["kind"]
+        when "retry" then transition.durable_retry?
+        when "batch" then transition.batches? && record["children"] == batch_count(transition)
+        else transition.awaits? && declared_fields(transition.await_config) == record.values_at("kind", "event", "into")
+        end
+      end
+
+      def declared_fields(config) = [config.fetch(:kind), config[:event], config[:into]].map { _1&.to_s }
+
+      def batch_count(transition) = transition.fanout? ? transition.fanout_config.fetch(:branches).length : 1
+
+      def nonempty_string?(value) = value.is_a?(String) && !value.empty?
+
+      def frozen(record) = record.transform_values { |value| value&.dup&.freeze }.freeze
+    end
+    private_constant :WaitRecord
+  end
+end

data/lib/smith/workflow.rb

@@ -28,6 +28,19 @@ module Smith
     include ArtifactIntegration
     include DataVolumePolicy
     include DeadlineEnforcement
+    include ActiveClock
+    include StepIO
+    include WaitCompletion
+    include WaitExecution
+    include RetryWait
+    include BatchExecution
+    include BatchWait
+    include BatchCharges
+    include BatchHolds
+    include BudgetPause
+    include RememberedCalls
+    include ChildRunResults
+    include ChildRunExecution
     include Execution
 
     DEFAULT_MAX_TRANSITIONS = 100
@@ -40,8 +53,21 @@ module Smith
 
     def initialize(context: {}, ledger: nil, created_at: nil)
       @state = self.class.initial_state
-      @context = context
+      # Symbol keys at every depth, the shape a restore produces, so a
+      # context reads the same live and after a checkpoint round trip.
+      @context = ContextKeys.normalize_context(context)
       @step_count = 0
+      # The wait the run is in (a persisted Hash of Strings), and the seconds
+      # its steps have spent running, which `wall_clock` budgets bound.
+      @awaiting = nil
+      @wait_mutex = Mutex.new
+      @wait_claimed = false
+      @active_seconds = 0.0
+      @active_step_started_at = nil
+      # The run's own identity, persisted with it: artifacts are scoped to
+      # it and idempotency keys derive from it, so it exists from the start
+      # and is in the checkpoint written before the first step.
+      @execution_namespace = SecureRandom.uuid
       @next_transition_name = nil
       @ledger = ledger || build_ledger
       @created_at = created_at || Time.now.utc.iso8601(6)
@@ -58,6 +84,9 @@ module Smith
       @usage_entries = []
       @usage_mutex = Mutex.new
       @last_output = nil
+      # The last output the run added to its transcript, which the next
+      # agent step reads as a handoff; an `output_to:` step's is not one.
+      @last_transcript_output = nil
       # Durable attribution of the most recent serial `execute :agent` step:
       # { model:, provider: } of the model that actually served it (post
       # fallback resolution), exposed to the following deterministic step as
@@ -98,8 +127,17 @@ module Smith
       # also: declaration so explicit input keys round-trip.
       @persisted_keys = ::Set.new(initial_persist_auto_seed)
       @persisted_keys_mutex = Mutex.new
+      # A nested workflow's parent's untrusted keys, in force before its
+      # seed builder runs (NestedExecution).
+      @inherited_untrusted_keys = StepIO.take_inherited_untrusted_keys
+      # A child run's parent's artifact namespaces, which it may read.
+      @inherited_artifact_namespaces = ArtifactIntegration.take_inherited_artifact_namespaces
       initialize_tool_result_state
       seed_initial_session_messages
+      initialize_budget_ledger(Budget::HostLedger.take_inherited)
+      # The attempts of the runs a nested workflow runs inside (ResponseMemo).
+      @memo_attempts = ResponseMemo.take_inherited
+      @budget_ledger_scope = capture_budget_ledger_scope
     end
 
     def persisted_keys
@@ -107,9 +145,12 @@ module Smith
     end
 
     def advance!
+      StepStream.instance_method(:deliver_streamed_output).bind_call(self)
+      return if waiting?
+
+      @step_work_started = false
       advance_claim = SplitStepPersistence.instance_method(:claim_split_step_advance!).bind_call(self)
       ensure_transition_budget!
-      @step_work_started = false
 
       transition = SplitStepPersistence.instance_method(:resolve_split_step_advance_transition).bind_call(self)
       transition = resolve_transition if transition.equal?(SplitStepPersistence::NO_SPLIT_TRANSITION)
@@ -121,10 +162,7 @@ module Smith
                     else
                       execute_step(transition)
                     end
-      @step_count += 1
-      @updated_at = Time.now.utc.iso8601(6)
-      record_step_snapshot(step_result)
-      step_result
+      count_step(step_result)
     rescue UnresolvedTransitionError => e
       step_result = GuardrailIntegration
                     .instance_method(:handle_unresolved_transition_failure)
@@ -132,25 +170,31 @@ module Smith
       record_step_snapshot(step_result)
       step_result
     ensure
+      clear_settled_pending_retry
       SplitStepPersistence.instance_method(:release_split_step_advance!).bind_call(self, advance_claim) if advance_claim
     end
 
     def run!
       SplitStepPersistence.instance_method(:ensure_split_step_execution_allowed!).bind_call(self)
       steps = []
-      until terminal?
+      until nothing_to_advance?
         step = advance!
         steps << step if step
       end
+      StepStream.instance_method(:deliver_streamed_output).bind_call(self)
       build_run_result(steps)
     end
 
+    # The awaited transition while the run waits.
     def pending_transition_name
+      return awaiting.transition if waiting?
+
       @next_transition_name || self.class.first_transition_from(@state)&.name
     end
 
+    # A waiting run is never terminal, whatever its state.
     def terminal?
-      !self.class.transition_from?(@state) && @next_transition_name.nil?
+      !waiting? && !self.class.transition_from?(@state) && @next_transition_name.nil?
     end
 
     def done?
@@ -181,6 +225,18 @@ module Smith
 
     private
 
+    # Entering a wait is not a step: nothing counts, and advance! returns nil.
+    def count_step(step_result)
+      @updated_at = Time.now.utc.iso8601(6)
+      return if wait_pending?(step_result)
+
+      @step_count += 1
+      record_step_snapshot(step_result)
+      step_result
+    end
+
+    def nothing_to_advance? = terminal? || waiting?
+
     def initial_persist_auto_seed
       manager = self.class.context_manager
       return [] unless manager.respond_to?(:persist_mode) && manager.persist_mode == :auto
@@ -218,6 +274,38 @@ module Smith
       end
     end
 
+    # A run's budget ledger is built at its first use; a nested workflow
+    # without one of its own charges `inherited`, the enclosing run's.
+    def initialize_budget_ledger(inherited = nil)
+      @inherited_budget_ledger = inherited
+      @budget_ledger = nil
+      @budget_ledger_built = false
+      @budget_ledger_mutex = Mutex.new
+    end
+
+    # The scope the run's ledger is built from, from its trusted context, so
+    # no untrusted key chooses whose budget is charged.
+    def capture_budget_ledger_scope
+      return unless self.class.budget_ledger
+
+      Budget::HostLedger.scope!(self.class.budget_ledger_scope_keys, trusted_context_view, untrusted_context_keys)
+    end
+
+    # The host's budget ledger this run's requests are held on: its own,
+    # built at first use from the scope it was created with, else the
+    # enclosing run's (a nested workflow's); nil for none.
+    def run_budget_ledger
+      return @inherited_budget_ledger&.call unless self.class.budget_ledger
+
+      @budget_ledger_mutex.synchronize do
+        unless @budget_ledger_built
+          @budget_ledger = Budget::HostLedger.build(self.class.budget_ledger, @budget_ledger_scope)
+          @budget_ledger_built = true
+        end
+        @budget_ledger
+      end
+    end
+
     def build_ledger
       config = self.class.budget
       return nil unless config
@@ -313,14 +401,14 @@ module Smith
     end
 
     # Fresh evaluation of the seed builder against the workflow's
-    # current @context. Returns nil when no builder is defined so
-    # callers (initialize, validate_seed_digest!) can distinguish
-    # "no builder" from "builder returned empty".
+    # current @context, without its untrusted keys. Returns nil when no
+    # builder is defined so callers (initialize, validate_seed_digest!)
+    # can distinguish "no builder" from "builder returned empty".
     def compute_seed_messages
       builder = self.class.seed_messages
       return nil unless builder
 
-      seeded = builder.arity == 1 ? builder.call(@context) : builder.call
+      seeded = builder.arity == 1 ? builder.call(trusted_context_view) : builder.call
       normalize_seed_messages(seeded)
     end
 
@@ -335,11 +423,12 @@ module Smith
     end
 
     def normalize_seed_messages(seeded)
-      return [] if seeded.nil?
-      return [seeded] if seeded.is_a?(Hash)
-      return seeded.to_a if seeded.respond_to?(:to_a)
-
-      raise WorkflowError, "seed_messages must return a message Hash or an Array of message Hashes"
+      messages = if seeded.nil? then []
+                 elsif seeded.is_a?(Hash) then [seeded]
+                 elsif seeded.respond_to?(:to_a) then seeded.to_a
+                 else raise WorkflowError, "seed_messages must return a message Hash or an Array of message Hashes"
+                 end
+      messages.map { |message| Message.parts?(message) ? Message.check_seed!(message) : message }
     end
 
     def snapshot_context
@@ -417,6 +506,13 @@ Smith::Workflow::SplitStepPersistence::SubclassBoundary.protect_execution_path!(
   Smith::Workflow::ArtifactIntegration,
   Smith::Workflow::DataVolumePolicy,
   Smith::Workflow::DeadlineEnforcement,
+  Smith::Workflow::ActiveClock,
+  Smith::Workflow::StepIO,
+  Smith::Workflow::WaitCompletion,
+  Smith::Workflow::ChildRunResults,
+  Smith::Workflow::ChildRunExecution,
+  Smith::Workflow::WaitExecution,
+  Smith::Workflow::RetryWait,
   Smith::Workflow::SplitStepPersistence::DefinitionBoundary,
   Smith::Workflow::SplitStepPersistence::ExecutionVerification,
   Smith::Workflow::SplitStepPersistence::ExecutionAuthorization,