aps-agent/server/agent_core/approval_store.py

from __future__ import annotations

import copy
import hashlib
import json
import math
import os
import tempfile
import threading
import time
import uuid
import warnings
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Protocol

_SCHEMA_VERSION = 1
_MAX_HISTORY = 2_000
_DATABASE_BACKEND = "database"
_FILE_BACKEND = "file"

RecordPredicate = Callable[[dict[str, Any]], bool]


class ApprovalBackend(Protocol):
    """Atomic approval operations shared by file and database backends."""

    backend: str

    def stage(self, record: dict[str, Any], *, now_epoch: float) -> bool: ...

    def decide(
        self,
        confirm_id: str,
        *,
        approve: bool,
        approver: dict[str, Any],
        allowed: RecordPredicate,
        now_epoch: float,
        grant_ttl_seconds: int,
        note: str | None = None,
    ) -> dict[str, Any] | None: ...

    def pending_record(
        self, confirm_id: str, *, allowed: RecordPredicate, now_epoch: float
    ) -> dict[str, Any] | None: ...

    def transfer(
        self,
        confirm_id: str,
        *,
        to_user_id: int,
        allowed: RecordPredicate,
        now_epoch: float,
        actor: dict[str, Any],
    ) -> dict[str, Any] | None: ...

    def consume_grant(
        self,
        grant: str,
        *,
        confirm_id: str,
        action: str,
        params_hash: str,
        allowed: RecordPredicate,
        decided_by: dict[str, Any],
        now_epoch: float,
    ) -> bool: ...

    def pending_items(
        self,
        *,
        tenant_uuid: str,
        project_id: str,
        allowed: RecordPredicate,
        now_epoch: float,
    ) -> list[dict[str, Any]]: ...

    def history_items(
        self,
        *,
        tenant_uuid: str,
        project_id: str,
        allowed: RecordPredicate,
        limit: int,
        now_epoch: float,
    ) -> list[dict[str, Any]]: ...

    def expire(self, *, now_epoch: float | None = None) -> bool: ...

    def clear(self) -> None: ...


def utc_now(epoch: float | None = None) -> str:
    value = time.time() if epoch is None else epoch
    return datetime.fromtimestamp(value, tz=UTC).isoformat()


def token_digest(token: str) -> str:
    return hashlib.sha256(token.encode("utf-8")).hexdigest()


def history_event(
    record: dict[str, Any],
    *,
    status: str,
    decided_by: dict[str, Any] | None = None,
    at_epoch: float | None = None,
    note: str | None = None,
) -> dict[str, Any]:
    decided_epoch = time.time() if at_epoch is None else at_epoch
    created_epoch = record.get("createdAtEpoch")
    dwell_ms = None
    if isinstance(created_epoch, (int, float)):
        dwell_ms = max(0, round((decided_epoch - float(created_epoch)) * 1000))
    event = {
        "confirmId": record.get("confirmId"),
        "sessionId": record.get("sessionId"),
        "action": record.get("action"),
        "power": record.get("power"),
        "status": status,
        "tenantUuid": record.get("tenantUuid"),
        "projectId": record.get("projectId"),
        "worldKey": record.get("worldKey"),
        "ownerUserId": record.get("ownerUserId"),
        "requester": record.get("requester"),
        "paramsHash": record.get("paramsHash"),
        "approvals": list(record.get("approvals") or []),
        "createdAt": record.get("createdAt"),
        "expiresAt": record.get("expiresAt"),
        "decidedAt": utc_now(decided_epoch),
        "dwellMs": dwell_ms,
    }
    if decided_by is not None:
        event["decidedBy"] = decided_by
    if note:
        event["note"] = note
    return event


def _approver_key(approver: dict[str, Any]) -> str:
    return str(approver.get("userId") or approver.get("username") or "")


def _grant_record(
    request: dict[str, Any],
    *,
    approvals: list[dict[str, Any]],
    now_epoch: float,
    grant_ttl_seconds: int,
) -> dict[str, Any]:
    return {
        "confirmId": request.get("confirmId"),
        "sessionId": request.get("sessionId"),
        "action": request.get("action"),
        "power": request.get("power"),
        "tenantUuid": request.get("tenantUuid"),
        "ownerUserId": request.get("ownerUserId"),
        "projectId": request.get("projectId"),
        "worldKey": request.get("worldKey"),
        "requester": request.get("requester"),
        "approvals": approvals,
        "paramsHash": request.get("paramsHash"),
        "createdAt": utc_now(now_epoch),
        "createdAtEpoch": now_epoch,
        "expiresAt": utc_now(now_epoch + grant_ttl_seconds),
        "expiresAtEpoch": now_epoch + grant_ttl_seconds,
    }


def default_approval_path() -> str:
    configured = os.environ.get("APS_APPROVAL_PATH")
    if configured:
        return str(Path(configured).expanduser().resolve())
    from server.aps_home import path_under_data

    return str(path_under_data("approvals.json"))


class ApprovalStore:
    """Persistent approval queue and one-time execution grants."""

    backend = _FILE_BACKEND

    def __init__(self, path: str | None = None) -> None:
        self.path = path or default_approval_path()
        self.lock_path = f"{self.path}.lock"
        self.lock = threading.RLock()
        self._transaction_state = threading.local()
        state = self._load()
        self.pending: dict[str, dict[str, Any]] = state["pending"]
        self.grants: dict[str, dict[str, Any]] = state["grants"]
        self.history: list[dict[str, Any]] = state["history"]

    def _empty(self) -> dict[str, Any]:
        return {
            "schemaVersion": _SCHEMA_VERSION,
            "pending": {},
            "grants": {},
            "history": [],
        }

    def _quarantine(self, reason: str) -> None:
        if not os.path.exists(self.path):
            return
        quarantined = f"{self.path}.corrupt-{int(time.time() * 1000)}-{os.getpid()}"
        try:
            os.replace(self.path, quarantined)
        except OSError as exc:
            warnings.warn(
                f"approval store is invalid ({reason}) and could not be quarantined: {exc}",
                RuntimeWarning,
                stacklevel=2,
            )
            return
        warnings.warn(
            f"approval store is invalid ({reason}); moved to {quarantined}",
            RuntimeWarning,
            stacklevel=2,
        )

    @staticmethod
    def _records_valid(records: dict[str, Any]) -> bool:
        for key, record in records.items():
            if not isinstance(key, str) or not isinstance(record, dict):
                return False
            try:
                expires_at = float(record.get("expiresAtEpoch"))
            except (TypeError, ValueError):
                return False
            if (
                not math.isfinite(expires_at)
                or expires_at <= 0
                or not isinstance(record.get("action"), str)
            ):
                return False
            if not isinstance(record.get("tenantUuid"), str):
                return False
        return True

    def _load(self) -> dict[str, Any]:
        try:
            with open(self.path, "r", encoding="utf-8") as handle:
                raw = json.load(handle)
        except FileNotFoundError:
            return self._empty()
        except (json.JSONDecodeError, OSError, TypeError) as exc:
            self._quarantine(type(exc).__name__)
            return self._empty()
        if not isinstance(raw, dict) or raw.get("schemaVersion") != _SCHEMA_VERSION:
            self._quarantine("schema")
            return self._empty()
        pending = raw.get("pending")
        grants = raw.get("grants")
        history = raw.get("history")
        if not isinstance(pending, dict) or not isinstance(grants, dict) or not isinstance(history, list):
            self._quarantine("shape")
            return self._empty()
        if not self._records_valid(pending) or not self._records_valid(grants):
            self._quarantine("record")
            return self._empty()
        if any(not isinstance(record, dict) for record in history):
            self._quarantine("history")
            return self._empty()
        return {
            "schemaVersion": _SCHEMA_VERSION,
            "pending": pending,
            "grants": grants,
            "history": history[-_MAX_HISTORY:],
        }

    def _refresh(self) -> None:
        state = self._load()
        self.pending.clear()
        self.pending.update(state["pending"])
        self.grants.clear()
        self.grants.update(state["grants"])
        self.history[:] = state["history"]

    @contextmanager
    def _process_lock(self) -> Iterator[None]:
        parent = os.path.dirname(self.lock_path) or "."
        os.makedirs(parent, exist_ok=True)
        with open(self.lock_path, "a+b") as handle:
            handle.seek(0, os.SEEK_END)
            if handle.tell() == 0:
                handle.write(b"\0")
                handle.flush()
            handle.seek(0)
            if os.name == "nt":
                import msvcrt

                msvcrt.locking(handle.fileno(), msvcrt.LK_LOCK, 1)
                try:
                    yield
                finally:
                    handle.seek(0)
                    msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1)
            else:
                import fcntl

                fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
                try:
                    yield
                finally:
                    fcntl.flock(handle.fileno(), fcntl.LOCK_UN)

    @contextmanager
    def transaction(self, *, refresh: bool = True) -> Iterator[None]:
        """Serialize a read-modify-write operation across threads and processes."""
        with self.lock:
            depth = int(getattr(self._transaction_state, "depth", 0))
            if depth:
                self._transaction_state.depth = depth + 1
                try:
                    yield
                finally:
                    self._transaction_state.depth = depth
                return
            with self._process_lock():
                if refresh:
                    self._refresh()
                self._transaction_state.depth = 1
                try:
                    yield
                finally:
                    self._transaction_state.depth = 0

    def _write(self) -> None:
        parent = os.path.dirname(self.path) or "."
        os.makedirs(parent, exist_ok=True)
        fd, temporary = tempfile.mkstemp(dir=parent, suffix=".tmp")
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as handle:
                json.dump(
                    {
                        "schemaVersion": _SCHEMA_VERSION,
                        "pending": self.pending,
                        "grants": self.grants,
                        "history": self.history[-_MAX_HISTORY:],
                    },
                    handle,
                    ensure_ascii=False,
                    separators=(",", ":"),
                )
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(temporary, self.path)
        except BaseException:
            if os.path.exists(temporary):
                os.unlink(temporary)
            raise

    def save(self) -> None:
        if int(getattr(self._transaction_state, "depth", 0)):
            self._write()
            return
        with self.transaction(refresh=False):
            self._write()

    def append_history(
        self,
        record: dict[str, Any],
        *,
        status: str,
        decided_by: dict[str, Any] | None = None,
        at_epoch: float | None = None,
        note: str | None = None,
    ) -> None:
        self.history.append(history_event(
            record,
            status=status,
            decided_by=decided_by,
            at_epoch=at_epoch,
            note=note,
        ))
        self.history = self.history[-_MAX_HISTORY:]

    def stage(self, record: dict[str, Any], *, now_epoch: float) -> bool:
        with self.transaction():
            self.expire(now_epoch=now_epoch)
            confirm_id = str(record.get("confirmId") or "")
            if not confirm_id or confirm_id in self.pending:
                return False
            self.pending[confirm_id] = copy.deepcopy(record)
            self.save()
            return True

    def decide(
        self,
        confirm_id: str,
        *,
        approve: bool,
        approver: dict[str, Any],
        allowed: RecordPredicate,
        now_epoch: float,
        grant_ttl_seconds: int,
        note: str | None = None,
    ) -> dict[str, Any] | None:
        with self.transaction():
            self.expire(now_epoch=now_epoch)
            pending = self.pending.get(confirm_id)
            if pending is None or not allowed(pending):
                return None
            if not approve:
                result = self.pending.pop(confirm_id)
                result["needsSecondConfirm"] = False
                self.append_history(
                    result, status="REJECTED", decided_by=approver, at_epoch=now_epoch,
                    note=note,
                )
                self.save()
                return result

            required = int(pending.get("requiredApprovals") or 1)
            approvals = list(pending.get("approvals") or [])
            if required > 1 and approvals and _approver_key(approvals[-1]) == _approver_key(approver):
                result = copy.deepcopy(pending)
                result["needsSecondConfirm"] = True
                result["approvalDenied"] = True
                result["separationRequired"] = True
                self.append_history(
                    result,
                    status="SOD_DENIED",
                    decided_by=approver,
                    at_epoch=now_epoch,
                    note="P3 second approval must be performed by a different user",
                )
                self.save()
                return result

            approval_entry = {**approver, "approvedAt": utc_now(now_epoch)}
            if note:
                approval_entry["note"] = note
            approvals.append(approval_entry)
            current = len(approvals)
            pending["approvals"] = approvals
            pending["approvalStep"] = current
            if current < required:
                result = copy.deepcopy(pending)
                result["needsSecondConfirm"] = True
                self.append_history(
                    result,
                    status="PARTIALLY_APPROVED",
                    decided_by=approver,
                    at_epoch=now_epoch,
                )
                self.save()
                return result

            result = self.pending.pop(confirm_id)
            result["approvalStep"] = required
            result["needsSecondConfirm"] = False
            if required > 1:
                grant = uuid.uuid4().hex
                self.grants[grant] = _grant_record(
                    result,
                    approvals=approvals,
                    now_epoch=now_epoch,
                    grant_ttl_seconds=grant_ttl_seconds,
                )
                result["executionGrant"] = grant
            self.append_history(
                result, status="APPROVED", decided_by=approver, at_epoch=now_epoch,
                note=note,
            )
            self.save()
            return result

    def pending_record(
        self, confirm_id: str, *, allowed: RecordPredicate, now_epoch: float
    ) -> dict[str, Any] | None:
        with self.transaction():
            self.expire(now_epoch=now_epoch)
            pending = self.pending.get(confirm_id)
            if pending is None or not allowed(pending):
                return None
            return copy.deepcopy(pending)

    def transfer(
        self,
        confirm_id: str,
        *,
        to_user_id: int,
        allowed: RecordPredicate,
        now_epoch: float,
        actor: dict[str, Any],
    ) -> dict[str, Any] | None:
        """把待确认项转派给指定用户（P2 单重审批；转派后 owner=目标用户）。

        仅当前可审批者（owner/被委托人）可转派；P3 因需双人职责分离拒绝转派。
        """
        with self.transaction():
            self.expire(now_epoch=now_epoch)
            pending = self.pending.get(confirm_id)
            if pending is None or not allowed(pending):
                return None
            if str(pending.get("power") or "") == "P3":
                return None                                  # P3 不转派（保 SOD）
            if to_user_id == int(pending.get("ownerUserId") or 0):
                return copy.deepcopy(pending)                # 转给自己 = no-op
            before_owner = pending.get("ownerUserId")
            before_delegate = pending.get("delegateUserId")
            pending["ownerUserId"] = to_user_id
            pending["delegateUserId"] = None
            pending["transferredAt"] = utc_now(now_epoch)
            result = copy.deepcopy(pending)
            self.append_history(
                result, status="TRANSFERRED", decided_by=actor, at_epoch=now_epoch,
                note=f"from_owner={before_owner};delegate={before_delegate};to={to_user_id}",
            )
            self.save()
            return result

    def consume_grant(
        self,
        grant: str,
        *,
        confirm_id: str,
        action: str,
        params_hash: str,
        allowed: RecordPredicate,
        decided_by: dict[str, Any],
        now_epoch: float,
    ) -> bool:
        with self.transaction():
            self.expire(now_epoch=now_epoch)
            record = self.grants.get(grant)
            if (
                record is None
                or record.get("confirmId") != confirm_id
                or record.get("action") != action
                or record.get("paramsHash") != params_hash
                or not allowed(record)
            ):
                return False
            self.grants.pop(grant, None)
            self.append_history(
                record, status="EXECUTED", decided_by=decided_by, at_epoch=now_epoch
            )
            self.save()
            return True

    def pending_items(
        self,
        *,
        tenant_uuid: str,
        project_id: str,
        allowed: RecordPredicate,
        now_epoch: float,
    ) -> list[dict[str, Any]]:
        del tenant_uuid, project_id
        with self.transaction():
            self.expire(now_epoch=now_epoch)
            return [
                {
                    "confirmId": confirm_id,
                    **copy.deepcopy(info),
                    "waitingMs": max(
                        0,
                        round((now_epoch - float(info.get("createdAtEpoch") or now_epoch)) * 1000),
                    ),
                }
                for confirm_id, info in self.pending.items()
                if allowed(info)
            ]

    def history_items(
        self,
        *,
        tenant_uuid: str,
        project_id: str,
        allowed: RecordPredicate,
        limit: int,
        now_epoch: float,
    ) -> list[dict[str, Any]]:
        del tenant_uuid, project_id
        with self.transaction():
            self.expire(now_epoch=now_epoch)
            visible = [copy.deepcopy(record) for record in self.history if allowed(record)]
            return visible[-limit:]

    def expire(self, *, now_epoch: float | None = None) -> bool:
        now = time.time() if now_epoch is None else now_epoch
        with self.transaction():
            changed = False
            for confirm_id, record in list(self.pending.items()):
                if float(record.get("expiresAtEpoch") or 0) <= now:
                    expired = self.pending.pop(confirm_id)
                    expired["confirmId"] = confirm_id
                    self.append_history(expired, status="EXPIRED", at_epoch=now)
                    changed = True
            for grant, record in list(self.grants.items()):
                if float(record.get("expiresAtEpoch") or 0) <= now:
                    expired = self.grants.pop(grant)
                    self.append_history(expired, status="GRANT_EXPIRED", at_epoch=now)
                    changed = True
            if changed:
                self.save()
            return changed

    def clear(self) -> None:
        with self.transaction():
            self.pending.clear()
            self.grants.clear()
            self.history.clear()
            self.save()


_store_lock = threading.Lock()
_store: ApprovalBackend | None = None
_store_key: tuple[str, ...] | None = None


def _backend_configuration() -> tuple[str, ...]:
    backend = (os.environ.get("APS_APPROVAL_BACKEND") or _FILE_BACKEND).strip().lower()
    if backend == _FILE_BACKEND:
        return (_FILE_BACKEND, default_approval_path())
    if backend != _DATABASE_BACKEND:
        raise RuntimeError("APS_APPROVAL_BACKEND must be 'file' or 'database'")
    if (os.environ.get("APS_APPROVAL_PATH") or "").strip():
        raise RuntimeError(
            "APS_APPROVAL_PATH cannot be combined with APS_APPROVAL_BACKEND=database"
        )
    database_url = (os.environ.get("APS_DATABASE_URL") or "").strip()
    if not database_url:
        raise RuntimeError(
            "APS_APPROVAL_BACKEND=database requires an explicit APS_DATABASE_URL"
        )
    from sqlalchemy.engine import make_url

    try:
        dialect = make_url(database_url).get_backend_name()
    except Exception as exc:
        raise RuntimeError("APS_DATABASE_URL is invalid") from exc
    allow_sqlite = os.environ.get("APS_APPROVAL_DATABASE_ALLOW_SQLITE") == "1"
    if dialect == "sqlite" and not allow_sqlite:
        raise RuntimeError(
            "SQLite approval storage is test-only; set "
            "APS_APPROVAL_DATABASE_ALLOW_SQLITE=1 explicitly"
        )
    if dialect not in {"mysql", "sqlite"}:
        raise RuntimeError("shared approval storage requires MySQL")
    return (_DATABASE_BACKEND, database_url, "sqlite-ok" if allow_sqlite else "mysql")


def _create_configured_store(configuration: tuple[str, ...]) -> ApprovalBackend:
    if configuration[0] == _FILE_BACKEND:
        return ApprovalStore(configuration[1])
    from server.agent_core.approval_db_store import DatabaseApprovalStore

    return DatabaseApprovalStore()


def get_approval_store() -> ApprovalBackend:
    global _store, _store_key
    configuration = _backend_configuration()
    with _store_lock:
        if _store is None or _store_key != configuration:
            _store = _create_configured_store(configuration)
            _store_key = configuration
        return _store


def set_approval_store(store: ApprovalBackend | None) -> ApprovalBackend:
    global _store, _store_key
    with _store_lock:
        if store is None:
            configuration = _backend_configuration()
            _store = _create_configured_store(configuration)
            _store_key = configuration
        else:
            _store = store
            if isinstance(store, ApprovalStore):
                _store_key = (_FILE_BACKEND, str(Path(store.path).resolve()))
            else:
                _store_key = (getattr(store, "backend", "explicit"), "explicit")
        return _store