harn_vm/stdlib/
long_running.rs

1use crate::value::VmDictExt;
2use std::collections::BTreeMap;
3use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
4use std::sync::{Arc, LazyLock, Mutex, OnceLock};
5use std::time::Instant;
6
7use crate::value::VmValue;
8
9static HANDLE_COUNTER: AtomicU64 = AtomicU64::new(1);
10static HANDLE_STORE: LazyLock<Mutex<BTreeMap<String, HandleEntry>>> =
11    LazyLock::new(|| Mutex::new(BTreeMap::new()));
12
13struct HandleEntry {
14    session_id: String,
15    cancel: Arc<AtomicBool>,
16}
17
18pub(crate) struct OperationHandleInfo {
19    pub handle_id: String,
20    pub started_at: String,
21    pub operation: String,
22    pub descriptor: String,
23}
24
25impl OperationHandleInfo {
26    pub(crate) fn into_vm_value(self) -> VmValue {
27        let mut dict = BTreeMap::new();
28        dict.put_str("handle_id", self.handle_id);
29        dict.put_str("started_at", self.started_at);
30        dict.insert("ended_at".to_string(), VmValue::Nil);
31        dict.insert("duration_ms".to_string(), VmValue::Int(0));
32        dict.put_str("status", "running");
33        dict.put_str("operation", self.operation);
34        dict.put_str("command_or_op_descriptor", self.descriptor);
35        VmValue::dict(dict)
36    }
37}
38
39pub(crate) fn spawn_json_operation<F>(
40    operation: impl Into<String>,
41    descriptor: impl Into<String>,
42    session_id: String,
43    run: F,
44) -> Result<OperationHandleInfo, String>
45where
46    F: FnOnce(Arc<AtomicBool>) -> Result<serde_json::Value, String> + Send + 'static,
47{
48    register_cleanup_hook();
49    let handle_id = format!(
50        "hso-{:x}-{}",
51        std::process::id(),
52        HANDLE_COUNTER.fetch_add(1, Ordering::SeqCst)
53    );
54    let started_at = chrono::Utc::now().to_rfc3339();
55    let operation = operation.into();
56    let descriptor = descriptor.into();
57    let cancel = Arc::new(AtomicBool::new(false));
58    {
59        let mut store = HANDLE_STORE
60            .lock()
61            .expect("stdlib long-running handle store poisoned");
62        store.insert(
63            handle_id.clone(),
64            HandleEntry {
65                session_id,
66                cancel: cancel.clone(),
67            },
68        );
69    }
70
71    let worker_handle_id = handle_id.clone();
72    let worker_started_at = started_at.clone();
73    let worker_operation = operation.clone();
74    let worker_descriptor = descriptor.clone();
75    std::thread::Builder::new()
76        .name(format!("hso-worker-{worker_handle_id}"))
77        .spawn(move || {
78            let started = Instant::now();
79            let result = run(cancel.clone());
80            let entry = {
81                let mut store = HANDLE_STORE
82                    .lock()
83                    .expect("stdlib long-running handle store poisoned");
84                store.remove(&worker_handle_id)
85            };
86            if cancel.load(Ordering::Acquire) {
87                return;
88            }
89            let Some(entry) = entry else {
90                return;
91            };
92            let mut payload = serde_json::Map::new();
93            payload.insert(
94                "handle_id".to_string(),
95                serde_json::Value::String(worker_handle_id),
96            );
97            payload.insert(
98                "operation".to_string(),
99                serde_json::Value::String(worker_operation),
100            );
101            payload.insert(
102                "command_or_op_descriptor".to_string(),
103                serde_json::Value::String(worker_descriptor),
104            );
105            payload.insert(
106                "started_at".to_string(),
107                serde_json::Value::String(worker_started_at),
108            );
109            payload.insert(
110                "ended_at".to_string(),
111                serde_json::Value::String(chrono::Utc::now().to_rfc3339()),
112            );
113            payload.insert(
114                "duration_ms".to_string(),
115                serde_json::Value::Number((started.elapsed().as_millis() as u64).into()),
116            );
117            match result {
118                Ok(value) => {
119                    payload.insert(
120                        "status".to_string(),
121                        serde_json::Value::String("completed".to_string()),
122                    );
123                    payload.insert("result".to_string(), value);
124                }
125                Err(error) => {
126                    payload.insert(
127                        "status".to_string(),
128                        serde_json::Value::String("failed".to_string()),
129                    );
130                    payload.insert("error".to_string(), serde_json::Value::String(error));
131                }
132            }
133            let content = serde_json::to_string(&payload).unwrap_or_default();
134            crate::orchestration::agent_inbox::push(
135                &entry.session_id,
136                "tool_result",
137                &content,
138                "stdlib.long_running",
139            );
140        })
141        .map_err(|error| {
142            let entry = {
143                let mut store = HANDLE_STORE
144                    .lock()
145                    .expect("stdlib long-running handle store poisoned");
146                store.remove(&handle_id)
147            };
148            if let Some(entry) = entry {
149                entry.cancel.store(true, Ordering::Release);
150            }
151            format!("failed to spawn stdlib long-running worker: {error}")
152        })?;
153
154    Ok(OperationHandleInfo {
155        handle_id,
156        started_at,
157        operation,
158        descriptor,
159    })
160}
161
162pub fn cancel_handle(handle_id: &str) -> bool {
163    let entry = {
164        let mut store = HANDLE_STORE
165            .lock()
166            .expect("stdlib long-running handle store poisoned");
167        store.remove(handle_id)
168    };
169    match entry {
170        Some(entry) => {
171            entry.cancel.store(true, Ordering::Release);
172            true
173        }
174        None => false,
175    }
176}
177
178pub(crate) fn cancel_session_handles(session_id: &str) {
179    let entries = {
180        let mut store = HANDLE_STORE
181            .lock()
182            .expect("stdlib long-running handle store poisoned");
183        let matching = store
184            .iter()
185            .filter(|(_, entry)| entry.session_id == session_id)
186            .map(|(handle_id, _)| handle_id.clone())
187            .collect::<Vec<_>>();
188        matching
189            .into_iter()
190            .filter_map(|handle_id| store.remove(&handle_id))
191            .collect::<Vec<_>>()
192    };
193    for entry in entries {
194        entry.cancel.store(true, Ordering::Release);
195    }
196}
197
198pub(crate) fn register_cleanup_hook() {
199    static REGISTERED: OnceLock<()> = OnceLock::new();
200    REGISTERED.get_or_init(|| {
201        let hook: Arc<dyn Fn(&str) + Send + Sync> = Arc::new(cancel_session_handles);
202        crate::llm::register_session_end_hook(hook);
203    });
204}
205
206/// Cancel every in-flight long-running handle and clear the registry.
207///
208/// Used only by the long-running fs/glob tests to start each test from
209/// a clean store. Production callers must NOT use this — it wipes
210/// process-global state shared across threads. Removed from
211/// `reset_stdlib_state` so other tests can't accidentally cancel a
212/// concurrent fs/glob worker thread mid-flight.
213#[cfg(test)]
214pub(crate) fn reset_state() {
215    let entries = {
216        let mut store = HANDLE_STORE
217            .lock()
218            .expect("stdlib long-running handle store poisoned");
219        std::mem::take(&mut *store)
220    };
221    for (_, entry) in entries {
222        entry.cancel.store(true, Ordering::Release);
223    }
224}
harn_vm/stdlib/long_running.rs

harn_vm/stdlib/
long_running.rs
