mayhem-cli 0.5.1__py3-none-any.whl

mayhem/agents/server.py

@@ -0,0 +1,281 @@
+"""Agent-side JSON-RPC server loop over ndjson stdio (ADR-0003).
+
+The agent never listens on sockets: the controller spawns it (locally or over
+``ssh … mayhem-agent serve``) and drives the session through its stdin/stdout.
+stderr is reserved for logs, never protocol frames.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import sys
+from typing import TYPE_CHECKING, Any
+
+from mayhem.agents import protocol as rpc
+from mayhem.agents.watchdog import DEFAULT_SWEEP_INTERVAL_S, AgentWatchdog
+from mayhem.domain.leases import FaultLease
+from mayhem.toolkit.fingerprint import interpreter_marker
+
+if TYPE_CHECKING:
+    from mayhem.agents.executors import FaultExecutor
+
+AGENT_VERSION = "0.1.0.dev0"
+
+
+class TaskRecord:
+    def __init__(self, task_id: str, params: dict[str, Any]) -> None:
+        self.task_id = task_id
+        self.params = params
+        self.state = "queued"
+        self.detail = ""
+        self.handle: asyncio.Task[dict[str, Any]] | None = None
+
+
+class AgentServer:
+    """Serves one controller session. One instance per spawned process."""
+
+    def __init__(
+        self,
+        roles: tuple[str, ...],
+        executors: tuple[FaultExecutor, ...] = (),
+        agent_id: str | None = None,
+        watchdog_interval_s: float = DEFAULT_SWEEP_INTERVAL_S,
+    ) -> None:
+        self.agent_id = agent_id or f"ag-{interpreter_marker()}"
+        self.roles = roles
+        self._executors = executors
+        self._tasks: dict[str, TaskRecord] = {}
+        self._task_counter = 0
+        self.watchdog = AgentWatchdog(interval_s=watchdog_interval_s)
+
+    # -- method dispatch -----------------------------------------------------
+
+    async def handle(self, request: rpc.RpcRequest) -> rpc.RpcResponse:
+        try:
+            rpc._require_context(request.params)
+        except rpc.ProtocolError as exc:
+            return rpc.error_response(request.id, exc.code, str(exc))
+        handler = {
+            rpc.METHOD_HANDSHAKE: self._on_handshake,
+            rpc.METHOD_CAPABILITIES_QUERY: self._on_capabilities,
+            rpc.METHOD_TASK_EXECUTE: self._on_task_execute,
+            rpc.METHOD_TASK_CANCEL: self._on_task_cancel,
+            rpc.METHOD_TASK_STATUS: self._on_task_status,
+            rpc.METHOD_HEALTH_PING: self._on_health_ping,
+        }.get(request.method)
+        if handler is None:
+            return rpc.error_response(
+                request.id,
+                rpc.RpcErrorCode.METHOD_NOT_FOUND,
+                f"no such method: {request.method}",
+            )
+        try:
+            result = await handler(request.params)
+        except Exception as exc:  # the wire must never see a traceback crash
+            return rpc.error_response(
+                request.id,
+                rpc.RpcErrorCode.INTERNAL_ERROR,
+                f"{type(exc).__name__}: {exc}",
+            )
+        return rpc.result_response(request.id, result)
+
+    async def _on_handshake(self, params: dict[str, Any]) -> dict[str, Any]:
+        return {
+            "agent_id": self.agent_id,
+            "protocol": rpc.PROTOCOL_VERSION,
+            "agent_version": AGENT_VERSION,
+            "roles": list(self.roles),
+        }
+
+    async def _on_capabilities(self, params: dict[str, Any]) -> dict[str, Any]:
+        return {
+            "roles": list(self.roles),
+            "faults": sorted({f for ex in self._executors for f in ex.capable_faults()}),
+        }
+
+    async def _on_health_ping(self, params: dict[str, Any]) -> dict[str, Any]:
+        return {
+            "pong": True,
+            "inflight": len(self._tasks),
+            "owned_leases": self.watchdog.active_count(),
+        }
+
+    async def _on_task_execute(self, params: dict[str, Any]) -> dict[str, Any]:
+        fault_id = params.get("fault_id")
+        if not isinstance(fault_id, str) or not fault_id:
+            raise ValueError("task.execute requires fault_id")
+        executor = next((ex for ex in self._executors if ex.supports(fault_id)), None)
+        if executor is None:
+            raise LookupError(f"agent has no executor for fault {fault_id}")
+        try:
+            lease = (
+                params["lease"]
+                if isinstance(params["lease"], FaultLease)
+                else FaultLease.model_validate(params["lease"])
+            )
+        except Exception as exc:
+            raise ValueError(f"invalid lease payload: {exc}") from exc
+        self._task_counter += 1
+        task_id = f"t-{self._task_counter:06d}"
+        record = TaskRecord(task_id, params)
+        record.state = "running"
+        self._tasks[task_id] = record
+        phase = params.get("phase", "inject")
+        try:
+            outcome = (
+                await asyncio.to_thread(executor.inject, lease)
+                if phase == "inject"
+                else await asyncio.to_thread(executor.undo, lease)
+            )
+        except Exception as exc:
+            record.state = "failed"
+            record.detail = f"{type(exc).__name__}: {exc}"
+            raise
+        if phase == "inject" and outcome.ok:
+            self.watchdog.register(
+                lease_id=str(lease.id),
+                lease=lease,
+                executor=executor,
+                ttl_seconds=float(lease.ttl_seconds),
+                task_id=task_id,
+            )
+        record.state = "done"
+        record.detail = outcome.detail
+        return {
+            "task_id": task_id,
+            "state": "done",
+            "ok": outcome.ok,
+            "detail": outcome.detail,
+        }
+
+    async def _on_task_cancel(self, params: dict[str, Any]) -> dict[str, Any]:
+        task_id = params.get("task_id")
+        record = self._tasks.get(task_id) if isinstance(task_id, str) else None
+        if record is None:
+            return {"task_id": task_id, "state": "unknown"}
+        if record.handle is not None and not record.handle.done():
+            record.handle.cancel()
+        record.state = "cancelled"
+        return {"task_id": task_id, "state": "cancelled"}
+
+    async def _on_task_status(self, params: dict[str, Any]) -> dict[str, Any]:
+        requested = params.get("task_id")
+        records = (
+            [r for r in self._tasks.values() if r.task_id == requested]
+            if isinstance(requested, str)
+            else list(self._tasks.values())
+        )
+        return {
+            "tasks": [
+                {"task_id": r.task_id, "state": r.state, "detail": r.detail} for r in records
+            ],
+            "watchdog": self.watchdog.snapshot(),
+            "compensated": self.watchdog.history(),
+        }
+
+    # -- io loop ---------------------------------------------------------------
+
+    async def serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        """Consume requests until EOF; every line gets exactly one response."""
+        while True:
+            raw = await reader.readline()
+            if not raw:
+                break
+            await self._serve_line(raw.strip(), writer)
+        if not writer.is_closing():
+            await writer.drain()
+            writer.close()
+
+    async def _serve_line(self, raw: bytes, writer: asyncio.StreamWriter) -> None:
+        try:
+            frame = rpc.decode(raw)
+        except rpc.ProtocolError as exc:
+            response = rpc.error_response(None, exc.code, str(exc))
+            writer.write(rpc.encode(response))
+            await writer.drain()
+            return
+        if isinstance(frame, rpc.RpcResponse):
+            return  # agents do not consume responses on their own session
+        if isinstance(frame, rpc.RpcNotification):
+            return  # event.emit / log.emit are accepted and dropped (logged via stderr)
+        response = await self.handle(frame)
+        writer.write(rpc.encode(response))
+        await writer.drain()
+
+
+def serve_sync(roles: tuple[str, ...], executors: tuple[FaultExecutor, ...]) -> None:
+    """Blocking stdio entrypoint used by ``mayhem-agent serve``.
+
+    Prefers a fully-async session so the ADR-0005 watchdog keeps sweeping even
+    while stdin is idle; falls back to per-line pumping where the platform
+    cannot wrap stdio as pipes.
+    """
+    server = AgentServer(roles=roles, executors=executors)
+    loop = asyncio.new_event_loop()
+    try:
+        loop.run_until_complete(_serve_async(server))
+    except (NotImplementedError, ValueError):
+        _serve_blocking(server, loop)
+    finally:
+        loop.close()
+
+
+async def _serve_async(server: AgentServer) -> None:
+    loop = asyncio.get_running_loop()
+    reader = asyncio.StreamReader()
+    await loop.connect_read_pipe(lambda: asyncio.StreamReaderProtocol(reader), sys.stdin)
+    transport, protocol = await loop.connect_write_pipe(
+        asyncio.streams.FlowControlMixin, sys.stdout
+    )
+    writer = asyncio.StreamWriter(transport, protocol, None, loop)
+    server.watchdog.start()
+    try:
+        await server.serve(reader, writer)
+    finally:
+        await server.watchdog.stop()
+
+
+def _serve_blocking(server: AgentServer, loop: asyncio.AbstractEventLoop) -> None:
+    """Per-line fallback: watchdog sweeps piggyback on request processing."""
+    while True:
+        raw = sys.stdin.readline()
+        if not raw:
+            break
+        line = raw.strip()
+        if not line:
+            continue
+        try:
+            frame = rpc.decode(line.encode("utf-8"))
+        except rpc.ProtocolError as exc:
+            _write_frame(rpc.error_response(None, exc.code, str(exc)))
+            continue
+        if isinstance(frame, rpc.RpcResponse):
+            continue
+        if isinstance(frame, rpc.RpcNotification):
+            continue  # accepted; logs/events surface on stderr
+        try:
+            response = loop.run_until_complete(
+                asyncio.gather(server.handle(frame), server.watchdog.sweep())
+            )[0]
+        except Exception as exc:  # session must survive handler bugs
+            response = rpc.error_response(
+                frame.id,
+                rpc.RpcErrorCode.INTERNAL_ERROR,
+                f"{type(exc).__name__}: {exc}",
+            )
+        _write_frame(response)
+
+
+def _write_frame(response: rpc.RpcResponse) -> None:
+    sys.stdout.write(rpc.encode(response).decode("utf-8"))
+    sys.stdout.flush()
+
+
+def frames_from_jsonl(text: str) -> list[dict[str, Any]]:
+    """Debug helper: pretty-print captured ndjson traffic (`| jq` equivalent)."""
+    out: list[dict[str, Any]] = []
+    for line in text.splitlines():
+        if line.strip():
+            out.append(json.loads(line))
+    return out

mayhem/agents/sinks.py

@@ -0,0 +1,60 @@
+"""Lease persistence boundary for the agent SDK.
+
+Agents never touch SQLite directly (the controller is the single writer,
+ADR-0007); they depend on this protocol. The controller binds a SQLite-backed
+implementation at wiring time; tests use the in-memory fake.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from mayhem.domain.leases import FaultLease
+
+
+class LeaseSink(Protocol):
+    def save(self, lease: FaultLease) -> None: ...
+
+    def load(self, lease_id: str) -> FaultLease | None: ...
+
+    def active_leases(self) -> tuple[FaultLease, ...]: ...
+
+    def next_sequence(self) -> int:
+        """Highest ``l-<n>`` sequence already persisted (0 when empty).
+
+        Lets a fresh run continue the counter instead of restarting from one,
+        so repeated runs against a shared DB never reuse a lease id
+        (``fault_leases.id`` / ``fault_invocations.lease_id`` are unique).
+        """
+        ...
+
+
+class InMemoryLeaseSink:
+    """Thread-hostile by design: the agent event loop is single-threaded."""
+
+    def __init__(self) -> None:
+        self._leases: dict[str, FaultLease] = {}
+
+    def save(self, lease: FaultLease) -> None:
+        self._leases[lease.id] = lease
+
+    def load(self, lease_id: str) -> FaultLease | None:
+        return self._leases.get(lease_id)
+
+    def active_leases(self) -> tuple[FaultLease, ...]:
+        return tuple(lease for lease in self._leases.values() if not lease.is_safe_terminal)
+
+    def next_sequence(self) -> int:
+        highest = 0
+        for lease_id in self._leases:
+            if not lease_id.startswith("l-"):
+                continue
+            try:
+                highest = max(highest, int(lease_id[2:]))
+            except ValueError:
+                continue
+        return highest
+
+    def all_leases(self) -> tuple[FaultLease, ...]:
+        return tuple(self._leases.values())

mayhem/agents/transports.py

@@ -0,0 +1,134 @@
+"""Controller-side transports for agent sessions (ADR-0003).
+
+Both transports spawn the same wire protocol; only the launcher differs.
+- ``LocalStdioTransport`` spawns ``mayhem-agent serve`` as a child process.
+- ``SSHTransport`` runs a persistent ``ssh … -- mayhem-agent serve`` exec channel
+  using the system OpenSSH client (no asyncssh dependency).
+Agents never listen on sockets — every session is controller-initiated.
+
+Reconnect policy: exponential backoff with jitter. Channel loss does NOT cancel
+in-flight faults; leases survive independently (ADR-0005) and the controller
+re-synchronizes via capabilities.query + task.status after reconnecting.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import random
+from dataclasses import dataclass, field
+from typing import Any
+
+RECONNECT_BASE_DELAY_S = 0.5
+RECONNECT_MAX_DELAY_S = 30.0
+
+
+@dataclass(frozen=True)
+class SessionSpec:
+    """Everything needed to (re)spawn one agent session."""
+
+    roles: tuple[str, ...] = ("proc", "fs")
+    agent_command: str = "mayhem-agent"  # overridable for tests / non-PATH installs
+
+
+class Transport:
+    async def connect(self) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+        raise NotImplementedError
+
+    def describe(self) -> str:
+        raise NotImplementedError
+
+
+class LocalStdioTransport(Transport):
+    def __init__(self, spec: SessionSpec | None = None) -> None:
+        self._spec = spec or SessionSpec()
+        self._process: asyncio.subprocess.Process | None = None
+
+    def _argv(self) -> list[str]:
+        return [
+            self._spec.agent_command,
+            "serve",
+            "--roles",
+            ",".join(self._spec.roles),
+        ]
+
+    async def connect(self) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+        self._process = await asyncio.create_subprocess_exec(
+            *self._argv(),
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=None,  # agent logs flow to our stderr for local debugging
+        )
+        assert self._process.stdin is not None and self._process.stdout is not None
+        return self._process.stdout, self._process.stdin
+
+    def describe(self) -> str:
+        return f"local:{self._spec.agent_command}"
+
+
+class SSHTransport(Transport):
+    """Persistent ``ssh -o ControlMaster=auto -o ControlPersist=600 host -- <agent> serve``."""
+
+    def __init__(self, target: str, spec: SessionSpec | None = None) -> None:
+        self._target = target  # user@host or host
+        self._spec = spec or SessionSpec()
+        self._process: asyncio.subprocess.Process | None = None
+
+    def argv(self) -> list[str]:
+        remote = f"{self._spec.agent_command} serve --roles {','.join(self._spec.roles)}"
+        return [
+            "ssh",
+            "-o",
+            "ControlMaster=auto",
+            "-o",
+            "ControlPersist=600",
+            self._target,
+            "--",
+            remote,
+        ]
+
+    async def connect(self) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+        self._process = await asyncio.create_subprocess_exec(
+            *self.argv(),
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=None,
+        )
+        assert self._process.stdin is not None and self._process.stdout is not None
+        return self._process.stdout, self._process.stdin
+
+    def describe(self) -> str:
+        return f"ssh:{self._target}"
+
+
+@dataclass
+class ReconnectingSession:
+    """Connect with exponential backoff + jitter until success or max attempts."""
+
+    transport: Transport
+    attempts: int = 8
+    base_delay_s: float = RECONNECT_BASE_DELAY_S
+    max_delay_s: float = RECONNECT_MAX_DELAY_S
+    _rng: random.Random = field(default_factory=random.Random)
+
+    async def open(self) -> tuple[asyncio.StreamReader, asyncio.StreamWriter, int]:
+        delay = self.base_delay_s
+        last_error: Exception | None = None
+        for attempt in range(1, self.attempts + 1):
+            try:
+                reader, writer = await self.transport.connect()
+                return reader, writer, attempt
+            except (OSError, ProcessLookupError) as exc:
+                last_error = exc
+                jittered = delay * self._rng.uniform(0.5, 1.5)
+                await asyncio.sleep(jittered)
+                delay = min(delay * 2, self.max_delay_s)
+        msg = f"could not reach {self.transport.describe()} after {self.attempts} attempts"
+        raise ConnectionError(msg) from last_error
+
+
+def session_for(target: str | None, roles: tuple[str, ...], **kwargs: Any) -> Transport:
+    """Factory: ``None`` → local child process; ``user@host`` → SSH exec channel."""
+    spec = SessionSpec(roles=roles, **kwargs)
+    if target is None:
+        return LocalStdioTransport(spec)
+    return SSHTransport(target, spec)

mayhem/agents/watchdog.py

@@ -0,0 +1,140 @@
+"""Agent-side watchdog (ADR-0005 §3).
+
+Each agent expires its own leases past TTL and compensates locally, so safety
+does not depend on the controller ever coming back. Runs inside the agent
+process's event loop; every compensation is reported on stderr and queryable
+via ``task.status`` — never written to stdout (protocol frames only).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import sys
+import time
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from mayhem.agents.executors import FaultExecutor
+
+DEFAULT_SWEEP_INTERVAL_S = 1.0
+
+
+class ActiveLease:
+    """A fault this agent has injected and still owns."""
+
+    def __init__(
+        self,
+        *,
+        lease: Any,
+        executor: FaultExecutor,
+        deadline_epoch_s: float,
+        task_id: str | None = None,
+    ) -> None:
+        self.lease = lease
+        self.executor = executor
+        self.deadline_epoch_s = deadline_epoch_s
+        self.task_id = task_id
+        self.state = "active"  # active -> compensating -> expired | dirty
+
+
+class AgentWatchdog:
+    """Expires owned leases past TTL and undoes them locally."""
+
+    def __init__(self, *, interval_s: float = DEFAULT_SWEEP_INTERVAL_S) -> None:
+        self.interval_s = interval_s
+        self._active: dict[str, ActiveLease] = {}
+        self._final_states: dict[str, str] = {}
+        self._runner: asyncio.Task[None] | None = None
+        self.clock: Callable[[], float] = time.time
+
+    def register(
+        self,
+        *,
+        lease_id: str,
+        lease: Any,
+        executor: FaultExecutor,
+        ttl_seconds: float,
+        now_epoch_s: float | None = None,
+        task_id: str | None = None,
+    ) -> float:
+        """Track an injected fault; returns its absolute expiry deadline."""
+        now = time.time() if now_epoch_s is None else now_epoch_s
+        entry = ActiveLease(
+            lease=lease,
+            executor=executor,
+            deadline_epoch_s=now + max(ttl_seconds, 0.0),
+            task_id=task_id,
+        )
+        self._active[lease_id] = entry
+        return entry.deadline_epoch_s
+
+    def active_count(self) -> int:
+        return len(self._active)
+
+    async def sweep(self, *, now_epoch_s: float | None = None) -> list[str]:
+        """Undo every lease past its TTL; returns ids compensated this sweep."""
+        now = time.time() if now_epoch_s is None else now_epoch_s
+        expired = [lid for lid, e in self._active.items() if e.deadline_epoch_s <= now]
+        for lease_id in expired:
+            entry = self._active.pop(lease_id)
+            entry.state = "compensating"
+            try:
+                outcome = await asyncio.to_thread(entry.executor.undo, entry.lease)
+            except Exception as exc:
+                entry.state = "dirty"
+                print(
+                    f"[watchdog] undo raised for {lease_id}: {type(exc).__name__}: {exc}",
+                    file=sys.stderr,
+                )
+                self._final_states[lease_id] = entry.state
+                continue
+            entry.state = "expired" if outcome.ok else "dirty"
+            self._final_states[lease_id] = entry.state
+            print(
+                f"[watchdog] {entry.state} {lease_id}: {outcome.detail}",
+                file=sys.stderr,
+            )
+        return expired
+
+    def final_state(self, lease_id: str) -> str | None:
+        """Terminal watchdog state for a compensated lease: expired or dirty."""
+        return self._final_states.get(lease_id)
+
+    def history(self) -> list[dict[str, Any]]:
+        """Leases this agent already self-compensated (expired|dirty)."""
+        return [{"lease_id": lid, "state": state} for lid, state in self._final_states.items()]
+
+    async def run_forever(self) -> None:
+        while True:
+            await asyncio.sleep(self.interval_s)
+            try:
+                await self.sweep()
+            except Exception as exc:  # a broken sweep must not kill the session
+                print(f"[watchdog] sweep error: {exc}", file=sys.stderr)
+
+    def start(self) -> asyncio.Task[None]:
+        if self._runner is None or self._runner.done():
+            self._runner = asyncio.get_running_loop().create_task(self.run_forever())
+        return self._runner
+
+    async def stop(self) -> None:
+        if self._runner is not None:
+            self._runner.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._runner
+            self._runner = None
+
+    def snapshot(self) -> list[dict[str, Any]]:
+        now = self.clock()
+        return [
+            {
+                "lease_id": lid,
+                "state": e.state,
+                "expires_in_s": round(e.deadline_epoch_s - now, 3),
+                "task_id": e.task_id,
+            }
+            for lid, e in self._active.items()
+        ]

mayhem/cli/__init__.py

@@ -0,0 +1,11 @@
+"""mayhem CLI package — Click-based, unique-prefix command resolution.
+
+Public surface: ``app`` (the root group, for tests and embedding) and
+``main`` (the console-script entry point with exit-code mapping).
+"""
+
+from __future__ import annotations
+
+from mayhem.cli.app import app, main
+
+__all__ = ["app", "main"]