""" Checkpoint Manager for Task Recovery Provides partial checkpointing and retry strategies for agent tasks. Enables recovery from failures without restarting entire workflows. """ import json import time from typing import Dict, List, Optional, Any from dataclasses import dataclass, asdict from datetime import datetime import redis import logging logger = logging.getLogger(__name__) @dataclass class Checkpoint: """Represents a task checkpoint""" task_id: str agent_id: str checkpoint_id: str timestamp: float phase: str # e.g., "analysis", "implementation", "testing" state: Dict[str, Any] # Serializable state data files_modified: List[str] git_commit: Optional[str] = None metadata: Optional[Dict[str, Any]] = None @dataclass class RetryPolicy: """Retry strategy configuration""" max_retries: int = 3 backoff_multiplier: float = 2.0 initial_delay_seconds: float = 1.0 max_delay_seconds: float = 60.0 retry_on_errors: List[str] = None # Error types to retry class CheckpointManager: """Manages task checkpoints and recovery""" def __init__(self, redis_client: redis.Redis): self.redis = redis_client self.checkpoint_ttl = 86400 * 7 # 7 days def save_checkpoint(self, checkpoint: Checkpoint) -> bool: """Save a checkpoint to Redis""" try: key = f"checkpoint:{checkpoint.task_id}:{checkpoint.checkpoint_id}" data = json.dumps(asdict(checkpoint)) self.redis.setex(key, self.checkpoint_ttl, data) # Add to task's checkpoint list list_key = f"checkpoints:{checkpoint.task_id}" self.redis.lpush(list_key, checkpoint.checkpoint_id) self.redis.expire(list_key, self.checkpoint_ttl) logger.info( f"Saved checkpoint {checkpoint.checkpoint_id} for task {checkpoint.task_id}" ) return True except Exception as e: logger.error(f"Failed to save checkpoint: {e}") return False def get_checkpoint(self, task_id: str, checkpoint_id: str) -> Optional[Checkpoint]: """Retrieve a specific checkpoint""" try: key = f"checkpoint:{task_id}:{checkpoint_id}" data = self.redis.get(key) if not data: return None checkpoint_dict = json.loads(data) return Checkpoint(**checkpoint_dict) except Exception as e: logger.error(f"Failed to retrieve checkpoint: {e}") return None def get_latest_checkpoint(self, task_id: str) -> Optional[Checkpoint]: """Get the most recent checkpoint for a task""" try: list_key = f"checkpoints:{task_id}" checkpoint_ids = self.redis.lrange(list_key, 0, 0) if not checkpoint_ids: return None checkpoint_id = checkpoint_ids[0].decode('utf-8') return self.get_checkpoint(task_id, checkpoint_id) except Exception as e: logger.error(f"Failed to get latest checkpoint: {e}") return None def list_checkpoints(self, task_id: str) -> List[str]: """List all checkpoint IDs for a task""" try: list_key = f"checkpoints:{task_id}" checkpoint_ids = self.redis.lrange(list_key, 0, -1) return [cid.decode('utf-8') for cid in checkpoint_ids] except Exception as e: logger.error(f"Failed to list checkpoints: {e}") return [] def delete_checkpoint(self, task_id: str, checkpoint_id: str) -> bool: """Delete a specific checkpoint""" try: key = f"checkpoint:{task_id}:{checkpoint_id}" self.redis.delete(key) # Remove from list list_key = f"checkpoints:{task_id}" self.redis.lrem(list_key, 0, checkpoint_id) logger.info(f"Deleted checkpoint {checkpoint_id} for task {task_id}") return True except Exception as e: logger.error(f"Failed to delete checkpoint: {e}") return False def cleanup_task_checkpoints(self, task_id: str) -> bool: """Delete all checkpoints for a task""" try: checkpoint_ids = self.list_checkpoints(task_id) for checkpoint_id in checkpoint_ids: self.delete_checkpoint(task_id, checkpoint_id) list_key = f"checkpoints:{task_id}" self.redis.delete(list_key) logger.info(f"Cleaned up all checkpoints for task {task_id}") return True except Exception as e: logger.error(f"Failed to cleanup checkpoints: {e}") return False class RetryManager: """Manages retry logic with exponential backoff""" def __init__(self, redis_client: redis.Redis): self.redis = redis_client def record_attempt(self, task_id: str, agent_id: str, error: Optional[str] = None): """Record a task attempt""" key = f"retry:{task_id}" attempt_data = { "agent_id": agent_id, "timestamp": time.time(), "error": error } self.redis.lpush(key, json.dumps(attempt_data)) self.redis.expire(key, 86400) # 24 hour TTL def get_attempt_count(self, task_id: str) -> int: """Get number of attempts for a task""" key = f"retry:{task_id}" return self.redis.llen(key) def should_retry(self, task_id: str, policy: RetryPolicy) -> bool: """Determine if task should be retried""" attempt_count = self.get_attempt_count(task_id) return attempt_count < policy.max_retries def get_retry_delay(self, task_id: str, policy: RetryPolicy) -> float: """Calculate delay before next retry (exponential backoff)""" attempt_count = self.get_attempt_count(task_id) delay = policy.initial_delay_seconds * (policy.backoff_multiplier ** attempt_count) return min(delay, policy.max_delay_seconds) def clear_attempts(self, task_id: str): """Clear retry history for a task""" key = f"retry:{task_id}" self.redis.delete(key) class RecoveryCoordinator: """Coordinates task recovery from checkpoints""" def __init__(self, checkpoint_manager: CheckpointManager, retry_manager: RetryManager): self.checkpoint_manager = checkpoint_manager self.retry_manager = retry_manager def recover_task( self, task_id: str, retry_policy: Optional[RetryPolicy] = None ) -> Optional[Dict[str, Any]]: """ Attempt to recover a failed task from its latest checkpoint Returns recovery instructions or None if recovery not possible """ if retry_policy is None: retry_policy = RetryPolicy() # Check if we should retry if not self.retry_manager.should_retry(task_id, retry_policy): logger.warning(f"Task {task_id} exceeded max retries") return None # Get latest checkpoint checkpoint = self.checkpoint_manager.get_latest_checkpoint(task_id) if not checkpoint: logger.warning(f"No checkpoint found for task {task_id}") return None # Calculate retry delay delay = self.retry_manager.get_retry_delay(task_id, retry_policy) recovery_plan = { "task_id": task_id, "checkpoint_id": checkpoint.checkpoint_id, "resume_phase": checkpoint.phase, "state": checkpoint.state, "files_modified": checkpoint.files_modified, "git_commit": checkpoint.git_commit, "retry_delay_seconds": delay, "attempt_number": self.retry_manager.get_attempt_count(task_id) + 1 } logger.info( f"Recovery plan created for task {task_id} from checkpoint {checkpoint.checkpoint_id}" ) return recovery_plan def mark_recovery_success(self, task_id: str): """Mark a task as successfully recovered""" self.retry_manager.clear_attempts(task_id) logger.info(f"Task {task_id} recovered successfully") def mark_recovery_failure(self, task_id: str, agent_id: str, error: str): """Record a failed recovery attempt""" self.retry_manager.record_attempt(task_id, agent_id, error) logger.warning(f"Recovery attempt failed for task {task_id}: {error}")