298 lines
10 KiB
Python
298 lines
10 KiB
Python
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from pydantic import ValidationError
|
|
|
|
from ..errors import ToolInputError
|
|
from .guards import AutomationDenied, max_active_per_owner, task_columns
|
|
from .schedule import Schedule, next_run, now_utc, to_iso
|
|
from .store import TaskStore
|
|
|
|
logger = logging.getLogger("devii.tasks.controller")
|
|
|
|
SCHEDULE_KEYS = (
|
|
"kind",
|
|
"run_at",
|
|
"delay_seconds",
|
|
"every_seconds",
|
|
"start_at",
|
|
"cron",
|
|
"max_runs",
|
|
"expires_at",
|
|
)
|
|
RESULT_PREVIEW_CHARS = 500
|
|
TRUTHY = {"1", "true", "yes", "on"}
|
|
|
|
|
|
def _as_bool(value: Any, default: bool = True) -> bool:
|
|
if isinstance(value, bool):
|
|
return value
|
|
if value is None:
|
|
return default
|
|
return str(value).strip().lower() in TRUTHY
|
|
|
|
|
|
def _serialize(row: dict[str, Any], preview: bool) -> dict[str, Any]:
|
|
view = {
|
|
"uid": row.get("uid"),
|
|
"label": row.get("label"),
|
|
"prompt": row.get("prompt"),
|
|
"enabled": bool(row.get("enabled")),
|
|
"status": row.get("status"),
|
|
"kind": row.get("kind"),
|
|
"next_run_at": row.get("next_run_at"),
|
|
"last_run_at": row.get("last_run_at"),
|
|
"run_count": row.get("run_count"),
|
|
"max_runs": row.get("max_runs"),
|
|
"expires_at": row.get("expires_at"),
|
|
"failure_count": int(row.get("failure_count") or 0),
|
|
"every_seconds": row.get("every_seconds"),
|
|
"cron": row.get("cron"),
|
|
"run_at": row.get("run_at"),
|
|
"notify": bool(row.get("notify")),
|
|
"tz": row.get("tz") or None,
|
|
}
|
|
last_error = row.get("last_error")
|
|
if last_error:
|
|
view["last_error"] = last_error
|
|
result = row.get("last_result")
|
|
if result:
|
|
view["last_result"] = result[:RESULT_PREVIEW_CHARS] if preview else result
|
|
return view
|
|
|
|
|
|
class TaskController:
|
|
def __init__(self, store: TaskStore) -> None:
|
|
self._store = store
|
|
|
|
async def dispatch(self, name: str, arguments: dict[str, Any]) -> str:
|
|
handlers = {
|
|
"current_time": self.current_time,
|
|
"create_task": self.create_task,
|
|
"list_tasks": self.list_tasks,
|
|
"get_task": self.get_task,
|
|
"update_task": self.update_task,
|
|
"delete_task": self.delete_task,
|
|
"run_task_now": self.run_task_now,
|
|
}
|
|
handler = handlers.get(name)
|
|
if handler is None:
|
|
raise ToolInputError(f"Unknown task tool: {name}")
|
|
return handler(arguments)
|
|
|
|
def current_time(self, arguments: dict[str, Any]) -> str:
|
|
return json.dumps({"utc": to_iso(now_utc())}, ensure_ascii=False)
|
|
|
|
def create_task(self, arguments: dict[str, Any]) -> str:
|
|
prompt = str(arguments.get("prompt", "")).strip()
|
|
if not prompt:
|
|
raise ToolInputError("create_task requires a non-empty prompt.")
|
|
self._require_creation()
|
|
self._require_capacity()
|
|
schedule = self._build_schedule(arguments)
|
|
reference = now_utc()
|
|
first = schedule.first_run(reference)
|
|
record: dict[str, Any] = {
|
|
"uid": uuid.uuid4().hex,
|
|
"label": (arguments.get("label") or "").strip() or None,
|
|
"prompt": prompt,
|
|
"enabled": True,
|
|
"status": "pending",
|
|
"created_at": to_iso(reference),
|
|
"next_run_at": to_iso(first),
|
|
"last_run_at": None,
|
|
"run_count": 0,
|
|
"last_result": None,
|
|
"last_error": None,
|
|
"failure_count": 0,
|
|
"notify": 1 if _as_bool(arguments.get("notify"), default=False) else 0,
|
|
"tz": (arguments.get("tz") or "").strip() or None,
|
|
**task_columns(schedule, reference),
|
|
}
|
|
try:
|
|
self._store.create(record)
|
|
except AutomationDenied as exc:
|
|
raise self._denied(exc) from exc
|
|
return json.dumps(
|
|
{"status": "created", "task": _serialize(record, preview=True)},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
def list_tasks(self, arguments: dict[str, Any]) -> str:
|
|
rows = self._store.list(
|
|
enabled_only=_as_bool(arguments.get("enabled_only"), default=False),
|
|
status=(arguments.get("status") or None),
|
|
)
|
|
rows.sort(key=lambda row: row.get("next_run_at") or "")
|
|
return json.dumps(
|
|
{
|
|
"count": len(rows),
|
|
"tasks": [_serialize(row, preview=True) for row in rows],
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
def get_task(self, arguments: dict[str, Any]) -> str:
|
|
row = self._require_task(arguments)
|
|
return json.dumps({"task": _serialize(row, preview=False)}, ensure_ascii=False)
|
|
|
|
def update_task(self, arguments: dict[str, Any]) -> str:
|
|
row = self._require_task(arguments)
|
|
self._require_scheduling()
|
|
reference = now_utc()
|
|
changes: dict[str, Any] = {}
|
|
|
|
if "prompt" in arguments and arguments["prompt"] is not None:
|
|
new_prompt = str(arguments["prompt"]).strip()
|
|
if not new_prompt:
|
|
raise ToolInputError("prompt cannot be empty.")
|
|
changes["prompt"] = new_prompt
|
|
if "label" in arguments:
|
|
changes["label"] = (arguments.get("label") or "").strip() or None
|
|
if "enabled" in arguments:
|
|
changes["enabled"] = _as_bool(arguments.get("enabled"))
|
|
if "notify" in arguments and arguments["notify"] is not None:
|
|
changes["notify"] = 1 if _as_bool(arguments.get("notify")) else 0
|
|
if "tz" in arguments and arguments["tz"] is not None:
|
|
changes["tz"] = (str(arguments.get("tz")) or "").strip() or None
|
|
|
|
if any(
|
|
key in arguments and arguments[key] is not None for key in SCHEDULE_KEYS
|
|
):
|
|
merged = {key: row.get(key) for key in SCHEDULE_KEYS}
|
|
for key in SCHEDULE_KEYS:
|
|
if key in arguments and arguments[key] is not None:
|
|
merged[key] = arguments[key]
|
|
schedule = self._build_schedule(merged)
|
|
changes.update(task_columns(schedule, reference))
|
|
changes["next_run_at"] = to_iso(schedule.first_run(reference))
|
|
changes["status"] = "pending"
|
|
changes["failure_count"] = 0
|
|
|
|
if changes.get("enabled") and row.get("status") in (
|
|
"done",
|
|
"disabled",
|
|
"error",
|
|
):
|
|
changes.setdefault("status", "pending")
|
|
if not changes.get("next_run_at") and not row.get("next_run_at"):
|
|
schedule = self._build_schedule(
|
|
{key: row.get(key) for key in SCHEDULE_KEYS}
|
|
)
|
|
changes["next_run_at"] = to_iso(schedule.first_run(reference))
|
|
if changes.get("enabled") is False:
|
|
changes["status"] = "disabled"
|
|
|
|
if not changes:
|
|
raise ToolInputError("No updatable fields supplied.")
|
|
|
|
self._store.update(row["uid"], changes)
|
|
return json.dumps(
|
|
{
|
|
"status": "updated",
|
|
"task": _serialize(self._store.get(row["uid"]), preview=True),
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
def delete_task(self, arguments: dict[str, Any]) -> str:
|
|
uid = self._require_uid(arguments)
|
|
deleted = self._store.delete(uid)
|
|
if not deleted:
|
|
raise ToolInputError(f"No task found with uid {uid}.")
|
|
return json.dumps({"status": "deleted", "uid": uid}, ensure_ascii=False)
|
|
|
|
def run_task_now(self, arguments: dict[str, Any]) -> str:
|
|
row = self._require_task(arguments)
|
|
self._require_scheduling()
|
|
self._store.update(
|
|
row["uid"],
|
|
{"enabled": True, "status": "pending", "next_run_at": to_iso(now_utc())},
|
|
)
|
|
return json.dumps(
|
|
{
|
|
"status": "queued",
|
|
"uid": row["uid"],
|
|
"note": "Will execute on the next scheduler tick.",
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
@staticmethod
|
|
def _denied(exc: AutomationDenied) -> ToolInputError:
|
|
if exc.retry_at is None:
|
|
return ToolInputError(f"Task refused: {exc.reason}.")
|
|
return ToolInputError(
|
|
f"Task refused: {exc.reason}. The next slot frees up at "
|
|
f"{to_iso(exc.retry_at)} UTC."
|
|
)
|
|
|
|
def _require_creation(self) -> None:
|
|
try:
|
|
self._store.require_creation_allowed()
|
|
except AutomationDenied as exc:
|
|
raise self._denied(exc) from exc
|
|
|
|
def _require_scheduling(self) -> None:
|
|
try:
|
|
self._store.require_scheduling_allowed()
|
|
except AutomationDenied as exc:
|
|
raise self._denied(exc) from exc
|
|
|
|
def _require_capacity(self) -> None:
|
|
limit = max_active_per_owner()
|
|
if limit > 0 and self._store.count_active() >= limit:
|
|
raise ToolInputError(
|
|
f"You already have {limit} active tasks. Delete or disable one first."
|
|
)
|
|
|
|
def _build_schedule(self, source: dict[str, Any]) -> Schedule:
|
|
payload = {
|
|
key: source.get(key) for key in SCHEDULE_KEYS if source.get(key) is not None
|
|
}
|
|
try:
|
|
return Schedule(**payload)
|
|
except ValidationError as exc:
|
|
raise ToolInputError(f"Invalid schedule: {exc.errors()[0]['msg']}") from exc
|
|
except ValueError as exc:
|
|
raise ToolInputError(f"Invalid schedule: {exc}") from exc
|
|
|
|
def _require_uid(self, arguments: dict[str, Any]) -> str:
|
|
uid = str(arguments.get("uid", "")).strip()
|
|
if not uid:
|
|
raise ToolInputError("This task tool requires a uid.")
|
|
return uid
|
|
|
|
def _require_task(self, arguments: dict[str, Any]) -> dict[str, Any]:
|
|
uid = self._require_uid(arguments)
|
|
row = self._store.get(uid)
|
|
if row is None:
|
|
raise ToolInputError(f"No task found with uid {uid}.")
|
|
return row
|
|
|
|
|
|
def compute_followup(
|
|
row: dict[str, Any], reference: Any
|
|
) -> tuple[dict[str, Any], bool]:
|
|
run_count = int(row.get("run_count") or 0) + 1
|
|
max_runs = row.get("max_runs")
|
|
upcoming = next_run(
|
|
row.get("kind"), row.get("every_seconds"), row.get("cron"), reference
|
|
)
|
|
|
|
changes: dict[str, Any] = {"run_count": run_count, "last_run_at": to_iso(reference)}
|
|
if upcoming is None or (max_runs is not None and run_count >= int(max_runs)):
|
|
changes["status"] = "done"
|
|
changes["enabled"] = False
|
|
changes["next_run_at"] = None
|
|
else:
|
|
changes["status"] = "pending"
|
|
changes["next_run_at"] = to_iso(upcoming)
|
|
return changes, changes["status"] == "done"
|