roboticus-agent 0.11.4

use std::collections::VecDeque;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReactState {
    Thinking,
    Acting,
    Observing,
    Persisting,
    Idle,
    Done,
}

#[derive(Debug, Clone)]
pub enum ReactAction {
    Think,
    Act { tool_name: String, params: String },
    Observe,
    Persist,
    NoOp,
    Finish,
}

const IDLE_THRESHOLD: usize = 3;
const LOOP_DETECTION_WINDOW: usize = 3;

pub struct AgentLoop {
    pub state: ReactState,
    pub turn_count: usize,
    pub max_turns: usize,
    idle_count: usize,
    recent_calls: VecDeque<(String, String)>,
    last_failed_call: Option<(String, String)>,
    last_error_message: Option<String>,
    suppressed_count: u8,
}

impl AgentLoop {
    pub fn new(max_turns: usize) -> Self {
        Self {
            state: ReactState::Idle,
            turn_count: 0,
            max_turns,
            idle_count: 0,
            recent_calls: VecDeque::with_capacity(LOOP_DETECTION_WINDOW + 1),
            last_failed_call: None,
            last_error_message: None,
            suppressed_count: 0,
        }
    }

    pub fn transition(&mut self, action: ReactAction) -> ReactState {
        match action {
            ReactAction::Think => {
                // Only count logical turns (Think phase starts a new turn).
                // Previously every transition incremented, inflating count 2-3x.
                self.turn_count += 1;
                if self.turn_count > self.max_turns {
                    self.state = ReactState::Done;
                    return self.state;
                }
                self.idle_count = 0;
                self.state = ReactState::Thinking;
            }
            ReactAction::Act { tool_name, params } => {
                self.idle_count = 0;
                // Evaluate against prior calls only; this avoids counting the
                // current call inside the detection window.
                if self.is_looping(&tool_name, &params) {
                    tracing::warn!(tool = %tool_name, "agent loop detected, forcing Done");
                    self.state = ReactState::Done;
                } else {
                    self.state = ReactState::Acting;
                }
                self.recent_calls
                    .push_back((tool_name.clone(), params.clone()));
                if self.recent_calls.len() > LOOP_DETECTION_WINDOW {
                    self.recent_calls.pop_front();
                }
            }
            ReactAction::Observe => {
                self.idle_count = 0;
                self.state = ReactState::Observing;
            }
            ReactAction::Persist => {
                self.idle_count = 0;
                self.state = ReactState::Persisting;
            }
            ReactAction::NoOp => {
                self.idle_count += 1;
                if self.idle_count >= IDLE_THRESHOLD {
                    self.state = ReactState::Idle;
                }
            }
            ReactAction::Finish => {
                self.state = ReactState::Done;
            }
        }

        self.state
    }

    pub fn is_idle(&self) -> bool {
        self.idle_count >= IDLE_THRESHOLD
    }

    /// Returns true if the same tool+params combination has appeared
    /// `LOOP_DETECTION_WINDOW` consecutive times.
    pub fn is_looping(&self, tool_name: &str, params: &str) -> bool {
        if self.recent_calls.len() < LOOP_DETECTION_WINDOW {
            return false;
        }

        self.recent_calls
            .iter()
            .all(|(t, p)| t == tool_name && p == params)
    }

    /// Record a tool call failure so subsequent identical calls can be suppressed.
    pub fn record_tool_error(&mut self, tool: &str, params: &str, error: &str) {
        self.last_failed_call = Some((tool.to_string(), params.to_string()));
        self.last_error_message = Some(error.to_string());
        self.suppressed_count = 0;
    }

    /// Returns true if the given tool+params matches the most recent failure.
    pub fn should_suppress_duplicate(&self, tool: &str, params: &str) -> bool {
        self.last_failed_call
            .as_ref()
            .map(|(t, p)| t == tool && p == params)
            .unwrap_or(false)
    }

    pub fn increment_suppressed(&mut self) {
        self.suppressed_count += 1;
    }

    /// After 2 suppressions the model is stuck — caller should abort the batch.
    pub fn should_abort_error_loop(&self) -> bool {
        self.suppressed_count >= 2
    }

    pub fn last_error(&self) -> Option<&str> {
        self.last_error_message.as_deref()
    }

    /// Clear error dedup state after a successful tool call.
    pub fn clear_error_state(&mut self) {
        self.last_failed_call = None;
        self.last_error_message = None;
        self.suppressed_count = 0;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn state_transitions() {
        let mut agent = AgentLoop::new(100);
        assert_eq!(agent.state, ReactState::Idle);

        let s = agent.transition(ReactAction::Think);
        assert_eq!(s, ReactState::Thinking);

        let s = agent.transition(ReactAction::Act {
            tool_name: "echo".into(),
            params: "{}".into(),
        });
        assert_eq!(s, ReactState::Acting);

        let s = agent.transition(ReactAction::Observe);
        assert_eq!(s, ReactState::Observing);

        let s = agent.transition(ReactAction::Persist);
        assert_eq!(s, ReactState::Persisting);

        let s = agent.transition(ReactAction::Finish);
        assert_eq!(s, ReactState::Done);
    }

    #[test]
    fn idle_detection() {
        let mut agent = AgentLoop::new(100);

        assert!(!agent.is_idle());
        agent.transition(ReactAction::NoOp);
        assert!(!agent.is_idle());
        agent.transition(ReactAction::NoOp);
        assert!(!agent.is_idle());
        agent.transition(ReactAction::NoOp);
        assert!(agent.is_idle());
        assert_eq!(agent.state, ReactState::Idle);

        agent.transition(ReactAction::Think);
        assert!(!agent.is_idle());
    }

    #[test]
    fn loop_detection() {
        let mut agent = AgentLoop::new(100);

        for _ in 0..3 {
            let s = agent.transition(ReactAction::Act {
                tool_name: "echo".into(),
                params: r#"{"msg":"hi"}"#.into(),
            });
            assert_eq!(s, ReactState::Acting);
        }

        assert!(agent.is_looping("echo", r#"{"msg":"hi"}"#));
        assert!(!agent.is_looping("echo", r#"{"msg":"bye"}"#));
        assert!(!agent.is_looping("other", r#"{"msg":"hi"}"#));

        let s = agent.transition(ReactAction::Act {
            tool_name: "echo".into(),
            params: r#"{"msg":"hi"}"#.into(),
        });
        assert_eq!(s, ReactState::Done);

        agent.transition(ReactAction::Act {
            tool_name: "read".into(),
            params: "{}".into(),
        });
        assert!(!agent.is_looping("echo", r#"{"msg":"hi"}"#));
    }

    #[test]
    fn agent_loop_detects_repeated_failed_tool_call() {
        let mut loop_state = AgentLoop::new(10);
        loop_state.record_tool_error(
            "compose-subagent",
            r#"{"name":"foo"}"#,
            "no skills available",
        );
        assert!(loop_state.should_suppress_duplicate("compose-subagent", r#"{"name":"foo"}"#));
        assert!(!loop_state.should_suppress_duplicate("compose-subagent", r#"{"name":"bar"}"#));
        assert!(!loop_state.should_suppress_duplicate("compose-skill", r#"{"name":"foo"}"#));
    }

    #[test]
    fn agent_loop_abort_after_repeated_suppression() {
        let mut loop_state = AgentLoop::new(10);
        loop_state.record_tool_error("compose-subagent", r#"{"name":"foo"}"#, "no skills");
        assert!(loop_state.should_suppress_duplicate("compose-subagent", r#"{"name":"foo"}"#));
        loop_state.increment_suppressed();
        assert!(loop_state.should_suppress_duplicate("compose-subagent", r#"{"name":"foo"}"#));
        loop_state.increment_suppressed();
        assert!(loop_state.should_abort_error_loop());
    }

    #[test]
    fn max_turns_forces_done() {
        let mut agent = AgentLoop::new(2);

        agent.transition(ReactAction::Think);
        // Non-Think transitions don't inflate the turn count
        agent.transition(ReactAction::Act {
            tool_name: "echo".into(),
            params: "{}".into(),
        });
        agent.transition(ReactAction::Observe);
        assert_eq!(agent.turn_count, 1);

        agent.transition(ReactAction::Think);
        assert_eq!(agent.turn_count, 2);

        // Third Think exceeds max_turns=2
        let s = agent.transition(ReactAction::Think);
        assert_eq!(s, ReactState::Done);
    }
}