smith-agents 0.4.3 → 0.4.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8b8a1ad3fc2cb40ef0287063f81e0bc01a6bca349b1625c4173daef763785202
-  data.tar.gz: b870358edea9686d15e164bb3d44e00edf8e50da8cb2d68b8a6a666bbc5d0f79
+  metadata.gz: 8a4cd092162200b5cb861d9b9cd48069eab2e00beef71abf5fbfb91705c235ee
+  data.tar.gz: 483180aec2d6b6d63138324edfb10f564ca2c4ff800603eeb3a43807f4b20f95
 SHA512:
-  metadata.gz: bd9dd710fb5008efc7896e950b52621d3c80dab9be902b323efdc76c3d4d12e4b1c066b2219e5e36c56edea5652add7406ecb1d740db79b320035eee9047cee6
-  data.tar.gz: 4d5e6b412d39269276096b697a399ad476dd1d4bc370792d8acfdd8f49041b567da2120e56a0fc1e090cbf4c351c9849201c8b0030fd5518f6c69495b175f534
+  metadata.gz: 240e30a347d30d06063e57bece46ef646ce68c31463192a123cffd47858311ee1a92d6f7c664c519347775ff5df861e63ae8103005806cf39e400826b3e04c75
+  data.tar.gz: e478845f5665afeea75fb2b905dd4a7c48aa4ff56ca748661a1dc8302f44dacb029fa60117d900acfeca5558a52a9fde5bc8c7552a7a8149808b92b4898d1c82

data/CHANGELOG.md

@@ -6,6 +6,93 @@ Format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). Version
 
 ## [Unreleased]
 
+## [0.4.5] - 2026-07-11
+
+Patch release for generic host-coordinated workflow step boundaries and
+fail-closed persistence correctness. Smith now exposes a bounded strict
+split-step protocol while leaving transactions, scheduling, lifecycle records,
+tools, and product policy under host ownership.
+
+### Fixed
+
+- Align `ActiveRecordStore#store_versioned` with Smith's persistence contract by
+  comparing `expected_version` to the stored payload's `persistence_version`.
+  Rails' optimistic-locking column remains an independent row-level CAS token,
+  so consecutive workflow persists no longer report a false conflict after the
+  initial insert.
+- Fail closed when an Active Record host model does not have optimistic locking
+  enabled on the adapter's configured `version_column`. Custom locking columns
+  remain host-owned and must be configured on the model explicitly.
+- Use Rails' native create-or-find savepoint for concurrent initial inserts,
+  preserving callback rollbacks and distinguishing key collisions from other
+  unique-constraint failures.
+- Keep malformed payload-version handling consistent across versioned adapters,
+  reject scalar state documents, fail closed when an explicit persisted version
+  is invalid, and never replay an uncertain versioned write below a host-owned
+  transaction boundary.
+- Resolve string-backed Active Record models on each operation so host framework
+  reloads cannot leave the adapter holding a stale class object.
+- Add a generic strict split-step persistence contract for hosts that coordinate
+  pre/post transition state without holding a transaction across provider or
+  tool execution. Mutable execution state and non-expiring persistence policy
+  are pinned for the boundary, subclass entry points remain guarded, and a
+  proven transaction rollback can retry from the exact committed preparation.
+  Same-transaction atomicity remains adapter- and host-owned.
+- Prevent Memory, Redis, and Active Record versioned adapters from recreating a
+  missing key when the caller expects a nonzero logical version. Missing state
+  now reports `PersistenceVersionConflict` with `actual: :missing`.
+- Bound split-step transition contract capture to cycle-aware `O(V + E)` traversal with explicit
+  node, byte, and depth limits; reject opaque mutable values; freeze supported
+  structured configuration; and preserve execution guards when workflow
+  subclasses receive later prepends.
+- Keep prepared transition execution on Smith's owned `advance!` path so host
+  wrappers cannot run as transition authority inside an active boundary.
+- Make the split-step aggregate own its internal require order so direct loading
+  does not depend on `smith.rb` preloading implementation files.
+- Require reconciliation before retrying an ambiguously acknowledged
+  checkpoint and retain a single checkpoint witness, keeping retry state in
+  constant space.
+- Make Memory expiry atomic with version comparison, isolate mutable payload
+  strings at its boundary, and pin Active Record column configuration.
+
+### Verification
+
+- Default suite: 1,045 examples, 0 failures.
+- Focused split-step and versioned-adapter suite: 107 examples, 0 failures;
+  changed files pass RuboCop and `git diff --check`.
+- Practical gem execution: 30 distinct 20-step workflow classes, 600 Memory
+  split steps, 1,000 Memory compare-and-swap writes, and 200 Active Record
+  split steps with restore after every committed checkpoint.
+- Smith Runtime host acceptance on Ruby 4.0.1 and Rails 8.1.3: 251 tests,
+  816 assertions, 0 failures; 20 practical signed-package compiles produced
+  valid Smith reports and cleaned every generated namespace.
+
+## [0.4.4] - 2026-07-10
+
+Patch release for provider-safe workflow handoffs. Smith keeps accepted agent
+outputs in durable session history while adapting only the next provider call
+when a completed workflow stage would otherwise look like an unsupported
+assistant prefill.
+
+### Fixed
+
+- Adapt workflow-prepared provider input that ends with an accepted assistant
+  result by adding a non-persisted user continuation for the next agent call.
+  This preserves Smith session history while avoiding unsupported assistant
+  prefilling on provider models that require a user turn before completion.
+  Provider preparation now also reads string-keyed roles and content restored
+  through JSON host persistence. Explicit assistant-prefill seed messages remain
+  unchanged unless they match Smith's recorded accepted workflow output.
+
+### Verification
+
+- Default suite: 932 examples, 0 failures.
+- Practical gem-level JSON persistence/restore probe with a 25-branch parallel
+  handoff, provider-safe message ordering, non-persisted continuation, and
+  explicit assistant-prefill preservation.
+- Smith Runtime host verification on Ruby 4.0.1 and Rails 8.1.3: 139 tests,
+  391 assertions, 0 failures, plus a process-level restored workflow run.
+
 ## [0.4.3] - 2026-07-05
 
 ### Documentation

data/README.md

@@ -17,7 +17,7 @@ environment before calling the slice complete.
 
 ```ruby
 # Gemfile
-gem "smith-agents", "~> 0.4.3", require: "smith"
+gem "smith-agents", "0.4.5", require: "smith"
 ```
 
 ```bash
@@ -211,7 +211,8 @@ Built-in adapters (all support TTL where the backend allows; `Redis`,
 - `:redis` — Redis client; uses WATCH/MULTI/EXEC for CAS
 - `:rails_cache`, `:solid_cache` — Rails cache backends
 - `:cache_store` — any object responding to `write/read/delete`
-- `:active_record` — keyed ActiveRecord model with `lock_version` column for CAS
+- `:active_record` — keyed ActiveRecord model with an enabled optimistic-locking
+  column for row-level CAS; Smith's logical version remains in the JSON payload
 
 See [`docs/PERSISTENCE.md`](docs/PERSISTENCE.md) for schema versioning, seed-drift validation, and the `idempotency_mode :strict` step-in-progress contract.
 
@@ -297,6 +298,11 @@ and external idempotency guarantees remain host-owned.
 
 ## Development
 
+Smith changes must not use brute-force graph search, validation, scheduling, or
+class-generation algorithms. Document the relevant time and space complexity,
+prefer indexed, linear, logarithmic, or otherwise bounded approaches, and treat
+an avoidable asymptotic regression as a release blocker.
+
 ```bash
 bundle install
 bundle exec rspec

data/docs/PERSISTENCE.md

@@ -120,6 +120,186 @@ workflow.clear_persisted!("ticket:T-1042")
 
 These helpers do not make Smith a job system or durable runtime. They only remove repetitive restore/checkpoint boilerplate around the configured persistence adapter while leaving queueing, projection, and recovery policy with the host app.
 
+### Host-Coordinated Step Boundaries
+
+Hosts can split one strict transition into explicit prepare, execute, and
+checkpoint phases:
+
+```ruby
+transition_name = nil
+
+ApplicationRecord.transaction do
+  transition_name = workflow.prepare_persisted_step!(key, adapter: adapter)
+  HostEvent.create!(name: "step.started", transition: transition_name)
+end
+
+workflow.confirm_prepared_step!
+step = workflow.execute_prepared_step! # provider/tool work; no host transaction
+
+ApplicationRecord.transaction do
+  workflow.persist!(key, adapter: adapter)
+  HostEvent.create!(name: "step.completed", transition: step.fetch(:transition))
+end
+
+workflow.complete_persisted_step!
+```
+
+This contract is available only to workflows using `idempotency_mode :strict`,
+a versioned persistence adapter whose `store_versioned` method accepts `ttl:`,
+and non-expiring workflow persistence. Smith explicitly pins `ttl: nil` for
+every boundary write, so a later global configuration change cannot turn an
+accepted checkpoint into expiring state.
+`prepare_persisted_step!` writes the normal `step_in_progress` marker and returns
+the pending transition name without consuming it. Duplicate preparation on the
+same workflow object is rejected. The boundary is pinned to the exact
+transition object, persistence key, and adapter instance selected during
+preparation. The key is copied into an immutable string and the transition
+contract is frozen. When an adapter can report an open transaction,
+`confirm_prepared_step!` must verify the committed preparation before execution;
+it refuses confirmation while that transaction is still open.
+`execute_prepared_step!` re-verifies the exact durable preparation and permits
+one execution attempt. Preparation also takes an O(S) defensive snapshot of
+mutable workflow execution state, where S is the serialized state size. This
+is the necessary ownership boundary: aliases held before preparation cannot
+change what the accepted transition later consumes, and public mutable-state
+readers and `to_state` return defensive snapshots while the boundary is active.
+Subclass execution entry points remain guarded, and subclass TTL helpers cannot
+override the pinned boundary policy. The
+in-memory marker remains armed so any serialization
+before committed completion still fails closed. Other workflow execution and
+checkpoint APIs are rejected until the host checkpoints the accepted state
+through `persist!`. After that transaction commits,
+`complete_persisted_step!` verifies the exact checkpoint through the adapter
+before clearing the marker and releasing the process-local boundary, and
+likewise refuses to run while the adapter reports an open transaction.
+`prepared_persisted_step?` exposes whether the one execution attempt remains
+available without revealing transition internals.
+
+Workflow classes are part of the execution contract once preparation begins.
+Hosts must not add or prepend methods to that class until the boundary is
+complete, and custom `inherited` hooks must call `super`. Ruby intentionally
+allows open-class mutation, so Smith treats post-preparation class mutation as
+an unsupported host lifecycle violation rather than attempting to sandbox it.
+
+The lifecycle row and Smith checkpoint are atomic only when the persistence
+adapter participates in the same transaction and database connection domain as
+the host record. `ActiveRecordStore` can provide that property when its model
+uses the same database connection; it also reports transaction state so Smith
+cannot confirm an uncommitted boundary. Memory, Redis, cache-backed,
+cross-database, and external adapters do not participate in
+`ApplicationRecord.transaction`; hosts using them must supply their own
+coordination protocol and must not claim the two writes are atomic.
+Custom transactional adapters should implement `transaction_open?` so Smith can
+enforce the same commit-aware confirmation rule.
+
+If execution raises or the process dies after external work, the durable marker
+remains set. A strict restore therefore fails closed with
+`StepInProgressOnRestore`; the host must reconcile operation results or classify
+the run as uncertain rather than blindly replaying the transition. The same
+in-memory workflow object cannot retry an attempted transition. If a host
+transaction rolls back after a successful `persist!`,
+`complete_persisted_step!` rejects the rolled-back checkpoint. When the adapter
+still contains the exact committed preparation, Smith restores that known
+version and permits the same in-memory object to retry the unchanged checkpoint.
+Any different or missing payload remains uncertain and non-retryable.
+
+Ambiguous persistence acknowledgements also fail closed. If preparation may
+have written before raising, the object cannot execute. If a post-step
+checkpoint may have written before raising, `complete_persisted_step!` can
+verify the exact attempted payload without replaying the write. The host must
+perform that reconciliation before retrying. If the adapter still contains the
+exact preparation, Smith permits one unchanged retry; any other result remains
+non-retryable. Smith retains one checkpoint witness, so retry bookkeeping is
+constant space rather than growing with the number of attempts.
+
+Transition contract capture uses cycle-aware `O(V + E)` traversal, with time
+and space bounded to 10,000 visits, 4 MiB of string data, and a maximum depth
+of 128. Strings are represented by SHA-256 digests rather than duplicated.
+Supported mutable data shapes are frozen, `Range` endpoints are traversed, and
+opaque mutable objects fail closed instead of receiving a misleading shallow
+signature. Callable identity is pinned for the boundary; Ruby closure state and
+method bodies remain host lifecycle concerns. Smith detects contract
+replacement but does not claim to sandbox Ruby's open object model. Subclass
+entry points remain guarded even when a host prepends modules after class
+creation, and prepared execution bypasses host `advance!` wrappers so they
+cannot become transition execution authority.
+
+## Active Record Optimistic Locking
+
+`ActiveRecordStore` keeps two version domains deliberately separate:
+
+- `persistence_version` inside the JSON payload is Smith's logical workflow
+  version and is compared with `expected_version`.
+- the host model's Active Record `locking_column` is Rails' row-level
+  compare-and-swap token and detects concurrent updates between load and save.
+
+Across every versioned adapter, a missing key can be created only with
+`expected_version: 0`. A nonzero expected version against missing state raises
+`PersistenceVersionConflict` with `actual: :missing`, preventing stale workflow
+objects from resurrecting host-deleted state.
+
+Legacy object-shaped payloads without `persistence_version` use version zero.
+Valid JSON scalars are not workflow-state documents, and explicit null,
+negative, or non-integer versions fail closed instead of being overwritten.
+
+The default host model needs a unique key, a JSON/text payload, and Rails'
+standard optimistic-locking column:
+
+```ruby
+create_table :workflow_states do |t|
+  t.string :key, null: false, index: { unique: true }
+  t.jsonb :payload, null: false
+  t.integer :lock_version, null: false, default: 0
+  t.timestamps
+end
+
+class WorkflowState < ApplicationRecord
+end
+
+Smith.configure do |config|
+  config.persistence_adapter = :active_record
+  config.persistence_options = { model: WorkflowState }
+end
+```
+
+For a custom locking column, configure both the host model and adapter. Smith
+validates this relationship but never mutates the host model:
+
+```ruby
+add_column :workflow_states, :workflow_revision, :integer,
+  null: false,
+  default: 0
+
+class WorkflowState < ApplicationRecord
+  self.locking_column = :workflow_revision
+end
+
+Smith.configure do |config|
+  config.persistence_adapter = :active_record
+  config.persistence_options = {
+    model: WorkflowState,
+    version_column: :workflow_revision
+  }
+end
+```
+
+The unique key must be enforced by the database. Initial creation uses Rails'
+native `create_or_find_by!` savepoint path, so a concurrent key insert does not
+invalidate a PostgreSQL caller transaction. A collision on another unique
+constraint is not misreported as a key-version conflict.
+
+Versioned writes are single-attempt. Smith translates known connection failures
+to `PersistenceIOError`, but never replays an uncertain write below the
+host-owned transaction boundary. The host must restore and reconcile persisted
+state before deciding whether to retry. Keep the host model persistence-focused:
+callbacks that raise after commit run after the write is permanent and therefore
+create an inherently uncertain outcome.
+
+The adapter participates in the caller's Active Record transaction and never
+commits it. Hosts remain responsible for outer transaction boundaries, queue
+claims, fencing, reconciliation, and recovery policy. See the
+[official Rails optimistic-locking documentation](https://api.rubyonrails.org/classes/ActiveRecord/Locking/Optimistic.html).
+
 ## Artifacts
 
 Use artifacts when outputs are too large to keep inline.
@@ -166,4 +346,3 @@ Why this matters:
 - large payloads can move out of the inline workflow result
 - refs are execution-scoped
 - nested workflows inherit artifact scope correctly
-

data/docs/workflow_claim.md

@@ -55,4 +55,7 @@ If the model has AASM and you want to skip events, call `.cas` explicitly; if yo
 
 ## Testing
 
-Specs that exercise the AR strategies are tagged `:ar` and excluded from the default suite. The full suite under `SMITH_AR_SPECS=1` boots an in-memory SQLite database and the `ClaimableRecord` fixture model. Default `bundle exec rspec` runs 816 examples (Claim load-hygiene only); `SMITH_AR_SPECS=1 bundle exec rspec` runs 831 (adds 15 AR-tagged Claim specs).
+The default spec suite boots an in-memory SQLite database and runs the tagged
+Active Record strategy examples against the `ClaimableRecord` fixture model.
+Package smoke tests separately verify that requiring Smith does not make Active
+Record a runtime dependency.

data/lib/smith/agent/lifecycle.rb

@@ -3,6 +3,10 @@
 module Smith
   class Agent
     module Lifecycle
+      WORKFLOW_CONTINUATION_MESSAGE =
+        "Use the preceding assistant result as input and perform your assigned workflow step."
+      private_constant :WORKFLOW_CONTINUATION_MESSAGE
+
       TRANSIENT_ERRORS = [
         RubyLLM::ServerError, RubyLLM::ServiceUnavailableError,
         RubyLLM::OverloadedError, RubyLLM::RateLimitError
@@ -111,6 +115,7 @@ module Smith
       def add_prepared_input(chat, prepared_input)
         return unless prepared_input
 
+        prepared_input = provider_safe_prepared_input(prepared_input)
         system_messages, other_messages = prepared_input.partition do |message|
           message_role(message) == :system
         end
@@ -119,6 +124,22 @@ module Smith
         other_messages.each { |message| chat.add_message(message) }
       end
 
+      def provider_safe_prepared_input(prepared_input)
+        messages = prepared_input.to_a
+        return messages unless workflow_handoff?(messages)
+
+        messages + [{ role: :user, content: WORKFLOW_CONTINUATION_MESSAGE }]
+      end
+
+      def workflow_handoff?(messages)
+        message = messages.reverse_each.find { |candidate| message_role(candidate) != :system }
+        return false unless message
+        return false unless message_role(message) == :assistant
+        return false unless defined?(@last_output) && !@last_output.nil?
+
+        message_content(message) == @last_output
+      end
+
       def merge_system_messages!(chat, prepared_system_messages)
         return prepared_system_messages.each { |message| chat.add_message(message) } unless chat.respond_to?(:messages)
 
@@ -143,19 +164,18 @@ module Smith
       end
 
       def message_role(message)
-        if message.respond_to?(:role)
-          message.role&.to_sym
-        else
-          message[:role]&.to_sym
-        end
+        message_attribute(message, :role)&.to_sym
       end
 
       def message_content(message)
-        if message.respond_to?(:content)
-          message.content
-        else
-          message[:content]
-        end
+        message_attribute(message, :content)
+      end
+
+      def message_attribute(message, name)
+        return message.public_send(name) if message.respond_to?(name)
+        return message[name] if message.respond_to?(:key?) && message.key?(name)
+
+        message[name.to_s]
       end
 
       def fallback_eligible?(error)

data/lib/smith/doctor/checks/persistence_capabilities.rb

@@ -58,6 +58,7 @@ module Smith
             message: "#{adapter.class.name} missing optional capabilities: #{missing.join(", ")}",
             detail: "Smith will fall back where possible: non-versioned writes when store_versioned " \
                     "is missing, and payload updated_at parsing when heartbeat methods are missing. " \
+                    "Without transaction_open?, commit-aware split-step confirmation is unavailable. " \
                     "Use RedisStore or Memory for full versioning and heartbeat coverage; " \
                     "ActiveRecordStore currently covers optimistic locking when lock_version is present."
           )

data/lib/smith/persistence_adapters/active_record_connection_errors.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Smith
+  module PersistenceAdapters
+    module ActiveRecordConnectionErrors
+      NAMES = %w[
+        ActiveRecord::ConnectionNotEstablished
+        ActiveRecord::ConnectionFailed
+        ActiveRecord::AdapterTimeout
+      ].freeze
+
+      module_function
+
+      def classes
+        NAMES.filter_map do |name|
+          Object.const_get(name)
+        rescue NameError
+          nil
+        end
+      end
+    end
+  end
+end

data/lib/smith/persistence_adapters/active_record_initial_write.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Smith
+  module PersistenceAdapters
+    module ActiveRecordInitialWrite
+      module_function
+
+      def call(model:, key_column:, payload_column:, key:, payload:)
+        candidate = nil
+        record = model.create_or_find_by!(key_column => key) do |new_record|
+          candidate = new_record
+          new_record.public_send(:"#{payload_column}=", payload)
+        end
+
+        created_record?(record, candidate)
+      rescue ::ActiveRecord::RecordInvalid => e
+        raise unless key_collision?(e.record, candidate, model, key_column, key)
+
+        false
+      rescue ::ActiveRecord::RecordNotFound => e
+        raise e.cause if e.cause.is_a?(::ActiveRecord::RecordNotUnique)
+
+        raise
+      end
+
+      def created_record?(record, candidate)
+        return false unless record.equal?(candidate)
+        return true if record.persisted?
+
+        raise ::ActiveRecord::RecordNotSaved, "Active Record rolled back the workflow-state insert"
+      end
+
+      def key_collision?(record, candidate, model, key_column, key)
+        record.equal?(candidate) &&
+          record.errors.of_kind?(key_column, :taken) &&
+          model.exists?(key_column => key)
+      end
+
+      private_class_method :created_record?, :key_collision?
+    end
+  end
+end

data/lib/smith/persistence_adapters/active_record_store.rb

@@ -6,32 +6,18 @@ module Smith
       # AR transient errors resolved via class-name guard so Smith
       # doesn't require activerecord at load time. Hosts that use this
       # adapter already have activerecord in their dep tree.
-      TRANSIENT_ERROR_NAMES = %w[
-        ActiveRecord::ConnectionNotEstablished
-        ActiveRecord::StatementInvalid
-        ActiveRecord::TransactionIsolationConflict
-      ].freeze
-
-      def self.transient_errors
-        TRANSIENT_ERROR_NAMES.filter_map do |name|
-          Object.const_get(name)
-        rescue NameError
-          nil
-        end
-      end
-
       def initialize(model:, key_column: :key, payload_column: :payload, version_column: :lock_version)
-        @model_source = model
-        @key_column = key_column
-        @payload_column = payload_column
-        @version_column = version_column
+        @model_source = model.is_a?(String) ? model.dup.freeze : model
+        @key_column = normalize_column(key_column)
+        @payload_column = normalize_column(payload_column)
+        @version_column = normalize_column(version_column)
       end
 
       def store(key, payload, ttl: nil) # rubocop:disable Lint/UnusedMethodArgument
         # TTL is deferred for ActiveRecordStore — would require an
         # `expires_at` column + a periodic sweeper job. Ignored here;
         # documented as a known limitation.
-        Retry.with_retries(operation: :store, transient: self.class.transient_errors) do
+        Retry.with_retries(operation: :store, transient: ActiveRecordConnectionErrors.classes) do
           record = model_class.find_or_initialize_by(@key_column => key)
           record.public_send(:"#{@payload_column}=", payload)
           record.save!
@@ -39,60 +25,114 @@ module Smith
       end
 
       def fetch(key)
-        Retry.with_retries(operation: :fetch, transient: self.class.transient_errors) do
+        Retry.with_retries(operation: :fetch, transient: ActiveRecordConnectionErrors.classes) do
           model_class.find_by(@key_column => key)&.public_send(@payload_column)
         end
       end
 
       def delete(key)
-        Retry.with_retries(operation: :delete, transient: self.class.transient_errors) do
+        Retry.with_retries(operation: :delete, transient: ActiveRecordConnectionErrors.classes) do
           model_class.where(@key_column => key).delete_all
         end
       end
 
+      def transaction_open?
+        model_class.connection.transaction_open?
+      end
+
       # Optimistic locking via Rails' built-in optimistic locking on the
       # `lock_version` column. Requires the AR model to have a
       # `lock_version` (or configured) integer column with default 0.
       # If absent, raises ArgumentError directing the host to migrate.
       def store_versioned(key, payload, expected_version:, ttl: nil) # rubocop:disable Lint/UnusedMethodArgument
-        unless model_class.column_names.include?(@version_column.to_s)
-          raise ArgumentError,
-                "ActiveRecordStore#store_versioned requires a #{@version_column} column on " \
-                "#{model_class.name}. Add via: " \
-                "add_column :#{model_class.table_name}, :#{@version_column}, :integer, default: 0"
-        end
+        ensure_version_column!
+        ensure_locking_configuration!
 
-        Retry.with_retries(operation: :store_versioned, transient: self.class.transient_errors) do
-          record = model_class.find_or_initialize_by(@key_column => key)
-          if record.persisted? && record.public_send(@version_column) != expected_version
-            raise Smith::PersistenceVersionConflict.new(
-              key: key, expected: expected_version, actual: record.public_send(@version_column)
-            )
-          end
-          record.public_send(:"#{@payload_column}=", payload)
-          record.save!
-        rescue defined?(::ActiveRecord::StaleObjectError) ? ::ActiveRecord::StaleObjectError : StandardError => e
-          raise unless defined?(::ActiveRecord::StaleObjectError) && e.is_a?(::ActiveRecord::StaleObjectError)
+        write_versioned(key, payload, expected_version)
+      rescue *ActiveRecordConnectionErrors.classes => e
+        raise Smith::PersistenceIOError.new(operation: :store_versioned, cause: e)
+      end
+
+      private
 
-          raise Smith::PersistenceVersionConflict.new(
-            key: key, expected: expected_version, actual: :concurrent
-          )
+      def normalize_column(column)
+        column.is_a?(String) ? column.dup.freeze : column
+      end
+
+      def ensure_version_column!
+        return if model_class.column_names.include?(@version_column.to_s)
+
+        raise ArgumentError,
+              "ActiveRecordStore#store_versioned requires a #{@version_column} column on " \
+              "#{model_class.name}. Add via: " \
+              "add_column :#{model_class.table_name}, :#{@version_column}, :integer, default: 0"
+      end
+
+      def ensure_locking_configuration!
+        locking_column = model_class.locking_column if model_class.respond_to?(:locking_column)
+        locking_enabled = model_class.locking_enabled? if model_class.respond_to?(:locking_enabled?)
+        return if locking_enabled && locking_column.to_s == @version_column.to_s
+
+        raise ArgumentError,
+              "ActiveRecordStore#store_versioned requires #{model_class.name}.locking_column " \
+              "to be #{@version_column.inspect} with optimistic locking enabled"
+      end
+
+      def write_versioned(key, payload, expected_version)
+        record = model_class.find_by(@key_column => key)
+        unless record
+          VersionExpectation.validate_missing!(key, expected_version)
+          return create_versioned_record(key, payload, expected_version)
         end
+
+        validate_payload_version!(record, key, expected_version)
+        record.public_send(:"#{@payload_column}=", payload)
+        record.save!
+      rescue ::ActiveRecord::StaleObjectError => e
+        raise unless e.record.equal?(record)
+
+        raise_concurrent_conflict(key, expected_version)
       end
 
-      private
+      def create_versioned_record(key, payload, expected_version)
+        created = ActiveRecordInitialWrite.call(
+          model: model_class,
+          key_column: @key_column,
+          payload_column: @payload_column,
+          key: key,
+          payload: payload
+        )
+        return true if created
+
+        raise_concurrent_conflict(key, expected_version)
+      end
+
+      def validate_payload_version!(record, key, expected_version)
+        current_version = payload_version(record)
+        return if current_version == expected_version
+
+        raise Smith::PersistenceVersionConflict.new(
+          key: key, expected: expected_version, actual: current_version
+        )
+      end
+
+      def payload_version(record) = PayloadVersion.call(record.public_send(@payload_column))
+
+      def raise_concurrent_conflict(key, expected_version)
+        raise Smith::PersistenceVersionConflict.new(
+          key: key, expected: expected_version, actual: :concurrent
+        )
+      end
 
       def model_class
-        @model_class ||= begin
-          case @model_source
-          when String
-            Object.const_get(@model_source)
-          else
-            @model_source
-          end
-        rescue NameError => e
-          raise ArgumentError, "ActiveRecord model #{@model_source.inspect} could not be resolved: #{e.message}"
+        case @model_source
+        when String
+          Object.const_get(@model_source)
+        else
+          @model_source
         end
+      rescue NameError => e
+        raise ArgumentError, "ActiveRecord model #{@model_source.inspect} could not be resolved: #{e.message}"
       end
     end
   end