selfware 0.2.2

#[derive(Debug, Clone)]
pub enum AgentState {
    Planning,
    Executing { step: usize },
    ErrorRecovery { error: String },
    Completed,
    Failed { reason: String },
}

pub struct AgentLoop {
    state: AgentState,
    max_iterations: usize,
    current_step: usize,
    iteration: usize,
}

impl AgentLoop {
    pub fn new(max_iterations: usize) -> Self {
        Self {
            state: AgentState::Planning,
            max_iterations,
            current_step: 0,
            iteration: 0,
        }
    }

    pub fn next_state(&mut self) -> Option<AgentState> {
        if self.iteration >= self.max_iterations {
            return Some(AgentState::Failed {
                reason: "Max iterations exceeded".to_string(),
            });
        }
        self.iteration += 1;
        Some(self.state.clone())
    }

    /// Returns a warning message when the loop is approaching the iteration
    /// limit.  Intended to be injected as a system message so the LLM can
    /// wrap up gracefully instead of being cut off abruptly.
    ///
    /// Returns `Some` at 80% and 90%+ of `max_iterations`, `None` otherwise.
    pub fn approaching_limit_warning(&self) -> Option<String> {
        if self.max_iterations == 0 {
            return None;
        }
        let remaining = self.max_iterations.saturating_sub(self.iteration);
        let pct_used = (self.iteration * 100) / self.max_iterations;

        if pct_used >= 90 {
            Some(format!(
                "[SYSTEM] Only {} iteration(s) remaining out of {}. Wrap up your current work and provide a final answer now.",
                remaining, self.max_iterations
            ))
        } else if pct_used >= 80 {
            Some(format!(
                "[SYSTEM] Approaching iteration limit: {} of {} iterations used ({} remaining). Start wrapping up.",
                self.iteration, self.max_iterations, remaining
            ))
        } else {
            None
        }
    }

    pub fn set_state(&mut self, state: AgentState) {
        self.state = state;
    }

    pub fn increment_step(&mut self) {
        self.current_step += 1;
        self.state = AgentState::Executing {
            step: self.current_step,
        };
    }

    pub fn current_step(&self) -> usize {
        self.current_step
    }

    pub fn current_iteration(&self) -> usize {
        self.iteration
    }

    pub fn current_state_label(&self) -> &'static str {
        match self.state {
            AgentState::Planning => "planning",
            AgentState::Executing { .. } => "executing",
            AgentState::ErrorRecovery { .. } => "error_recovery",
            AgentState::Completed => "completed",
            AgentState::Failed { .. } => "failed",
        }
    }

    /// Restore loop progress from persisted state.
    pub fn restore_progress(&mut self, step: usize, iteration: usize) {
        self.current_step = step;
        self.iteration = iteration;
        self.state = AgentState::Executing { step };
    }

    /// Reset loop state for a new task, preserving max_iterations.
    ///
    /// Without this, queued tasks share the iteration counter from the previous
    /// task and may hit the max-iterations limit prematurely.
    pub fn reset_for_task(&mut self) {
        self.state = AgentState::Planning;
        self.current_step = 0;
        self.iteration = 0;
    }
}

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

    #[test]
    fn test_agent_loop_new() {
        let loop_ctrl = AgentLoop::new(100);
        assert_eq!(loop_ctrl.max_iterations, 100);
        assert_eq!(loop_ctrl.current_step, 0);
        assert_eq!(loop_ctrl.iteration, 0);
    }

    #[test]
    fn test_agent_loop_initial_state_is_planning() {
        let mut loop_ctrl = AgentLoop::new(100);
        let state = loop_ctrl.next_state();
        assert!(matches!(state, Some(AgentState::Planning)));
    }

    #[test]
    fn test_agent_loop_set_state() {
        let mut loop_ctrl = AgentLoop::new(100);
        loop_ctrl.set_state(AgentState::Executing { step: 0 });
        let state = loop_ctrl.next_state();
        assert!(matches!(state, Some(AgentState::Executing { step: 0 })));
    }

    #[test]
    fn test_agent_loop_increment_step() {
        let mut loop_ctrl = AgentLoop::new(100);
        assert_eq!(loop_ctrl.current_step(), 0);
        loop_ctrl.increment_step();
        assert_eq!(loop_ctrl.current_step(), 1);
        loop_ctrl.increment_step();
        assert_eq!(loop_ctrl.current_step(), 2);
    }

    #[test]
    fn test_agent_loop_max_iterations_exceeded() {
        let mut loop_ctrl = AgentLoop::new(3);

        // First 3 iterations should work
        assert!(loop_ctrl.next_state().is_some());
        assert!(loop_ctrl.next_state().is_some());
        assert!(loop_ctrl.next_state().is_some());

        // 4th should fail
        let state = loop_ctrl.next_state();
        assert!(
            matches!(state, Some(AgentState::Failed { reason }) if reason == "Max iterations exceeded")
        );
    }

    #[test]
    fn test_agent_state_error_recovery() {
        let mut loop_ctrl = AgentLoop::new(100);
        loop_ctrl.set_state(AgentState::ErrorRecovery {
            error: "Test error".to_string(),
        });

        let state = loop_ctrl.next_state();
        match state {
            Some(AgentState::ErrorRecovery { error }) => {
                assert_eq!(error, "Test error");
            }
            _ => panic!("Expected ErrorRecovery state"),
        }
    }

    #[test]
    fn test_agent_state_failed() {
        let mut loop_ctrl = AgentLoop::new(100);
        loop_ctrl.set_state(AgentState::Failed {
            reason: "Something went wrong".to_string(),
        });

        let state = loop_ctrl.next_state();
        match state {
            Some(AgentState::Failed { reason }) => {
                assert_eq!(reason, "Something went wrong");
            }
            _ => panic!("Expected Failed state"),
        }
    }

    #[test]
    fn test_executing_state_tracks_step() {
        let state = AgentState::Executing { step: 5 };
        match state {
            AgentState::Executing { step } => assert_eq!(step, 5),
            _ => panic!("Expected Executing state"),
        }
    }

    #[test]
    fn test_increment_step_updates_state() {
        let mut loop_ctrl = AgentLoop::new(100);
        loop_ctrl.increment_step();

        // After increment, state should be Executing with current step
        match &loop_ctrl.state {
            AgentState::Executing { step } => assert_eq!(*step, 1),
            _ => panic!("Expected Executing state after increment"),
        }
    }

    #[test]
    fn test_reset_for_task() {
        let mut loop_ctrl = AgentLoop::new(10);

        // Simulate several iterations
        loop_ctrl.next_state();
        loop_ctrl.next_state();
        loop_ctrl.next_state();
        loop_ctrl.increment_step();
        loop_ctrl.increment_step();

        assert_eq!(loop_ctrl.iteration, 3);
        assert_eq!(loop_ctrl.current_step(), 2);

        // Reset for a new task
        loop_ctrl.reset_for_task();
        assert_eq!(loop_ctrl.iteration, 0);
        assert_eq!(loop_ctrl.current_step(), 0);
        assert!(matches!(loop_ctrl.state, AgentState::Planning));

        // Should be able to iterate fully again
        for _ in 0..10 {
            let state = loop_ctrl.next_state();
            assert!(!matches!(state, Some(AgentState::Failed { .. })));
        }
        // 11th should fail
        let state = loop_ctrl.next_state();
        assert!(matches!(state, Some(AgentState::Failed { .. })));
    }

    #[test]
    fn test_restore_progress() {
        let mut loop_ctrl = AgentLoop::new(10);
        loop_ctrl.restore_progress(3, 7);

        assert_eq!(loop_ctrl.current_step(), 3);
        assert_eq!(loop_ctrl.current_iteration(), 7);
        assert!(matches!(loop_ctrl.state, AgentState::Executing { step: 3 }));
    }

    #[test]
    fn test_approaching_limit_warning_none_early() {
        let mut loop_ctrl = AgentLoop::new(100);
        // Advance to 50% — no warning
        for _ in 0..50 {
            loop_ctrl.next_state();
        }
        assert!(loop_ctrl.approaching_limit_warning().is_none());
    }

    #[test]
    fn test_approaching_limit_warning_at_80_pct() {
        let mut loop_ctrl = AgentLoop::new(100);
        for _ in 0..80 {
            loop_ctrl.next_state();
        }
        let warning = loop_ctrl.approaching_limit_warning();
        assert!(warning.is_some());
        assert!(warning.unwrap().contains("wrapping up"));
    }

    #[test]
    fn test_approaching_limit_warning_at_90_pct() {
        let mut loop_ctrl = AgentLoop::new(100);
        for _ in 0..90 {
            loop_ctrl.next_state();
        }
        let warning = loop_ctrl.approaching_limit_warning();
        assert!(warning.is_some());
        assert!(warning.unwrap().contains("final answer"));
    }

    #[test]
    fn test_approaching_limit_warning_zero_max() {
        let loop_ctrl = AgentLoop::new(0);
        assert!(loop_ctrl.approaching_limit_warning().is_none());
    }
}