instrument pre-completion chat availability

Author: randy-concepcion

src/mlpa/core/auth/authorize.py

@@ -13,6 +13,8 @@
     ServiceType,
 )
 from mlpa.core.config import env
+from mlpa.core.metrics import record_chat_availability_for
+from mlpa.core.prometheus_metrics import AvailabilityReason
 from mlpa.core.routers.appattest import app_attest_auth
 from mlpa.core.utils import extract_user_from_play_integrity_jwt, parse_app_attest_jwt
 
@@ -123,27 +125,57 @@ async def authorize_chat_request(
     )
 
     if not is_service_type_valid:
+        record_chat_availability_for(
+            AvailabilityReason.INVALID_SERVICE_TYPE_FOR_MODEL,
+            model=chat_request.model,
+            service_type=service_type.value,
+            purpose="",
+        )
         raise HTTPException(
             status_code=400,
             detail=f"Invalid service-type value for model {chat_request.model}. Should be one of {env.forced_model_service_type_pairs.get(chat_request.model)}",
         )
 
-    return await _authorize_common_request(
-        request=request,
-        build_authorized_request=lambda user, purpose_value: AuthorizedChatRequest(
-            user=user,
+    try:
+        return await _authorize_common_request(
+            request=request,
+            build_authorized_request=lambda user, purpose_value: AuthorizedChatRequest(
+                user=user,
+                service_type=service_type.value,
+                purpose=purpose_value,
+                **chat_request.model_dump(exclude_unset=True),
+            ),
+            authorization=authorization,
+            service_type=service_type,
+            purpose=purpose,
+            x_dev_authorization=x_dev_authorization,
+            use_app_attest=use_app_attest,
+            use_qa_certificates=use_qa_certificates,
+            use_play_integrity=use_play_integrity,
+        )
+    except HTTPException as exc:
+        # Record only the terminal auth dispositions the dependency intentionally
+        # emits. 401 is an expected or normalized auth rejection. 400 is a
+        # client-error request to the auth layer (invalid purpose, or malformed App
+        # Attest base64 decoded in app_attest_auth). Any other status, including App
+        # Attest's explicit 500, is re-raised unrecorded; auth-system-failure capture
+        # is left to a follow-on auth backend change. Non-HTTPException errors are not
+        # caught here and propagate unrecorded. Purpose is unresolved at these
+        # dispositions, so a stable "" placeholder is used, never the unvalidated
+        # header value.
+        if exc.status_code == 401:
+            reason = AvailabilityReason.AUTH_REJECTED
+        elif exc.status_code == 400:
+            reason = AvailabilityReason.INVALID_AUTH_REQUEST
+        else:
+            raise
+        record_chat_availability_for(
+            reason,
+            model=chat_request.model,
             service_type=service_type.value,
-            purpose=purpose_value,
-            **chat_request.model_dump(exclude_unset=True),
-        ),
-        authorization=authorization,
-        service_type=service_type,
-        purpose=purpose,
-        x_dev_authorization=x_dev_authorization,
-        use_app_attest=use_app_attest,
-        use_qa_certificates=use_qa_certificates,
-        use_play_integrity=use_play_integrity,
-    )
+            purpose="",
+        )
+        raise
 
 
 async def authorize_search_request(

src/mlpa/core/completions.py

@@ -54,20 +54,29 @@ def _build_litellm_body(req: AuthorizedChatRequest, *, stream: bool) -> dict:
 async def get_or_create_user_for_completion(
     user_id: str, req: AuthorizedChatRequest | AuthorizedSearchRequest
 ):
-    """Wraps get_or_create_user and records a signup-cap rejection metric if applicable."""
+    """Wraps get_or_create_user, recording the signup-cap rejection metric and the
+    pre-completion availability disposition (signup cap, or a 5xx provisioning
+    failure) for chat requests."""
     try:
         return await get_or_create_user(user_id)
     except HTTPException as exc:
-        if (
-            exc.status_code == 403
-            and isinstance(exc.detail, dict)
-            and exc.detail.get("error") == ERROR_CODE_MAX_USERS_REACHED
-            and isinstance(req, AuthorizedChatRequest)
-        ):
-            record_chat_request_rejection(
-                req,
-                PrometheusRejectionReason.SIGNUP_CAP_EXCEEDED,
-            )
+        if isinstance(req, AuthorizedChatRequest):
+            if (
+                exc.status_code == 403
+                and isinstance(exc.detail, dict)
+                and exc.detail.get("error") == ERROR_CODE_MAX_USERS_REACHED
+            ):
+                record_chat_request_rejection(
+                    req,
+                    PrometheusRejectionReason.SIGNUP_CAP_EXCEEDED,
+                )
+                record_chat_availability(req, AvailabilityReason.SIGNUP_CAP_EXCEEDED)
+            elif exc.status_code >= 500:
+                # User-resolution / provisioning system failure for an eligible
+                # request. Non-signup-cap, non-5xx dispositions are left unrecorded;
+                # a future client-side 4xx should get its own classification rather
+                # than counting as an availability failure.
+                record_chat_availability(req, AvailabilityReason.PROVISIONING_FAILURE)
         raise
 
 

src/mlpa/core/metrics.py

@@ -25,16 +25,33 @@ def record_chat_request_rejection(
     metrics.chat_request_rejections.labels(reason=reason, **_chat_labels(req)).inc()
 
 
-def record_chat_availability(
-    req: AuthorizedChatRequest, reason: AvailabilityReason
+def record_chat_availability_for(
+    reason: AvailabilityReason,
+    *,
+    model: str,
+    service_type: str,
+    purpose: str,
 ) -> None:
     metrics.chat_availability.labels(
         outcome=availability_outcome_for(reason),
         reason=reason,
-        **_chat_labels(req),
+        model=model,
+        service_type=service_type,
+        purpose=purpose,
     ).inc()
 
 
+def record_chat_availability(
+    req: AuthorizedChatRequest, reason: AvailabilityReason
+) -> None:
+    record_chat_availability_for(
+        reason,
+        model=req.model,
+        service_type=req.service_type,
+        purpose=req.purpose,
+    )
+
+
 def record_completion_latency(
     req: AuthorizedChatRequest,
     result: PrometheusResult,

src/mlpa/core/prometheus_metrics.py

@@ -44,9 +44,19 @@ class AvailabilityReason(StrEnum):
     INVALID_REQUEST = "invalid_request"
     # abort
     CLIENT_DISCONNECT = "client_disconnect"
-    # TODO: add pre-completion reasons (auth_rejected, auth_system_failure,
-    # signup_cap_exceeded, blocked, provisioning_failure, invalid_purpose,
-    # invalid_service_type_for_model) when those paths are instrumented.
+    # pre-completion (auth dependency and route body, before the completion path)
+    # excluded
+    AUTH_REJECTED = "auth_rejected"
+    INVALID_AUTH_REQUEST = "invalid_auth_request"
+    INVALID_SERVICE_TYPE_FOR_MODEL = "invalid_service_type_for_model"
+    SIGNUP_CAP_EXCEEDED = "signup_cap_exceeded"
+    BLOCKED = "blocked"
+    # failure
+    PROVISIONING_FAILURE = "provisioning_failure"
+    # Defined but not emitted yet: auth backends normalize system failures to 401
+    # (indistinguishable from expected rejections), so capturing this is left to a
+    # follow-on auth backend change that surfaces a real disposition.
+    AUTH_SYSTEM_FAILURE = "auth_system_failure"
 
 
 _AVAILABILITY_OUTCOME_BY_REASON: dict[AvailabilityReason, AvailabilityOutcome] = {
@@ -60,6 +70,13 @@ class AvailabilityReason(StrEnum):
     AvailabilityReason.INVALID_MODEL_NAME: AvailabilityOutcome.EXCLUDED,
     AvailabilityReason.INVALID_REQUEST: AvailabilityOutcome.EXCLUDED,
     AvailabilityReason.CLIENT_DISCONNECT: AvailabilityOutcome.ABORT,
+    AvailabilityReason.AUTH_REJECTED: AvailabilityOutcome.EXCLUDED,
+    AvailabilityReason.INVALID_AUTH_REQUEST: AvailabilityOutcome.EXCLUDED,
+    AvailabilityReason.INVALID_SERVICE_TYPE_FOR_MODEL: AvailabilityOutcome.EXCLUDED,
+    AvailabilityReason.SIGNUP_CAP_EXCEEDED: AvailabilityOutcome.EXCLUDED,
+    AvailabilityReason.BLOCKED: AvailabilityOutcome.EXCLUDED,
+    AvailabilityReason.PROVISIONING_FAILURE: AvailabilityOutcome.FAILURE,
+    AvailabilityReason.AUTH_SYSTEM_FAILURE: AvailabilityOutcome.FAILURE,
 }
 
 

src/mlpa/run.py

@@ -26,9 +26,11 @@
 )
 from mlpa.core.http_client import close_http_client, get_http_client
 from mlpa.core.logger import logger, setup_logger
+from mlpa.core.metrics import record_chat_availability
 from mlpa.core.middleware import register_middleware
 from mlpa.core.openapi import customize_openapi
 from mlpa.core.pg_services.services import app_attest_pg, litellm_pg
+from mlpa.core.prometheus_metrics import AvailabilityReason
 from mlpa.core.routers.appattest import appattest_router
 from mlpa.core.routers.health import health_router
 from mlpa.core.routers.mock import mock_router
@@ -207,6 +209,7 @@ async def chat_completion(
         )
     user, _ = await get_or_create_user_for_completion(user_id, authorized_chat_request)
     if user.get("blocked"):
+        record_chat_availability(authorized_chat_request, AvailabilityReason.BLOCKED)
         raise HTTPException(status_code=403, detail={"error": "User is blocked."})
 
     if authorized_chat_request.stream: