ai_agent/utils/
idle_timeout.rs1use crate::constants::env::ai_code;
7use std::sync::{Arc, Mutex};
8use std::time::Duration;
9
10pub 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
34pub 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 pub fn start(&self) {
44 let delay_ms = match self.delay_ms {
45 Some(d) => d,
46 None => return, };
48
49 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 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}