Codex Handbook
tui/src/chatwidget/turn_lifecycle.rs 97 lines
//! Agent-turn lifecycle state for `ChatWidget`.use std::collections::HashSet;use std::time::Instant;use codex_utils_sleep_inhibitor::SleepInhibitor;#[derive(Debug)]pub(super) struct TurnLifecycleState {    pub(super) sleep_inhibitor: SleepInhibitor,    /// Tracks whether codex-core currently considers an agent turn to be in progress.    pub(super) agent_turn_running: bool,    pub(super) last_turn_id: Option<String>,    pub(super) budget_limited_turn_ids: HashSet<String>,    pub(super) goal_status_active_turn_started_at: Option<Instant>,}impl TurnLifecycleState {    pub(super) fn new(prevent_idle_sleep: bool) -> Self {        Self {            sleep_inhibitor: SleepInhibitor::new(prevent_idle_sleep),            agent_turn_running: false,            last_turn_id: None,            budget_limited_turn_ids: HashSet::new(),            goal_status_active_turn_started_at: None,        }    }    pub(super) fn start(&mut self, now: Instant) {        self.agent_turn_running = true;        self.goal_status_active_turn_started_at = Some(now);        self.sleep_inhibitor.set_turn_running(/*turn_running*/ true);    }    pub(super) fn finish(&mut self) {        self.agent_turn_running = false;        self.goal_status_active_turn_started_at = None;        self.sleep_inhibitor            .set_turn_running(/*turn_running*/ false);    }    pub(super) fn restore_running(&mut self, running: bool, now: Instant) {        self.agent_turn_running = running;        self.goal_status_active_turn_started_at = running.then_some(now);        self.sleep_inhibitor.set_turn_running(running);    }    pub(super) fn reset_thread(&mut self) {        self.finish();        self.last_turn_id = None;        self.budget_limited_turn_ids.clear();    }    pub(super) fn set_prevent_idle_sleep(&mut self, enabled: bool) {        self.sleep_inhibitor = SleepInhibitor::new(enabled);        self.sleep_inhibitor            .set_turn_running(self.agent_turn_running);    }    pub(super) fn mark_budget_limited(&mut self, turn_id: String) {        self.budget_limited_turn_ids.insert(turn_id);    }    pub(super) fn take_budget_limited(&mut self, turn_id: &str) -> bool {        self.budget_limited_turn_ids.remove(turn_id)    }}#[cfg(test)]mod tests {    use super::*;    #[test]    fn start_and_finish_update_running_state() {        let mut state = TurnLifecycleState::new(/*prevent_idle_sleep*/ false);        state.start(Instant::now());        assert!(state.agent_turn_running);        assert!(state.goal_status_active_turn_started_at.is_some());        assert!(state.sleep_inhibitor.is_turn_running());        state.finish();        assert!(!state.agent_turn_running);        assert!(state.goal_status_active_turn_started_at.is_none());        assert!(!state.sleep_inhibitor.is_turn_running());    }    #[test]    fn budget_limited_turn_ids_are_consumed() {        let mut state = TurnLifecycleState::new(/*prevent_idle_sleep*/ false);        state.mark_budget_limited("turn-1".to_string());        assert!(state.take_budget_limited("turn-1"));        assert!(!state.take_budget_limited("turn-1"));    }}