ironflow_store/memory/

run_store.rs

use std::cmp::Reverse;
use std::collections::{BTreeMap, HashMap, HashSet};

use chrono::{DateTime, Duration, Utc};
use rust_decimal::Decimal;
use uuid::Uuid;

use crate::entities::{
    ApiKey, ConcurrencyGroupBacklog, IDEMPOTENCY_WINDOW, LeaseRequest, LeaseUpdate, NewRun,
    NewStep, NewStepDependency, Page, PurgePolicy, PurgeReason, PurgeableRun, ReapedRun, Run,
    RunActor, RunCreation, RunFilter, RunStats, RunStatus, RunUpdate, StatsHistoryBucket,
    StatsHistoryFilter, Step, StepApproval, StepDependency, StepStatus, StepUpdate, TriggerKind,
    User, WorkerCapabilities, normalize_worker_tags, validate_concurrency_limits,
    validate_priority, validate_worker_tags,
};
use crate::error::StoreError;
use crate::store::{LEASE_EXPIRED_ERROR, RunStore, StoreFuture};

use super::descendants::active_descendants;
use super::stats_history::aggregate_history_buckets;
use super::{InMemoryStore, State};

/// Resolve [`Run::created_by_label`] from the current users and API keys.
///
/// Mirrors the `LEFT JOIN` the PostgreSQL store performs at read time, so a
/// renamed API key or a renamed user is reflected immediately.
fn resolve_created_by_label(
    actor: Option<&RunActor>,
    users: &HashMap<Uuid, User>,
    api_keys: &HashMap<Uuid, ApiKey>,
) -> Option<String> {
    let actor = actor?;
    let username = users.get(&actor.user_id()).map(|u| u.username.clone());

    match actor.api_key_id() {
        None => username,
        Some(api_key_id) => {
            let key_name = api_keys.get(&api_key_id).map(|k| k.name.clone())?;
            Some(match username {
                Some(username) => format!("{key_name} ({username})"),
                None => key_name,
            })
        }
    }
}

/// Return a clone of `run` with its display label resolved against `state`.
fn run_with_label(run: &Run, state: &State) -> Run {
    let mut run = run.clone();
    run.created_by_label =
        resolve_created_by_label(run.created_by.as_ref(), &state.users, &state.api_keys);
    run
}

/// Drop the worker lease held on a run.
fn clear_lease(run: &mut Run) {
    run.worker_id = None;
    run.lease_expires_at = None;
}

/// Number of root runs in `Running` that carry `group`.
///
/// Sub-workflow runs execute inside their parent's slot and are never counted.
fn running_count(runs: &HashMap<Uuid, Run>, group: &str) -> u64 {
    runs.values()
        .filter(|r| {
            r.status.state == RunStatus::Running
                && !matches!(r.trigger, TriggerKind::Workflow)
                && r.concurrency_limits.iter().any(|l| l.group == group)
        })
        .count() as u64
}

/// Whether at least one of the run's concurrency groups is saturated for the
/// run's own limit.
fn is_blocked(run: &Run, runs: &HashMap<Uuid, Run>) -> bool {
    run.concurrency_limits
        .iter()
        .any(|l| running_count(runs, &l.group) >= u64::from(l.limit))
}

/// Whether a run waits for a pick: pending or retrying, and due.
fn is_due(run: &Run, now: DateTime<Utc>) -> bool {
    matches!(run.status.state, RunStatus::Pending | RunStatus::Retrying)
        && run.scheduled_at.is_none_or(|at| at <= now)
}

fn run_matches_filter(run: &Run, filter: &RunFilter, steps: &HashMap<Uuid, Step>) -> bool {
    if let Some(ref wf) = filter.workflow_name
        && !run
            .workflow_name
            .to_lowercase()
            .contains(&wf.to_lowercase())
    {
        return false;
    }
    if let Some(ref status) = filter.status
        && &run.status.state != status
    {
        return false;
    }
    if let Some(after) = filter.created_after
        && run.created_at < after
    {
        return false;
    }
    if let Some(before) = filter.created_before
        && run.created_at > before
    {
        return false;
    }
    if let Some(has_steps) = filter.has_steps
        && matches!(
            run.status.state,
            RunStatus::Completed | RunStatus::Cancelled
        )
    {
        let run_has_steps = steps.values().any(|s| s.run_id == run.id);
        if has_steps != run_has_steps {
            return false;
        }
    }
    if let Some(ref labels) = filter.labels {
        for (key, value) in labels {
            if run.labels.get(key) != Some(value) {
                return false;
            }
        }
    }
    if let Some(user_id) = filter.created_by_user_id
        && run.created_by.as_ref().map(RunActor::user_id) != Some(user_id)
    {
        return false;
    }
    if let Some(ref group) = filter.concurrency_group
        && !run.concurrency_limits.iter().any(|l| &l.group == group)
    {
        return false;
    }
    if let Some(priority) = filter.priority
        && run.priority != priority
    {
        return false;
    }
    if let Some(ref caps) = filter.eligible_for
        && !caps.can_take(&run.workflow_name, &run.worker_tags)
    {
        return false;
    }
    true
}

/// Insert a run into the locked state, honouring its idempotency and
/// concurrency keys. Writes nothing when it returns an error.
///
/// Shared by [`RunStore::create_run`] and the schedule firing, which must
/// create the run under the same lock as the schedule update.
pub(super) fn insert_run(state: &mut State, req: NewRun) -> Result<RunCreation, StoreError> {
    validate_concurrency_limits(&req.concurrency_limits)?;
    // Mirrors the CHECK constraint of the PostgreSQL column.
    validate_priority(req.priority).map_err(StoreError::Database)?;
    validate_worker_tags(&req.worker_tags)?;
    let now = Utc::now();

    if let Some(ref key) = req.idempotency_key
        && let Some(existing) = state
            .idempotency_keys
            .get(key)
            .and_then(|id| state.runs.get(id))
    {
        if now - existing.created_at < IDEMPOTENCY_WINDOW {
            return Ok(RunCreation::Existing(run_with_label(existing, state)));
        }
        // The key outlived its window: release it from the stale run.
        let stale_id = existing.id;
        state.idempotency_keys.remove(key);
        if let Some(stale) = state.runs.get_mut(&stale_id) {
            stale.idempotency_key = None;
        }
    }

    if let Some(ref key) = req.concurrency_key
        && let Some(holder) = state
            .runs
            .values()
            .filter(|r| {
                r.concurrency_key.as_deref() == Some(key.as_str()) && !r.status.state.is_terminal()
            })
            .min_by_key(|r| r.created_at)
    {
        return Err(StoreError::ConcurrencyConflict {
            key: key.clone(),
            run_id: holder.id,
        });
    }

    let run = Run {
        id: Uuid::now_v7(),
        workflow_name: req.workflow_name,
        status: crate::entities::FsmState::new(RunStatus::Pending, Uuid::now_v7()),
        trigger: req.trigger,
        payload: req.payload,
        error: None,
        retry_count: 0,
        max_retries: req.max_retries,
        cost_usd: Decimal::ZERO,
        duration_ms: 0,
        created_at: now,
        updated_at: now,
        started_at: None,
        completed_at: None,
        handler_version: req.handler_version,
        labels: req.labels,
        scheduled_at: req.scheduled_at,
        created_by: req.created_by,
        created_by_label: None,
        idempotency_key: req.idempotency_key.clone(),
        concurrency_key: req.concurrency_key,
        priority: req.priority,
        concurrency_limits: req.concurrency_limits,
        max_cost_usd: req.max_cost_usd,
        worker_id: None,
        lease_expires_at: None,
        output: None,
        lease_recoveries: 0,
        capacity_wait_kind: None,
        worker_tags: normalize_worker_tags(req.worker_tags),
    };

    if let Some(key) = req.idempotency_key {
        state.idempotency_keys.insert(key, run.id);
    }
    state.runs.insert(run.id, run.clone());
    Ok(RunCreation::Created(run_with_label(&run, state)))
}

impl RunStore for InMemoryStore {
    fn create_run(&self, req: NewRun) -> StoreFuture<'_, RunCreation> {
        Box::pin(async move {
            // Single critical section: the key lookup and the insert cannot be
            // interleaved by a concurrent call sharing the same key.
            let mut state = self.state.write().await;
            insert_run(&mut state, req)
        })
    }

    fn find_run_by_idempotency_key(&self, key: &str) -> StoreFuture<'_, Option<Run>> {
        let key = key.to_string();
        Box::pin(async move {
            let now = Utc::now();
            let state = self.state.read().await;
            Ok(state
                .idempotency_keys
                .get(&key)
                .and_then(|id| state.runs.get(id))
                .filter(|run| now - run.created_at < IDEMPOTENCY_WINDOW)
                .map(|run| run_with_label(run, &state)))
        })
    }

    fn get_run(&self, id: Uuid) -> StoreFuture<'_, Option<Run>> {
        Box::pin(async move {
            let state = self.state.read().await;
            Ok(state.runs.get(&id).map(|r| run_with_label(r, &state)))
        })
    }

    fn list_runs(&self, filter: RunFilter, page: u32, per_page: u32) -> StoreFuture<'_, Page<Run>> {
        Box::pin(async move {
            let state = self.state.read().await;

            let mut runs: Vec<&Run> = state
                .runs
                .values()
                .filter(|r| run_matches_filter(r, &filter, &state.steps))
                .collect();

            // Sort newest first.
            runs.sort_by_key(|r| Reverse(r.created_at));

            let total = runs.len() as u64;
            let page = page.max(1);
            let per_page = per_page.clamp(1, 100);
            let offset = ((page - 1) * per_page) as usize;
            let items: Vec<Run> = runs
                .into_iter()
                .skip(offset)
                .take(per_page as usize)
                .map(|r| run_with_label(r, &state))
                .collect();

            Ok(Page {
                items,
                total,
                page,
                per_page,
            })
        })
    }

    fn update_run_status(&self, id: Uuid, new_status: RunStatus) -> StoreFuture<'_, ()> {
        Box::pin(async move {
            let mut state = self.state.write().await;
            let run = state.runs.get_mut(&id).ok_or(StoreError::RunNotFound(id))?;

            if !run.status.state.can_transition_to(&new_status) {
                return Err(StoreError::InvalidTransition {
                    from: run.status.state,
                    to: new_status,
                });
            }

            if run.status.state == new_status && new_status.is_terminal() {
                return Ok(());
            }

            let now = Utc::now();
            run.status.state = new_status;
            run.updated_at = now;

            if new_status == RunStatus::Running && run.started_at.is_none() {
                run.started_at = Some(now);
            }
            if new_status.is_terminal() {
                run.completed_at = Some(now);
            }
            if new_status != RunStatus::Running {
                clear_lease(run);
            }
            if new_status != RunStatus::Sleeping {
                run.capacity_wait_kind = None;
            }

            Ok(())
        })
    }

    fn update_run(&self, id: Uuid, update: RunUpdate) -> StoreFuture<'_, ()> {
        Box::pin(async move {
            let mut state = self.state.write().await;
            let run = state.runs.get_mut(&id).ok_or(StoreError::RunNotFound(id))?;

            let now = Utc::now();

            if let Some(status) = update.status {
                if !run.status.state.can_transition_to(&status) {
                    return Err(StoreError::InvalidTransition {
                        from: run.status.state,
                        to: status,
                    });
                }
                if !(run.status.state == status && status.is_terminal()) {
                    run.status.state = status;
                    if status == RunStatus::Running && run.started_at.is_none() {
                        run.started_at = Some(now);
                    }
                    if status.is_terminal() {
                        run.completed_at = Some(now);
                    }
                    if status != RunStatus::Running {
                        clear_lease(run);
                    }
                }
                // The kind only means something while the run sleeps on capacity.
                run.capacity_wait_kind = if status == RunStatus::Sleeping {
                    update.capacity_wait_kind.clone()
                } else {
                    None
                };
            }

            // After the status block, so `Running` + `Set` ends with a lease.
            match update.lease {
                Some(LeaseUpdate::Set {
                    worker_id,
                    expires_at,
                }) => {
                    run.worker_id = Some(worker_id);
                    run.lease_expires_at = Some(expires_at);
                }
                Some(LeaseUpdate::Release) => clear_lease(run),
                None => {}
            }

            if let Some(error) = update.error {
                run.error = Some(error);
            }
            if update.increment_retry {
                run.retry_count += 1;
            }
            if let Some(cost) = update.cost_usd {
                run.cost_usd = cost;
            }
            if let Some(dur) = update.duration_ms {
                run.duration_ms = dur;
            }
            if let Some(started) = update.started_at {
                run.started_at = Some(started);
            }
            if let Some(completed) = update.completed_at {
                run.completed_at = Some(completed);
            }
            if let Some(scheduled) = update.scheduled_at {
                run.scheduled_at = Some(scheduled);
            }
            if let Some(output) = update.output {
                run.output = Some(output);
            }

            run.updated_at = now;
            Ok(())
        })
    }

    fn list_active_descendants(&self, run_id: Uuid) -> StoreFuture<'_, Vec<Run>> {
        Box::pin(async move {
            let state = self.state.read().await;
            Ok(active_descendants(&state.runs, run_id))
        })
    }

    fn pick_next_pending_for(
        &self,
        lease: Option<LeaseRequest>,
        capabilities: Option<WorkerCapabilities>,
    ) -> StoreFuture<'_, Option<Run>> {
        Box::pin(async move {
            let mut state = self.state.write().await;
            let now = Utc::now();

            // Find the highest priority run waiting for execution whose
            // scheduled_at has passed (or is None), the oldest first among
            // equal priorities. `Retrying` runs are runs whose automatic retry
            // backoff has been armed: they become eligible again once
            // `scheduled_at` has passed. Runs held back by a saturated
            // concurrency group are skipped, so they never block younger runs.
            // The group check and the transition below happen under the same
            // write lock, so concurrent pickers cannot overshoot a limit.
            // Runs the worker cannot take (unregistered workflow, missing tag)
            // are skipped the same way.
            let oldest_id = state
                .runs
                .values()
                .filter(|r| is_due(r, now) && !is_blocked(r, &state.runs))
                .filter(|r| {
                    capabilities
                        .as_ref()
                        .is_none_or(|c| c.can_take(&r.workflow_name, &r.worker_tags))
                })
                .min_by_key(|r| (Reverse(r.priority), r.created_at))
                .map(|r| r.id);

            let Some(id) = oldest_id else {
                return Ok(None);
            };

            // Transition to Running and attach the lease atomically.
            let run = state.runs.get_mut(&id).expect("run exists");
            let now = Utc::now();
            run.status.state = RunStatus::Running;
            run.started_at = Some(now);
            run.updated_at = now;
            match lease {
                Some(lease) => {
                    run.lease_expires_at = Some(lease.expires_at(now));
                    run.worker_id = Some(lease.worker_id);
                }
                None => clear_lease(run),
            }
            let run = run.clone();

            Ok(Some(run_with_label(&run, &state)))
        })
    }

    fn count_blocked_runs_by_group(&self) -> StoreFuture<'_, Vec<ConcurrencyGroupBacklog>> {
        Box::pin(async move {
            let state = self.state.read().await;
            let now = Utc::now();

            let mut blocked: BTreeMap<String, u64> = BTreeMap::new();
            for run in state.runs.values().filter(|r| is_due(r, now)) {
                for limit in &run.concurrency_limits {
                    if running_count(&state.runs, &limit.group) >= u64::from(limit.limit) {
                        *blocked.entry(limit.group.clone()).or_default() += 1;
                    }
                }
            }

            Ok(blocked
                .into_iter()
                .map(|(group, blocked_runs)| ConcurrencyGroupBacklog {
                    group,
                    blocked_runs,
                })
                .collect())
        })
    }

    fn renew_lease(&self, id: Uuid, lease: LeaseRequest) -> StoreFuture<'_, DateTime<Utc>> {
        Box::pin(async move {
            let mut state = self.state.write().await;
            let run = state.runs.get_mut(&id).ok_or(StoreError::RunNotFound(id))?;

            if run.status.state != RunStatus::Running
                || run.worker_id.as_deref() != Some(lease.worker_id.as_str())
            {
                return Err(StoreError::LeaseLost {
                    run_id: id,
                    held_by: run.worker_id.clone(),
                });
            }

            let now = Utc::now();
            let expires_at = lease.expires_at(now);
            run.lease_expires_at = Some(expires_at);
            run.updated_at = now;

            Ok(expires_at)
        })
    }

    fn reap_expired_leases(&self, limit: u32) -> StoreFuture<'_, Vec<ReapedRun>> {
        Box::pin(async move {
            let mut state = self.state.write().await;
            let now = Utc::now();

            let mut expired: Vec<Uuid> = state
                .runs
                .values()
                .filter(|r| {
                    r.status.state == RunStatus::Running
                        && r.lease_expires_at.is_some_and(|at| at < now)
                })
                .map(|r| r.id)
                .collect();
            expired.sort_unstable();
            expired.truncate(limit as usize);

            let mut reaped = Vec::with_capacity(expired.len());
            for id in expired {
                let run = state.runs.get_mut(&id).expect("run exists");
                run.lease_recoveries += 1;
                clear_lease(run);
                run.updated_at = now;

                let to = if run.lease_recoveries > run.max_retries {
                    run.status.state = RunStatus::Failed;
                    run.error = Some(LEASE_EXPIRED_ERROR.to_string());
                    run.completed_at = Some(now);
                    RunStatus::Failed
                } else {
                    run.status.state = RunStatus::Pending;
                    RunStatus::Pending
                };

                reaped.push(ReapedRun {
                    run: run.clone(),
                    from: RunStatus::Running,
                    to,
                });
            }

            Ok(reaped)
        })
    }

    fn claim_due_approval_deadlines(&self, limit: u32) -> StoreFuture<'_, Vec<Step>> {
        Box::pin(async move {
            let mut state = self.state.write().await;
            let now = Utc::now();

            let mut due: Vec<(DateTime<Utc>, Uuid)> = state
                .steps
                .values()
                .filter(|s| {
                    s.status.state == StepStatus::AwaitingApproval
                        && s.approval_deadline_at.is_some_and(|at| at <= now)
                })
                .map(|s| (s.approval_deadline_at.expect("deadline is set"), s.id))
                .collect();
            due.sort_unstable();
            due.truncate(limit as usize);

            let mut claimed = Vec::with_capacity(due.len());
            for (_, id) in due {
                let step = state.steps.get_mut(&id).expect("step exists");
                // Clone before clearing so the caller still sees the deadline
                // that fired.
                claimed.push(step.clone());
                step.approval_deadline_at = None;
                step.updated_at = now;
            }

            Ok(claimed)
        })
    }

    fn claim_due_sleeping_runs(&self, limit: u32) -> StoreFuture<'_, Vec<Run>> {
        Box::pin(async move {
            let mut state = self.state.write().await;
            let now = Utc::now();

            let mut due: Vec<(DateTime<Utc>, Uuid)> = state
                .runs
                .values()
                .filter(|r| r.status.state == RunStatus::Sleeping)
                .filter_map(|r| r.scheduled_at.filter(|at| *at <= now).map(|at| (at, r.id)))
                .collect();
            due.sort_unstable();
            due.truncate(limit as usize);

            let mut woken = Vec::with_capacity(due.len());
            for (_, id) in due {
                let run = state.runs.get_mut(&id).expect("run exists");
                run.status.state = RunStatus::Pending;
                run.scheduled_at = None;
                run.capacity_wait_kind = None;
                run.updated_at = now;
                let run = run.clone();
                woken.push(run_with_label(&run, &state));
            }

            Ok(woken)
        })
    }

    fn list_purgeable_runs(
        &self,
        policy: &PurgePolicy,
        batch_size: u32,
    ) -> StoreFuture<'_, Vec<PurgeableRun>> {
        let max_age_days = policy.max_age_days;
        let max_runs_per_workflow = policy.max_runs_per_workflow;
        Box::pin(async move {
            let state = self.state.read().await;
            let cutoff = Utc::now() - Duration::days(i64::from(max_age_days));
            let mut result: Vec<PurgeableRun> = Vec::new();
            let mut seen: HashSet<Uuid> = HashSet::new();

            for run in state.runs.values() {
                if run.status.state.is_terminal() && run.created_at < cutoff {
                    seen.insert(run.id);
                    result.push(PurgeableRun {
                        run_id: run.id,
                        workflow_name: run.workflow_name.clone(),
                        reason: PurgeReason::TooOld,
                    });
                }
            }

            let mut by_workflow: HashMap<&str, Vec<&Run>> = HashMap::new();
            for run in state.runs.values() {
                if run.status.state.is_terminal() {
                    by_workflow.entry(&run.workflow_name).or_default().push(run);
                }
            }
            for (_, mut runs) in by_workflow {
                if runs.len() > max_runs_per_workflow as usize {
                    runs.sort_by_key(|r| r.created_at);
                    let excess = runs.len() - max_runs_per_workflow as usize;
                    for run in runs.into_iter().take(excess) {
                        if seen.insert(run.id) {
                            result.push(PurgeableRun {
                                run_id: run.id,
                                workflow_name: run.workflow_name.clone(),
                                reason: PurgeReason::ExceedsWorkflowLimit,
                            });
                        }
                    }
                }
            }

            result.sort_by_key(|p| p.run_id);
            result.truncate(batch_size as usize);
            Ok(result)
        })
    }

    fn delete_run(&self, id: Uuid) -> StoreFuture<'_, Vec<String>> {
        Box::pin(async move {
            let mut state = self.state.write().await;

            if !state.runs.contains_key(&id) {
                return Err(StoreError::RunNotFound(id));
            }

            // Collect artifact storage keys before removing them.
            let storage_keys: Vec<String> = state
                .artifacts
                .values()
                .filter(|a| a.run_id == id)
                .map(|a| a.storage_key.clone())
                .collect();

            // Remove artifacts.
            state.artifacts.retain(|_, a| a.run_id != id);

            // Remove step dependencies (both sides).
            let step_ids: Vec<Uuid> = state
                .steps
                .values()
                .filter(|s| s.run_id == id)
                .map(|s| s.id)
                .collect();
            state
                .step_dependencies
                .retain(|d| !step_ids.contains(&d.step_id) && !step_ids.contains(&d.depends_on));

            // Remove steps.
            state.steps.retain(|_, s| s.run_id != id);

            // Remove idempotency key pointing to this run.
            state.idempotency_keys.retain(|_, &mut run_id| run_id != id);

            // Remove the run itself.
            state.runs.remove(&id);

            Ok(storage_keys)
        })
    }

    fn create_step(&self, req: NewStep) -> StoreFuture<'_, Step> {
        Box::pin(async move {
            let mut state = self.state.write().await;

            let attempt = state
                .runs
                .get(&req.run_id)
                .ok_or(StoreError::RunNotFound(req.run_id))?
                .retry_count
                + 1;

            let now = Utc::now();
            let step = Step {
                id: Uuid::now_v7(),
                trace_id: req.trace_id,
                run_id: req.run_id,
                name: req.name,
                kind: req.kind,
                position: req.position,
                status: crate::entities::FsmState::new(StepStatus::Pending, Uuid::now_v7()),
                attempt,
                input: req.input,
                output: None,
                error: None,
                duration_ms: 0,
                cost_usd: Decimal::ZERO,
                input_tokens: None,
                cache_read_input_tokens: None,
                cache_creation_input_tokens: None,
                output_tokens: None,
                created_at: now,
                updated_at: now,
                started_at: None,
                completed_at: None,
                debug_messages: None,
                is_error_handler: req.is_error_handler,
                approval_deadline_at: None,
                approval_stage: 0,
                approval_assignee: None,
                approval_requirement: None,
                approvals: Vec::new(),
                account_id: None,
                environment_id: None,
                session_id: None,
            };

            state.steps.insert(step.id, step.clone());
            Ok(step)
        })
    }

    fn update_step(&self, id: Uuid, update: StepUpdate) -> StoreFuture<'_, ()> {
        Box::pin(async move {
            let mut state = self.state.write().await;
            let step = state
                .steps
                .get_mut(&id)
                .ok_or(StoreError::StepNotFound(id))?;

            let now = Utc::now();

            if let Some(status) = update.status {
                if !matches!(
                    (step.status.state, status),
                    (StepStatus::Pending, StepStatus::Running)
                        | (StepStatus::Pending, StepStatus::Skipped)
                        | (StepStatus::Running, StepStatus::Completed)
                        | (StepStatus::Running, StepStatus::Failed)
                        | (StepStatus::Running, StepStatus::AwaitingApproval)
                        | (StepStatus::AwaitingApproval, StepStatus::Running)
                        | (StepStatus::AwaitingApproval, StepStatus::Completed)
                        | (StepStatus::AwaitingApproval, StepStatus::Failed)
                        | (StepStatus::AwaitingApproval, StepStatus::Rejected)
                ) {
                    return Err(StoreError::Database(format!(
                        "invalid step status transition: {:?} -> {:?}",
                        step.status.state, status
                    )));
                }
                step.status.state = status;
            }
            if let Some(output) = update.output {
                step.output = Some(output);
            }
            if let Some(error) = update.error {
                step.error = Some(error);
            }
            if let Some(dur) = update.duration_ms {
                step.duration_ms = dur;
            }
            if let Some(cost) = update.cost_usd {
                step.cost_usd = cost;
            }
            if let Some(tokens) = update.input_tokens {
                step.input_tokens = Some(tokens);
            }
            if let Some(tokens) = update.cache_read_input_tokens {
                step.cache_read_input_tokens = Some(tokens);
            }
            if let Some(tokens) = update.cache_creation_input_tokens {
                step.cache_creation_input_tokens = Some(tokens);
            }
            if let Some(account_id) = update.account_id {
                step.account_id = Some(account_id);
            }
            if let Some(environment_id) = update.environment_id {
                step.environment_id = Some(environment_id);
            }
            if let Some(session_id) = update.session_id {
                step.session_id = Some(session_id);
            }
            if let Some(tokens) = update.output_tokens {
                step.output_tokens = Some(tokens);
            }
            if let Some(started) = update.started_at {
                step.started_at = Some(started);
            }
            if let Some(completed) = update.completed_at {
                step.completed_at = Some(completed);
            }
            if let Some(debug_msgs) = update.debug_messages {
                step.debug_messages = Some(debug_msgs);
            }
            // Clearing wins over setting: an update that resolves the gate and
            // reschedules at once would otherwise leave a live timer behind.
            if update.clear_approval_deadline {
                step.approval_deadline_at = None;
            } else if let Some(deadline) = update.approval_deadline_at {
                step.approval_deadline_at = Some(deadline);
            }
            if let Some(stage) = update.approval_stage {
                step.approval_stage = stage;
            }
            if let Some(assignee) = update.approval_assignee {
                step.approval_assignee = Some(assignee);
            }
            if let Some(requirement) = update.approval_requirement {
                step.approval_requirement = Some(requirement);
            }

            step.updated_at = now;
            Ok(())
        })
    }

    fn get_step(&self, id: Uuid) -> StoreFuture<'_, Option<Step>> {
        Box::pin(async move {
            let state = self.state.read().await;
            Ok(state.steps.get(&id).cloned())
        })
    }

    fn record_step_approval(&self, step_id: Uuid, approval: StepApproval) -> StoreFuture<'_, Step> {
        Box::pin(async move {
            let mut state = self.state.write().await;
            let step = state
                .steps
                .get_mut(&step_id)
                .ok_or(StoreError::StepNotFound(step_id))?;

            if !step.approvals.iter().any(|a| a.user_id == approval.user_id) {
                step.approvals.push(approval);
                step.updated_at = Utc::now();
            }
            Ok(step.clone())
        })
    }

    fn list_steps(&self, run_id: Uuid) -> StoreFuture<'_, Vec<Step>> {
        Box::pin(async move {
            let state = self.state.read().await;
            let mut steps: Vec<Step> = state
                .steps
                .values()
                .filter(|s| s.run_id == run_id)
                .cloned()
                .collect();
            steps.sort_by_key(|s| s.position);
            Ok(steps)
        })
    }

    fn get_stats(&self, filter: RunFilter) -> StoreFuture<'_, RunStats> {
        Box::pin(async move {
            let state = self.state.read().await;

            let mut total_cost_usd = Decimal::ZERO;
            let mut total_duration_ms = 0u64;
            let mut total_runs = 0u64;
            let mut completed_runs = 0u64;
            let mut failed_runs = 0u64;
            let mut cancelled_runs = 0u64;
            let mut active_runs = 0u64;
            let mut awaiting_approval_runs = 0u64;

            for run in state.runs.values() {
                if !run_matches_filter(run, &filter, &state.steps) {
                    continue;
                }

                total_cost_usd += run.cost_usd;
                total_duration_ms += run.duration_ms;
                total_runs += 1;

                match run.status.state {
                    RunStatus::Completed | RunStatus::Warning => completed_runs += 1,
                    RunStatus::Failed => failed_runs += 1,
                    RunStatus::Cancelled => cancelled_runs += 1,
                    RunStatus::AwaitingApproval => {
                        active_runs += 1;
                        awaiting_approval_runs += 1;
                    }
                    RunStatus::Pending
                    | RunStatus::Running
                    | RunStatus::Retrying
                    | RunStatus::Sleeping => {
                        active_runs += 1;
                    }
                }
            }

            Ok(RunStats {
                total_runs,
                completed_runs,
                failed_runs,
                cancelled_runs,
                active_runs,
                awaiting_approval_runs,
                total_cost_usd,
                total_duration_ms,
            })
        })
    }

    fn get_stats_history(
        &self,
        filter: StatsHistoryFilter,
    ) -> StoreFuture<'_, Vec<StatsHistoryBucket>> {
        Box::pin(async move {
            let state = self.state.read().await;
            let now = Utc::now();
            let start = now - Duration::hours(filter.period.hours());
            let run_filter = filter.to_run_filter();
            let buckets = aggregate_history_buckets(
                state
                    .runs
                    .values()
                    .filter(|r| run_matches_filter(r, &run_filter, &state.steps)),
                start,
                now,
                filter.granularity,
            );
            Ok(buckets)
        })
    }

    fn create_step_dependencies(&self, deps: Vec<NewStepDependency>) -> StoreFuture<'_, ()> {
        Box::pin(async move {
            let mut state = self.state.write().await;

            for dep in deps {
                if !state.steps.contains_key(&dep.step_id) {
                    return Err(StoreError::StepNotFound(dep.step_id));
                }
                if !state.steps.contains_key(&dep.depends_on) {
                    return Err(StoreError::StepNotFound(dep.depends_on));
                }

                let already_exists = state
                    .step_dependencies
                    .iter()
                    .any(|d| d.step_id == dep.step_id && d.depends_on == dep.depends_on);

                if !already_exists {
                    state.step_dependencies.push(StepDependency {
                        step_id: dep.step_id,
                        depends_on: dep.depends_on,
                        created_at: Utc::now(),
                    });
                }
            }

            Ok(())
        })
    }

    fn list_step_dependencies(&self, run_id: Uuid) -> StoreFuture<'_, Vec<StepDependency>> {
        Box::pin(async move {
            let state = self.state.read().await;

            let run_step_ids: std::collections::HashSet<Uuid> = state
                .steps
                .values()
                .filter(|s| s.run_id == run_id)
                .map(|s| s.id)
                .collect();

            let mut deps: Vec<StepDependency> = state
                .step_dependencies
                .iter()
                .filter(|d| run_step_ids.contains(&d.step_id))
                .cloned()
                .collect();

            deps.sort_by_key(|d| d.created_at);
            Ok(deps)
        })
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::time::Duration;

    use chrono::TimeDelta;
    use serde_json::json;
    use tokio::spawn;
    use tokio::time::sleep;

    use super::*;
    use crate::api_key_store::ApiKeyStore;
    use crate::entities::{
        ApiKeyScope, ApiKeyUpdate, ApprovalRequirement, NewApiKey, NewUser, TriggerKind,
    };
    use crate::user_store::UserStore;

    use crate::memory::tests::{create_terminal_run, new_run_req};
    use crate::store::RunStore;

    use crate::entities::{StepKind, step_trace_id};

    fn new_step_req(run_id: Uuid, name: &str, position: u32) -> NewStep {
        NewStep {
            run_id,
            trace_id: step_trace_id(run_id, name, position),
            name: name.to_string(),
            kind: StepKind::Shell,
            position,
            input: None,
            is_error_handler: false,
        }
    }

    // ---- create_run ----

    #[tokio::test]
    async fn create_run_returns_pending_status() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();
        assert_eq!(run.status.state, RunStatus::Pending);
        assert_eq!(run.workflow_name, "test");
        assert_eq!(run.retry_count, 0);
        assert_eq!(run.max_retries, 3);
    }

    #[tokio::test]
    async fn create_run_generates_unique_ids() {
        let store = InMemoryStore::new();
        let r1 = store.create_run(new_run_req("a")).await.unwrap().into_run();
        let r2 = store.create_run(new_run_req("b")).await.unwrap().into_run();
        assert_ne!(r1.id, r2.id);
    }

    // ---- get_run ----

    #[tokio::test]
    async fn get_run_returns_created_run() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();
        let fetched = store.get_run(run.id).await.unwrap();
        assert!(fetched.is_some());
        assert_eq!(fetched.unwrap().id, run.id);
    }

    #[tokio::test]
    async fn get_run_returns_none_for_missing() {
        let store = InMemoryStore::new();
        let fetched = store.get_run(Uuid::nil()).await.unwrap();
        assert!(fetched.is_none());
    }

    // ---- update_run_status ----

    #[tokio::test]
    async fn update_run_status_valid_transition() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();

        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(fetched.status.state, RunStatus::Running);
        assert!(fetched.started_at.is_some());
    }

    #[tokio::test]
    async fn update_run_status_invalid_transition_returns_error() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let result = store.update_run_status(run.id, RunStatus::Completed).await;
        assert!(result.is_err());

        let err = result.unwrap_err();
        assert!(matches!(err, StoreError::InvalidTransition { .. }));
    }

    #[tokio::test]
    async fn update_run_status_not_found() {
        let store = InMemoryStore::new();
        let result = store
            .update_run_status(Uuid::nil(), RunStatus::Running)
            .await;
        assert!(matches!(result.unwrap_err(), StoreError::RunNotFound(_)));
    }

    #[tokio::test]
    async fn update_run_status_terminal_sets_completed_at() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run_status(run.id, RunStatus::Completed)
            .await
            .unwrap();

        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(fetched.status.state, RunStatus::Completed);
        assert!(fetched.completed_at.is_some());
    }

    #[tokio::test]
    async fn update_run_status_terminal_to_same_is_idempotent() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run_status(run.id, RunStatus::Failed)
            .await
            .unwrap();

        let before = store.get_run(run.id).await.unwrap().unwrap();
        let completed_at_before = before.completed_at;

        store
            .update_run_status(run.id, RunStatus::Failed)
            .await
            .unwrap();

        let after = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(after.status.state, RunStatus::Failed);
        assert_eq!(after.completed_at, completed_at_before);
    }

    #[tokio::test]
    async fn update_run_terminal_to_same_via_update_run_is_idempotent() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Failed),
                    error: Some("first failure".to_string()),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let before = store.get_run(run.id).await.unwrap().unwrap();

        store
            .update_run(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Failed),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let after = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(after.status.state, RunStatus::Failed);
        assert_eq!(after.completed_at, before.completed_at);
        assert_eq!(after.error, Some("first failure".to_string()));
    }

    // ---- list_runs ----

    #[tokio::test]
    async fn list_runs_empty_store() {
        let store = InMemoryStore::new();
        let page = store.list_runs(RunFilter::default(), 1, 20).await.unwrap();
        assert_eq!(page.total, 0);
        assert!(page.items.is_empty());
    }

    #[tokio::test]
    async fn list_runs_with_workflow_filter() {
        let store = InMemoryStore::new();
        store
            .create_run(new_run_req("deploy"))
            .await
            .unwrap()
            .into_run();
        store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();
        store
            .create_run(new_run_req("deploy"))
            .await
            .unwrap()
            .into_run();

        let filter = RunFilter {
            workflow_name: Some("deploy".to_string()),
            ..RunFilter::default()
        };
        let page = store.list_runs(filter, 1, 20).await.unwrap();
        assert_eq!(page.total, 2);
        assert!(page.items.iter().all(|r| r.workflow_name == "deploy"));
    }

    #[tokio::test]
    async fn list_runs_with_status_filter() {
        let store = InMemoryStore::new();
        let run = store.create_run(new_run_req("a")).await.unwrap().into_run();
        store.create_run(new_run_req("b")).await.unwrap().into_run();

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();

        let filter = RunFilter {
            status: Some(RunStatus::Running),
            ..RunFilter::default()
        };
        let page = store.list_runs(filter, 1, 20).await.unwrap();
        assert_eq!(page.total, 1);
        assert_eq!(page.items[0].id, run.id);
    }

    #[tokio::test]
    async fn list_runs_pagination() {
        let store = InMemoryStore::new();
        for i in 0..5 {
            store
                .create_run(new_run_req(&format!("wf-{i}")))
                .await
                .unwrap()
                .into_run();
        }

        let page1 = store.list_runs(RunFilter::default(), 1, 2).await.unwrap();
        assert_eq!(page1.total, 5);
        assert_eq!(page1.items.len(), 2);
        assert_eq!(page1.page, 1);
        assert_eq!(page1.per_page, 2);

        let page2 = store.list_runs(RunFilter::default(), 2, 2).await.unwrap();
        assert_eq!(page2.items.len(), 2);

        let page3 = store.list_runs(RunFilter::default(), 3, 2).await.unwrap();
        assert_eq!(page3.items.len(), 1);
    }

    // ---- worker lease ----

    fn lease(worker_id: &str, ttl_secs: u64) -> Option<LeaseRequest> {
        Some(LeaseRequest {
            worker_id: worker_id.to_string(),
            ttl: Duration::from_secs(ttl_secs),
        })
    }

    /// Pick a run holding a lease that is already in the past.
    ///
    /// The short sleep matters: `Utc::now()` has microsecond resolution, so a
    /// sub-microsecond TTL can still read as "not yet expired" in the same tick.
    async fn pick_with_expired_lease(store: &InMemoryStore, max_retries: u32) -> Run {
        let mut req = new_run_req("test");
        req.max_retries = max_retries;
        store.create_run(req).await.unwrap();
        let picked = expire_now(store).await;
        picked.expect("a pending run was just created")
    }

    /// Pick the next pending run with a lease that is expired by the time this
    /// returns.
    async fn expire_now(store: &InMemoryStore) -> Option<Run> {
        let picked = store
            .pick_next_pending(Some(LeaseRequest {
                worker_id: "worker-1".to_string(),
                ttl: Duration::from_nanos(1),
            }))
            .await
            .unwrap();
        sleep(Duration::from_millis(2)).await;
        picked
    }

    #[tokio::test]
    async fn pick_next_pending_attaches_lease() {
        let store = InMemoryStore::new();
        store.create_run(new_run_req("test")).await.unwrap();

        let picked = store
            .pick_next_pending(lease("worker-1", 90))
            .await
            .unwrap()
            .unwrap();

        assert_eq!(picked.worker_id.as_deref(), Some("worker-1"));
        let expires = picked.lease_expires_at.expect("lease set");
        assert!(expires > Utc::now());
        assert!(expires <= Utc::now() + TimeDelta::seconds(91));
    }

    #[tokio::test]
    async fn pick_next_pending_without_lease_leaves_run_unowned() {
        let store = InMemoryStore::new();
        store.create_run(new_run_req("test")).await.unwrap();

        let picked = store.pick_next_pending(None).await.unwrap().unwrap();

        assert!(picked.worker_id.is_none());
        assert!(picked.lease_expires_at.is_none());
    }

    #[tokio::test]
    async fn renew_lease_extends_expiry_for_owner() {
        let store = InMemoryStore::new();
        store.create_run(new_run_req("test")).await.unwrap();
        let picked = store
            .pick_next_pending(lease("worker-1", 1))
            .await
            .unwrap()
            .unwrap();

        let renewed = store
            .renew_lease(picked.id, lease("worker-1", 90).unwrap())
            .await
            .unwrap();

        assert!(renewed > picked.lease_expires_at.unwrap());
        let after = store.get_run(picked.id).await.unwrap().unwrap();
        assert_eq!(after.lease_expires_at, Some(renewed));
    }

    #[tokio::test]
    async fn renew_lease_rejects_other_worker() {
        let store = InMemoryStore::new();
        store.create_run(new_run_req("test")).await.unwrap();
        let picked = store
            .pick_next_pending(lease("worker-1", 90))
            .await
            .unwrap()
            .unwrap();

        let err = store
            .renew_lease(picked.id, lease("worker-2", 90).unwrap())
            .await
            .unwrap_err();

        assert!(matches!(
            err,
            StoreError::LeaseLost { held_by: Some(ref w), .. } if w == "worker-1"
        ));
    }

    #[tokio::test]
    async fn renew_lease_rejects_run_that_left_running() {
        let store = InMemoryStore::new();
        store.create_run(new_run_req("test")).await.unwrap();
        let picked = store
            .pick_next_pending(lease("worker-1", 90))
            .await
            .unwrap()
            .unwrap();
        store
            .update_run_status(picked.id, RunStatus::Cancelled)
            .await
            .unwrap();

        let err = store
            .renew_lease(picked.id, lease("worker-1", 90).unwrap())
            .await
            .unwrap_err();

        assert!(matches!(err, StoreError::LeaseLost { .. }));
    }

    #[tokio::test]
    async fn renew_lease_on_unknown_run_is_not_found() {
        let store = InMemoryStore::new();

        let err = store
            .renew_lease(Uuid::now_v7(), lease("worker-1", 90).unwrap())
            .await
            .unwrap_err();

        assert!(matches!(err, StoreError::RunNotFound(_)));
    }

    #[tokio::test]
    async fn leaving_running_clears_the_lease() {
        for target in [
            RunStatus::Completed,
            RunStatus::Retrying,
            RunStatus::AwaitingApproval,
        ] {
            let store = InMemoryStore::new();
            store.create_run(new_run_req("test")).await.unwrap();
            let picked = store
                .pick_next_pending(lease("worker-1", 90))
                .await
                .unwrap()
                .unwrap();

            store.update_run_status(picked.id, target).await.unwrap();

            let after = store.get_run(picked.id).await.unwrap().unwrap();
            assert!(after.worker_id.is_none(), "worker_id kept for {target}");
            assert!(
                after.lease_expires_at.is_none(),
                "lease_expires_at kept for {target}"
            );
        }
    }

    #[tokio::test]
    async fn update_run_to_terminal_clears_the_lease() {
        let store = InMemoryStore::new();
        store.create_run(new_run_req("test")).await.unwrap();
        let picked = store
            .pick_next_pending(lease("worker-1", 90))
            .await
            .unwrap()
            .unwrap();

        store
            .update_run(
                picked.id,
                RunUpdate {
                    status: Some(RunStatus::Failed),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let after = store.get_run(picked.id).await.unwrap().unwrap();
        assert!(after.worker_id.is_none());
        assert!(after.lease_expires_at.is_none());
    }

    #[tokio::test]
    async fn update_run_lease_set_with_running_status_attaches_lease() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();
        let expires_at = Utc::now() + TimeDelta::seconds(60);

        store
            .update_run(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Running),
                    lease: Some(LeaseUpdate::Set {
                        worker_id: "worker-1".to_string(),
                        expires_at,
                    }),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let after = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(after.status.state, RunStatus::Running);
        assert_eq!(after.worker_id.as_deref(), Some("worker-1"));
        assert_eq!(after.lease_expires_at, Some(expires_at));
        // The transferred lease is renewable by its holder like a picked one.
        store
            .renew_lease(run.id, lease("worker-1", 90).unwrap())
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn update_run_lease_release_keeps_run_running_without_lease() {
        let store = InMemoryStore::new();
        store.create_run(new_run_req("test")).await.unwrap();
        let picked = store
            .pick_next_pending(lease("worker-1", 90))
            .await
            .unwrap()
            .unwrap();

        store
            .update_run(
                picked.id,
                RunUpdate {
                    lease: Some(LeaseUpdate::Release),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let after = store.get_run(picked.id).await.unwrap().unwrap();
        assert_eq!(after.status.state, RunStatus::Running);
        assert!(after.worker_id.is_none());
        assert!(after.lease_expires_at.is_none());
        let err = store
            .renew_lease(picked.id, lease("worker-1", 90).unwrap())
            .await
            .unwrap_err();
        assert!(matches!(err, StoreError::LeaseLost { .. }));
    }

    #[tokio::test]
    async fn update_run_lease_none_leaves_lease_untouched() {
        let store = InMemoryStore::new();
        store.create_run(new_run_req("test")).await.unwrap();
        let picked = store
            .pick_next_pending(lease("worker-1", 90))
            .await
            .unwrap()
            .unwrap();

        store
            .update_run(
                picked.id,
                RunUpdate {
                    cost_usd: Some(Decimal::new(150, 2)),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let after = store.get_run(picked.id).await.unwrap().unwrap();
        assert_eq!(after.worker_id.as_deref(), Some("worker-1"));
        assert_eq!(after.lease_expires_at, picked.lease_expires_at);
    }

    // ---- reap_expired_leases ----

    #[tokio::test]
    async fn reap_expired_leases_empty_store() {
        let store = InMemoryStore::new();
        assert!(store.reap_expired_leases(100).await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn reap_expired_leases_requeues_run() {
        let store = InMemoryStore::new();
        let picked = pick_with_expired_lease(&store, 3).await;

        let reaped = store.reap_expired_leases(100).await.unwrap();

        assert_eq!(reaped.len(), 1);
        assert_eq!(reaped[0].from, RunStatus::Running);
        assert_eq!(reaped[0].to, RunStatus::Pending);

        let after = store.get_run(picked.id).await.unwrap().unwrap();
        assert_eq!(after.status.state, RunStatus::Pending);
        assert_eq!(after.retry_count, 0);
        assert_eq!(after.lease_recoveries, 1);
        assert!(after.error.is_none());
        assert!(after.worker_id.is_none());
        assert!(after.lease_expires_at.is_none());
    }

    #[tokio::test]
    async fn reap_expired_leases_requeued_run_is_pickable_again() {
        let store = InMemoryStore::new();
        let picked = pick_with_expired_lease(&store, 3).await;
        store.reap_expired_leases(100).await.unwrap();

        let repicked = store
            .pick_next_pending(lease("worker-2", 90))
            .await
            .unwrap()
            .unwrap();

        assert_eq!(repicked.id, picked.id);
        assert_eq!(repicked.worker_id.as_deref(), Some("worker-2"));
    }

    #[tokio::test]
    async fn reap_expired_leases_ignores_valid_lease() {
        let store = InMemoryStore::new();
        store.create_run(new_run_req("test")).await.unwrap();
        let picked = store
            .pick_next_pending(lease("worker-1", 90))
            .await
            .unwrap()
            .unwrap();

        assert!(store.reap_expired_leases(100).await.unwrap().is_empty());

        let after = store.get_run(picked.id).await.unwrap().unwrap();
        assert_eq!(after.status.state, RunStatus::Running);
        assert_eq!(after.retry_count, 0);
    }

    #[tokio::test]
    async fn reap_expired_leases_ignores_run_without_lease() {
        let store = InMemoryStore::new();
        store.create_run(new_run_req("test")).await.unwrap();
        let picked = store.pick_next_pending(None).await.unwrap().unwrap();

        assert!(store.reap_expired_leases(100).await.unwrap().is_empty());

        let after = store.get_run(picked.id).await.unwrap().unwrap();
        assert_eq!(after.status.state, RunStatus::Running);
    }

    #[tokio::test]
    async fn reap_expired_leases_fails_run_when_retries_exhausted() {
        let store = InMemoryStore::new();
        let picked = pick_with_expired_lease(&store, 0).await;

        let reaped = store.reap_expired_leases(100).await.unwrap();

        assert_eq!(reaped[0].to, RunStatus::Failed);
        let after = store.get_run(picked.id).await.unwrap().unwrap();
        assert_eq!(after.status.state, RunStatus::Failed);
        assert_eq!(after.error.as_deref(), Some(LEASE_EXPIRED_ERROR));
        assert!(after.completed_at.is_some());
    }

    #[tokio::test]
    async fn reap_expired_leases_fails_after_max_retries_recoveries() {
        let store = InMemoryStore::new();
        let picked = pick_with_expired_lease(&store, 2).await;

        // Recovery 1 and 2 requeue, recovery 3 gives up.
        for expected in [RunStatus::Pending, RunStatus::Pending, RunStatus::Failed] {
            let reaped = store.reap_expired_leases(100).await.unwrap();
            assert_eq!(reaped[0].to, expected);
            if expected == RunStatus::Pending {
                expire_now(&store).await;
            }
        }

        let after = store.get_run(picked.id).await.unwrap().unwrap();
        assert_eq!(after.lease_recoveries, 3);
        assert_eq!(after.retry_count, 0);
        assert_eq!(after.error.as_deref(), Some(LEASE_EXPIRED_ERROR));
    }

    #[tokio::test]
    async fn reap_expired_leases_keeps_the_attempt_number() {
        let store = InMemoryStore::new();
        let picked = pick_with_expired_lease(&store, 3).await;
        let before = store
            .create_step(new_step_req(picked.id, "build", 0))
            .await
            .unwrap();

        store.reap_expired_leases(100).await.unwrap();
        let repicked = store
            .pick_next_pending(lease("worker-2", 90))
            .await
            .unwrap()
            .unwrap();
        let after = store
            .create_step(new_step_req(repicked.id, "build", 0))
            .await
            .unwrap();

        assert_eq!(before.attempt, 1);
        assert_eq!(after.attempt, before.attempt);
        assert_eq!(repicked.retry_count, 0);
        assert_eq!(repicked.lease_recoveries, 1);
    }

    #[tokio::test]
    async fn reap_expired_leases_counts_apart_from_handler_retries() {
        let store = InMemoryStore::new();
        let picked = pick_with_expired_lease(&store, 1).await;
        store
            .update_run(
                picked.id,
                RunUpdate {
                    increment_retry: true,
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        // One handler retry already consumed max_retries: the lease recovery
        // budget is separate, so the first reap still requeues.
        let reaped = store.reap_expired_leases(100).await.unwrap();

        assert_eq!(reaped[0].to, RunStatus::Pending);
        assert_eq!(reaped[0].run.retry_count, 1);
        assert_eq!(reaped[0].run.lease_recoveries, 1);
    }

    #[tokio::test]
    async fn reap_expired_leases_respects_limit() {
        let store = InMemoryStore::new();
        for _ in 0..3 {
            pick_with_expired_lease(&store, 3).await;
        }

        let reaped = store.reap_expired_leases(2).await.unwrap();
        assert_eq!(reaped.len(), 2);

        let rest = store.reap_expired_leases(100).await.unwrap();
        assert_eq!(rest.len(), 1);
    }

    // ---- pick_next_pending ----

    #[tokio::test]
    async fn pick_next_pending_empty_store() {
        let store = InMemoryStore::new();
        let result = store.pick_next_pending(None).await.unwrap();
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn pick_next_pending_returns_oldest_and_transitions_to_running() {
        let store = InMemoryStore::new();
        let r1 = store
            .create_run(new_run_req("first"))
            .await
            .unwrap()
            .into_run();
        let _r2 = store
            .create_run(new_run_req("second"))
            .await
            .unwrap()
            .into_run();

        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(picked.id, r1.id);
        assert_eq!(picked.status.state, RunStatus::Running);
        assert!(picked.started_at.is_some());

        // Verify it's Running in the store too.
        let fetched = store.get_run(r1.id).await.unwrap().unwrap();
        assert_eq!(fetched.status.state, RunStatus::Running);
    }

    #[tokio::test]
    async fn pick_next_pending_skips_non_pending() {
        let store = InMemoryStore::new();
        let r1 = store.create_run(new_run_req("a")).await.unwrap().into_run();
        let r2 = store.create_run(new_run_req("b")).await.unwrap().into_run();

        // Transition r1 to Running.
        store
            .update_run_status(r1.id, RunStatus::Running)
            .await
            .unwrap();

        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(picked.id, r2.id);
    }

    // ---- create_step ----

    #[tokio::test]
    async fn create_step_returns_pending() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let step = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "build", 0),
                name: "build".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: Some(json!({"command": "cargo build"})),
                is_error_handler: false,
            })
            .await
            .unwrap();

        assert_eq!(step.status.state, StepStatus::Pending);
        assert_eq!(step.name, "build");
        assert_eq!(step.run_id, run.id);
        assert_eq!(step.position, 0);
    }

    #[tokio::test]
    async fn create_step_for_missing_run_returns_error() {
        let store = InMemoryStore::new();
        let result = store
            .create_step(NewStep {
                run_id: Uuid::nil(),
                trace_id: step_trace_id(Uuid::nil(), "build", 0),
                name: "build".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await;
        assert!(matches!(result.unwrap_err(), StoreError::RunNotFound(_)));
    }

    // ---- update_step ----

    #[tokio::test]
    async fn update_step_applies_partial_update() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let step = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "build", 0),
                name: "build".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        // Transition Pending → Running first
        store
            .update_step(
                step.id,
                StepUpdate {
                    status: Some(StepStatus::Running),
                    ..StepUpdate::default()
                },
            )
            .await
            .unwrap();

        // Then Running → Completed
        store
            .update_step(
                step.id,
                StepUpdate {
                    status: Some(StepStatus::Completed),
                    output: Some(json!({"stdout": "ok"})),
                    duration_ms: Some(150),
                    ..StepUpdate::default()
                },
            )
            .await
            .unwrap();

        let steps = store.list_steps(run.id).await.unwrap();
        assert_eq!(steps.len(), 1);
        assert_eq!(steps[0].status.state, StepStatus::Completed);
        assert_eq!(steps[0].duration_ms, 150);
        assert!(steps[0].output.is_some());
    }

    #[tokio::test]
    async fn update_step_not_found() {
        let store = InMemoryStore::new();
        let result = store.update_step(Uuid::nil(), StepUpdate::default()).await;
        assert!(matches!(result.unwrap_err(), StoreError::StepNotFound(_)));
    }

    // ---- list_steps ----

    #[tokio::test]
    async fn list_steps_ordered_by_position() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        // Insert out of order.
        store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "deploy", 2),
                name: "deploy".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 2,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();
        store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "build", 0),
                name: "build".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();
        store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "test", 1),
                name: "test".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 1,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let steps = store.list_steps(run.id).await.unwrap();
        assert_eq!(steps.len(), 3);
        assert_eq!(steps[0].name, "build");
        assert_eq!(steps[1].name, "test");
        assert_eq!(steps[2].name, "deploy");
    }

    #[tokio::test]
    async fn list_steps_empty_for_run_without_steps() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();
        let steps = store.list_steps(run.id).await.unwrap();
        assert!(steps.is_empty());
    }

    // ---- update_run ----

    #[tokio::test]
    async fn update_run_applies_cost_and_duration() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run(
                run.id,
                RunUpdate {
                    cost_usd: Some(Decimal::new(123, 2)),
                    duration_ms: Some(5000),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(fetched.cost_usd, Decimal::new(123, 2));
        assert_eq!(fetched.duration_ms, 5000);
    }

    #[tokio::test]
    async fn update_run_increment_retry() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();
        assert_eq!(run.retry_count, 0);

        store
            .update_run(
                run.id,
                RunUpdate {
                    increment_retry: true,
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(fetched.retry_count, 1);
    }

    #[tokio::test]
    async fn update_run_not_found() {
        let store = InMemoryStore::new();
        let result = store.update_run(Uuid::nil(), RunUpdate::default()).await;
        assert!(matches!(result.unwrap_err(), StoreError::RunNotFound(_)));
    }

    // ---- concurrent access ----

    #[tokio::test]
    async fn concurrent_pick_next_pending_no_double_pick() {
        let store = InMemoryStore::new();

        // Create 10 pending runs.
        for i in 0..10 {
            store
                .create_run(new_run_req(&format!("wf-{i}")))
                .await
                .unwrap()
                .into_run();
        }

        // Concurrently pick from multiple tasks.
        let mut handles = Vec::new();
        for _ in 0..10 {
            let s = store.clone();
            handles.push(spawn(async move { s.pick_next_pending(None).await }));
        }

        let mut picked_ids = Vec::new();
        for h in handles {
            if let Ok(Ok(Some(run))) = h.await {
                picked_ids.push(run.id);
            }
        }

        // Each run should be picked at most once.
        let unique: std::collections::HashSet<_> = picked_ids.iter().collect();
        assert_eq!(unique.len(), picked_ids.len());
    }

    // ---- get_stats ----

    #[tokio::test]
    async fn get_stats_empty_store() {
        let store = InMemoryStore::new();
        let stats = store.get_stats(RunFilter::default()).await.unwrap();
        assert_eq!(stats.total_runs, 0);
        assert_eq!(stats.completed_runs, 0);
        assert_eq!(stats.failed_runs, 0);
        assert_eq!(stats.cancelled_runs, 0);
        assert_eq!(stats.active_runs, 0);
        assert_eq!(stats.total_cost_usd, Decimal::ZERO);
        assert_eq!(stats.total_duration_ms, 0);
    }

    #[tokio::test]
    async fn get_stats_aggregates_counts_and_totals() {
        let store = InMemoryStore::new();

        // Create runs in various states.
        let r1 = store
            .create_run(new_run_req("wf1"))
            .await
            .unwrap()
            .into_run();
        let r2 = store
            .create_run(new_run_req("wf2"))
            .await
            .unwrap()
            .into_run();
        let r3 = store
            .create_run(new_run_req("wf3"))
            .await
            .unwrap()
            .into_run();
        let _r4 = store
            .create_run(new_run_req("wf4"))
            .await
            .unwrap()
            .into_run();

        // Transition r1 to Running, then Completed.
        store
            .update_run_status(r1.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run_status(r1.id, RunStatus::Completed)
            .await
            .unwrap();

        // Transition r2 to Running, then Failed.
        store
            .update_run_status(r2.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run_status(r2.id, RunStatus::Failed)
            .await
            .unwrap();

        // Transition r3 to Cancelled.
        store
            .update_run_status(r3.id, RunStatus::Cancelled)
            .await
            .unwrap();

        // r4 remains Pending (active).

        // Update cost and duration.
        store
            .update_run(
                r1.id,
                RunUpdate {
                    cost_usd: Some(Decimal::new(1000, 2)),
                    duration_ms: Some(1000),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        store
            .update_run(
                r2.id,
                RunUpdate {
                    cost_usd: Some(Decimal::new(500, 2)),
                    duration_ms: Some(500),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let stats = store.get_stats(RunFilter::default()).await.unwrap();
        assert_eq!(stats.total_runs, 4);
        assert_eq!(stats.completed_runs, 1);
        assert_eq!(stats.failed_runs, 1);
        assert_eq!(stats.cancelled_runs, 1);
        assert_eq!(stats.active_runs, 1); // r4 is Pending
        assert_eq!(stats.total_cost_usd, Decimal::new(1500, 2));
        assert_eq!(stats.total_duration_ms, 1500);
    }

    #[tokio::test]
    async fn update_run_status_running_to_retrying() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();

        store
            .update_run_status(run.id, RunStatus::Retrying)
            .await
            .unwrap();

        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(fetched.status.state, RunStatus::Retrying);
        assert!(!fetched.status.state.is_terminal());
        assert!(fetched.completed_at.is_none()); // Not a terminal state
    }

    #[tokio::test]
    async fn update_run_status_retrying_to_running_allowed() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run_status(run.id, RunStatus::Retrying)
            .await
            .unwrap();

        // Retrying → Running should be allowed
        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();

        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(fetched.status.state, RunStatus::Running);
    }

    #[tokio::test]
    async fn update_run_with_invalid_status_transition_errors() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        // Try to apply invalid status transition via update_run
        let result = store
            .update_run(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Completed), // invalid from Pending
                    ..RunUpdate::default()
                },
            )
            .await;

        assert!(result.is_err());
    }

    #[tokio::test]
    async fn create_step_with_complex_input() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let complex_input = json!({
            "command": "cargo build",
            "env": {
                "RUST_LOG": "debug",
                "CUSTOM": "value"
            },
            "timeout": 60,
            "retry_policy": {
                "max_attempts": 3,
                "backoff": "exponential"
            }
        });

        let step = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "build", 0),
                name: "build".to_string(),
                kind: crate::entities::StepKind::Agent,
                position: 0,
                input: Some(complex_input.clone()),
                is_error_handler: false,
            })
            .await
            .unwrap();

        assert_eq!(step.input, Some(complex_input));
    }

    #[tokio::test]
    async fn update_step_with_error_message() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let step = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "build", 0),
                name: "build".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        store
            .update_step(
                step.id,
                StepUpdate {
                    status: Some(StepStatus::Running),
                    ..StepUpdate::default()
                },
            )
            .await
            .unwrap();

        store
            .update_step(
                step.id,
                StepUpdate {
                    status: Some(StepStatus::Failed),
                    error: Some("Connection timeout after 30s".to_string()),
                    duration_ms: Some(30000),
                    ..StepUpdate::default()
                },
            )
            .await
            .unwrap();

        let steps = store.list_steps(run.id).await.unwrap();
        assert_eq!(steps[0].status.state, StepStatus::Failed);
        assert_eq!(
            steps[0].error,
            Some("Connection timeout after 30s".to_string())
        );
        assert_eq!(steps[0].duration_ms, 30000);
    }

    #[tokio::test]
    async fn list_steps_for_nonexistent_run_returns_empty() {
        let store = InMemoryStore::new();
        let steps = store.list_steps(Uuid::nil()).await.unwrap();
        assert!(steps.is_empty());
    }

    #[tokio::test]
    async fn update_step_pending_to_skipped() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let step = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "build", 0),
                name: "build".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        // Pending → Skipped is allowed (when prior step failed)
        store
            .update_step(
                step.id,
                StepUpdate {
                    status: Some(StepStatus::Skipped),
                    ..StepUpdate::default()
                },
            )
            .await
            .unwrap();

        let steps = store.list_steps(run.id).await.unwrap();
        assert_eq!(steps[0].status.state, StepStatus::Skipped);
    }

    #[tokio::test]
    async fn list_runs_with_combined_filters() {
        let store = InMemoryStore::new();

        let r1 = store
            .create_run(new_run_req("deploy"))
            .await
            .unwrap()
            .into_run();
        let r2 = store
            .create_run(new_run_req("deploy"))
            .await
            .unwrap()
            .into_run();
        let _r3 = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        // r1: Pending → Running → Completed
        store
            .update_run_status(r1.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run_status(r1.id, RunStatus::Completed)
            .await
            .unwrap();

        // r2: Pending → Running
        store
            .update_run_status(r2.id, RunStatus::Running)
            .await
            .unwrap();

        // Filter by workflow AND status
        let filter = RunFilter {
            workflow_name: Some("deploy".to_string()),
            status: Some(RunStatus::Running),
            ..RunFilter::default()
        };

        let page = store.list_runs(filter, 1, 100).await.unwrap();
        assert_eq!(page.total, 1);
        assert_eq!(page.items[0].id, r2.id);
    }

    #[tokio::test]
    async fn list_runs_workflow_filter_is_case_insensitive_partial_match() {
        let store = InMemoryStore::new();
        store
            .create_run(new_run_req("weather-report"))
            .await
            .unwrap()
            .into_run();
        store
            .create_run(new_run_req("deploy-prod"))
            .await
            .unwrap()
            .into_run();

        // Partial match
        let filter = RunFilter {
            workflow_name: Some("weather".to_string()),
            ..RunFilter::default()
        };
        let page = store.list_runs(filter, 1, 100).await.unwrap();
        assert_eq!(page.total, 1);
        assert_eq!(page.items[0].workflow_name, "weather-report");

        // Case-insensitive
        let filter = RunFilter {
            workflow_name: Some("Weather-REPORT".to_string()),
            ..RunFilter::default()
        };
        let page = store.list_runs(filter, 1, 100).await.unwrap();
        assert_eq!(page.total, 1);
        assert_eq!(page.items[0].workflow_name, "weather-report");

        // Suffix match
        let filter = RunFilter {
            workflow_name: Some("report".to_string()),
            ..RunFilter::default()
        };
        let page = store.list_runs(filter, 1, 100).await.unwrap();
        assert_eq!(page.total, 1);
        assert_eq!(page.items[0].workflow_name, "weather-report");

        // No match
        let filter = RunFilter {
            workflow_name: Some("build".to_string()),
            ..RunFilter::default()
        };
        let page = store.list_runs(filter, 1, 100).await.unwrap();
        assert_eq!(page.total, 0);
    }

    #[tokio::test]
    async fn list_runs_has_steps_true_only_filters_completed_and_cancelled() {
        let store = InMemoryStore::new();
        let run_with = create_terminal_run(&store, "with-steps", RunStatus::Completed).await;
        let _run_without = create_terminal_run(&store, "without-steps", RunStatus::Completed).await;

        store
            .create_step(NewStep {
                run_id: run_with.id,
                trace_id: step_trace_id(run_with.id, "build", 0),
                name: "build".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let filter = RunFilter {
            has_steps: Some(true),
            ..RunFilter::default()
        };
        let page = store.list_runs(filter, 1, 100).await.unwrap();
        assert_eq!(page.total, 1);
        assert_eq!(page.items[0].id, run_with.id);
    }

    #[tokio::test]
    async fn list_runs_has_steps_false_only_filters_completed_and_cancelled() {
        let store = InMemoryStore::new();
        let run_with = create_terminal_run(&store, "with-steps", RunStatus::Cancelled).await;
        let run_without = create_terminal_run(&store, "without-steps", RunStatus::Cancelled).await;

        store
            .create_step(NewStep {
                run_id: run_with.id,
                trace_id: step_trace_id(run_with.id, "build", 0),
                name: "build".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let filter = RunFilter {
            has_steps: Some(false),
            ..RunFilter::default()
        };
        let page = store.list_runs(filter, 1, 100).await.unwrap();
        assert_eq!(page.total, 1);
        assert_eq!(page.items[0].id, run_without.id);
    }

    #[tokio::test]
    async fn list_runs_has_steps_none_returns_all() {
        let store = InMemoryStore::new();
        let run_with = store
            .create_run(new_run_req("with-steps"))
            .await
            .unwrap()
            .into_run();
        let _run_without = store
            .create_run(new_run_req("without-steps"))
            .await
            .unwrap()
            .into_run();

        store
            .create_step(NewStep {
                run_id: run_with.id,
                trace_id: step_trace_id(run_with.id, "build", 0),
                name: "build".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let filter = RunFilter {
            has_steps: None,
            ..RunFilter::default()
        };
        let page = store.list_runs(filter, 1, 100).await.unwrap();
        assert_eq!(page.total, 2);
    }

    #[tokio::test]
    async fn list_runs_has_steps_true_does_not_filter_non_terminal_runs() {
        let store = InMemoryStore::new();
        let pending_run = store
            .create_run(new_run_req("pending-empty"))
            .await
            .unwrap()
            .into_run();
        let running_run = store
            .create_run(new_run_req("running-empty"))
            .await
            .unwrap()
            .into_run();
        store
            .update_run_status(running_run.id, RunStatus::Running)
            .await
            .unwrap();

        let filter = RunFilter {
            has_steps: Some(true),
            ..RunFilter::default()
        };
        let page = store.list_runs(filter, 1, 100).await.unwrap();
        assert_eq!(page.total, 2);
        let ids: Vec<_> = page.items.iter().map(|r| r.id).collect();
        assert!(ids.contains(&pending_run.id));
        assert!(ids.contains(&running_run.id));
    }

    #[tokio::test]
    async fn get_stats_with_mixed_active_statuses() {
        let store = InMemoryStore::new();

        let _r1 = store
            .create_run(new_run_req("wf"))
            .await
            .unwrap()
            .into_run(); // Pending
        let r2 = store
            .create_run(new_run_req("wf"))
            .await
            .unwrap()
            .into_run();
        let r3 = store
            .create_run(new_run_req("wf"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run_status(r2.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run_status(r3.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run_status(r3.id, RunStatus::Retrying)
            .await
            .unwrap();

        let r4 = store
            .create_run(new_run_req("wf"))
            .await
            .unwrap()
            .into_run();
        store
            .update_run_status(r4.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run_status(r4.id, RunStatus::AwaitingApproval)
            .await
            .unwrap();

        let r5 = store
            .create_run(new_run_req("wf"))
            .await
            .unwrap()
            .into_run();
        store
            .update_run_status(r5.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run_status(r5.id, RunStatus::Sleeping)
            .await
            .unwrap();

        let stats = store.get_stats(RunFilter::default()).await.unwrap();
        // _r1 (Pending), r2 (Running), r3 (Retrying), r4 (AwaitingApproval), r5 (Sleeping)
        assert_eq!(stats.active_runs, 5);
        assert_eq!(stats.awaiting_approval_runs, 1);
    }

    #[tokio::test]
    async fn run_with_different_trigger_kinds() {
        let store = InMemoryStore::new();

        let r1 = store
            .create_run(NewRun {
                created_by: None,
                workflow_name: "test".to_string(),
                trigger: TriggerKind::Manual,
                payload: json!({}),
                max_retries: 1,
                handler_version: None,
                labels: HashMap::new(),
                scheduled_at: None,
                idempotency_key: None,
                concurrency_key: None,
                priority: 0,
                concurrency_limits: Vec::new(),
                max_cost_usd: None,
                worker_tags: Vec::new(),
            })
            .await
            .unwrap()
            .into_run();

        let r2 = store
            .create_run(NewRun {
                created_by: None,
                workflow_name: "test".to_string(),
                trigger: TriggerKind::Webhook {
                    path: "/hooks/github".to_string(),
                },
                payload: json!({}),
                max_retries: 1,
                handler_version: None,
                labels: HashMap::new(),
                scheduled_at: None,
                idempotency_key: None,
                concurrency_key: None,
                priority: 0,
                concurrency_limits: Vec::new(),
                max_cost_usd: None,
                worker_tags: Vec::new(),
            })
            .await
            .unwrap()
            .into_run();

        let r3 = store
            .create_run(NewRun {
                created_by: None,
                workflow_name: "test".to_string(),
                trigger: TriggerKind::Cron {
                    schedule: "0 0 * * *".to_string(),
                },
                payload: json!({}),
                max_retries: 1,
                handler_version: None,
                labels: HashMap::new(),
                scheduled_at: None,
                idempotency_key: None,
                concurrency_key: None,
                priority: 0,
                concurrency_limits: Vec::new(),
                max_cost_usd: None,
                worker_tags: Vec::new(),
            })
            .await
            .unwrap()
            .into_run();

        let r4 = store
            .create_run(NewRun {
                created_by: None,
                workflow_name: "test".to_string(),
                trigger: TriggerKind::Api,
                payload: json!({}),
                max_retries: 1,
                handler_version: None,
                labels: HashMap::new(),
                scheduled_at: None,
                idempotency_key: None,
                concurrency_key: None,
                priority: 0,
                concurrency_limits: Vec::new(),
                max_cost_usd: None,
                worker_tags: Vec::new(),
            })
            .await
            .unwrap()
            .into_run();

        let r5 = store
            .create_run(NewRun {
                created_by: None,
                workflow_name: "test".to_string(),
                trigger: TriggerKind::Retry {
                    parent_run_id: Uuid::nil(),
                },
                payload: json!({}),
                max_retries: 1,
                handler_version: None,
                labels: HashMap::new(),
                scheduled_at: None,
                idempotency_key: None,
                concurrency_key: None,
                priority: 0,
                concurrency_limits: Vec::new(),
                max_cost_usd: None,
                worker_tags: Vec::new(),
            })
            .await
            .unwrap()
            .into_run();

        assert_eq!(r1.trigger, TriggerKind::Manual);
        assert!(matches!(r2.trigger, TriggerKind::Webhook { .. }));
        assert!(matches!(r3.trigger, TriggerKind::Cron { .. }));
        assert_eq!(r4.trigger, TriggerKind::Api);
        assert!(matches!(r5.trigger, TriggerKind::Retry { .. }));
    }

    // ---- create_step_dependencies ----

    #[tokio::test]
    async fn create_step_dependencies_stores_dependencies() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let step1 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step1", 0),
                name: "step1".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let step2 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step2", 1),
                name: "step2".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 1,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let result = store
            .create_step_dependencies(vec![NewStepDependency {
                step_id: step2.id,
                depends_on: step1.id,
            }])
            .await;

        assert!(result.is_ok());

        let deps = store.list_step_dependencies(run.id).await.unwrap();
        assert_eq!(deps.len(), 1);
        assert_eq!(deps[0].step_id, step2.id);
        assert_eq!(deps[0].depends_on, step1.id);
    }

    #[tokio::test]
    async fn create_step_dependencies_duplicate_dependencies_are_idempotent() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let step1 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step1", 0),
                name: "step1".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let step2 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step2", 1),
                name: "step2".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 1,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let dep = NewStepDependency {
            step_id: step2.id,
            depends_on: step1.id,
        };

        store
            .create_step_dependencies(vec![dep.clone()])
            .await
            .unwrap();
        store.create_step_dependencies(vec![dep]).await.unwrap();

        let deps = store.list_step_dependencies(run.id).await.unwrap();
        assert_eq!(deps.len(), 1);
    }

    #[tokio::test]
    async fn create_step_dependencies_missing_step_id_returns_error() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let step1 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step1", 0),
                name: "step1".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let result = store
            .create_step_dependencies(vec![NewStepDependency {
                step_id: Uuid::nil(),
                depends_on: step1.id,
            }])
            .await;

        assert!(matches!(result.unwrap_err(), StoreError::StepNotFound(_)));
    }

    #[tokio::test]
    async fn create_step_dependencies_missing_depends_on_returns_error() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let step1 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step1", 0),
                name: "step1".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let result = store
            .create_step_dependencies(vec![NewStepDependency {
                step_id: step1.id,
                depends_on: Uuid::nil(),
            }])
            .await;

        assert!(matches!(result.unwrap_err(), StoreError::StepNotFound(_)));
    }

    #[tokio::test]
    async fn create_step_dependencies_multiple_dependencies() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let step1 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step1", 0),
                name: "step1".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let step2 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step2", 1),
                name: "step2".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 1,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let step3 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step3", 2),
                name: "step3".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 2,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let result = store
            .create_step_dependencies(vec![
                NewStepDependency {
                    step_id: step2.id,
                    depends_on: step1.id,
                },
                NewStepDependency {
                    step_id: step3.id,
                    depends_on: step2.id,
                },
            ])
            .await;

        assert!(result.is_ok());

        let deps = store.list_step_dependencies(run.id).await.unwrap();
        assert_eq!(deps.len(), 2);
    }

    // ---- list_step_dependencies ----

    #[tokio::test]
    async fn list_step_dependencies_empty_for_run_with_no_dependencies() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step1", 0),
                name: "step1".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let deps = store.list_step_dependencies(run.id).await.unwrap();
        assert!(deps.is_empty());
    }

    #[tokio::test]
    async fn list_step_dependencies_returns_only_deps_for_given_run() {
        let store = InMemoryStore::new();
        let run1 = store
            .create_run(new_run_req("test1"))
            .await
            .unwrap()
            .into_run();
        let run2 = store
            .create_run(new_run_req("test2"))
            .await
            .unwrap()
            .into_run();

        let step1_run1 = store
            .create_step(NewStep {
                run_id: run1.id,
                trace_id: step_trace_id(run1.id, "step1", 0),
                name: "step1".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let step2_run1 = store
            .create_step(NewStep {
                run_id: run1.id,
                trace_id: step_trace_id(run1.id, "step2", 1),
                name: "step2".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 1,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let step1_run2 = store
            .create_step(NewStep {
                run_id: run2.id,
                trace_id: step_trace_id(run2.id, "step1", 0),
                name: "step1".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let step2_run2 = store
            .create_step(NewStep {
                run_id: run2.id,
                trace_id: step_trace_id(run2.id, "step2", 1),
                name: "step2".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 1,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        store
            .create_step_dependencies(vec![
                NewStepDependency {
                    step_id: step2_run1.id,
                    depends_on: step1_run1.id,
                },
                NewStepDependency {
                    step_id: step2_run2.id,
                    depends_on: step1_run2.id,
                },
            ])
            .await
            .unwrap();

        let deps_run1 = store.list_step_dependencies(run1.id).await.unwrap();
        let deps_run2 = store.list_step_dependencies(run2.id).await.unwrap();

        assert_eq!(deps_run1.len(), 1);
        assert_eq!(deps_run1[0].step_id, step2_run1.id);
        assert_eq!(deps_run1[0].depends_on, step1_run1.id);

        assert_eq!(deps_run2.len(), 1);
        assert_eq!(deps_run2[0].step_id, step2_run2.id);
        assert_eq!(deps_run2[0].depends_on, step1_run2.id);
    }

    #[tokio::test]
    async fn list_step_dependencies_returns_empty_for_nonexistent_run() {
        let store = InMemoryStore::new();
        let deps = store.list_step_dependencies(Uuid::nil()).await.unwrap();
        assert!(deps.is_empty());
    }

    #[tokio::test]
    async fn list_step_dependencies_sorted_by_created_at() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let step1 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step1", 0),
                name: "step1".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let step2 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step2", 1),
                name: "step2".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 1,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let step3 = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "step3", 2),
                name: "step3".to_string(),
                kind: crate::entities::StepKind::Shell,
                position: 2,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        store
            .create_step_dependencies(vec![NewStepDependency {
                step_id: step2.id,
                depends_on: step1.id,
            }])
            .await
            .unwrap();

        store
            .create_step_dependencies(vec![NewStepDependency {
                step_id: step3.id,
                depends_on: step1.id,
            }])
            .await
            .unwrap();

        let deps = store.list_step_dependencies(run.id).await.unwrap();
        assert_eq!(deps.len(), 2);
        assert!(deps[0].created_at <= deps[1].created_at);
    }

    // ---- update_run_returning ----

    #[tokio::test]
    async fn update_run_returning_applies_and_returns() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        // Transition Pending -> Running first
        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();

        let updated = store
            .update_run_returning(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Completed),
                    cost_usd: Some(Decimal::new(4200, 2)),
                    duration_ms: Some(1500),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        assert_eq!(updated.id, run.id);
        assert_eq!(updated.status.state, RunStatus::Completed);
        assert_eq!(updated.cost_usd, Decimal::new(4200, 2));
        assert_eq!(updated.duration_ms, 1500);
        assert!(updated.completed_at.is_some());
    }

    #[tokio::test]
    async fn update_run_returning_not_found() {
        let store = InMemoryStore::new();
        let result = store
            .update_run_returning(
                Uuid::nil(),
                RunUpdate {
                    status: Some(RunStatus::Running),
                    ..RunUpdate::default()
                },
            )
            .await;

        assert!(matches!(result, Err(StoreError::RunNotFound(_))));
    }

    #[tokio::test]
    async fn update_run_returning_invalid_transition() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        let result = store
            .update_run_returning(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Completed),
                    ..RunUpdate::default()
                },
            )
            .await;

        assert!(matches!(result, Err(StoreError::InvalidTransition { .. })));
    }

    // ---- retry scheduling ----

    #[tokio::test]
    async fn create_step_stamps_the_current_attempt() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("retry-wf"))
            .await
            .unwrap()
            .into_run();

        let first = store
            .create_step(new_step_req(run.id, "build", 0))
            .await
            .unwrap();
        assert_eq!(first.attempt, 1);

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Retrying),
                    increment_retry: true,
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let second = store
            .create_step(new_step_req(run.id, "build", 0))
            .await
            .unwrap();
        assert_eq!(second.attempt, 2);
    }

    #[tokio::test]
    async fn pick_next_pending_ignores_retrying_run_before_its_backoff() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("retry-wf"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Retrying),
                    increment_retry: true,
                    scheduled_at: Some(Utc::now() + TimeDelta::seconds(60)),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        assert!(store.pick_next_pending(None).await.unwrap().is_none());
    }

    #[tokio::test]
    async fn pick_next_pending_resumes_retrying_run_after_its_backoff() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("retry-wf"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Retrying),
                    increment_retry: true,
                    scheduled_at: Some(Utc::now() - TimeDelta::seconds(1)),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(picked.id, run.id);
        assert_eq!(picked.status.state, RunStatus::Running);
        assert_eq!(picked.retry_count, 1);
    }

    // ---- priority ----

    async fn create_with_priority(store: &InMemoryStore, name: &str, priority: i16) -> Run {
        let run = store
            .create_run(NewRun {
                priority,
                ..new_run_req(name)
            })
            .await
            .unwrap()
            .into_run();
        // Distinct created_at values, so FIFO among equal priorities is observable.
        sleep(Duration::from_millis(2)).await;
        run
    }

    #[tokio::test]
    async fn pick_next_pending_priority_serves_higher_priority_first() {
        let store = InMemoryStore::new();
        let low = create_with_priority(&store, "low", 0).await;
        let high = create_with_priority(&store, "high", 10).await;

        let first = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(first.id, high.id);
        assert_eq!(first.priority, 10);
        let second = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(second.id, low.id);
    }

    #[tokio::test]
    async fn pick_next_pending_priority_is_fifo_among_equal_priorities() {
        let store = InMemoryStore::new();
        let older = create_with_priority(&store, "older", 5).await;
        let younger = create_with_priority(&store, "younger", 5).await;

        let first = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(first.id, older.id);
        let second = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(second.id, younger.id);
    }

    #[tokio::test]
    async fn pick_next_pending_priority_negative_runs_after_default() {
        let store = InMemoryStore::new();
        let negative = create_with_priority(&store, "background", -50).await;
        let default = create_with_priority(&store, "default", 0).await;

        let first = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(first.id, default.id);
        let second = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(second.id, negative.id);
    }

    #[tokio::test]
    async fn pick_next_pending_priority_skips_high_priority_run_not_yet_due() {
        let store = InMemoryStore::new();
        let later = store
            .create_run(NewRun {
                priority: 100,
                scheduled_at: Some(Utc::now() + TimeDelta::seconds(3600)),
                ..new_run_req("later")
            })
            .await
            .unwrap()
            .into_run();
        let now = create_with_priority(&store, "now", 0).await;

        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(picked.id, now.id);
        assert!(store.pick_next_pending(None).await.unwrap().is_none());
        let later = store.get_run(later.id).await.unwrap().unwrap();
        assert_eq!(later.status.state, RunStatus::Pending);
    }

    #[tokio::test]
    async fn pick_next_pending_priority_kept_after_retry() {
        let store = InMemoryStore::new();
        let run = create_with_priority(&store, "retry-wf", 42).await;

        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Retrying),
                    increment_retry: true,
                    scheduled_at: Some(Utc::now() - TimeDelta::seconds(1)),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();
        let fresh = create_with_priority(&store, "fresh", 0).await;

        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(picked.id, run.id);
        assert_eq!(picked.priority, 42);
        assert_eq!(picked.retry_count, 1);
        let next = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(next.id, fresh.id);
    }

    #[tokio::test]
    async fn create_run_priority_out_of_range_is_rejected() {
        let store = InMemoryStore::new();
        for priority in [101, -101] {
            let err = store
                .create_run(NewRun {
                    priority,
                    ..new_run_req("out-of-range")
                })
                .await
                .unwrap_err();
            assert!(matches!(err, StoreError::Database(_)), "{err:?}");
        }
        let page = store.list_runs(RunFilter::default(), 1, 10).await.unwrap();
        assert_eq!(page.total, 0);
    }

    #[tokio::test]
    async fn list_runs_priority_filter_is_exact_match() {
        let store = InMemoryStore::new();
        let urgent = create_with_priority(&store, "urgent", 10).await;
        create_with_priority(&store, "default", 0).await;
        create_with_priority(&store, "more-urgent", 20).await;

        let page = store
            .list_runs(
                RunFilter {
                    priority: Some(10),
                    ..RunFilter::default()
                },
                1,
                10,
            )
            .await
            .unwrap();
        assert_eq!(page.total, 1);
        assert_eq!(page.items[0].id, urgent.id);

        let all = store.list_runs(RunFilter::default(), 1, 10).await.unwrap();
        assert_eq!(all.total, 3);
    }

    #[tokio::test]
    async fn update_run_persists_scheduled_at() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();
        let when = Utc::now() + TimeDelta::seconds(30);

        store
            .update_run(
                run.id,
                RunUpdate {
                    scheduled_at: Some(when),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(fetched.scheduled_at, Some(when));
    }

    // ---- output ----

    #[tokio::test]
    async fn new_run_has_no_output() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        assert!(run.output.is_none());
        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert!(fetched.output.is_none());
    }

    #[tokio::test]
    async fn update_run_sets_output() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run(
                run.id,
                RunUpdate {
                    output: Some(json!({"verdict": "approved"})),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(fetched.output, Some(json!({"verdict": "approved"})));
    }

    #[tokio::test]
    async fn update_run_without_output_keeps_previous_output() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        store
            .update_run(
                run.id,
                RunUpdate {
                    output: Some(json!({"verdict": "approved"})),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();
        store
            .update_run(
                run.id,
                RunUpdate {
                    error: Some("boom".to_string()),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();

        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(fetched.output, Some(json!({"verdict": "approved"})));
        assert_eq!(fetched.error.as_deref(), Some("boom"));
    }

    #[tokio::test]
    async fn update_run_output_last_write_wins() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        for verdict in ["first", "second"] {
            store
                .update_run(
                    run.id,
                    RunUpdate {
                        output: Some(json!({ "verdict": verdict })),
                        ..RunUpdate::default()
                    },
                )
                .await
                .unwrap();
        }

        let fetched = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(fetched.output, Some(json!({"verdict": "second"})));
    }

    // ---- created_by ----

    async fn seed_user(store: &InMemoryStore, username: &str) -> Uuid {
        store
            .create_user(NewUser {
                email: format!("{username}@example.com"),
                username: username.to_string(),
                password_hash: "hash".to_string(),
                is_admin: Some(false),
            })
            .await
            .unwrap()
            .id
    }

    async fn seed_api_key(store: &InMemoryStore, user_id: Uuid, name: &str) -> Uuid {
        store
            .create_api_key(NewApiKey {
                user_id,
                name: name.to_string(),
                key_hash: "hash".to_string(),
                key_prefix: "irfl_0000".to_string(),
                scopes: vec![ApiKeyScope::RunsWrite],
                expires_at: None,
                rate_limit_override: None,
            })
            .await
            .unwrap()
            .id
    }

    fn run_req_by(actor: RunActor) -> NewRun {
        NewRun {
            created_by: Some(actor),
            ..new_run_req("test")
        }
    }

    #[tokio::test]
    async fn create_run_without_actor_has_no_author() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();

        assert!(run.created_by.is_none());
        assert!(run.created_by_label.is_none());
    }

    #[tokio::test]
    async fn create_run_by_user_resolves_username_as_label() {
        let store = InMemoryStore::new();
        let user_id = seed_user(&store, "alice").await;

        let run = store
            .create_run(run_req_by(RunActor::User { user_id }))
            .await
            .unwrap()
            .into_run();

        assert_eq!(run.created_by, Some(RunActor::User { user_id }));
        assert_eq!(run.created_by_label.as_deref(), Some("alice"));
    }

    #[tokio::test]
    async fn create_run_by_api_key_resolves_key_and_owner_as_label() {
        let store = InMemoryStore::new();
        let user_id = seed_user(&store, "alice").await;
        let api_key_id = seed_api_key(&store, user_id, "ci-deploy").await;

        let run = store
            .create_run(run_req_by(RunActor::ApiKey {
                api_key_id,
                user_id,
            }))
            .await
            .unwrap()
            .into_run();

        assert_eq!(run.created_by_label.as_deref(), Some("ci-deploy (alice)"));
    }

    #[tokio::test]
    async fn label_follows_api_key_rename() {
        let store = InMemoryStore::new();
        let user_id = seed_user(&store, "alice").await;
        let api_key_id = seed_api_key(&store, user_id, "ci-deploy").await;
        let run = store
            .create_run(run_req_by(RunActor::ApiKey {
                api_key_id,
                user_id,
            }))
            .await
            .unwrap()
            .into_run();

        store
            .update_api_key(
                api_key_id,
                ApiKeyUpdate {
                    name: Some("ci-release".to_string()),
                    ..ApiKeyUpdate::default()
                },
            )
            .await
            .unwrap();

        let reread = store.get_run(run.id).await.unwrap().unwrap();
        assert_eq!(
            reread.created_by_label.as_deref(),
            Some("ci-release (alice)")
        );
    }

    #[tokio::test]
    async fn label_is_none_when_user_is_unknown() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(run_req_by(RunActor::User {
                user_id: Uuid::now_v7(),
            }))
            .await
            .unwrap()
            .into_run();

        assert!(run.created_by.is_some());
        assert!(run.created_by_label.is_none());
    }

    #[tokio::test]
    async fn label_is_key_name_only_when_owner_is_unknown() {
        let store = InMemoryStore::new();
        let owner = seed_user(&store, "alice").await;
        let api_key_id = seed_api_key(&store, owner, "ci-deploy").await;

        // An actor pointing at a user that was never created.
        let run = store
            .create_run(run_req_by(RunActor::ApiKey {
                api_key_id,
                user_id: Uuid::now_v7(),
            }))
            .await
            .unwrap()
            .into_run();

        assert_eq!(run.created_by_label.as_deref(), Some("ci-deploy"));
    }

    #[tokio::test]
    async fn list_runs_filters_by_author() {
        let store = InMemoryStore::new();
        let alice = seed_user(&store, "alice").await;
        let bob = seed_user(&store, "bob").await;

        store
            .create_run(run_req_by(RunActor::User { user_id: alice }))
            .await
            .unwrap()
            .into_run();
        store
            .create_run(run_req_by(RunActor::User { user_id: bob }))
            .await
            .unwrap()
            .into_run();
        store.create_run(new_run_req("anonymous")).await.unwrap();

        let page = store
            .list_runs(
                RunFilter {
                    created_by_user_id: Some(alice),
                    ..RunFilter::default()
                },
                1,
                20,
            )
            .await
            .unwrap();

        assert_eq!(page.total, 1);
        assert_eq!(page.items[0].created_by_label.as_deref(), Some("alice"));
    }

    #[tokio::test]
    async fn list_runs_author_filter_matches_runs_from_the_users_api_keys() {
        let store = InMemoryStore::new();
        let alice = seed_user(&store, "alice").await;
        let api_key_id = seed_api_key(&store, alice, "ci-deploy").await;

        store
            .create_run(run_req_by(RunActor::ApiKey {
                api_key_id,
                user_id: alice,
            }))
            .await
            .unwrap()
            .into_run();

        let page = store
            .list_runs(
                RunFilter {
                    created_by_user_id: Some(alice),
                    ..RunFilter::default()
                },
                1,
                20,
            )
            .await
            .unwrap();

        assert_eq!(page.total, 1);
    }

    #[tokio::test]
    async fn list_runs_author_filter_excludes_unrelated_users() {
        let store = InMemoryStore::new();
        let alice = seed_user(&store, "alice").await;

        store
            .create_run(run_req_by(RunActor::User { user_id: alice }))
            .await
            .unwrap()
            .into_run();

        let page = store
            .list_runs(
                RunFilter {
                    created_by_user_id: Some(Uuid::now_v7()),
                    ..RunFilter::default()
                },
                1,
                20,
            )
            .await
            .unwrap();

        assert_eq!(page.total, 0);
    }

    #[tokio::test]
    async fn list_runs_without_author_filter_returns_every_run() {
        let store = InMemoryStore::new();
        let alice = seed_user(&store, "alice").await;

        store
            .create_run(run_req_by(RunActor::User { user_id: alice }))
            .await
            .unwrap()
            .into_run();
        store.create_run(new_run_req("anonymous")).await.unwrap();

        let page = store.list_runs(RunFilter::default(), 1, 20).await.unwrap();
        assert_eq!(page.total, 2);
    }

    #[tokio::test]
    async fn pick_next_pending_resolves_author_label() {
        let store = InMemoryStore::new();
        let user_id = seed_user(&store, "alice").await;
        store
            .create_run(run_req_by(RunActor::User { user_id }))
            .await
            .unwrap()
            .into_run();

        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
        assert_eq!(picked.created_by_label.as_deref(), Some("alice"));
    }

    // ---- list_purgeable_runs ----

    #[tokio::test]
    async fn list_purgeable_runs_returns_old_terminal_runs() {
        let store = InMemoryStore::new();
        let old = create_terminal_run(&store, "deploy", RunStatus::Completed).await;
        store
            .set_run_created_at(old.id, Utc::now() - chrono::Duration::days(100))
            .await;

        let policy = PurgePolicy {
            max_age_days: 90,
            max_runs_per_workflow: 10000,
            dry_run: false,
        };
        let result = store.list_purgeable_runs(&policy, 100).await.unwrap();

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].run_id, old.id);
        assert_eq!(result[0].reason, PurgeReason::TooOld);
    }

    #[tokio::test]
    async fn list_purgeable_runs_ignores_non_terminal_states() {
        let store = InMemoryStore::new();

        // Pending
        let pending = store
            .create_run(new_run_req("deploy"))
            .await
            .unwrap()
            .into_run();
        store
            .set_run_created_at(pending.id, Utc::now() - chrono::Duration::days(200))
            .await;

        // Running
        let running = store
            .create_run(new_run_req("deploy"))
            .await
            .unwrap()
            .into_run();
        store
            .update_run_status(running.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .set_run_created_at(running.id, Utc::now() - chrono::Duration::days(200))
            .await;

        let policy = PurgePolicy {
            max_age_days: 90,
            max_runs_per_workflow: 1,
            dry_run: false,
        };
        let result = store.list_purgeable_runs(&policy, 100).await.unwrap();
        assert!(result.is_empty());
    }

    #[tokio::test]
    async fn list_purgeable_runs_returns_excess_per_workflow() {
        let store = InMemoryStore::new();
        let r1 = create_terminal_run(&store, "deploy", RunStatus::Completed).await;
        store
            .set_run_created_at(r1.id, Utc::now() - chrono::Duration::days(10))
            .await;
        let r2 = create_terminal_run(&store, "deploy", RunStatus::Completed).await;
        store
            .set_run_created_at(r2.id, Utc::now() - chrono::Duration::days(5))
            .await;
        let _r3 = create_terminal_run(&store, "deploy", RunStatus::Completed).await;

        let policy = PurgePolicy {
            max_age_days: 365,
            max_runs_per_workflow: 2,
            dry_run: false,
        };
        let result = store.list_purgeable_runs(&policy, 100).await.unwrap();

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].run_id, r1.id);
        assert_eq!(result[0].reason, PurgeReason::ExceedsWorkflowLimit);
    }

    // ---- delete_run ----

    #[tokio::test]
    async fn delete_run_removes_run_and_associated_data() {
        use crate::artifact_store::ArtifactStore;
        use crate::entities::{NewStep, StepKind};

        let store = InMemoryStore::new();
        let run = create_terminal_run(&store, "deploy", RunStatus::Completed).await;
        let step = store
            .create_step(NewStep {
                run_id: run.id,
                trace_id: step_trace_id(run.id, "build", 0),
                name: "build".to_string(),
                kind: StepKind::Shell,
                position: 0,
                input: None,
                is_error_handler: false,
            })
            .await
            .unwrap();

        let artifact_id = Uuid::now_v7();
        store
            .create_artifact(crate::entities::NewArtifact {
                id: artifact_id,
                run_id: run.id,
                step_id: step.id,
                name: "report.html".to_string(),
                storage_key: format!("artifacts/{}/{}/{}", run.id, step.id, artifact_id),
                content_type: "text/html".to_string(),
                size_bytes: 42,
                sha256: "0".repeat(64),
            })
            .await
            .unwrap();

        let keys = store.delete_run(run.id).await.unwrap();

        assert_eq!(keys.len(), 1);
        assert!(keys[0].contains(&artifact_id.to_string()));
        assert!(store.get_run(run.id).await.unwrap().is_none());
        assert!(store.list_steps(run.id).await.unwrap().is_empty());
        assert!(
            store
                .list_artifacts_for_run(run.id)
                .await
                .unwrap()
                .is_empty()
        );
    }

    #[tokio::test]
    async fn delete_run_not_found() {
        let store = InMemoryStore::new();
        let err = store.delete_run(Uuid::now_v7()).await.unwrap_err();
        assert!(matches!(err, StoreError::RunNotFound(_)));
    }

    // ---- record_step_approval ----

    fn vote(user_id: Uuid, name: &str) -> StepApproval {
        StepApproval {
            user_id,
            approved_by: name.to_string(),
            at: Utc::now(),
        }
    }

    #[tokio::test]
    async fn record_step_approval_appends_distinct_voters() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();
        let step = store
            .create_step(new_step_req(run.id, "gate", 0))
            .await
            .unwrap();
        assert!(step.approvals.is_empty());
        assert!(step.approval_requirement.is_none());

        let alice = Uuid::now_v7();
        let bob = Uuid::now_v7();
        let after_first = store
            .record_step_approval(step.id, vote(alice, "alice"))
            .await
            .unwrap();
        assert_eq!(after_first.approvals.len(), 1);

        let after_second = store
            .record_step_approval(step.id, vote(bob, "bob"))
            .await
            .unwrap();
        assert_eq!(after_second.approvals.len(), 2);
        assert_eq!(after_second.approvals[0].user_id, alice);
        assert_eq!(after_second.approvals[1].user_id, bob);
    }

    #[tokio::test]
    async fn record_step_approval_ignores_same_user() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();
        let step = store
            .create_step(new_step_req(run.id, "gate", 0))
            .await
            .unwrap();

        let alice = Uuid::now_v7();
        store
            .record_step_approval(step.id, vote(alice, "alice"))
            .await
            .unwrap();
        let again = store
            .record_step_approval(step.id, vote(alice, "alice-key"))
            .await
            .unwrap();

        assert_eq!(again.approvals.len(), 1);
        assert_eq!(again.approvals[0].approved_by, "alice");
    }

    #[tokio::test]
    async fn record_step_approval_unknown_step_is_not_found() {
        let store = InMemoryStore::new();
        let err = store
            .record_step_approval(Uuid::now_v7(), vote(Uuid::now_v7(), "alice"))
            .await
            .unwrap_err();
        assert!(matches!(err, StoreError::StepNotFound(_)));
    }

    #[tokio::test]
    async fn update_step_sets_approval_requirement() {
        let store = InMemoryStore::new();
        let run = store
            .create_run(new_run_req("test"))
            .await
            .unwrap()
            .into_run();
        let step = store
            .create_step(new_step_req(run.id, "gate", 0))
            .await
            .unwrap();
        let requirement = ApprovalRequirement {
            required_approvers: 3,
            ..ApprovalRequirement::default()
        };

        store
            .update_step(
                step.id,
                StepUpdate {
                    approval_requirement: Some(requirement.clone()),
                    ..StepUpdate::default()
                },
            )
            .await
            .unwrap();

        let fetched = store.get_step(step.id).await.unwrap().unwrap();
        assert_eq!(fetched.approval_requirement, Some(requirement));
    }

    async fn sleeping_run(store: &InMemoryStore, scheduled_at: DateTime<Utc>) -> Run {
        let run = store
            .create_run(new_run_req("sleepy"))
            .await
            .unwrap()
            .into_run();
        store
            .update_run_status(run.id, RunStatus::Running)
            .await
            .unwrap();
        store
            .update_run(
                run.id,
                RunUpdate {
                    status: Some(RunStatus::Sleeping),
                    scheduled_at: Some(scheduled_at),
                    ..RunUpdate::default()
                },
            )
            .await
            .unwrap();
        store.get_run(run.id).await.unwrap().unwrap()
    }

    #[tokio::test]
    async fn claim_due_sleeping_runs_requeues_due_runs() {
        let store = InMemoryStore::new();
        let due = sleeping_run(&store, Utc::now() - TimeDelta::seconds(5)).await;

        let woken = store.claim_due_sleeping_runs(10).await.unwrap();
        assert_eq!(woken.len(), 1);
        assert_eq!(woken[0].id, due.id);
        assert_eq!(woken[0].status.state, RunStatus::Pending);
        assert!(woken[0].scheduled_at.is_none());

        let fetched = store.get_run(due.id).await.unwrap().unwrap();
        assert_eq!(fetched.status.state, RunStatus::Pending);
        assert!(fetched.scheduled_at.is_none());

        // Exactly once: a second tick finds nothing.
        assert!(store.claim_due_sleeping_runs(10).await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn claim_due_sleeping_runs_skips_future_runs() {
        let store = InMemoryStore::new();
        let future = sleeping_run(&store, Utc::now() + TimeDelta::hours(1)).await;

        assert!(store.claim_due_sleeping_runs(10).await.unwrap().is_empty());
        let fetched = store.get_run(future.id).await.unwrap().unwrap();
        assert_eq!(fetched.status.state, RunStatus::Sleeping);
    }

    #[tokio::test]
    async fn claim_due_sleeping_runs_honours_limit_oldest_first() {
        let store = InMemoryStore::new();
        let older = sleeping_run(&store, Utc::now() - TimeDelta::seconds(20)).await;
        let newer = sleeping_run(&store, Utc::now() - TimeDelta::seconds(10)).await;

        let woken = store.claim_due_sleeping_runs(1).await.unwrap();
        assert_eq!(woken.len(), 1);
        assert_eq!(woken[0].id, older.id);

        let woken = store.claim_due_sleeping_runs(1).await.unwrap();
        assert_eq!(woken[0].id, newer.id);
    }
}