425 lines
13 KiB
Python
Executable File
425 lines
13 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Task Runner for Allegro-Primus
|
|
Executes tasks in isolated worktrees with result capture
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from dataclasses import dataclass, asdict
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Callable, Any
|
|
from contextlib import contextmanager
|
|
|
|
|
|
@dataclass
|
|
class TaskResult:
|
|
"""Result of a task execution"""
|
|
task_id: str
|
|
issue_number: str
|
|
success: bool
|
|
return_code: int
|
|
stdout: str
|
|
stderr: str
|
|
execution_time: float
|
|
timestamp: str
|
|
artifacts: List[str]
|
|
commit_hash: Optional[str] = None
|
|
error: Optional[str] = None
|
|
|
|
def to_dict(self) -> Dict:
|
|
return asdict(self)
|
|
|
|
def to_json(self, indent: int = 2) -> str:
|
|
return json.dumps(self.to_dict(), indent=indent)
|
|
|
|
|
|
class TaskRunner:
|
|
"""Runs tasks in isolated worktree environments"""
|
|
|
|
def __init__(self, worktree_path: str, logs_dir: str = None):
|
|
self.worktree_path = Path(worktree_path)
|
|
self.logs_dir = Path(logs_dir) if logs_dir else self.worktree_path / ".ap-logs"
|
|
self.logs_dir.mkdir(exist_ok=True)
|
|
self.task_history: List[TaskResult] = []
|
|
|
|
def _run_git(self, args: List[str], cwd: str = None) -> subprocess.CompletedProcess:
|
|
"""Run git command in worktree"""
|
|
return subprocess.run(
|
|
["git"] + args,
|
|
capture_output=True,
|
|
text=True,
|
|
cwd=cwd or str(self.worktree_path)
|
|
)
|
|
|
|
def _write_log(self, task_id: str, content: str):
|
|
"""Write content to task log file"""
|
|
log_file = self.logs_dir / f"{task_id}.log"
|
|
log_file.write_text(content)
|
|
|
|
@contextmanager
|
|
def _capture_context(self):
|
|
"""Context manager for capturing output"""
|
|
yield
|
|
|
|
def run_command(self,
|
|
task_id: str,
|
|
command: List[str],
|
|
issue_number: str,
|
|
env: Dict[str, str] = None,
|
|
timeout: int = 300,
|
|
capture_output: bool = True) -> TaskResult:
|
|
"""
|
|
Run a shell command in the worktree
|
|
|
|
Args:
|
|
task_id: Unique identifier for this task
|
|
command: Command and arguments as list
|
|
issue_number: Associated issue number
|
|
env: Additional environment variables
|
|
timeout: Command timeout in seconds
|
|
capture_output: Whether to capture stdout/stderr
|
|
|
|
Returns:
|
|
TaskResult with execution details
|
|
"""
|
|
start_time = time.time()
|
|
timestamp = datetime.now().isoformat()
|
|
|
|
# Prepare environment
|
|
run_env = os.environ.copy()
|
|
run_env["AP_WORKTREE"] = str(self.worktree_path)
|
|
run_env["AP_ISSUE"] = issue_number
|
|
run_env["AP_TASK_ID"] = task_id
|
|
if env:
|
|
run_env.update(env)
|
|
|
|
try:
|
|
# Run command
|
|
result = subprocess.run(
|
|
command,
|
|
capture_output=capture_output,
|
|
text=True,
|
|
cwd=str(self.worktree_path),
|
|
env=run_env,
|
|
timeout=timeout
|
|
)
|
|
|
|
execution_time = time.time() - start_time
|
|
|
|
# Create result
|
|
task_result = TaskResult(
|
|
task_id=task_id,
|
|
issue_number=issue_number,
|
|
success=result.returncode == 0,
|
|
return_code=result.returncode,
|
|
stdout=result.stdout,
|
|
stderr=result.stderr,
|
|
execution_time=execution_time,
|
|
timestamp=timestamp,
|
|
artifacts=[]
|
|
)
|
|
|
|
# Write log
|
|
log_content = f"""Task: {task_id}
|
|
Issue: {issue_number}
|
|
Command: {' '.join(command)}
|
|
Timestamp: {timestamp}
|
|
Duration: {execution_time:.2f}s
|
|
Return Code: {result.returncode}
|
|
|
|
=== STDOUT ===
|
|
{result.stdout}
|
|
|
|
=== STDERR ===
|
|
{result.stderr}
|
|
"""
|
|
self._write_log(task_id, log_content)
|
|
|
|
except subprocess.TimeoutExpired as e:
|
|
execution_time = time.time() - start_time
|
|
task_result = TaskResult(
|
|
task_id=task_id,
|
|
issue_number=issue_number,
|
|
success=False,
|
|
return_code=-1,
|
|
stdout=e.stdout if e.stdout else "",
|
|
stderr=f"Timeout after {timeout}s",
|
|
execution_time=execution_time,
|
|
timestamp=timestamp,
|
|
artifacts=[],
|
|
error="Timeout"
|
|
)
|
|
except Exception as e:
|
|
execution_time = time.time() - start_time
|
|
task_result = TaskResult(
|
|
task_id=task_id,
|
|
issue_number=issue_number,
|
|
success=False,
|
|
return_code=-1,
|
|
stdout="",
|
|
stderr=str(e),
|
|
execution_time=execution_time,
|
|
timestamp=timestamp,
|
|
artifacts=[],
|
|
error=traceback.format_exc()
|
|
)
|
|
|
|
self.task_history.append(task_result)
|
|
return task_result
|
|
|
|
def run_script(self,
|
|
task_id: str,
|
|
script_content: str,
|
|
issue_number: str,
|
|
interpreter: str = "python3",
|
|
timeout: int = 300) -> TaskResult:
|
|
"""
|
|
Run a script in the worktree
|
|
|
|
Args:
|
|
task_id: Task identifier
|
|
script_content: Script to execute
|
|
issue_number: Associated issue
|
|
interpreter: Script interpreter
|
|
timeout: Execution timeout
|
|
"""
|
|
# Write script to temporary file
|
|
script_path = self.worktree_path / f".ap-script-{task_id}.py"
|
|
script_path.write_text(script_content)
|
|
|
|
try:
|
|
result = self.run_command(
|
|
task_id=task_id,
|
|
command=[interpreter, str(script_path)],
|
|
issue_number=issue_number,
|
|
timeout=timeout
|
|
)
|
|
finally:
|
|
# Cleanup script
|
|
script_path.unlink(missing_ok=True)
|
|
|
|
return result
|
|
|
|
def run_python_function(self,
|
|
task_id: str,
|
|
func: Callable,
|
|
issue_number: str,
|
|
args: tuple = (),
|
|
kwargs: Dict = None) -> TaskResult:
|
|
"""
|
|
Execute a Python function in the worktree context
|
|
|
|
Args:
|
|
task_id: Task identifier
|
|
func: Function to execute
|
|
issue_number: Associated issue
|
|
args: Positional arguments
|
|
kwargs: Keyword arguments
|
|
"""
|
|
start_time = time.time()
|
|
timestamp = datetime.now().isoformat()
|
|
kwargs = kwargs or {}
|
|
|
|
# Store original cwd
|
|
original_cwd = os.getcwd()
|
|
|
|
try:
|
|
# Change to worktree
|
|
os.chdir(self.worktree_path)
|
|
|
|
# Set environment markers
|
|
os.environ["AP_WORKTREE"] = str(self.worktree_path)
|
|
os.environ["AP_ISSUE"] = issue_number
|
|
os.environ["AP_TASK_ID"] = task_id
|
|
|
|
# Execute function
|
|
result = func(*args, **kwargs)
|
|
|
|
execution_time = time.time() - start_time
|
|
|
|
task_result = TaskResult(
|
|
task_id=task_id,
|
|
issue_number=issue_number,
|
|
success=True,
|
|
return_code=0,
|
|
stdout=str(result) if result else "",
|
|
stderr="",
|
|
execution_time=execution_time,
|
|
timestamp=timestamp,
|
|
artifacts=[]
|
|
)
|
|
|
|
except Exception as e:
|
|
execution_time = time.time() - start_time
|
|
task_result = TaskResult(
|
|
task_id=task_id,
|
|
issue_number=issue_number,
|
|
success=False,
|
|
return_code=-1,
|
|
stdout="",
|
|
stderr=str(e),
|
|
execution_time=execution_time,
|
|
timestamp=timestamp,
|
|
artifacts=[],
|
|
error=traceback.format_exc()
|
|
)
|
|
finally:
|
|
os.chdir(original_cwd)
|
|
|
|
self.task_history.append(task_result)
|
|
return task_result
|
|
|
|
def auto_commit(self,
|
|
message: str,
|
|
files: List[str] = None,
|
|
allow_empty: bool = False) -> Optional[str]:
|
|
"""
|
|
Automatically commit changes in the worktree
|
|
|
|
Args:
|
|
message: Commit message
|
|
files: Specific files to commit (None = all changes)
|
|
allow_empty: Allow empty commits
|
|
|
|
Returns:
|
|
Commit hash if successful
|
|
"""
|
|
# Add files
|
|
if files:
|
|
for f in files:
|
|
self._run_git(["add", f])
|
|
else:
|
|
self._run_git(["add", "-A"])
|
|
|
|
# Check if there are changes
|
|
status = self._run_git(["status", "--porcelain"])
|
|
if not status.stdout.strip() and not allow_empty:
|
|
return None
|
|
|
|
# Commit
|
|
cmd = ["commit", "-m", f"[AP] {message}"]
|
|
if allow_empty:
|
|
cmd.append("--allow-empty")
|
|
|
|
result = self._run_git(cmd)
|
|
|
|
if result.returncode == 0:
|
|
# Get commit hash
|
|
log_result = self._run_git(["log", "-1", "--format=%H"])
|
|
return log_result.stdout.strip()
|
|
|
|
return None
|
|
|
|
def get_logs(self, task_id: str = None) -> str:
|
|
"""Get logs for a task or all tasks"""
|
|
if task_id:
|
|
log_file = self.logs_dir / f"{task_id}.log"
|
|
if log_file.exists():
|
|
return log_file.read_text()
|
|
return ""
|
|
|
|
# Return all logs
|
|
all_logs = []
|
|
for log_file in sorted(self.logs_dir.glob("*.log")):
|
|
all_logs.append(f"=== {log_file.stem} ===")
|
|
all_logs.append(log_file.read_text())
|
|
all_logs.append("")
|
|
|
|
return "\n".join(all_logs)
|
|
|
|
def generate_report(self, output_path: str = None) -> str:
|
|
"""Generate a report of all task executions"""
|
|
report = {
|
|
"worktree": str(self.worktree_path),
|
|
"total_tasks": len(self.task_history),
|
|
"successful": sum(1 for t in self.task_history if t.success),
|
|
"failed": sum(1 for t in self.task_history if not t.success),
|
|
"tasks": [t.to_dict() for t in self.task_history]
|
|
}
|
|
|
|
report_json = json.dumps(report, indent=2)
|
|
|
|
if output_path:
|
|
Path(output_path).write_text(report_json)
|
|
|
|
return report_json
|
|
|
|
|
|
class TaskBatch:
|
|
"""Batch multiple tasks for sequential execution"""
|
|
|
|
def __init__(self, runner: TaskRunner, issue_number: str):
|
|
self.runner = runner
|
|
self.issue_number = issue_number
|
|
self.tasks: List[Dict] = []
|
|
self.results: List[TaskResult] = []
|
|
|
|
def add_command(self, task_id: str, command: List[str], **kwargs):
|
|
"""Add a command task to the batch"""
|
|
self.tasks.append({
|
|
"type": "command",
|
|
"task_id": task_id,
|
|
"command": command,
|
|
"kwargs": kwargs
|
|
})
|
|
|
|
def add_script(self, task_id: str, script: str, **kwargs):
|
|
"""Add a script task to the batch"""
|
|
self.tasks.append({
|
|
"type": "script",
|
|
"task_id": task_id,
|
|
"script": script,
|
|
"kwargs": kwargs
|
|
})
|
|
|
|
def execute(self, stop_on_error: bool = True) -> List[TaskResult]:
|
|
"""Execute all tasks in the batch"""
|
|
self.results = []
|
|
|
|
for task in self.tasks:
|
|
if task["type"] == "command":
|
|
result = self.runner.run_command(
|
|
task_id=task["task_id"],
|
|
command=task["command"],
|
|
issue_number=self.issue_number,
|
|
**task["kwargs"]
|
|
)
|
|
elif task["type"] == "script":
|
|
result = self.runner.run_script(
|
|
task_id=task["task_id"],
|
|
script_content=task["script"],
|
|
issue_number=self.issue_number,
|
|
**task["kwargs"]
|
|
)
|
|
|
|
self.results.append(result)
|
|
|
|
if stop_on_error and not result.success:
|
|
break
|
|
|
|
return self.results
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Test
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
runner = TaskRunner(tmpdir)
|
|
|
|
# Test command execution
|
|
result = runner.run_command(
|
|
task_id="test-1",
|
|
command=["echo", "Hello from AP"],
|
|
issue_number="123"
|
|
)
|
|
|
|
print(f"Success: {result.success}")
|
|
print(f"Output: {result.stdout}")
|