"""Durable task lifecycle for long-running tools and agent operations.""" from __future__ import annotations import fcntl import hashlib import json import os import re import secrets import signal import subprocess import sys import time from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Literal from .events import EventLog from .sandbox import control_root, ensure_control_root, sandbox_environment TaskStatus = Literal["working", "input_required", "completed", "failed", "cancelled"] @dataclass(frozen=True) class TaskRecord: task_id: str kind: str status: TaskStatus session_id_sha256: str request_sha256: str created_unix_ms: int updated_unix_ms: int status_message: str = "" progress: dict[str, Any] | None = None result: dict[str, Any] | None = None error: str = "" process_id: int = 0 def to_dict(self) -> dict[str, Any]: return asdict(self) class TaskStore: def __init__(self, workspace: str | Path) -> None: self.workspace = ensure_control_root(workspace) root = control_root(self.workspace) / "tasks" root.mkdir(parents=True, exist_ok=True, mode=0o700) os.chmod(root, 0o700) self.root = root def _path(self, task_id: str) -> Path: if not re.fullmatch(r"task_[a-f0-9]{32}", task_id): raise ValueError("task id is invalid") return self.root / f"{task_id}.json" @staticmethod def _write(path: Path, record: TaskRecord) -> None: temporary = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp") temporary.write_text( json.dumps(record.to_dict(), sort_keys=True, separators=(",", ":")), encoding="utf-8", ) temporary.chmod(0o600) with temporary.open("rb") as handle: os.fsync(handle.fileno()) os.replace(temporary, path) directory_fd = os.open(path.parent, os.O_RDONLY) try: os.fsync(directory_fd) finally: os.close(directory_fd) @staticmethod def _load(path: Path) -> TaskRecord: payload = json.loads(path.read_text(encoding="utf-8")) if not isinstance(payload, dict): raise RuntimeError("task record is invalid") return TaskRecord(**payload) def create( self, *, kind: str, session_id: str, request: dict[str, Any], process_id: int = 0, status_message: str = "", ) -> TaskRecord: canonical = json.dumps( request, sort_keys=True, default=str, separators=(",", ":") ).encode("utf-8") now = int(time.time() * 1000) record = TaskRecord( task_id="task_" + secrets.token_hex(16), kind=kind, status="working", session_id_sha256=hashlib.sha256(session_id.encode("utf-8")).hexdigest(), request_sha256=hashlib.sha256(canonical).hexdigest(), created_unix_ms=now, updated_unix_ms=now, process_id=process_id, status_message=status_message, ) self._write(self._path(record.task_id), record) return record def get(self, task_id: str) -> TaskRecord: return self._load(self._path(task_id)) @staticmethod def _session_sha256(session_id: str) -> str: return hashlib.sha256(session_id.encode("utf-8")).hexdigest() @classmethod def _require_session(cls, record: TaskRecord, session_id: str) -> TaskRecord: if session_id and not secrets.compare_digest( record.session_id_sha256, cls._session_sha256(session_id) ): raise PermissionError("task does not belong to this session") return record def status(self, task_id: str, *, session_id: str = "") -> TaskRecord: record = self.get(task_id) self._require_session(record, session_id) if record.status != "working" or record.process_id <= 0: return record try: os.kill(record.process_id, 0) return record except ProcessLookupError: failed = self.update( task_id, status="failed", status_message="execution worker exited without a final receipt", error="execution worker unavailable", process_id=0, ) EventLog(self.workspace).append( "task_failed", session_id_sha256=failed.session_id_sha256, task_id=failed.task_id, status="failed", detail={"reason": "worker_unavailable"}, ) return failed def update( self, task_id: str, *, status: TaskStatus | None = None, status_message: str | None = None, progress: dict[str, Any] | None = None, result: dict[str, Any] | None = None, error: str | None = None, process_id: int | None = None, ) -> TaskRecord: path = self._path(task_id) lock_path = path.with_suffix(".lock") with lock_path.open("a+b") as lock_handle: lock_path.chmod(0o600) fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX) try: current = self._load(path) if current.status in {"completed", "failed", "cancelled"}: if status is None or status != current.status: return current payload = current.to_dict() if status is not None: payload["status"] = status if status_message is not None: payload["status_message"] = status_message if progress is not None: payload["progress"] = progress if result is not None: payload["result"] = result if error is not None: payload["error"] = error if process_id is not None: payload["process_id"] = process_id payload["updated_unix_ms"] = int(time.time() * 1000) updated = TaskRecord(**payload) self._write(path, updated) return updated finally: fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN) def list(self, *, session_id: str = "") -> tuple[TaskRecord, ...]: session_hash = ( hashlib.sha256(session_id.encode("utf-8")).hexdigest() if session_id else "" ) records: list[TaskRecord] = [] for path in sorted(self.root.glob("task_*.json")): try: record = self._load(path) except (OSError, TypeError, ValueError, RuntimeError): continue if session_hash and record.session_id_sha256 != session_hash: continue record = self.status(record.task_id) records.append(record) return tuple(records) def start_terminal( self, *, session_id: str, command: str, workspace: str | Path, ) -> TaskRecord: """Start a durable terminal task without imposing a host time limit.""" workspace_root = ensure_control_root(workspace) if workspace_root != self.workspace: raise ValueError("task workspace does not match its state store") if not command.strip(): raise ValueError("terminal task command is required") record = self.create( kind="terminal", session_id=session_id, request={"command": command, "workspace": str(workspace_root)}, status_message="starting", ) source_root = str(Path(__file__).resolve().parents[2]) environment = sandbox_environment(pythonpath=source_root) environment["NEXUM_HOME"] = str(control_root(self.workspace)) process = subprocess.Popen( [ sys.executable, "-m", "nexum_runtime.tooling.task_worker", "--task-root", str(self.root), "--task-id", record.task_id, "--workspace", str(workspace_root), "--command", command, ], cwd=workspace_root, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True, close_fds=True, env=environment, ) updated = self.update( record.task_id, process_id=process.pid, status_message="running", ) EventLog(workspace_root).append( "task_started", session_id=session_id, task_id=updated.task_id, status=updated.status, detail={"kind": updated.kind, "process_id": updated.process_id}, ) return updated def cancel(self, task_id: str, *, session_id: str = "") -> TaskRecord: """Cancel a running task and retain its durable state for inspection.""" record = self.status(task_id, session_id=session_id) if record.status != "working": return record updated = self.update( task_id, status="cancelled", status_message="cancelled by caller", ) if record.process_id > 0: try: os.killpg(record.process_id, signal.SIGTERM) except ProcessLookupError: pass EventLog(self.workspace).append( "task_cancelled", session_id_sha256=record.session_id_sha256, task_id=record.task_id, status="cancelled", ) return updated __all__ = ["TaskRecord", "TaskStatus", "TaskStore"]