Skip to main content

ralph_workflow/
interrupt.rs

1//! Interrupt signal handling for graceful checkpoint save.
2//!
3//! This module provides signal handling for the Ralph pipeline, ensuring
4//! clean shutdown when the user interrupts with Ctrl+C.
5//!
6//! When an interrupt is received, the handler will:
7//! 1. Save a checkpoint with the `Interrupted` phase
8//! 2. Clean up temporary files
9//! 3. Exit gracefully
10
11use std::sync::Mutex;
12
13use crate::workspace::Workspace;
14
15/// Global interrupt context for checkpoint saving on interrupt.
16///
17/// This is set during pipeline initialization and used by the interrupt
18/// handler to save a checkpoint when the user presses Ctrl+C.
19static INTERRUPT_CONTEXT: Mutex<Option<InterruptContext>> = Mutex::new(None);
20
21/// Context needed to save a checkpoint when interrupted.
22///
23/// This structure holds references to all the state needed to create
24/// a checkpoint when the user interrupts the pipeline with Ctrl+C.
25pub struct InterruptContext {
26    /// Current pipeline phase
27    pub phase: crate::checkpoint::PipelinePhase,
28    /// Current iteration number
29    pub iteration: u32,
30    /// Total iterations configured
31    pub total_iterations: u32,
32    /// Current reviewer pass number
33    pub reviewer_pass: u32,
34    /// Total reviewer passes configured
35    pub total_reviewer_passes: u32,
36    /// Run context for tracking execution lineage
37    pub run_context: crate::checkpoint::RunContext,
38    /// Execution history tracking
39    pub execution_history: crate::checkpoint::ExecutionHistory,
40    /// Prompt history for deterministic resume
41    pub prompt_history: std::collections::HashMap<String, String>,
42    /// Workspace for checkpoint persistence
43    pub workspace: std::sync::Arc<dyn Workspace>,
44}
45
46/// Set the global interrupt context.
47///
48/// This function should be called during pipeline initialization to
49/// provide the interrupt handler with the context needed to save
50/// a checkpoint when interrupted.
51///
52/// # Arguments
53///
54/// * `context` - The interrupt context to store
55///
56/// # Note
57///
58/// This function is typically called at the start of `run_pipeline()`
59/// to ensure the interrupt handler has the most up-to-date context.
60pub fn set_interrupt_context(context: InterruptContext) {
61    let mut ctx = INTERRUPT_CONTEXT.lock().unwrap_or_else(|poison| {
62        // If mutex is poisoned, recover the guard and clear the state
63        poison.into_inner()
64    });
65    *ctx = Some(context);
66}
67
68/// Clear the global interrupt context.
69///
70/// This should be called when the pipeline completes successfully
71/// to prevent saving an interrupt checkpoint after normal completion.
72pub fn clear_interrupt_context() {
73    let mut ctx = INTERRUPT_CONTEXT.lock().unwrap_or_else(|poison| {
74        // If mutex is poisoned, recover the guard and clear the state
75        poison.into_inner()
76    });
77    *ctx = None;
78}
79
80/// Set up the interrupt handler for graceful shutdown with checkpoint saving.
81///
82/// This function registers a SIGINT handler that will:
83/// 1. Save a checkpoint with the current pipeline state
84/// 2. Clean up generated files
85/// 3. Exit gracefully
86///
87/// Call this early in main() after initializing the pipeline context.
88pub fn setup_interrupt_handler() {
89    ctrlc::set_handler(|| {
90        eprintln!("\nāœ‹ Interrupt received! Saving checkpoint...");
91
92        // Try to save checkpoint if context is available
93        let ctx = INTERRUPT_CONTEXT.lock().unwrap_or_else(|poison| {
94            // If mutex is poisoned, recover the guard (context may be inconsistent)
95            poison.into_inner()
96        });
97        if let Some(ref context) = *ctx {
98            if let Err(e) = save_interrupt_checkpoint(context) {
99                eprintln!("Warning: Failed to save checkpoint: {}", e);
100            } else {
101                eprintln!("āœ“ Checkpoint saved. Resume with: ralph --resume");
102            }
103        }
104        drop(ctx); // Release lock before cleanup
105
106        eprintln!("Cleaning up...");
107        crate::git_helpers::cleanup_agent_phase_silent();
108        std::process::exit(130); // Standard exit code for SIGINT
109    })
110    .ok(); // Ignore errors if handler can't be set
111}
112
113/// Save a checkpoint when the pipeline is interrupted.
114///
115/// This function creates a checkpoint with the `Interrupted` phase,
116/// which has the highest phase rank so resuming will run the full pipeline.
117///
118/// The original phase information is preserved for display purposes
119/// by updating the checkpoint with current progress information.
120///
121/// # Arguments
122///
123/// * `context` - The interrupt context containing the current pipeline state
124fn save_interrupt_checkpoint(context: &InterruptContext) -> anyhow::Result<()> {
125    use crate::checkpoint::PipelinePhase;
126    use crate::checkpoint::{load_checkpoint_with_workspace, save_checkpoint_with_workspace};
127
128    // Read checkpoint from file if exists, update it with Interrupted phase
129    if let Ok(Some(mut checkpoint)) = load_checkpoint_with_workspace(&*context.workspace) {
130        // Store the original phase for reference (it will be overwritten below)
131        let _original_phase = context.phase;
132
133        // Update existing checkpoint to Interrupted phase with current progress
134        checkpoint.phase = PipelinePhase::Interrupted;
135        checkpoint.iteration = context.iteration;
136        checkpoint.total_iterations = context.total_iterations;
137        checkpoint.reviewer_pass = context.reviewer_pass;
138        checkpoint.total_reviewer_passes = context.total_reviewer_passes;
139        checkpoint.actual_developer_runs = context.run_context.actual_developer_runs;
140        checkpoint.actual_reviewer_runs = context.run_context.actual_reviewer_runs;
141        checkpoint.execution_history = Some(context.execution_history.clone());
142        checkpoint.prompt_history = Some(context.prompt_history.clone());
143        save_checkpoint_with_workspace(&*context.workspace, &checkpoint)?;
144    } else {
145        // No checkpoint exists yet - this is early interruption
146        // We can't save a full checkpoint without config/registry access
147        // Just save a minimal checkpoint marker
148        eprintln!("Note: Interrupted before first checkpoint. Minimal state saved.");
149    }
150
151    Ok(())
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use crate::workspace::MemoryWorkspace;
158    use std::sync::Arc;
159
160    #[test]
161    fn test_interrupt_context_creation() {
162        let workspace = Arc::new(MemoryWorkspace::new_test());
163        let context = InterruptContext {
164            phase: crate::checkpoint::PipelinePhase::Development,
165            iteration: 2,
166            total_iterations: 5,
167            reviewer_pass: 0,
168            total_reviewer_passes: 2,
169            run_context: crate::checkpoint::RunContext::new(),
170            execution_history: crate::checkpoint::ExecutionHistory::new(),
171            prompt_history: std::collections::HashMap::new(),
172            workspace,
173        };
174
175        assert_eq!(context.phase, crate::checkpoint::PipelinePhase::Development);
176        assert_eq!(context.iteration, 2);
177        assert_eq!(context.total_iterations, 5);
178    }
179
180    #[test]
181    fn test_set_and_clear_interrupt_context() {
182        let workspace = Arc::new(MemoryWorkspace::new_test());
183        let context = InterruptContext {
184            phase: crate::checkpoint::PipelinePhase::Planning,
185            iteration: 1,
186            total_iterations: 3,
187            reviewer_pass: 0,
188            total_reviewer_passes: 1,
189            run_context: crate::checkpoint::RunContext::new(),
190            execution_history: crate::checkpoint::ExecutionHistory::new(),
191            prompt_history: std::collections::HashMap::new(),
192            workspace,
193        };
194
195        set_interrupt_context(context);
196        {
197            let ctx = INTERRUPT_CONTEXT.lock().unwrap();
198            assert!(ctx.is_some());
199            assert_eq!(
200                ctx.as_ref().unwrap().phase,
201                crate::checkpoint::PipelinePhase::Planning
202            );
203        }
204
205        clear_interrupt_context();
206        let ctx = INTERRUPT_CONTEXT.lock().unwrap();
207        assert!(ctx.is_none());
208    }
209}