Skip to main content

ironflow_store/memory/
run_store.rs

1use std::cmp::Reverse;
2use std::collections::{BTreeMap, HashMap, HashSet};
3
4use chrono::{DateTime, Duration, Utc};
5use rust_decimal::Decimal;
6use uuid::Uuid;
7
8use crate::entities::{
9    AgentRoleStats, ApiKey, Assignee, ConcurrencyGroupBacklog, IDEMPOTENCY_WINDOW, LeaseRequest,
10    LeaseUpdate, NewRun, NewStep, NewStepDependency, Page, PurgePolicy, PurgeReason, PurgeableRun,
11    ReapedRun, Run, RunActor, RunCreation, RunFilter, RunStats, RunStatus, RunUpdate,
12    StatsHistoryBucket, StatsHistoryFilter, Step, StepApproval, StepDependency, StepStatus,
13    StepUpdate, TriggerKind, User, WorkerCapabilities, WorkflowPause, agent_role_from_input,
14    normalize_worker_tags, validate_concurrency_limits, validate_priority, validate_worker_tags,
15};
16use crate::error::StoreError;
17use crate::store::{LEASE_EXPIRED_ERROR, RunStore, StoreFuture};
18
19use super::descendants::active_descendants;
20use super::stats_history::aggregate_history_buckets;
21use super::{InMemoryStore, State};
22
23/// Resolve [`Run::created_by_label`] from the current users and API keys.
24///
25/// Mirrors the `LEFT JOIN` the PostgreSQL store performs at read time, so a
26/// renamed API key or a renamed user is reflected immediately.
27fn resolve_created_by_label(
28    actor: Option<&RunActor>,
29    users: &HashMap<Uuid, User>,
30    api_keys: &HashMap<Uuid, ApiKey>,
31) -> Option<String> {
32    let actor = actor?;
33    let username = users.get(&actor.user_id()).map(|u| u.username.clone());
34
35    match actor.api_key_id() {
36        None => username,
37        Some(api_key_id) => {
38            let key_name = api_keys.get(&api_key_id).map(|k| k.name.clone())?;
39            Some(match username {
40                Some(username) => format!("{key_name} ({username})"),
41                None => key_name,
42            })
43        }
44    }
45}
46
47/// Return a clone of `run` with its display label resolved against `state`.
48fn run_with_label(run: &Run, state: &State) -> Run {
49    let mut run = run.clone();
50    run.created_by_label =
51        resolve_created_by_label(run.created_by.as_ref(), &state.users, &state.api_keys);
52    run
53}
54
55/// Drop the worker lease held on a run.
56fn clear_lease(run: &mut Run) {
57    run.worker_id = None;
58    run.lease_expires_at = None;
59}
60
61/// Record the state a run is paused from, or forget it once the run leaves
62/// `Paused`. Called before `run.status.state` takes `next`.
63fn track_resume_status(run: &mut Run, next: RunStatus) {
64    if next == RunStatus::Paused {
65        run.resume_status = Some(run.status.state);
66    } else if run.status.state == RunStatus::Paused {
67        run.resume_status = None;
68    }
69}
70
71/// Number of root runs in `Running` that carry `group`.
72///
73/// Sub-workflow runs execute inside their parent's slot and are never counted.
74fn running_count(runs: &HashMap<Uuid, Run>, group: &str) -> u64 {
75    runs.values()
76        .filter(|r| {
77            r.status.state == RunStatus::Running
78                && !matches!(r.trigger, TriggerKind::Workflow)
79                && r.concurrency_limits.iter().any(|l| l.group == group)
80        })
81        .count() as u64
82}
83
84/// Whether at least one of the run's concurrency groups is saturated for the
85/// run's own limit.
86fn is_blocked(run: &Run, runs: &HashMap<Uuid, Run>) -> bool {
87    run.concurrency_limits
88        .iter()
89        .any(|l| running_count(runs, &l.group) >= u64::from(l.limit))
90}
91
92/// Whether a run waits for a pick: pending or retrying, and due.
93fn is_due(run: &Run, now: DateTime<Utc>) -> bool {
94    matches!(run.status.state, RunStatus::Pending | RunStatus::Retrying)
95        && run.scheduled_at.is_none_or(|at| at <= now)
96}
97
98fn run_matches_filter(run: &Run, filter: &RunFilter, steps: &HashMap<Uuid, Step>) -> bool {
99    if let Some(ref wf) = filter.workflow_name
100        && !run
101            .workflow_name
102            .to_lowercase()
103            .contains(&wf.to_lowercase())
104    {
105        return false;
106    }
107    if let Some(ref status) = filter.status
108        && &run.status.state != status
109    {
110        return false;
111    }
112    if let Some(after) = filter.created_after
113        && run.created_at < after
114    {
115        return false;
116    }
117    if let Some(before) = filter.created_before
118        && run.created_at > before
119    {
120        return false;
121    }
122    if let Some(has_steps) = filter.has_steps
123        && matches!(
124            run.status.state,
125            RunStatus::Completed | RunStatus::Cancelled
126        )
127    {
128        let run_has_steps = steps.values().any(|s| s.run_id == run.id);
129        if has_steps != run_has_steps {
130            return false;
131        }
132    }
133    if let Some(ref labels) = filter.labels {
134        for (key, value) in labels {
135            if run.labels.get(key) != Some(value) {
136                return false;
137            }
138        }
139    }
140    if let Some(user_id) = filter.created_by_user_id
141        && run.created_by.as_ref().map(RunActor::user_id) != Some(user_id)
142    {
143        return false;
144    }
145    if let Some(ref group) = filter.concurrency_group
146        && !run.concurrency_limits.iter().any(|l| &l.group == group)
147    {
148        return false;
149    }
150    if let Some(priority) = filter.priority
151        && run.priority != priority
152    {
153        return false;
154    }
155    if let Some(ref caps) = filter.eligible_for
156        && !caps.can_take(&run.workflow_name, &run.worker_tags)
157    {
158        return false;
159    }
160    true
161}
162
163/// Insert a run into the locked state, honouring its idempotency and
164/// concurrency keys. Writes nothing when it returns an error.
165///
166/// Shared by [`RunStore::create_run`] and the schedule firing, which must
167/// create the run under the same lock as the schedule update.
168pub(super) fn insert_run(state: &mut State, req: NewRun) -> Result<RunCreation, StoreError> {
169    validate_concurrency_limits(&req.concurrency_limits)?;
170    // Mirrors the CHECK constraint of the PostgreSQL column.
171    validate_priority(req.priority).map_err(StoreError::Database)?;
172    validate_worker_tags(&req.worker_tags)?;
173    let now = Utc::now();
174
175    if let Some(ref key) = req.idempotency_key
176        && let Some(existing) = state
177            .idempotency_keys
178            .get(key)
179            .and_then(|id| state.runs.get(id))
180    {
181        if now - existing.created_at < IDEMPOTENCY_WINDOW {
182            return Ok(RunCreation::Existing(run_with_label(existing, state)));
183        }
184        // The key outlived its window: release it from the stale run.
185        let stale_id = existing.id;
186        state.idempotency_keys.remove(key);
187        if let Some(stale) = state.runs.get_mut(&stale_id) {
188            stale.idempotency_key = None;
189        }
190    }
191
192    if let Some(ref key) = req.concurrency_key
193        && let Some(holder) = state
194            .runs
195            .values()
196            .filter(|r| {
197                r.concurrency_key.as_deref() == Some(key.as_str()) && !r.status.state.is_terminal()
198            })
199            .min_by_key(|r| r.created_at)
200    {
201        return Err(StoreError::ConcurrencyConflict {
202            key: key.clone(),
203            run_id: holder.id,
204        });
205    }
206
207    let run = Run {
208        id: Uuid::now_v7(),
209        workflow_name: req.workflow_name,
210        status: crate::entities::FsmState::new(RunStatus::Pending, Uuid::now_v7()),
211        trigger: req.trigger,
212        payload: req.payload,
213        error: None,
214        retry_count: 0,
215        max_retries: req.max_retries,
216        cost_usd: Decimal::ZERO,
217        duration_ms: 0,
218        created_at: now,
219        updated_at: now,
220        started_at: None,
221        completed_at: None,
222        handler_version: req.handler_version,
223        labels: req.labels,
224        scheduled_at: req.scheduled_at,
225        created_by: req.created_by,
226        created_by_label: None,
227        idempotency_key: req.idempotency_key.clone(),
228        concurrency_key: req.concurrency_key,
229        priority: req.priority,
230        concurrency_limits: req.concurrency_limits,
231        max_cost_usd: req.max_cost_usd,
232        worker_id: None,
233        lease_expires_at: None,
234        output: None,
235        lease_recoveries: 0,
236        capacity_wait_kind: None,
237        resume_status: None,
238        worker_tags: normalize_worker_tags(req.worker_tags),
239    };
240
241    if let Some(key) = req.idempotency_key {
242        state.idempotency_keys.insert(key, run.id);
243    }
244    state.runs.insert(run.id, run.clone());
245    Ok(RunCreation::Created(run_with_label(&run, state)))
246}
247
248impl RunStore for InMemoryStore {
249    fn create_run(&self, req: NewRun) -> StoreFuture<'_, RunCreation> {
250        Box::pin(async move {
251            // Single critical section: the key lookup and the insert cannot be
252            // interleaved by a concurrent call sharing the same key.
253            let mut state = self.state.write().await;
254            insert_run(&mut state, req)
255        })
256    }
257
258    fn find_run_by_idempotency_key(&self, key: &str) -> StoreFuture<'_, Option<Run>> {
259        let key = key.to_string();
260        Box::pin(async move {
261            let now = Utc::now();
262            let state = self.state.read().await;
263            Ok(state
264                .idempotency_keys
265                .get(&key)
266                .and_then(|id| state.runs.get(id))
267                .filter(|run| now - run.created_at < IDEMPOTENCY_WINDOW)
268                .map(|run| run_with_label(run, &state)))
269        })
270    }
271
272    fn get_run(&self, id: Uuid) -> StoreFuture<'_, Option<Run>> {
273        Box::pin(async move {
274            let state = self.state.read().await;
275            Ok(state.runs.get(&id).map(|r| run_with_label(r, &state)))
276        })
277    }
278
279    fn list_runs(&self, filter: RunFilter, page: u32, per_page: u32) -> StoreFuture<'_, Page<Run>> {
280        Box::pin(async move {
281            let state = self.state.read().await;
282
283            let mut runs: Vec<&Run> = state
284                .runs
285                .values()
286                .filter(|r| run_matches_filter(r, &filter, &state.steps))
287                .collect();
288
289            // Sort newest first.
290            runs.sort_by_key(|r| Reverse(r.created_at));
291
292            let total = runs.len() as u64;
293            let page = page.max(1);
294            let per_page = per_page.clamp(1, 100);
295            let offset = ((page - 1) * per_page) as usize;
296            let items: Vec<Run> = runs
297                .into_iter()
298                .skip(offset)
299                .take(per_page as usize)
300                .map(|r| run_with_label(r, &state))
301                .collect();
302
303            Ok(Page {
304                items,
305                total,
306                page,
307                per_page,
308            })
309        })
310    }
311
312    fn update_run_status(&self, id: Uuid, new_status: RunStatus) -> StoreFuture<'_, ()> {
313        Box::pin(async move {
314            let mut state = self.state.write().await;
315            let run = state.runs.get_mut(&id).ok_or(StoreError::RunNotFound(id))?;
316
317            if !run.status.state.can_transition_to(&new_status) {
318                return Err(StoreError::InvalidTransition {
319                    from: run.status.state,
320                    to: new_status,
321                });
322            }
323
324            if run.status.state == new_status && new_status.is_terminal() {
325                return Ok(());
326            }
327
328            let now = Utc::now();
329            track_resume_status(run, new_status);
330            run.status.state = new_status;
331            run.updated_at = now;
332
333            if new_status == RunStatus::Running && run.started_at.is_none() {
334                run.started_at = Some(now);
335            }
336            if new_status.is_terminal() {
337                run.completed_at = Some(now);
338            }
339            if new_status != RunStatus::Running {
340                clear_lease(run);
341            }
342            // A paused run keeps its capacity wait kind for the resume.
343            if new_status != RunStatus::Sleeping && new_status != RunStatus::Paused {
344                run.capacity_wait_kind = None;
345            }
346
347            Ok(())
348        })
349    }
350
351    fn update_run(&self, id: Uuid, update: RunUpdate) -> StoreFuture<'_, ()> {
352        Box::pin(async move {
353            let mut state = self.state.write().await;
354            let run = state.runs.get_mut(&id).ok_or(StoreError::RunNotFound(id))?;
355
356            let now = Utc::now();
357
358            // Only a paused run carries a state to return to.
359            if let (None, Some(target)) = (update.status, update.resume_status) {
360                if run.status.state != RunStatus::Paused {
361                    return Err(StoreError::InvalidTransition {
362                        from: run.status.state,
363                        to: target,
364                    });
365                }
366                run.resume_status = Some(target);
367            }
368
369            if let Some(status) = update.status {
370                if !run.status.state.can_transition_to(&status) {
371                    return Err(StoreError::InvalidTransition {
372                        from: run.status.state,
373                        to: status,
374                    });
375                }
376                if !(run.status.state == status && status.is_terminal()) {
377                    track_resume_status(run, status);
378                    run.status.state = status;
379                    if status == RunStatus::Running && run.started_at.is_none() {
380                        run.started_at = Some(now);
381                    }
382                    if status.is_terminal() {
383                        run.completed_at = Some(now);
384                    }
385                    if status != RunStatus::Running {
386                        clear_lease(run);
387                    }
388                }
389                // The kind only means something while the run sleeps on capacity;
390                // a pause keeps it for the resume.
391                if status == RunStatus::Sleeping {
392                    run.capacity_wait_kind = update.capacity_wait_kind.clone();
393                } else if status != RunStatus::Paused {
394                    run.capacity_wait_kind = None;
395                }
396            }
397
398            // After the status block, so `Running` + `Set` ends with a lease.
399            match update.lease {
400                Some(LeaseUpdate::Set {
401                    worker_id,
402                    expires_at,
403                }) => {
404                    run.worker_id = Some(worker_id);
405                    run.lease_expires_at = Some(expires_at);
406                }
407                Some(LeaseUpdate::Release) => clear_lease(run),
408                None => {}
409            }
410
411            if let Some(error) = update.error {
412                run.error = Some(error);
413            }
414            if update.increment_retry {
415                run.retry_count += 1;
416            }
417            if let Some(cost) = update.cost_usd {
418                run.cost_usd = cost;
419            }
420            if let Some(dur) = update.duration_ms {
421                run.duration_ms = dur;
422            }
423            if let Some(started) = update.started_at {
424                run.started_at = Some(started);
425            }
426            if let Some(completed) = update.completed_at {
427                run.completed_at = Some(completed);
428            }
429            if let Some(scheduled) = update.scheduled_at {
430                run.scheduled_at = Some(scheduled);
431            }
432            if let Some(output) = update.output {
433                run.output = Some(output);
434            }
435
436            run.updated_at = now;
437            Ok(())
438        })
439    }
440
441    fn list_active_descendants(&self, run_id: Uuid) -> StoreFuture<'_, Vec<Run>> {
442        Box::pin(async move {
443            let state = self.state.read().await;
444            Ok(active_descendants(&state.runs, run_id))
445        })
446    }
447
448    fn pick_next_pending_for(
449        &self,
450        lease: Option<LeaseRequest>,
451        capabilities: Option<WorkerCapabilities>,
452    ) -> StoreFuture<'_, Option<Run>> {
453        Box::pin(async move {
454            let mut state = self.state.write().await;
455            let now = Utc::now();
456
457            // Find the highest priority run waiting for execution whose
458            // scheduled_at has passed (or is None), the oldest first among
459            // equal priorities. `Retrying` runs are runs whose automatic retry
460            // backoff has been armed: they become eligible again once
461            // `scheduled_at` has passed. Runs held back by a saturated
462            // concurrency group are skipped, so they never block younger runs.
463            // The group check and the transition below happen under the same
464            // write lock, so concurrent pickers cannot overshoot a limit.
465            // Runs the worker cannot take (unregistered workflow, missing tag)
466            // are skipped the same way.
467            let oldest_id = state
468                .runs
469                .values()
470                .filter(|r| is_due(r, now) && !is_blocked(r, &state.runs))
471                .filter(|r| !state.workflow_pauses.contains_key(&r.workflow_name))
472                .filter(|r| {
473                    capabilities
474                        .as_ref()
475                        .is_none_or(|c| c.can_take(&r.workflow_name, &r.worker_tags))
476                })
477                .min_by_key(|r| (Reverse(r.priority), r.created_at))
478                .map(|r| r.id);
479
480            let Some(id) = oldest_id else {
481                return Ok(None);
482            };
483
484            // Transition to Running and attach the lease atomically.
485            let run = state.runs.get_mut(&id).expect("run exists");
486            let now = Utc::now();
487            run.status.state = RunStatus::Running;
488            run.started_at = Some(now);
489            run.updated_at = now;
490            match lease {
491                Some(lease) => {
492                    run.lease_expires_at = Some(lease.expires_at(now));
493                    run.worker_id = Some(lease.worker_id);
494                }
495                None => clear_lease(run),
496            }
497            let run = run.clone();
498
499            Ok(Some(run_with_label(&run, &state)))
500        })
501    }
502
503    fn count_blocked_runs_by_group(&self) -> StoreFuture<'_, Vec<ConcurrencyGroupBacklog>> {
504        Box::pin(async move {
505            let state = self.state.read().await;
506            let now = Utc::now();
507
508            let mut blocked: BTreeMap<String, u64> = BTreeMap::new();
509            for run in state
510                .runs
511                .values()
512                .filter(|r| is_due(r, now) && !state.workflow_pauses.contains_key(&r.workflow_name))
513            {
514                for limit in &run.concurrency_limits {
515                    if running_count(&state.runs, &limit.group) >= u64::from(limit.limit) {
516                        *blocked.entry(limit.group.clone()).or_default() += 1;
517                    }
518                }
519            }
520
521            Ok(blocked
522                .into_iter()
523                .map(|(group, blocked_runs)| ConcurrencyGroupBacklog {
524                    group,
525                    blocked_runs,
526                })
527                .collect())
528        })
529    }
530
531    fn renew_lease(&self, id: Uuid, lease: LeaseRequest) -> StoreFuture<'_, DateTime<Utc>> {
532        Box::pin(async move {
533            let mut state = self.state.write().await;
534            let run = state.runs.get_mut(&id).ok_or(StoreError::RunNotFound(id))?;
535
536            if run.status.state != RunStatus::Running
537                || run.worker_id.as_deref() != Some(lease.worker_id.as_str())
538            {
539                return Err(StoreError::LeaseLost {
540                    run_id: id,
541                    held_by: run.worker_id.clone(),
542                });
543            }
544
545            let now = Utc::now();
546            let expires_at = lease.expires_at(now);
547            run.lease_expires_at = Some(expires_at);
548            run.updated_at = now;
549
550            Ok(expires_at)
551        })
552    }
553
554    fn reap_expired_leases(&self, limit: u32) -> StoreFuture<'_, Vec<ReapedRun>> {
555        Box::pin(async move {
556            let mut state = self.state.write().await;
557            let now = Utc::now();
558
559            let mut expired: Vec<Uuid> = state
560                .runs
561                .values()
562                .filter(|r| {
563                    r.status.state == RunStatus::Running
564                        && r.lease_expires_at.is_some_and(|at| at < now)
565                })
566                .map(|r| r.id)
567                .collect();
568            expired.sort_unstable();
569            expired.truncate(limit as usize);
570
571            let mut reaped = Vec::with_capacity(expired.len());
572            for id in expired {
573                let run = state.runs.get_mut(&id).expect("run exists");
574                run.lease_recoveries += 1;
575                clear_lease(run);
576                run.updated_at = now;
577
578                let to = if run.lease_recoveries > run.max_retries {
579                    run.status.state = RunStatus::Failed;
580                    run.error = Some(LEASE_EXPIRED_ERROR.to_string());
581                    run.completed_at = Some(now);
582                    RunStatus::Failed
583                } else {
584                    run.status.state = RunStatus::Pending;
585                    RunStatus::Pending
586                };
587
588                reaped.push(ReapedRun {
589                    run: run.clone(),
590                    from: RunStatus::Running,
591                    to,
592                });
593            }
594
595            Ok(reaped)
596        })
597    }
598
599    fn claim_due_approval_deadlines(&self, limit: u32) -> StoreFuture<'_, Vec<Step>> {
600        Box::pin(async move {
601            let mut state = self.state.write().await;
602            let now = Utc::now();
603
604            let mut due: Vec<(DateTime<Utc>, Uuid)> = state
605                .steps
606                .values()
607                .filter(|s| {
608                    s.status.state == StepStatus::AwaitingApproval
609                        && s.approval_deadline_at.is_some_and(|at| at <= now)
610                })
611                .map(|s| (s.approval_deadline_at.expect("deadline is set"), s.id))
612                .collect();
613            due.sort_unstable();
614            due.truncate(limit as usize);
615
616            let mut claimed = Vec::with_capacity(due.len());
617            for (_, id) in due {
618                let step = state.steps.get_mut(&id).expect("step exists");
619                // Clone before clearing so the caller still sees the deadline
620                // that fired.
621                claimed.push(step.clone());
622                step.approval_deadline_at = None;
623                step.updated_at = now;
624            }
625
626            Ok(claimed)
627        })
628    }
629
630    fn claim_due_sleeping_runs(&self, limit: u32) -> StoreFuture<'_, Vec<Run>> {
631        Box::pin(async move {
632            let mut state = self.state.write().await;
633            let now = Utc::now();
634
635            let mut due: Vec<(DateTime<Utc>, Uuid)> = state
636                .runs
637                .values()
638                .filter(|r| r.status.state == RunStatus::Sleeping)
639                .filter_map(|r| r.scheduled_at.filter(|at| *at <= now).map(|at| (at, r.id)))
640                .collect();
641            due.sort_unstable();
642            due.truncate(limit as usize);
643
644            let mut woken = Vec::with_capacity(due.len());
645            for (_, id) in due {
646                let run = state.runs.get_mut(&id).expect("run exists");
647                run.status.state = RunStatus::Pending;
648                run.scheduled_at = None;
649                run.capacity_wait_kind = None;
650                run.updated_at = now;
651                let run = run.clone();
652                woken.push(run_with_label(&run, &state));
653            }
654
655            Ok(woken)
656        })
657    }
658
659    fn pause_workflow(
660        &self,
661        workflow_name: &str,
662        paused_by: Option<Uuid>,
663    ) -> StoreFuture<'_, WorkflowPause> {
664        let workflow_name = workflow_name.to_string();
665        Box::pin(async move {
666            let mut state = self.state.write().await;
667            let pause = state
668                .workflow_pauses
669                .entry(workflow_name.clone())
670                .or_insert_with(|| WorkflowPause {
671                    workflow_name,
672                    paused_at: Utc::now(),
673                    paused_by,
674                });
675            Ok(pause.clone())
676        })
677    }
678
679    fn resume_workflow(&self, workflow_name: &str) -> StoreFuture<'_, bool> {
680        let workflow_name = workflow_name.to_string();
681        Box::pin(async move {
682            let mut state = self.state.write().await;
683            Ok(state.workflow_pauses.remove(&workflow_name).is_some())
684        })
685    }
686
687    fn list_workflow_pauses(&self) -> StoreFuture<'_, Vec<WorkflowPause>> {
688        Box::pin(async move {
689            let state = self.state.read().await;
690            let mut pauses: Vec<WorkflowPause> = state.workflow_pauses.values().cloned().collect();
691            pauses.sort_by(|a, b| a.workflow_name.cmp(&b.workflow_name));
692            Ok(pauses)
693        })
694    }
695
696    fn list_purgeable_runs(
697        &self,
698        policy: &PurgePolicy,
699        batch_size: u32,
700    ) -> StoreFuture<'_, Vec<PurgeableRun>> {
701        let max_age_days = policy.max_age_days;
702        let max_runs_per_workflow = policy.max_runs_per_workflow;
703        Box::pin(async move {
704            let state = self.state.read().await;
705            let cutoff = Utc::now() - Duration::days(i64::from(max_age_days));
706            let mut result: Vec<PurgeableRun> = Vec::new();
707            let mut seen: HashSet<Uuid> = HashSet::new();
708
709            for run in state.runs.values() {
710                if run.status.state.is_terminal() && run.created_at < cutoff {
711                    seen.insert(run.id);
712                    result.push(PurgeableRun {
713                        run_id: run.id,
714                        workflow_name: run.workflow_name.clone(),
715                        reason: PurgeReason::TooOld,
716                    });
717                }
718            }
719
720            let mut by_workflow: HashMap<&str, Vec<&Run>> = HashMap::new();
721            for run in state.runs.values() {
722                if run.status.state.is_terminal() {
723                    by_workflow.entry(&run.workflow_name).or_default().push(run);
724                }
725            }
726            for (_, mut runs) in by_workflow {
727                if runs.len() > max_runs_per_workflow as usize {
728                    runs.sort_by_key(|r| r.created_at);
729                    let excess = runs.len() - max_runs_per_workflow as usize;
730                    for run in runs.into_iter().take(excess) {
731                        if seen.insert(run.id) {
732                            result.push(PurgeableRun {
733                                run_id: run.id,
734                                workflow_name: run.workflow_name.clone(),
735                                reason: PurgeReason::ExceedsWorkflowLimit,
736                            });
737                        }
738                    }
739                }
740            }
741
742            result.sort_by_key(|p| p.run_id);
743            result.truncate(batch_size as usize);
744            Ok(result)
745        })
746    }
747
748    fn delete_run(&self, id: Uuid) -> StoreFuture<'_, Vec<String>> {
749        Box::pin(async move {
750            let mut state = self.state.write().await;
751
752            if !state.runs.contains_key(&id) {
753                return Err(StoreError::RunNotFound(id));
754            }
755
756            // Collect artifact storage keys before removing them.
757            let storage_keys: Vec<String> = state
758                .artifacts
759                .values()
760                .filter(|a| a.run_id == id)
761                .map(|a| a.storage_key.clone())
762                .collect();
763
764            // Remove artifacts.
765            state.artifacts.retain(|_, a| a.run_id != id);
766
767            // Remove step dependencies (both sides).
768            let step_ids: Vec<Uuid> = state
769                .steps
770                .values()
771                .filter(|s| s.run_id == id)
772                .map(|s| s.id)
773                .collect();
774            state
775                .step_dependencies
776                .retain(|d| !step_ids.contains(&d.step_id) && !step_ids.contains(&d.depends_on));
777
778            // Remove steps.
779            state.steps.retain(|_, s| s.run_id != id);
780
781            // Remove idempotency key pointing to this run.
782            state.idempotency_keys.retain(|_, &mut run_id| run_id != id);
783
784            // Remove the run itself.
785            state.runs.remove(&id);
786
787            Ok(storage_keys)
788        })
789    }
790
791    fn create_step(&self, req: NewStep) -> StoreFuture<'_, Step> {
792        Box::pin(async move {
793            let mut state = self.state.write().await;
794
795            let attempt = state
796                .runs
797                .get(&req.run_id)
798                .ok_or(StoreError::RunNotFound(req.run_id))?
799                .retry_count
800                + 1;
801
802            let agent_role = agent_role_from_input(&req.kind, req.input.as_ref());
803            let now = Utc::now();
804            let step = Step {
805                id: Uuid::now_v7(),
806                trace_id: req.trace_id,
807                run_id: req.run_id,
808                name: req.name,
809                kind: req.kind,
810                position: req.position,
811                status: crate::entities::FsmState::new(StepStatus::Pending, Uuid::now_v7()),
812                attempt,
813                input: req.input,
814                output: None,
815                error: None,
816                duration_ms: 0,
817                cost_usd: Decimal::ZERO,
818                input_tokens: None,
819                cache_read_input_tokens: None,
820                cache_creation_input_tokens: None,
821                output_tokens: None,
822                created_at: now,
823                updated_at: now,
824                started_at: None,
825                completed_at: None,
826                debug_messages: None,
827                is_error_handler: req.is_error_handler,
828                approval_deadline_at: None,
829                approval_stage: 0,
830                approval_assignee: None,
831                approval_assignee_user_id: None,
832                approval_requirement: None,
833                approvals: Vec::new(),
834                account_id: None,
835                environment_id: None,
836                session_id: None,
837                agent_role,
838            };
839
840            state.steps.insert(step.id, step.clone());
841            Ok(step)
842        })
843    }
844
845    fn update_step(&self, id: Uuid, update: StepUpdate) -> StoreFuture<'_, ()> {
846        Box::pin(async move {
847            let mut state = self.state.write().await;
848            // Resolved when the gate opens, like the PostgreSQL store: an account created
849            // later under the same username is never bound to it.
850            let assignee_user_id =
851                update
852                    .approval_assignee
853                    .as_ref()
854                    .map(|assignee| match assignee {
855                        Assignee::User(name) => state
856                            .users
857                            .values()
858                            .find(|u| u.username == *name)
859                            .map(|u| u.id),
860                        Assignee::Group(_) => None,
861                    });
862            let step = state
863                .steps
864                .get_mut(&id)
865                .ok_or(StoreError::StepNotFound(id))?;
866
867            let now = Utc::now();
868
869            if let Some(status) = update.status {
870                if !matches!(
871                    (step.status.state, status),
872                    (StepStatus::Pending, StepStatus::Running)
873                        | (StepStatus::Pending, StepStatus::Skipped)
874                        | (StepStatus::Running, StepStatus::Completed)
875                        | (StepStatus::Running, StepStatus::Failed)
876                        | (StepStatus::Running, StepStatus::AwaitingApproval)
877                        | (StepStatus::AwaitingApproval, StepStatus::Running)
878                        | (StepStatus::AwaitingApproval, StepStatus::Completed)
879                        | (StepStatus::AwaitingApproval, StepStatus::Failed)
880                        | (StepStatus::AwaitingApproval, StepStatus::Rejected)
881                ) {
882                    return Err(StoreError::Database(format!(
883                        "invalid step status transition: {:?} -> {:?}",
884                        step.status.state, status
885                    )));
886                }
887                step.status.state = status;
888            }
889            if let Some(output) = update.output {
890                step.output = Some(output);
891            }
892            if let Some(error) = update.error {
893                step.error = Some(error);
894            }
895            if let Some(dur) = update.duration_ms {
896                step.duration_ms = dur;
897            }
898            if let Some(cost) = update.cost_usd {
899                step.cost_usd = cost;
900            }
901            if let Some(tokens) = update.input_tokens {
902                step.input_tokens = Some(tokens);
903            }
904            if let Some(tokens) = update.cache_read_input_tokens {
905                step.cache_read_input_tokens = Some(tokens);
906            }
907            if let Some(tokens) = update.cache_creation_input_tokens {
908                step.cache_creation_input_tokens = Some(tokens);
909            }
910            if let Some(account_id) = update.account_id {
911                step.account_id = Some(account_id);
912            }
913            if let Some(environment_id) = update.environment_id {
914                step.environment_id = Some(environment_id);
915            }
916            if let Some(session_id) = update.session_id {
917                step.session_id = Some(session_id);
918            }
919            if let Some(tokens) = update.output_tokens {
920                step.output_tokens = Some(tokens);
921            }
922            if let Some(started) = update.started_at {
923                step.started_at = Some(started);
924            }
925            if let Some(completed) = update.completed_at {
926                step.completed_at = Some(completed);
927            }
928            if let Some(debug_msgs) = update.debug_messages {
929                step.debug_messages = Some(debug_msgs);
930            }
931            // Clearing wins over setting: an update that resolves the gate and
932            // reschedules at once would otherwise leave a live timer behind.
933            if update.clear_approval_deadline {
934                step.approval_deadline_at = None;
935            } else if let Some(deadline) = update.approval_deadline_at {
936                step.approval_deadline_at = Some(deadline);
937            }
938            if let Some(stage) = update.approval_stage {
939                step.approval_stage = stage;
940            }
941            if let Some(assignee) = update.approval_assignee {
942                step.approval_assignee = Some(assignee);
943                step.approval_assignee_user_id = assignee_user_id.flatten();
944            }
945            if let Some(requirement) = update.approval_requirement {
946                step.approval_requirement = Some(requirement);
947            }
948
949            step.updated_at = now;
950            Ok(())
951        })
952    }
953
954    fn get_step(&self, id: Uuid) -> StoreFuture<'_, Option<Step>> {
955        Box::pin(async move {
956            let state = self.state.read().await;
957            Ok(state.steps.get(&id).cloned())
958        })
959    }
960
961    fn record_step_approval(&self, step_id: Uuid, approval: StepApproval) -> StoreFuture<'_, Step> {
962        Box::pin(async move {
963            let mut state = self.state.write().await;
964            let step = state
965                .steps
966                .get_mut(&step_id)
967                .ok_or(StoreError::StepNotFound(step_id))?;
968
969            if !step.approvals.iter().any(|a| a.user_id == approval.user_id) {
970                step.approvals.push(approval);
971                step.updated_at = Utc::now();
972            }
973            Ok(step.clone())
974        })
975    }
976
977    fn list_steps(&self, run_id: Uuid) -> StoreFuture<'_, Vec<Step>> {
978        Box::pin(async move {
979            let state = self.state.read().await;
980            let mut steps: Vec<Step> = state
981                .steps
982                .values()
983                .filter(|s| s.run_id == run_id)
984                .cloned()
985                .collect();
986            steps.sort_by_key(|s| s.position);
987            Ok(steps)
988        })
989    }
990
991    fn get_stats(&self, filter: RunFilter) -> StoreFuture<'_, RunStats> {
992        Box::pin(async move {
993            let state = self.state.read().await;
994
995            let mut total_cost_usd = Decimal::ZERO;
996            let mut total_duration_ms = 0u64;
997            let mut total_runs = 0u64;
998            let mut completed_runs = 0u64;
999            let mut failed_runs = 0u64;
1000            let mut cancelled_runs = 0u64;
1001            let mut active_runs = 0u64;
1002            let mut awaiting_approval_runs = 0u64;
1003
1004            for run in state.runs.values() {
1005                if !run_matches_filter(run, &filter, &state.steps) {
1006                    continue;
1007                }
1008
1009                total_cost_usd += run.cost_usd;
1010                total_duration_ms += run.duration_ms;
1011                total_runs += 1;
1012
1013                match run.status.state {
1014                    RunStatus::Completed | RunStatus::Warning => completed_runs += 1,
1015                    RunStatus::Failed => failed_runs += 1,
1016                    RunStatus::Cancelled => cancelled_runs += 1,
1017                    RunStatus::AwaitingApproval => {
1018                        active_runs += 1;
1019                        awaiting_approval_runs += 1;
1020                    }
1021                    RunStatus::Pending
1022                    | RunStatus::Running
1023                    | RunStatus::Retrying
1024                    | RunStatus::Sleeping
1025                    | RunStatus::Paused => {
1026                        active_runs += 1;
1027                    }
1028                }
1029            }
1030
1031            Ok(RunStats {
1032                total_runs,
1033                completed_runs,
1034                failed_runs,
1035                cancelled_runs,
1036                active_runs,
1037                awaiting_approval_runs,
1038                total_cost_usd,
1039                total_duration_ms,
1040            })
1041        })
1042    }
1043
1044    fn get_stats_history(
1045        &self,
1046        filter: StatsHistoryFilter,
1047    ) -> StoreFuture<'_, Vec<StatsHistoryBucket>> {
1048        Box::pin(async move {
1049            let state = self.state.read().await;
1050            let now = Utc::now();
1051            let start = now - Duration::hours(filter.period.hours());
1052            let run_filter = filter.to_run_filter();
1053            let buckets = aggregate_history_buckets(
1054                state
1055                    .runs
1056                    .values()
1057                    .filter(|r| run_matches_filter(r, &run_filter, &state.steps)),
1058                start,
1059                now,
1060                filter.granularity,
1061            );
1062            Ok(buckets)
1063        })
1064    }
1065
1066    fn list_steps_by_agent_role<'a>(
1067        &'a self,
1068        agent_role: &'a str,
1069        page: u32,
1070        per_page: u32,
1071    ) -> StoreFuture<'a, Page<Step>> {
1072        Box::pin(async move {
1073            let state = self.state.read().await;
1074
1075            let mut steps: Vec<&Step> = state
1076                .steps
1077                .values()
1078                .filter(|s| s.agent_role.as_deref() == Some(agent_role))
1079                .collect();
1080            steps.sort_by_key(|s| Reverse((s.created_at, s.id)));
1081
1082            let total = steps.len() as u64;
1083            let page = page.max(1);
1084            let per_page = per_page.clamp(1, 100);
1085            // Computed in u64: a large page times per_page overflows u32.
1086            let offset =
1087                usize::try_from(u64::from(page - 1) * u64::from(per_page)).unwrap_or(usize::MAX);
1088            let items: Vec<Step> = steps
1089                .into_iter()
1090                .skip(offset)
1091                .take(per_page as usize)
1092                .cloned()
1093                .collect();
1094
1095            Ok(Page {
1096                items,
1097                total,
1098                page,
1099                per_page,
1100            })
1101        })
1102    }
1103
1104    fn get_agent_role_stats(
1105        &self,
1106        created_after: DateTime<Utc>,
1107    ) -> StoreFuture<'_, Vec<AgentRoleStats>> {
1108        Box::pin(async move {
1109            let state = self.state.read().await;
1110
1111            let mut by_role: BTreeMap<String, AgentRoleStats> = BTreeMap::new();
1112            for step in state.steps.values() {
1113                if step.created_at < created_after {
1114                    continue;
1115                }
1116                let Some(role) = step.agent_role.as_deref() else {
1117                    continue;
1118                };
1119                let stats = by_role
1120                    .entry(role.to_string())
1121                    .or_insert_with(|| AgentRoleStats {
1122                        agent_role: role.to_string(),
1123                        ..AgentRoleStats::default()
1124                    });
1125                stats.executions += 1;
1126                if step.status.state == StepStatus::Failed {
1127                    stats.failed += 1;
1128                }
1129                stats.total_cost_usd += step.cost_usd;
1130                stats.total_duration_ms += step.duration_ms;
1131            }
1132
1133            let mut roles: Vec<AgentRoleStats> = by_role.into_values().collect();
1134            roles.sort_by(|a, b| {
1135                b.total_cost_usd
1136                    .cmp(&a.total_cost_usd)
1137                    .then_with(|| a.agent_role.cmp(&b.agent_role))
1138            });
1139            Ok(roles)
1140        })
1141    }
1142
1143    fn create_step_dependencies(&self, deps: Vec<NewStepDependency>) -> StoreFuture<'_, ()> {
1144        Box::pin(async move {
1145            let mut state = self.state.write().await;
1146
1147            for dep in deps {
1148                if !state.steps.contains_key(&dep.step_id) {
1149                    return Err(StoreError::StepNotFound(dep.step_id));
1150                }
1151                if !state.steps.contains_key(&dep.depends_on) {
1152                    return Err(StoreError::StepNotFound(dep.depends_on));
1153                }
1154
1155                let already_exists = state
1156                    .step_dependencies
1157                    .iter()
1158                    .any(|d| d.step_id == dep.step_id && d.depends_on == dep.depends_on);
1159
1160                if !already_exists {
1161                    state.step_dependencies.push(StepDependency {
1162                        step_id: dep.step_id,
1163                        depends_on: dep.depends_on,
1164                        created_at: Utc::now(),
1165                    });
1166                }
1167            }
1168
1169            Ok(())
1170        })
1171    }
1172
1173    fn list_step_dependencies(&self, run_id: Uuid) -> StoreFuture<'_, Vec<StepDependency>> {
1174        Box::pin(async move {
1175            let state = self.state.read().await;
1176
1177            let run_step_ids: std::collections::HashSet<Uuid> = state
1178                .steps
1179                .values()
1180                .filter(|s| s.run_id == run_id)
1181                .map(|s| s.id)
1182                .collect();
1183
1184            let mut deps: Vec<StepDependency> = state
1185                .step_dependencies
1186                .iter()
1187                .filter(|d| run_step_ids.contains(&d.step_id))
1188                .cloned()
1189                .collect();
1190
1191            deps.sort_by_key(|d| d.created_at);
1192            Ok(deps)
1193        })
1194    }
1195}
1196
1197#[cfg(test)]
1198mod tests {
1199    use std::collections::HashMap;
1200    use std::time::Duration;
1201
1202    use chrono::TimeDelta;
1203    use serde_json::json;
1204    use tokio::spawn;
1205    use tokio::time::sleep;
1206
1207    use super::*;
1208    use crate::api_key_store::ApiKeyStore;
1209    use crate::entities::{
1210        ApiKeyScope, ApiKeyUpdate, ApprovalRequirement, NewApiKey, NewUser, TriggerKind,
1211    };
1212    use crate::user_store::UserStore;
1213
1214    use crate::memory::tests::{create_terminal_run, new_run_req};
1215    use crate::store::RunStore;
1216
1217    use crate::entities::{StepKind, step_trace_id};
1218
1219    fn new_step_req(run_id: Uuid, name: &str, position: u32) -> NewStep {
1220        NewStep {
1221            run_id,
1222            trace_id: step_trace_id(run_id, name, position),
1223            name: name.to_string(),
1224            kind: StepKind::Shell,
1225            position,
1226            input: None,
1227            is_error_handler: false,
1228        }
1229    }
1230
1231    // ---- create_run ----
1232
1233    #[tokio::test]
1234    async fn create_run_returns_pending_status() {
1235        let store = InMemoryStore::new();
1236        let run = store
1237            .create_run(new_run_req("test"))
1238            .await
1239            .unwrap()
1240            .into_run();
1241        assert_eq!(run.status.state, RunStatus::Pending);
1242        assert_eq!(run.workflow_name, "test");
1243        assert_eq!(run.retry_count, 0);
1244        assert_eq!(run.max_retries, 3);
1245    }
1246
1247    #[tokio::test]
1248    async fn create_run_generates_unique_ids() {
1249        let store = InMemoryStore::new();
1250        let r1 = store.create_run(new_run_req("a")).await.unwrap().into_run();
1251        let r2 = store.create_run(new_run_req("b")).await.unwrap().into_run();
1252        assert_ne!(r1.id, r2.id);
1253    }
1254
1255    // ---- get_run ----
1256
1257    #[tokio::test]
1258    async fn get_run_returns_created_run() {
1259        let store = InMemoryStore::new();
1260        let run = store
1261            .create_run(new_run_req("test"))
1262            .await
1263            .unwrap()
1264            .into_run();
1265        let fetched = store.get_run(run.id).await.unwrap();
1266        assert!(fetched.is_some());
1267        assert_eq!(fetched.unwrap().id, run.id);
1268    }
1269
1270    #[tokio::test]
1271    async fn get_run_returns_none_for_missing() {
1272        let store = InMemoryStore::new();
1273        let fetched = store.get_run(Uuid::nil()).await.unwrap();
1274        assert!(fetched.is_none());
1275    }
1276
1277    // ---- update_run_status ----
1278
1279    #[tokio::test]
1280    async fn update_run_status_valid_transition() {
1281        let store = InMemoryStore::new();
1282        let run = store
1283            .create_run(new_run_req("test"))
1284            .await
1285            .unwrap()
1286            .into_run();
1287
1288        store
1289            .update_run_status(run.id, RunStatus::Running)
1290            .await
1291            .unwrap();
1292
1293        let fetched = store.get_run(run.id).await.unwrap().unwrap();
1294        assert_eq!(fetched.status.state, RunStatus::Running);
1295        assert!(fetched.started_at.is_some());
1296    }
1297
1298    #[tokio::test]
1299    async fn update_run_status_invalid_transition_returns_error() {
1300        let store = InMemoryStore::new();
1301        let run = store
1302            .create_run(new_run_req("test"))
1303            .await
1304            .unwrap()
1305            .into_run();
1306
1307        let result = store.update_run_status(run.id, RunStatus::Completed).await;
1308        assert!(result.is_err());
1309
1310        let err = result.unwrap_err();
1311        assert!(matches!(err, StoreError::InvalidTransition { .. }));
1312    }
1313
1314    #[tokio::test]
1315    async fn update_run_status_not_found() {
1316        let store = InMemoryStore::new();
1317        let result = store
1318            .update_run_status(Uuid::nil(), RunStatus::Running)
1319            .await;
1320        assert!(matches!(result.unwrap_err(), StoreError::RunNotFound(_)));
1321    }
1322
1323    #[tokio::test]
1324    async fn update_run_status_terminal_sets_completed_at() {
1325        let store = InMemoryStore::new();
1326        let run = store
1327            .create_run(new_run_req("test"))
1328            .await
1329            .unwrap()
1330            .into_run();
1331
1332        store
1333            .update_run_status(run.id, RunStatus::Running)
1334            .await
1335            .unwrap();
1336        store
1337            .update_run_status(run.id, RunStatus::Completed)
1338            .await
1339            .unwrap();
1340
1341        let fetched = store.get_run(run.id).await.unwrap().unwrap();
1342        assert_eq!(fetched.status.state, RunStatus::Completed);
1343        assert!(fetched.completed_at.is_some());
1344    }
1345
1346    #[tokio::test]
1347    async fn update_run_status_terminal_to_same_is_idempotent() {
1348        let store = InMemoryStore::new();
1349        let run = store
1350            .create_run(new_run_req("test"))
1351            .await
1352            .unwrap()
1353            .into_run();
1354
1355        store
1356            .update_run_status(run.id, RunStatus::Running)
1357            .await
1358            .unwrap();
1359        store
1360            .update_run_status(run.id, RunStatus::Failed)
1361            .await
1362            .unwrap();
1363
1364        let before = store.get_run(run.id).await.unwrap().unwrap();
1365        let completed_at_before = before.completed_at;
1366
1367        store
1368            .update_run_status(run.id, RunStatus::Failed)
1369            .await
1370            .unwrap();
1371
1372        let after = store.get_run(run.id).await.unwrap().unwrap();
1373        assert_eq!(after.status.state, RunStatus::Failed);
1374        assert_eq!(after.completed_at, completed_at_before);
1375    }
1376
1377    #[tokio::test]
1378    async fn update_run_terminal_to_same_via_update_run_is_idempotent() {
1379        let store = InMemoryStore::new();
1380        let run = store
1381            .create_run(new_run_req("test"))
1382            .await
1383            .unwrap()
1384            .into_run();
1385
1386        store
1387            .update_run_status(run.id, RunStatus::Running)
1388            .await
1389            .unwrap();
1390        store
1391            .update_run(
1392                run.id,
1393                RunUpdate {
1394                    status: Some(RunStatus::Failed),
1395                    error: Some("first failure".to_string()),
1396                    ..RunUpdate::default()
1397                },
1398            )
1399            .await
1400            .unwrap();
1401
1402        let before = store.get_run(run.id).await.unwrap().unwrap();
1403
1404        store
1405            .update_run(
1406                run.id,
1407                RunUpdate {
1408                    status: Some(RunStatus::Failed),
1409                    ..RunUpdate::default()
1410                },
1411            )
1412            .await
1413            .unwrap();
1414
1415        let after = store.get_run(run.id).await.unwrap().unwrap();
1416        assert_eq!(after.status.state, RunStatus::Failed);
1417        assert_eq!(after.completed_at, before.completed_at);
1418        assert_eq!(after.error, Some("first failure".to_string()));
1419    }
1420
1421    // ---- list_runs ----
1422
1423    #[tokio::test]
1424    async fn list_runs_empty_store() {
1425        let store = InMemoryStore::new();
1426        let page = store.list_runs(RunFilter::default(), 1, 20).await.unwrap();
1427        assert_eq!(page.total, 0);
1428        assert!(page.items.is_empty());
1429    }
1430
1431    #[tokio::test]
1432    async fn list_runs_with_workflow_filter() {
1433        let store = InMemoryStore::new();
1434        store
1435            .create_run(new_run_req("deploy"))
1436            .await
1437            .unwrap()
1438            .into_run();
1439        store
1440            .create_run(new_run_req("test"))
1441            .await
1442            .unwrap()
1443            .into_run();
1444        store
1445            .create_run(new_run_req("deploy"))
1446            .await
1447            .unwrap()
1448            .into_run();
1449
1450        let filter = RunFilter {
1451            workflow_name: Some("deploy".to_string()),
1452            ..RunFilter::default()
1453        };
1454        let page = store.list_runs(filter, 1, 20).await.unwrap();
1455        assert_eq!(page.total, 2);
1456        assert!(page.items.iter().all(|r| r.workflow_name == "deploy"));
1457    }
1458
1459    #[tokio::test]
1460    async fn list_runs_with_status_filter() {
1461        let store = InMemoryStore::new();
1462        let run = store.create_run(new_run_req("a")).await.unwrap().into_run();
1463        store.create_run(new_run_req("b")).await.unwrap().into_run();
1464
1465        store
1466            .update_run_status(run.id, RunStatus::Running)
1467            .await
1468            .unwrap();
1469
1470        let filter = RunFilter {
1471            status: Some(RunStatus::Running),
1472            ..RunFilter::default()
1473        };
1474        let page = store.list_runs(filter, 1, 20).await.unwrap();
1475        assert_eq!(page.total, 1);
1476        assert_eq!(page.items[0].id, run.id);
1477    }
1478
1479    #[tokio::test]
1480    async fn list_runs_pagination() {
1481        let store = InMemoryStore::new();
1482        for i in 0..5 {
1483            store
1484                .create_run(new_run_req(&format!("wf-{i}")))
1485                .await
1486                .unwrap()
1487                .into_run();
1488        }
1489
1490        let page1 = store.list_runs(RunFilter::default(), 1, 2).await.unwrap();
1491        assert_eq!(page1.total, 5);
1492        assert_eq!(page1.items.len(), 2);
1493        assert_eq!(page1.page, 1);
1494        assert_eq!(page1.per_page, 2);
1495
1496        let page2 = store.list_runs(RunFilter::default(), 2, 2).await.unwrap();
1497        assert_eq!(page2.items.len(), 2);
1498
1499        let page3 = store.list_runs(RunFilter::default(), 3, 2).await.unwrap();
1500        assert_eq!(page3.items.len(), 1);
1501    }
1502
1503    // ---- worker lease ----
1504
1505    fn lease(worker_id: &str, ttl_secs: u64) -> Option<LeaseRequest> {
1506        Some(LeaseRequest {
1507            worker_id: worker_id.to_string(),
1508            ttl: Duration::from_secs(ttl_secs),
1509        })
1510    }
1511
1512    /// Pick a run holding a lease that is already in the past.
1513    ///
1514    /// The short sleep matters: `Utc::now()` has microsecond resolution, so a
1515    /// sub-microsecond TTL can still read as "not yet expired" in the same tick.
1516    async fn pick_with_expired_lease(store: &InMemoryStore, max_retries: u32) -> Run {
1517        let mut req = new_run_req("test");
1518        req.max_retries = max_retries;
1519        store.create_run(req).await.unwrap();
1520        let picked = expire_now(store).await;
1521        picked.expect("a pending run was just created")
1522    }
1523
1524    /// Pick the next pending run with a lease that is expired by the time this
1525    /// returns.
1526    async fn expire_now(store: &InMemoryStore) -> Option<Run> {
1527        let picked = store
1528            .pick_next_pending(Some(LeaseRequest {
1529                worker_id: "worker-1".to_string(),
1530                ttl: Duration::from_nanos(1),
1531            }))
1532            .await
1533            .unwrap();
1534        sleep(Duration::from_millis(2)).await;
1535        picked
1536    }
1537
1538    #[tokio::test]
1539    async fn pick_next_pending_attaches_lease() {
1540        let store = InMemoryStore::new();
1541        store.create_run(new_run_req("test")).await.unwrap();
1542
1543        let picked = store
1544            .pick_next_pending(lease("worker-1", 90))
1545            .await
1546            .unwrap()
1547            .unwrap();
1548
1549        assert_eq!(picked.worker_id.as_deref(), Some("worker-1"));
1550        let expires = picked.lease_expires_at.expect("lease set");
1551        assert!(expires > Utc::now());
1552        assert!(expires <= Utc::now() + TimeDelta::seconds(91));
1553    }
1554
1555    #[tokio::test]
1556    async fn pick_next_pending_without_lease_leaves_run_unowned() {
1557        let store = InMemoryStore::new();
1558        store.create_run(new_run_req("test")).await.unwrap();
1559
1560        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
1561
1562        assert!(picked.worker_id.is_none());
1563        assert!(picked.lease_expires_at.is_none());
1564    }
1565
1566    #[tokio::test]
1567    async fn renew_lease_extends_expiry_for_owner() {
1568        let store = InMemoryStore::new();
1569        store.create_run(new_run_req("test")).await.unwrap();
1570        let picked = store
1571            .pick_next_pending(lease("worker-1", 1))
1572            .await
1573            .unwrap()
1574            .unwrap();
1575
1576        let renewed = store
1577            .renew_lease(picked.id, lease("worker-1", 90).unwrap())
1578            .await
1579            .unwrap();
1580
1581        assert!(renewed > picked.lease_expires_at.unwrap());
1582        let after = store.get_run(picked.id).await.unwrap().unwrap();
1583        assert_eq!(after.lease_expires_at, Some(renewed));
1584    }
1585
1586    #[tokio::test]
1587    async fn renew_lease_rejects_other_worker() {
1588        let store = InMemoryStore::new();
1589        store.create_run(new_run_req("test")).await.unwrap();
1590        let picked = store
1591            .pick_next_pending(lease("worker-1", 90))
1592            .await
1593            .unwrap()
1594            .unwrap();
1595
1596        let err = store
1597            .renew_lease(picked.id, lease("worker-2", 90).unwrap())
1598            .await
1599            .unwrap_err();
1600
1601        assert!(matches!(
1602            err,
1603            StoreError::LeaseLost { held_by: Some(ref w), .. } if w == "worker-1"
1604        ));
1605    }
1606
1607    #[tokio::test]
1608    async fn renew_lease_rejects_run_that_left_running() {
1609        let store = InMemoryStore::new();
1610        store.create_run(new_run_req("test")).await.unwrap();
1611        let picked = store
1612            .pick_next_pending(lease("worker-1", 90))
1613            .await
1614            .unwrap()
1615            .unwrap();
1616        store
1617            .update_run_status(picked.id, RunStatus::Cancelled)
1618            .await
1619            .unwrap();
1620
1621        let err = store
1622            .renew_lease(picked.id, lease("worker-1", 90).unwrap())
1623            .await
1624            .unwrap_err();
1625
1626        assert!(matches!(err, StoreError::LeaseLost { .. }));
1627    }
1628
1629    #[tokio::test]
1630    async fn renew_lease_on_unknown_run_is_not_found() {
1631        let store = InMemoryStore::new();
1632
1633        let err = store
1634            .renew_lease(Uuid::now_v7(), lease("worker-1", 90).unwrap())
1635            .await
1636            .unwrap_err();
1637
1638        assert!(matches!(err, StoreError::RunNotFound(_)));
1639    }
1640
1641    #[tokio::test]
1642    async fn leaving_running_clears_the_lease() {
1643        for target in [
1644            RunStatus::Completed,
1645            RunStatus::Retrying,
1646            RunStatus::AwaitingApproval,
1647        ] {
1648            let store = InMemoryStore::new();
1649            store.create_run(new_run_req("test")).await.unwrap();
1650            let picked = store
1651                .pick_next_pending(lease("worker-1", 90))
1652                .await
1653                .unwrap()
1654                .unwrap();
1655
1656            store.update_run_status(picked.id, target).await.unwrap();
1657
1658            let after = store.get_run(picked.id).await.unwrap().unwrap();
1659            assert!(after.worker_id.is_none(), "worker_id kept for {target}");
1660            assert!(
1661                after.lease_expires_at.is_none(),
1662                "lease_expires_at kept for {target}"
1663            );
1664        }
1665    }
1666
1667    #[tokio::test]
1668    async fn update_run_to_terminal_clears_the_lease() {
1669        let store = InMemoryStore::new();
1670        store.create_run(new_run_req("test")).await.unwrap();
1671        let picked = store
1672            .pick_next_pending(lease("worker-1", 90))
1673            .await
1674            .unwrap()
1675            .unwrap();
1676
1677        store
1678            .update_run(
1679                picked.id,
1680                RunUpdate {
1681                    status: Some(RunStatus::Failed),
1682                    ..RunUpdate::default()
1683                },
1684            )
1685            .await
1686            .unwrap();
1687
1688        let after = store.get_run(picked.id).await.unwrap().unwrap();
1689        assert!(after.worker_id.is_none());
1690        assert!(after.lease_expires_at.is_none());
1691    }
1692
1693    #[tokio::test]
1694    async fn update_run_lease_set_with_running_status_attaches_lease() {
1695        let store = InMemoryStore::new();
1696        let run = store
1697            .create_run(new_run_req("test"))
1698            .await
1699            .unwrap()
1700            .into_run();
1701        let expires_at = Utc::now() + TimeDelta::seconds(60);
1702
1703        store
1704            .update_run(
1705                run.id,
1706                RunUpdate {
1707                    status: Some(RunStatus::Running),
1708                    lease: Some(LeaseUpdate::Set {
1709                        worker_id: "worker-1".to_string(),
1710                        expires_at,
1711                    }),
1712                    ..RunUpdate::default()
1713                },
1714            )
1715            .await
1716            .unwrap();
1717
1718        let after = store.get_run(run.id).await.unwrap().unwrap();
1719        assert_eq!(after.status.state, RunStatus::Running);
1720        assert_eq!(after.worker_id.as_deref(), Some("worker-1"));
1721        assert_eq!(after.lease_expires_at, Some(expires_at));
1722        // The transferred lease is renewable by its holder like a picked one.
1723        store
1724            .renew_lease(run.id, lease("worker-1", 90).unwrap())
1725            .await
1726            .unwrap();
1727    }
1728
1729    #[tokio::test]
1730    async fn update_run_lease_release_keeps_run_running_without_lease() {
1731        let store = InMemoryStore::new();
1732        store.create_run(new_run_req("test")).await.unwrap();
1733        let picked = store
1734            .pick_next_pending(lease("worker-1", 90))
1735            .await
1736            .unwrap()
1737            .unwrap();
1738
1739        store
1740            .update_run(
1741                picked.id,
1742                RunUpdate {
1743                    lease: Some(LeaseUpdate::Release),
1744                    ..RunUpdate::default()
1745                },
1746            )
1747            .await
1748            .unwrap();
1749
1750        let after = store.get_run(picked.id).await.unwrap().unwrap();
1751        assert_eq!(after.status.state, RunStatus::Running);
1752        assert!(after.worker_id.is_none());
1753        assert!(after.lease_expires_at.is_none());
1754        let err = store
1755            .renew_lease(picked.id, lease("worker-1", 90).unwrap())
1756            .await
1757            .unwrap_err();
1758        assert!(matches!(err, StoreError::LeaseLost { .. }));
1759    }
1760
1761    #[tokio::test]
1762    async fn update_run_lease_none_leaves_lease_untouched() {
1763        let store = InMemoryStore::new();
1764        store.create_run(new_run_req("test")).await.unwrap();
1765        let picked = store
1766            .pick_next_pending(lease("worker-1", 90))
1767            .await
1768            .unwrap()
1769            .unwrap();
1770
1771        store
1772            .update_run(
1773                picked.id,
1774                RunUpdate {
1775                    cost_usd: Some(Decimal::new(150, 2)),
1776                    ..RunUpdate::default()
1777                },
1778            )
1779            .await
1780            .unwrap();
1781
1782        let after = store.get_run(picked.id).await.unwrap().unwrap();
1783        assert_eq!(after.worker_id.as_deref(), Some("worker-1"));
1784        assert_eq!(after.lease_expires_at, picked.lease_expires_at);
1785    }
1786
1787    // ---- reap_expired_leases ----
1788
1789    #[tokio::test]
1790    async fn reap_expired_leases_empty_store() {
1791        let store = InMemoryStore::new();
1792        assert!(store.reap_expired_leases(100).await.unwrap().is_empty());
1793    }
1794
1795    #[tokio::test]
1796    async fn reap_expired_leases_requeues_run() {
1797        let store = InMemoryStore::new();
1798        let picked = pick_with_expired_lease(&store, 3).await;
1799
1800        let reaped = store.reap_expired_leases(100).await.unwrap();
1801
1802        assert_eq!(reaped.len(), 1);
1803        assert_eq!(reaped[0].from, RunStatus::Running);
1804        assert_eq!(reaped[0].to, RunStatus::Pending);
1805
1806        let after = store.get_run(picked.id).await.unwrap().unwrap();
1807        assert_eq!(after.status.state, RunStatus::Pending);
1808        assert_eq!(after.retry_count, 0);
1809        assert_eq!(after.lease_recoveries, 1);
1810        assert!(after.error.is_none());
1811        assert!(after.worker_id.is_none());
1812        assert!(after.lease_expires_at.is_none());
1813    }
1814
1815    #[tokio::test]
1816    async fn reap_expired_leases_requeued_run_is_pickable_again() {
1817        let store = InMemoryStore::new();
1818        let picked = pick_with_expired_lease(&store, 3).await;
1819        store.reap_expired_leases(100).await.unwrap();
1820
1821        let repicked = store
1822            .pick_next_pending(lease("worker-2", 90))
1823            .await
1824            .unwrap()
1825            .unwrap();
1826
1827        assert_eq!(repicked.id, picked.id);
1828        assert_eq!(repicked.worker_id.as_deref(), Some("worker-2"));
1829    }
1830
1831    #[tokio::test]
1832    async fn reap_expired_leases_ignores_valid_lease() {
1833        let store = InMemoryStore::new();
1834        store.create_run(new_run_req("test")).await.unwrap();
1835        let picked = store
1836            .pick_next_pending(lease("worker-1", 90))
1837            .await
1838            .unwrap()
1839            .unwrap();
1840
1841        assert!(store.reap_expired_leases(100).await.unwrap().is_empty());
1842
1843        let after = store.get_run(picked.id).await.unwrap().unwrap();
1844        assert_eq!(after.status.state, RunStatus::Running);
1845        assert_eq!(after.retry_count, 0);
1846    }
1847
1848    #[tokio::test]
1849    async fn reap_expired_leases_ignores_run_without_lease() {
1850        let store = InMemoryStore::new();
1851        store.create_run(new_run_req("test")).await.unwrap();
1852        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
1853
1854        assert!(store.reap_expired_leases(100).await.unwrap().is_empty());
1855
1856        let after = store.get_run(picked.id).await.unwrap().unwrap();
1857        assert_eq!(after.status.state, RunStatus::Running);
1858    }
1859
1860    #[tokio::test]
1861    async fn reap_expired_leases_fails_run_when_retries_exhausted() {
1862        let store = InMemoryStore::new();
1863        let picked = pick_with_expired_lease(&store, 0).await;
1864
1865        let reaped = store.reap_expired_leases(100).await.unwrap();
1866
1867        assert_eq!(reaped[0].to, RunStatus::Failed);
1868        let after = store.get_run(picked.id).await.unwrap().unwrap();
1869        assert_eq!(after.status.state, RunStatus::Failed);
1870        assert_eq!(after.error.as_deref(), Some(LEASE_EXPIRED_ERROR));
1871        assert!(after.completed_at.is_some());
1872    }
1873
1874    #[tokio::test]
1875    async fn reap_expired_leases_fails_after_max_retries_recoveries() {
1876        let store = InMemoryStore::new();
1877        let picked = pick_with_expired_lease(&store, 2).await;
1878
1879        // Recovery 1 and 2 requeue, recovery 3 gives up.
1880        for expected in [RunStatus::Pending, RunStatus::Pending, RunStatus::Failed] {
1881            let reaped = store.reap_expired_leases(100).await.unwrap();
1882            assert_eq!(reaped[0].to, expected);
1883            if expected == RunStatus::Pending {
1884                expire_now(&store).await;
1885            }
1886        }
1887
1888        let after = store.get_run(picked.id).await.unwrap().unwrap();
1889        assert_eq!(after.lease_recoveries, 3);
1890        assert_eq!(after.retry_count, 0);
1891        assert_eq!(after.error.as_deref(), Some(LEASE_EXPIRED_ERROR));
1892    }
1893
1894    #[tokio::test]
1895    async fn reap_expired_leases_keeps_the_attempt_number() {
1896        let store = InMemoryStore::new();
1897        let picked = pick_with_expired_lease(&store, 3).await;
1898        let before = store
1899            .create_step(new_step_req(picked.id, "build", 0))
1900            .await
1901            .unwrap();
1902
1903        store.reap_expired_leases(100).await.unwrap();
1904        let repicked = store
1905            .pick_next_pending(lease("worker-2", 90))
1906            .await
1907            .unwrap()
1908            .unwrap();
1909        let after = store
1910            .create_step(new_step_req(repicked.id, "build", 0))
1911            .await
1912            .unwrap();
1913
1914        assert_eq!(before.attempt, 1);
1915        assert_eq!(after.attempt, before.attempt);
1916        assert_eq!(repicked.retry_count, 0);
1917        assert_eq!(repicked.lease_recoveries, 1);
1918    }
1919
1920    #[tokio::test]
1921    async fn reap_expired_leases_counts_apart_from_handler_retries() {
1922        let store = InMemoryStore::new();
1923        let picked = pick_with_expired_lease(&store, 1).await;
1924        store
1925            .update_run(
1926                picked.id,
1927                RunUpdate {
1928                    increment_retry: true,
1929                    ..RunUpdate::default()
1930                },
1931            )
1932            .await
1933            .unwrap();
1934
1935        // One handler retry already consumed max_retries: the lease recovery
1936        // budget is separate, so the first reap still requeues.
1937        let reaped = store.reap_expired_leases(100).await.unwrap();
1938
1939        assert_eq!(reaped[0].to, RunStatus::Pending);
1940        assert_eq!(reaped[0].run.retry_count, 1);
1941        assert_eq!(reaped[0].run.lease_recoveries, 1);
1942    }
1943
1944    #[tokio::test]
1945    async fn reap_expired_leases_respects_limit() {
1946        let store = InMemoryStore::new();
1947        for _ in 0..3 {
1948            pick_with_expired_lease(&store, 3).await;
1949        }
1950
1951        let reaped = store.reap_expired_leases(2).await.unwrap();
1952        assert_eq!(reaped.len(), 2);
1953
1954        let rest = store.reap_expired_leases(100).await.unwrap();
1955        assert_eq!(rest.len(), 1);
1956    }
1957
1958    // ---- pick_next_pending ----
1959
1960    #[tokio::test]
1961    async fn pick_next_pending_empty_store() {
1962        let store = InMemoryStore::new();
1963        let result = store.pick_next_pending(None).await.unwrap();
1964        assert!(result.is_none());
1965    }
1966
1967    #[tokio::test]
1968    async fn pick_next_pending_returns_oldest_and_transitions_to_running() {
1969        let store = InMemoryStore::new();
1970        let r1 = store
1971            .create_run(new_run_req("first"))
1972            .await
1973            .unwrap()
1974            .into_run();
1975        let _r2 = store
1976            .create_run(new_run_req("second"))
1977            .await
1978            .unwrap()
1979            .into_run();
1980
1981        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
1982        assert_eq!(picked.id, r1.id);
1983        assert_eq!(picked.status.state, RunStatus::Running);
1984        assert!(picked.started_at.is_some());
1985
1986        // Verify it's Running in the store too.
1987        let fetched = store.get_run(r1.id).await.unwrap().unwrap();
1988        assert_eq!(fetched.status.state, RunStatus::Running);
1989    }
1990
1991    #[tokio::test]
1992    async fn pick_next_pending_skips_non_pending() {
1993        let store = InMemoryStore::new();
1994        let r1 = store.create_run(new_run_req("a")).await.unwrap().into_run();
1995        let r2 = store.create_run(new_run_req("b")).await.unwrap().into_run();
1996
1997        // Transition r1 to Running.
1998        store
1999            .update_run_status(r1.id, RunStatus::Running)
2000            .await
2001            .unwrap();
2002
2003        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
2004        assert_eq!(picked.id, r2.id);
2005    }
2006
2007    // ---- create_step ----
2008
2009    #[tokio::test]
2010    async fn create_step_returns_pending() {
2011        let store = InMemoryStore::new();
2012        let run = store
2013            .create_run(new_run_req("test"))
2014            .await
2015            .unwrap()
2016            .into_run();
2017
2018        let step = store
2019            .create_step(NewStep {
2020                run_id: run.id,
2021                trace_id: step_trace_id(run.id, "build", 0),
2022                name: "build".to_string(),
2023                kind: crate::entities::StepKind::Shell,
2024                position: 0,
2025                input: Some(json!({"command": "cargo build"})),
2026                is_error_handler: false,
2027            })
2028            .await
2029            .unwrap();
2030
2031        assert_eq!(step.status.state, StepStatus::Pending);
2032        assert_eq!(step.name, "build");
2033        assert_eq!(step.run_id, run.id);
2034        assert_eq!(step.position, 0);
2035    }
2036
2037    #[tokio::test]
2038    async fn create_step_for_missing_run_returns_error() {
2039        let store = InMemoryStore::new();
2040        let result = store
2041            .create_step(NewStep {
2042                run_id: Uuid::nil(),
2043                trace_id: step_trace_id(Uuid::nil(), "build", 0),
2044                name: "build".to_string(),
2045                kind: crate::entities::StepKind::Shell,
2046                position: 0,
2047                input: None,
2048                is_error_handler: false,
2049            })
2050            .await;
2051        assert!(matches!(result.unwrap_err(), StoreError::RunNotFound(_)));
2052    }
2053
2054    // ---- update_step ----
2055
2056    #[tokio::test]
2057    async fn update_step_applies_partial_update() {
2058        let store = InMemoryStore::new();
2059        let run = store
2060            .create_run(new_run_req("test"))
2061            .await
2062            .unwrap()
2063            .into_run();
2064
2065        let step = store
2066            .create_step(NewStep {
2067                run_id: run.id,
2068                trace_id: step_trace_id(run.id, "build", 0),
2069                name: "build".to_string(),
2070                kind: crate::entities::StepKind::Shell,
2071                position: 0,
2072                input: None,
2073                is_error_handler: false,
2074            })
2075            .await
2076            .unwrap();
2077
2078        // Transition Pending → Running first
2079        store
2080            .update_step(
2081                step.id,
2082                StepUpdate {
2083                    status: Some(StepStatus::Running),
2084                    ..StepUpdate::default()
2085                },
2086            )
2087            .await
2088            .unwrap();
2089
2090        // Then Running → Completed
2091        store
2092            .update_step(
2093                step.id,
2094                StepUpdate {
2095                    status: Some(StepStatus::Completed),
2096                    output: Some(json!({"stdout": "ok"})),
2097                    duration_ms: Some(150),
2098                    ..StepUpdate::default()
2099                },
2100            )
2101            .await
2102            .unwrap();
2103
2104        let steps = store.list_steps(run.id).await.unwrap();
2105        assert_eq!(steps.len(), 1);
2106        assert_eq!(steps[0].status.state, StepStatus::Completed);
2107        assert_eq!(steps[0].duration_ms, 150);
2108        assert!(steps[0].output.is_some());
2109    }
2110
2111    #[tokio::test]
2112    async fn update_step_not_found() {
2113        let store = InMemoryStore::new();
2114        let result = store.update_step(Uuid::nil(), StepUpdate::default()).await;
2115        assert!(matches!(result.unwrap_err(), StoreError::StepNotFound(_)));
2116    }
2117
2118    // ---- list_steps ----
2119
2120    #[tokio::test]
2121    async fn list_steps_ordered_by_position() {
2122        let store = InMemoryStore::new();
2123        let run = store
2124            .create_run(new_run_req("test"))
2125            .await
2126            .unwrap()
2127            .into_run();
2128
2129        // Insert out of order.
2130        store
2131            .create_step(NewStep {
2132                run_id: run.id,
2133                trace_id: step_trace_id(run.id, "deploy", 2),
2134                name: "deploy".to_string(),
2135                kind: crate::entities::StepKind::Shell,
2136                position: 2,
2137                input: None,
2138                is_error_handler: false,
2139            })
2140            .await
2141            .unwrap();
2142        store
2143            .create_step(NewStep {
2144                run_id: run.id,
2145                trace_id: step_trace_id(run.id, "build", 0),
2146                name: "build".to_string(),
2147                kind: crate::entities::StepKind::Shell,
2148                position: 0,
2149                input: None,
2150                is_error_handler: false,
2151            })
2152            .await
2153            .unwrap();
2154        store
2155            .create_step(NewStep {
2156                run_id: run.id,
2157                trace_id: step_trace_id(run.id, "test", 1),
2158                name: "test".to_string(),
2159                kind: crate::entities::StepKind::Shell,
2160                position: 1,
2161                input: None,
2162                is_error_handler: false,
2163            })
2164            .await
2165            .unwrap();
2166
2167        let steps = store.list_steps(run.id).await.unwrap();
2168        assert_eq!(steps.len(), 3);
2169        assert_eq!(steps[0].name, "build");
2170        assert_eq!(steps[1].name, "test");
2171        assert_eq!(steps[2].name, "deploy");
2172    }
2173
2174    #[tokio::test]
2175    async fn list_steps_empty_for_run_without_steps() {
2176        let store = InMemoryStore::new();
2177        let run = store
2178            .create_run(new_run_req("test"))
2179            .await
2180            .unwrap()
2181            .into_run();
2182        let steps = store.list_steps(run.id).await.unwrap();
2183        assert!(steps.is_empty());
2184    }
2185
2186    // ---- update_run ----
2187
2188    #[tokio::test]
2189    async fn update_run_applies_cost_and_duration() {
2190        let store = InMemoryStore::new();
2191        let run = store
2192            .create_run(new_run_req("test"))
2193            .await
2194            .unwrap()
2195            .into_run();
2196
2197        store
2198            .update_run(
2199                run.id,
2200                RunUpdate {
2201                    cost_usd: Some(Decimal::new(123, 2)),
2202                    duration_ms: Some(5000),
2203                    ..RunUpdate::default()
2204                },
2205            )
2206            .await
2207            .unwrap();
2208
2209        let fetched = store.get_run(run.id).await.unwrap().unwrap();
2210        assert_eq!(fetched.cost_usd, Decimal::new(123, 2));
2211        assert_eq!(fetched.duration_ms, 5000);
2212    }
2213
2214    #[tokio::test]
2215    async fn update_run_increment_retry() {
2216        let store = InMemoryStore::new();
2217        let run = store
2218            .create_run(new_run_req("test"))
2219            .await
2220            .unwrap()
2221            .into_run();
2222        assert_eq!(run.retry_count, 0);
2223
2224        store
2225            .update_run(
2226                run.id,
2227                RunUpdate {
2228                    increment_retry: true,
2229                    ..RunUpdate::default()
2230                },
2231            )
2232            .await
2233            .unwrap();
2234
2235        let fetched = store.get_run(run.id).await.unwrap().unwrap();
2236        assert_eq!(fetched.retry_count, 1);
2237    }
2238
2239    #[tokio::test]
2240    async fn update_run_not_found() {
2241        let store = InMemoryStore::new();
2242        let result = store.update_run(Uuid::nil(), RunUpdate::default()).await;
2243        assert!(matches!(result.unwrap_err(), StoreError::RunNotFound(_)));
2244    }
2245
2246    // ---- concurrent access ----
2247
2248    #[tokio::test]
2249    async fn concurrent_pick_next_pending_no_double_pick() {
2250        let store = InMemoryStore::new();
2251
2252        // Create 10 pending runs.
2253        for i in 0..10 {
2254            store
2255                .create_run(new_run_req(&format!("wf-{i}")))
2256                .await
2257                .unwrap()
2258                .into_run();
2259        }
2260
2261        // Concurrently pick from multiple tasks.
2262        let mut handles = Vec::new();
2263        for _ in 0..10 {
2264            let s = store.clone();
2265            handles.push(spawn(async move { s.pick_next_pending(None).await }));
2266        }
2267
2268        let mut picked_ids = Vec::new();
2269        for h in handles {
2270            if let Ok(Ok(Some(run))) = h.await {
2271                picked_ids.push(run.id);
2272            }
2273        }
2274
2275        // Each run should be picked at most once.
2276        let unique: std::collections::HashSet<_> = picked_ids.iter().collect();
2277        assert_eq!(unique.len(), picked_ids.len());
2278    }
2279
2280    // ---- get_stats ----
2281
2282    #[tokio::test]
2283    async fn get_stats_empty_store() {
2284        let store = InMemoryStore::new();
2285        let stats = store.get_stats(RunFilter::default()).await.unwrap();
2286        assert_eq!(stats.total_runs, 0);
2287        assert_eq!(stats.completed_runs, 0);
2288        assert_eq!(stats.failed_runs, 0);
2289        assert_eq!(stats.cancelled_runs, 0);
2290        assert_eq!(stats.active_runs, 0);
2291        assert_eq!(stats.total_cost_usd, Decimal::ZERO);
2292        assert_eq!(stats.total_duration_ms, 0);
2293    }
2294
2295    #[tokio::test]
2296    async fn get_stats_aggregates_counts_and_totals() {
2297        let store = InMemoryStore::new();
2298
2299        // Create runs in various states.
2300        let r1 = store
2301            .create_run(new_run_req("wf1"))
2302            .await
2303            .unwrap()
2304            .into_run();
2305        let r2 = store
2306            .create_run(new_run_req("wf2"))
2307            .await
2308            .unwrap()
2309            .into_run();
2310        let r3 = store
2311            .create_run(new_run_req("wf3"))
2312            .await
2313            .unwrap()
2314            .into_run();
2315        let _r4 = store
2316            .create_run(new_run_req("wf4"))
2317            .await
2318            .unwrap()
2319            .into_run();
2320
2321        // Transition r1 to Running, then Completed.
2322        store
2323            .update_run_status(r1.id, RunStatus::Running)
2324            .await
2325            .unwrap();
2326        store
2327            .update_run_status(r1.id, RunStatus::Completed)
2328            .await
2329            .unwrap();
2330
2331        // Transition r2 to Running, then Failed.
2332        store
2333            .update_run_status(r2.id, RunStatus::Running)
2334            .await
2335            .unwrap();
2336        store
2337            .update_run_status(r2.id, RunStatus::Failed)
2338            .await
2339            .unwrap();
2340
2341        // Transition r3 to Cancelled.
2342        store
2343            .update_run_status(r3.id, RunStatus::Cancelled)
2344            .await
2345            .unwrap();
2346
2347        // r4 remains Pending (active).
2348
2349        // Update cost and duration.
2350        store
2351            .update_run(
2352                r1.id,
2353                RunUpdate {
2354                    cost_usd: Some(Decimal::new(1000, 2)),
2355                    duration_ms: Some(1000),
2356                    ..RunUpdate::default()
2357                },
2358            )
2359            .await
2360            .unwrap();
2361
2362        store
2363            .update_run(
2364                r2.id,
2365                RunUpdate {
2366                    cost_usd: Some(Decimal::new(500, 2)),
2367                    duration_ms: Some(500),
2368                    ..RunUpdate::default()
2369                },
2370            )
2371            .await
2372            .unwrap();
2373
2374        let stats = store.get_stats(RunFilter::default()).await.unwrap();
2375        assert_eq!(stats.total_runs, 4);
2376        assert_eq!(stats.completed_runs, 1);
2377        assert_eq!(stats.failed_runs, 1);
2378        assert_eq!(stats.cancelled_runs, 1);
2379        assert_eq!(stats.active_runs, 1); // r4 is Pending
2380        assert_eq!(stats.total_cost_usd, Decimal::new(1500, 2));
2381        assert_eq!(stats.total_duration_ms, 1500);
2382    }
2383
2384    #[tokio::test]
2385    async fn update_run_status_running_to_retrying() {
2386        let store = InMemoryStore::new();
2387        let run = store
2388            .create_run(new_run_req("test"))
2389            .await
2390            .unwrap()
2391            .into_run();
2392
2393        store
2394            .update_run_status(run.id, RunStatus::Running)
2395            .await
2396            .unwrap();
2397
2398        store
2399            .update_run_status(run.id, RunStatus::Retrying)
2400            .await
2401            .unwrap();
2402
2403        let fetched = store.get_run(run.id).await.unwrap().unwrap();
2404        assert_eq!(fetched.status.state, RunStatus::Retrying);
2405        assert!(!fetched.status.state.is_terminal());
2406        assert!(fetched.completed_at.is_none()); // Not a terminal state
2407    }
2408
2409    #[tokio::test]
2410    async fn update_run_status_retrying_to_running_allowed() {
2411        let store = InMemoryStore::new();
2412        let run = store
2413            .create_run(new_run_req("test"))
2414            .await
2415            .unwrap()
2416            .into_run();
2417
2418        store
2419            .update_run_status(run.id, RunStatus::Running)
2420            .await
2421            .unwrap();
2422        store
2423            .update_run_status(run.id, RunStatus::Retrying)
2424            .await
2425            .unwrap();
2426
2427        // Retrying → Running should be allowed
2428        store
2429            .update_run_status(run.id, RunStatus::Running)
2430            .await
2431            .unwrap();
2432
2433        let fetched = store.get_run(run.id).await.unwrap().unwrap();
2434        assert_eq!(fetched.status.state, RunStatus::Running);
2435    }
2436
2437    #[tokio::test]
2438    async fn update_run_with_invalid_status_transition_errors() {
2439        let store = InMemoryStore::new();
2440        let run = store
2441            .create_run(new_run_req("test"))
2442            .await
2443            .unwrap()
2444            .into_run();
2445
2446        // Try to apply invalid status transition via update_run
2447        let result = store
2448            .update_run(
2449                run.id,
2450                RunUpdate {
2451                    status: Some(RunStatus::Completed), // invalid from Pending
2452                    ..RunUpdate::default()
2453                },
2454            )
2455            .await;
2456
2457        assert!(result.is_err());
2458    }
2459
2460    #[tokio::test]
2461    async fn create_step_with_complex_input() {
2462        let store = InMemoryStore::new();
2463        let run = store
2464            .create_run(new_run_req("test"))
2465            .await
2466            .unwrap()
2467            .into_run();
2468
2469        let complex_input = json!({
2470            "command": "cargo build",
2471            "env": {
2472                "RUST_LOG": "debug",
2473                "CUSTOM": "value"
2474            },
2475            "timeout": 60,
2476            "retry_policy": {
2477                "max_attempts": 3,
2478                "backoff": "exponential"
2479            }
2480        });
2481
2482        let step = store
2483            .create_step(NewStep {
2484                run_id: run.id,
2485                trace_id: step_trace_id(run.id, "build", 0),
2486                name: "build".to_string(),
2487                kind: crate::entities::StepKind::Agent,
2488                position: 0,
2489                input: Some(complex_input.clone()),
2490                is_error_handler: false,
2491            })
2492            .await
2493            .unwrap();
2494
2495        assert_eq!(step.input, Some(complex_input));
2496    }
2497
2498    #[tokio::test]
2499    async fn update_step_with_error_message() {
2500        let store = InMemoryStore::new();
2501        let run = store
2502            .create_run(new_run_req("test"))
2503            .await
2504            .unwrap()
2505            .into_run();
2506
2507        let step = store
2508            .create_step(NewStep {
2509                run_id: run.id,
2510                trace_id: step_trace_id(run.id, "build", 0),
2511                name: "build".to_string(),
2512                kind: crate::entities::StepKind::Shell,
2513                position: 0,
2514                input: None,
2515                is_error_handler: false,
2516            })
2517            .await
2518            .unwrap();
2519
2520        store
2521            .update_step(
2522                step.id,
2523                StepUpdate {
2524                    status: Some(StepStatus::Running),
2525                    ..StepUpdate::default()
2526                },
2527            )
2528            .await
2529            .unwrap();
2530
2531        store
2532            .update_step(
2533                step.id,
2534                StepUpdate {
2535                    status: Some(StepStatus::Failed),
2536                    error: Some("Connection timeout after 30s".to_string()),
2537                    duration_ms: Some(30000),
2538                    ..StepUpdate::default()
2539                },
2540            )
2541            .await
2542            .unwrap();
2543
2544        let steps = store.list_steps(run.id).await.unwrap();
2545        assert_eq!(steps[0].status.state, StepStatus::Failed);
2546        assert_eq!(
2547            steps[0].error,
2548            Some("Connection timeout after 30s".to_string())
2549        );
2550        assert_eq!(steps[0].duration_ms, 30000);
2551    }
2552
2553    #[tokio::test]
2554    async fn list_steps_for_nonexistent_run_returns_empty() {
2555        let store = InMemoryStore::new();
2556        let steps = store.list_steps(Uuid::nil()).await.unwrap();
2557        assert!(steps.is_empty());
2558    }
2559
2560    #[tokio::test]
2561    async fn update_step_pending_to_skipped() {
2562        let store = InMemoryStore::new();
2563        let run = store
2564            .create_run(new_run_req("test"))
2565            .await
2566            .unwrap()
2567            .into_run();
2568
2569        let step = store
2570            .create_step(NewStep {
2571                run_id: run.id,
2572                trace_id: step_trace_id(run.id, "build", 0),
2573                name: "build".to_string(),
2574                kind: crate::entities::StepKind::Shell,
2575                position: 0,
2576                input: None,
2577                is_error_handler: false,
2578            })
2579            .await
2580            .unwrap();
2581
2582        // Pending → Skipped is allowed (when prior step failed)
2583        store
2584            .update_step(
2585                step.id,
2586                StepUpdate {
2587                    status: Some(StepStatus::Skipped),
2588                    ..StepUpdate::default()
2589                },
2590            )
2591            .await
2592            .unwrap();
2593
2594        let steps = store.list_steps(run.id).await.unwrap();
2595        assert_eq!(steps[0].status.state, StepStatus::Skipped);
2596    }
2597
2598    #[tokio::test]
2599    async fn list_runs_with_combined_filters() {
2600        let store = InMemoryStore::new();
2601
2602        let r1 = store
2603            .create_run(new_run_req("deploy"))
2604            .await
2605            .unwrap()
2606            .into_run();
2607        let r2 = store
2608            .create_run(new_run_req("deploy"))
2609            .await
2610            .unwrap()
2611            .into_run();
2612        let _r3 = store
2613            .create_run(new_run_req("test"))
2614            .await
2615            .unwrap()
2616            .into_run();
2617
2618        // r1: Pending → Running → Completed
2619        store
2620            .update_run_status(r1.id, RunStatus::Running)
2621            .await
2622            .unwrap();
2623        store
2624            .update_run_status(r1.id, RunStatus::Completed)
2625            .await
2626            .unwrap();
2627
2628        // r2: Pending → Running
2629        store
2630            .update_run_status(r2.id, RunStatus::Running)
2631            .await
2632            .unwrap();
2633
2634        // Filter by workflow AND status
2635        let filter = RunFilter {
2636            workflow_name: Some("deploy".to_string()),
2637            status: Some(RunStatus::Running),
2638            ..RunFilter::default()
2639        };
2640
2641        let page = store.list_runs(filter, 1, 100).await.unwrap();
2642        assert_eq!(page.total, 1);
2643        assert_eq!(page.items[0].id, r2.id);
2644    }
2645
2646    #[tokio::test]
2647    async fn list_runs_workflow_filter_is_case_insensitive_partial_match() {
2648        let store = InMemoryStore::new();
2649        store
2650            .create_run(new_run_req("weather-report"))
2651            .await
2652            .unwrap()
2653            .into_run();
2654        store
2655            .create_run(new_run_req("deploy-prod"))
2656            .await
2657            .unwrap()
2658            .into_run();
2659
2660        // Partial match
2661        let filter = RunFilter {
2662            workflow_name: Some("weather".to_string()),
2663            ..RunFilter::default()
2664        };
2665        let page = store.list_runs(filter, 1, 100).await.unwrap();
2666        assert_eq!(page.total, 1);
2667        assert_eq!(page.items[0].workflow_name, "weather-report");
2668
2669        // Case-insensitive
2670        let filter = RunFilter {
2671            workflow_name: Some("Weather-REPORT".to_string()),
2672            ..RunFilter::default()
2673        };
2674        let page = store.list_runs(filter, 1, 100).await.unwrap();
2675        assert_eq!(page.total, 1);
2676        assert_eq!(page.items[0].workflow_name, "weather-report");
2677
2678        // Suffix match
2679        let filter = RunFilter {
2680            workflow_name: Some("report".to_string()),
2681            ..RunFilter::default()
2682        };
2683        let page = store.list_runs(filter, 1, 100).await.unwrap();
2684        assert_eq!(page.total, 1);
2685        assert_eq!(page.items[0].workflow_name, "weather-report");
2686
2687        // No match
2688        let filter = RunFilter {
2689            workflow_name: Some("build".to_string()),
2690            ..RunFilter::default()
2691        };
2692        let page = store.list_runs(filter, 1, 100).await.unwrap();
2693        assert_eq!(page.total, 0);
2694    }
2695
2696    #[tokio::test]
2697    async fn list_runs_has_steps_true_only_filters_completed_and_cancelled() {
2698        let store = InMemoryStore::new();
2699        let run_with = create_terminal_run(&store, "with-steps", RunStatus::Completed).await;
2700        let _run_without = create_terminal_run(&store, "without-steps", RunStatus::Completed).await;
2701
2702        store
2703            .create_step(NewStep {
2704                run_id: run_with.id,
2705                trace_id: step_trace_id(run_with.id, "build", 0),
2706                name: "build".to_string(),
2707                kind: crate::entities::StepKind::Shell,
2708                position: 0,
2709                input: None,
2710                is_error_handler: false,
2711            })
2712            .await
2713            .unwrap();
2714
2715        let filter = RunFilter {
2716            has_steps: Some(true),
2717            ..RunFilter::default()
2718        };
2719        let page = store.list_runs(filter, 1, 100).await.unwrap();
2720        assert_eq!(page.total, 1);
2721        assert_eq!(page.items[0].id, run_with.id);
2722    }
2723
2724    #[tokio::test]
2725    async fn list_runs_has_steps_false_only_filters_completed_and_cancelled() {
2726        let store = InMemoryStore::new();
2727        let run_with = create_terminal_run(&store, "with-steps", RunStatus::Cancelled).await;
2728        let run_without = create_terminal_run(&store, "without-steps", RunStatus::Cancelled).await;
2729
2730        store
2731            .create_step(NewStep {
2732                run_id: run_with.id,
2733                trace_id: step_trace_id(run_with.id, "build", 0),
2734                name: "build".to_string(),
2735                kind: crate::entities::StepKind::Shell,
2736                position: 0,
2737                input: None,
2738                is_error_handler: false,
2739            })
2740            .await
2741            .unwrap();
2742
2743        let filter = RunFilter {
2744            has_steps: Some(false),
2745            ..RunFilter::default()
2746        };
2747        let page = store.list_runs(filter, 1, 100).await.unwrap();
2748        assert_eq!(page.total, 1);
2749        assert_eq!(page.items[0].id, run_without.id);
2750    }
2751
2752    #[tokio::test]
2753    async fn list_runs_has_steps_none_returns_all() {
2754        let store = InMemoryStore::new();
2755        let run_with = store
2756            .create_run(new_run_req("with-steps"))
2757            .await
2758            .unwrap()
2759            .into_run();
2760        let _run_without = store
2761            .create_run(new_run_req("without-steps"))
2762            .await
2763            .unwrap()
2764            .into_run();
2765
2766        store
2767            .create_step(NewStep {
2768                run_id: run_with.id,
2769                trace_id: step_trace_id(run_with.id, "build", 0),
2770                name: "build".to_string(),
2771                kind: crate::entities::StepKind::Shell,
2772                position: 0,
2773                input: None,
2774                is_error_handler: false,
2775            })
2776            .await
2777            .unwrap();
2778
2779        let filter = RunFilter {
2780            has_steps: None,
2781            ..RunFilter::default()
2782        };
2783        let page = store.list_runs(filter, 1, 100).await.unwrap();
2784        assert_eq!(page.total, 2);
2785    }
2786
2787    #[tokio::test]
2788    async fn list_runs_has_steps_true_does_not_filter_non_terminal_runs() {
2789        let store = InMemoryStore::new();
2790        let pending_run = store
2791            .create_run(new_run_req("pending-empty"))
2792            .await
2793            .unwrap()
2794            .into_run();
2795        let running_run = store
2796            .create_run(new_run_req("running-empty"))
2797            .await
2798            .unwrap()
2799            .into_run();
2800        store
2801            .update_run_status(running_run.id, RunStatus::Running)
2802            .await
2803            .unwrap();
2804
2805        let filter = RunFilter {
2806            has_steps: Some(true),
2807            ..RunFilter::default()
2808        };
2809        let page = store.list_runs(filter, 1, 100).await.unwrap();
2810        assert_eq!(page.total, 2);
2811        let ids: Vec<_> = page.items.iter().map(|r| r.id).collect();
2812        assert!(ids.contains(&pending_run.id));
2813        assert!(ids.contains(&running_run.id));
2814    }
2815
2816    #[tokio::test]
2817    async fn get_stats_with_mixed_active_statuses() {
2818        let store = InMemoryStore::new();
2819
2820        let _r1 = store
2821            .create_run(new_run_req("wf"))
2822            .await
2823            .unwrap()
2824            .into_run(); // Pending
2825        let r2 = store
2826            .create_run(new_run_req("wf"))
2827            .await
2828            .unwrap()
2829            .into_run();
2830        let r3 = store
2831            .create_run(new_run_req("wf"))
2832            .await
2833            .unwrap()
2834            .into_run();
2835
2836        store
2837            .update_run_status(r2.id, RunStatus::Running)
2838            .await
2839            .unwrap();
2840        store
2841            .update_run_status(r3.id, RunStatus::Running)
2842            .await
2843            .unwrap();
2844        store
2845            .update_run_status(r3.id, RunStatus::Retrying)
2846            .await
2847            .unwrap();
2848
2849        let r4 = store
2850            .create_run(new_run_req("wf"))
2851            .await
2852            .unwrap()
2853            .into_run();
2854        store
2855            .update_run_status(r4.id, RunStatus::Running)
2856            .await
2857            .unwrap();
2858        store
2859            .update_run_status(r4.id, RunStatus::AwaitingApproval)
2860            .await
2861            .unwrap();
2862
2863        let r5 = store
2864            .create_run(new_run_req("wf"))
2865            .await
2866            .unwrap()
2867            .into_run();
2868        store
2869            .update_run_status(r5.id, RunStatus::Running)
2870            .await
2871            .unwrap();
2872        store
2873            .update_run_status(r5.id, RunStatus::Sleeping)
2874            .await
2875            .unwrap();
2876
2877        let stats = store.get_stats(RunFilter::default()).await.unwrap();
2878        // _r1 (Pending), r2 (Running), r3 (Retrying), r4 (AwaitingApproval), r5 (Sleeping)
2879        assert_eq!(stats.active_runs, 5);
2880        assert_eq!(stats.awaiting_approval_runs, 1);
2881    }
2882
2883    #[tokio::test]
2884    async fn run_with_different_trigger_kinds() {
2885        let store = InMemoryStore::new();
2886
2887        let r1 = store
2888            .create_run(NewRun {
2889                created_by: None,
2890                workflow_name: "test".to_string(),
2891                trigger: TriggerKind::Manual,
2892                payload: json!({}),
2893                max_retries: 1,
2894                handler_version: None,
2895                labels: HashMap::new(),
2896                scheduled_at: None,
2897                idempotency_key: None,
2898                concurrency_key: None,
2899                priority: 0,
2900                concurrency_limits: Vec::new(),
2901                max_cost_usd: None,
2902                worker_tags: Vec::new(),
2903            })
2904            .await
2905            .unwrap()
2906            .into_run();
2907
2908        let r2 = store
2909            .create_run(NewRun {
2910                created_by: None,
2911                workflow_name: "test".to_string(),
2912                trigger: TriggerKind::Webhook {
2913                    path: "/hooks/github".to_string(),
2914                },
2915                payload: json!({}),
2916                max_retries: 1,
2917                handler_version: None,
2918                labels: HashMap::new(),
2919                scheduled_at: None,
2920                idempotency_key: None,
2921                concurrency_key: None,
2922                priority: 0,
2923                concurrency_limits: Vec::new(),
2924                max_cost_usd: None,
2925                worker_tags: Vec::new(),
2926            })
2927            .await
2928            .unwrap()
2929            .into_run();
2930
2931        let r3 = store
2932            .create_run(NewRun {
2933                created_by: None,
2934                workflow_name: "test".to_string(),
2935                trigger: TriggerKind::Cron {
2936                    schedule: "0 0 * * *".to_string(),
2937                    schedule_id: None,
2938                    scheduled_for: None,
2939                },
2940                payload: json!({}),
2941                max_retries: 1,
2942                handler_version: None,
2943                labels: HashMap::new(),
2944                scheduled_at: None,
2945                idempotency_key: None,
2946                concurrency_key: None,
2947                priority: 0,
2948                concurrency_limits: Vec::new(),
2949                max_cost_usd: None,
2950                worker_tags: Vec::new(),
2951            })
2952            .await
2953            .unwrap()
2954            .into_run();
2955
2956        let r4 = store
2957            .create_run(NewRun {
2958                created_by: None,
2959                workflow_name: "test".to_string(),
2960                trigger: TriggerKind::Api,
2961                payload: json!({}),
2962                max_retries: 1,
2963                handler_version: None,
2964                labels: HashMap::new(),
2965                scheduled_at: None,
2966                idempotency_key: None,
2967                concurrency_key: None,
2968                priority: 0,
2969                concurrency_limits: Vec::new(),
2970                max_cost_usd: None,
2971                worker_tags: Vec::new(),
2972            })
2973            .await
2974            .unwrap()
2975            .into_run();
2976
2977        let r5 = store
2978            .create_run(NewRun {
2979                created_by: None,
2980                workflow_name: "test".to_string(),
2981                trigger: TriggerKind::Retry {
2982                    parent_run_id: Uuid::nil(),
2983                },
2984                payload: json!({}),
2985                max_retries: 1,
2986                handler_version: None,
2987                labels: HashMap::new(),
2988                scheduled_at: None,
2989                idempotency_key: None,
2990                concurrency_key: None,
2991                priority: 0,
2992                concurrency_limits: Vec::new(),
2993                max_cost_usd: None,
2994                worker_tags: Vec::new(),
2995            })
2996            .await
2997            .unwrap()
2998            .into_run();
2999
3000        assert_eq!(r1.trigger, TriggerKind::Manual);
3001        assert!(matches!(r2.trigger, TriggerKind::Webhook { .. }));
3002        assert!(matches!(r3.trigger, TriggerKind::Cron { .. }));
3003        assert_eq!(r4.trigger, TriggerKind::Api);
3004        assert!(matches!(r5.trigger, TriggerKind::Retry { .. }));
3005    }
3006
3007    // ---- create_step_dependencies ----
3008
3009    #[tokio::test]
3010    async fn create_step_dependencies_stores_dependencies() {
3011        let store = InMemoryStore::new();
3012        let run = store
3013            .create_run(new_run_req("test"))
3014            .await
3015            .unwrap()
3016            .into_run();
3017
3018        let step1 = store
3019            .create_step(NewStep {
3020                run_id: run.id,
3021                trace_id: step_trace_id(run.id, "step1", 0),
3022                name: "step1".to_string(),
3023                kind: crate::entities::StepKind::Shell,
3024                position: 0,
3025                input: None,
3026                is_error_handler: false,
3027            })
3028            .await
3029            .unwrap();
3030
3031        let step2 = store
3032            .create_step(NewStep {
3033                run_id: run.id,
3034                trace_id: step_trace_id(run.id, "step2", 1),
3035                name: "step2".to_string(),
3036                kind: crate::entities::StepKind::Shell,
3037                position: 1,
3038                input: None,
3039                is_error_handler: false,
3040            })
3041            .await
3042            .unwrap();
3043
3044        let result = store
3045            .create_step_dependencies(vec![NewStepDependency {
3046                step_id: step2.id,
3047                depends_on: step1.id,
3048            }])
3049            .await;
3050
3051        assert!(result.is_ok());
3052
3053        let deps = store.list_step_dependencies(run.id).await.unwrap();
3054        assert_eq!(deps.len(), 1);
3055        assert_eq!(deps[0].step_id, step2.id);
3056        assert_eq!(deps[0].depends_on, step1.id);
3057    }
3058
3059    #[tokio::test]
3060    async fn create_step_dependencies_duplicate_dependencies_are_idempotent() {
3061        let store = InMemoryStore::new();
3062        let run = store
3063            .create_run(new_run_req("test"))
3064            .await
3065            .unwrap()
3066            .into_run();
3067
3068        let step1 = store
3069            .create_step(NewStep {
3070                run_id: run.id,
3071                trace_id: step_trace_id(run.id, "step1", 0),
3072                name: "step1".to_string(),
3073                kind: crate::entities::StepKind::Shell,
3074                position: 0,
3075                input: None,
3076                is_error_handler: false,
3077            })
3078            .await
3079            .unwrap();
3080
3081        let step2 = store
3082            .create_step(NewStep {
3083                run_id: run.id,
3084                trace_id: step_trace_id(run.id, "step2", 1),
3085                name: "step2".to_string(),
3086                kind: crate::entities::StepKind::Shell,
3087                position: 1,
3088                input: None,
3089                is_error_handler: false,
3090            })
3091            .await
3092            .unwrap();
3093
3094        let dep = NewStepDependency {
3095            step_id: step2.id,
3096            depends_on: step1.id,
3097        };
3098
3099        store
3100            .create_step_dependencies(vec![dep.clone()])
3101            .await
3102            .unwrap();
3103        store.create_step_dependencies(vec![dep]).await.unwrap();
3104
3105        let deps = store.list_step_dependencies(run.id).await.unwrap();
3106        assert_eq!(deps.len(), 1);
3107    }
3108
3109    #[tokio::test]
3110    async fn create_step_dependencies_missing_step_id_returns_error() {
3111        let store = InMemoryStore::new();
3112        let run = store
3113            .create_run(new_run_req("test"))
3114            .await
3115            .unwrap()
3116            .into_run();
3117
3118        let step1 = store
3119            .create_step(NewStep {
3120                run_id: run.id,
3121                trace_id: step_trace_id(run.id, "step1", 0),
3122                name: "step1".to_string(),
3123                kind: crate::entities::StepKind::Shell,
3124                position: 0,
3125                input: None,
3126                is_error_handler: false,
3127            })
3128            .await
3129            .unwrap();
3130
3131        let result = store
3132            .create_step_dependencies(vec![NewStepDependency {
3133                step_id: Uuid::nil(),
3134                depends_on: step1.id,
3135            }])
3136            .await;
3137
3138        assert!(matches!(result.unwrap_err(), StoreError::StepNotFound(_)));
3139    }
3140
3141    #[tokio::test]
3142    async fn create_step_dependencies_missing_depends_on_returns_error() {
3143        let store = InMemoryStore::new();
3144        let run = store
3145            .create_run(new_run_req("test"))
3146            .await
3147            .unwrap()
3148            .into_run();
3149
3150        let step1 = store
3151            .create_step(NewStep {
3152                run_id: run.id,
3153                trace_id: step_trace_id(run.id, "step1", 0),
3154                name: "step1".to_string(),
3155                kind: crate::entities::StepKind::Shell,
3156                position: 0,
3157                input: None,
3158                is_error_handler: false,
3159            })
3160            .await
3161            .unwrap();
3162
3163        let result = store
3164            .create_step_dependencies(vec![NewStepDependency {
3165                step_id: step1.id,
3166                depends_on: Uuid::nil(),
3167            }])
3168            .await;
3169
3170        assert!(matches!(result.unwrap_err(), StoreError::StepNotFound(_)));
3171    }
3172
3173    #[tokio::test]
3174    async fn create_step_dependencies_multiple_dependencies() {
3175        let store = InMemoryStore::new();
3176        let run = store
3177            .create_run(new_run_req("test"))
3178            .await
3179            .unwrap()
3180            .into_run();
3181
3182        let step1 = store
3183            .create_step(NewStep {
3184                run_id: run.id,
3185                trace_id: step_trace_id(run.id, "step1", 0),
3186                name: "step1".to_string(),
3187                kind: crate::entities::StepKind::Shell,
3188                position: 0,
3189                input: None,
3190                is_error_handler: false,
3191            })
3192            .await
3193            .unwrap();
3194
3195        let step2 = store
3196            .create_step(NewStep {
3197                run_id: run.id,
3198                trace_id: step_trace_id(run.id, "step2", 1),
3199                name: "step2".to_string(),
3200                kind: crate::entities::StepKind::Shell,
3201                position: 1,
3202                input: None,
3203                is_error_handler: false,
3204            })
3205            .await
3206            .unwrap();
3207
3208        let step3 = store
3209            .create_step(NewStep {
3210                run_id: run.id,
3211                trace_id: step_trace_id(run.id, "step3", 2),
3212                name: "step3".to_string(),
3213                kind: crate::entities::StepKind::Shell,
3214                position: 2,
3215                input: None,
3216                is_error_handler: false,
3217            })
3218            .await
3219            .unwrap();
3220
3221        let result = store
3222            .create_step_dependencies(vec![
3223                NewStepDependency {
3224                    step_id: step2.id,
3225                    depends_on: step1.id,
3226                },
3227                NewStepDependency {
3228                    step_id: step3.id,
3229                    depends_on: step2.id,
3230                },
3231            ])
3232            .await;
3233
3234        assert!(result.is_ok());
3235
3236        let deps = store.list_step_dependencies(run.id).await.unwrap();
3237        assert_eq!(deps.len(), 2);
3238    }
3239
3240    // ---- list_step_dependencies ----
3241
3242    #[tokio::test]
3243    async fn list_step_dependencies_empty_for_run_with_no_dependencies() {
3244        let store = InMemoryStore::new();
3245        let run = store
3246            .create_run(new_run_req("test"))
3247            .await
3248            .unwrap()
3249            .into_run();
3250
3251        store
3252            .create_step(NewStep {
3253                run_id: run.id,
3254                trace_id: step_trace_id(run.id, "step1", 0),
3255                name: "step1".to_string(),
3256                kind: crate::entities::StepKind::Shell,
3257                position: 0,
3258                input: None,
3259                is_error_handler: false,
3260            })
3261            .await
3262            .unwrap();
3263
3264        let deps = store.list_step_dependencies(run.id).await.unwrap();
3265        assert!(deps.is_empty());
3266    }
3267
3268    #[tokio::test]
3269    async fn list_step_dependencies_returns_only_deps_for_given_run() {
3270        let store = InMemoryStore::new();
3271        let run1 = store
3272            .create_run(new_run_req("test1"))
3273            .await
3274            .unwrap()
3275            .into_run();
3276        let run2 = store
3277            .create_run(new_run_req("test2"))
3278            .await
3279            .unwrap()
3280            .into_run();
3281
3282        let step1_run1 = store
3283            .create_step(NewStep {
3284                run_id: run1.id,
3285                trace_id: step_trace_id(run1.id, "step1", 0),
3286                name: "step1".to_string(),
3287                kind: crate::entities::StepKind::Shell,
3288                position: 0,
3289                input: None,
3290                is_error_handler: false,
3291            })
3292            .await
3293            .unwrap();
3294
3295        let step2_run1 = store
3296            .create_step(NewStep {
3297                run_id: run1.id,
3298                trace_id: step_trace_id(run1.id, "step2", 1),
3299                name: "step2".to_string(),
3300                kind: crate::entities::StepKind::Shell,
3301                position: 1,
3302                input: None,
3303                is_error_handler: false,
3304            })
3305            .await
3306            .unwrap();
3307
3308        let step1_run2 = store
3309            .create_step(NewStep {
3310                run_id: run2.id,
3311                trace_id: step_trace_id(run2.id, "step1", 0),
3312                name: "step1".to_string(),
3313                kind: crate::entities::StepKind::Shell,
3314                position: 0,
3315                input: None,
3316                is_error_handler: false,
3317            })
3318            .await
3319            .unwrap();
3320
3321        let step2_run2 = store
3322            .create_step(NewStep {
3323                run_id: run2.id,
3324                trace_id: step_trace_id(run2.id, "step2", 1),
3325                name: "step2".to_string(),
3326                kind: crate::entities::StepKind::Shell,
3327                position: 1,
3328                input: None,
3329                is_error_handler: false,
3330            })
3331            .await
3332            .unwrap();
3333
3334        store
3335            .create_step_dependencies(vec![
3336                NewStepDependency {
3337                    step_id: step2_run1.id,
3338                    depends_on: step1_run1.id,
3339                },
3340                NewStepDependency {
3341                    step_id: step2_run2.id,
3342                    depends_on: step1_run2.id,
3343                },
3344            ])
3345            .await
3346            .unwrap();
3347
3348        let deps_run1 = store.list_step_dependencies(run1.id).await.unwrap();
3349        let deps_run2 = store.list_step_dependencies(run2.id).await.unwrap();
3350
3351        assert_eq!(deps_run1.len(), 1);
3352        assert_eq!(deps_run1[0].step_id, step2_run1.id);
3353        assert_eq!(deps_run1[0].depends_on, step1_run1.id);
3354
3355        assert_eq!(deps_run2.len(), 1);
3356        assert_eq!(deps_run2[0].step_id, step2_run2.id);
3357        assert_eq!(deps_run2[0].depends_on, step1_run2.id);
3358    }
3359
3360    #[tokio::test]
3361    async fn list_step_dependencies_returns_empty_for_nonexistent_run() {
3362        let store = InMemoryStore::new();
3363        let deps = store.list_step_dependencies(Uuid::nil()).await.unwrap();
3364        assert!(deps.is_empty());
3365    }
3366
3367    #[tokio::test]
3368    async fn list_step_dependencies_sorted_by_created_at() {
3369        let store = InMemoryStore::new();
3370        let run = store
3371            .create_run(new_run_req("test"))
3372            .await
3373            .unwrap()
3374            .into_run();
3375
3376        let step1 = store
3377            .create_step(NewStep {
3378                run_id: run.id,
3379                trace_id: step_trace_id(run.id, "step1", 0),
3380                name: "step1".to_string(),
3381                kind: crate::entities::StepKind::Shell,
3382                position: 0,
3383                input: None,
3384                is_error_handler: false,
3385            })
3386            .await
3387            .unwrap();
3388
3389        let step2 = store
3390            .create_step(NewStep {
3391                run_id: run.id,
3392                trace_id: step_trace_id(run.id, "step2", 1),
3393                name: "step2".to_string(),
3394                kind: crate::entities::StepKind::Shell,
3395                position: 1,
3396                input: None,
3397                is_error_handler: false,
3398            })
3399            .await
3400            .unwrap();
3401
3402        let step3 = store
3403            .create_step(NewStep {
3404                run_id: run.id,
3405                trace_id: step_trace_id(run.id, "step3", 2),
3406                name: "step3".to_string(),
3407                kind: crate::entities::StepKind::Shell,
3408                position: 2,
3409                input: None,
3410                is_error_handler: false,
3411            })
3412            .await
3413            .unwrap();
3414
3415        store
3416            .create_step_dependencies(vec![NewStepDependency {
3417                step_id: step2.id,
3418                depends_on: step1.id,
3419            }])
3420            .await
3421            .unwrap();
3422
3423        store
3424            .create_step_dependencies(vec![NewStepDependency {
3425                step_id: step3.id,
3426                depends_on: step1.id,
3427            }])
3428            .await
3429            .unwrap();
3430
3431        let deps = store.list_step_dependencies(run.id).await.unwrap();
3432        assert_eq!(deps.len(), 2);
3433        assert!(deps[0].created_at <= deps[1].created_at);
3434    }
3435
3436    // ---- update_run_returning ----
3437
3438    #[tokio::test]
3439    async fn update_run_returning_applies_and_returns() {
3440        let store = InMemoryStore::new();
3441        let run = store
3442            .create_run(new_run_req("test"))
3443            .await
3444            .unwrap()
3445            .into_run();
3446
3447        // Transition Pending -> Running first
3448        store
3449            .update_run_status(run.id, RunStatus::Running)
3450            .await
3451            .unwrap();
3452
3453        let updated = store
3454            .update_run_returning(
3455                run.id,
3456                RunUpdate {
3457                    status: Some(RunStatus::Completed),
3458                    cost_usd: Some(Decimal::new(4200, 2)),
3459                    duration_ms: Some(1500),
3460                    ..RunUpdate::default()
3461                },
3462            )
3463            .await
3464            .unwrap();
3465
3466        assert_eq!(updated.id, run.id);
3467        assert_eq!(updated.status.state, RunStatus::Completed);
3468        assert_eq!(updated.cost_usd, Decimal::new(4200, 2));
3469        assert_eq!(updated.duration_ms, 1500);
3470        assert!(updated.completed_at.is_some());
3471    }
3472
3473    #[tokio::test]
3474    async fn update_run_returning_not_found() {
3475        let store = InMemoryStore::new();
3476        let result = store
3477            .update_run_returning(
3478                Uuid::nil(),
3479                RunUpdate {
3480                    status: Some(RunStatus::Running),
3481                    ..RunUpdate::default()
3482                },
3483            )
3484            .await;
3485
3486        assert!(matches!(result, Err(StoreError::RunNotFound(_))));
3487    }
3488
3489    #[tokio::test]
3490    async fn update_run_returning_invalid_transition() {
3491        let store = InMemoryStore::new();
3492        let run = store
3493            .create_run(new_run_req("test"))
3494            .await
3495            .unwrap()
3496            .into_run();
3497
3498        let result = store
3499            .update_run_returning(
3500                run.id,
3501                RunUpdate {
3502                    status: Some(RunStatus::Completed),
3503                    ..RunUpdate::default()
3504                },
3505            )
3506            .await;
3507
3508        assert!(matches!(result, Err(StoreError::InvalidTransition { .. })));
3509    }
3510
3511    // ---- retry scheduling ----
3512
3513    #[tokio::test]
3514    async fn create_step_stamps_the_current_attempt() {
3515        let store = InMemoryStore::new();
3516        let run = store
3517            .create_run(new_run_req("retry-wf"))
3518            .await
3519            .unwrap()
3520            .into_run();
3521
3522        let first = store
3523            .create_step(new_step_req(run.id, "build", 0))
3524            .await
3525            .unwrap();
3526        assert_eq!(first.attempt, 1);
3527
3528        store
3529            .update_run_status(run.id, RunStatus::Running)
3530            .await
3531            .unwrap();
3532        store
3533            .update_run(
3534                run.id,
3535                RunUpdate {
3536                    status: Some(RunStatus::Retrying),
3537                    increment_retry: true,
3538                    ..RunUpdate::default()
3539                },
3540            )
3541            .await
3542            .unwrap();
3543
3544        let second = store
3545            .create_step(new_step_req(run.id, "build", 0))
3546            .await
3547            .unwrap();
3548        assert_eq!(second.attempt, 2);
3549    }
3550
3551    #[tokio::test]
3552    async fn pick_next_pending_ignores_retrying_run_before_its_backoff() {
3553        let store = InMemoryStore::new();
3554        let run = store
3555            .create_run(new_run_req("retry-wf"))
3556            .await
3557            .unwrap()
3558            .into_run();
3559
3560        store
3561            .update_run_status(run.id, RunStatus::Running)
3562            .await
3563            .unwrap();
3564        store
3565            .update_run(
3566                run.id,
3567                RunUpdate {
3568                    status: Some(RunStatus::Retrying),
3569                    increment_retry: true,
3570                    scheduled_at: Some(Utc::now() + TimeDelta::seconds(60)),
3571                    ..RunUpdate::default()
3572                },
3573            )
3574            .await
3575            .unwrap();
3576
3577        assert!(store.pick_next_pending(None).await.unwrap().is_none());
3578    }
3579
3580    #[tokio::test]
3581    async fn pick_next_pending_resumes_retrying_run_after_its_backoff() {
3582        let store = InMemoryStore::new();
3583        let run = store
3584            .create_run(new_run_req("retry-wf"))
3585            .await
3586            .unwrap()
3587            .into_run();
3588
3589        store
3590            .update_run_status(run.id, RunStatus::Running)
3591            .await
3592            .unwrap();
3593        store
3594            .update_run(
3595                run.id,
3596                RunUpdate {
3597                    status: Some(RunStatus::Retrying),
3598                    increment_retry: true,
3599                    scheduled_at: Some(Utc::now() - TimeDelta::seconds(1)),
3600                    ..RunUpdate::default()
3601                },
3602            )
3603            .await
3604            .unwrap();
3605
3606        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
3607        assert_eq!(picked.id, run.id);
3608        assert_eq!(picked.status.state, RunStatus::Running);
3609        assert_eq!(picked.retry_count, 1);
3610    }
3611
3612    // ---- priority ----
3613
3614    async fn create_with_priority(store: &InMemoryStore, name: &str, priority: i16) -> Run {
3615        let run = store
3616            .create_run(NewRun {
3617                priority,
3618                ..new_run_req(name)
3619            })
3620            .await
3621            .unwrap()
3622            .into_run();
3623        // Distinct created_at values, so FIFO among equal priorities is observable.
3624        sleep(Duration::from_millis(2)).await;
3625        run
3626    }
3627
3628    #[tokio::test]
3629    async fn pick_next_pending_priority_serves_higher_priority_first() {
3630        let store = InMemoryStore::new();
3631        let low = create_with_priority(&store, "low", 0).await;
3632        let high = create_with_priority(&store, "high", 10).await;
3633
3634        let first = store.pick_next_pending(None).await.unwrap().unwrap();
3635        assert_eq!(first.id, high.id);
3636        assert_eq!(first.priority, 10);
3637        let second = store.pick_next_pending(None).await.unwrap().unwrap();
3638        assert_eq!(second.id, low.id);
3639    }
3640
3641    #[tokio::test]
3642    async fn pick_next_pending_priority_is_fifo_among_equal_priorities() {
3643        let store = InMemoryStore::new();
3644        let older = create_with_priority(&store, "older", 5).await;
3645        let younger = create_with_priority(&store, "younger", 5).await;
3646
3647        let first = store.pick_next_pending(None).await.unwrap().unwrap();
3648        assert_eq!(first.id, older.id);
3649        let second = store.pick_next_pending(None).await.unwrap().unwrap();
3650        assert_eq!(second.id, younger.id);
3651    }
3652
3653    #[tokio::test]
3654    async fn pick_next_pending_priority_negative_runs_after_default() {
3655        let store = InMemoryStore::new();
3656        let negative = create_with_priority(&store, "background", -50).await;
3657        let default = create_with_priority(&store, "default", 0).await;
3658
3659        let first = store.pick_next_pending(None).await.unwrap().unwrap();
3660        assert_eq!(first.id, default.id);
3661        let second = store.pick_next_pending(None).await.unwrap().unwrap();
3662        assert_eq!(second.id, negative.id);
3663    }
3664
3665    #[tokio::test]
3666    async fn pick_next_pending_priority_skips_high_priority_run_not_yet_due() {
3667        let store = InMemoryStore::new();
3668        let later = store
3669            .create_run(NewRun {
3670                priority: 100,
3671                scheduled_at: Some(Utc::now() + TimeDelta::seconds(3600)),
3672                ..new_run_req("later")
3673            })
3674            .await
3675            .unwrap()
3676            .into_run();
3677        let now = create_with_priority(&store, "now", 0).await;
3678
3679        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
3680        assert_eq!(picked.id, now.id);
3681        assert!(store.pick_next_pending(None).await.unwrap().is_none());
3682        let later = store.get_run(later.id).await.unwrap().unwrap();
3683        assert_eq!(later.status.state, RunStatus::Pending);
3684    }
3685
3686    #[tokio::test]
3687    async fn pick_next_pending_priority_kept_after_retry() {
3688        let store = InMemoryStore::new();
3689        let run = create_with_priority(&store, "retry-wf", 42).await;
3690
3691        store
3692            .update_run_status(run.id, RunStatus::Running)
3693            .await
3694            .unwrap();
3695        store
3696            .update_run(
3697                run.id,
3698                RunUpdate {
3699                    status: Some(RunStatus::Retrying),
3700                    increment_retry: true,
3701                    scheduled_at: Some(Utc::now() - TimeDelta::seconds(1)),
3702                    ..RunUpdate::default()
3703                },
3704            )
3705            .await
3706            .unwrap();
3707        let fresh = create_with_priority(&store, "fresh", 0).await;
3708
3709        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
3710        assert_eq!(picked.id, run.id);
3711        assert_eq!(picked.priority, 42);
3712        assert_eq!(picked.retry_count, 1);
3713        let next = store.pick_next_pending(None).await.unwrap().unwrap();
3714        assert_eq!(next.id, fresh.id);
3715    }
3716
3717    #[tokio::test]
3718    async fn create_run_priority_out_of_range_is_rejected() {
3719        let store = InMemoryStore::new();
3720        for priority in [101, -101] {
3721            let err = store
3722                .create_run(NewRun {
3723                    priority,
3724                    ..new_run_req("out-of-range")
3725                })
3726                .await
3727                .unwrap_err();
3728            assert!(matches!(err, StoreError::Database(_)), "{err:?}");
3729        }
3730        let page = store.list_runs(RunFilter::default(), 1, 10).await.unwrap();
3731        assert_eq!(page.total, 0);
3732    }
3733
3734    #[tokio::test]
3735    async fn list_runs_priority_filter_is_exact_match() {
3736        let store = InMemoryStore::new();
3737        let urgent = create_with_priority(&store, "urgent", 10).await;
3738        create_with_priority(&store, "default", 0).await;
3739        create_with_priority(&store, "more-urgent", 20).await;
3740
3741        let page = store
3742            .list_runs(
3743                RunFilter {
3744                    priority: Some(10),
3745                    ..RunFilter::default()
3746                },
3747                1,
3748                10,
3749            )
3750            .await
3751            .unwrap();
3752        assert_eq!(page.total, 1);
3753        assert_eq!(page.items[0].id, urgent.id);
3754
3755        let all = store.list_runs(RunFilter::default(), 1, 10).await.unwrap();
3756        assert_eq!(all.total, 3);
3757    }
3758
3759    #[tokio::test]
3760    async fn update_run_persists_scheduled_at() {
3761        let store = InMemoryStore::new();
3762        let run = store
3763            .create_run(new_run_req("test"))
3764            .await
3765            .unwrap()
3766            .into_run();
3767        let when = Utc::now() + TimeDelta::seconds(30);
3768
3769        store
3770            .update_run(
3771                run.id,
3772                RunUpdate {
3773                    scheduled_at: Some(when),
3774                    ..RunUpdate::default()
3775                },
3776            )
3777            .await
3778            .unwrap();
3779
3780        let fetched = store.get_run(run.id).await.unwrap().unwrap();
3781        assert_eq!(fetched.scheduled_at, Some(when));
3782    }
3783
3784    // ---- output ----
3785
3786    #[tokio::test]
3787    async fn new_run_has_no_output() {
3788        let store = InMemoryStore::new();
3789        let run = store
3790            .create_run(new_run_req("test"))
3791            .await
3792            .unwrap()
3793            .into_run();
3794
3795        assert!(run.output.is_none());
3796        let fetched = store.get_run(run.id).await.unwrap().unwrap();
3797        assert!(fetched.output.is_none());
3798    }
3799
3800    #[tokio::test]
3801    async fn update_run_sets_output() {
3802        let store = InMemoryStore::new();
3803        let run = store
3804            .create_run(new_run_req("test"))
3805            .await
3806            .unwrap()
3807            .into_run();
3808
3809        store
3810            .update_run(
3811                run.id,
3812                RunUpdate {
3813                    output: Some(json!({"verdict": "approved"})),
3814                    ..RunUpdate::default()
3815                },
3816            )
3817            .await
3818            .unwrap();
3819
3820        let fetched = store.get_run(run.id).await.unwrap().unwrap();
3821        assert_eq!(fetched.output, Some(json!({"verdict": "approved"})));
3822    }
3823
3824    #[tokio::test]
3825    async fn update_run_without_output_keeps_previous_output() {
3826        let store = InMemoryStore::new();
3827        let run = store
3828            .create_run(new_run_req("test"))
3829            .await
3830            .unwrap()
3831            .into_run();
3832
3833        store
3834            .update_run(
3835                run.id,
3836                RunUpdate {
3837                    output: Some(json!({"verdict": "approved"})),
3838                    ..RunUpdate::default()
3839                },
3840            )
3841            .await
3842            .unwrap();
3843        store
3844            .update_run(
3845                run.id,
3846                RunUpdate {
3847                    error: Some("boom".to_string()),
3848                    ..RunUpdate::default()
3849                },
3850            )
3851            .await
3852            .unwrap();
3853
3854        let fetched = store.get_run(run.id).await.unwrap().unwrap();
3855        assert_eq!(fetched.output, Some(json!({"verdict": "approved"})));
3856        assert_eq!(fetched.error.as_deref(), Some("boom"));
3857    }
3858
3859    #[tokio::test]
3860    async fn update_run_output_last_write_wins() {
3861        let store = InMemoryStore::new();
3862        let run = store
3863            .create_run(new_run_req("test"))
3864            .await
3865            .unwrap()
3866            .into_run();
3867
3868        for verdict in ["first", "second"] {
3869            store
3870                .update_run(
3871                    run.id,
3872                    RunUpdate {
3873                        output: Some(json!({ "verdict": verdict })),
3874                        ..RunUpdate::default()
3875                    },
3876                )
3877                .await
3878                .unwrap();
3879        }
3880
3881        let fetched = store.get_run(run.id).await.unwrap().unwrap();
3882        assert_eq!(fetched.output, Some(json!({"verdict": "second"})));
3883    }
3884
3885    // ---- created_by ----
3886
3887    async fn seed_user(store: &InMemoryStore, username: &str) -> Uuid {
3888        store
3889            .create_user(NewUser {
3890                email: format!("{username}@example.com"),
3891                username: username.to_string(),
3892                password_hash: "hash".to_string(),
3893                is_admin: Some(false),
3894            })
3895            .await
3896            .unwrap()
3897            .id
3898    }
3899
3900    async fn seed_api_key(store: &InMemoryStore, user_id: Uuid, name: &str) -> Uuid {
3901        store
3902            .create_api_key(NewApiKey {
3903                user_id,
3904                name: name.to_string(),
3905                key_hash: "hash".to_string(),
3906                key_prefix: "irfl_0000".to_string(),
3907                scopes: vec![ApiKeyScope::RunsWrite],
3908                expires_at: None,
3909                rate_limit_override: None,
3910            })
3911            .await
3912            .unwrap()
3913            .id
3914    }
3915
3916    fn run_req_by(actor: RunActor) -> NewRun {
3917        NewRun {
3918            created_by: Some(actor),
3919            ..new_run_req("test")
3920        }
3921    }
3922
3923    #[tokio::test]
3924    async fn create_run_without_actor_has_no_author() {
3925        let store = InMemoryStore::new();
3926        let run = store
3927            .create_run(new_run_req("test"))
3928            .await
3929            .unwrap()
3930            .into_run();
3931
3932        assert!(run.created_by.is_none());
3933        assert!(run.created_by_label.is_none());
3934    }
3935
3936    #[tokio::test]
3937    async fn create_run_by_user_resolves_username_as_label() {
3938        let store = InMemoryStore::new();
3939        let user_id = seed_user(&store, "alice").await;
3940
3941        let run = store
3942            .create_run(run_req_by(RunActor::User { user_id }))
3943            .await
3944            .unwrap()
3945            .into_run();
3946
3947        assert_eq!(run.created_by, Some(RunActor::User { user_id }));
3948        assert_eq!(run.created_by_label.as_deref(), Some("alice"));
3949    }
3950
3951    #[tokio::test]
3952    async fn create_run_by_api_key_resolves_key_and_owner_as_label() {
3953        let store = InMemoryStore::new();
3954        let user_id = seed_user(&store, "alice").await;
3955        let api_key_id = seed_api_key(&store, user_id, "ci-deploy").await;
3956
3957        let run = store
3958            .create_run(run_req_by(RunActor::ApiKey {
3959                api_key_id,
3960                user_id,
3961            }))
3962            .await
3963            .unwrap()
3964            .into_run();
3965
3966        assert_eq!(run.created_by_label.as_deref(), Some("ci-deploy (alice)"));
3967    }
3968
3969    #[tokio::test]
3970    async fn label_follows_api_key_rename() {
3971        let store = InMemoryStore::new();
3972        let user_id = seed_user(&store, "alice").await;
3973        let api_key_id = seed_api_key(&store, user_id, "ci-deploy").await;
3974        let run = store
3975            .create_run(run_req_by(RunActor::ApiKey {
3976                api_key_id,
3977                user_id,
3978            }))
3979            .await
3980            .unwrap()
3981            .into_run();
3982
3983        store
3984            .update_api_key(
3985                api_key_id,
3986                ApiKeyUpdate {
3987                    name: Some("ci-release".to_string()),
3988                    ..ApiKeyUpdate::default()
3989                },
3990            )
3991            .await
3992            .unwrap();
3993
3994        let reread = store.get_run(run.id).await.unwrap().unwrap();
3995        assert_eq!(
3996            reread.created_by_label.as_deref(),
3997            Some("ci-release (alice)")
3998        );
3999    }
4000
4001    #[tokio::test]
4002    async fn label_is_none_when_user_is_unknown() {
4003        let store = InMemoryStore::new();
4004        let run = store
4005            .create_run(run_req_by(RunActor::User {
4006                user_id: Uuid::now_v7(),
4007            }))
4008            .await
4009            .unwrap()
4010            .into_run();
4011
4012        assert!(run.created_by.is_some());
4013        assert!(run.created_by_label.is_none());
4014    }
4015
4016    #[tokio::test]
4017    async fn label_is_key_name_only_when_owner_is_unknown() {
4018        let store = InMemoryStore::new();
4019        let owner = seed_user(&store, "alice").await;
4020        let api_key_id = seed_api_key(&store, owner, "ci-deploy").await;
4021
4022        // An actor pointing at a user that was never created.
4023        let run = store
4024            .create_run(run_req_by(RunActor::ApiKey {
4025                api_key_id,
4026                user_id: Uuid::now_v7(),
4027            }))
4028            .await
4029            .unwrap()
4030            .into_run();
4031
4032        assert_eq!(run.created_by_label.as_deref(), Some("ci-deploy"));
4033    }
4034
4035    #[tokio::test]
4036    async fn list_runs_filters_by_author() {
4037        let store = InMemoryStore::new();
4038        let alice = seed_user(&store, "alice").await;
4039        let bob = seed_user(&store, "bob").await;
4040
4041        store
4042            .create_run(run_req_by(RunActor::User { user_id: alice }))
4043            .await
4044            .unwrap()
4045            .into_run();
4046        store
4047            .create_run(run_req_by(RunActor::User { user_id: bob }))
4048            .await
4049            .unwrap()
4050            .into_run();
4051        store.create_run(new_run_req("anonymous")).await.unwrap();
4052
4053        let page = store
4054            .list_runs(
4055                RunFilter {
4056                    created_by_user_id: Some(alice),
4057                    ..RunFilter::default()
4058                },
4059                1,
4060                20,
4061            )
4062            .await
4063            .unwrap();
4064
4065        assert_eq!(page.total, 1);
4066        assert_eq!(page.items[0].created_by_label.as_deref(), Some("alice"));
4067    }
4068
4069    #[tokio::test]
4070    async fn list_runs_author_filter_matches_runs_from_the_users_api_keys() {
4071        let store = InMemoryStore::new();
4072        let alice = seed_user(&store, "alice").await;
4073        let api_key_id = seed_api_key(&store, alice, "ci-deploy").await;
4074
4075        store
4076            .create_run(run_req_by(RunActor::ApiKey {
4077                api_key_id,
4078                user_id: alice,
4079            }))
4080            .await
4081            .unwrap()
4082            .into_run();
4083
4084        let page = store
4085            .list_runs(
4086                RunFilter {
4087                    created_by_user_id: Some(alice),
4088                    ..RunFilter::default()
4089                },
4090                1,
4091                20,
4092            )
4093            .await
4094            .unwrap();
4095
4096        assert_eq!(page.total, 1);
4097    }
4098
4099    #[tokio::test]
4100    async fn list_runs_author_filter_excludes_unrelated_users() {
4101        let store = InMemoryStore::new();
4102        let alice = seed_user(&store, "alice").await;
4103
4104        store
4105            .create_run(run_req_by(RunActor::User { user_id: alice }))
4106            .await
4107            .unwrap()
4108            .into_run();
4109
4110        let page = store
4111            .list_runs(
4112                RunFilter {
4113                    created_by_user_id: Some(Uuid::now_v7()),
4114                    ..RunFilter::default()
4115                },
4116                1,
4117                20,
4118            )
4119            .await
4120            .unwrap();
4121
4122        assert_eq!(page.total, 0);
4123    }
4124
4125    #[tokio::test]
4126    async fn list_runs_without_author_filter_returns_every_run() {
4127        let store = InMemoryStore::new();
4128        let alice = seed_user(&store, "alice").await;
4129
4130        store
4131            .create_run(run_req_by(RunActor::User { user_id: alice }))
4132            .await
4133            .unwrap()
4134            .into_run();
4135        store.create_run(new_run_req("anonymous")).await.unwrap();
4136
4137        let page = store.list_runs(RunFilter::default(), 1, 20).await.unwrap();
4138        assert_eq!(page.total, 2);
4139    }
4140
4141    #[tokio::test]
4142    async fn pick_next_pending_resolves_author_label() {
4143        let store = InMemoryStore::new();
4144        let user_id = seed_user(&store, "alice").await;
4145        store
4146            .create_run(run_req_by(RunActor::User { user_id }))
4147            .await
4148            .unwrap()
4149            .into_run();
4150
4151        let picked = store.pick_next_pending(None).await.unwrap().unwrap();
4152        assert_eq!(picked.created_by_label.as_deref(), Some("alice"));
4153    }
4154
4155    // ---- list_purgeable_runs ----
4156
4157    #[tokio::test]
4158    async fn list_purgeable_runs_returns_old_terminal_runs() {
4159        let store = InMemoryStore::new();
4160        let old = create_terminal_run(&store, "deploy", RunStatus::Completed).await;
4161        store
4162            .set_run_created_at(old.id, Utc::now() - chrono::Duration::days(100))
4163            .await;
4164
4165        let policy = PurgePolicy {
4166            max_age_days: 90,
4167            max_runs_per_workflow: 10000,
4168            dry_run: false,
4169        };
4170        let result = store.list_purgeable_runs(&policy, 100).await.unwrap();
4171
4172        assert_eq!(result.len(), 1);
4173        assert_eq!(result[0].run_id, old.id);
4174        assert_eq!(result[0].reason, PurgeReason::TooOld);
4175    }
4176
4177    #[tokio::test]
4178    async fn list_purgeable_runs_ignores_non_terminal_states() {
4179        let store = InMemoryStore::new();
4180
4181        // Pending
4182        let pending = store
4183            .create_run(new_run_req("deploy"))
4184            .await
4185            .unwrap()
4186            .into_run();
4187        store
4188            .set_run_created_at(pending.id, Utc::now() - chrono::Duration::days(200))
4189            .await;
4190
4191        // Running
4192        let running = store
4193            .create_run(new_run_req("deploy"))
4194            .await
4195            .unwrap()
4196            .into_run();
4197        store
4198            .update_run_status(running.id, RunStatus::Running)
4199            .await
4200            .unwrap();
4201        store
4202            .set_run_created_at(running.id, Utc::now() - chrono::Duration::days(200))
4203            .await;
4204
4205        let policy = PurgePolicy {
4206            max_age_days: 90,
4207            max_runs_per_workflow: 1,
4208            dry_run: false,
4209        };
4210        let result = store.list_purgeable_runs(&policy, 100).await.unwrap();
4211        assert!(result.is_empty());
4212    }
4213
4214    #[tokio::test]
4215    async fn list_purgeable_runs_returns_excess_per_workflow() {
4216        let store = InMemoryStore::new();
4217        let r1 = create_terminal_run(&store, "deploy", RunStatus::Completed).await;
4218        store
4219            .set_run_created_at(r1.id, Utc::now() - chrono::Duration::days(10))
4220            .await;
4221        let r2 = create_terminal_run(&store, "deploy", RunStatus::Completed).await;
4222        store
4223            .set_run_created_at(r2.id, Utc::now() - chrono::Duration::days(5))
4224            .await;
4225        let _r3 = create_terminal_run(&store, "deploy", RunStatus::Completed).await;
4226
4227        let policy = PurgePolicy {
4228            max_age_days: 365,
4229            max_runs_per_workflow: 2,
4230            dry_run: false,
4231        };
4232        let result = store.list_purgeable_runs(&policy, 100).await.unwrap();
4233
4234        assert_eq!(result.len(), 1);
4235        assert_eq!(result[0].run_id, r1.id);
4236        assert_eq!(result[0].reason, PurgeReason::ExceedsWorkflowLimit);
4237    }
4238
4239    // ---- delete_run ----
4240
4241    #[tokio::test]
4242    async fn delete_run_removes_run_and_associated_data() {
4243        use crate::artifact_store::ArtifactStore;
4244        use crate::entities::{NewStep, StepKind};
4245
4246        let store = InMemoryStore::new();
4247        let run = create_terminal_run(&store, "deploy", RunStatus::Completed).await;
4248        let step = store
4249            .create_step(NewStep {
4250                run_id: run.id,
4251                trace_id: step_trace_id(run.id, "build", 0),
4252                name: "build".to_string(),
4253                kind: StepKind::Shell,
4254                position: 0,
4255                input: None,
4256                is_error_handler: false,
4257            })
4258            .await
4259            .unwrap();
4260
4261        let artifact_id = Uuid::now_v7();
4262        store
4263            .create_artifact(crate::entities::NewArtifact {
4264                id: artifact_id,
4265                run_id: run.id,
4266                step_id: step.id,
4267                name: "report.html".to_string(),
4268                storage_key: format!("artifacts/{}/{}/{}", run.id, step.id, artifact_id),
4269                content_type: "text/html".to_string(),
4270                size_bytes: 42,
4271                sha256: "0".repeat(64),
4272            })
4273            .await
4274            .unwrap();
4275
4276        let keys = store.delete_run(run.id).await.unwrap();
4277
4278        assert_eq!(keys.len(), 1);
4279        assert!(keys[0].contains(&artifact_id.to_string()));
4280        assert!(store.get_run(run.id).await.unwrap().is_none());
4281        assert!(store.list_steps(run.id).await.unwrap().is_empty());
4282        assert!(
4283            store
4284                .list_artifacts_for_run(run.id)
4285                .await
4286                .unwrap()
4287                .is_empty()
4288        );
4289    }
4290
4291    #[tokio::test]
4292    async fn delete_run_not_found() {
4293        let store = InMemoryStore::new();
4294        let err = store.delete_run(Uuid::now_v7()).await.unwrap_err();
4295        assert!(matches!(err, StoreError::RunNotFound(_)));
4296    }
4297
4298    // ---- record_step_approval ----
4299
4300    fn vote(user_id: Uuid, name: &str) -> StepApproval {
4301        StepApproval {
4302            user_id,
4303            approved_by: name.to_string(),
4304            at: Utc::now(),
4305        }
4306    }
4307
4308    #[tokio::test]
4309    async fn record_step_approval_appends_distinct_voters() {
4310        let store = InMemoryStore::new();
4311        let run = store
4312            .create_run(new_run_req("test"))
4313            .await
4314            .unwrap()
4315            .into_run();
4316        let step = store
4317            .create_step(new_step_req(run.id, "gate", 0))
4318            .await
4319            .unwrap();
4320        assert!(step.approvals.is_empty());
4321        assert!(step.approval_requirement.is_none());
4322
4323        let alice = Uuid::now_v7();
4324        let bob = Uuid::now_v7();
4325        let after_first = store
4326            .record_step_approval(step.id, vote(alice, "alice"))
4327            .await
4328            .unwrap();
4329        assert_eq!(after_first.approvals.len(), 1);
4330
4331        let after_second = store
4332            .record_step_approval(step.id, vote(bob, "bob"))
4333            .await
4334            .unwrap();
4335        assert_eq!(after_second.approvals.len(), 2);
4336        assert_eq!(after_second.approvals[0].user_id, alice);
4337        assert_eq!(after_second.approvals[1].user_id, bob);
4338    }
4339
4340    #[tokio::test]
4341    async fn record_step_approval_ignores_same_user() {
4342        let store = InMemoryStore::new();
4343        let run = store
4344            .create_run(new_run_req("test"))
4345            .await
4346            .unwrap()
4347            .into_run();
4348        let step = store
4349            .create_step(new_step_req(run.id, "gate", 0))
4350            .await
4351            .unwrap();
4352
4353        let alice = Uuid::now_v7();
4354        store
4355            .record_step_approval(step.id, vote(alice, "alice"))
4356            .await
4357            .unwrap();
4358        let again = store
4359            .record_step_approval(step.id, vote(alice, "alice-key"))
4360            .await
4361            .unwrap();
4362
4363        assert_eq!(again.approvals.len(), 1);
4364        assert_eq!(again.approvals[0].approved_by, "alice");
4365    }
4366
4367    #[tokio::test]
4368    async fn record_step_approval_unknown_step_is_not_found() {
4369        let store = InMemoryStore::new();
4370        let err = store
4371            .record_step_approval(Uuid::now_v7(), vote(Uuid::now_v7(), "alice"))
4372            .await
4373            .unwrap_err();
4374        assert!(matches!(err, StoreError::StepNotFound(_)));
4375    }
4376
4377    #[tokio::test]
4378    async fn update_step_sets_approval_requirement() {
4379        let store = InMemoryStore::new();
4380        let run = store
4381            .create_run(new_run_req("test"))
4382            .await
4383            .unwrap()
4384            .into_run();
4385        let step = store
4386            .create_step(new_step_req(run.id, "gate", 0))
4387            .await
4388            .unwrap();
4389        let requirement = ApprovalRequirement {
4390            required_approvers: 3,
4391            ..ApprovalRequirement::default()
4392        };
4393
4394        store
4395            .update_step(
4396                step.id,
4397                StepUpdate {
4398                    approval_requirement: Some(requirement.clone()),
4399                    ..StepUpdate::default()
4400                },
4401            )
4402            .await
4403            .unwrap();
4404
4405        let fetched = store.get_step(step.id).await.unwrap().unwrap();
4406        assert_eq!(fetched.approval_requirement, Some(requirement));
4407    }
4408
4409    async fn sleeping_run(store: &InMemoryStore, scheduled_at: DateTime<Utc>) -> Run {
4410        let run = store
4411            .create_run(new_run_req("sleepy"))
4412            .await
4413            .unwrap()
4414            .into_run();
4415        store
4416            .update_run_status(run.id, RunStatus::Running)
4417            .await
4418            .unwrap();
4419        store
4420            .update_run(
4421                run.id,
4422                RunUpdate {
4423                    status: Some(RunStatus::Sleeping),
4424                    scheduled_at: Some(scheduled_at),
4425                    ..RunUpdate::default()
4426                },
4427            )
4428            .await
4429            .unwrap();
4430        store.get_run(run.id).await.unwrap().unwrap()
4431    }
4432
4433    #[tokio::test]
4434    async fn claim_due_sleeping_runs_requeues_due_runs() {
4435        let store = InMemoryStore::new();
4436        let due = sleeping_run(&store, Utc::now() - TimeDelta::seconds(5)).await;
4437
4438        let woken = store.claim_due_sleeping_runs(10).await.unwrap();
4439        assert_eq!(woken.len(), 1);
4440        assert_eq!(woken[0].id, due.id);
4441        assert_eq!(woken[0].status.state, RunStatus::Pending);
4442        assert!(woken[0].scheduled_at.is_none());
4443
4444        let fetched = store.get_run(due.id).await.unwrap().unwrap();
4445        assert_eq!(fetched.status.state, RunStatus::Pending);
4446        assert!(fetched.scheduled_at.is_none());
4447
4448        // Exactly once: a second tick finds nothing.
4449        assert!(store.claim_due_sleeping_runs(10).await.unwrap().is_empty());
4450    }
4451
4452    #[tokio::test]
4453    async fn claim_due_sleeping_runs_skips_future_runs() {
4454        let store = InMemoryStore::new();
4455        let future = sleeping_run(&store, Utc::now() + TimeDelta::hours(1)).await;
4456
4457        assert!(store.claim_due_sleeping_runs(10).await.unwrap().is_empty());
4458        let fetched = store.get_run(future.id).await.unwrap().unwrap();
4459        assert_eq!(fetched.status.state, RunStatus::Sleeping);
4460    }
4461
4462    #[tokio::test]
4463    async fn claim_due_sleeping_runs_honours_limit_oldest_first() {
4464        let store = InMemoryStore::new();
4465        let older = sleeping_run(&store, Utc::now() - TimeDelta::seconds(20)).await;
4466        let newer = sleeping_run(&store, Utc::now() - TimeDelta::seconds(10)).await;
4467
4468        let woken = store.claim_due_sleeping_runs(1).await.unwrap();
4469        assert_eq!(woken.len(), 1);
4470        assert_eq!(woken[0].id, older.id);
4471
4472        let woken = store.claim_due_sleeping_runs(1).await.unwrap();
4473        assert_eq!(woken[0].id, newer.id);
4474    }
4475
4476    // ---- approval assignee binding ----
4477
4478    async fn open_gate(store: &InMemoryStore, assignee: Assignee) -> Uuid {
4479        let run = store
4480            .create_run(new_run_req("gated"))
4481            .await
4482            .unwrap()
4483            .into_run();
4484        let step = store
4485            .create_step(new_step_req(run.id, "approve", 0))
4486            .await
4487            .unwrap();
4488        store
4489            .update_step(
4490                step.id,
4491                StepUpdate {
4492                    approval_assignee: Some(assignee),
4493                    ..StepUpdate::default()
4494                },
4495            )
4496            .await
4497            .unwrap();
4498        step.id
4499    }
4500
4501    #[tokio::test]
4502    async fn gate_assignee_is_bound_when_the_gate_opens() {
4503        let store = InMemoryStore::new();
4504        let alice = seed_user(&store, "alice").await;
4505
4506        let step_id = open_gate(&store, Assignee::user("alice")).await;
4507
4508        let step = store.get_step(step_id).await.unwrap().unwrap();
4509        assert_eq!(step.approval_assignee_user_id, Some(alice));
4510    }
4511
4512    #[tokio::test]
4513    async fn account_created_after_the_gate_opened_is_not_bound() {
4514        let store = InMemoryStore::new();
4515
4516        let step_id = open_gate(&store, Assignee::user("alice")).await;
4517        seed_user(&store, "alice").await;
4518
4519        let step = store.get_step(step_id).await.unwrap().unwrap();
4520        assert_eq!(step.approval_assignee, Some(Assignee::user("alice")));
4521        assert!(step.approval_assignee_user_id.is_none());
4522    }
4523
4524    #[tokio::test]
4525    async fn reassigning_to_a_group_clears_the_bound_user() {
4526        let store = InMemoryStore::new();
4527        seed_user(&store, "alice").await;
4528        let step_id = open_gate(&store, Assignee::user("alice")).await;
4529
4530        store
4531            .update_step(
4532                step_id,
4533                StepUpdate {
4534                    approval_assignee: Some(Assignee::group("sre-oncall")),
4535                    ..StepUpdate::default()
4536                },
4537            )
4538            .await
4539            .unwrap();
4540
4541        let step = store.get_step(step_id).await.unwrap().unwrap();
4542        assert_eq!(step.approval_assignee, Some(Assignee::group("sre-oncall")));
4543        assert!(step.approval_assignee_user_id.is_none());
4544    }
4545}