fleetforge_runtime/

replay.rs

//! Time-travel replay executor for deterministic verification and evals.

use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;

use anyhow::{anyhow, bail, Context, Result};
use bytes::Bytes;
use hex;
use serde_json::{json, Map, Value};
use tracing::{info, warn};
use uuid::Uuid;

use crate::model::{RunId, StepId};
use fleetforge_storage::{models, ObjectStoreConfig, Storage, StorageConfig};

/// Replay strategy for deterministic time-travel.
#[derive(Debug, Clone, Copy)]
pub enum ReplayStrategy {
    /// Rehydrate outputs from recorded snapshots only.
    Mocked,
    /// Execute live using registered gateways/executors.
    Live,
}

/// Tolerances for drift detection when comparing recorded vs current outputs.
#[derive(Debug, Clone, Copy)]
pub struct DriftTolerance {
    pub tokens: f64,
    pub cost: f64,
}

impl Default for DriftTolerance {
    fn default() -> Self {
        Self {
            tokens: 0.0,
            cost: 0.0,
        }
    }
}

/// Replay outcome including rehydrated snapshots and drift findings.
#[derive(Debug)]
pub struct ReplayOutcome {
    pub run_id: RunId,
    pub strategy: ReplayStrategy,
    pub steps: Vec<StepReplay>,
    pub all_within_tolerance: bool,
}

/// Rehydrated step snapshot with optional drift metadata.
#[derive(Debug)]
pub struct StepReplay {
    pub step_id: StepId,
    pub snapshot: Value,
    pub drift: Option<StepDrift>,
}

/// Drift comparison between recorded snapshot and current output.
#[derive(Debug)]
pub struct StepDrift {
    pub recorded: Value,
    pub current: Value,
    pub token_delta: Option<f64>,
    pub cost_delta: Option<f64>,
    pub within_tolerance: bool,
}

/// Orchestrates mocked/live replays by loading past step snapshots.
pub struct ReplayExecutor {
    storage: Arc<Storage>,
}

impl ReplayExecutor {
    /// Builds a replay executor backed by the shared storage handle.
    pub fn new(storage: Arc<Storage>) -> Self {
        Self { storage }
    }

    /// Replays every step in a run and returns drift diagnostics.
    pub async fn replay_run(
        &self,
        run_id: RunId,
        strategy: ReplayStrategy,
        tolerance: DriftTolerance,
    ) -> Result<ReplayOutcome> {
        let steps = self.storage.fetch_steps_for_run(Uuid::from(run_id)).await?;
        let attempt_rows = self
            .storage
            .fetch_step_attempts_for_run(Uuid::from(run_id))
            .await?;
        let mut attempts_by_step: HashMap<Uuid, Vec<models::StepAttempt>> = HashMap::new();
        for attempt in attempt_rows {
            attempts_by_step
                .entry(attempt.step_id)
                .or_default()
                .push(attempt);
        }
        for attempts in attempts_by_step.values_mut() {
            attempts.sort_by_key(|attempt| attempt.attempt);
        }

        match strategy {
            ReplayStrategy::Mocked => {
                self.replay_mocked(run_id, steps, attempts_by_step, tolerance)
                    .await
            }
            ReplayStrategy::Live => {
                warn!(
                    run = %run_id,
                    "live replay not yet implemented; falling back to mocked replay"
                );
                self.replay_mocked(run_id, steps, attempts_by_step, tolerance)
                    .await
            }
        }
    }
}

impl ReplayExecutor {
    async fn replay_mocked(
        &self,
        run_id: RunId,
        steps: Vec<models::Step>,
        attempts_by_step: HashMap<Uuid, Vec<models::StepAttempt>>,
        tolerance: DriftTolerance,
    ) -> Result<ReplayOutcome> {
        let mut attempts_by_step = attempts_by_step;
        let mut replays = Vec::with_capacity(steps.len());
        let mut all_within = true;

        for step in steps {
            info!(
                run = %run_id,
                step = %step.step_id,
                "mock replay using recorded snapshot"
            );

            let attempt_list = attempts_by_step
                .get(&step.step_id)
                .ok_or_else(|| anyhow!("no attempts recorded for step {}", step.step_id))?;
            let (baseline, latest) = select_attempt_pair(attempt_list).ok_or_else(|| {
                anyhow!(
                    "step {} missing output snapshots; cannot run mocked replay",
                    step.step_id
                )
            })?;

            let recorded_snapshot = baseline.output_snapshot.clone().ok_or_else(|| {
                anyhow!(
                    "attempt {} for step {} missing output_snapshot; cannot replay",
                    baseline.attempt,
                    step.step_id
                )
            })?;
            let current_snapshot = latest.output_snapshot.clone().ok_or_else(|| {
                anyhow!(
                    "latest attempt {} for step {} missing output_snapshot; cannot replay",
                    latest.attempt,
                    step.step_id
                )
            })?;

            let mut drift = compute_drift(&recorded_snapshot, &current_snapshot, tolerance);

            let recorded_attestations = extract_attestation_ids(&recorded_snapshot);
            let current_attestations = extract_attestation_ids(&current_snapshot);
            if recorded_attestations != current_attestations {
                let diff_payload = build_attestation_diff_payload(
                    run_id,
                    StepId(step.step_id),
                    &recorded_attestations,
                    &current_attestations,
                );
                let diff_sha = self
                    .persist_replay_diff(
                        run_id,
                        StepId(step.step_id),
                        "replay_attestation_diff",
                        diff_payload,
                    )
                    .await?;
                warn!(
                    run = %run_id,
                    step = %step.step_id,
                    recorded = ?recorded_attestations,
                    current = ?current_attestations,
                    diff_sha256 = %diff_sha,
                    "attestation ids diverged during replay"
                );
                bail!(
                    "attestation ids diverged during replay for step {} (kind=replay_attestation_diff, diff_sha256={})",
                    step.step_id,
                    diff_sha
                );
            }

            if !tool_events_identical(baseline.tool_events.as_ref(), latest.tool_events.as_ref()) {
                let diff_payload = build_tool_event_diff_payload(
                    run_id,
                    StepId(step.step_id),
                    baseline.tool_events.as_ref(),
                    latest.tool_events.as_ref(),
                );
                let diff_sha = self
                    .persist_replay_diff(
                        run_id,
                        StepId(step.step_id),
                        "replay_tool_diff",
                        diff_payload,
                    )
                    .await?;
                warn!(
                    run = %run_id,
                    step = %step.step_id,
                    diff_sha256 = %diff_sha,
                    "tool events diverged during replay"
                );
                bail!(
                    "tool events diverged during replay for step {} (kind=replay_tool_diff, diff_sha256={})",
                    step.step_id,
                    diff_sha
                );
            }

            if !drift.within_tolerance {
                all_within = false;
            }

            replays.push(StepReplay {
                step_id: StepId(step.step_id),
                snapshot: recorded_snapshot,
                drift: Some(drift),
            });
        }

        Ok(ReplayOutcome {
            run_id,
            strategy: ReplayStrategy::Mocked,
            steps: replays,
            all_within_tolerance: all_within,
        })
    }
}

fn compute_drift(recorded: &Value, current: &Value, tolerance: DriftTolerance) -> StepDrift {
    let token_delta = diff_numeric(
        recorded,
        current,
        &["tokens", "token_count", "total_tokens"],
    );
    let cost_delta = diff_numeric(recorded, current, &["cost", "cost_usd", "total_cost"]);

    let recorded_tokens = find_numeric(recorded, &["tokens", "token_count", "total_tokens"]);
    let recorded_cost = find_numeric(recorded, &["cost", "cost_usd", "total_cost"]);
    let sanitized_recorded = strip_keys(
        recorded,
        &[
            "tokens",
            "token_count",
            "total_tokens",
            "cost",
            "cost_usd",
            "total_cost",
        ],
    );
    let sanitized_current = strip_keys(
        current,
        &[
            "tokens",
            "token_count",
            "total_tokens",
            "cost",
            "cost_usd",
            "total_cost",
        ],
    );

    let tokens_within = match (token_delta, recorded_tokens) {
        (Some(delta), Some(base)) if base.abs() > f64::EPSILON => {
            (delta / base).abs() <= tolerance.tokens
        }
        (Some(delta), Some(_)) => delta.abs() <= tolerance.tokens,
        (Some(delta), None) => delta.abs() <= tolerance.tokens,
        _ => true,
    };
    let cost_within = match (cost_delta, recorded_cost) {
        (Some(delta), Some(base)) if base.abs() > f64::EPSILON => {
            (delta / base).abs() <= tolerance.cost
        }
        (Some(delta), Some(_)) => delta.abs() <= tolerance.cost,
        (Some(delta), None) => delta.abs() <= tolerance.cost,
        _ => true,
    };

    let structure_equal = sanitized_recorded == sanitized_current;

    StepDrift {
        recorded: recorded.clone(),
        current: current.clone(),
        token_delta,
        cost_delta,
        within_tolerance: tokens_within && cost_within && structure_equal,
    }
}

fn diff_numeric(recorded: &Value, current: &Value, keys: &[&str]) -> Option<f64> {
    let recorded_value = find_numeric(recorded, keys)?;
    let current_value = find_numeric(current, keys)?;
    Some(current_value - recorded_value)
}

fn find_numeric(value: &Value, keys: &[&str]) -> Option<f64> {
    match value {
        Value::Object(map) => {
            for key in keys {
                if let Some(v) = map.get(*key) {
                    if let Some(num) = v.as_f64() {
                        return Some(num);
                    }
                }
            }
            for v in map.values() {
                if let Some(found) = find_numeric(v, keys) {
                    return Some(found);
                }
            }
            None
        }
        Value::Array(items) => {
            for item in items {
                if let Some(found) = find_numeric(item, keys) {
                    return Some(found);
                }
            }
            None
        }
        _ => None,
    }
}

fn strip_keys(value: &Value, keys: &[&str]) -> Value {
    match value {
        Value::Object(map) => {
            let mut new_map = serde_json::Map::new();
            for (k, v) in map {
                if keys.iter().any(|key| key == k) {
                    continue;
                }
                new_map.insert(k.clone(), strip_keys(v, keys));
            }
            Value::Object(new_map)
        }
        Value::Array(items) => Value::Array(items.iter().map(|v| strip_keys(v, keys)).collect()),
        _ => value.clone(),
    }
}

fn select_attempt_pair(
    attempts: &[models::StepAttempt],
) -> Option<(&models::StepAttempt, &models::StepAttempt)> {
    let baseline = attempts
        .iter()
        .find(|attempt| attempt.output_snapshot.is_some())?;
    let latest = attempts
        .iter()
        .rev()
        .find(|attempt| attempt.output_snapshot.is_some())
        .unwrap_or(baseline);
    Some((baseline, latest))
}

fn extract_attestation_ids(value: &Value) -> Vec<String> {
    match value {
        Value::Object(map) => map
            .get("attestation_ids")
            .and_then(Value::as_array)
            .map(|arr| {
                arr.iter()
                    .filter_map(Value::as_str)
                    .map(|s| s.to_string())
                    .collect()
            })
            .unwrap_or_default(),
        _ => Vec::new(),
    }
}

fn tool_events_identical(recorded: Option<&Value>, current: Option<&Value>) -> bool {
    match (recorded, current) {
        (None, None) => true,
        (Some(a), Some(b)) => canonicalize(a) == canonicalize(b),
        _ => false,
    }
}

fn canonicalize(value: &Value) -> Value {
    match value {
        Value::Object(map) => {
            let mut ordered = BTreeMap::new();
            for (key, val) in map {
                ordered.insert(key.clone(), canonicalize(val));
            }
            let mut object = Map::new();
            for (key, val) in ordered {
                object.insert(key, val);
            }
            Value::Object(object)
        }
        Value::Array(items) => Value::Array(items.iter().map(|item| canonicalize(item)).collect()),
        _ => value.clone(),
    }
}

impl ReplayExecutor {
    async fn persist_replay_diff(
        &self,
        run_id: RunId,
        step_id: StepId,
        kind: &'static str,
        payload: Value,
    ) -> Result<String> {
        let metadata = json!({
            "kind": kind,
            "run_id": Uuid::from(run_id).to_string(),
            "step_id": Uuid::from(step_id).to_string(),
        });
        let serialized =
            serde_json::to_vec(&payload).context("failed to serialize replay diff payload")?;
        let artifact = self
            .storage
            .artifacts()
            .put(Bytes::from(serialized), "application/json", metadata)
            .await
            .context("failed to persist replay diff artifact")?;
        Ok(hex::encode(artifact.sha256))
    }
}

fn build_attestation_diff_payload(
    run_id: RunId,
    step_id: StepId,
    recorded: &[String],
    current: &[String],
) -> Value {
    let (missing, unexpected) = diff_attestations(recorded, current);
    json!({
        "run_id": Uuid::from(run_id).to_string(),
        "step_id": Uuid::from(step_id).to_string(),
        "recorded": recorded,
        "current": current,
        "missing": missing,
        "unexpected": unexpected,
    })
}

fn diff_attestations(recorded: &[String], current: &[String]) -> (Vec<String>, Vec<String>) {
    (
        recorded
            .iter()
            .filter(|value| !current.contains(value))
            .cloned()
            .collect(),
        current
            .iter()
            .filter(|value| !recorded.contains(value))
            .cloned()
            .collect(),
    )
}

fn build_tool_event_diff_payload(
    run_id: RunId,
    step_id: StepId,
    recorded: Option<&Value>,
    current: Option<&Value>,
) -> Value {
    json!({
        "run_id": Uuid::from(run_id).to_string(),
        "step_id": Uuid::from(step_id).to_string(),
        "recorded": recorded.unwrap_or(&Value::Null),
        "current": current.unwrap_or(&Value::Null),
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use std::sync::Arc;
    use std::time::Duration;
    use testcontainers::clients::Cli;
    use testcontainers::images::postgres::Postgres;
    use tokio::time::sleep;

    #[tokio::test]
    async fn mocked_replay_parity() -> Result<()> {
        if std::fs::metadata("/var/run/docker.sock").is_err()
            && std::env::var("DOCKER_HOST").is_err()
        {
            warn!("Skipping replay test because Docker is unavailable");
            return Ok(());
        }

        let docker = Cli::default();
        let container = docker.run(Postgres::default());
        let port = container.get_host_port_ipv4(5432);
        let database_url = format!("postgres://postgres:postgres@127.0.0.1:{port}/postgres");

        // Postgres inside the container can take a moment to accept connections.
        sleep(Duration::from_millis(250)).await;

        let storage = Arc::new(
            Storage::connect(StorageConfig {
                database_url,
                max_connections: 2,
                connect_timeout: Duration::from_secs(10),
                object_store: ObjectStoreConfig::InMemory,
            })
            .await?,
        );

        let run_id = Uuid::new_v4();
        let step_id = Uuid::new_v4();
        let spec_json = json!({
            "id": step_id,
            "type": "map",
            "inputs": {"task": "demo"},
            "policy": {"budget": {"tokens": 600, "cost": 0.002}}
        });

        let _ = storage
            .insert_run_with_steps(
                models::NewRun {
                    run_id,
                    status: models::RunStatus::Pending,
                    dag_json: json!({"steps": [spec_json.clone()], "edges": []}),
                    input_ctx: json!({}),
                    seed: 42,
                    idempotency_key: None,
                },
                &[models::NewStep {
                    run_id,
                    step_id,
                    idx: 0,
                    priority: 5,
                    spec_json: spec_json.clone(),
                    status: models::StepStatus::Queued,
                    attempt: 0,
                    input_snapshot: None,
                    provider: Some("fleet.map".to_string()),
                    provider_version: None,
                    pending_dependencies: 0,
                    total_dependencies: 0,
                    estimated_tokens: 600,
                    estimated_cost: 0.002,
                    actual_tokens: None,
                    actual_cost: None,
                    ..models::NewStep::default()
                }],
                &[],
            )
            .await?;

        let output = json!({
            "message": "Weather summary",
            "metrics": {
                "tokens": 600,
                "cost": 0.002
            }
        });

        let mut tx = storage.pool().begin().await?;
        storage
            .update_step_status_tx(
                &mut tx,
                run_id,
                step_id,
                models::StepStatus::Succeeded,
                Some(1),
                Some(output.clone()),
                None,
                None,
                Some(output.clone()),
                None,
                Some("fleet.map"),
                None,
            )
            .await?;
        storage
            .update_step_usage_tx(&mut tx, run_id, step_id, Some(600), Some(0.002))
            .await?;
        storage
            .insert_step_attempt_tx(
                &mut tx,
                models::NewStepAttempt {
                    run_id,
                    step_id,
                    attempt: 1,
                    status: models::StepStatus::Succeeded,
                    latency_ms: None,
                    spec_snapshot: spec_json.clone(),
                    input_snapshot: None,
                    output_snapshot: Some(output.clone()),
                    error_snapshot: None,
                    checkpoint: None,
                    budget_snapshot: None,
                    guardrail_events: None,
                    tool_events: None,
                    artifacts: None,
                },
            )
            .await?;
        tx.commit().await?;

        let executor = ReplayExecutor::new(storage.clone());
        let result = executor
            .replay_run(
                RunId(run_id),
                ReplayStrategy::Mocked,
                DriftTolerance {
                    tokens: 0.0,
                    cost: 0.0,
                },
            )
            .await?;

        assert!(result.all_within_tolerance);
        assert_eq!(result.steps.len(), 1);
        assert!(result.steps[0].drift.as_ref().unwrap().within_tolerance);

        drop(container);
        Ok(())
    }

    #[tokio::test]
    async fn mocked_replay_within_one_percent() -> Result<()> {
        if std::fs::metadata("/var/run/docker.sock").is_err()
            && std::env::var("DOCKER_HOST").is_err()
        {
            warn!("Skipping replay test because Docker is unavailable");
            return Ok(());
        }

        let docker = Cli::default();
        let container = docker.run(Postgres::default());
        let port = container.get_host_port_ipv4(5432);
        let database_url = format!("postgres://postgres:postgres@127.0.0.1:{port}/postgres");
        sleep(Duration::from_millis(250)).await;

        let storage = Arc::new(
            Storage::connect(StorageConfig {
                database_url,
                max_connections: 2,
                connect_timeout: Duration::from_secs(10),
                object_store: ObjectStoreConfig::InMemory,
            })
            .await?,
        );

        let run_id = Uuid::new_v4();
        let step_id = Uuid::new_v4();

        let recorded_snapshot = json!({
            "message": "Original output",
            "metrics": {
                "tokens": 100,
                "cost": 1.00
            }
        });

        let actual_output = json!({
            "message": "Original output",
            "metrics": {
                "tokens": 100.5,
                "cost": 1.009
            }
        });

        let _ = storage
            .insert_run_with_steps(
                models::NewRun {
                    run_id,
                    status: models::RunStatus::Pending,
                    dag_json: json!({
                        "steps": [{
                            "id": step_id,
                            "type": "tool",
                            "inputs": {"prompt": "demo"},
                            "policy": {"budget": {"tokens": 100, "cost": 1.0}}
                        }],
                        "edges": []
                    }),
                    input_ctx: json!({}),
                    seed: 1,
                    idempotency_key: None,
                },
                &[models::NewStep {
                    run_id,
                    step_id,
                    idx: 0,
                    priority: 0,
                    spec_json: json!({
                        "id": step_id,
                        "type": "tool",
                        "inputs": {"prompt": "demo"},
                        "policy": {"budget": {"tokens": 100, "cost": 1.0}}
                    }),
                    status: models::StepStatus::Succeeded,
                    attempt: 1,
                    input_snapshot: Some(json!({"prompt": "demo"})),
                    provider: None,
                    provider_version: None,
                    pending_dependencies: 0,
                    total_dependencies: 0,
                    estimated_tokens: 100,
                    estimated_cost: 1.0,
                    actual_tokens: Some(101),
                    actual_cost: Some(1.02),
                    ..models::NewStep::default()
                }],
                &[],
            )
            .await?;

        let mut tx = storage.pool().begin().await?;
        storage
            .insert_step_attempt_tx(
                &mut tx,
                models::NewStepAttempt {
                    run_id,
                    step_id,
                    attempt: 1,
                    status: models::StepStatus::Succeeded,
                    latency_ms: None,
                    spec_snapshot: json!({
                        "id": step_id,
                        "type": "tool",
                        "inputs": {"prompt": "demo"},
                        "policy": {"budget": {"tokens": 100, "cost": 1.0}}
                    }),
                    input_snapshot: Some(json!({"prompt": "demo"})),
                    output_snapshot: Some(recorded_snapshot.clone()),
                    error_snapshot: None,
                    checkpoint: None,
                    budget_snapshot: None,
                    guardrail_events: None,
                    tool_events: None,
                    artifacts: None,
                },
            )
            .await?;
        storage
            .insert_step_attempt_tx(
                &mut tx,
                models::NewStepAttempt {
                    run_id,
                    step_id,
                    attempt: 2,
                    status: models::StepStatus::Succeeded,
                    latency_ms: None,
                    spec_snapshot: json!({
                        "id": step_id,
                        "type": "tool",
                        "inputs": {"prompt": "demo"},
                        "policy": {"budget": {"tokens": 100, "cost": 1.0}}
                    }),
                    input_snapshot: Some(json!({"prompt": "demo"})),
                    output_snapshot: Some(actual_output.clone()),
                    error_snapshot: None,
                    checkpoint: None,
                    budget_snapshot: None,
                    guardrail_events: None,
                    tool_events: None,
                    artifacts: None,
                },
            )
            .await?;
        storage
            .update_step_status_tx(
                &mut tx,
                run_id,
                step_id,
                models::StepStatus::Succeeded,
                Some(2),
                Some(actual_output.clone()),
                None,
                Some(json!({"prompt": "demo"})),
                Some(actual_output.clone()),
                None,
                None,
                None,
            )
            .await?;
        storage
            .update_step_usage_tx(&mut tx, run_id, step_id, Some(101), Some(1.02))
            .await?;
        tx.commit().await?;

        let executor = ReplayExecutor::new(storage.clone());
        let outcome = executor
            .replay_run(
                RunId(run_id),
                ReplayStrategy::Mocked,
                DriftTolerance {
                    tokens: 0.01,
                    cost: 0.01,
                },
            )
            .await?;

        assert!(
            outcome.all_within_tolerance,
            "replay outcome exceeded tolerance"
        );
        let drift = outcome
            .steps
            .first()
            .and_then(|step| step.drift.as_ref())
            .expect("drift metadata present");
        assert!(
            drift.within_tolerance,
            "drift should be within tolerance: {drift:?}"
        );

        drop(container);
        Ok(())
    }

    #[tokio::test]
    async fn mocked_replay_detects_attestation_mismatch() -> Result<()> {
        if std::fs::metadata("/var/run/docker.sock").is_err()
            && std::env::var("DOCKER_HOST").is_err()
        {
            warn!("Skipping replay test because Docker is unavailable");
            return Ok(());
        }

        let docker = Cli::default();
        let container = docker.run(Postgres::default());
        let port = container.get_host_port_ipv4(5432);
        let database_url = format!("postgres://postgres:postgres@127.0.0.1:{port}/postgres");
        sleep(Duration::from_millis(250)).await;

        let storage = Arc::new(
            Storage::connect(StorageConfig {
                database_url,
                max_connections: 2,
                connect_timeout: Duration::from_secs(10),
                object_store: ObjectStoreConfig::InMemory,
            })
            .await?,
        );

        let run_id = Uuid::new_v4();
        let step_id = Uuid::new_v4();

        let baseline_output = json!({
            "result": "ok",
            "attestation_ids": [Uuid::new_v4().to_string()],
        });
        let latest_output = json!({
            "result": "ok",
            "attestation_ids": [Uuid::new_v4().to_string()],
        });

        let _ = storage
            .insert_run_with_steps(
                models::NewRun {
                    run_id,
                    status: models::RunStatus::Pending,
                    dag_json: json!({"steps": [{"id": step_id, "type": "tool", "inputs": {}}], "edges": []}),
                    input_ctx: json!({}),
                    seed: 11,
                    idempotency_key: None,
                },
                &[models::NewStep {
                    run_id,
                    step_id,
                    idx: 0,
                    priority: 0,
                    spec_json: json!({"id": step_id, "type": "tool", "inputs": {}}),
                    status: models::StepStatus::Succeeded,
                    attempt: 1,
                    input_snapshot: None,
                    provider: None,
                    provider_version: None,
                    pending_dependencies: 0,
                    total_dependencies: 0,
                    estimated_tokens: 0,
                    estimated_cost: 0.0,
                    actual_tokens: None,
                    actual_cost: None,
                    ..models::NewStep::default()
                }],
                &[],
            )
            .await?;

        let mut tx = storage.pool().begin().await?;
        storage
            .insert_step_attempt_tx(
                &mut tx,
                models::NewStepAttempt {
                    run_id,
                    step_id,
                    attempt: 1,
                    status: models::StepStatus::Succeeded,
                    latency_ms: None,
                    spec_snapshot: json!({"id": step_id, "type": "tool", "inputs": {}}),
                    input_snapshot: None,
                    output_snapshot: Some(baseline_output.clone()),
                    error_snapshot: None,
                    checkpoint: None,
                    budget_snapshot: None,
                    guardrail_events: None,
                    tool_events: None,
                    artifacts: None,
                },
            )
            .await?;
        storage
            .insert_step_attempt_tx(
                &mut tx,
                models::NewStepAttempt {
                    run_id,
                    step_id,
                    attempt: 2,
                    status: models::StepStatus::Succeeded,
                    latency_ms: None,
                    spec_snapshot: json!({"id": step_id, "type": "tool", "inputs": {}}),
                    input_snapshot: None,
                    output_snapshot: Some(latest_output.clone()),
                    error_snapshot: None,
                    checkpoint: None,
                    budget_snapshot: None,
                    guardrail_events: None,
                    tool_events: None,
                    artifacts: None,
                },
            )
            .await?;
        storage
            .update_step_status_tx(
                &mut tx,
                run_id,
                step_id,
                models::StepStatus::Succeeded,
                Some(2),
                Some(latest_output.clone()),
                None,
                None,
                Some(latest_output.clone()),
                None,
                None,
                None,
            )
            .await?;
        tx.commit().await?;

        let executor = ReplayExecutor::new(storage.clone());
        let err = executor
            .replay_run(
                RunId(run_id),
                ReplayStrategy::Mocked,
                DriftTolerance {
                    tokens: 0.01,
                    cost: 0.01,
                },
            )
            .await
            .expect_err("attestation mismatch must fail replay");
        assert!(
            err.to_string().contains("replay_attestation_diff"),
            "error should reference replay_attestation_diff artifact: {err:?}"
        );

        drop(container);
        Ok(())
    }

    #[tokio::test]
    async fn mocked_replay_detects_tool_event_mismatch() -> Result<()> {
        if std::fs::metadata("/var/run/docker.sock").is_err()
            && std::env::var("DOCKER_HOST").is_err()
        {
            warn!("Skipping replay test because Docker is unavailable");
            return Ok(());
        }

        let docker = Cli::default();
        let container = docker.run(Postgres::default());
        let port = container.get_host_port_ipv4(5432);
        let database_url = format!("postgres://postgres:postgres@127.0.0.1:{port}/postgres");
        sleep(Duration::from_millis(250)).await;

        let storage = Arc::new(
            Storage::connect(StorageConfig {
                database_url,
                max_connections: 2,
                connect_timeout: Duration::from_secs(10),
                object_store: ObjectStoreConfig::InMemory,
            })
            .await?,
        );

        let run_id = Uuid::new_v4();
        let step_id = Uuid::new_v4();

        let output = json!({
            "result": "value",
            "attestation_ids": [Uuid::new_v4().to_string()],
        });

        let baseline_tool_events = json!([{
            "kind": "tool.run",
            "payload": {"stdout": "hello"},
            "trace": {"span_id": "abc"}
        }]);
        let latest_tool_events = json!([{
            "kind": "tool.run",
            "payload": {"stdout": "goodbye"},
            "trace": {"span_id": "def"}
        }]);

        let _ = storage
            .insert_run_with_steps(
                models::NewRun {
                    run_id,
                    status: models::RunStatus::Pending,
                    dag_json: json!({"steps": [{"id": step_id, "type": "tool", "inputs": {}}], "edges": []}),
                    input_ctx: json!({}),
                    seed: 21,
                    idempotency_key: None,
                },
                &[models::NewStep {
                    run_id,
                    step_id,
                    idx: 0,
                    priority: 0,
                    spec_json: json!({"id": step_id, "type": "tool", "inputs": {}}),
                    status: models::StepStatus::Succeeded,
                    attempt: 1,
                    input_snapshot: None,
                    provider: None,
                    provider_version: None,
                    pending_dependencies: 0,
                    total_dependencies: 0,
                    estimated_tokens: 0,
                    estimated_cost: 0.0,
                    actual_tokens: None,
                    actual_cost: None,
                    ..models::NewStep::default()
                }],
                &[],
            )
            .await?;

        let mut tx = storage.pool().begin().await?;
        storage
            .insert_step_attempt_tx(
                &mut tx,
                models::NewStepAttempt {
                    run_id,
                    step_id,
                    attempt: 1,
                    status: models::StepStatus::Succeeded,
                    latency_ms: None,
                    spec_snapshot: json!({"id": step_id, "type": "tool", "inputs": {}}),
                    input_snapshot: None,
                    output_snapshot: Some(output.clone()),
                    error_snapshot: None,
                    checkpoint: None,
                    budget_snapshot: None,
                    guardrail_events: None,
                    tool_events: Some(baseline_tool_events.clone()),
                    artifacts: None,
                },
            )
            .await?;
        storage
            .insert_step_attempt_tx(
                &mut tx,
                models::NewStepAttempt {
                    run_id,
                    step_id,
                    attempt: 2,
                    status: models::StepStatus::Succeeded,
                    latency_ms: None,
                    spec_snapshot: json!({"id": step_id, "type": "tool", "inputs": {}}),
                    input_snapshot: None,
                    output_snapshot: Some(output.clone()),
                    error_snapshot: None,
                    checkpoint: None,
                    budget_snapshot: None,
                    guardrail_events: None,
                    tool_events: Some(latest_tool_events.clone()),
                    artifacts: None,
                },
            )
            .await?;
        storage
            .update_step_status_tx(
                &mut tx,
                run_id,
                step_id,
                models::StepStatus::Succeeded,
                Some(2),
                Some(output.clone()),
                None,
                None,
                Some(output.clone()),
                None,
                None,
                None,
            )
            .await?;
        tx.commit().await?;

        let executor = ReplayExecutor::new(storage.clone());
        let err = executor
            .replay_run(
                RunId(run_id),
                ReplayStrategy::Mocked,
                DriftTolerance {
                    tokens: 0.0,
                    cost: 0.0,
                },
            )
            .await
            .expect_err("tool event mismatch must fail replay");
        assert!(
            err.to_string().contains("replay_tool_diff"),
            "error should reference replay_tool_diff artifact: {err:?}"
        );

        drop(container);
        Ok(())
    }

    #[tokio::test]
    async fn mocked_replay_requires_snapshots() {
        if std::fs::metadata("/var/run/docker.sock").is_err()
            && std::env::var("DOCKER_HOST").is_err()
        {
            warn!("Skipping replay test because Docker is unavailable");
            return;
        }

        let docker = Cli::default();
        let container = docker.run(Postgres::default());
        let port = container.get_host_port_ipv4(5432);
        let database_url = format!("postgres://postgres:postgres@127.0.0.1:{port}/postgres");
        sleep(Duration::from_millis(250)).await;

        let storage = Arc::new(
            Storage::connect(StorageConfig {
                database_url,
                max_connections: 2,
                connect_timeout: Duration::from_secs(10),
                object_store: ObjectStoreConfig::InMemory,
            })
            .await
            .expect("storage connection"),
        );

        let run_id = Uuid::new_v4();
        let step_id = Uuid::new_v4();

        let _ = storage
            .insert_run_with_steps(
                models::NewRun {
                    run_id,
                    status: models::RunStatus::Pending,
                    dag_json: json!({"steps": [], "edges": []}),
                    input_ctx: json!({}),
                    seed: 7,
                    idempotency_key: None,
                },
                &[models::NewStep {
                    run_id,
                    step_id,
                    idx: 0,
                    priority: 0,
                    spec_json: json!({"id": step_id, "type": "map", "inputs": {}}),
                    status: models::StepStatus::Queued,
                    attempt: 0,
                    input_snapshot: None,
                    provider: None,
                    provider_version: None,
                    pending_dependencies: 0,
                    total_dependencies: 0,
                    estimated_tokens: 0,
                    estimated_cost: 0.0,
                    actual_tokens: None,
                    actual_cost: None,
                    ..models::NewStep::default()
                }],
                &[],
            )
            .await
            .expect("seed run");

        let mut tx = storage.pool().begin().await.expect("begin tx");
        storage
            .insert_step_attempt_tx(
                &mut tx,
                models::NewStepAttempt {
                    run_id,
                    step_id,
                    attempt: 1,
                    status: models::StepStatus::Queued,
                    latency_ms: None,
                    spec_snapshot: json!({"id": step_id, "type": "map", "inputs": {}}),
                    input_snapshot: None,
                    output_snapshot: None,
                    error_snapshot: None,
                    checkpoint: None,
                    budget_snapshot: None,
                    guardrail_events: None,
                    tool_events: None,
                    artifacts: None,
                },
            )
            .await
            .expect("insert attempt");
        tx.commit().await.expect("commit");

        let executor = ReplayExecutor::new(storage);
        let err = executor
            .replay_run(
                RunId(run_id),
                ReplayStrategy::Mocked,
                DriftTolerance::default(),
            )
            .await
            .expect_err("missing snapshots should error");

        assert!(err
            .to_string()
            .contains("missing output snapshots; cannot run mocked replay"));
        drop(container);
    }
}