microclaw_engine/
run_control.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
3use std::sync::{Arc, LazyLock};
4
5use tokio::sync::{Mutex, Notify};
6
7type RunKey = (String, i64);
8
9#[derive(Clone)]
10struct ActiveRun {
11    run_id: u64,
12    source_message_id: Option<String>,
13    cancelled: Arc<AtomicBool>,
14    notify: Arc<Notify>,
15}
16
17static NEXT_RUN_ID: AtomicU64 = AtomicU64::new(1);
18static ACTIVE_RUNS: LazyLock<Mutex<HashMap<RunKey, Vec<ActiveRun>>>> =
19    LazyLock::new(|| Mutex::new(HashMap::new()));
20static ABORTED_SOURCE_MESSAGE_IDS: LazyLock<Mutex<HashMap<RunKey, HashSet<String>>>> =
21    LazyLock::new(|| Mutex::new(HashMap::new()));
22
23pub const STOPPED_TEXT: &str = "Current run aborted.";
24
25pub async fn register_run(
26    channel: &str,
27    chat_id: i64,
28    source_message_id: Option<String>,
29) -> (u64, Arc<AtomicBool>, Arc<Notify>) {
30    let run_id = NEXT_RUN_ID.fetch_add(1, Ordering::Relaxed);
31    let cancelled = Arc::new(AtomicBool::new(false));
32    let notify = Arc::new(Notify::new());
33    let run = ActiveRun {
34        run_id,
35        source_message_id,
36        cancelled: cancelled.clone(),
37        notify: notify.clone(),
38    };
39    let mut map = ACTIVE_RUNS.lock().await;
40    map.entry((channel.to_string(), chat_id))
41        .or_default()
42        .push(run);
43    (run_id, cancelled, notify)
44}
45
46pub async fn unregister_run(channel: &str, chat_id: i64, run_id: u64) {
47    let mut map = ACTIVE_RUNS.lock().await;
48    let key = (channel.to_string(), chat_id);
49    if let Some(runs) = map.get_mut(&key) {
50        runs.retain(|r| r.run_id != run_id);
51        if runs.is_empty() {
52            map.remove(&key);
53        }
54    }
55}
56
57pub async fn abort_runs(channel: &str, chat_id: i64) -> usize {
58    let key = (channel.to_string(), chat_id);
59    let runs = {
60        let mut map = ACTIVE_RUNS.lock().await;
61        map.remove(&key).unwrap_or_default()
62    };
63    let aborted_source_ids: Vec<String> = runs
64        .iter()
65        .filter_map(|r| r.source_message_id.clone())
66        .collect();
67    let count = runs.len();
68    tracing::info!(
69        target: "run_control",
70        channel = %channel,
71        chat_id = %chat_id,
72        run_count = %count,
73        aborted_source_ids = ?aborted_source_ids,
74        "abort_runs firing - notifying waiters and marking cancelled"
75    );
76    for run in runs {
77        run.cancelled.store(true, Ordering::SeqCst);
78        run.notify.notify_waiters();
79    }
80    if !aborted_source_ids.is_empty() {
81        let mut guard = ABORTED_SOURCE_MESSAGE_IDS.lock().await;
82        let ids = guard.entry(key).or_default();
83        for id in aborted_source_ids {
84            ids.insert(id);
85        }
86    }
87    count
88}
89
90pub fn is_cancelled(flag: &AtomicBool) -> bool {
91    flag.load(Ordering::SeqCst)
92}
93
94pub async fn is_aborted_source_message(channel: &str, chat_id: i64, message_id: &str) -> bool {
95    let key = (channel.to_string(), chat_id);
96    let guard = ABORTED_SOURCE_MESSAGE_IDS.lock().await;
97    guard
98        .get(&key)
99        .map(|ids| ids.contains(message_id))
100        .unwrap_or(false)
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[tokio::test]
108    async fn test_abort_runs_marks_cancelled() {
109        let channel = "test.abort";
110        let chat_id = 101;
111        let (run_id, cancelled, _notify) =
112            register_run(channel, chat_id, Some("u1".to_string())).await;
113        assert!(!is_cancelled(&cancelled));
114
115        let aborted = abort_runs(channel, chat_id).await;
116        assert_eq!(aborted, 1);
117        assert!(is_cancelled(&cancelled));
118        assert!(is_aborted_source_message(channel, chat_id, "u1").await);
119
120        unregister_run(channel, chat_id, run_id).await;
121    }
122
123    #[tokio::test]
124    async fn test_abort_runs_without_active_returns_zero() {
125        let aborted = abort_runs("test.none", 999).await;
126        assert_eq!(aborted, 0);
127    }
128}
microclaw_engine/run_control.rs

microclaw_engine/
run_control.rs
