ai_agents_state/
machine.rs

1use chrono::Utc;
2use parking_lot::RwLock;
3use std::sync::atomic::{AtomicU64, Ordering};
4
5use ai_agents_core::{AgentError, Result, StateMachineSnapshot, StateTransitionEvent};
6
7use super::config::{StateConfig, StateDefinition, ToolRef, Transition};
8
9pub struct StateMachine {
10    config: StateConfig,
11    /// Serializes state-authority snapshots with transitions, reset, and restore.
12    state_guard: RwLock<()>,
13    /// Monotonic generation for changes that can invalidate state-scoped authorization.
14    generation: AtomicU64,
15    current: RwLock<String>,
16    previous: RwLock<Option<String>>,
17    turn_count: RwLock<u32>,
18    no_transition_count: RwLock<u32>,
19    history: RwLock<Vec<StateTransitionEvent>>,
20}
21
22impl StateMachine {
23    pub fn new(config: StateConfig) -> Result<Self> {
24        config.validate()?;
25        let initial = Self::resolve_initial_state(&config)?;
26        Ok(Self {
27            config,
28            state_guard: RwLock::new(()),
29            generation: AtomicU64::new(1),
30            current: RwLock::new(initial),
31            previous: RwLock::new(None),
32            turn_count: RwLock::new(0),
33            no_transition_count: RwLock::new(0),
34            history: RwLock::new(Vec::new()),
35        })
36    }
37
38    fn resolve_initial_state(config: &StateConfig) -> Result<String> {
39        let mut path = config.initial.clone();
40        let mut current_def = config.states.get(&config.initial);
41
42        while let Some(def) = current_def {
43            if let (Some(initial_sub), Some(sub_states)) = (&def.initial, &def.states) {
44                path = format!("{}.{}", path, initial_sub);
45                current_def = sub_states.get(initial_sub);
46            } else {
47                break;
48            }
49        }
50
51        Ok(path)
52    }
53
54    pub fn current(&self) -> String {
55        self.current.read().clone()
56    }
57
58    pub fn previous(&self) -> Option<String> {
59        self.previous.read().clone()
60    }
61
62    pub fn current_definition(&self) -> Option<StateDefinition> {
63        let current = self.current.read();
64        self.config.get_state(&current).cloned()
65    }
66
67    pub fn get_definition(&self, state: &str) -> Option<&StateDefinition> {
68        self.config.get_state(state)
69    }
70
71    pub fn get_parent_definition(&self) -> Option<StateDefinition> {
72        let current = self.current.read();
73        let parts: Vec<&str> = current.split('.').collect();
74        if parts.len() <= 1 {
75            return None;
76        }
77        let parent_path = parts[..parts.len() - 1].join(".");
78        self.config.get_state(&parent_path).cloned()
79    }
80
81    /// Returns the generation and every explicit tool scope from root to current state as one authority snapshot.
82    pub fn current_tool_scope_snapshot(&self) -> (u64, Vec<Vec<ToolRef>>) {
83        let _guard = self.state_guard.read();
84        let current = self.current.read().clone();
85        let mut path = String::new();
86        let mut scopes = Vec::new();
87
88        for part in current.split('.') {
89            if !path.is_empty() {
90                path.push('.');
91            }
92            path.push_str(part);
93            if let Some(tools) = self
94                .config
95                .get_state(&path)
96                .and_then(|definition| definition.tools.clone())
97            {
98                scopes.push(tools);
99            }
100        }
101
102        (self.generation.load(Ordering::SeqCst), scopes)
103    }
104
105    /// Returns the current state-authority generation after any active mutation completes.
106    pub fn generation(&self) -> u64 {
107        let _guard = self.state_guard.read();
108        self.generation.load(Ordering::SeqCst)
109    }
110
111    /// Applies one valid transition and advances the authority generation after all state fields are committed.
112    pub fn transition_to(&self, state: &str, reason: &str) -> Result<()> {
113        let _guard = self.state_guard.write();
114        let current_path = self.current.read().clone();
115        let resolved_path = self.config.resolve_full_path(&current_path, state);
116
117        if self.config.get_state(&resolved_path).is_none() {
118            return Err(AgentError::InvalidSpec(format!(
119                "Unknown state: {} (resolved from {})",
120                resolved_path, state
121            )));
122        }
123
124        let final_path = self.resolve_to_leaf_state(&resolved_path)?;
125
126        let from = {
127            let mut current = self.current.write();
128            let mut previous = self.previous.write();
129            let from = current.clone();
130            *previous = Some(from.clone());
131            *current = final_path.clone();
132            from
133        };
134
135        *self.turn_count.write() = 0;
136        *self.no_transition_count.write() = 0;
137
138        let event = StateTransitionEvent {
139            from,
140            to: final_path,
141            reason: reason.to_string(),
142            timestamp: Utc::now(),
143        };
144        self.history.write().push(event);
145        self.generation.fetch_add(1, Ordering::SeqCst);
146
147        Ok(())
148    }
149
150    fn resolve_to_leaf_state(&self, path: &str) -> Result<String> {
151        let mut current_path = path.to_string();
152
153        loop {
154            let def = self.config.get_state(&current_path).ok_or_else(|| {
155                AgentError::InvalidSpec(format!("State not found: {}", current_path))
156            })?;
157
158            if let (Some(initial_sub), Some(sub_states)) = (&def.initial, &def.states)
159                && sub_states.contains_key(initial_sub)
160            {
161                current_path = format!("{}.{}", current_path, initial_sub);
162                continue;
163            }
164            break;
165        }
166
167        Ok(current_path)
168    }
169
170    pub fn available_transitions(&self) -> Vec<Transition> {
171        let mut transitions = Vec::new();
172
173        if let Some(def) = self.current_definition() {
174            transitions.extend(def.transitions.clone());
175        }
176
177        transitions.extend(self.config.global_transitions.clone());
178
179        transitions.sort_by_key(|transition| std::cmp::Reverse(transition.priority));
180        transitions
181    }
182
183    pub fn auto_transitions(&self) -> Vec<Transition> {
184        self.available_transitions()
185            .into_iter()
186            .filter(|t| t.auto)
187            .collect()
188    }
189
190    pub fn history(&self) -> Vec<StateTransitionEvent> {
191        self.history.read().clone()
192    }
193
194    pub fn increment_turn(&self) {
195        let _guard = self.state_guard.write();
196        *self.turn_count.write() += 1;
197    }
198
199    pub fn turn_count(&self) -> u32 {
200        *self.turn_count.read()
201    }
202
203    pub fn increment_no_transition(&self) {
204        let _guard = self.state_guard.write();
205        *self.no_transition_count.write() += 1;
206    }
207
208    pub fn no_transition_count(&self) -> u32 {
209        *self.no_transition_count.read()
210    }
211
212    pub fn reset_no_transition(&self) {
213        let _guard = self.state_guard.write();
214        *self.no_transition_count.write() = 0;
215    }
216
217    pub fn check_fallback(&self) -> Option<String> {
218        if let Some(max) = self.config.max_no_transition
219            && self.no_transition_count() >= max
220        {
221            return self.config.fallback.clone();
222        }
223        None
224    }
225
226    /// Restores the initial state and advances generation only when reset changes machine state.
227    pub fn reset(&self) {
228        let _guard = self.state_guard.write();
229        let initial =
230            Self::resolve_initial_state(&self.config).unwrap_or(self.config.initial.clone());
231        let changed = *self.current.read() != initial
232            || self.previous.read().is_some()
233            || *self.turn_count.read() != 0
234            || *self.no_transition_count.read() != 0
235            || !self.history.read().is_empty();
236        if !changed {
237            return;
238        }
239
240        *self.current.write() = initial;
241        *self.previous.write() = None;
242        *self.turn_count.write() = 0;
243        *self.no_transition_count.write() = 0;
244        self.history.write().clear();
245        self.generation.fetch_add(1, Ordering::SeqCst);
246    }
247
248    /// Captures a consistent persistence snapshot after any active state mutation completes.
249    pub fn snapshot(&self) -> StateMachineSnapshot {
250        let _guard = self.state_guard.read();
251        StateMachineSnapshot {
252            current_state: self.current.read().clone(),
253            previous_state: self.previous.read().clone(),
254            turn_count: *self.turn_count.read(),
255            no_transition_count: *self.no_transition_count.read(),
256            history: self.history.read().clone(),
257        }
258    }
259
260    /// Restores a valid snapshot and advances generation only when persisted machine state differs.
261    pub fn restore(&self, snapshot: StateMachineSnapshot) -> Result<()> {
262        if self.config.get_state(&snapshot.current_state).is_none() {
263            return Err(AgentError::InvalidSpec(format!(
264                "Snapshot contains unknown state: {}",
265                snapshot.current_state
266            )));
267        }
268
269        let _guard = self.state_guard.write();
270        let current_history = self.history.read();
271        let history_matches = current_history.len() == snapshot.history.len()
272            && current_history
273                .iter()
274                .zip(&snapshot.history)
275                .all(|(left, right)| {
276                    left.from == right.from
277                        && left.to == right.to
278                        && left.reason == right.reason
279                        && left.timestamp == right.timestamp
280                });
281        let changed = *self.current.read() != snapshot.current_state
282            || *self.previous.read() != snapshot.previous_state
283            || *self.turn_count.read() != snapshot.turn_count
284            || *self.no_transition_count.read() != snapshot.no_transition_count
285            || !history_matches;
286        drop(current_history);
287        if !changed {
288            return Ok(());
289        }
290
291        *self.current.write() = snapshot.current_state;
292        *self.previous.write() = snapshot.previous_state;
293        *self.turn_count.write() = snapshot.turn_count;
294        *self.no_transition_count.write() = snapshot.no_transition_count;
295        *self.history.write() = snapshot.history;
296        self.generation.fetch_add(1, Ordering::SeqCst);
297        Ok(())
298    }
299
300    pub fn config(&self) -> &StateConfig {
301        &self.config
302    }
303
304    pub fn check_timeout(&self) -> Option<String> {
305        let def = self.current_definition()?;
306        let max_turns = def.max_turns?;
307        let timeout_to = def.timeout_to.as_ref()?;
308        if self.turn_count() >= max_turns {
309            let current_path = self.current.read().clone();
310            Some(self.config.resolve_full_path(&current_path, timeout_to))
311        } else {
312            None
313        }
314    }
315
316    pub fn current_depth(&self) -> usize {
317        self.current.read().split('.').count()
318    }
319
320    pub fn is_in_sub_state(&self) -> bool {
321        self.current_depth() > 1
322    }
323
324    pub fn parent_state(&self) -> Option<String> {
325        let current = self.current.read();
326        let parts: Vec<&str> = current.split('.').collect();
327        if parts.len() > 1 {
328            Some(parts[..parts.len() - 1].join("."))
329        } else {
330            None
331        }
332    }
333
334    pub fn root_state(&self) -> String {
335        let current = self.current.read();
336        current.split('.').next().unwrap_or(&current).to_string()
337    }
338
339    /// Check if a transition target is on cooldown (was recently transitioned to).
340    pub fn is_on_cooldown(&self, target: &str, cooldown_turns: u32) -> bool {
341        let history = self.history.read();
342        let total_transitions = history.len();
343        if total_transitions == 0 || cooldown_turns == 0 {
344            return false;
345        }
346        // Look at the last `cooldown_turns` transitions
347        let start = total_transitions.saturating_sub(cooldown_turns as usize);
348        history[start..].iter().any(|e| e.to == target)
349    }
350}
351
352#[cfg(test)]
353mod tests {
354    use super::super::config::TransitionTiming;
355    use super::*;
356    use std::collections::HashMap;
357
358    fn create_test_config() -> StateConfig {
359        let mut states = HashMap::new();
360        states.insert(
361            "greeting".into(),
362            StateDefinition {
363                prompt: Some("Welcome!".into()),
364                transitions: vec![Transition {
365                    to: "support".into(),
366                    when: "needs help".into(),
367                    guard: None,
368                    intent: None,
369                    auto: true,
370                    priority: 10,
371                    cooldown_turns: None,
372                    timing: TransitionTiming::PostResponse,
373                    requires_response: false,
374                    run_extractors: false,
375                }],
376                ..Default::default()
377            },
378        );
379        states.insert(
380            "support".into(),
381            StateDefinition {
382                prompt: Some("How can I help?".into()),
383                max_turns: Some(3),
384                timeout_to: Some("escalation".into()),
385                ..Default::default()
386            },
387        );
388        states.insert(
389            "escalation".into(),
390            StateDefinition {
391                prompt: Some("Escalating...".into()),
392                ..Default::default()
393            },
394        );
395        StateConfig {
396            initial: "greeting".into(),
397            states,
398            global_transitions: vec![],
399            fallback: None,
400            max_no_transition: None,
401            regenerate_on_transition: true,
402        }
403    }
404
405    fn create_hierarchical_config() -> StateConfig {
406        let mut sub_states = HashMap::new();
407        sub_states.insert(
408            "gathering_info".into(),
409            StateDefinition {
410                prompt: Some("Gathering info".into()),
411                transitions: vec![Transition {
412                    to: "proposing".into(),
413                    when: "understood".into(),
414                    guard: None,
415                    intent: None,
416                    auto: true,
417                    priority: 0,
418                    cooldown_turns: None,
419                    timing: TransitionTiming::PostResponse,
420                    requires_response: false,
421                    run_extractors: false,
422                }],
423                ..Default::default()
424            },
425        );
426        sub_states.insert(
427            "proposing".into(),
428            StateDefinition {
429                prompt: Some("Proposing solution".into()),
430                transitions: vec![Transition {
431                    to: "^closing".into(),
432                    when: "resolved".into(),
433                    guard: None,
434                    intent: None,
435                    auto: true,
436                    priority: 0,
437                    cooldown_turns: None,
438                    timing: TransitionTiming::PostResponse,
439                    requires_response: false,
440                    run_extractors: false,
441                }],
442                ..Default::default()
443            },
444        );
445
446        let mut states = HashMap::new();
447        states.insert(
448            "problem_solving".into(),
449            StateDefinition {
450                prompt: Some("Problem solving".into()),
451                initial: Some("gathering_info".into()),
452                states: Some(sub_states),
453                ..Default::default()
454            },
455        );
456        states.insert(
457            "closing".into(),
458            StateDefinition {
459                prompt: Some("Thank you".into()),
460                ..Default::default()
461            },
462        );
463
464        StateConfig {
465            initial: "problem_solving".into(),
466            states,
467            global_transitions: vec![],
468            fallback: None,
469            max_no_transition: None,
470            regenerate_on_transition: true,
471        }
472    }
473
474    #[test]
475    fn test_new_state_machine() {
476        let config = create_test_config();
477        let sm = StateMachine::new(config).unwrap();
478        assert_eq!(sm.current(), "greeting");
479        assert!(sm.previous().is_none());
480        assert_eq!(sm.turn_count(), 0);
481    }
482
483    #[test]
484    fn test_transition() {
485        let config = create_test_config();
486        let sm = StateMachine::new(config).unwrap();
487        sm.transition_to("support", "user asked for help").unwrap();
488        assert_eq!(sm.current(), "support");
489        assert_eq!(sm.previous(), Some("greeting".into()));
490        assert_eq!(sm.history().len(), 1);
491    }
492
493    #[test]
494    fn test_turn_counting() {
495        let config = create_test_config();
496        let sm = StateMachine::new(config).unwrap();
497        assert_eq!(sm.turn_count(), 0);
498        sm.increment_turn();
499        sm.increment_turn();
500        assert_eq!(sm.turn_count(), 2);
501        sm.transition_to("support", "reason").unwrap();
502        assert_eq!(sm.turn_count(), 0);
503    }
504
505    #[test]
506    fn test_timeout_check() {
507        let config = create_test_config();
508        let sm = StateMachine::new(config).unwrap();
509        sm.transition_to("support", "needs help").unwrap();
510        assert!(sm.check_timeout().is_none());
511        sm.increment_turn();
512        sm.increment_turn();
513        sm.increment_turn();
514        assert_eq!(sm.check_timeout(), Some("escalation".into()));
515    }
516
517    #[test]
518    fn test_snapshot_restore() {
519        let config = create_test_config();
520        let sm = StateMachine::new(config.clone()).unwrap();
521        sm.transition_to("support", "reason").unwrap();
522        sm.increment_turn();
523
524        let snapshot = sm.snapshot();
525        assert_eq!(snapshot.current_state, "support");
526        assert_eq!(snapshot.turn_count, 1);
527
528        let sm2 = StateMachine::new(config).unwrap();
529        sm2.restore(snapshot).unwrap();
530        assert_eq!(sm2.current(), "support");
531        assert_eq!(sm2.turn_count(), 1);
532    }
533
534    #[test]
535    fn test_reset() {
536        let config = create_test_config();
537        let sm = StateMachine::new(config).unwrap();
538        sm.transition_to("support", "reason").unwrap();
539        sm.increment_turn();
540        sm.reset();
541        assert_eq!(sm.current(), "greeting");
542        assert!(sm.previous().is_none());
543        assert_eq!(sm.turn_count(), 0);
544        assert!(sm.history().is_empty());
545    }
546
547    #[test]
548    fn generation_tracks_only_successful_transition_reset_and_restore_changes() {
549        let sm = StateMachine::new(create_test_config()).unwrap();
550        let initial_generation = sm.generation();
551        let initial_snapshot = sm.snapshot();
552
553        sm.reset();
554        sm.restore(initial_snapshot.clone()).unwrap();
555        assert_eq!(sm.generation(), initial_generation);
556
557        assert!(sm.transition_to("missing", "invalid").is_err());
558        assert_eq!(sm.generation(), initial_generation);
559
560        sm.transition_to("support", "valid").unwrap();
561        let transition_generation = sm.generation();
562        assert_eq!(transition_generation, initial_generation + 1);
563        let transitioned_snapshot = sm.snapshot();
564
565        sm.restore(transitioned_snapshot.clone()).unwrap();
566        assert_eq!(sm.generation(), transition_generation);
567        sm.increment_turn();
568        assert_eq!(sm.generation(), transition_generation);
569
570        sm.reset();
571        let reset_generation = sm.generation();
572        assert_eq!(reset_generation, transition_generation + 1);
573        sm.reset();
574        assert_eq!(sm.generation(), reset_generation);
575
576        sm.restore(transitioned_snapshot).unwrap();
577        assert_eq!(sm.generation(), reset_generation + 1);
578        sm.restore(initial_snapshot).unwrap();
579        assert_eq!(sm.generation(), reset_generation + 2);
580    }
581
582    #[test]
583    fn current_tool_scope_snapshot_keeps_complete_ancestor_chain() {
584        let yaml = r#"
585initial: root
586states:
587  root:
588    tools: [root_tool]
589    initial: middle
590    states:
591      middle:
592        initial: leaf
593        states:
594          leaf:
595            tools: []
596"#;
597        let config: StateConfig = serde_yaml::from_str(yaml).unwrap();
598        let sm = StateMachine::new(config).unwrap();
599
600        let (generation, scopes) = sm.current_tool_scope_snapshot();
601
602        assert_eq!(generation, sm.generation());
603        assert_eq!(scopes.len(), 2);
604        assert_eq!(scopes[0][0].id(), "root_tool");
605        assert!(scopes[1].is_empty());
606    }
607
608    #[test]
609    fn test_hierarchical_initial_state() {
610        let config = create_hierarchical_config();
611        let sm = StateMachine::new(config).unwrap();
612        assert_eq!(sm.current(), "problem_solving.gathering_info");
613    }
614
615    #[test]
616    fn test_hierarchical_transition_sibling() {
617        let config = create_hierarchical_config();
618        let sm = StateMachine::new(config).unwrap();
619        assert_eq!(sm.current(), "problem_solving.gathering_info");
620
621        sm.transition_to("proposing", "understood").unwrap();
622        assert_eq!(sm.current(), "problem_solving.proposing");
623    }
624
625    #[test]
626    fn test_hierarchical_transition_parent() {
627        let config = create_hierarchical_config();
628        let sm = StateMachine::new(config).unwrap();
629        sm.transition_to("proposing", "understood").unwrap();
630        sm.transition_to("^closing", "resolved").unwrap();
631        assert_eq!(sm.current(), "closing");
632    }
633
634    #[test]
635    fn test_current_depth() {
636        let config = create_hierarchical_config();
637        let sm = StateMachine::new(config).unwrap();
638        assert_eq!(sm.current_depth(), 2);
639        assert!(sm.is_in_sub_state());
640
641        sm.transition_to("^closing", "done").unwrap();
642        assert_eq!(sm.current_depth(), 1);
643        assert!(!sm.is_in_sub_state());
644    }
645
646    #[test]
647    fn test_parent_state() {
648        let config = create_hierarchical_config();
649        let sm = StateMachine::new(config).unwrap();
650        assert_eq!(sm.parent_state(), Some("problem_solving".into()));
651
652        sm.transition_to("^closing", "done").unwrap();
653        assert!(sm.parent_state().is_none());
654    }
655
656    #[test]
657    fn test_root_state() {
658        let config = create_hierarchical_config();
659        let sm = StateMachine::new(config).unwrap();
660        assert_eq!(sm.root_state(), "problem_solving");
661
662        sm.transition_to("^closing", "done").unwrap();
663        assert_eq!(sm.root_state(), "closing");
664    }
665
666    #[test]
667    fn test_no_transition_count() {
668        let config = create_test_config();
669        let sm = StateMachine::new(config).unwrap();
670
671        assert_eq!(sm.no_transition_count(), 0);
672        sm.increment_no_transition();
673        sm.increment_no_transition();
674        assert_eq!(sm.no_transition_count(), 2);
675
676        sm.reset_no_transition();
677        assert_eq!(sm.no_transition_count(), 0);
678    }
679
680    #[test]
681    fn test_fallback() {
682        let mut config = create_test_config();
683        config.fallback = Some("escalation".into());
684        config.max_no_transition = Some(3);
685
686        let sm = StateMachine::new(config).unwrap();
687        assert!(sm.check_fallback().is_none());
688
689        sm.increment_no_transition();
690        sm.increment_no_transition();
691        sm.increment_no_transition();
692        assert_eq!(sm.check_fallback(), Some("escalation".into()));
693    }
694
695    #[test]
696    fn test_global_transitions() {
697        let mut config = create_test_config();
698        config.global_transitions = vec![Transition {
699            to: "escalation".into(),
700            when: "user is angry".into(),
701            guard: None,
702            intent: None,
703            auto: true,
704            priority: 100,
705            cooldown_turns: None,
706            timing: TransitionTiming::PostResponse,
707            requires_response: false,
708            run_extractors: false,
709        }];
710
711        let sm = StateMachine::new(config).unwrap();
712        let transitions = sm.available_transitions();
713
714        assert!(
715            transitions
716                .iter()
717                .any(|t| t.to == "escalation" && t.priority == 100)
718        );
719        assert_eq!(transitions[0].to, "escalation");
720    }
721
722    #[test]
723    fn test_get_parent_definition() {
724        let config = create_hierarchical_config();
725        let sm = StateMachine::new(config).unwrap();
726
727        let parent = sm.get_parent_definition();
728        assert!(parent.is_some());
729        assert_eq!(parent.unwrap().prompt, Some("Problem solving".into()));
730
731        sm.transition_to("^closing", "done").unwrap();
732        assert!(sm.get_parent_definition().is_none());
733    }
734}
ai_agents_state/machine.rs

ai_agents_state/
machine.rs
