Skip to main content

mika_a2a/
state_machine.rs

1use crate::types::TaskState;
2
3/// Validates task state transitions according to the A2A protocol.
4pub struct TaskStateMachine;
5
6impl TaskStateMachine {
7    /// Check if a state is terminal (no further transitions allowed).
8    pub fn is_terminal(state: TaskState) -> bool {
9        matches!(
10            state,
11            TaskState::Completed | TaskState::Failed | TaskState::Canceled | TaskState::Rejected
12        )
13    }
14
15    /// Check if a transition from one state to another is valid.
16    pub fn can_transition(from: TaskState, to: TaskState) -> bool {
17        matches!(
18            (from, to),
19            // From Submitted
20            (TaskState::Submitted, TaskState::Working)
21                | (TaskState::Submitted, TaskState::Canceled)
22                | (TaskState::Submitted, TaskState::Rejected)
23                // From Working
24                | (TaskState::Working, TaskState::Completed)
25                | (TaskState::Working, TaskState::Failed)
26                | (TaskState::Working, TaskState::Canceled)
27                | (TaskState::Working, TaskState::InputRequired)
28                | (TaskState::Working, TaskState::AuthRequired)
29                // From InputRequired
30                | (TaskState::InputRequired, TaskState::Working)
31                | (TaskState::InputRequired, TaskState::Canceled)
32                | (TaskState::InputRequired, TaskState::Failed)
33                // From AuthRequired
34                | (TaskState::AuthRequired, TaskState::Working)
35                | (TaskState::AuthRequired, TaskState::Canceled)
36                | (TaskState::AuthRequired, TaskState::Failed)
37        )
38    }
39}
40
41#[cfg(test)]
42mod tests {
43    use super::*;
44
45    #[test]
46    fn test_terminal_states() {
47        assert!(TaskStateMachine::is_terminal(TaskState::Completed));
48        assert!(TaskStateMachine::is_terminal(TaskState::Failed));
49        assert!(TaskStateMachine::is_terminal(TaskState::Canceled));
50        assert!(TaskStateMachine::is_terminal(TaskState::Rejected));
51        assert!(!TaskStateMachine::is_terminal(TaskState::Submitted));
52        assert!(!TaskStateMachine::is_terminal(TaskState::Working));
53        assert!(!TaskStateMachine::is_terminal(TaskState::InputRequired));
54        assert!(!TaskStateMachine::is_terminal(TaskState::AuthRequired));
55        assert!(!TaskStateMachine::is_terminal(TaskState::Unknown));
56    }
57
58    #[test]
59    fn test_valid_transitions() {
60        // Submitted transitions
61        assert!(TaskStateMachine::can_transition(
62            TaskState::Submitted,
63            TaskState::Working
64        ));
65        assert!(TaskStateMachine::can_transition(
66            TaskState::Submitted,
67            TaskState::Canceled
68        ));
69        assert!(TaskStateMachine::can_transition(
70            TaskState::Submitted,
71            TaskState::Rejected
72        ));
73
74        // Working transitions
75        assert!(TaskStateMachine::can_transition(
76            TaskState::Working,
77            TaskState::Completed
78        ));
79        assert!(TaskStateMachine::can_transition(
80            TaskState::Working,
81            TaskState::Failed
82        ));
83        assert!(TaskStateMachine::can_transition(
84            TaskState::Working,
85            TaskState::Canceled
86        ));
87        assert!(TaskStateMachine::can_transition(
88            TaskState::Working,
89            TaskState::InputRequired
90        ));
91        assert!(TaskStateMachine::can_transition(
92            TaskState::Working,
93            TaskState::AuthRequired
94        ));
95
96        // InputRequired transitions
97        assert!(TaskStateMachine::can_transition(
98            TaskState::InputRequired,
99            TaskState::Working
100        ));
101        assert!(TaskStateMachine::can_transition(
102            TaskState::InputRequired,
103            TaskState::Canceled
104        ));
105        assert!(TaskStateMachine::can_transition(
106            TaskState::InputRequired,
107            TaskState::Failed
108        ));
109
110        // AuthRequired transitions
111        assert!(TaskStateMachine::can_transition(
112            TaskState::AuthRequired,
113            TaskState::Working
114        ));
115        assert!(TaskStateMachine::can_transition(
116            TaskState::AuthRequired,
117            TaskState::Canceled
118        ));
119        assert!(TaskStateMachine::can_transition(
120            TaskState::AuthRequired,
121            TaskState::Failed
122        ));
123    }
124
125    #[test]
126    fn test_invalid_transitions() {
127        // Cannot transition from terminal states
128        assert!(!TaskStateMachine::can_transition(
129            TaskState::Completed,
130            TaskState::Working
131        ));
132        assert!(!TaskStateMachine::can_transition(
133            TaskState::Failed,
134            TaskState::Working
135        ));
136        assert!(!TaskStateMachine::can_transition(
137            TaskState::Canceled,
138            TaskState::Working
139        ));
140        assert!(!TaskStateMachine::can_transition(
141            TaskState::Rejected,
142            TaskState::Working
143        ));
144
145        // Cannot skip states
146        assert!(!TaskStateMachine::can_transition(
147            TaskState::Submitted,
148            TaskState::Completed
149        ));
150        assert!(!TaskStateMachine::can_transition(
151            TaskState::Submitted,
152            TaskState::InputRequired
153        ));
154
155        // Cannot go backward
156        assert!(!TaskStateMachine::can_transition(
157            TaskState::Working,
158            TaskState::Submitted
159        ));
160    }
161}