Skip to main content

start_command/
signal_handler.rs

1//! Signal handling for graceful cleanup on process interruption
2//!
3//! This module provides signal handlers that update execution status when the process
4//! is interrupted by signals like SIGINT (Ctrl+C), SIGTERM (kill), or SIGHUP (terminal close).
5
6use std::sync::atomic::{AtomicBool, AtomicI32, Ordering};
7use std::sync::Mutex;
8
9use crate::execution_store::{ExecutionRecord, ExecutionStore};
10
11// Global state for signal handling cleanup
12// These are used to update execution status when the process is interrupted
13static SIGNAL_RECEIVED: AtomicBool = AtomicBool::new(false);
14static SIGNAL_EXIT_CODE: AtomicI32 = AtomicI32::new(0);
15static CURRENT_EXECUTION: Mutex<Option<(ExecutionRecord, ExecutionStore)>> = Mutex::new(None);
16// When set, a signal only records its exit code; the running code polls
17// `was_signal_received()` and finishes itself (`--use-formal-ai`, issue #255,
18// which must stop its tool's process group and its runtime before exiting).
19static DEFERRED: AtomicBool = AtomicBool::new(false);
20
21/// Make signals only record their exit code (true) or exit at once (false).
22pub fn set_deferred_signals(deferred: bool) {
23    DEFERRED.store(deferred, Ordering::SeqCst);
24}
25
26/// Check if a signal has been received
27#[allow(dead_code)]
28pub fn was_signal_received() -> bool {
29    SIGNAL_RECEIVED.load(Ordering::SeqCst)
30}
31
32/// Get the exit code from the received signal
33#[allow(dead_code)]
34pub fn get_signal_exit_code() -> i32 {
35    SIGNAL_EXIT_CODE.load(Ordering::SeqCst)
36}
37
38/// Set up signal handlers for graceful cleanup on interruption
39#[cfg(unix)]
40pub fn setup_signal_handlers() {
41    use std::sync::Once;
42    static INIT: Once = Once::new();
43
44    INIT.call_once(|| {
45        unsafe {
46            // SIGINT (Ctrl+C) - exit code 130 (128 + 2)
47            libc::signal(libc::SIGINT, signal_handler as *const () as usize);
48            // SIGTERM (kill command) - exit code 143 (128 + 15)
49            libc::signal(libc::SIGTERM, signal_handler as *const () as usize);
50            // SIGHUP (terminal closed) - exit code 129 (128 + 1)
51            libc::signal(libc::SIGHUP, signal_handler as *const () as usize);
52        }
53    });
54}
55
56#[cfg(not(unix))]
57pub fn setup_signal_handlers() {
58    // Signal handling not supported on non-Unix platforms
59}
60
61/// Signal handler function
62#[cfg(unix)]
63extern "C" fn signal_handler(sig: i32) {
64    // Calculate exit code based on signal (128 + signal number)
65    let exit_code = 128 + sig;
66    SIGNAL_EXIT_CODE.store(exit_code, Ordering::SeqCst);
67    SIGNAL_RECEIVED.store(true, Ordering::SeqCst);
68    if DEFERRED.load(Ordering::SeqCst) {
69        return;
70    }
71
72    // Try to clean up the current execution record
73    cleanup_execution_on_signal(sig, exit_code);
74
75    // Exit with the appropriate code
76    std::process::exit(exit_code);
77}
78
79/// Clean up execution record when a signal is received
80#[cfg(unix)]
81fn cleanup_execution_on_signal(signal: i32, exit_code: i32) {
82    if let Ok(mut guard) = CURRENT_EXECUTION.lock() {
83        if let Some((ref mut record, ref store)) = *guard {
84            // Mark as completed with signal exit code
85            record.complete(exit_code);
86            if let Err(e) = store.save(record) {
87                // Log error to stderr (can't easily check config here)
88                eprintln!(
89                    "\n[Tracking] Warning: Could not save execution record on signal {}: {}",
90                    signal, e
91                );
92            }
93            // Clear the record to prevent double cleanup
94            *guard = None;
95        }
96    }
97}
98
99/// Set the current execution record for signal cleanup
100pub fn set_current_execution(record: ExecutionRecord, store: ExecutionStore) {
101    if let Ok(mut guard) = CURRENT_EXECUTION.lock() {
102        *guard = Some((record, store));
103    }
104}
105
106/// Clear the current execution record (call after normal completion)
107pub fn clear_current_execution() {
108    if let Ok(mut guard) = CURRENT_EXECUTION.lock() {
109        *guard = None;
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[test]
118    fn test_signal_received_initially_false() {
119        // Note: This test might be flaky if run after other tests that set the flag
120        // In a fresh process, SIGNAL_RECEIVED should be false
121        // We can't reliably test this after setup_signal_handlers is called
122    }
123
124    #[test]
125    fn test_set_and_clear_current_execution() {
126        // Just verify the functions don't panic
127        clear_current_execution();
128    }
129}