AareDAQ/src/aare/daq/server_exception_handler.py

# aare/common/server_error_handler.py
"""Server-side exception handlers for the DAQ FastAPI app.

Per the exception-handling redesign (EXCEPTION_REDESIGN_PLAN.md §2), error
handling is collapsed into a small number of handlers that route on the new
``AareException`` hierarchy:

- ``AutomationError``: 503 if ``critical``, else 422
- ``AareUserError``:    400
- ``AareAuthError``:    401 (auth) or 403 (rights)
- ``HTTPException``:    pass-through status, with the new body shape
- ``Exception``:        500, marked critical (fail-safe for unknown errors)

All handlers produce the unified response body defined in §3:

.. code-block:: json

   {
     "critical": true,
     "code": "TELL_COMMUNICATION_ERROR",
     "exception_class": "TellCommunicationError",
     "message": "Tell did not respond within 5s",
     "context": {"endpoint": "/state", "operation": "GET"}
   }
"""
from __future__ import annotations

import json
from typing import Any

from fastapi import HTTPException
from fastapi import status as api_status
from starlette.requests import Request
from starlette.responses import JSONResponse

from aare.common.logger_config import setup_logger
from aare.common.error_codes import AareErrorCode, code_for_exception_class
from aare.common.exception_handler import (
    AareException,
    AutomationError,
    AareUserError,
    AareAuthError,
    AuthenticationException,
)

logger = setup_logger("aareDAQ")


# Attributes that live on exceptions for routing/transport, NOT for the
# context payload.
_NON_CONTEXT_ATTRS = frozenset({
    "args",
    "message",
    "critical",
    "code",
    "headers",
    "status_code",
    "exception",  # BECCommunicationError stores a Python exception here -- not JSON-safe
})


def _extract_context(exc: Exception) -> dict[str, Any]:
    """Pull JSON-serializable structured fields off an exception for the
    response ``context`` field.

    Allow-list approach via ``_NON_CONTEXT_ATTRS`` exclusions: walks
    ``vars(exc)`` and includes anything that's not private, not in the
    excluded set, and JSON-serializable. Per plan §3 the schema is
    per-exception; clients tolerate unknown keys."""
    ctx: dict[str, Any] = {}
    for k, v in vars(exc).items():
        if k.startswith("_"):
            continue
        if k in _NON_CONTEXT_ATTRS:
            continue
        if v is None:
            continue
        try:
            json.dumps(v)
        except (TypeError, ValueError):
            continue
        ctx[k] = v
    return ctx


def _error_body(exc: Exception, *, code_override: str | None = None,
                critical_override: bool | None = None) -> dict[str, Any]:
    """Build the unified error response body for any exception.

    ``code_override`` is used by the auth handler to surface the finer-grained
    ``AuthErrorCode`` carried on the instance, instead of the class-derived
    code. ``critical_override`` is used by the bare-Exception fallback to
    force critical=True regardless of the exception's own flag."""
    if code_override is not None:
        code = code_override
    else:
        code = code_for_exception_class(type(exc).__name__)

    if critical_override is not None:
        critical = critical_override
    else:
        # Default for non-AareException is critical=True (fail-safe).
        critical = bool(getattr(exc, "critical", True))

    return {
        "critical": critical,
        "code": code,
        "exception_class": type(exc).__name__,
        "message": str(exc) or type(exc).__name__,
        "context": _extract_context(exc),
    }


def register_exception_handlers(app) -> None:
    """
    Register the four routing-root handlers (plus HTTPException + Exception
    fallback) on the given FastAPI app. Call once right after
    ``app = FastAPI()``.
    """

    @app.exception_handler(AutomationError)
    async def _handle_automation_error(request: Request, exc: AutomationError) -> JSONResponse:
        status = (
            api_status.HTTP_503_SERVICE_UNAVAILABLE
            if exc.critical
            else api_status.HTTP_422_UNPROCESSABLE_ENTITY
        )
        return JSONResponse(status_code=status, content=_error_body(exc))

    @app.exception_handler(AareUserError)
    async def _handle_user_error(request: Request, exc: AareUserError) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_400_BAD_REQUEST,
            content=_error_body(exc),
        )

    @app.exception_handler(AareAuthError)
    async def _handle_auth_error(request: Request, exc: AareAuthError) -> JSONResponse:
        # Auth flow needs the finer-grained AuthErrorCode (e.g. INVALID_TOKEN
        # vs SESSION_ALREADY_ACTIVE) so the GUI can route between re-auth /
        # claim-session / login dialogs. Fall back to the class-derived code
        # when the instance carries no `code` attribute.
        instance_code = getattr(exc, "code", None)
        code_override = str(instance_code) if instance_code is not None else None
        status = getattr(exc, "status_code", None)
        if status is None:
            status = (
                api_status.HTTP_401_UNAUTHORIZED
                if isinstance(exc, AuthenticationException)
                else api_status.HTTP_403_FORBIDDEN
            )
        return JSONResponse(
            status_code=status,
            content=_error_body(exc, code_override=code_override),
            headers=getattr(exc, "headers", None),
        )

    @app.exception_handler(HTTPException)
    async def _handle_http_exception(request: Request, exc: HTTPException) -> JSONResponse:
        """HTTPException isn't in the AareException hierarchy but FastAPI raises
        it internally (e.g. validation, route-not-found) and routes raise it
        manually. Adapt to the new body shape so clients see one schema."""
        detail = exc.detail
        # Preserve legacy callers that passed a {code, message} dict
        if isinstance(detail, dict) and "code" in detail and "message" in detail:
            legacy_code = str(detail.get("code"))
            legacy_message = str(detail.get("message"))
            return JSONResponse(
                status_code=exc.status_code,
                content={
                    "critical": exc.status_code >= 500,
                    "code": legacy_code,
                    "exception_class": "HTTPException",
                    "message": legacy_message,
                    "context": {k: v for k, v in detail.items() if k not in ("code", "message")},
                },
                headers=exc.headers,
            )
        return JSONResponse(
            status_code=exc.status_code,
            content={
                "critical": exc.status_code >= 500,
                "code": "HTTP_ERROR",
                "exception_class": "HTTPException",
                "message": str(detail) or "HTTP error",
                "context": {},
            },
            headers=exc.headers,
        )

    @app.exception_handler(Exception)
    async def _handle_unknown(request: Request, exc: Exception) -> JSONResponse:
        """Fail-safe: any unrecognized exception is treated as critical."""
        logger.exception("Unhandled exception in DAQ", exc_info=exc)
        return JSONResponse(
            status_code=api_status.HTTP_500_INTERNAL_SERVER_ERROR,
            content=_error_body(
                exc,
                code_override=str(AareErrorCode.INTERNAL_ERROR),
                critical_override=True,
            ),
        )