559 lines
17 KiB
Python
559 lines
17 KiB
Python
"""
|
|
Work Scheduler for Allegro-Primus
|
|
Intelligent task scheduling based on Gitea issues and priorities.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import logging
|
|
import asyncio
|
|
from typing import List, Dict, Any, Optional, Callable
|
|
from dataclasses import dataclass, asdict
|
|
from datetime import datetime, timedelta
|
|
from enum import Enum
|
|
import time
|
|
|
|
from gitea_client import GiteaClient, Issue, IssueState, IssuePriority
|
|
from repo_manager import RepoManager
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TaskStatus(str, Enum):
|
|
PENDING = "pending"
|
|
IN_PROGRESS = "in_progress"
|
|
BLOCKED = "blocked"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
|
|
|
|
@dataclass
|
|
class WorkTask:
|
|
"""Represents a scheduled work task."""
|
|
issue: Issue
|
|
owner: str
|
|
repo: str
|
|
status: TaskStatus
|
|
priority_score: float
|
|
estimated_hours: Optional[float] = None
|
|
assigned_to: Optional[str] = None
|
|
branch_name: Optional[str] = None
|
|
pr_number: Optional[int] = None
|
|
started_at: Optional[str] = None
|
|
completed_at: Optional[str] = None
|
|
notes: List[str] = None
|
|
|
|
def __post_init__(self):
|
|
if self.notes is None:
|
|
self.notes = []
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
data = asdict(self)
|
|
data['issue'] = asdict(self.issue)
|
|
return data
|
|
|
|
|
|
@dataclass
|
|
class WorkCycle:
|
|
"""Represents a scheduled work cycle."""
|
|
cycle_id: str
|
|
started_at: str
|
|
tasks: List[WorkTask]
|
|
status: str
|
|
completed_tasks: int = 0
|
|
|
|
|
|
class WorkScheduler:
|
|
"""
|
|
Intelligent work scheduler for autonomous issue resolution.
|
|
Queries Gitea, prioritizes tasks, and manages work cycles.
|
|
"""
|
|
|
|
# Priority weights for scoring
|
|
PRIORITY_WEIGHTS = {
|
|
"P0": 1000, # Critical
|
|
"P1": 100, # High
|
|
"P2": 10, # Medium
|
|
"P3": 1 # Low
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
gitea_client: Optional[GiteaClient] = None,
|
|
repo_manager: Optional[RepoManager] = None,
|
|
work_dir: str = "./work",
|
|
max_concurrent: int = 3
|
|
):
|
|
"""
|
|
Initialize WorkScheduler.
|
|
|
|
Args:
|
|
gitea_client: Gitea API client
|
|
repo_manager: Repository manager for git ops
|
|
work_dir: Directory for work state
|
|
max_concurrent: Maximum concurrent tasks
|
|
"""
|
|
self.gitea = gitea_client or GiteaClient()
|
|
self.repo_manager = repo_manager or RepoManager()
|
|
self.work_dir = work_dir
|
|
self.max_concurrent = max_concurrent
|
|
self.current_cycle: Optional[WorkCycle] = None
|
|
self.task_handlers: Dict[str, Callable] = {}
|
|
|
|
os.makedirs(work_dir, exist_ok=True)
|
|
logger.info(f"WorkScheduler initialized (max_concurrent={max_concurrent})")
|
|
|
|
def register_task_handler(
|
|
self,
|
|
task_type: str,
|
|
handler: Callable[[WorkTask], Any]
|
|
) -> None:
|
|
"""
|
|
Register a handler function for specific task types.
|
|
|
|
Args:
|
|
task_type: Type identifier (e.g., "bug", "feature", "refactor")
|
|
handler: Function to handle tasks of this type
|
|
"""
|
|
self.task_handlers[task_type] = handler
|
|
logger.info(f"Registered handler for task type: {task_type}")
|
|
|
|
def _calculate_priority_score(self, issue: Issue) -> float:
|
|
"""
|
|
Calculate priority score for an issue.
|
|
|
|
Higher score = higher priority.
|
|
Considers:
|
|
- Priority label (P0, P1, P2, P3)
|
|
- Age of issue
|
|
- Assignee status
|
|
- Comment activity
|
|
"""
|
|
score = 0.0
|
|
|
|
# Base priority from label
|
|
priority = issue.priority or "P3"
|
|
score += self.PRIORITY_WEIGHTS.get(priority, 1)
|
|
|
|
# Age factor: older issues get slight boost
|
|
try:
|
|
created = datetime.fromisoformat(issue.created_at.replace('Z', '+00:00'))
|
|
age_days = (datetime.now(created.tzinfo) - created).days
|
|
score += min(age_days * 0.1, 10) # Max 10 points for age
|
|
except:
|
|
pass
|
|
|
|
# Unassigned issues get priority
|
|
if not issue.assignee:
|
|
score += 5
|
|
|
|
return score
|
|
|
|
def fetch_priority_issues(
|
|
self,
|
|
owner: str,
|
|
repo: str,
|
|
priorities: List[str] = None
|
|
) -> List[Issue]:
|
|
"""
|
|
Fetch high-priority open issues from Gitea.
|
|
|
|
Args:
|
|
owner: Repository owner
|
|
repo: Repository name
|
|
priorities: Priority labels to fetch (default: P0, P1)
|
|
|
|
Returns:
|
|
List of priority issues sorted by score
|
|
"""
|
|
if priorities is None:
|
|
priorities = ["P0", "P1"]
|
|
|
|
logger.info(f"Fetching issues with priorities: {priorities}")
|
|
|
|
try:
|
|
issues = self.gitea.get_priority_issues(
|
|
owner, repo,
|
|
priorities=priorities,
|
|
state=IssueState.OPEN
|
|
)
|
|
logger.info(f"Found {len(issues)} priority issues")
|
|
return issues
|
|
except Exception as e:
|
|
logger.error(f"Failed to fetch issues: {e}")
|
|
return []
|
|
|
|
def prioritize_tasks(
|
|
self,
|
|
issues: List[Issue],
|
|
owner: str,
|
|
repo: str
|
|
) -> List[WorkTask]:
|
|
"""
|
|
Convert issues to prioritized work tasks.
|
|
|
|
Args:
|
|
issues: List of issues from Gitea
|
|
owner: Repository owner
|
|
repo: Repository name
|
|
|
|
Returns:
|
|
List of WorkTask sorted by priority score
|
|
"""
|
|
tasks = []
|
|
|
|
for issue in issues:
|
|
score = self._calculate_priority_score(issue)
|
|
task = WorkTask(
|
|
issue=issue,
|
|
owner=owner,
|
|
repo=repo,
|
|
status=TaskStatus.PENDING,
|
|
priority_score=score
|
|
)
|
|
tasks.append(task)
|
|
|
|
# Sort by priority score (descending)
|
|
tasks.sort(key=lambda t: t.priority_score, reverse=True)
|
|
|
|
logger.info(f"Prioritized {len(tasks)} tasks")
|
|
return tasks
|
|
|
|
def create_work_cycle(
|
|
self,
|
|
owner: str,
|
|
repo: str,
|
|
cycle_name: Optional[str] = None
|
|
) -> WorkCycle:
|
|
"""
|
|
Create a new work cycle with prioritized tasks.
|
|
|
|
Args:
|
|
owner: Repository owner
|
|
repo: Repository name
|
|
cycle_name: Optional cycle identifier
|
|
|
|
Returns:
|
|
WorkCycle with scheduled tasks
|
|
"""
|
|
# Fetch and prioritize issues
|
|
issues = self.fetch_priority_issues(owner, repo)
|
|
tasks = self.prioritize_tasks(issues, owner, repo)
|
|
|
|
# Create cycle
|
|
cycle_id = cycle_name or f"cycle-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
|
|
cycle = WorkCycle(
|
|
cycle_id=cycle_id,
|
|
started_at=datetime.now().isoformat(),
|
|
tasks=tasks,
|
|
status="created"
|
|
)
|
|
|
|
self.current_cycle = cycle
|
|
self._save_cycle_state(cycle)
|
|
|
|
logger.info(f"Created work cycle '{cycle_id}' with {len(tasks)} tasks")
|
|
return cycle
|
|
|
|
def start_task(self, task: WorkTask) -> bool:
|
|
"""
|
|
Mark a task as in-progress and update Gitea.
|
|
|
|
Args:
|
|
task: WorkTask to start
|
|
|
|
Returns:
|
|
True if started successfully
|
|
"""
|
|
try:
|
|
task.status = TaskStatus.IN_PROGRESS
|
|
task.started_at = datetime.now().isoformat()
|
|
|
|
# Create work branch
|
|
branch = self.repo_manager.create_work_branch(
|
|
task.owner,
|
|
task.repo,
|
|
task.issue.number,
|
|
task.issue.title[:20]
|
|
)
|
|
task.branch_name = branch
|
|
|
|
# Update issue in Gitea - add comment and assign
|
|
self.gitea.add_issue_comment(
|
|
task.owner,
|
|
task.repo,
|
|
task.issue.number,
|
|
f"🤖 Allegro-Primus is working on this issue.\n"
|
|
f"Branch: `{branch}`\n"
|
|
f"Started: {task.started_at}"
|
|
)
|
|
|
|
# Try to self-assign if not assigned
|
|
if not task.issue.assignee:
|
|
try:
|
|
user = self.gitea.get_user()
|
|
self.gitea.update_issue(
|
|
task.owner,
|
|
task.repo,
|
|
task.issue.number,
|
|
assignees=[user['login']]
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Could not self-assign: {e}")
|
|
|
|
task.notes.append(f"Started work on branch {branch}")
|
|
logger.info(f"Started task for issue #{task.issue.number}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
task.status = TaskStatus.FAILED
|
|
task.notes.append(f"Failed to start: {e}")
|
|
logger.error(f"Failed to start task: {e}")
|
|
return False
|
|
|
|
def complete_task(
|
|
self,
|
|
task: WorkTask,
|
|
success: bool = True,
|
|
pr_number: Optional[int] = None
|
|
) -> None:
|
|
"""
|
|
Mark a task as completed.
|
|
|
|
Args:
|
|
task: WorkTask to complete
|
|
success: Whether task completed successfully
|
|
pr_number: PR number if created
|
|
"""
|
|
task.completed_at = datetime.now().isoformat()
|
|
task.pr_number = pr_number
|
|
|
|
if success:
|
|
task.status = TaskStatus.COMPLETED
|
|
|
|
# Update issue
|
|
comment = f"✅ Allegro-Primus has completed work on this issue."
|
|
if pr_number:
|
|
comment += f"\nPull Request: #{pr_number}"
|
|
|
|
self.gitea.add_issue_comment(
|
|
task.owner,
|
|
task.repo,
|
|
task.issue.number,
|
|
comment
|
|
)
|
|
|
|
logger.info(f"Completed task for issue #{task.issue.number}")
|
|
else:
|
|
task.status = TaskStatus.FAILED
|
|
self.gitea.add_issue_comment(
|
|
task.owner,
|
|
task.repo,
|
|
task.issue.number,
|
|
f"❌ Task failed. Check logs for details."
|
|
)
|
|
logger.error(f"Task failed for issue #{task.issue.number}")
|
|
|
|
def run_work_cycle(
|
|
self,
|
|
owner: str,
|
|
repo: str,
|
|
task_processor: Optional[Callable[[WorkTask], bool]] = None
|
|
) -> WorkCycle:
|
|
"""
|
|
Execute a complete work cycle.
|
|
|
|
Args:
|
|
owner: Repository owner
|
|
repo: Repository name
|
|
task_processor: Function to process each task
|
|
|
|
Returns:
|
|
Completed WorkCycle
|
|
"""
|
|
cycle = self.create_work_cycle(owner, repo)
|
|
cycle.status = "running"
|
|
|
|
if not cycle.tasks:
|
|
logger.info("No tasks to process")
|
|
cycle.status = "completed"
|
|
return cycle
|
|
|
|
# Process tasks up to max_concurrent
|
|
active_tasks = cycle.tasks[:self.max_concurrent]
|
|
|
|
for task in active_tasks:
|
|
# Start task
|
|
if not self.start_task(task):
|
|
continue
|
|
|
|
# Process if handler provided
|
|
if task_processor:
|
|
try:
|
|
success = task_processor(task)
|
|
self.complete_task(task, success=success)
|
|
except Exception as e:
|
|
logger.error(f"Task processor error: {e}")
|
|
self.complete_task(task, success=False)
|
|
|
|
cycle.completed_tasks += 1
|
|
self._save_cycle_state(cycle)
|
|
|
|
cycle.status = "completed"
|
|
self._save_cycle_state(cycle)
|
|
|
|
logger.info(f"Work cycle completed. {cycle.completed_tasks}/{len(cycle.tasks)} tasks done.")
|
|
return cycle
|
|
|
|
def _save_cycle_state(self, cycle: WorkCycle) -> None:
|
|
"""Save work cycle state to disk."""
|
|
filepath = os.path.join(self.work_dir, f"{cycle.cycle_id}.json")
|
|
with open(filepath, 'w') as f:
|
|
json.dump({
|
|
"cycle_id": cycle.cycle_id,
|
|
"started_at": cycle.started_at,
|
|
"status": cycle.status,
|
|
"completed_tasks": cycle.completed_tasks,
|
|
"tasks": [t.to_dict() for t in cycle.tasks]
|
|
}, f, indent=2)
|
|
|
|
def load_cycle(self, cycle_id: str) -> Optional[WorkCycle]:
|
|
"""Load a work cycle from disk."""
|
|
filepath = os.path.join(self.work_dir, f"{cycle_id}.json")
|
|
if not os.path.exists(filepath):
|
|
return None
|
|
|
|
with open(filepath, 'r') as f:
|
|
data = json.load(f)
|
|
|
|
# Note: Full reconstruction would require Issue objects
|
|
# This is a simplified load for monitoring
|
|
return WorkCycle(
|
|
cycle_id=data["cycle_id"],
|
|
started_at=data["started_at"],
|
|
tasks=[], # Would need to reconstruct from data
|
|
status=data["status"],
|
|
completed_tasks=data.get("completed_tasks", 0)
|
|
)
|
|
|
|
def get_cycle_history(self) -> List[Dict[str, Any]]:
|
|
"""Get list of all saved work cycles."""
|
|
cycles = []
|
|
for filename in os.listdir(self.work_dir):
|
|
if filename.endswith('.json'):
|
|
cycle_id = filename[:-5]
|
|
cycle = self.load_cycle(cycle_id)
|
|
if cycle:
|
|
cycles.append({
|
|
"cycle_id": cycle.cycle_id,
|
|
"started_at": cycle.started_at,
|
|
"status": cycle.status,
|
|
"completed_tasks": cycle.completed_tasks
|
|
})
|
|
return sorted(cycles, key=lambda x: x["started_at"], reverse=True)
|
|
|
|
def get_pending_tasks(self, owner: str, repo: str) -> List[WorkTask]:
|
|
"""Get list of pending high-priority tasks."""
|
|
issues = self.fetch_priority_issues(owner, repo)
|
|
return self.prioritize_tasks(issues, owner, repo)
|
|
|
|
def auto_schedule(
|
|
self,
|
|
owner: str,
|
|
repo: str,
|
|
interval_minutes: int = 60
|
|
):
|
|
"""
|
|
Run continuous scheduling loop.
|
|
|
|
Args:
|
|
owner: Repository owner
|
|
repo: Repository name
|
|
interval_minutes: Minutes between cycles
|
|
"""
|
|
logger.info(f"Starting auto-scheduler (interval={interval_minutes}min)")
|
|
|
|
while True:
|
|
try:
|
|
self.run_work_cycle(owner, repo)
|
|
logger.info(f"Sleeping for {interval_minutes} minutes...")
|
|
time.sleep(interval_minutes * 60)
|
|
except KeyboardInterrupt:
|
|
logger.info("Scheduler stopped by user")
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Scheduler error: {e}")
|
|
time.sleep(60) # Short sleep on error
|
|
|
|
def update_issue_status(
|
|
self,
|
|
owner: str,
|
|
repo: str,
|
|
issue_number: int,
|
|
status: str,
|
|
comment: Optional[str] = None
|
|
) -> None:
|
|
"""
|
|
Update issue status with optional comment.
|
|
|
|
Args:
|
|
owner: Repository owner
|
|
repo: Repository name
|
|
issue_number: Issue number
|
|
status: New status label
|
|
comment: Optional comment to add
|
|
"""
|
|
try:
|
|
# Get current issue
|
|
issue = self.gitea.get_issue(owner, repo, issue_number)
|
|
|
|
# Update labels to reflect status
|
|
current_labels = [l["name"] for l in issue.labels]
|
|
|
|
# Remove old status labels
|
|
status_labels = ["in-progress", "pending-review", "blocked", "completed"]
|
|
new_labels = [l for l in current_labels if l not in status_labels]
|
|
|
|
# Add new status
|
|
if status:
|
|
new_labels.append(status)
|
|
|
|
self.gitea.update_issue(
|
|
owner, repo, issue_number,
|
|
labels=list(set(new_labels))
|
|
)
|
|
|
|
# Add comment
|
|
if comment:
|
|
self.gitea.add_issue_comment(owner, repo, issue_number, comment)
|
|
|
|
logger.info(f"Updated issue #{issue_number} status to {status}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to update issue status: {e}")
|
|
|
|
|
|
# ==================== Example Usage ====================
|
|
|
|
if __name__ == "__main__":
|
|
# Initialize scheduler
|
|
scheduler = WorkScheduler()
|
|
|
|
# Example: Create and run a work cycle
|
|
# cycle = scheduler.create_work_cycle("myuser", "myrepo")
|
|
# print(f"Scheduled {len(cycle.tasks)} tasks")
|
|
|
|
# Example: Process a specific task
|
|
# def my_processor(task: WorkTask) -> bool:
|
|
# # Your task processing logic here
|
|
# # Return True for success, False for failure
|
|
# return True
|
|
#
|
|
# scheduler.run_work_cycle("myuser", "myrepo", task_processor=my_processor)
|
|
|
|
# Example: Auto-schedule with continuous monitoring
|
|
# scheduler.auto_schedule("myuser", "myrepo", interval_minutes=30)
|
|
|
|
print("WorkScheduler example - configure with your repo details")
|