202 lines
8.7 KiB
Python
202 lines
8.7 KiB
Python
"""Fail-closed Agent Pull Runner primitives for AIOS autonomy v1.
|
|
|
|
The runner owns task discovery and reporting only. It never accepts shell
|
|
strings: executors receive argv, an allowlisted cwd and an allowlisted env.
|
|
Durable state is intentionally small so restart recovery and duplicate
|
|
suppression can be verified without coupling the runner to one database.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import random
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Mapping, Protocol, Sequence
|
|
|
|
|
|
class RunnerError(RuntimeError):
|
|
"""Raised when a runner contract would be violated."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RunnerConfig:
|
|
worker_id: str
|
|
capabilities: tuple[str, ...] = ()
|
|
long_poll_seconds: int = 30
|
|
fallback_seconds: int = 60
|
|
jitter_seconds: int = 5
|
|
max_backoff_seconds: int = 300
|
|
executable_allowlist: tuple[str, ...] = ()
|
|
working_directory_allowlist: tuple[str, ...] = ()
|
|
environment_allowlist: tuple[str, ...] = ()
|
|
|
|
|
|
@dataclass
|
|
class RunnerState:
|
|
inflight: dict[str, dict[str, Any]] = field(default_factory=dict)
|
|
terminal: dict[str, dict[str, Any]] = field(default_factory=dict)
|
|
|
|
|
|
class RunnerTransport(Protocol):
|
|
def register_worker(self, worker_id: str, capabilities: Sequence[str]) -> Mapping[str, Any]: ...
|
|
def heartbeat(self, worker_id: str) -> Mapping[str, Any]: ...
|
|
def long_poll(self, worker_id: str, timeout_seconds: int) -> Mapping[str, Any] | None: ...
|
|
def poll(self, worker_id: str) -> Mapping[str, Any] | None: ...
|
|
def claim(self, worker_id: str, task: Mapping[str, Any]) -> Mapping[str, Any]: ...
|
|
def renew(self, worker_id: str, task_id: str) -> Mapping[str, Any]: ...
|
|
def report_progress(self, worker_id: str, task_id: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: ...
|
|
def report_result(self, worker_id: str, task_id: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: ...
|
|
def report_evidence(self, worker_id: str, task_id: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: ...
|
|
|
|
|
|
class JsonStateStore:
|
|
"""Atomic JSON state store for restart recovery."""
|
|
|
|
def __init__(self, path: str | os.PathLike[str]) -> None:
|
|
self.path = Path(path)
|
|
|
|
def load(self) -> RunnerState:
|
|
if not self.path.exists():
|
|
return RunnerState()
|
|
value = json.loads(self.path.read_text(encoding="utf-8"))
|
|
if not isinstance(value, dict):
|
|
raise RunnerError("runner state must be an object")
|
|
return RunnerState(
|
|
inflight=dict(value.get("inflight", {})),
|
|
terminal=dict(value.get("terminal", {})),
|
|
)
|
|
|
|
def save(self, state: RunnerState) -> None:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
fd, name = tempfile.mkstemp(prefix=self.path.name + ".", dir=self.path.parent)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
|
json.dump({"inflight": state.inflight, "terminal": state.terminal}, handle, sort_keys=True)
|
|
handle.write("\n")
|
|
os.chmod(name, 0o600)
|
|
os.replace(name, self.path)
|
|
finally:
|
|
if os.path.exists(name):
|
|
os.unlink(name)
|
|
|
|
|
|
class AgentPullRunner:
|
|
"""Pull, claim, execute and report one task at a time."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: RunnerConfig,
|
|
transport: RunnerTransport,
|
|
state_store: JsonStateStore | None = None,
|
|
*,
|
|
random_fn: Callable[[], float] = random.random,
|
|
) -> None:
|
|
if not config.worker_id.strip():
|
|
raise RunnerError("worker_id is required")
|
|
self.config = config
|
|
self.transport = transport
|
|
self.state_store = state_store
|
|
self.random_fn = random_fn
|
|
self.state = state_store.load() if state_store else RunnerState()
|
|
|
|
def _save(self) -> None:
|
|
if self.state_store:
|
|
self.state_store.save(self.state)
|
|
|
|
def register(self) -> Mapping[str, Any]:
|
|
return self.transport.register_worker(self.config.worker_id, self.config.capabilities)
|
|
|
|
def heartbeat(self) -> Mapping[str, Any]:
|
|
return self.transport.heartbeat(self.config.worker_id)
|
|
|
|
def discovery_delay(self, failures: int) -> float:
|
|
failures = max(0, int(failures))
|
|
base = min(self.config.max_backoff_seconds, self.config.fallback_seconds * (2**failures))
|
|
jitter = self.random_fn() * max(0, self.config.jitter_seconds)
|
|
return min(self.config.max_backoff_seconds, base + jitter)
|
|
|
|
def discover_once(self) -> Mapping[str, Any] | None:
|
|
try:
|
|
task = self.transport.long_poll(self.config.worker_id, self.config.long_poll_seconds)
|
|
except Exception:
|
|
task = None
|
|
if task is None:
|
|
task = self.transport.poll(self.config.worker_id)
|
|
if not task:
|
|
return None
|
|
task_id = str(task.get("task_id", ""))
|
|
if not task_id:
|
|
raise RunnerError("task_id is required")
|
|
if task_id in self.state.terminal:
|
|
return {"status": "duplicate_suppressed", "task_id": task_id}
|
|
claimed = dict(self.transport.claim(self.config.worker_id, task))
|
|
self.state.inflight[task_id] = claimed
|
|
self._save()
|
|
return claimed
|
|
|
|
def renew(self, task_id: str) -> Mapping[str, Any]:
|
|
if task_id not in self.state.inflight:
|
|
raise RunnerError("task is not inflight")
|
|
return self.transport.renew(self.config.worker_id, task_id)
|
|
|
|
def report_progress(self, task_id: str, payload: Mapping[str, Any]) -> Mapping[str, Any]:
|
|
if task_id not in self.state.inflight:
|
|
raise RunnerError("task is not inflight")
|
|
return self.transport.report_progress(self.config.worker_id, task_id, payload)
|
|
|
|
def report_evidence(self, task_id: str, payload: Mapping[str, Any]) -> Mapping[str, Any]:
|
|
if task_id not in self.state.inflight:
|
|
raise RunnerError("task is not inflight")
|
|
return self.transport.report_evidence(self.config.worker_id, task_id, payload)
|
|
|
|
def report_result(self, task_id: str, payload: Mapping[str, Any]) -> Mapping[str, Any]:
|
|
if task_id in self.state.terminal:
|
|
return {"status": "duplicate_terminal_suppressed", "task_id": task_id}
|
|
if task_id not in self.state.inflight:
|
|
raise RunnerError("task is not inflight")
|
|
result = dict(self.transport.report_result(self.config.worker_id, task_id, payload))
|
|
self.state.terminal[task_id] = {"result": result, "reported_at": time.time()}
|
|
self.state.inflight.pop(task_id, None)
|
|
self._save()
|
|
return result
|
|
|
|
def recover_after_restart(self) -> list[Mapping[str, Any]]:
|
|
recovered: list[Mapping[str, Any]] = []
|
|
for task_id in tuple(self.state.inflight):
|
|
recovered.append(self.renew(task_id))
|
|
return recovered
|
|
|
|
def validate_execution(self, argv: Sequence[str], cwd: str | os.PathLike[str], env: Mapping[str, str]) -> None:
|
|
if not argv or any(not isinstance(item, str) or not item for item in argv):
|
|
raise RunnerError("argv_only_execution_required")
|
|
if any(ch in argv[0] for ch in (";", "|", "&", "`", "\n")):
|
|
raise RunnerError("shell_syntax_denied")
|
|
executable = shutil.which(argv[0])
|
|
if not executable:
|
|
raise RunnerError("executable_not_found")
|
|
if self.config.executable_allowlist and not any(
|
|
os.path.realpath(executable).startswith(os.path.realpath(prefix).rstrip(os.sep) + os.sep)
|
|
or os.path.realpath(executable) == os.path.realpath(prefix)
|
|
for prefix in self.config.executable_allowlist
|
|
):
|
|
raise RunnerError("executable_not_allowlisted")
|
|
real_cwd = os.path.realpath(os.fspath(cwd))
|
|
if self.config.working_directory_allowlist and not any(
|
|
real_cwd == os.path.realpath(prefix)
|
|
or real_cwd.startswith(os.path.realpath(prefix).rstrip(os.sep) + os.sep)
|
|
for prefix in self.config.working_directory_allowlist
|
|
):
|
|
raise RunnerError("working_directory_not_allowlisted")
|
|
if self.config.environment_allowlist and any(key not in self.config.environment_allowlist for key in env):
|
|
raise RunnerError("environment_key_not_allowlisted")
|
|
|
|
def execute_argv(self, argv: Sequence[str], cwd: str | os.PathLike[str], env: Mapping[str, str]) -> Mapping[str, Any]:
|
|
self.validate_execution(argv, cwd, env)
|
|
completed = subprocess.run(list(argv), cwd=os.fspath(cwd), env=dict(env), shell=False, text=True, capture_output=True, check=False)
|
|
return {"returncode": completed.returncode, "stdout": completed.stdout, "stderr": completed.stderr}
|