Skip to main content

ai_agent/utils/
idle_timeout.rs

1//! Idle timeout management utilities
2//!
3//! Creates an idle timeout manager for SDK mode.
4//! Automatically exits the process after the specified idle duration.
5
6use crate::constants::env::ai_code;
7use std::sync::{Arc, Mutex};
8use std::time::Duration;
9
10/// Creates an idle timeout manager for SDK mode.
11///
12/// # Arguments
13/// * `is_idle` - Function that returns true if the system is currently idle
14///
15/// # Returns
16/// Object with start/stop methods to control the idle timer
17pub fn create_idle_timeout_manager<F>(is_idle: F) -> IdleTimeoutManager
18where
19    F: Fn() -> bool + Send + Sync + 'static,
20{
21    let exit_after_stop_delay = std::env::var(ai_code::EXIT_AFTER_STOP_DELAY)
22        .ok()
23        .and_then(|v| v.parse::<u64>().ok());
24
25    let delay_ms = exit_after_stop_delay.filter(|&d| d > 0);
26
27    IdleTimeoutManager {
28        is_idle: Arc::new(is_idle),
29        delay_ms,
30        timer: Mutex::new(None),
31    }
32}
33
34/// Idle timeout manager
35pub struct IdleTimeoutManager {
36    is_idle: Arc<dyn Fn() -> bool + Send + Sync + 'static>,
37    delay_ms: Option<u64>,
38    timer: Mutex<Option<tokio::task::JoinHandle<()>>>,
39}
40
41impl IdleTimeoutManager {
42    /// Start the idle timer
43    pub fn start(&self) {
44        let delay_ms = match self.delay_ms {
45            Some(d) => d,
46            None => return, // No delay configured
47        };
48
49        // Clear any existing timer
50        self.stop();
51
52        let is_idle_fn = Arc::clone(&self.is_idle);
53        let start_time = std::time::SystemTime::now()
54            .duration_since(std::time::UNIX_EPOCH)
55            .unwrap()
56            .as_millis() as u64;
57
58        let timer = tokio::spawn(async move {
59            tokio::time::sleep(Duration::from_millis(delay_ms)).await;
60
61            let idle_duration = std::time::SystemTime::now()
62                .duration_since(std::time::UNIX_EPOCH)
63                .unwrap()
64                .as_millis() as u64
65                - start_time;
66
67            if is_idle_fn() && idle_duration >= delay_ms {
68                log::info!("Exiting after {}ms of idle time", delay_ms);
69                std::process::exit(0);
70            }
71        });
72
73        *self.timer.lock().unwrap() = Some(timer);
74    }
75
76    /// Stop the idle timer
77    pub fn stop(&self) {
78        if let Some(timer) = self.timer.lock().unwrap().take() {
79            timer.abort();
80        }
81    }
82}
83
84impl Drop for IdleTimeoutManager {
85    fn drop(&mut self) {
86        self.stop();
87    }
88}