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