| """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"] |
|
|
