Magent/backend/app/services/operation_progress.py

from __future__ import annotations

from contextvars import ContextVar, Token
from copy import deepcopy
from datetime import datetime, timezone
import re
import threading
import time
import uuid
from typing import Any, Dict, Optional


_OPERATION_ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]{8,80}$")
_OPERATION_TTL_SECONDS = 15 * 60
_MAX_OPERATIONS = 500
_MAX_EVENTS = 60
_current_operation_id: ContextVar[Optional[str]] = ContextVar(
    "magent_operation_id", default=None
)
_operations: Dict[str, Dict[str, Any]] = {}
_lock = threading.Lock()


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def normalize_operation_id(value: Optional[str]) -> Optional[str]:
    if not isinstance(value, str):
        return None
    normalized = value.strip()
    return normalized if _OPERATION_ID_PATTERN.fullmatch(normalized) else None


def _prune_locked(now_monotonic: float) -> None:
    expired = [
        operation_id
        for operation_id, operation in _operations.items()
        if now_monotonic - float(operation.get("updated_monotonic") or 0) > _OPERATION_TTL_SECONDS
    ]
    for operation_id in expired:
        _operations.pop(operation_id, None)
    if len(_operations) <= _MAX_OPERATIONS:
        return
    oldest = sorted(
        _operations,
        key=lambda operation_id: float(_operations[operation_id].get("updated_monotonic") or 0),
    )
    for operation_id in oldest[: len(_operations) - _MAX_OPERATIONS]:
        _operations.pop(operation_id, None)


def begin_operation(operation_id: str, *, label: Optional[str], path: str) -> Token:
    now_monotonic = time.monotonic()
    now_iso = _now_iso()
    normalized_label = str(label or "Requested action").strip()[:120] or "Requested action"
    with _lock:
        _prune_locked(now_monotonic)
        _operations[operation_id] = {
            "id": operation_id,
            "label": normalized_label,
            "path": path,
            "status": "running",
            "started_at": now_iso,
            "updated_at": now_iso,
            "updated_monotonic": now_monotonic,
            "duration_ms": None,
            "events": [
                {
                    "id": uuid.uuid4().hex,
                    "service": "Magent",
                    "state": "complete",
                    "message": "Magent received the action.",
                    "started_at": now_iso,
                    "finished_at": now_iso,
                    "duration_ms": 0,
                    "status_code": None,
                }
            ],
        }
    return _current_operation_id.set(operation_id)


def reset_operation(token: Token) -> None:
    _current_operation_id.reset(token)


def start_remote_call(service: str, message: Optional[str] = None) -> Optional[str]:
    operation_id = _current_operation_id.get()
    if not operation_id:
        return None
    event_id = uuid.uuid4().hex
    now_iso = _now_iso()
    now_monotonic = time.monotonic()
    with _lock:
        operation = _operations.get(operation_id)
        if not operation:
            return None
        operation["events"].append(
            {
                "id": event_id,
                "service": service,
                "state": "active",
                "message": message or f"Contacting {service}…",
                "started_at": now_iso,
                "finished_at": None,
                "duration_ms": None,
                "status_code": None,
                "started_monotonic": now_monotonic,
            }
        )
        operation["events"] = operation["events"][-_MAX_EVENTS:]
        operation["updated_at"] = now_iso
        operation["updated_monotonic"] = now_monotonic
    return event_id


def finish_remote_call(
    event_id: Optional[str],
    *,
    success: bool,
    status_code: Optional[int] = None,
    message: Optional[str] = None,
) -> None:
    operation_id = _current_operation_id.get()
    if not operation_id or not event_id:
        return
    now_iso = _now_iso()
    now_monotonic = time.monotonic()
    with _lock:
        operation = _operations.get(operation_id)
        if not operation:
            return
        event = next(
            (candidate for candidate in operation["events"] if candidate.get("id") == event_id),
            None,
        )
        if not event:
            return
        started_monotonic = float(event.pop("started_monotonic", now_monotonic))
        event["state"] = "complete" if success else "error"
        event["finished_at"] = now_iso
        event["duration_ms"] = round((now_monotonic - started_monotonic) * 1000, 1)
        event["status_code"] = status_code
        event["message"] = message or (
            f"{event['service']} responded successfully."
            if success
            else f"{event['service']} returned an error."
        )
        operation["updated_at"] = now_iso
        operation["updated_monotonic"] = now_monotonic


def finish_operation(operation_id: str, *, success: bool, status_code: Optional[int]) -> None:
    now_iso = _now_iso()
    now_monotonic = time.monotonic()
    with _lock:
        operation = _operations.get(operation_id)
        if not operation:
            return
        for event in operation["events"]:
            if event.get("state") == "active":
                started_monotonic = float(event.pop("started_monotonic", now_monotonic))
                event["state"] = "error"
                event["finished_at"] = now_iso
                event["duration_ms"] = round((now_monotonic - started_monotonic) * 1000, 1)
                event["message"] = f"{event.get('service') or 'Remote service'} did not complete."
        started = datetime.fromisoformat(str(operation["started_at"]))
        duration_ms = (datetime.now(timezone.utc) - started).total_seconds() * 1000
        operation["status"] = "complete" if success else "error"
        operation["status_code"] = status_code
        operation["duration_ms"] = round(duration_ms, 1)
        operation["updated_at"] = now_iso
        operation["updated_monotonic"] = now_monotonic
        operation["events"].append(
            {
                "id": uuid.uuid4().hex,
                "service": "Magent",
                "state": "complete" if success else "error",
                "message": (
                    "Magent finished processing the action."
                    if success
                    else "Magent could not complete the action."
                ),
                "started_at": now_iso,
                "finished_at": now_iso,
                "duration_ms": 0,
                "status_code": status_code,
            }
        )
        operation["events"] = operation["events"][-_MAX_EVENTS:]


def get_operation(operation_id: str) -> Optional[Dict[str, Any]]:
    normalized = normalize_operation_id(operation_id)
    if not normalized:
        return None
    with _lock:
        operation = _operations.get(normalized)
        if not operation:
            return None
        result = deepcopy(operation)
    result.pop("updated_monotonic", None)
    for event in result.get("events", []):
        event.pop("started_monotonic", None)
    return result