ForcePilot/backend/package/yuxi/channels/adapters/signal/exec_auth.py

from __future__ import annotations

import asyncio
import logging
from enum import StrEnum
from typing import Any
from collections.abc import Awaitable, Callable

logger = logging.getLogger(__name__)

EXEC_AUTH_TIMEOUT = 30.0


class ExecAuthResult(StrEnum):
    APPROVED = "approved"
    DENIED = "denied"
    PENDING = "pending"


class ExecAuthAdapter:
    def __init__(self, auto_approve: bool = False):
        self._auto_approve = auto_approve
        self._pending_requests: dict[str, dict] = {}
        self._request_events: dict[str, asyncio.Event] = {}
        self._request_counter = 0
        self._on_approval_request: Callable[[dict], Awaitable[None]] | None = None

    @property
    def auto_approve(self) -> bool:
        return self._auto_approve

    def on_approval_request(self, handler: Callable[[dict], Awaitable[None]]) -> None:
        self._on_approval_request = handler

    async def request_approval(
        self, action: str, params: dict[str, Any], user_id: str, timeout: float = EXEC_AUTH_TIMEOUT
    ) -> ExecAuthResult:
        if self._auto_approve:
            return ExecAuthResult.APPROVED

        self._request_counter += 1
        request_id = f"exec_auth_{self._request_counter}"
        event = asyncio.Event()
        request = {
            "id": request_id,
            "action": action,
            "params": params,
            "user_id": user_id,
            "result": ExecAuthResult.PENDING,
        }
        self._pending_requests[request_id] = request
        self._request_events[request_id] = event

        if self._on_approval_request:
            await self._on_approval_request(request)

        try:
            await asyncio.wait_for(event.wait(), timeout=timeout)
        except TimeoutError:
            logger.warning(f"Exec auth request {request_id} timed out after {timeout}s")
            request["result"] = ExecAuthResult.DENIED
            self._cleanup_request(request_id)
            return ExecAuthResult.DENIED

        result = request.get("result", ExecAuthResult.DENIED)
        self._cleanup_request(request_id)
        return ExecAuthResult(result)

    def approve(self, request_id: str) -> ExecAuthResult:
        request = self._pending_requests.get(request_id)
        if not request:
            return ExecAuthResult.DENIED
        request["result"] = ExecAuthResult.APPROVED
        event = self._request_events.get(request_id)
        if event:
            event.set()
        return ExecAuthResult.APPROVED

    def deny(self, request_id: str) -> ExecAuthResult:
        request = self._pending_requests.get(request_id)
        if not request:
            return ExecAuthResult.DENIED
        request["result"] = ExecAuthResult.DENIED
        event = self._request_events.get(request_id)
        if event:
            event.set()
        return ExecAuthResult.DENIED

    def get_pending(self) -> list[dict]:
        return [r for r in self._pending_requests.values() if r["result"] == ExecAuthResult.PENDING]

    def _cleanup_request(self, request_id: str) -> None:
        self._pending_requests.pop(request_id, None)
        self._request_events.pop(request_id, None)