//! 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<Utc>,
}
/// A single progress entry from the progress log.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProgressEntry {
/// Timestamp of the entry.
pub timestamp: DateTime<Utc>,
/// 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<CheckpointInfo>,
/// Recent progress entries from progress.log.
pub recent_progress: Vec<ProgressEntry>,
/// 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<String> {
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<CheckpointInfo> {
// 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<ProgressEntry> {
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<I>) -> Option<ProgressEntry>
where
I: Iterator<Item = &'a str>,
{
// 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"));
}
}
