299 lines
9.8 KiB
Python
299 lines
9.8 KiB
Python
|
"""
|
||
|
|
Task and state management for AI agents.
|
||
|
|
"""
|
||
|
|
import json
|
||
|
|
import fcntl
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Dict, List, Optional, Any
|
||
|
|
from dataclasses import dataclass, asdict
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class Task:
|
||
|
|
"""Represents a single task."""
|
||
|
|
id: str
|
||
|
|
title: str
|
||
|
|
priority: str # 'high', 'medium', 'low'
|
||
|
|
status: str # 'pending', 'in_progress', 'completed', 'failed', 'skipped'
|
||
|
|
estimated_hours: float
|
||
|
|
description: str
|
||
|
|
acceptance_criteria: List[str]
|
||
|
|
dependencies: List[str]
|
||
|
|
files_to_modify: List[str]
|
||
|
|
component: str = "web_v2"
|
||
|
|
|
||
|
|
def to_dict(self) -> Dict[str, Any]:
|
||
|
|
return asdict(self)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_dict(cls, data: Dict[str, Any]) -> "Task":
|
||
|
|
return cls(**data)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class AgentState:
|
||
|
|
"""Represents the current state of an agent."""
|
||
|
|
component: str
|
||
|
|
current_task_id: Optional[str]
|
||
|
|
iteration_count: int
|
||
|
|
max_iterations: int
|
||
|
|
started_at: Optional[str]
|
||
|
|
last_commit: Optional[str]
|
||
|
|
retry_count: int
|
||
|
|
max_retries: int
|
||
|
|
status: str # 'idle', 'in_progress', 'completed', 'error'
|
||
|
|
errors: List[str]
|
||
|
|
completed_tasks: List[str]
|
||
|
|
skipped_tasks: List[str]
|
||
|
|
|
||
|
|
def to_dict(self) -> Dict[str, Any]:
|
||
|
|
return asdict(self)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_dict(cls, data: Dict[str, Any]) -> "AgentState":
|
||
|
|
return cls(**data)
|
||
|
|
|
||
|
|
|
||
|
|
class TaskManager:
|
||
|
|
"""Manages tasks and agent state."""
|
||
|
|
|
||
|
|
def __init__(self, task_list_path: Path, agent_state_path: Path):
|
||
|
|
self.task_list_path = task_list_path
|
||
|
|
self.agent_state_path = agent_state_path
|
||
|
|
|
||
|
|
def _read_json_file(self, path: Path) -> Dict[str, Any]:
|
||
|
|
"""Read and parse JSON file with file locking."""
|
||
|
|
with open(path, 'r') as f:
|
||
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_SH)
|
||
|
|
try:
|
||
|
|
return json.load(f)
|
||
|
|
finally:
|
||
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
||
|
|
|
||
|
|
def _write_json_file(self, path: Path, data: Dict[str, Any]) -> None:
|
||
|
|
"""Write JSON file with file locking."""
|
||
|
|
with open(path, 'w') as f:
|
||
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
|
||
|
|
try:
|
||
|
|
json.dump(data, f, indent=2)
|
||
|
|
finally:
|
||
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
||
|
|
|
||
|
|
def load_tasks(self) -> List[Task]:
|
||
|
|
"""Load tasks from task list file."""
|
||
|
|
data = self._read_json_file(self.task_list_path)
|
||
|
|
return [Task.from_dict(task_data) for task_data in data.get("tasks", [])]
|
||
|
|
|
||
|
|
def save_tasks(self, tasks: List[Task]) -> None:
|
||
|
|
"""Save tasks to task list file."""
|
||
|
|
data = self._read_json_file(self.task_list_path)
|
||
|
|
data["tasks"] = [task.to_dict() for task in tasks]
|
||
|
|
self._write_json_file(self.task_list_path, data)
|
||
|
|
|
||
|
|
def load_state(self) -> AgentState:
|
||
|
|
"""Load agent state from file."""
|
||
|
|
data = self._read_json_file(self.agent_state_path)
|
||
|
|
return AgentState.from_dict(data)
|
||
|
|
|
||
|
|
def save_state(self, state: AgentState) -> None:
|
||
|
|
"""Save agent state to file."""
|
||
|
|
self._write_json_file(self.agent_state_path, state.to_dict())
|
||
|
|
|
||
|
|
def get_next_task(self) -> Optional[Task]:
|
||
|
|
"""Get next pending task by priority."""
|
||
|
|
tasks = self.load_tasks()
|
||
|
|
state = self.load_state()
|
||
|
|
|
||
|
|
# Filter pending tasks
|
||
|
|
pending_tasks = [
|
||
|
|
task for task in tasks
|
||
|
|
if task.status == "pending" and task.id not in state.skipped_tasks
|
||
|
|
]
|
||
|
|
|
||
|
|
if not pending_tasks:
|
||
|
|
return None
|
||
|
|
|
||
|
|
# Check dependencies
|
||
|
|
completed_task_ids = set(state.completed_tasks)
|
||
|
|
|
||
|
|
def can_start_task(task: Task) -> bool:
|
||
|
|
"""Check if task dependencies are satisfied."""
|
||
|
|
if not task.dependencies:
|
||
|
|
return True
|
||
|
|
|
||
|
|
# Parse dependencies (format: "component:TASK-ID" or "TASK-ID")
|
||
|
|
for dep in task.dependencies:
|
||
|
|
if ":" in dep:
|
||
|
|
dep_component, dep_id = dep.split(":", 1)
|
||
|
|
# TODO: Check other component's state
|
||
|
|
# For now, only check current component
|
||
|
|
if dep_component == state.component and dep_id not in completed_task_ids:
|
||
|
|
return False
|
||
|
|
else:
|
||
|
|
if dep not in completed_task_ids:
|
||
|
|
return False
|
||
|
|
|
||
|
|
return True
|
||
|
|
|
||
|
|
# Filter by dependencies
|
||
|
|
ready_tasks = [task for task in pending_tasks if can_start_task(task)]
|
||
|
|
|
||
|
|
if not ready_tasks:
|
||
|
|
return None
|
||
|
|
|
||
|
|
# Sort by priority: high > medium > low
|
||
|
|
priority_order = {"high": 0, "medium": 1, "low": 2}
|
||
|
|
ready_tasks.sort(key=lambda t: priority_order.get(t.priority, 3))
|
||
|
|
|
||
|
|
return ready_tasks[0]
|
||
|
|
|
||
|
|
def update_task_status(self, task_id: str, status: str) -> None:
|
||
|
|
"""Update task status."""
|
||
|
|
tasks = self.load_tasks()
|
||
|
|
for task in tasks:
|
||
|
|
if task.id == task_id:
|
||
|
|
task.status = status
|
||
|
|
break
|
||
|
|
self.save_tasks(tasks)
|
||
|
|
|
||
|
|
def mark_task_completed(self, task_id: str) -> None:
|
||
|
|
"""Mark task as completed."""
|
||
|
|
self.update_task_status(task_id, "completed")
|
||
|
|
|
||
|
|
state = self.load_state()
|
||
|
|
if task_id not in state.completed_tasks:
|
||
|
|
state.completed_tasks.append(task_id)
|
||
|
|
state.current_task_id = None
|
||
|
|
state.retry_count = 0
|
||
|
|
self.save_state(state)
|
||
|
|
|
||
|
|
def mark_task_failed(self, task_id: str, error: str) -> None:
|
||
|
|
"""Mark task as failed."""
|
||
|
|
self.update_task_status(task_id, "failed")
|
||
|
|
|
||
|
|
state = self.load_state()
|
||
|
|
state.errors.append(f"{task_id}: {error}")
|
||
|
|
state.current_task_id = None
|
||
|
|
self.save_state(state)
|
||
|
|
|
||
|
|
def mark_task_skipped(self, task_id: str, reason: str) -> None:
|
||
|
|
"""Mark task as skipped."""
|
||
|
|
self.update_task_status(task_id, "skipped")
|
||
|
|
|
||
|
|
state = self.load_state()
|
||
|
|
if task_id not in state.skipped_tasks:
|
||
|
|
state.skipped_tasks.append(task_id)
|
||
|
|
state.errors.append(f"{task_id} skipped: {reason}")
|
||
|
|
state.current_task_id = None
|
||
|
|
self.save_state(state)
|
||
|
|
|
||
|
|
def start_task(self, task: Task) -> None:
|
||
|
|
"""Mark task as started."""
|
||
|
|
self.update_task_status(task.id, "in_progress")
|
||
|
|
|
||
|
|
state = self.load_state()
|
||
|
|
state.current_task_id = task.id
|
||
|
|
state.status = "in_progress"
|
||
|
|
if not state.started_at:
|
||
|
|
state.started_at = datetime.now(timezone.utc).isoformat()
|
||
|
|
state.iteration_count += 1
|
||
|
|
self.save_state(state)
|
||
|
|
|
||
|
|
def increment_retry(self) -> int:
|
||
|
|
"""Increment retry count and return new value."""
|
||
|
|
state = self.load_state()
|
||
|
|
state.retry_count += 1
|
||
|
|
self.save_state(state)
|
||
|
|
return state.retry_count
|
||
|
|
|
||
|
|
def is_agent_running(self) -> bool:
|
||
|
|
"""Check if agent is currently running."""
|
||
|
|
state = self.load_state()
|
||
|
|
|
||
|
|
if state.status != "in_progress":
|
||
|
|
return False
|
||
|
|
|
||
|
|
# Check if started too long ago (stale lock)
|
||
|
|
if state.started_at:
|
||
|
|
started = datetime.fromisoformat(state.started_at)
|
||
|
|
now = datetime.now(timezone.utc)
|
||
|
|
hours_elapsed = (now - started).total_seconds() / 3600
|
||
|
|
|
||
|
|
if hours_elapsed > 2: # Consider stale after 2 hours
|
||
|
|
return False
|
||
|
|
|
||
|
|
return True
|
||
|
|
|
||
|
|
def reset_state(self) -> None:
|
||
|
|
"""Reset agent state to idle."""
|
||
|
|
state = self.load_state()
|
||
|
|
state.status = "idle"
|
||
|
|
state.current_task_id = None
|
||
|
|
state.started_at = None
|
||
|
|
state.iteration_count = 0
|
||
|
|
state.retry_count = 0
|
||
|
|
self.save_state(state)
|
||
|
|
|
||
|
|
def all_tasks_completed(self) -> bool:
|
||
|
|
"""Check if all tasks are completed."""
|
||
|
|
tasks = self.load_tasks()
|
||
|
|
return all(task.status in ["completed", "skipped"] for task in tasks)
|
||
|
|
|
||
|
|
|
||
|
|
class GlobalLock:
|
||
|
|
"""Manages global lock for shared resource access."""
|
||
|
|
|
||
|
|
def __init__(self, lock_file: Path):
|
||
|
|
self.lock_file = lock_file
|
||
|
|
|
||
|
|
def _read_lock(self) -> Dict[str, Any]:
|
||
|
|
"""Read lock file."""
|
||
|
|
with open(self.lock_file, 'r') as f:
|
||
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_SH)
|
||
|
|
try:
|
||
|
|
return json.load(f)
|
||
|
|
finally:
|
||
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
||
|
|
|
||
|
|
def _write_lock(self, data: Dict[str, Any]) -> None:
|
||
|
|
"""Write lock file."""
|
||
|
|
with open(self.lock_file, 'w') as f:
|
||
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
|
||
|
|
try:
|
||
|
|
json.dump(data, f, indent=2)
|
||
|
|
finally:
|
||
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
||
|
|
|
||
|
|
def acquire(self, component: str, reason: str = "Working on shared resources") -> bool:
|
||
|
|
"""Acquire global lock."""
|
||
|
|
lock_data = self._read_lock()
|
||
|
|
|
||
|
|
if lock_data.get("locked"):
|
||
|
|
return False
|
||
|
|
|
||
|
|
lock_data["locked"] = True
|
||
|
|
lock_data["locked_by"] = component
|
||
|
|
lock_data["locked_at"] = datetime.now(timezone.utc).isoformat()
|
||
|
|
lock_data["reason"] = reason
|
||
|
|
|
||
|
|
self._write_lock(lock_data)
|
||
|
|
return True
|
||
|
|
|
||
|
|
def release(self, component: str) -> None:
|
||
|
|
"""Release global lock."""
|
||
|
|
lock_data = self._read_lock()
|
||
|
|
|
||
|
|
if lock_data.get("locked_by") == component:
|
||
|
|
lock_data["locked"] = False
|
||
|
|
lock_data["locked_by"] = None
|
||
|
|
lock_data["locked_at"] = None
|
||
|
|
lock_data["reason"] = None
|
||
|
|
self._write_lock(lock_data)
|
||
|
|
|
||
|
|
def is_locked(self) -> bool:
|
||
|
|
"""Check if global lock is held."""
|
||
|
|
lock_data = self._read_lock()
|
||
|
|
return lock_data.get("locked", False)
|
||
|
|
|