new-qwen/serv/approvals.py

from __future__ import annotations

import json
import threading
import time
import uuid
from pathlib import Path
from typing import Any


class ApprovalStore:
    def __init__(self, base_dir: Path, retention_seconds: int) -> None:
        self.base_dir = base_dir
        self.base_dir.mkdir(parents=True, exist_ok=True)
        self.retention_seconds = retention_seconds
        self._approvals: dict[str, dict[str, Any]] = {}
        self._conditions: dict[str, threading.Condition] = {}
        self._lock = threading.RLock()
        self._load_existing()
        self.cleanup()

    def _path(self, approval_id: str) -> Path:
        return self.base_dir / f"{approval_id}.json"

    def _save(self, approval: dict[str, Any]) -> None:
        self._path(approval["approval_id"]).write_text(
            json.dumps(approval, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

    def _load_existing(self) -> None:
        now = time.time()
        for path in sorted(self.base_dir.glob("*.json")):
            try:
                approval = json.loads(path.read_text(encoding="utf-8"))
            except (OSError, json.JSONDecodeError):
                continue
            updated_at = float(approval.get("updated_at") or approval.get("created_at") or now)
            if now - updated_at > self.retention_seconds:
                try:
                    path.unlink()
                except OSError:
                    pass
                continue
            if approval.get("status") == "pending":
                approval["status"] = "rejected"
                approval["reason"] = "Server restarted while waiting for approval"
                approval["updated_at"] = time.time()
                path.write_text(
                    json.dumps(approval, ensure_ascii=False, indent=2),
                    encoding="utf-8",
                )
            self._approvals[approval["approval_id"]] = approval

    def create(
        self,
        *,
        job_id: str,
        tool_name: str,
        arguments: dict[str, Any],
    ) -> dict[str, Any]:
        approval_id = uuid.uuid4().hex
        approval = {
            "approval_id": approval_id,
            "job_id": job_id,
            "tool_name": tool_name,
            "arguments": arguments,
            "status": "pending",
            "created_at": time.time(),
            "updated_at": time.time(),
            "actor": None,
            "reason": None,
        }
        with self._lock:
            self._approvals[approval_id] = approval
            self._conditions[approval_id] = threading.Condition(self._lock)
            self._save(approval)
        return approval.copy()

    def get(self, approval_id: str) -> dict[str, Any] | None:
        with self._lock:
            approval = self._approvals.get(approval_id)
            return approval.copy() if approval else None

    def list_pending(self) -> list[dict[str, Any]]:
        with self._lock:
            pending = [
                approval.copy()
                for approval in self._approvals.values()
                if approval.get("status") == "pending"
            ]
        pending.sort(key=lambda item: item.get("created_at", 0))
        return pending

    def respond(
        self,
        approval_id: str,
        *,
        approved: bool,
        actor: str,
    ) -> dict[str, Any]:
        with self._lock:
            approval = self._approvals.get(approval_id)
            if not approval:
                raise KeyError("Unknown approval_id")
            if approval["status"] != "pending":
                return approval.copy()
            approval["status"] = "approved" if approved else "rejected"
            approval["actor"] = actor
            approval["updated_at"] = time.time()
            approval["reason"] = None if approved else "Rejected by operator"
            self._save(approval)
            condition = self._conditions.get(approval_id)
            if condition:
                condition.notify_all()
            return approval.copy()

    def wait(self, approval_id: str, timeout_seconds: float = 3600.0) -> dict[str, Any]:
        with self._lock:
            approval = self._approvals.get(approval_id)
            if not approval:
                raise KeyError("Unknown approval_id")
            if approval["status"] != "pending":
                return approval.copy()
            condition = self._conditions.setdefault(
                approval_id,
                threading.Condition(self._lock),
            )
            deadline = time.time() + timeout_seconds
            while approval["status"] == "pending":
                remaining = deadline - time.time()
                if remaining <= 0:
                    approval["status"] = "rejected"
                    approval["reason"] = "Approval timeout"
                    approval["updated_at"] = time.time()
                    self._save(approval)
                    break
                condition.wait(timeout=remaining)
            return approval.copy()

    def reject_pending_for_job(
        self,
        job_id: str,
        *,
        actor: str,
        reason: str,
    ) -> list[dict[str, Any]]:
        rejected: list[dict[str, Any]] = []
        with self._lock:
            for approval_id, approval in self._approvals.items():
                if approval.get("job_id") != job_id or approval.get("status") != "pending":
                    continue
                approval["status"] = "rejected"
                approval["actor"] = actor
                approval["reason"] = reason
                approval["updated_at"] = time.time()
                self._save(approval)
                condition = self._conditions.get(approval_id)
                if condition:
                    condition.notify_all()
                rejected.append(approval.copy())
        return rejected

    def cleanup(self) -> None:
        now = time.time()
        with self._lock:
            expired = [
                approval_id
                for approval_id, approval in self._approvals.items()
                if now
                - float(approval.get("updated_at") or approval.get("created_at") or now)
                > self.retention_seconds
            ]
            for approval_id in expired:
                self._approvals.pop(approval_id, None)
                self._conditions.pop(approval_id, None)
                try:
                    self._path(approval_id).unlink()
                except OSError:
                    pass