awaken_runtime/agent/state/
run_lifecycle.rs

1use crate::state::StateKey;
2use awaken_contract::contract::lifecycle::RunStatus;
3use serde::{Deserialize, Serialize};
4
5/// Run lifecycle state stored in the state engine.
6#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
7pub struct RunLifecycleState {
8    /// Current run id.
9    pub run_id: String,
10    /// Coarse lifecycle status.
11    pub status: RunStatus,
12    /// Reason string for the current status (set when Done or Waiting, None when Running).
13    #[serde(
14        default,
15        skip_serializing_if = "Option::is_none",
16        alias = "done_reason",
17        alias = "pause_reason"
18    )]
19    pub status_reason: Option<String>,
20    /// Last update timestamp (unix millis).
21    pub updated_at: u64,
22    /// Total steps completed.
23    pub step_count: u32,
24}
25
26/// Update for the run lifecycle state key.
27pub enum RunLifecycleUpdate {
28    Start {
29        run_id: String,
30        updated_at: u64,
31    },
32    StepCompleted {
33        updated_at: u64,
34    },
35    SetWaiting {
36        updated_at: u64,
37        pause_reason: String,
38    },
39    SetRunning {
40        updated_at: u64,
41    },
42    Done {
43        done_reason: String,
44        updated_at: u64,
45    },
46}
47
48impl RunLifecycleUpdate {
49    /// The target `RunStatus` this update will produce.
50    pub fn target_status(&self) -> RunStatus {
51        match self {
52            Self::Start { .. } | Self::StepCompleted { .. } | Self::SetRunning { .. } => {
53                RunStatus::Running
54            }
55            Self::SetWaiting { .. } => RunStatus::Waiting,
56            Self::Done { .. } => RunStatus::Done,
57        }
58    }
59}
60
61/// State key for run lifecycle tracking.
62pub struct RunLifecycle;
63
64impl StateKey for RunLifecycle {
65    const KEY: &'static str = "__runtime.run_lifecycle";
66
67    type Value = RunLifecycleState;
68    type Update = RunLifecycleUpdate;
69
70    fn apply(value: &mut Self::Value, update: Self::Update) {
71        let target_status = update.target_status();
72        if !value.status.can_transition_to(target_status) {
73            tracing::error!(
74                from = ?value.status,
75                to = ?target_status,
76                "invalid lifecycle transition — skipping update"
77            );
78            return;
79        }
80        match update {
81            RunLifecycleUpdate::Start { run_id, updated_at } => {
82                value.run_id = run_id;
83                value.status = RunStatus::Running;
84                value.status_reason = None;
85                value.updated_at = updated_at;
86                value.step_count = 0;
87            }
88            RunLifecycleUpdate::StepCompleted { updated_at } => {
89                value.step_count += 1;
90                value.updated_at = updated_at;
91            }
92            RunLifecycleUpdate::SetWaiting {
93                updated_at,
94                pause_reason,
95            } => {
96                value.status = RunStatus::Waiting;
97                value.status_reason = Some(pause_reason);
98                value.updated_at = updated_at;
99            }
100            RunLifecycleUpdate::SetRunning { updated_at } => {
101                value.status = RunStatus::Running;
102                value.status_reason = None;
103                value.updated_at = updated_at;
104            }
105            RunLifecycleUpdate::Done {
106                done_reason,
107                updated_at,
108            } => {
109                value.status = RunStatus::Done;
110                value.status_reason = Some(done_reason);
111                value.updated_at = updated_at;
112            }
113        }
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    fn now_ms() -> u64 {
122        std::time::SystemTime::now()
123            .duration_since(std::time::UNIX_EPOCH)
124            .unwrap()
125            .as_millis() as u64
126    }
127
128    #[test]
129    fn run_lifecycle_start_sets_running() {
130        let mut state = RunLifecycleState::default();
131        RunLifecycle::apply(
132            &mut state,
133            RunLifecycleUpdate::Start {
134                run_id: "r1".into(),
135                updated_at: 100,
136            },
137        );
138        assert_eq!(state.run_id, "r1");
139        assert_eq!(state.status, RunStatus::Running);
140        assert_eq!(state.step_count, 0);
141    }
142
143    #[test]
144    fn run_lifecycle_step_completed_increments() {
145        let mut state = RunLifecycleState::default();
146        RunLifecycle::apply(
147            &mut state,
148            RunLifecycleUpdate::Start {
149                run_id: "r1".into(),
150                updated_at: 100,
151            },
152        );
153        RunLifecycle::apply(
154            &mut state,
155            RunLifecycleUpdate::StepCompleted { updated_at: 200 },
156        );
157        RunLifecycle::apply(
158            &mut state,
159            RunLifecycleUpdate::StepCompleted { updated_at: 300 },
160        );
161        assert_eq!(state.step_count, 2);
162        assert_eq!(state.updated_at, 300);
163    }
164
165    #[test]
166    fn run_lifecycle_done_sets_terminal() {
167        let mut state = RunLifecycleState::default();
168        RunLifecycle::apply(
169            &mut state,
170            RunLifecycleUpdate::Start {
171                run_id: "r1".into(),
172                updated_at: 100,
173            },
174        );
175        RunLifecycle::apply(
176            &mut state,
177            RunLifecycleUpdate::Done {
178                done_reason: "natural".into(),
179                updated_at: 200,
180            },
181        );
182        assert_eq!(state.status, RunStatus::Done);
183        assert_eq!(state.status_reason.as_deref(), Some("natural"));
184        assert!(state.status.is_terminal());
185    }
186
187    #[test]
188    fn run_lifecycle_waiting_transition() {
189        let mut state = RunLifecycleState::default();
190        RunLifecycle::apply(
191            &mut state,
192            RunLifecycleUpdate::Start {
193                run_id: "r1".into(),
194                updated_at: 100,
195            },
196        );
197        RunLifecycle::apply(
198            &mut state,
199            RunLifecycleUpdate::SetWaiting {
200                updated_at: 150,
201                pause_reason: "suspended".into(),
202            },
203        );
204        assert_eq!(state.status, RunStatus::Waiting);
205        assert_eq!(state.status_reason.as_deref(), Some("suspended"));
206    }
207
208    #[test]
209    fn run_lifecycle_status_reason_set_and_cleared() {
210        let mut state = RunLifecycleState::default();
211        RunLifecycle::apply(
212            &mut state,
213            RunLifecycleUpdate::Start {
214                run_id: "r1".into(),
215                updated_at: 100,
216            },
217        );
218        assert!(state.status_reason.is_none());
219
220        // SetWaiting stores status_reason
221        RunLifecycle::apply(
222            &mut state,
223            RunLifecycleUpdate::SetWaiting {
224                updated_at: 150,
225                pause_reason: "awaiting_tasks".into(),
226            },
227        );
228        assert_eq!(state.status_reason.as_deref(), Some("awaiting_tasks"));
229
230        // SetRunning clears status_reason
231        RunLifecycle::apply(
232            &mut state,
233            RunLifecycleUpdate::SetRunning { updated_at: 200 },
234        );
235        assert!(state.status_reason.is_none());
236
237        // SetWaiting again with different reason
238        RunLifecycle::apply(
239            &mut state,
240            RunLifecycleUpdate::SetWaiting {
241                updated_at: 250,
242                pause_reason: "user_input_required".into(),
243            },
244        );
245        assert_eq!(state.status_reason.as_deref(), Some("user_input_required"));
246
247        // Done sets status_reason
248        RunLifecycle::apply(
249            &mut state,
250            RunLifecycleUpdate::Done {
251                done_reason: "finished".into(),
252                updated_at: 300,
253            },
254        );
255        assert_eq!(state.status_reason.as_deref(), Some("finished"));
256    }
257
258    #[test]
259    fn run_lifecycle_status_reason_cleared_on_start() {
260        let mut state = RunLifecycleState {
261            run_id: "r1".into(),
262            status: RunStatus::Waiting,
263            status_reason: Some("old_reason".into()),
264            updated_at: 100,
265            step_count: 1,
266        };
267        RunLifecycle::apply(
268            &mut state,
269            RunLifecycleUpdate::Start {
270                run_id: "r2".into(),
271                updated_at: 200,
272            },
273        );
274        assert!(state.status_reason.is_none());
275    }
276
277    #[test]
278    fn run_lifecycle_full_sequence() {
279        let mut state = RunLifecycleState::default();
280        let t = now_ms();
281
282        // Start
283        RunLifecycle::apply(
284            &mut state,
285            RunLifecycleUpdate::Start {
286                run_id: "run-42".into(),
287                updated_at: t,
288            },
289        );
290        assert_eq!(state.status, RunStatus::Running);
291
292        // Steps
293        RunLifecycle::apply(
294            &mut state,
295            RunLifecycleUpdate::StepCompleted { updated_at: t + 1 },
296        );
297        RunLifecycle::apply(
298            &mut state,
299            RunLifecycleUpdate::StepCompleted { updated_at: t + 2 },
300        );
301        RunLifecycle::apply(
302            &mut state,
303            RunLifecycleUpdate::StepCompleted { updated_at: t + 3 },
304        );
305        assert_eq!(state.step_count, 3);
306
307        // Done
308        RunLifecycle::apply(
309            &mut state,
310            RunLifecycleUpdate::Done {
311                done_reason: "stopped:max_turns".into(),
312                updated_at: t + 4,
313            },
314        );
315        assert_eq!(state.status, RunStatus::Done);
316        assert_eq!(state.step_count, 3);
317    }
318
319    #[test]
320    fn run_lifecycle_rejects_done_to_running() {
321        let mut state = RunLifecycleState {
322            run_id: "r1".into(),
323            status: RunStatus::Done,
324            status_reason: Some("natural".into()),
325            updated_at: 100,
326            step_count: 1,
327        };
328        // Done -> Running should be rejected (state unchanged)
329        RunLifecycle::apply(
330            &mut state,
331            RunLifecycleUpdate::Start {
332                run_id: "r2".into(),
333                updated_at: 200,
334            },
335        );
336        assert_eq!(state.status, RunStatus::Done);
337        assert_eq!(state.run_id, "r1");
338        assert_eq!(state.updated_at, 100);
339    }
340
341    #[test]
342    fn run_lifecycle_rejects_done_to_waiting() {
343        let mut state = RunLifecycleState {
344            run_id: "r1".into(),
345            status: RunStatus::Done,
346            status_reason: Some("natural".into()),
347            updated_at: 100,
348            step_count: 1,
349        };
350        // Done -> Waiting should be rejected (state unchanged)
351        RunLifecycle::apply(
352            &mut state,
353            RunLifecycleUpdate::SetWaiting {
354                updated_at: 200,
355                pause_reason: "suspended".into(),
356            },
357        );
358        assert_eq!(state.status, RunStatus::Done);
359        assert_eq!(state.updated_at, 100);
360    }
361
362    #[test]
363    fn run_lifecycle_rejects_done_to_step_completed() {
364        let mut state = RunLifecycleState {
365            run_id: "r1".into(),
366            status: RunStatus::Done,
367            status_reason: Some("natural".into()),
368            updated_at: 100,
369            step_count: 1,
370        };
371        // Done -> Running (via StepCompleted) should be rejected (state unchanged)
372        RunLifecycle::apply(
373            &mut state,
374            RunLifecycleUpdate::StepCompleted { updated_at: 200 },
375        );
376        assert_eq!(state.status, RunStatus::Done);
377        assert_eq!(state.step_count, 1);
378        assert_eq!(state.updated_at, 100);
379    }
380
381    #[test]
382    fn run_lifecycle_allows_waiting_to_running_via_start() {
383        let mut state = RunLifecycleState {
384            run_id: "r1".into(),
385            status: RunStatus::Waiting,
386            status_reason: Some("suspended".into()),
387            updated_at: 100,
388            step_count: 1,
389        };
390        // Waiting -> Running is valid
391        RunLifecycle::apply(
392            &mut state,
393            RunLifecycleUpdate::Start {
394                run_id: "r2".into(),
395                updated_at: 200,
396            },
397        );
398        assert_eq!(state.status, RunStatus::Running);
399    }
400
401    #[test]
402    fn run_lifecycle_state_serde_roundtrip() {
403        let state = RunLifecycleState {
404            run_id: "r1".into(),
405            status: RunStatus::Done,
406            status_reason: Some("natural".into()),
407            updated_at: 12345,
408            step_count: 3,
409        };
410        let json = serde_json::to_string(&state).unwrap();
411        let parsed: RunLifecycleState = serde_json::from_str(&json).unwrap();
412        assert_eq!(parsed, state);
413    }
414
415    // -----------------------------------------------------------------------
416    // Migrated from uncarve: additional lifecycle tests
417    // -----------------------------------------------------------------------
418
419    #[test]
420    fn run_lifecycle_default_state() {
421        let state = RunLifecycleState::default();
422        assert!(state.run_id.is_empty());
423        assert_eq!(state.status, RunStatus::default());
424        assert!(state.status_reason.is_none());
425        assert_eq!(state.step_count, 0);
426        assert_eq!(state.updated_at, 0);
427    }
428
429    #[test]
430    fn run_lifecycle_multiple_steps_then_done() {
431        let mut state = RunLifecycleState::default();
432        RunLifecycle::apply(
433            &mut state,
434            RunLifecycleUpdate::Start {
435                run_id: "r1".into(),
436                updated_at: 100,
437            },
438        );
439
440        for step in 1..=10u64 {
441            RunLifecycle::apply(
442                &mut state,
443                RunLifecycleUpdate::StepCompleted {
444                    updated_at: 100 + step,
445                },
446            );
447        }
448        assert_eq!(state.step_count, 10);
449        assert_eq!(state.status, RunStatus::Running);
450
451        RunLifecycle::apply(
452            &mut state,
453            RunLifecycleUpdate::Done {
454                done_reason: "max_rounds".into(),
455                updated_at: 200,
456            },
457        );
458        assert_eq!(state.status, RunStatus::Done);
459        assert_eq!(state.step_count, 10);
460    }
461
462    #[test]
463    fn run_lifecycle_waiting_to_done() {
464        let mut state = RunLifecycleState::default();
465        RunLifecycle::apply(
466            &mut state,
467            RunLifecycleUpdate::Start {
468                run_id: "r1".into(),
469                updated_at: 100,
470            },
471        );
472        RunLifecycle::apply(
473            &mut state,
474            RunLifecycleUpdate::SetWaiting {
475                updated_at: 150,
476                pause_reason: "suspended".into(),
477            },
478        );
479        assert_eq!(state.status, RunStatus::Waiting);
480
481        RunLifecycle::apply(
482            &mut state,
483            RunLifecycleUpdate::Done {
484                done_reason: "cancelled".into(),
485                updated_at: 200,
486            },
487        );
488        assert_eq!(state.status, RunStatus::Done);
489    }
490
491    #[test]
492    fn run_lifecycle_waiting_to_running_to_waiting() {
493        let mut state = RunLifecycleState::default();
494        RunLifecycle::apply(
495            &mut state,
496            RunLifecycleUpdate::Start {
497                run_id: "r1".into(),
498                updated_at: 100,
499            },
500        );
501        RunLifecycle::apply(
502            &mut state,
503            RunLifecycleUpdate::SetWaiting {
504                updated_at: 150,
505                pause_reason: "suspended".into(),
506            },
507        );
508        RunLifecycle::apply(
509            &mut state,
510            RunLifecycleUpdate::SetRunning { updated_at: 200 },
511        );
512        assert_eq!(state.status, RunStatus::Running);
513        RunLifecycle::apply(
514            &mut state,
515            RunLifecycleUpdate::SetWaiting {
516                updated_at: 250,
517                pause_reason: "suspended".into(),
518            },
519        );
520        assert_eq!(state.status, RunStatus::Waiting);
521    }
522
523    #[test]
524    fn run_lifecycle_update_target_status() {
525        assert_eq!(
526            RunLifecycleUpdate::Start {
527                run_id: "r".into(),
528                updated_at: 0,
529            }
530            .target_status(),
531            RunStatus::Running
532        );
533        assert_eq!(
534            RunLifecycleUpdate::StepCompleted { updated_at: 0 }.target_status(),
535            RunStatus::Running
536        );
537        assert_eq!(
538            RunLifecycleUpdate::SetWaiting {
539                updated_at: 0,
540                pause_reason: "test".into()
541            }
542            .target_status(),
543            RunStatus::Waiting
544        );
545        assert_eq!(
546            RunLifecycleUpdate::SetRunning { updated_at: 0 }.target_status(),
547            RunStatus::Running
548        );
549        assert_eq!(
550            RunLifecycleUpdate::Done {
551                done_reason: "done".into(),
552                updated_at: 0,
553            }
554            .target_status(),
555            RunStatus::Done
556        );
557    }
558
559    #[test]
560    fn run_lifecycle_start_resets_step_count() {
561        let mut state = RunLifecycleState::default();
562        RunLifecycle::apply(
563            &mut state,
564            RunLifecycleUpdate::Start {
565                run_id: "r1".into(),
566                updated_at: 100,
567            },
568        );
569        RunLifecycle::apply(
570            &mut state,
571            RunLifecycleUpdate::StepCompleted { updated_at: 200 },
572        );
573        assert_eq!(state.step_count, 1);
574
575        // Transition to waiting, then start a new run
576        RunLifecycle::apply(
577            &mut state,
578            RunLifecycleUpdate::SetWaiting {
579                updated_at: 250,
580                pause_reason: "suspended".into(),
581            },
582        );
583        RunLifecycle::apply(
584            &mut state,
585            RunLifecycleUpdate::Start {
586                run_id: "r2".into(),
587                updated_at: 300,
588            },
589        );
590        assert_eq!(state.step_count, 0, "step count should reset on new start");
591        assert_eq!(state.run_id, "r2");
592    }
593
594    #[test]
595    fn run_lifecycle_done_preserves_step_count() {
596        let mut state = RunLifecycleState::default();
597        RunLifecycle::apply(
598            &mut state,
599            RunLifecycleUpdate::Start {
600                run_id: "r1".into(),
601                updated_at: 100,
602            },
603        );
604        RunLifecycle::apply(
605            &mut state,
606            RunLifecycleUpdate::StepCompleted { updated_at: 200 },
607        );
608        RunLifecycle::apply(
609            &mut state,
610            RunLifecycleUpdate::StepCompleted { updated_at: 300 },
611        );
612        RunLifecycle::apply(
613            &mut state,
614            RunLifecycleUpdate::Done {
615                done_reason: "finished".into(),
616                updated_at: 400,
617            },
618        );
619        assert_eq!(state.step_count, 2, "done should not reset step count");
620    }
621
622    #[test]
623    fn run_lifecycle_state_equality() {
624        let s1 = RunLifecycleState {
625            run_id: "r1".into(),
626            status: RunStatus::Running,
627            status_reason: None,
628            updated_at: 100,
629            step_count: 3,
630        };
631        let s2 = s1.clone();
632        assert_eq!(s1, s2);
633
634        let s3 = RunLifecycleState {
635            step_count: 4,
636            ..s1.clone()
637        };
638        assert_ne!(s1, s3);
639    }
640
641    #[test]
642    fn run_lifecycle_status_is_terminal() {
643        let mut state = RunLifecycleState::default();
644        RunLifecycle::apply(
645            &mut state,
646            RunLifecycleUpdate::Start {
647                run_id: "r1".into(),
648                updated_at: 100,
649            },
650        );
651        assert!(!state.status.is_terminal());
652
653        RunLifecycle::apply(
654            &mut state,
655            RunLifecycleUpdate::Done {
656                done_reason: "done".into(),
657                updated_at: 200,
658            },
659        );
660        assert!(state.status.is_terminal());
661    }
662
663    // -----------------------------------------------------------------------
664    // Continuation semantics: SetRunning preserves step_count
665    // -----------------------------------------------------------------------
666
667    #[test]
668    fn continuation_set_running_preserves_step_count() {
669        let mut state = RunLifecycleState::default();
670        RunLifecycle::apply(
671            &mut state,
672            RunLifecycleUpdate::Start {
673                run_id: "r1".into(),
674                updated_at: 100,
675            },
676        );
677        RunLifecycle::apply(
678            &mut state,
679            RunLifecycleUpdate::StepCompleted { updated_at: 200 },
680        );
681        RunLifecycle::apply(
682            &mut state,
683            RunLifecycleUpdate::StepCompleted { updated_at: 300 },
684        );
685        assert_eq!(state.step_count, 2);
686
687        // Simulate awaiting_tasks → continuation
688        RunLifecycle::apply(
689            &mut state,
690            RunLifecycleUpdate::SetWaiting {
691                updated_at: 400,
692                pause_reason: "awaiting_tasks".into(),
693            },
694        );
695        // Continuation: SetRunning instead of Start → step_count preserved
696        RunLifecycle::apply(
697            &mut state,
698            RunLifecycleUpdate::SetRunning { updated_at: 500 },
699        );
700        assert_eq!(state.status, RunStatus::Running);
701        assert_eq!(state.step_count, 2, "continuation must preserve step_count");
702        assert!(state.status_reason.is_none());
703    }
704
705    #[test]
706    fn new_start_resets_step_count_after_waiting() {
707        let mut state = RunLifecycleState::default();
708        RunLifecycle::apply(
709            &mut state,
710            RunLifecycleUpdate::Start {
711                run_id: "r1".into(),
712                updated_at: 100,
713            },
714        );
715        RunLifecycle::apply(
716            &mut state,
717            RunLifecycleUpdate::StepCompleted { updated_at: 200 },
718        );
719        RunLifecycle::apply(
720            &mut state,
721            RunLifecycleUpdate::SetWaiting {
722                updated_at: 300,
723                pause_reason: "awaiting_tasks".into(),
724            },
725        );
726        // New start (not continuation) → step_count resets
727        RunLifecycle::apply(
728            &mut state,
729            RunLifecycleUpdate::Start {
730                run_id: "r2".into(),
731                updated_at: 400,
732            },
733        );
734        assert_eq!(state.step_count, 0, "new Start must reset step_count");
735        assert_eq!(state.run_id, "r2");
736    }
737
738    #[test]
739    fn status_reason_serde_backward_compat_missing() {
740        // Old serialized state without any reason field
741        let json = r#"{"run_id":"r1","status":"waiting","updated_at":100,"step_count":0}"#;
742        let parsed: RunLifecycleState = serde_json::from_str(json).unwrap();
743        assert!(
744            parsed.status_reason.is_none(),
745            "missing status_reason should deserialize as None"
746        );
747    }
748
749    #[test]
750    fn status_reason_serde_backward_compat_done_reason_alias() {
751        // Old serialized state with done_reason field
752        let json = r#"{"run_id":"r1","status":"done","done_reason":"natural","updated_at":100,"step_count":1}"#;
753        let parsed: RunLifecycleState = serde_json::from_str(json).unwrap();
754        assert_eq!(parsed.status_reason.as_deref(), Some("natural"));
755    }
756
757    #[test]
758    fn status_reason_serde_backward_compat_pause_reason_alias() {
759        // Old serialized state with pause_reason field
760        let json = r#"{"run_id":"r1","status":"waiting","pause_reason":"awaiting_tasks","updated_at":100,"step_count":2}"#;
761        let parsed: RunLifecycleState = serde_json::from_str(json).unwrap();
762        assert_eq!(parsed.status_reason.as_deref(), Some("awaiting_tasks"));
763    }
764
765    #[test]
766    fn status_reason_included_in_serde() {
767        let state = RunLifecycleState {
768            run_id: "r1".into(),
769            status: RunStatus::Waiting,
770            status_reason: Some("awaiting_tasks".into()),
771            updated_at: 100,
772            step_count: 2,
773        };
774        let json = serde_json::to_string(&state).unwrap();
775        assert!(json.contains("awaiting_tasks"));
776        let parsed: RunLifecycleState = serde_json::from_str(&json).unwrap();
777        assert_eq!(parsed.status_reason.as_deref(), Some("awaiting_tasks"));
778    }
779}
awaken_runtime/agent/state/run_lifecycle.rs

awaken_runtime/agent/state/
run_lifecycle.rs
