use std::{ collections::{BTreeMap, BTreeSet}, process::Command, thread, time::Duration, }; use anyhow::{Context, Result, anyhow}; use chrono::{DateTime, Datelike, FixedOffset, SecondsFormat, TimeZone, Utc}; use clap::Parser; use reqwest::header::{CONNECTION, HeaderMap, HeaderValue}; use reqwest::{ Method, StatusCode, blocking::{Client, Response}, }; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; const DEFAULT_AW_URL: &str = "http://127.0.0.1:5600"; const DEFAULT_HOST: &str = "HOST-EXAMPLE"; const AUTOHEAL_SOURCE: &str = "aw-worktime-autoheal"; #[derive(Debug, Parser)] #[command(about = "AW Worktime autoheal")] struct Cli { #[arg(long)] dry_run: bool, #[arg(long)] json: bool, } #[derive(Debug, Clone)] struct Config { aw_url: String, host: String, worktime_health_url: String, worktime_report_timeout_seconds: f64, management_warm_enabled: bool, management_warm_url: String, management_warm_timeout_seconds: f64, today_probe_enabled: bool, today_probe_url: String, today_probe_timeout_seconds: f64, session_freshness_seconds: f64, aw_timeout_seconds: f64, aw_post_chunk_size: usize, } #[derive(Debug, Clone, Deserialize, Serialize)] struct AwEvent { timestamp: Option, #[serde(default)] duration: Option, #[serde(default)] data: Value, } #[derive(Debug, Serialize)] struct RunSummary { ok: bool, dry_run: bool, host: String, worktime_ok: bool, api_restarted: bool, management_warm_ok: Option, need_heal: bool, ui_bridge_started: bool, reset_buckets: bool, backfill_afk: usize, backfill_window: usize, reason: String, } struct AwClient { base_url: String, client: Client, } fn env(name: &str, fallback: &str) -> String { std::env::var(name) .ok() .map(|value| value.trim().to_string()) .filter(|value| !value.is_empty()) .unwrap_or_else(|| fallback.to_string()) } fn env_bool_legacy(name: &str, fallback: bool) -> bool { match std::env::var(name) .ok() .map(|value| value.trim().to_ascii_lowercase()) { Some(value) if !value.is_empty() => matches!(value.as_str(), "1" | "true" | "yes" | "on"), _ => fallback, } } fn env_f64(name: &str, fallback: f64) -> f64 { std::env::var(name) .ok() .and_then(|value| value.trim().parse().ok()) .filter(|value| *value > 0.0) .unwrap_or(fallback) } fn env_usize(name: &str, fallback: usize) -> usize { std::env::var(name) .ok() .and_then(|value| value.trim().parse().ok()) .filter(|value| *value > 0) .unwrap_or(fallback) } fn load_config() -> Config { Config { aw_url: env("AW_URL", DEFAULT_AW_URL) .trim_end_matches('/') .to_string(), host: env("AW_WORKTIME_HOST", DEFAULT_HOST), worktime_health_url: env("WORKTIME_HEALTH_URL", "http://127.0.0.1:5610/health"), worktime_report_timeout_seconds: env_f64("WORKTIME_REPORT_TIMEOUT_SECONDS", 20.0), management_warm_enabled: env_bool_legacy("WORKTIME_MANAGEMENT_WARM_ENABLED", true), management_warm_url: env( "WORKTIME_MANAGEMENT_WARM_URL", "http://127.0.0.1:5610/reports/worktime/management?day=today&format=json", ), management_warm_timeout_seconds: env_f64("WORKTIME_MANAGEMENT_WARM_TIMEOUT_SECONDS", 60.0), today_probe_enabled: env_bool_legacy("WORKTIME_TODAY_PROBE_ENABLED", true), today_probe_url: env( "WORKTIME_TODAY_PROBE_URL", "http://127.0.0.1:5610/reports/worktime/today?day=today&format=json", ), today_probe_timeout_seconds: env_f64("WORKTIME_TODAY_PROBE_TIMEOUT_SECONDS", 20.0), session_freshness_seconds: env_f64("WORKTIME_SESSION_FRESHNESS_SECONDS", 600.0), aw_timeout_seconds: env_f64("WORKTIME_AUTOHEAL_AW_TIMEOUT_SECONDS", 30.0), aw_post_chunk_size: env_usize("WORKTIME_AUTOHEAL_AW_POST_CHUNK_SIZE", 500), } } impl Config { fn window_bucket(&self) -> String { format!("aw-rdp-window_{}", self.host) } fn afk_bucket(&self) -> String { format!("aw-rdp-afk_{}", self.host) } fn session_bucket(&self) -> String { format!("aw-worktime-sessions_{}", self.host) } } impl AwClient { fn new(config: &Config) -> Result { let mut headers = HeaderMap::new(); headers.insert(CONNECTION, HeaderValue::from_static("close")); let client = Client::builder() .timeout(Duration::from_secs_f64( config.aw_timeout_seconds.max(0.001), )) .no_proxy() .pool_max_idle_per_host(0) .default_headers(headers) .build() .context("build ActivityWatch HTTP client")?; Ok(Self { base_url: config.aw_url.clone(), client, }) } fn url(&self, path: &str) -> String { format!("{}{}", self.base_url, path) } fn send_retry(&self, method: Method, path: &str, payload: Option) -> Result { let mut last_error = None; for attempt in 1..=3 { let mut request = self.client.request(method.clone(), self.url(path)); if let Some(payload) = payload.clone() { request = request.json(&payload); } match request.send() { Ok(response) => return Ok(response), Err(error) => { last_error = Some(error); if attempt < 3 { log(&format!("warn aw request={path} retry_after_send_error")); thread::sleep(Duration::from_millis(750)); } } } } Err(last_error.expect("error exists after failed request")) .with_context(|| format!("request ActivityWatch {path}")) } fn request_json( &self, method: Method, path: &str, payload: Option, ignore_not_found: bool, ) -> Result> { let response = self.send_retry(method, path, payload)?; let status = response.status(); if ignore_not_found && status == StatusCode::NOT_FOUND { return Ok(None); } if !status.is_success() { return Err(anyhow!("ActivityWatch {path} returned HTTP {status}")); } let bytes = response.bytes().context("read ActivityWatch response")?; if bytes.is_empty() { Ok(None) } else { serde_json::from_slice(&bytes) .map(Some) .with_context(|| format!("decode ActivityWatch JSON from {path}")) } } fn get_events(&self, bucket_id: &str, limit: usize) -> Result> { let path = format!("/api/0/buckets/{bucket_id}/events?limit={limit}"); match self.request_json(Method::GET, &path, None, true)? { Some(value) => serde_json::from_value(value).context("decode AW events"), None => Ok(Vec::new()), } } fn delete_bucket(&self, bucket_id: &str) -> Result<()> { let path = format!("/api/0/buckets/{bucket_id}"); let response = self.send_retry(Method::DELETE, &path, None)?; if response.status().is_success() || response.status() == StatusCode::NOT_FOUND { Ok(()) } else { Err(anyhow!( "delete bucket {bucket_id} returned HTTP {}", response.status() )) } } fn ensure_bucket( &self, bucket_id: &str, event_type: &str, client_name: &str, host: &str, ) -> Result<()> { let path = format!("/api/0/buckets/{bucket_id}"); let payload = json!({ "client": client_name, "type": event_type, "hostname": host, }); let response = self.send_retry(Method::POST, &path, Some(payload))?; let status = response.status(); if status.is_success() || status == StatusCode::NOT_MODIFIED || status == StatusCode::CONFLICT { Ok(()) } else { Err(anyhow!("ensure bucket {bucket_id} returned HTTP {status}")) } } fn post_events_chunked( &self, bucket_id: &str, events: &[AwEvent], chunk_size: usize, ) -> Result<()> { for chunk in events.chunks(chunk_size.max(1)) { let path = format!("/api/0/buckets/{bucket_id}/events"); self.request_json(Method::POST, &path, Some(json!(chunk)), false)?; } Ok(()) } } fn log(message: &str) { println!("{} {}", chrono::Local::now().format("%F %T"), message); } fn build_probe_client(timeout_seconds: f64) -> Result { Client::builder() .timeout(Duration::from_secs_f64(timeout_seconds.max(0.001))) .no_proxy() .build() .context("build probe HTTP client") } fn probe_url(url: &str, timeout_seconds: f64) -> bool { let Ok(client) = build_probe_client(timeout_seconds) else { return false; }; client .get(url) .send() .map(|response| response.status().is_success()) .unwrap_or(false) } fn probe_reports(config: &Config) -> bool { if !probe_url( &config.worktime_health_url, config.worktime_report_timeout_seconds, ) { return false; } if config.today_probe_enabled && !probe_url(&config.today_probe_url, config.today_probe_timeout_seconds) { return false; } true } fn run_systemctl(args: &[&str], dry_run: bool) -> bool { if dry_run { log(&format!("dry-run systemctl {}", args.join(" "))); return true; } Command::new("systemctl") .args(args) .status() .map(|status| status.success()) .unwrap_or(false) } fn parse_iso_utc(ts: &str) -> Result> { DateTime::parse_from_rfc3339(ts) .map(|parsed| parsed.with_timezone(&Utc)) .with_context(|| format!("parse timestamp {ts}")) } fn to_iso_utc(dt: DateTime) -> String { dt.to_rfc3339_opts(SecondsFormat::Secs, true) } fn moscow_day_start_utc() -> DateTime { let msk = FixedOffset::east_opt(3 * 3600).expect("valid Moscow offset"); let today = Utc::now().with_timezone(&msk); msk.with_ymd_and_hms(today.year(), today.month(), today.day(), 0, 0, 0) .single() .expect("valid day start") .with_timezone(&Utc) } fn value_string(data: &Value, key: &str) -> String { match data.get(key) { Some(Value::String(value)) => value.trim().to_string(), Some(Value::Number(value)) => value.to_string(), Some(Value::Bool(value)) => value.to_string(), Some(Value::Null) | None => String::new(), Some(value) => value.to_string().trim_matches('"').trim().to_string(), } } fn value_i64(data: &Value, key: &str) -> Option { match data.get(key) { Some(Value::Number(value)) => value.as_i64(), Some(Value::String(value)) => value.trim().parse().ok(), _ => None, } } fn is_session_active(data: &Value) -> bool { if data.get("active").and_then(Value::as_bool).unwrap_or(false) { return true; } let state = value_string(data, "state").to_lowercase(); if state == "active" || state == "активно" { return true; } if state == "unknown" { let sid = value_i64(data, "sessionId").unwrap_or(-1); let user = value_string(data, "username").to_lowercase(); let session_name = value_string(data, "sessionName").to_lowercase(); return sid > 0 && !user.is_empty() && !user.ends_with('$') && (session_name.starts_with("rdp-") || session_name == "console"); } false } fn active_users_at_latest_session(events: &[AwEvent]) -> (Option>, BTreeSet) { let mut latest_ts = None; let mut active_users = BTreeSet::new(); for event in events { let Some(ts) = event.timestamp.as_deref() else { continue; }; let Ok(cur) = parse_iso_utc(ts) else { continue; }; if latest_ts.is_none_or(|latest| cur > latest) { latest_ts = Some(cur); active_users.clear(); } if latest_ts == Some(cur) { let user = value_string(&event.data, "username"); if !user.is_empty() && is_session_active(&event.data) { active_users.insert(user); } } } (latest_ts, active_users) } fn active_window_duration_today(window_events: &[AwEvent], start: DateTime) -> f64 { window_events .iter() .filter(|event| { event .timestamp .as_deref() .and_then(|ts| parse_iso_utc(ts).ok()) .is_some_and(|ts| ts >= start) }) .filter(|event| { value_string(&event.data, "title") .to_lowercase() .contains("rdp active") }) .map(|event| event.duration.unwrap_or(0.0).max(0.0)) .sum() } fn should_heal( aw: &AwClient, config: &Config, now: DateTime, start: DateTime, ) -> Result<(bool, String)> { let window_events = match aw.get_events(&config.window_bucket(), 12_000) { Ok(events) => events, Err(error) => { return Ok((true, format!("window_read_failed: {error:#}"))); } }; let session_events = match aw.get_events(&config.session_bucket(), 12_000) { Ok(events) => events, Err(error) => { return Ok((true, format!("session_read_failed: {error:#}"))); } }; let (latest_ts, active_users) = active_users_at_latest_session(&session_events); let Some(latest_ts) = latest_ts else { return Ok((false, "no_sessions".to_string())); }; let age = (now - latest_ts).num_milliseconds() as f64 / 1000.0; if age > config.session_freshness_seconds { return Ok((false, format!("sessions_stale age_seconds={age:.0}"))); } if active_users.is_empty() { return Ok((false, "no_active_users".to_string())); } let active = active_window_duration_today(&window_events, start); if active <= 0.0 { Ok(( true, format!("zero_activity active_users={}", active_users.len()), )) } else { Ok(( false, format!("activity_present active_seconds={active:.0}"), )) } } fn build_window_title(users: &[String]) -> String { if users.is_empty() { "RDP idle".to_string() } else { format!("RDP active ({}): {}", users.len(), users.join(", ")) } } fn build_backfill( session_events: Vec, start: DateTime, ) -> (Vec, Vec) { let mut grouped: BTreeMap, Vec> = BTreeMap::new(); for event in session_events { let Some(ts) = event.timestamp.as_deref() else { continue; }; let Ok(parsed) = parse_iso_utc(ts) else { continue; }; if parsed >= start { grouped.entry(parsed).or_default().push(event); } } let keys: Vec> = grouped.keys().copied().collect(); let mut out_afk = Vec::new(); let mut out_win = Vec::new(); for (idx, ts) in keys.iter().enumerate() { let rows = &grouped[ts]; let mut duration = keys .get(idx + 1) .map(|next| ((*next - *ts).num_milliseconds() as f64 / 1000.0).max(0.0)) .unwrap_or(10.0); if duration <= 0.0 { duration = 10.0; } duration = duration.min(30.0); let users: Vec = rows .iter() .filter_map(|event| { let user = value_string(&event.data, "username"); if !user.is_empty() && is_session_active(&event.data) { Some(user) } else { None } }) .collect::>() .into_iter() .collect(); let active = !users.is_empty(); let ts = to_iso_utc(*ts); out_afk.push(AwEvent { timestamp: Some(ts.clone()), duration: Some(duration), data: json!({ "status": if active { "not-afk" } else { "afk" }, "source": AUTOHEAL_SOURCE, }), }); out_win.push(AwEvent { timestamp: Some(ts), duration: Some(duration), data: json!({ "app": "RDP", "title": build_window_title(&users), "source": AUTOHEAL_SOURCE, }), }); } (out_afk, out_win) } fn reset_bucket( aw: &AwClient, bucket_id: &str, event_type: &str, client_name: &str, host: &str, ) -> Result<()> { if let Err(error) = aw.delete_bucket(bucket_id) { log(&format!( "warn delete bucket {bucket_id} ignored: {error:#}" )); } aw.ensure_bucket(bucket_id, event_type, client_name, host) } fn run(cli: &Cli) -> Result { let config = load_config(); let aw = AwClient::new(&config)?; let mut api_restarted = false; let mut worktime_ok = true; if !probe_reports(&config) { log("worktime API probe failed, restarting aw-worktime-api.service"); api_restarted = true; run_systemctl(&["restart", "aw-worktime-api.service"], cli.dry_run); thread::sleep(Duration::from_secs(2)); if !probe_reports(&config) { log("worktime API still degraded after restart"); worktime_ok = false; } else { log("worktime API recovered after restart"); } } if !worktime_ok { log("skip warm/heal because worktime API is still unavailable"); return Ok(RunSummary { ok: true, dry_run: cli.dry_run, host: config.host, worktime_ok, api_restarted, management_warm_ok: None, need_heal: false, ui_bridge_started: false, reset_buckets: false, backfill_afk: 0, backfill_window: 0, reason: "worktime_unavailable".to_string(), }); } let management_warm_ok = if config.management_warm_enabled { let ok = probe_url( &config.management_warm_url, config.management_warm_timeout_seconds, ); if ok { log("management cache warm ok"); } else { log("management cache warm failed"); } Some(ok) } else { None }; let start = moscow_day_start_utc(); let now = Utc::now(); let (need_heal, reason) = should_heal(&aw, &config, now, start)?; if !need_heal { log(&format!( "health ok: {reason} for {}, no action", config.host )); return Ok(RunSummary { ok: true, dry_run: cli.dry_run, host: config.host, worktime_ok, api_restarted, management_warm_ok, need_heal: false, ui_bridge_started: false, reset_buckets: false, backfill_afk: 0, backfill_window: 0, reason, }); } log(&format!( "detected zero activity for {}, running heal", config.host )); let mut ui_bridge_started = run_systemctl(&["restart", "aw-worktime-ui-bridge.timer"], cli.dry_run); ui_bridge_started &= run_systemctl(&["start", "aw-worktime-ui-bridge.service"], cli.dry_run); let session_events = aw.get_events(&config.session_bucket(), 12_000)?; let (afk_events, window_events) = build_backfill(session_events, start); if afk_events.is_empty() || window_events.is_empty() { log(&format!( "heal skipped for {}, no source sessions", config.host )); return Ok(RunSummary { ok: true, dry_run: cli.dry_run, host: config.host, worktime_ok, api_restarted, management_warm_ok, need_heal: true, ui_bridge_started, reset_buckets: false, backfill_afk: 0, backfill_window: 0, reason: "no_source_sessions".to_string(), }); } if !cli.dry_run { reset_bucket( &aw, &config.afk_bucket(), "afkstatus", "aw-worktime-ui-bridge", &config.host, )?; reset_bucket( &aw, &config.window_bucket(), "currentwindow", "aw-worktime-ui-bridge", &config.host, )?; aw.post_events_chunked(&config.afk_bucket(), &afk_events, config.aw_post_chunk_size)?; aw.post_events_chunked( &config.window_bucket(), &window_events, config.aw_post_chunk_size, )?; } println!( "autoheal backfill posted afk={} win={}", afk_events.len(), window_events.len() ); log(&format!("heal completed for {}", config.host)); Ok(RunSummary { ok: true, dry_run: cli.dry_run, host: config.host, worktime_ok, api_restarted, management_warm_ok, need_heal: true, ui_bridge_started, reset_buckets: !cli.dry_run, backfill_afk: afk_events.len(), backfill_window: window_events.len(), reason, }) } fn main() -> Result<()> { let cli = Cli::parse(); let summary = run(&cli)?; if cli.json { println!("{}", serde_json::to_string_pretty(&summary)?); } Ok(()) } #[cfg(test)] mod tests { use super::*; fn event(ts: &str, duration: f64, data: Value) -> AwEvent { AwEvent { timestamp: Some(ts.to_string()), duration: Some(duration), data, } } #[test] fn unknown_rdp_session_counts_as_active() { assert!(is_session_active(&json!({ "sessionId": 5, "state": "Unknown", "username": "user5", "sessionName": "rdp-tcp#0" }))); } #[test] fn active_users_use_latest_sample_only() { let events = vec![ event( "2026-05-27T07:59:25Z", 0.0, json!({"sessionId": 2, "state": "Активно", "username": "admin", "sessionName": "console"}), ), event( "2026-05-27T07:59:30Z", 0.0, json!({"sessionId": 3, "state": "Диск", "username": "old", "sessionName": ""}), ), event( "2026-05-27T07:59:30Z", 0.0, json!({"sessionId": 4, "state": "Активно", "username": "user5", "sessionName": "rdp-tcp#0"}), ), ]; let (latest, users) = active_users_at_latest_session(&events); assert_eq!( latest.map(to_iso_utc).as_deref(), Some("2026-05-27T07:59:30Z") ); assert_eq!(users, BTreeSet::from(["user5".to_string()])); } #[test] fn active_window_duration_counts_today_rdp_active_titles() { let start = parse_iso_utc("2026-06-01T00:00:00Z").unwrap(); let events = vec![ event( "2026-05-31T23:59:59Z", 30.0, json!({"title": "RDP active (1): user"}), ), event( "2026-06-01T00:00:00Z", 10.0, json!({"title": "RDP active (1): user"}), ), event("2026-06-01T00:01:00Z", 20.0, json!({"title": "RDP idle"})), ]; assert_eq!(active_window_duration_today(&events, start), 10.0); } #[test] fn build_backfill_caps_duration_and_marks_active() { let start = parse_iso_utc("2026-06-01T00:00:00Z").unwrap(); let events = vec![ event( "2026-06-01T00:00:00Z", 0.0, json!({"sessionId": 3, "state": "Активно", "username": "user5", "sessionName": "rdp-tcp#0"}), ), event( "2026-06-01T00:00:45Z", 0.0, json!({"sessionId": 3, "state": "Диск", "username": "user5", "sessionName": "rdp-tcp#0"}), ), ]; let (afk, win) = build_backfill(events, start); assert_eq!(afk.len(), 2); assert_eq!(win[0].duration, Some(30.0)); assert_eq!(value_string(&afk[0].data, "status"), "not-afk"); assert_eq!(value_string(&win[0].data, "title"), "RDP active (1): user5"); assert_eq!(value_string(&afk[1].data, "status"), "afk"); } }