Skip to main content

harn_hostlib/tools/
long_running.rs

1//! Long-running tool handle machinery.
2//!
3//! When a caller passes `long_running: true` to `run_command`, `run_test`, or
4//! `run_build_command`, the builtin spawns the child process without waiting,
5//! registers it here, and returns a handle dict immediately:
6//!
7//! ```json
8//! {
9//!   "handle_id": "hto-<pid-hex>-<n>",
10//!   "started_at": "...",
11//!   "command_or_op_descriptor": "..."
12//! }
13//! ```
14//!
15//! A background thread waits for the child and, when it exits, pushes a
16//! `tool_result` entry into the active session's `agent_inbox` via
17//! `harn_vm::orchestration::agent_inbox::push(...)` so the agent-loop's
18//! next turn-preflight (or post-compaction drain) picks it up.
19//!
20//! ### Cancellation
21//!
22//! `cancel_handle(handle_id)` kills the spawned process (SIGKILL) within
23//! 2 seconds. The session-end hook registered on startup kills every
24//! in-flight handle associated with the ending session.
25//!
26//! #### PID-based signaling
27//!
28//! The waiter thread takes ownership of the `Child` object to drain
29//! stdout/stderr and call `wait()`. To keep cancellation possible even
30//! after the waiter has taken the `Child`, we store the raw OS process ID
31//! in the entry and kill by PID when needed. On Unix we call `kill(2)`
32//! directly via an `extern "C"` declaration (no `libc` crate required).
33//! A shared `cancelled` flag suppresses the feedback push when the waiter
34//! sees an exit caused by cancellation. Callers that need artifact-stable
35//! cancellation can opt into waiting for the waiter result through
36//! `cancel_handle`.
37
38use std::collections::BTreeMap;
39use std::io::{Read, Write};
40use std::path::PathBuf;
41use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
42use std::sync::{Arc, LazyLock, Mutex, OnceLock};
43use std::time::Duration;
44
45use harn_vm::VmValue;
46
47use crate::error::HostlibError;
48use crate::json::vm_dict_to_json;
49use crate::process::{self as process_handle, ProcessHandle, ProcessKiller, SpawnSpec};
50use crate::tools::args::to_agent_path;
51use crate::tools::proc::{self, CaptureConfig, CommandStatus, EnvMode};
52
53/// Atomic counter for generating unique handle IDs within this process.
54static HANDLE_COUNTER: AtomicU64 = AtomicU64::new(1);
55
56/// Shared cancellation state between the store entry and its waiter thread.
57struct CancelState {
58    /// Set to `true` when `cancel_handle` / `cancel_session_handles` runs.
59    /// The waiter checks this before pushing feedback.
60    cancelled: AtomicBool,
61    /// Set by cancellation paths that represent a timeout rather than a
62    /// user-requested kill. The waiter uses this for the returned result
63    /// status while still suppressing inbox feedback.
64    timed_out: AtomicBool,
65}
66
67#[derive(Default)]
68struct OutputState {
69    stdout: Vec<u8>,
70    stderr: Vec<u8>,
71}
72
73/// Shared state for a single in-flight child process.
74struct HandleEntry {
75    /// The process handle. `None` after the waiter thread takes ownership.
76    handle: Option<Box<dyn ProcessHandle>>,
77    /// Killer that works even after the waiter took `handle`.
78    killer: Arc<dyn ProcessKiller>,
79    session_id: String,
80    /// Shared with the waiter thread.
81    cancel_state: Arc<CancelState>,
82    /// Sender used by the waiter thread to signal that the post-exit
83    /// feedback push is complete. `None` if the test-side hasn't asked
84    /// to be notified.
85    completion_tx: Option<std::sync::mpsc::SyncSender<()>>,
86    /// Optional one-shot result channel installed by `cancel_handle` when a
87    /// caller wants cancellation to wait until artifacts have been drained.
88    result_tx: Option<std::sync::mpsc::SyncSender<VmValue>>,
89    /// Opaque verification snapshot binding provided by the caller.
90    snapshot_binding: Option<harn_vm::value::DictMap>,
91}
92
93#[derive(Default)]
94struct HandleStore {
95    entries: BTreeMap<String, HandleEntry>,
96}
97
98static HANDLE_STORE: LazyLock<Mutex<HandleStore>> =
99    LazyLock::new(|| Mutex::new(HandleStore::default()));
100
101/// Metadata returned to the caller immediately when a long-running spawn
102/// succeeds. Serialised as a response dict by the calling builtin.
103pub struct LongRunningHandleInfo {
104    /// Command identifier shared with foreground command responses.
105    pub command_id: String,
106    /// Opaque handle identifier, e.g. `"hto-<pid-hex>-<n>"`.
107    pub handle_id: String,
108    /// RFC 3339 timestamp of the spawn.
109    pub started_at: String,
110    /// Raw child process id reported by the platform.
111    pub pid: u32,
112    /// Child process group id when the platform exposes it.
113    pub process_group_id: Option<u32>,
114    /// Human-readable display form of the argv (space-joined).
115    pub command_display: String,
116    /// Opaque verification snapshot binding provided by the caller.
117    pub snapshot_binding: Option<harn_vm::value::DictMap>,
118}
119
120pub(crate) struct LongRunningSpawnOptions {
121    pub(crate) env_mode: EnvMode,
122    pub(crate) env_remove: Vec<String>,
123    pub(crate) capture: CaptureConfig,
124    pub(crate) session_id: String,
125    pub(crate) progress_interval: Option<Duration>,
126    pub(crate) progress_max_inline_bytes: usize,
127    pub(crate) snapshot_binding: Option<harn_vm::value::DictMap>,
128}
129
130struct WaiterContext {
131    command_id: String,
132    handle_id: String,
133    session_id: String,
134    started_at: String,
135    process_group_id: Option<u32>,
136    command_display: String,
137    progress_interval: Option<Duration>,
138    progress_max_inline_bytes: usize,
139    snapshot_binding: Option<harn_vm::value::DictMap>,
140}
141
142struct ProgressThreadContext {
143    command_id: String,
144    handle_id: String,
145    session_id: String,
146    started_at: String,
147    command_display: String,
148    process_group_id: Option<u32>,
149    output_path: PathBuf,
150    stdout_path: PathBuf,
151    stderr_path: PathBuf,
152    output_state: Arc<Mutex<OutputState>>,
153    cancel_state: Arc<CancelState>,
154    done: Arc<AtomicBool>,
155    started: std::time::Instant,
156    interval: Duration,
157    max_inline_bytes: usize,
158    snapshot_binding: Option<harn_vm::value::DictMap>,
159}
160
161impl LongRunningHandleInfo {
162    /// Convert into the standard handle response dict returned to the agent.
163    pub fn into_handle_response(self) -> VmValue {
164        let Self {
165            command_id,
166            handle_id,
167            started_at,
168            pid,
169            process_group_id,
170            command_display,
171            snapshot_binding,
172        } = self;
173        proc::running_response(
174            command_id,
175            handle_id,
176            pid,
177            process_group_id,
178            started_at,
179            command_display,
180            snapshot_binding.as_ref(),
181        )
182    }
183}
184
185/// Spawn the argv as a long-running child process and return a handle.
186///
187/// The background waiter pushes a `tool_result` entry into the active
188/// session's `agent_inbox` when the process exits so the next
189/// agent-loop turn sees the result.
190pub fn spawn_long_running(
191    builtin: &'static str,
192    program: String,
193    args: Vec<String>,
194    cwd: Option<PathBuf>,
195    env: BTreeMap<String, String>,
196    session_id: String,
197) -> Result<LongRunningHandleInfo, HostlibError> {
198    spawn_long_running_with_options(
199        builtin,
200        program,
201        args,
202        cwd,
203        env,
204        LongRunningSpawnOptions {
205            env_mode: EnvMode::InheritClean,
206            env_remove: Vec::new(),
207            capture: CaptureConfig::default(),
208            session_id,
209            progress_interval: None,
210            progress_max_inline_bytes: CaptureConfig::default().max_inline_bytes,
211            snapshot_binding: None,
212        },
213    )
214}
215
216pub(crate) fn spawn_long_running_with_options(
217    builtin: &'static str,
218    program: String,
219    args: Vec<String>,
220    cwd: Option<PathBuf>,
221    env: BTreeMap<String, String>,
222    options: LongRunningSpawnOptions,
223) -> Result<LongRunningHandleInfo, HostlibError> {
224    let mut env = env;
225    proc::apply_toolchain_path(cwd.as_deref(), &mut env, options.env_mode);
226    let spec = SpawnSpec {
227        builtin,
228        program: program.clone(),
229        args: args.clone(),
230        cwd,
231        env,
232        env_remove: options.env_remove.clone(),
233        env_mode: options.env_mode,
234        use_stdin: false,
235        configure_process_group: true,
236    };
237    let handle = process_handle::spawn_process(spec)
238        .map_err(|e| proc::process_error_to_hostlib(builtin, e))?;
239
240    let pid = handle.pid().unwrap_or(0);
241    let process_group_id = handle.process_group_id();
242    let killer = handle.killer();
243    let id = HANDLE_COUNTER.fetch_add(1, Ordering::SeqCst);
244    let handle_id = format!("hto-{:x}-{id}", std::process::id());
245    let command_id = proc::next_command_id();
246    let started_at = proc::now_rfc3339();
247    let _artifacts = proc::register_live_artifacts(&command_id, Some(&handle_id))?;
248
249    let mut all_argv = vec![program];
250    all_argv.extend(args.iter().cloned());
251    let command_display = all_argv.join(" ");
252
253    let cancel_state = Arc::new(CancelState {
254        cancelled: AtomicBool::new(false),
255        timed_out: AtomicBool::new(false),
256    });
257
258    {
259        let mut store = HANDLE_STORE
260            .lock()
261            .expect("long-running handle store poisoned");
262        store.entries.insert(
263            handle_id.clone(),
264            HandleEntry {
265                handle: Some(handle),
266                killer,
267                session_id: options.session_id.clone(),
268                cancel_state: cancel_state.clone(),
269                completion_tx: None,
270                result_tx: None,
271                snapshot_binding: options.snapshot_binding.clone(),
272            },
273        );
274    }
275
276    let waiter_context = WaiterContext {
277        command_id: command_id.clone(),
278        handle_id: handle_id.clone(),
279        session_id: options.session_id,
280        started_at: started_at.clone(),
281        process_group_id,
282        command_display: command_display.clone(),
283        progress_interval: options.progress_interval,
284        progress_max_inline_bytes: options.progress_max_inline_bytes,
285        snapshot_binding: options.snapshot_binding.clone(),
286    };
287    let waiter_thread_name = waiter_context.handle_id.clone();
288    let capture = options.capture;
289    std::thread::Builder::new()
290        .name(format!("hto-waiter-{waiter_thread_name}"))
291        .spawn(move || {
292            waiter_thread(waiter_context, cancel_state, capture);
293        })
294        .map_err(|e| HostlibError::Backend {
295            builtin,
296            message: format!("failed to spawn waiter thread: {e}"),
297        })?;
298
299    Ok(LongRunningHandleInfo {
300        command_id,
301        handle_id,
302        started_at,
303        pid,
304        process_group_id,
305        command_display,
306        snapshot_binding: options.snapshot_binding,
307    })
308}
309
310/// Background thread that waits for a child process and fires feedback.
311fn waiter_thread(context: WaiterContext, cancel_state: Arc<CancelState>, capture: CaptureConfig) {
312    let waiter_start = std::time::Instant::now();
313
314    // Take the handle out of the store. If the entry is already gone (i.e.
315    // cancel_handle ran and removed it before us), exit without action.
316    let mut handle = {
317        let mut store = HANDLE_STORE
318            .lock()
319            .expect("long-running handle store poisoned");
320        match store.entries.get_mut(&context.handle_id) {
321            Some(entry) => match entry.handle.take() {
322                Some(h) => h,
323                None => return, // already cancelled before we ran
324            },
325            None => return, // entry removed (cancelled before store insert — shouldn't happen)
326        }
327    };
328
329    let output_state = Arc::new(Mutex::new(OutputState::default()));
330    let done = Arc::new(AtomicBool::new(false));
331    let planned = proc::planned_artifact_paths(&context.command_id);
332    if let Some(parent) = planned.output_path.parent() {
333        let _ = std::fs::create_dir_all(parent);
334    }
335    let _ = std::fs::File::create(&planned.stdout_path);
336    let _ = std::fs::File::create(&planned.stderr_path);
337    let combined_file = std::fs::File::create(&planned.output_path)
338        .ok()
339        .map(|file| Arc::new(Mutex::new(file)));
340
341    let stdout_thread = handle.take_stdout().map(|out| {
342        spawn_output_drain(
343            out,
344            output_state.clone(),
345            planned.stdout_path.clone(),
346            combined_file.clone(),
347            true,
348        )
349    });
350    let stderr_thread = handle.take_stderr().map(|err| {
351        spawn_output_drain(
352            err,
353            output_state.clone(),
354            planned.stderr_path.clone(),
355            combined_file.clone(),
356            false,
357        )
358    });
359
360    let progress_thread = context
361        .progress_interval
362        .filter(|interval| !interval.is_zero())
363        .map(|interval| {
364            spawn_progress_thread(ProgressThreadContext {
365                command_id: context.command_id.clone(),
366                handle_id: context.handle_id.clone(),
367                session_id: context.session_id.clone(),
368                started_at: context.started_at.clone(),
369                command_display: context.command_display.clone(),
370                process_group_id: context.process_group_id,
371                output_path: planned.output_path.clone(),
372                stdout_path: planned.stdout_path.clone(),
373                stderr_path: planned.stderr_path.clone(),
374                output_state: output_state.clone(),
375                cancel_state: cancel_state.clone(),
376                done: done.clone(),
377                started: waiter_start,
378                interval,
379                max_inline_bytes: context.progress_max_inline_bytes,
380                snapshot_binding: context.snapshot_binding.clone(),
381            })
382        });
383
384    let status = handle.wait().ok();
385
386    if let Some(thread) = stdout_thread {
387        let _ = thread.join();
388    }
389    if let Some(thread) = stderr_thread {
390        let _ = thread.join();
391    }
392    done.store(true, Ordering::Release);
393    drop(progress_thread);
394    let (stdout, stderr) = {
395        let state = output_state
396            .lock()
397            .unwrap_or_else(|poison| poison.into_inner());
398        (state.stdout.clone(), state.stderr.clone())
399    };
400
401    // Remove our entry from the store, taking notifiers on the way out so we
402    // can signal them after the feedback/result path completes.
403    let (completion_tx, result_tx) = {
404        let mut store = HANDLE_STORE
405            .lock()
406            .expect("long-running handle store poisoned");
407        let entry = store
408            .entries
409            .remove(&context.handle_id)
410            .map(|mut e| (e.completion_tx.take(), e.result_tx.take()));
411        entry.unwrap_or((None, None))
412    };
413
414    let signal_done = move || {
415        if let Some(tx) = completion_tx {
416            let _ = tx.try_send(());
417        }
418    };
419
420    let cancelled = cancel_state.cancelled.load(Ordering::Acquire);
421    let timed_out = cancelled && cancel_state.timed_out.load(Ordering::Acquire);
422
423    let (exit_code, signal_name) = match status {
424        Some(s) => decode_exit_status(s),
425        // wait() itself failed — treat as killed (extremely unusual).
426        None => (-1, Some("SIGKILL".to_string())),
427    };
428    let command_status = if timed_out {
429        CommandStatus::TimedOut
430    } else if cancelled {
431        CommandStatus::Killed
432    } else {
433        CommandStatus::Completed
434    };
435    let duration = waiter_start.elapsed();
436    let duration_ms = duration.as_millis() as i64;
437    let artifacts = match proc::persist_artifacts(
438        &context.command_id,
439        &stdout,
440        &stderr,
441        Some(&context.handle_id),
442    ) {
443        Ok(artifacts) => artifacts,
444        Err(_) => return,
445    };
446    let (inline_stdout, inline_stderr) = proc::inline_output(&stdout, &stderr, capture);
447
448    let mut payload = serde_json::Map::new();
449    payload.insert(
450        "command_id".into(),
451        serde_json::Value::String(context.command_id.clone()),
452    );
453    payload.insert(
454        "status".into(),
455        serde_json::Value::String(command_status.as_str().to_string()),
456    );
457    payload.insert(
458        "handle_id".into(),
459        serde_json::Value::String(context.handle_id),
460    );
461    payload.insert(
462        "command_or_op_descriptor".into(),
463        serde_json::Value::String(context.command_display),
464    );
465    payload.insert(
466        "started_at".into(),
467        serde_json::Value::String(context.started_at),
468    );
469    payload.insert(
470        "ended_at".into(),
471        serde_json::Value::String(proc::now_rfc3339()),
472    );
473    payload.insert(
474        "duration_ms".into(),
475        serde_json::Value::Number(duration_ms.into()),
476    );
477    payload.insert(
478        "exit_code".into(),
479        serde_json::Value::Number(exit_code.into()),
480    );
481    payload.insert("timed_out".into(), serde_json::Value::Bool(timed_out));
482    payload.insert("stdout".into(), serde_json::Value::String(inline_stdout));
483    payload.insert("stderr".into(), serde_json::Value::String(inline_stderr));
484    payload.insert(
485        "output_path".into(),
486        serde_json::Value::String(to_agent_path(&artifacts.output_path)),
487    );
488    payload.insert(
489        "stdout_path".into(),
490        serde_json::Value::String(to_agent_path(&artifacts.stdout_path)),
491    );
492    payload.insert(
493        "stderr_path".into(),
494        serde_json::Value::String(to_agent_path(&artifacts.stderr_path)),
495    );
496    payload.insert(
497        "line_count".into(),
498        serde_json::Value::Number(artifacts.line_count.into()),
499    );
500    payload.insert(
501        "byte_count".into(),
502        serde_json::Value::Number(artifacts.byte_count.into()),
503    );
504    payload.insert(
505        "output_sha256".into(),
506        serde_json::Value::String(artifacts.output_sha256),
507    );
508    if let Some(pgid) = context.process_group_id {
509        payload.insert(
510            "process_group_id".into(),
511            serde_json::Value::Number((pgid as u64).into()),
512        );
513    }
514    if let Some(sig) = signal_name {
515        payload.insert("signal".into(), serde_json::Value::String(sig));
516    } else {
517        payload.insert("signal".into(), serde_json::Value::Null);
518    }
519    if let Some(snapshot_binding) = context.snapshot_binding.as_ref() {
520        payload.insert("snapshot_binding".into(), vm_dict_to_json(snapshot_binding));
521    }
522
523    if let Some(tx) = result_tx {
524        let value = serde_json::Value::Object(payload.clone());
525        let _ = tx.try_send(harn_vm::json_to_vm_value(&value));
526    }
527    if !cancelled {
528        let content = serde_json::to_string(&payload).unwrap_or_default();
529        harn_vm::orchestration::agent_inbox::push(
530            &context.session_id,
531            "tool_result",
532            &content,
533            "hostlib.long_running.exit",
534        );
535    }
536    signal_done();
537}
538
539fn spawn_output_drain(
540    mut reader: Box<dyn Read + Send>,
541    state: Arc<Mutex<OutputState>>,
542    path: std::path::PathBuf,
543    combined_file: Option<Arc<Mutex<std::fs::File>>>,
544    stdout: bool,
545) -> std::thread::JoinHandle<()> {
546    std::thread::spawn(move || {
547        let mut file = std::fs::File::create(path).ok();
548        let mut buf = [0_u8; 8192];
549        loop {
550            let read = match reader.read(&mut buf) {
551                Ok(0) => break,
552                Ok(read) => read,
553                Err(_) => break,
554            };
555            let chunk = &buf[..read];
556            if let Some(file) = file.as_mut() {
557                let _ = file.write_all(chunk);
558            }
559            if let Some(combined) = combined_file.as_ref() {
560                if let Ok(mut combined) = combined.lock() {
561                    let _ = combined.write_all(chunk);
562                }
563            }
564            if let Ok(mut state) = state.lock() {
565                if stdout {
566                    state.stdout.extend_from_slice(chunk);
567                } else {
568                    state.stderr.extend_from_slice(chunk);
569                }
570            }
571        }
572    })
573}
574
575fn spawn_progress_thread(context: ProgressThreadContext) -> std::thread::JoinHandle<()> {
576    std::thread::spawn(move || {
577        while !context.done.load(Ordering::Acquire)
578            && !context.cancel_state.cancelled.load(Ordering::Acquire)
579        {
580            std::thread::sleep(context.interval);
581            if context.done.load(Ordering::Acquire)
582                || context.cancel_state.cancelled.load(Ordering::Acquire)
583            {
584                break;
585            }
586            let (stdout, stderr) = {
587                let state = context
588                    .output_state
589                    .lock()
590                    .unwrap_or_else(|poison| poison.into_inner());
591                (state.stdout.clone(), state.stderr.clone())
592            };
593            let capture = CaptureConfig {
594                max_inline_bytes: context.max_inline_bytes,
595                ..CaptureConfig::default()
596            };
597            let (inline_stdout, inline_stderr) = proc::inline_output(&stdout, &stderr, capture);
598            let byte_count = stdout.len().saturating_add(stderr.len());
599            let mut payload = serde_json::json!({
600                "command_id": &context.command_id,
601                "handle_id": &context.handle_id,
602                "status": CommandStatus::Running.as_str(),
603                "command_or_op_descriptor": &context.command_display,
604                "started_at": &context.started_at,
605                "ended_at": null,
606                "duration_ms": context.started.elapsed().as_millis() as i64,
607                "exit_code": null,
608                "signal": null,
609                "stdout": inline_stdout,
610                "stderr": inline_stderr,
611                "output_path": to_agent_path(&context.output_path),
612                "stdout_path": to_agent_path(&context.stdout_path),
613                "stderr_path": to_agent_path(&context.stderr_path),
614                "byte_count": byte_count as i64,
615                "line_count": stdout.iter().chain(stderr.iter()).filter(|byte| **byte == b'\n').count() as i64,
616                "process_group_id": context.process_group_id,
617            });
618            if let (Some(object), Some(snapshot_binding)) =
619                (payload.as_object_mut(), context.snapshot_binding.as_ref())
620            {
621                object.insert(
622                    "snapshot_binding".to_string(),
623                    vm_dict_to_json(snapshot_binding),
624                );
625            }
626            harn_vm::orchestration::agent_inbox::push(
627                &context.session_id,
628                "tool_progress",
629                &payload.to_string(),
630                "hostlib.long_running.progress",
631            );
632        }
633    })
634}
635
636pub(crate) struct CancelOptions {
637    pub(crate) timed_out: bool,
638    pub(crate) wait_result: Option<Duration>,
639}
640
641pub(crate) struct CancelOutcome {
642    pub(crate) cancelled: bool,
643    pub(crate) result: Option<VmValue>,
644}
645
646/// Cancel a specific in-flight long-running handle. Kills the process and lets
647/// the waiter drain output/artifacts. Returns `true` if the handle was found
648/// and cancellation was newly requested.
649pub fn cancel_handle(handle_id: &str) -> bool {
650    cancel_handle_with_options(
651        handle_id,
652        CancelOptions {
653            timed_out: false,
654            wait_result: None,
655        },
656    )
657    .cancelled
658}
659
660pub(crate) fn snapshot_binding_for_handle(handle_id: &str) -> Option<harn_vm::value::DictMap> {
661    let store = HANDLE_STORE
662        .lock()
663        .expect("long-running handle store poisoned");
664    store
665        .entries
666        .get(handle_id)
667        .and_then(|entry| entry.snapshot_binding.clone())
668}
669
670pub(crate) fn cancel_handle_with_options(handle_id: &str, options: CancelOptions) -> CancelOutcome {
671    let (killer, cancel_state, result_rx) = {
672        let mut store = HANDLE_STORE
673            .lock()
674            .expect("long-running handle store poisoned");
675        let Some(entry) = store.entries.get_mut(handle_id) else {
676            return CancelOutcome {
677                cancelled: false,
678                result: None,
679            };
680        };
681        if entry.cancel_state.cancelled.swap(true, Ordering::AcqRel) {
682            return CancelOutcome {
683                cancelled: false,
684                result: None,
685            };
686        }
687        entry
688            .cancel_state
689            .timed_out
690            .store(options.timed_out, Ordering::Release);
691        let result_rx = options.wait_result.map(|_| {
692            let (tx, rx) = std::sync::mpsc::sync_channel::<VmValue>(1);
693            entry.result_tx = Some(tx);
694            rx
695        });
696        (entry.killer.clone(), entry.cancel_state.clone(), result_rx)
697    };
698    do_kill(killer, cancel_state);
699    let result = match (options.wait_result, result_rx) {
700        (Some(timeout), Some(rx)) => rx.recv_timeout(timeout).ok(),
701        _ => None,
702    };
703    CancelOutcome {
704        cancelled: true,
705        result,
706    }
707}
708
709/// Tuple shape used by `cancel_session_handles` to drain entries while
710/// holding the store lock for as little as possible. Boxed-trait fields
711/// make it noisy to inline as an unnamed type.
712type SessionKillEntry = (Arc<dyn ProcessKiller>, Arc<CancelState>);
713
714/// Cancel all in-flight handles for a given session. Called by the
715/// session-end hook to avoid orphaned processes.
716pub fn cancel_session_handles(session_id: &str) {
717    let to_kill: Vec<SessionKillEntry> = {
718        let store = HANDLE_STORE
719            .lock()
720            .expect("long-running handle store poisoned");
721        let matching: Vec<String> = store
722            .entries
723            .iter()
724            .filter(|(_, e)| e.session_id == session_id)
725            .map(|(id, _)| id.clone())
726            .collect();
727        matching
728            .into_iter()
729            .filter_map(|id| {
730                let entry = store.entries.get(&id)?;
731                if entry.cancel_state.cancelled.swap(true, Ordering::AcqRel) {
732                    return None;
733                }
734                entry.cancel_state.timed_out.store(false, Ordering::Release);
735                Some((entry.killer.clone(), entry.cancel_state.clone()))
736            })
737            .collect()
738    };
739    for (killer, cancel_state) in to_kill {
740        do_kill(killer, cancel_state);
741    }
742}
743
744/// Set the cancellation flag and kill the process. Used by both `cancel_handle`
745/// and `cancel_session_handles`.
746fn do_kill(killer: Arc<dyn ProcessKiller>, cancel_state: Arc<CancelState>) {
747    // Kill via the handle's killer (works whether or not we still own
748    // the handle). The waiter owns process reaping and artifact finalization.
749    killer.kill();
750    cancel_state.cancelled.store(true, Ordering::Release);
751}
752
753/// Register the session-cleanup hook with harn-vm. Uses a `OnceLock` so the
754/// hook is registered exactly once even if `register_builtins` is called
755/// multiple times (e.g. in tests).
756pub(crate) fn register_cleanup_hook() {
757    static REGISTERED: OnceLock<()> = OnceLock::new();
758    REGISTERED.get_or_init(|| {
759        let hook: Arc<dyn Fn(&str) + Send + Sync> = Arc::new(|session_id: &str| {
760            cancel_session_handles(session_id);
761        });
762        harn_vm::register_session_end_hook(hook);
763    });
764}
765
766fn decode_exit_status(status: process_handle::ExitStatus) -> (i32, Option<String>) {
767    if let Some(code) = status.code {
768        return (code, None);
769    }
770    if let Some(sig) = status.signal {
771        return (-1, Some(format!("SIG{sig}")));
772    }
773    (-1, None)
774}
775
776/// Register a completion notifier for `handle_id`. The waiter thread sends
777/// `()` on the returned receiver after it pushes the feedback item to the
778/// global queue. Returns `None` if the handle is no longer in the store
779/// (e.g. already cancelled or completed). Used by tests to await waiter
780/// completion deterministically — no polling, no `thread::sleep`.
781pub fn register_completion_notifier(handle_id: &str) -> Option<std::sync::mpsc::Receiver<()>> {
782    let (tx, rx) = std::sync::mpsc::sync_channel::<()>(1);
783    let mut store = HANDLE_STORE
784        .lock()
785        .expect("long-running handle store poisoned");
786    let entry = store.entries.get_mut(handle_id)?;
787    entry.completion_tx = Some(tx);
788    Some(rx)
789}