# retoor from __future__ import annotations import json import logging import time import uuid_utils from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request from starlette.responses import Response from devplacepy_services.base.auth import validate_internal_key from devplacepy_services.base.errors import error_response, sanitize_detail from devplacepy_services.base.http import INTERNAL_CALLS, REQUEST_ID, current_internal_calls from devplacepy_services.base.stats import COUNTERS logger = logging.getLogger(__name__) PUBLIC_PATHS = {"/health"} class RequestStatsMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): COUNTERS.record_request() response = await call_next(request) if response.status_code >= 500: COUNTERS.record_error() return response class InternalCallCounterMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): token_calls = INTERNAL_CALLS.set(0) response = await call_next(request) service_name = getattr(request.app.state, "service_name", "") if service_name in {"web", "gateway"}: response.headers["X-Internal-Calls"] = str(current_internal_calls()) INTERNAL_CALLS.reset(token_calls) return response class InternalAuthMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): path = request.url.path if path in PUBLIC_PATHS: return await call_next(request) if not validate_internal_key(request.headers.get("X-Internal-Key")): return error_response(401, "Unauthorized", "unauthorized") return await call_next(request) class SanitizeErrorsMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): request_id = request.headers.get("X-Request-Id") or uuid_utils.uuid7().hex token_id = REQUEST_ID.set(request_id) started = time.perf_counter() try: response = await call_next(request) except Exception: COUNTERS.record_error() logger.exception( "unhandled service error", extra={ "request_id": request_id, "service": getattr(request.app.state, "service_name", ""), "path": request.url.path, }, ) response = error_response(500, "Internal server error", "internal_error") finally: REQUEST_ID.reset(token_id) duration_ms = int((time.perf_counter() - started) * 1000) service_name = getattr(request.app.state, "service_name", "") logger.info( "request", extra={ "request_id": request_id, "service": service_name, "path": request.url.path, "status": response.status_code, "duration_ms": duration_ms, "internal_calls": current_internal_calls(), }, ) if response.status_code < 400: return response if not hasattr(response, "body_iterator"): return response content_type = response.headers.get("content-type", "") if "application/json" not in content_type: if response.status_code == 404: return error_response(404, "Not found", "not_found") if response.status_code == 401: return error_response(401, "Unauthorized", "unauthorized") if response.status_code >= 500: return error_response(500, "Internal server error", "internal_error") return error_response(response.status_code, "Request failed", "error") body = b"" async for chunk in response.body_iterator: body += chunk try: parsed = json.loads(body) except json.JSONDecodeError: return error_response(response.status_code, "Request failed", "error") if isinstance(parsed, dict) and "error" in parsed and "code" in parsed: return Response( content=body, status_code=response.status_code, media_type="application/json", ) if isinstance(parsed, dict) and isinstance(parsed.get("error"), dict) and "message" in parsed["error"]: return Response( content=body, status_code=response.status_code, media_type="application/json", ) if isinstance(parsed, dict) and "detail" in parsed: message, code = sanitize_detail(parsed["detail"], response.status_code) return error_response(response.status_code, message, code) return error_response(response.status_code, "Request failed", "error")