agents/engineering_state.py

"""Versioned, bounded engineering lifecycle state for the unified agent loop.

The module is deliberately dependency-free.  It mirrors the legacy lifecycle without
being authoritative for recovery when the rollout mode is enabled, and it never stores
raw prompts, credentials, or arbitrary tool output.
"""
from __future__ import annotations

import hashlib
import os
import re
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Mapping

SCHEMA_VERSION = 1
MAX_HISTORY = 64
MAX_DIAGNOSTICS = 24
MAX_PREVIEW_CHARS = 256
MAX_ID_CHARS = 180

_SECRET_PATTERNS = (
    re.compile(r"(?i)(bearer\s+)[A-Za-z0-9._~+/=-]{8,}"),
    re.compile(r"(?i)(api[_-]?key\s*[:=]\s*)[^\s,;]+"),
    re.compile(r"(?i)(token\s*[:=]\s*)[^\s,;]+"),
    re.compile(r"(?i)\b(?:ghp|gho|github_pat|hf|sk|xoxb|xapp|r8)_[A-Za-z0-9_-]{8,}\b"),
    re.compile(r"\beyJ[A-Za-z0-9_-]{20,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\b"),
)


class EngineeringStateMode(str, Enum):
    OFF = "off"
    SHADOW = "shadow"
    CANARY = "canary"
    AUTHORITATIVE = "authoritative"


@dataclass(frozen=True)
class EngineeringStateConfig:
    """Conservative rollout configuration read once per run."""

    mode: EngineeringStateMode = EngineeringStateMode.OFF
    canary_rate: float = 0.0

    @classmethod
    def from_env(cls) -> "EngineeringStateConfig":
        raw_mode = os.getenv("ENGINEERING_STATE_MODE", "authoritative").strip().lower()  # P1 default; off remains an explicit rollback mode
        try:
            mode = EngineeringStateMode(raw_mode)
        except ValueError:
            mode = EngineeringStateMode.OFF
        try:
            rate = float(os.getenv("ENGINEERING_STATE_CANARY_RATE", "0"))
        except (TypeError, ValueError):
            rate = 0.0
        return cls(mode=mode, canary_rate=max(0.0, min(rate, 1.0)))

    @property
    def enabled(self) -> bool:
        return self.mode is not EngineeringStateMode.OFF

    def selects_canary(self, run_id: str, session_id: str) -> bool:
        if self.mode is not EngineeringStateMode.CANARY or not session_id:
            return False
        if self.canary_rate >= 1.0:
            return True
        if self.canary_rate <= 0.0:
            return False
        digest = hashlib.sha256(f"{run_id}:{session_id}".encode()).digest()
        bucket = int.from_bytes(digest[:8], "big") / float(2**64)
        return bucket < self.canary_rate


def _bounded_id(value: str | None) -> str:
    return re.sub(r"[^A-Za-z0-9_.:/-]", "_", str(value or ""))[:MAX_ID_CHARS]


def redact_text(value: object, max_chars: int = MAX_PREVIEW_CHARS) -> str:
    """Redact common credential forms before anything reaches a checkpoint."""
    text = str(value or "")[: max_chars * 4]
    for pattern in _SECRET_PATTERNS:
        if pattern.groups:
            text = pattern.sub(lambda match: f"{match.group(1)}[REDACTED]", text)
        else:
            text = pattern.sub("[REDACTED]", text)
    return text[:max_chars]


_ALLOWED_TRANSITIONS: dict[str, frozenset[str]] = {
    "IDLE": frozenset({"CLASSIFYING", "FAILED"}),
    "CLASSIFYING": frozenset({"TOOL_EXECUTING", "THINKING", "COMPLETED", "FAILED"}),
    "TOOL_EXECUTING": frozenset({"THINKING", "COMPLETED", "FAILED"}),
    "THINKING": frozenset({"COMPLETED", "FAILED"}),
    "FAILED": frozenset({"IDLE"}),
    "COMPLETED": frozenset({"IDLE", "FAILED"}),
}


@dataclass
class EngineeringState:
    """Bounded state envelope that can be persisted and safely restored."""

    run_id: str
    session_id: str
    checkpoint_id: str
    goal_digest: str
    goal_preview: str
    current_state: str = "IDLE"
    history: list[dict[str, Any]] = field(default_factory=list)
    diagnostics: list[str] = field(default_factory=list)
    revision: int = 0
    sequence: int = 0
    created_at_ms: int = field(default_factory=lambda: int(time.time() * 1000))
    updated_at_ms: int = field(default_factory=lambda: int(time.time() * 1000))

    @classmethod
    def start(
        cls,
        goal: str,
        *,
        run_id: str,
        session_id: str = "",
        checkpoint_id: str | None = None,
        now_ms: int | None = None,
    ) -> "EngineeringState":
        now = int(time.time() * 1000) if now_ms is None else int(now_ms)
        normalized_goal = str(goal or "")
        return cls(
            run_id=_bounded_id(run_id),
            session_id=_bounded_id(session_id),
            checkpoint_id=_bounded_id(checkpoint_id or session_id or run_id),
            goal_digest=hashlib.sha256(normalized_goal.encode("utf-8", "replace")).hexdigest(),
            goal_preview=redact_text(normalized_goal),
            created_at_ms=now,
            updated_at_ms=now,
        )

    @property
    def status(self) -> str:
        if self.current_state == "COMPLETED":
            return "completed"
        if self.current_state == "FAILED":
            return "failed"
        return "active"

    def transition(self, next_state: str, *, now_ms: int | None = None) -> bool:
        """Apply an idempotent transition; reject illegal transitions deterministically."""
        target = str(next_state)
        if target == self.current_state:
            return False
        allowed = _ALLOWED_TRANSITIONS.get(self.current_state, frozenset())
        if target not in allowed:
            raise ValueError(f"Invalid EngineeringState transition: {self.current_state} -> {target}")
        now = int(time.time() * 1000) if now_ms is None else int(now_ms)
        self.sequence += 1
        self.revision += 1
        self.history.append({
            "sequence": self.sequence,
            "from_state": self.current_state,
            "to_state": target,
            "at_ms": now,
        })
        if len(self.history) > MAX_HISTORY:
            del self.history[:-MAX_HISTORY]
        self.current_state = target
        self.updated_at_ms = now
        return True

    def prepare_for_resume(self) -> None:
        """Normalize a restored snapshot before a new loop execution."""
        if self.current_state != "IDLE":
            self.current_state = "IDLE"
            self.revision += 1
            self.updated_at_ms = int(time.time() * 1000)
            self.diagnostic("resume normalized state to IDLE")

    def diagnostic(self, message: str) -> None:
        value = redact_text(message, 180)
        if not value or value in self.diagnostics:
            return
        self.diagnostics.append(value)
        if len(self.diagnostics) > MAX_DIAGNOSTICS:
            del self.diagnostics[:-MAX_DIAGNOSTICS]
        self.revision += 1
        self.updated_at_ms = int(time.time() * 1000)

    def snapshot(self) -> dict[str, Any]:
        """Return a bounded JSON-compatible envelope; never expose the raw goal."""
        return {
            "schema_version": SCHEMA_VERSION,
            "run_id": self.run_id,
            "session_id": self.session_id,
            "checkpoint_id": self.checkpoint_id,
            "goal_digest": self.goal_digest,
            "goal_preview": self.goal_preview,
            "status": self.status,
            "current_state": self.current_state,
            "revision": self.revision,
            "sequence": self.sequence,
            "history": list(self.history[-MAX_HISTORY:]),
            "diagnostics": list(self.diagnostics[-MAX_DIAGNOSTICS:]),
            "created_at_ms": self.created_at_ms,
            "updated_at_ms": self.updated_at_ms,
        }

    def projection(self) -> dict[str, Any]:
        """Small read-only view safe for API/SSE consumers."""
        return {
            "schema_version": SCHEMA_VERSION,
            "status": self.status,
            "current_state": self.current_state,
            "revision": self.revision,
            "sequence": self.sequence,
            "checkpoint_id": self.checkpoint_id,
            "history": [dict(item) for item in self.history[-16:]],
            "diagnostics": list(self.diagnostics[-MAX_DIAGNOSTICS:]),
        }

    @classmethod
    def from_snapshot(cls, payload: Mapping[str, Any]) -> "EngineeringState":
        if not isinstance(payload, Mapping):
            raise ValueError("engineering state must be an object")
        if int(payload.get("schema_version", -1)) != SCHEMA_VERSION:
            raise ValueError("unsupported engineering state schema")
        history = payload.get("history", [])
        diagnostics = payload.get("diagnostics", [])
        if not isinstance(history, list) or len(history) > MAX_HISTORY:
            raise ValueError("invalid engineering state history")
        if not isinstance(diagnostics, list) or len(diagnostics) > MAX_DIAGNOSTICS:
            raise ValueError("invalid engineering state diagnostics")
        current = str(payload.get("current_state", ""))
        if current not in _ALLOWED_TRANSITIONS:
            raise ValueError("invalid engineering state current state")
        revision = int(payload.get("revision", -1))
        sequence = int(payload.get("sequence", -1))
        if revision < 0 or sequence < 0 or revision < sequence:
            raise ValueError("invalid engineering state revision")
        state = cls(
            run_id=_bounded_id(str(payload.get("run_id", ""))),
            session_id=_bounded_id(str(payload.get("session_id", ""))),
            checkpoint_id=_bounded_id(str(payload.get("checkpoint_id", ""))),
            goal_digest=str(payload.get("goal_digest", "")),
            goal_preview=redact_text(payload.get("goal_preview", "")),
            current_state=current,
            history=[dict(item) for item in history if isinstance(item, Mapping)],
            diagnostics=[redact_text(item, 180) for item in diagnostics],
            revision=revision,
            sequence=sequence,
            created_at_ms=int(payload.get("created_at_ms", 0)),
            updated_at_ms=int(payload.get("updated_at_ms", 0)),
        )
        if len(state.goal_digest) != 64 or not re.fullmatch(r"[0-9a-f]{64}", state.goal_digest):
            raise ValueError("invalid engineering state goal digest")
        return state