//! Context recovery for task resumption and restart. //! //! This module provides functionality to rebuild context when tasks resume or restart, //! ensuring Claude can quickly orient itself with the current state of the worktree. use std::path::Path; use std::process::Command; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; /// Summary of git status information. #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct GitStatusSummary { /// Number of staged files. pub staged: u32, /// Number of modified (unstaged) files. pub modified: u32, /// Number of untracked files. pub untracked: u32, /// Number of deleted files. pub deleted: u32, /// Brief description (e.g., "3 modified, 1 staged"). pub summary: String, } /// Information about a checkpoint (git commit). #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CheckpointInfo { /// Commit SHA (short form). pub sha: String, /// Commit message (first line). pub message: String, /// Commit timestamp. pub timestamp: DateTime, } /// A single progress entry from the progress log. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ProgressEntry { /// Timestamp of the entry. pub timestamp: DateTime, /// Entry type (e.g., "ITERATION", "TASK_START", "CHECKPOINT"). pub entry_type: String, /// Entry message or summary. pub message: String, } /// Context recovery data for a task. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ContextRecovery { /// Current git branch name. pub current_branch: String, /// Git status summary. pub git_status: GitStatusSummary, /// Last checkpoint information (most recent commit). pub last_checkpoint: Option, /// Recent progress entries from progress.log. pub recent_progress: Vec, /// Current phase of the task. pub current_phase: String, } impl ContextRecovery { /// Format the context recovery as a markdown header suitable for injection into prompts. pub fn to_markdown(&self) -> String { let mut output = String::new(); output.push_str("## Context Recovery\n"); output.push_str(&format!("- Current branch: {}\n", self.current_branch)); output.push_str(&format!("- Git status: {}\n", self.git_status.summary)); if let Some(ref checkpoint) = self.last_checkpoint { output.push_str(&format!( "- Last checkpoint: {} - {} ({})\n", checkpoint.sha, checkpoint.message, checkpoint.timestamp.format("%Y-%m-%d %H:%M:%S UTC") )); } else { output.push_str("- Last checkpoint: none\n"); } if self.recent_progress.is_empty() { output.push_str("- Progress log (recent): none\n"); } else { output.push_str("- Progress log (recent):\n"); for entry in &self.recent_progress { output.push_str(&format!( " - [{}] {}: {}\n", entry.timestamp.format("%H:%M:%S"), entry.entry_type, entry.message )); } } output.push_str(&format!("- Current phase: {}\n", self.current_phase)); output.push_str("\n---\n\n"); output } } /// Build context recovery information for a task worktree. /// /// # Arguments /// * `worktree_path` - Path to the task's worktree directory /// * `phase` - Current task phase (e.g., "execute", "review") /// * `max_progress_entries` - Maximum number of recent progress entries to include /// /// # Returns /// A formatted markdown string containing context recovery information. pub fn build_context_recovery( worktree_path: &Path, phase: &str, max_progress_entries: usize, ) -> String { let context = build_context_recovery_data(worktree_path, phase, max_progress_entries); context.to_markdown() } /// Build context recovery data structure for a task worktree. /// /// # Arguments /// * `worktree_path` - Path to the task's worktree directory /// * `phase` - Current task phase (e.g., "execute", "review") /// * `max_progress_entries` - Maximum number of recent progress entries to include /// /// # Returns /// A ContextRecovery struct containing all gathered information. pub fn build_context_recovery_data( worktree_path: &Path, phase: &str, max_progress_entries: usize, ) -> ContextRecovery { let current_branch = get_current_branch(worktree_path).unwrap_or_else(|| "unknown".to_string()); let git_status = get_git_status_summary(worktree_path); let last_checkpoint = get_last_checkpoint(worktree_path); let recent_progress = get_recent_progress(worktree_path, max_progress_entries); ContextRecovery { current_branch, git_status, last_checkpoint, recent_progress, current_phase: phase.to_string(), } } /// Get the current git branch name. fn get_current_branch(worktree_path: &Path) -> Option { let output = Command::new("git") .args(["rev-parse", "--abbrev-ref", "HEAD"]) .current_dir(worktree_path) .output() .ok()?; if output.status.success() { let branch = String::from_utf8_lossy(&output.stdout).trim().to_string(); if branch.is_empty() || branch == "HEAD" { // Detached HEAD state - try to get the commit SHA let sha_output = Command::new("git") .args(["rev-parse", "--short", "HEAD"]) .current_dir(worktree_path) .output() .ok()?; if sha_output.status.success() { Some(format!("(detached at {})", String::from_utf8_lossy(&sha_output.stdout).trim())) } else { Some("(detached)".to_string()) } } else { Some(branch) } } else { None } } /// Get a summary of the git status. fn get_git_status_summary(worktree_path: &Path) -> GitStatusSummary { let output = Command::new("git") .args(["status", "--porcelain"]) .current_dir(worktree_path) .output(); let output = match output { Ok(o) if o.status.success() => o, _ => return GitStatusSummary { summary: "unable to get git status".to_string(), ..Default::default() }, }; let status_output = String::from_utf8_lossy(&output.stdout); let mut staged = 0u32; let mut modified = 0u32; let mut untracked = 0u32; let mut deleted = 0u32; for line in status_output.lines() { if line.len() < 2 { continue; } let index_status = line.chars().next().unwrap_or(' '); let worktree_status = line.chars().nth(1).unwrap_or(' '); // Count staged changes (index column) match index_status { 'A' | 'M' | 'R' | 'C' => staged += 1, 'D' => { staged += 1; deleted += 1; } _ => {} } // Count unstaged changes (worktree column) match worktree_status { 'M' => modified += 1, 'D' => { modified += 1; deleted += 1; } '?' => untracked += 1, _ => {} } } // Build summary string let mut parts = Vec::new(); if staged > 0 { parts.push(format!("{} staged", staged)); } if modified > 0 { parts.push(format!("{} modified", modified)); } if untracked > 0 { parts.push(format!("{} untracked", untracked)); } let summary = if parts.is_empty() { "clean".to_string() } else { parts.join(", ") }; GitStatusSummary { staged, modified, untracked, deleted, summary, } } /// Get information about the last checkpoint (most recent commit). fn get_last_checkpoint(worktree_path: &Path) -> Option { // Get the most recent commit info let output = Command::new("git") .args([ "log", "-1", "--format=%h|%s|%aI", // short sha | subject | ISO 8601 date ]) .current_dir(worktree_path) .output() .ok()?; if !output.status.success() { return None; } let log_output = String::from_utf8_lossy(&output.stdout).trim().to_string(); let parts: Vec<&str> = log_output.splitn(3, '|').collect(); if parts.len() < 3 { return None; } let sha = parts[0].to_string(); let message = parts[1].to_string(); let timestamp = DateTime::parse_from_rfc3339(parts[2]) .ok()? .with_timezone(&Utc); Some(CheckpointInfo { sha, message, timestamp, }) } /// Get recent progress entries from the progress.log file. fn get_recent_progress(worktree_path: &Path, max_entries: usize) -> Vec { let progress_log_path = worktree_path.join("progress.log"); let content = match std::fs::read_to_string(&progress_log_path) { Ok(c) => c, Err(_) => return Vec::new(), }; let mut entries = Vec::new(); // Parse the progress.log file // Expected format: // [2024-01-15T10:30:00Z] ENTRY_TYPE // key: value // key: value // (blank line) let mut lines = content.lines().peekable(); while let Some(line) = lines.next() { // Look for timestamp lines if line.starts_with('[') && line.contains(']') { if let Some(entry) = parse_progress_entry(line, &mut lines) { entries.push(entry); } } } // Return the most recent entries (up to max_entries) let start = if entries.len() > max_entries { entries.len() - max_entries } else { 0 }; entries[start..].to_vec() } /// Parse a single progress entry from the log. fn parse_progress_entry<'a, I>(header_line: &str, lines: &mut std::iter::Peekable) -> Option where I: Iterator, { // Parse header: [2024-01-15T10:30:00Z] ENTRY_TYPE let close_bracket = header_line.find(']')?; let timestamp_str = &header_line[1..close_bracket]; let entry_type = header_line[close_bracket + 1..].trim().to_string(); let timestamp = DateTime::parse_from_rfc3339(timestamp_str) .ok()? .with_timezone(&Utc); // Collect message from following lines until blank line or next entry let mut message_parts = Vec::new(); while let Some(&next_line) = lines.peek() { if next_line.is_empty() || next_line.starts_with('#') { lines.next(); // consume blank/comment line break; } if next_line.starts_with('[') { // Next entry starts, don't consume break; } let line = lines.next().unwrap(); // Extract key-value or just the line if let Some(colon_pos) = line.find(':') { let key = line[..colon_pos].trim(); let value = line[colon_pos + 1..].trim(); if key == "progress" || key == "message" || key == "reason" { message_parts.push(value.to_string()); } } } let message = if message_parts.is_empty() { entry_type.clone() } else { message_parts.join("; ") }; Some(ProgressEntry { timestamp, entry_type, message, }) } #[cfg(test)] mod tests { use super::*; use std::fs; use tempfile::TempDir; fn setup_git_repo() -> TempDir { let temp_dir = TempDir::new().unwrap(); let path = temp_dir.path(); // Initialize git repo Command::new("git") .args(["init"]) .current_dir(path) .output() .unwrap(); // Configure git for commits Command::new("git") .args(["config", "user.email", "test@example.com"]) .current_dir(path) .output() .unwrap(); Command::new("git") .args(["config", "user.name", "Test User"]) .current_dir(path) .output() .unwrap(); // Create initial commit fs::write(path.join("README.md"), "# Test").unwrap(); Command::new("git") .args(["add", "README.md"]) .current_dir(path) .output() .unwrap(); Command::new("git") .args(["commit", "-m", "Initial commit"]) .current_dir(path) .output() .unwrap(); temp_dir } #[test] fn test_get_current_branch() { let temp_dir = setup_git_repo(); let branch = get_current_branch(temp_dir.path()); // Default branch could be main or master depending on git config assert!(branch.is_some()); let branch_name = branch.unwrap(); assert!(branch_name == "main" || branch_name == "master"); } #[test] fn test_git_status_clean() { let temp_dir = setup_git_repo(); let status = get_git_status_summary(temp_dir.path()); assert_eq!(status.staged, 0); assert_eq!(status.modified, 0); assert_eq!(status.untracked, 0); assert_eq!(status.summary, "clean"); } #[test] fn test_git_status_with_changes() { let temp_dir = setup_git_repo(); let path = temp_dir.path(); // Create untracked file fs::write(path.join("new_file.txt"), "new content").unwrap(); // Modify tracked file fs::write(path.join("README.md"), "# Modified").unwrap(); let status = get_git_status_summary(path); assert_eq!(status.modified, 1); assert_eq!(status.untracked, 1); assert!(status.summary.contains("modified")); assert!(status.summary.contains("untracked")); } #[test] fn test_get_last_checkpoint() { let temp_dir = setup_git_repo(); let checkpoint = get_last_checkpoint(temp_dir.path()); assert!(checkpoint.is_some()); let checkpoint = checkpoint.unwrap(); assert_eq!(checkpoint.message, "Initial commit"); assert!(!checkpoint.sha.is_empty()); } #[test] fn test_build_context_recovery() { let temp_dir = setup_git_repo(); let markdown = build_context_recovery(temp_dir.path(), "execute", 5); assert!(markdown.contains("## Context Recovery")); assert!(markdown.contains("Current branch:")); assert!(markdown.contains("Git status:")); assert!(markdown.contains("Last checkpoint:")); assert!(markdown.contains("Current phase: execute")); } #[test] fn test_progress_log_parsing() { let temp_dir = setup_git_repo(); let path = temp_dir.path(); // Create a mock progress.log let progress_content = r#"# progress.log # Auto-generated by Makima - DO NOT EDIT MANUALLY [2024-01-15T10:30:00Z] TASK_START task_id: test-123 task_name: Test Task [2024-01-15T10:35:00Z] ITERATION 1 progress: Analyzed codebase structure files_modified: 2 [2024-01-15T10:40:00Z] ITERATION 2 progress: Implemented first feature files_modified: 3 "#; fs::write(path.join("progress.log"), progress_content).unwrap(); let entries = get_recent_progress(path, 5); assert_eq!(entries.len(), 3); assert_eq!(entries[0].entry_type, "TASK_START"); assert_eq!(entries[1].entry_type, "ITERATION 1"); assert_eq!(entries[2].entry_type, "ITERATION 2"); } #[test] fn test_context_recovery_to_markdown() { let context = ContextRecovery { current_branch: "feature/test".to_string(), git_status: GitStatusSummary { staged: 1, modified: 2, untracked: 0, deleted: 0, summary: "1 staged, 2 modified".to_string(), }, last_checkpoint: Some(CheckpointInfo { sha: "abc1234".to_string(), message: "Added feature X".to_string(), timestamp: Utc::now(), }), recent_progress: vec![ ProgressEntry { timestamp: Utc::now(), entry_type: "ITERATION 1".to_string(), message: "Started implementation".to_string(), }, ], current_phase: "execute".to_string(), }; let markdown = context.to_markdown(); assert!(markdown.contains("## Context Recovery")); assert!(markdown.contains("feature/test")); assert!(markdown.contains("1 staged, 2 modified")); assert!(markdown.contains("abc1234")); assert!(markdown.contains("Added feature X")); assert!(markdown.contains("ITERATION 1")); assert!(markdown.contains("execute")); } }