feat: runner task cancellation via asyncio.Task.cancel(), api_server cancel_current support
- TaskQueue.cancel(): mark any task (pending or running) as cancelled in SQLite - AutonomousRunner._execute_task(): self-registers via asyncio.current_task() so cancel_current() works whether called from start() or directly in tests - AutonomousRunner.cancel_current(task_id): cancels running asyncio.Task by ID, returns True/False; re-raises CancelledError after writing DB record - AgentAPIServer: add autonomous_runner param (self._autonomous_runner) to avoid naming conflict with self._runner (web.AppRunner); _handle_cancel_task() now handles both pending (queue.cancel) and running (cancel_current) states Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Claude Sonnet 4.6
@@ -0,0 +1,221 @@
|
||||
"""Lightweight HTTP API for the KVM Agent daemon.
|
||||
|
||||
Exposes the SQLite task queue and agent config over REST,
|
||||
so the web frontend can submit tasks and poll status.
|
||||
|
||||
Runs inside the same asyncio event loop as AutonomousRunner.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .config import AgentConfig
|
||||
from .runner import TaskQueue
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .runner import AutonomousRunner
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CORS_HEADERS = {
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "GET, POST, PATCH, DELETE, OPTIONS",
|
||||
"Access-Control-Allow-Headers": "Content-Type, Authorization",
|
||||
}
|
||||
|
||||
|
||||
def _cors(resp: web.Response) -> web.Response:
|
||||
resp.headers.update(CORS_HEADERS)
|
||||
return resp
|
||||
|
||||
|
||||
def _json(data, status: int = 200) -> web.Response:
|
||||
return _cors(web.json_response(data, status=status))
|
||||
|
||||
|
||||
def _error(msg: str, status: int = 400) -> web.Response:
|
||||
return _json({"error": msg}, status=status)
|
||||
|
||||
|
||||
def _task_dict(t) -> dict:
|
||||
return {
|
||||
"id": t.task_id,
|
||||
"task_type": t.task_type,
|
||||
"description": t.description,
|
||||
"variables": t.variables,
|
||||
"status": t.status,
|
||||
"created_at": t.created_at,
|
||||
"started_at": t.started_at,
|
||||
"finished_at": t.finished_at,
|
||||
"result": t.result,
|
||||
"success": t.success,
|
||||
}
|
||||
|
||||
|
||||
class AgentAPIServer:
|
||||
"""aiohttp-based API server for the KVM Agent daemon."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
queue: TaskQueue,
|
||||
config: AgentConfig,
|
||||
autonomous_runner: AutonomousRunner | None = None,
|
||||
host: str = "0.0.0.0",
|
||||
port: int = 8890,
|
||||
):
|
||||
self._queue = queue
|
||||
self._config = config
|
||||
self._autonomous_runner = autonomous_runner
|
||||
self._host = host
|
||||
self._port = port
|
||||
self._runner: Optional[web.AppRunner] = None
|
||||
|
||||
async def start(self):
|
||||
app = web.Application()
|
||||
app.router.add_route("OPTIONS", "/{path:.*}", self._handle_options)
|
||||
app.router.add_get("/api/v1/agent/status", self._handle_status)
|
||||
app.router.add_get("/api/v1/agent/tasks", self._handle_list_tasks)
|
||||
app.router.add_post("/api/v1/agent/tasks", self._handle_submit_task)
|
||||
app.router.add_get("/api/v1/agent/tasks/{id}", self._handle_get_task)
|
||||
app.router.add_delete("/api/v1/agent/tasks/{id}", self._handle_cancel_task)
|
||||
app.router.add_get("/api/v1/agent/config", self._handle_get_config)
|
||||
app.router.add_patch("/api/v1/agent/config", self._handle_patch_config)
|
||||
|
||||
self._runner = web.AppRunner(app)
|
||||
await self._runner.setup()
|
||||
site = web.TCPSite(self._runner, self._host, self._port)
|
||||
await site.start()
|
||||
logger.info("Agent API server listening on %s:%d", self._host, self._port)
|
||||
|
||||
async def stop(self):
|
||||
if self._runner:
|
||||
await self._runner.cleanup()
|
||||
logger.info("Agent API server stopped")
|
||||
|
||||
# ── Handlers ──────────────────────────────────────────────
|
||||
|
||||
async def _handle_options(self, request: web.Request) -> web.Response:
|
||||
return _cors(web.Response(status=204))
|
||||
|
||||
async def _handle_status(self, request: web.Request) -> web.Response:
|
||||
tasks = self._queue.list_tasks(status="running", limit=1)
|
||||
current = None
|
||||
if tasks:
|
||||
t = tasks[0]
|
||||
current = {
|
||||
"id": t.task_id,
|
||||
"description": t.description,
|
||||
"started_at": t.started_at,
|
||||
}
|
||||
|
||||
pending = self._queue.list_tasks(status="pending", limit=100)
|
||||
return _json({
|
||||
"running": True,
|
||||
"queue_depth": len(pending),
|
||||
"current_task": current,
|
||||
})
|
||||
|
||||
async def _handle_list_tasks(self, request: web.Request) -> web.Response:
|
||||
status_filter = request.query.get("status")
|
||||
limit = int(request.query.get("limit", "20"))
|
||||
tasks = self._queue.list_tasks(
|
||||
status=status_filter or None,
|
||||
limit=limit,
|
||||
)
|
||||
return _json({
|
||||
"tasks": [_task_dict(t) for t in tasks],
|
||||
"total": len(tasks),
|
||||
})
|
||||
|
||||
async def _handle_submit_task(self, request: web.Request) -> web.Response:
|
||||
try:
|
||||
body = await request.json()
|
||||
except Exception:
|
||||
return _error("Invalid JSON body")
|
||||
|
||||
description = body.get("description", "").strip()
|
||||
if not description:
|
||||
return _error("description is required")
|
||||
|
||||
task_type = body.get("task_type", "task")
|
||||
variables = body.get("variables")
|
||||
|
||||
task_id = self._queue.enqueue(
|
||||
task_type=task_type,
|
||||
description=description,
|
||||
variables=variables,
|
||||
)
|
||||
return _json({"task_id": task_id}, status=201)
|
||||
|
||||
async def _handle_get_task(self, request: web.Request) -> web.Response:
|
||||
try:
|
||||
task_id = int(request.match_info["id"])
|
||||
except ValueError:
|
||||
return _error("Invalid task ID")
|
||||
|
||||
task = self._queue.get_task(task_id)
|
||||
if not task:
|
||||
return _error("Task not found", status=404)
|
||||
|
||||
return _json(_task_dict(task))
|
||||
|
||||
async def _handle_cancel_task(self, request: web.Request) -> web.Response:
|
||||
try:
|
||||
task_id = int(request.match_info["id"])
|
||||
except ValueError:
|
||||
return _error("Invalid task ID")
|
||||
|
||||
task = self._queue.get_task(task_id)
|
||||
if not task:
|
||||
return _error("Task not found", status=404)
|
||||
|
||||
if task.status == "pending":
|
||||
self._queue.cancel(task_id, result="Cancelled by user")
|
||||
return _json({"cancelled": True})
|
||||
|
||||
if task.status == "running":
|
||||
if self._autonomous_runner and self._autonomous_runner.cancel_current(task_id):
|
||||
return _json({"cancelled": True})
|
||||
return _error("Task is running but runner not available", status=503)
|
||||
|
||||
return _error(f"Cannot cancel task in '{task.status}' status")
|
||||
|
||||
async def _handle_get_config(self, request: web.Request) -> web.Response:
|
||||
cfg = self._config
|
||||
return _json({
|
||||
"kvm_url": cfg.kvm_url,
|
||||
"max_steps": cfg.max_steps,
|
||||
"step_delay": cfg.step_delay,
|
||||
"click_delay": cfg.click_delay,
|
||||
"llm_model": cfg.llm_model,
|
||||
"llm_temperature": cfg.llm_temperature,
|
||||
"llm_max_tokens": cfg.llm_max_tokens,
|
||||
"memory_enabled": cfg.memory_enabled,
|
||||
"template_enabled": cfg.template_enabled,
|
||||
"screen_width": cfg.screen_width,
|
||||
"screen_height": cfg.screen_height,
|
||||
"daemon_poll_interval": cfg.daemon_poll_interval,
|
||||
})
|
||||
|
||||
async def _handle_patch_config(self, request: web.Request) -> web.Response:
|
||||
try:
|
||||
body = await request.json()
|
||||
except Exception:
|
||||
return _error("Invalid JSON body")
|
||||
|
||||
allowed = {"llm_model", "max_steps", "llm_temperature", "step_delay"}
|
||||
updated = {}
|
||||
for key, value in body.items():
|
||||
if key in allowed:
|
||||
setattr(self._config, key, value)
|
||||
updated[key] = value
|
||||
|
||||
if not updated:
|
||||
return _error("No valid fields to update")
|
||||
|
||||
return _json({"updated": updated})
|
||||
@@ -0,0 +1,276 @@
|
||||
"""Autonomous runner — daemon process with SQLite task queue.
|
||||
|
||||
Runs as a systemd service. External tools (CLI, API, cron) enqueue tasks;
|
||||
the runner dequeues and executes them one at a time via the KVM Agent.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import signal
|
||||
import sqlite3
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TaskStatus(Enum):
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueuedTask:
|
||||
"""A task in the queue."""
|
||||
|
||||
task_id: int
|
||||
task_type: str # "task" or "workflow"
|
||||
description: str
|
||||
variables: dict
|
||||
status: str
|
||||
created_at: float
|
||||
started_at: Optional[float] = None
|
||||
finished_at: Optional[float] = None
|
||||
result: str = ""
|
||||
success: bool = False
|
||||
|
||||
|
||||
class TaskQueue:
|
||||
"""SQLite-backed persistent task queue."""
|
||||
|
||||
def __init__(self, db_path: str = "task_queue.db"):
|
||||
self._db_path = db_path
|
||||
self._conn: Optional[sqlite3.Connection] = None
|
||||
|
||||
def open(self):
|
||||
self._conn = sqlite3.connect(self._db_path)
|
||||
self._conn.execute("PRAGMA journal_mode=WAL")
|
||||
self._conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS tasks (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
task_type TEXT NOT NULL DEFAULT 'task',
|
||||
description TEXT NOT NULL,
|
||||
variables TEXT NOT NULL DEFAULT '{}',
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
created_at REAL NOT NULL,
|
||||
started_at REAL,
|
||||
finished_at REAL,
|
||||
result TEXT DEFAULT '',
|
||||
success INTEGER DEFAULT 0
|
||||
)
|
||||
""")
|
||||
self._conn.commit()
|
||||
|
||||
def close(self):
|
||||
if self._conn:
|
||||
self._conn.close()
|
||||
self._conn = None
|
||||
|
||||
def enqueue(
|
||||
self,
|
||||
task_type: str,
|
||||
description: str,
|
||||
variables: Optional[dict] = None,
|
||||
) -> int:
|
||||
"""Add a task to the queue. Returns task ID."""
|
||||
assert self._conn is not None
|
||||
cur = self._conn.execute(
|
||||
"INSERT INTO tasks (task_type, description, variables, status, created_at) "
|
||||
"VALUES (?, ?, ?, 'pending', ?)",
|
||||
(task_type, description, json.dumps(variables or {}), time.time()),
|
||||
)
|
||||
self._conn.commit()
|
||||
return cur.lastrowid or 0
|
||||
|
||||
def dequeue(self) -> Optional[QueuedTask]:
|
||||
"""Get the next pending task and mark it as running."""
|
||||
assert self._conn is not None
|
||||
row = self._conn.execute(
|
||||
"SELECT id, task_type, description, variables, status, created_at "
|
||||
"FROM tasks WHERE status = 'pending' ORDER BY id LIMIT 1"
|
||||
).fetchone()
|
||||
if not row:
|
||||
return None
|
||||
|
||||
now = time.time()
|
||||
self._conn.execute(
|
||||
"UPDATE tasks SET status = 'running', started_at = ? WHERE id = ?",
|
||||
(now, row[0]),
|
||||
)
|
||||
self._conn.commit()
|
||||
return QueuedTask(
|
||||
task_id=row[0],
|
||||
task_type=row[1],
|
||||
description=row[2],
|
||||
variables=json.loads(row[3]),
|
||||
status="running",
|
||||
created_at=row[4],
|
||||
started_at=now,
|
||||
)
|
||||
|
||||
def complete(self, task_id: int, result: str, success: bool):
|
||||
"""Mark a task as completed or failed."""
|
||||
assert self._conn is not None
|
||||
status = "completed" if success else "failed"
|
||||
self._conn.execute(
|
||||
"UPDATE tasks SET status = ?, finished_at = ?, result = ?, success = ? "
|
||||
"WHERE id = ?",
|
||||
(status, time.time(), result, int(success), task_id),
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def cancel(self, task_id: int, result: str = "Cancelled by user") -> None:
|
||||
"""Mark a task as cancelled regardless of its current status."""
|
||||
assert self._conn is not None
|
||||
self._conn.execute(
|
||||
"UPDATE tasks SET status = 'cancelled', finished_at = ?, result = ? WHERE id = ?",
|
||||
(time.time(), result, task_id),
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def get_task(self, task_id: int) -> Optional[QueuedTask]:
|
||||
"""Get a task by ID."""
|
||||
assert self._conn is not None
|
||||
row = self._conn.execute(
|
||||
"SELECT id, task_type, description, variables, status, created_at, "
|
||||
"started_at, finished_at, result, success FROM tasks WHERE id = ?",
|
||||
(task_id,),
|
||||
).fetchone()
|
||||
if not row:
|
||||
return None
|
||||
return QueuedTask(
|
||||
task_id=row[0], task_type=row[1], description=row[2],
|
||||
variables=json.loads(row[3]), status=row[4], created_at=row[5],
|
||||
started_at=row[6], finished_at=row[7], result=row[8],
|
||||
success=bool(row[9]),
|
||||
)
|
||||
|
||||
def list_tasks(self, status: Optional[str] = None, limit: int = 20) -> list[QueuedTask]:
|
||||
"""List tasks, optionally filtered by status."""
|
||||
assert self._conn is not None
|
||||
if status:
|
||||
rows = self._conn.execute(
|
||||
"SELECT id, task_type, description, variables, status, created_at, "
|
||||
"started_at, finished_at, result, success "
|
||||
"FROM tasks WHERE status = ? ORDER BY id DESC LIMIT ?",
|
||||
(status, limit),
|
||||
).fetchall()
|
||||
else:
|
||||
rows = self._conn.execute(
|
||||
"SELECT id, task_type, description, variables, status, created_at, "
|
||||
"started_at, finished_at, result, success "
|
||||
"FROM tasks ORDER BY id DESC LIMIT ?",
|
||||
(limit,),
|
||||
).fetchall()
|
||||
return [
|
||||
QueuedTask(
|
||||
task_id=r[0], task_type=r[1], description=r[2],
|
||||
variables=json.loads(r[3]), status=r[4], created_at=r[5],
|
||||
started_at=r[6], finished_at=r[7], result=r[8],
|
||||
success=bool(r[9]),
|
||||
)
|
||||
for r in rows
|
||||
]
|
||||
|
||||
|
||||
class AutonomousRunner:
|
||||
"""Daemon that dequeues and executes tasks from the TaskQueue.
|
||||
|
||||
Args:
|
||||
agent: KVMAgent instance for executing tasks.
|
||||
queue: TaskQueue instance.
|
||||
poll_interval: Seconds between queue polls when idle.
|
||||
"""
|
||||
|
||||
def __init__(self, agent, queue: TaskQueue, poll_interval: float = 5.0):
|
||||
self._agent = agent
|
||||
self._queue = queue
|
||||
self._poll_interval = poll_interval
|
||||
self._shutdown = False
|
||||
self._current_task_id: Optional[int] = None
|
||||
self._running_coro: Optional[asyncio.Task] = None
|
||||
|
||||
async def start(self):
|
||||
"""Main loop — poll queue, execute tasks, repeat."""
|
||||
logger.info("Autonomous runner started (poll=%.1fs)", self._poll_interval)
|
||||
|
||||
# Handle graceful shutdown
|
||||
loop = asyncio.get_event_loop()
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
loop.add_signal_handler(sig, self._handle_signal)
|
||||
|
||||
while not self._shutdown:
|
||||
task = self._queue.dequeue()
|
||||
if task:
|
||||
logger.info(
|
||||
"Dequeued task #%d: [%s] %s",
|
||||
task.task_id, task.task_type, task.description,
|
||||
)
|
||||
self._running_coro = asyncio.create_task(self._execute_task(task))
|
||||
try:
|
||||
await self._running_coro
|
||||
except asyncio.CancelledError:
|
||||
pass # handled inside _execute_task
|
||||
finally:
|
||||
self._running_coro = None
|
||||
self._current_task_id = None
|
||||
else:
|
||||
await asyncio.sleep(self._poll_interval)
|
||||
|
||||
logger.info("Autonomous runner stopped")
|
||||
|
||||
async def _execute_task(self, task: QueuedTask):
|
||||
"""Execute a single queued task."""
|
||||
self._current_task_id = task.task_id
|
||||
# Track the asyncio.Task wrapping this coroutine so cancel_current()
|
||||
# works even when _execute_task is called directly (e.g. in tests).
|
||||
current = asyncio.current_task()
|
||||
if current is not None and self._running_coro is None:
|
||||
self._running_coro = current
|
||||
try:
|
||||
result = await self._agent.run_task(task.description)
|
||||
self._queue.complete(
|
||||
task.task_id,
|
||||
result=result.final_reason,
|
||||
success=result.success,
|
||||
)
|
||||
logger.info(
|
||||
"Task #%d %s: %s",
|
||||
task.task_id,
|
||||
"succeeded" if result.success else "failed",
|
||||
result.final_reason,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
self._queue.cancel(task.task_id, result="Cancelled by user")
|
||||
logger.info("Task #%d cancelled", task.task_id)
|
||||
raise # re-raise so asyncio knows task was cancelled
|
||||
except Exception as e:
|
||||
logger.error("Task #%d crashed: %s", task.task_id, e)
|
||||
self._queue.complete(task.task_id, result=str(e), success=False)
|
||||
finally:
|
||||
# Only clear if we set it ourselves (direct-call path)
|
||||
if self._running_coro is current:
|
||||
self._running_coro = None
|
||||
self._current_task_id = None
|
||||
|
||||
def cancel_current(self, task_id: int) -> bool:
|
||||
"""Cancel the currently running task by asyncio cancellation.
|
||||
|
||||
Returns True if the task was found and cancellation was requested,
|
||||
False if no matching task is currently running.
|
||||
"""
|
||||
if self._current_task_id == task_id and self._running_coro is not None:
|
||||
self._running_coro.cancel()
|
||||
return True
|
||||
return False
|
||||
|
||||
def _handle_signal(self):
|
||||
logger.info("Shutdown signal received")
|
||||
self._shutdown = True
|
||||
self._agent.stop()
|
||||
@@ -0,0 +1,84 @@
|
||||
"""Tests for runner task cancellation via asyncio.Task.cancel()."""
|
||||
import asyncio
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from services.kvm_agent.runner import TaskQueue, AutonomousRunner
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def queue(tmp_path):
|
||||
q = TaskQueue(str(tmp_path / "test.db"))
|
||||
q.open()
|
||||
yield q
|
||||
q.close()
|
||||
|
||||
|
||||
def test_queue_cancel_pending(queue):
|
||||
"""TaskQueue.cancel() marks pending task as cancelled."""
|
||||
task_id = queue.enqueue("task", "test task")
|
||||
queue.cancel(task_id, "User cancelled")
|
||||
task = queue.get_task(task_id)
|
||||
assert task.status == "cancelled"
|
||||
assert task.result == "User cancelled"
|
||||
|
||||
|
||||
def test_queue_cancel_running(queue):
|
||||
"""TaskQueue.cancel() marks running task as cancelled."""
|
||||
task_id = queue.enqueue("task", "test task")
|
||||
queue.dequeue() # marks as running
|
||||
queue.cancel(task_id, "User cancelled")
|
||||
task = queue.get_task(task_id)
|
||||
assert task.status == "cancelled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_cancel_current_task(queue):
|
||||
"""cancel_current() cancels running task and marks it cancelled in DB."""
|
||||
agent = MagicMock()
|
||||
|
||||
async def slow_run(description):
|
||||
await asyncio.sleep(10) # will be cancelled
|
||||
return MagicMock(success=True, final_reason="done")
|
||||
|
||||
agent.run_task = slow_run
|
||||
|
||||
runner = AutonomousRunner(agent, queue, poll_interval=0.1)
|
||||
task_id = queue.enqueue("task", "slow task")
|
||||
queued = queue.dequeue() # mark as running
|
||||
|
||||
# Start _execute_task in background
|
||||
exec_task = asyncio.create_task(runner._execute_task(queued))
|
||||
await asyncio.sleep(0.05) # let it start sleeping
|
||||
|
||||
# Cancel it
|
||||
success = runner.cancel_current(task_id)
|
||||
assert success is True
|
||||
|
||||
# Wait for it to finish handling cancellation
|
||||
await asyncio.gather(exec_task, return_exceptions=True)
|
||||
|
||||
task = queue.get_task(task_id)
|
||||
assert task.status == "cancelled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_cancel_wrong_id_returns_false(queue):
|
||||
"""cancel_current() with wrong task_id returns False."""
|
||||
agent = MagicMock()
|
||||
runner = AutonomousRunner(agent, queue, poll_interval=0.1)
|
||||
|
||||
result = runner.cancel_current(999)
|
||||
assert result is False
|
||||
|
||||
|
||||
def test_api_cancel_pending_task(queue):
|
||||
"""api_server cancel handler can cancel pending tasks via queue.cancel()."""
|
||||
task_id = queue.enqueue("task", "pending task")
|
||||
task = queue.get_task(task_id)
|
||||
assert task.status == "pending"
|
||||
|
||||
# Simulate what _handle_cancel_task does for pending tasks
|
||||
queue.cancel(task_id, "Cancelled by user")
|
||||
task = queue.get_task(task_id)
|
||||
assert task.status == "cancelled"
|
||||
Reference in New Issue
Block a user