#!/usr/bin/env python3 """Durable local session/job state paths, validation, and publication.""" from __future__ import annotations import contextlib import copy import json import os import re from collections.abc import Mapping from pathlib import Path from typing import Any from mmo_util import ( SAFE_JOB_ID, append_jsonl, atomic_write_json, is_within, process_alive, process_group_alive, process_matches, read_json_object, state_root, terminate_process, terminate_process_group, utc_now, ) from mmo_version import MMO_SCHEMA_VERSION, PACKAGE_VERSION ACTIVE_JOB_STATUSES = { "queued", "starting", "running", "waiting", "paused", "detached", "recovering", "finalizing", "cancelling", } RECOVERABLE_JOB_STATUSES = ACTIVE_JOB_STATUSES | {"suspended"} ADMITTING_JOB_STATUSES = {"queued", "starting", "running", "waiting", "paused", "detached"} TERMINAL_JOB_STATUSES = { "completed", "completed_with_warnings", "failed", "stopped", "cancelled", } ADMITTING_SESSION_STATUSES = {"starting", "running", "detached", "paused", "suspended"} ACTIVE_SESSION_STATUSES = { "starting", "running", "detached", "paused", "suspended", "finishing", "stopping", "cancelling", } TERMINAL_SESSION_STATUSES = {"completed", "stopped", "failed", "cancelled"} ROOT_EXECUTION_HOSTS = frozenset({"app_server"}) RETIRED_SESSION_FIELDS = frozenset( { "root_execution_policy_enforced", "root_rollout_path", } ) RETIRED_JOB_FIELDS = frozenset( { "execution_policy", "renewal_quantum_seconds", "max_active_work_seconds", "active_work_cap_seconds", "active_work_seconds", "automatic_renewal_enabled", "app_server_rollout_path", } ) def _validate_switchyard_identity(data: Mapping[str, Any], label: str) -> None: if "switchyard_version" not in data: raise ValueError(f"{label} has no Switchyard version identity") version = data.get("switchyard_version") if data.get("gateway_base_url") is None: if version is not None: raise ValueError(f"{label} has a Switchyard version without a gateway") return if not isinstance(version, str) or re.fullmatch(r"[0-9]+\.[0-9]+\.[0-9]+", version) is None: raise ValueError(f"{label} Switchyard version is invalid") def root_mcp_token(session_id: str) -> str: """Read the canonical persisted root MCP capability for an active session.""" try: data = read_json_object( session_dir(session_id) / "capabilities.json", label="session capabilities", ) except FileNotFoundError as exc: raise RuntimeError("root MCP caller capability is unavailable for this session") from exc token = data.get("root") if not isinstance(token, str) or not token: raise RuntimeError("root MCP caller capability is invalid") return token def store_session_capabilities( directory: Path, *, root_token: str, native_tokens: Mapping[str, str], ) -> None: """Persist host-scoped MCP identities needed by detached app-server processes.""" if not root_token or not all( isinstance(value, str) and value for value in native_tokens.values() ): raise ValueError("session capabilities must be non-empty strings") atomic_write_json( directory / "capabilities.json", { "schema_version": MMO_SCHEMA_VERSION, "root": root_token, "native": dict(native_tokens), }, 0o600, ) def load_session_capabilities(directory: Path) -> tuple[str, dict[str, str]]: data = read_json_object(directory / "capabilities.json", label="session capabilities") root_token = data.get("root") native_tokens = data.get("native") if ( data.get("schema_version") != MMO_SCHEMA_VERSION or not isinstance(root_token, str) or not root_token or not isinstance(native_tokens, Mapping) or not all( isinstance(key, str) and isinstance(value, str) and value for key, value in native_tokens.items() ) ): raise ValueError("session capabilities are invalid") return root_token, dict(native_tokens) def root_mcp_capability_environment(session: Mapping[str, Any]) -> dict[str, str]: """Return the authenticated Agent-MCP identity for a root-owned process.""" session_id = str(session["session_id"]) return { "MMO_ROOT_SESSION_ID": session_id, "MMO_RUN_ID": str(session["current_run_id"]), "MMO_CALLER_AGENT": str(session["root_agent"]), "MMO_CALLER_TOKEN": root_mcp_token(session_id), } def revoke_session_capabilities(session_id: str) -> None: """Revoke the bearer identities for a terminal session.""" # A terminal session retains its immutable transcript and evidence, but its # bearer credentials no longer authorize any control-plane operation. (session_dir(session_id) / "capabilities.json").unlink(missing_ok=True) def sessions_root() -> Path: root = state_root() / "sessions" root.mkdir(parents=True, exist_ok=True, mode=0o700) return root def jobs_root() -> Path: root = state_root() / "jobs" root.mkdir(parents=True, exist_ok=True, mode=0o700) return root def runtime_lock_path() -> Path: return state_root() / ".runtime.lock" def session_dir(session_id: str) -> Path: if not SAFE_JOB_ID.fullmatch(session_id): raise ValueError("invalid session id") path = sessions_root() / session_id if not path.is_dir(): raise FileNotFoundError(f"unknown session: {session_id}") return path def job_dir(job_id: str) -> Path: if not SAFE_JOB_ID.fullmatch(job_id): raise ValueError("invalid job id") path = jobs_root() / job_id if not path.is_dir(): raise FileNotFoundError(f"unknown agent job: {job_id}") return path def session_state_path(directory: Path) -> Path: return directory / "session.json" def job_state_path(directory: Path) -> Path: return directory / "metadata.json" def ensure_runs_root(directory: Path) -> Path: root = directory / "runs" if root.is_symlink(): raise RuntimeError("session runs directory cannot be a symlink") root.mkdir(parents=True, exist_ok=True, mode=0o700) if not is_within(root.resolve(), directory.resolve()): raise RuntimeError("session runs directory escapes its logical session") return root def _run_path(directory: Path, run_id: str) -> Path: if not SAFE_JOB_ID.fullmatch(run_id): raise ValueError("invalid run id") root = directory / "runs" if root.is_symlink(): raise RuntimeError("session runs directory cannot be a symlink") path = root / run_id / "run.json" if not path.is_file(): raise FileNotFoundError(f"unknown session run: {run_id}") return path def session_lifecycle_lock_path(directory: Path) -> Path: path = directory / "lifecycle.lock" if directory.is_symlink() or path.is_symlink(): raise RuntimeError("persistent session lifecycle lock cannot traverse a symlink") if not is_within(path.resolve(), directory.resolve()): raise RuntimeError("persistent session lifecycle lock escapes its logical session") return path def job_control_lock_path(directory: Path) -> Path: path = directory / "control.lock" if directory.is_symlink() or path.is_symlink(): raise RuntimeError("persistent job control lock cannot traverse a symlink") if not is_within(path.resolve(), directory.resolve()): raise RuntimeError("persistent job control lock escapes its logical job") return path def _read_run_state(path: Path) -> dict[str, Any]: session_directory = path.parent.parent.parent if path.parent.parent.is_symlink() or path.parent.is_symlink() or path.is_symlink(): raise ValueError("session run state cannot traverse a symlink") if not is_within(path.resolve(), session_directory.resolve()): raise ValueError("session run state escapes its logical session") data = read_json_object(path, label="session run state") if data.get("schema_version") != MMO_SCHEMA_VERSION or isinstance( data.get("schema_version"), bool ): raise ValueError(f"unsupported session run schema; expected {MMO_SCHEMA_VERSION}") if data.get("package_version") != PACKAGE_VERSION: raise ValueError(f"persistent session run package must be {PACKAGE_VERSION}") sequence = data.get("sequence") if not isinstance(sequence, int) or isinstance(sequence, bool) or sequence != 1: raise ValueError("immutable session run sequence must be exactly 1") if data.get("kind") != "initial": raise ValueError("immutable session run kind must be initial") if data.get("status") not in ACTIVE_SESSION_STATUSES | TERMINAL_SESSION_STATUSES: raise ValueError("session run status is invalid") _validate_switchyard_identity(data, "persistent session run") expected_id = path.parent.name if data.get("run_id") != expected_id: raise ValueError( "session run identity mismatch: " f"directory is {expected_id!r}, record is {data.get('run_id')!r}" ) if data.get("session_id") != path.parent.parent.parent.name: raise ValueError("session run belongs to another logical session") return data def load_session_run(session_id: str, run_id: str) -> dict[str, Any]: return _read_run_state(_run_path(session_dir(session_id), run_id)) def iter_session_runs(session_id: str) -> list[dict[str, Any]]: directory = session_dir(session_id) root = directory / "runs" if root.is_symlink(): raise RuntimeError("session runs directory cannot be a symlink") if not root.is_dir(): return [] results: list[dict[str, Any]] = [] for candidate in sorted(root.iterdir(), reverse=True): path = candidate / "run.json" if not candidate.is_dir() or not path.is_file(): continue try: results.append(_read_run_state(path)) except (OSError, ValueError, json.JSONDecodeError) as exc: raise RuntimeError(f"invalid session run state: {path}: {exc}") from exc return results def public_run(data: Mapping[str, Any]) -> dict[str, Any]: keys = ( "run_id", "package_version", "session_id", "sequence", "kind", "status", "created_at", "started_at", "finished_at", "root_pid", "exit_code", "gateway_base_url", "switchyard_version", "root_execution_host", "root_goal_status", "root_goal_tokens_used", "root_goal_token_budget", "error", ) return {key: data.get(key) for key in keys if data.get(key) is not None} _RUN_MIRROR_FIELDS = ( "status", "started_at", "finished_at", "root_pid", "root_pgid", "root_start_token", "root_process_group_isolated", "exit_code", "error", "cancel_requested_at", "transition_started_at", "requested_terminal_status", "gateway_base_url", "gateway_pid", "gateway_hash", "gateway_routing_log_path", "switchyard_version", "route_availability", "root_mcp_token_hash", "native_token_hashes", "root_execution_host", "root_execution_mode", "root_thread_id", "root_thread_generation", "root_thread_lineage", "root_thread_transition", "root_goal_status", "root_goal_objective", "root_goal_bootstrap_pending", "root_goal_token_budget", "root_max_goal_token_budget", "root_goal_tokens_used", "root_goal_time_used_seconds", "root_stall_warning_seconds", "root_last_progress_at", "root_finalization_grace_seconds", "root_finalization_started_at", "root_app_server_lifecycle_timeout_seconds", "root_pending_request_count", "root_last_turn_id", "root_turn_start_pending", "root_finalizing", "root_completion_deferred", "root_completion_deferred_jobs", ) def _validate_root_thread_lineage(data: Mapping[str, Any]) -> None: generation = data.get("root_thread_generation") lineage = data.get("root_thread_lineage") transition = data.get("root_thread_transition") if not isinstance(generation, int) or isinstance(generation, bool) or generation < 0: raise ValueError("persistent session root thread generation is invalid") if not isinstance(lineage, list): raise ValueError("persistent session root thread lineage must be a list") if len(lineage) != generation: raise ValueError("persistent session root thread lineage disagrees with its generation") seen: set[str] = set() for index, raw in enumerate(lineage, 1): if not isinstance(raw, Mapping): raise ValueError("persistent session root thread lineage entry is invalid") thread_id = raw.get("thread_id") if ( raw.get("generation") != index or isinstance(raw.get("generation"), bool) or not isinstance(thread_id, str) or not thread_id or thread_id in seen or not isinstance(raw.get("codex_session_id"), str) or not raw.get("codex_session_id") or not isinstance(raw.get("adopted_at"), str) or not raw.get("adopted_at") or not isinstance(raw.get("reason"), str) or not raw.get("reason") ): raise ValueError("persistent session root thread lineage entry is invalid") seen.add(thread_id) superseded_at = raw.get("superseded_at") successor_id = raw.get("successor_thread_id") terminal_entry = index == generation if terminal_entry: if superseded_at is not None or successor_id is not None: raise ValueError("current root thread lineage entry cannot be superseded") else: successor = lineage[index] if ( not isinstance(successor, Mapping) or not isinstance(superseded_at, str) or not superseded_at or not isinstance(successor_id, str) or successor_id != successor.get("thread_id") ): raise ValueError("superseded root thread lineage entry is incomplete") root_thread_id = data.get("root_thread_id") if generation == 0: if root_thread_id is not None: raise ValueError("uninitialized root thread lineage has a current thread") elif root_thread_id != lineage[-1].get("thread_id"): raise ValueError("persistent session current root thread disagrees with its lineage") if transition is not None: if not isinstance(transition, Mapping): raise ValueError("persistent session root thread transition is invalid") if ( transition.get("from_thread_id") != root_thread_id or not isinstance(transition.get("to_thread_id"), str) or not transition.get("to_thread_id") or transition.get("to_thread_id") in seen or transition.get("generation") != generation + 1 or isinstance(transition.get("generation"), bool) or not isinstance(transition.get("observed_at"), str) or not transition.get("observed_at") or not isinstance(transition.get("reason"), str) or not transition.get("reason") ): raise ValueError("persistent session root thread transition is inconsistent") def mirror_active_run(directory: Path, session: Mapping[str, Any]) -> None: run_id = session.get("current_run_id") if not isinstance(run_id, str): return path = _run_path(directory, run_id) run = _read_run_state(path) for key in _RUN_MIRROR_FIELDS: if key in session: run[key] = copy.deepcopy(session[key]) else: run.pop(key, None) atomic_write_json(path, run) def _validate_session_record(path: Path) -> dict[str, Any]: data = read_json_object(path, label="session state") retired = sorted(RETIRED_SESSION_FIELDS & set(data)) if retired: raise ValueError("persistent session contains retired fields: " + ", ".join(retired)) if data.get("schema_version") != MMO_SCHEMA_VERSION or isinstance( data.get("schema_version"), bool ): raise ValueError(f"unsupported session state schema; expected {MMO_SCHEMA_VERSION}") if data.get("package_version") != PACKAGE_VERSION: raise ValueError(f"persistent session package must be {PACKAGE_VERSION}") if data.get("profile_version") != PACKAGE_VERSION: raise ValueError(f"persistent session profile must be {PACKAGE_VERSION}") if data.get("session_kind") not in {"interactive", "noninteractive"}: raise ValueError("persistent session kind must be interactive or noninteractive") if data.get("root_execution_host") not in ROOT_EXECUTION_HOSTS: raise ValueError("persistent session root execution host is invalid") _validate_switchyard_identity(data, "persistent session") _validate_root_thread_lineage(data) socket_path = data.get("root_app_server_socket") if not isinstance(socket_path, str) or not Path(socket_path).is_absolute(): raise ValueError("persistent session app-server socket path is invalid") run_sequence = data.get("run_sequence") if not isinstance(run_sequence, int) or isinstance(run_sequence, bool) or run_sequence != 1: raise ValueError("immutable session run sequence must be exactly 1") for key in ("current_run_id", "last_run_id"): run_id = data.get(key) if run_id is not None and ( not isinstance(run_id, str) or not SAFE_JOB_ID.fullmatch(run_id) ): raise ValueError(f"persistent session {key} is invalid") expected_id = path.parent.name if data.get("session_id") != expected_id: raise ValueError( "session state identity mismatch: " f"directory is {expected_id!r}, record is {data.get('session_id')!r}" ) return data def _validate_job_record(path: Path) -> dict[str, Any]: data = read_json_object(path, label="job state") retired = sorted(RETIRED_JOB_FIELDS & set(data)) if retired: raise ValueError("persistent job contains retired fields: " + ", ".join(retired)) if data.get("schema_version") != MMO_SCHEMA_VERSION or isinstance( data.get("schema_version"), bool ): raise ValueError(f"unsupported job state schema; expected {MMO_SCHEMA_VERSION}") if data.get("package_version") != PACKAGE_VERSION: raise ValueError(f"persistent job package must be {PACKAGE_VERSION}") expected_id = path.parent.name if data.get("job_id") != expected_id: raise ValueError( "job state identity mismatch: " f"directory is {expected_id!r}, record is {data.get('job_id')!r}" ) return data def read_session_record(directory: Path) -> dict[str, Any]: """Read and validate the session record owned by ``directory``.""" return _validate_session_record(session_state_path(directory)) def read_job_record(directory: Path) -> dict[str, Any]: """Read and validate the worker-job record owned by ``directory``.""" return _validate_job_record(job_state_path(directory)) def publish_session_record( directory: Path, session: Mapping[str, Any], *, mirror_run: bool, ) -> None: """Atomically publish session state and optionally refresh its active-run mirror.""" atomic_write_json(session_state_path(directory), dict(session)) if mirror_run: mirror_active_run(directory, session) def publish_initial_session_records( directory: Path, session: Mapping[str, Any], run: Mapping[str, Any], ) -> None: """Publish a new run before making its canonical session record visible.""" run_id = run.get("run_id") if not isinstance(run_id, str) or not SAFE_JOB_ID.fullmatch(run_id): raise ValueError("initial session run id is invalid") if run.get("session_id") != session.get("session_id"): raise ValueError("initial session and run identities disagree") run_directory = ensure_runs_root(directory) / run_id if run_directory.is_symlink(): raise RuntimeError("initial session run directory cannot be a symlink") run_directory.mkdir(mode=0o700) atomic_write_json(run_directory / "run.json", dict(run)) publish_session_record(directory, session, mirror_run=False) def publish_job_record(directory: Path, job: Mapping[str, Any]) -> None: """Atomically publish the canonical worker-job record.""" atomic_write_json(job_state_path(directory), dict(job)) def terminate_recorded_process_group( data: Mapping[str, Any], *, prefix: str, grace_seconds: float, ) -> None: pid = data.get(f"{prefix}_pid") start_token = data.get(f"{prefix}_start_token") if ( not isinstance(pid, int) or isinstance(pid, bool) or pid <= 1 or pid == os.getpid() or not isinstance(start_token, str) ): return if bool(data.get(f"{prefix}_process_group_isolated", True)): pgid = data.get(f"{prefix}_pgid", pid) if not isinstance(pgid, int) or isinstance(pgid, bool) or pgid <= 1 or pgid != pid: return if process_matches(pid, start_token) or ( not process_alive(pid) and process_group_alive(pgid) ): terminate_process_group(pgid, grace_seconds=grace_seconds) if process_group_alive(pgid): raise RuntimeError(f"{prefix} process group {pgid} did not terminate") return if process_matches(pid, start_token): terminate_process(pid, grace_seconds=grace_seconds) if process_matches(pid, start_token): raise RuntimeError(f"{prefix} process {pid} did not terminate") def iter_session_records(*, strict: bool = True) -> list[dict[str, Any]]: """Enumerate validated session records without lifecycle side effects.""" results: list[dict[str, Any]] = [] for directory in sorted(sessions_root().iterdir(), reverse=True): path = session_state_path(directory) if not directory.is_dir() or not path.is_file(): continue try: results.append(_validate_session_record(path)) except (OSError, ValueError, json.JSONDecodeError) as exc: if strict: raise RuntimeError( f"invalid session state blocks safe accounting: {path}: {exc}" ) from exc continue return results def iter_job_records(*, strict: bool = True) -> list[dict[str, Any]]: """Enumerate validated worker-job records without lifecycle side effects.""" results: list[dict[str, Any]] = [] for directory in sorted(jobs_root().iterdir(), reverse=True): path = job_state_path(directory) if not directory.is_dir() or not path.is_file(): continue try: results.append(_validate_job_record(path)) except (OSError, ValueError, json.JSONDecodeError) as exc: if strict: raise RuntimeError( f"invalid job state blocks safe accounting: {path}: {exc}" ) from exc continue return results def append_audit(session_id: str, event: str, **data: Any) -> None: directory = session_dir(session_id) if not isinstance(data.get("run_id"), str): with contextlib.suppress(OSError, ValueError, json.JSONDecodeError): session = read_session_record(directory) run_id = session.get("current_run_id") or session.get("last_run_id") if isinstance(run_id, str): data["run_id"] = run_id append_jsonl( directory / "audit.jsonl", {"timestamp": utc_now(), "event": event, "session_id": session_id, **data}, )