feat(B.13-B.17+B.16): Error-Handling-Infra, Observability, Graceful Shutdown, Cost-Cap, API Versioning

B.13 Error-Handling-Infrastruktur: - ErrorCategory Enum (TRANSIENT/PERMANENT/PARTIAL), ApiError erweitert - Einheitliches Error-Response-Format: {code, detail, field, trace_id, retryable, category} - 3 FastAPI Exception-Handler (ApiError, HTTPException, unhandled) - classify_exception() Helper, 6 neue Error-Codes - 28 Tests in test_error_handling.py B.14 Observability & trace_id-Korrelation: - trace_id pro Request (UUID4 short) in structlog contextvars - X-Trace-Id Response-Header - Sensitive Fields structlog processor - llm_complete()/llm_embed() akzeptieren trace_id kwarg - 12 Tests in test_observability.py B.15 Graceful Shutdown & Connection Draining: - _shutdown_event + _inflight_requests Tracking in main.py - drain_all_connections() in ws_helpers.py - Worker on_shutdown pausiert WorkflowInstances (status=paused) - 8 Tests in test_graceful_shutdown.py B.16 API Versioning Strategie: - Plugin-Dev-Guide Kapitel 30: URL-basiertes Versioning, Breaking Change Prozess B.17 Cost Overrun Protection: - llm_monthly_budget_usd + llm_hard_cutoff Settings - _check_tenant_budget() vor jedem LLM-Call - _track_tenant_cost() in Redis (INCRBYFLOAT) - _check_cost_alerts() bei 50%/80%/100% -> post_system_message() - 20 Tests in test_cost_protection.py Total: 68 neue Tests, alle grün. Keine Regressionen.
2026-08-13 21:33:14 +02:00
parent 1baa9481a2
commit b5546ea7bd
13 changed files with 1461 additions and 31 deletions
@@ -21,6 +21,7 @@ import json
 import logging
 import os
 import uuid
+from datetime import datetime, timezone
 from typing import Any, TYPE_CHECKING

 import litellm
@@ -94,6 +95,175 @@ _PERMANENT_KEYWORDS = frozenset(
 # ──────────────────────────────────────────────────────────────────────────


+# ── Cost Overrun Protection (B.17) ───────────────────────────────────────────
+
+
+def _get_cost_key(tenant_id: uuid.UUID | str) -> str:
+    """Build the Redis cost-tracking key for the current month."""
+    now = datetime.now(timezone.utc)
+    month_str = now.strftime("%Y-%m")
+    return f"cost:tenant:{tenant_id}:month:{month_str}"
+
+
+async def get_tenant_monthly_cost(tenant_id: uuid.UUID | str) -> float:
+    """Get the accumulated LLM cost for a tenant in the current month.
+
+    Reads from Redis key ``cost:tenant:{tenant_id}:month:{YYYY-MM}``.
+    Returns 0.0 if Redis is unavailable or the key doesn't exist.
+    """
+    try:
+        from app.core.auth import get_redis
+        r = get_redis()
+        key = _get_cost_key(tenant_id)
+        val = await r.get(key)
+        return float(val) if val else 0.0
+    except Exception:
+        logger.debug("Failed to get tenant monthly cost from Redis")
+        return 0.0
+
+
+async def _check_tenant_budget(
+    tenant_id: uuid.UUID | str | None,
+    db: AsyncSession | None = None,
+) -> None:
+    """Check if tenant has exceeded their LLM budget.
+
+    Raises ``ValueError("Tenant LLM budget exceeded")`` if the tenant's
+    accumulated monthly cost exceeds the configured budget and
+    ``llm_hard_cutoff`` is enabled.
+    """
+    if tenant_id is None:
+        return
+
+    from app.config import get_settings
+    settings = get_settings()
+
+    if settings.llm_monthly_budget_usd <= 0:
+        return  # No budget limit configured
+
+    current_cost = await get_tenant_monthly_cost(tenant_id)
+    if current_cost >= settings.llm_monthly_budget_usd:
+        if settings.llm_hard_cutoff:
+            logger.warning(
+                "Tenant %s LLM budget exceeded: $%.2f >= $%.2f (hard cutoff)",
+                tenant_id, current_cost, settings.llm_monthly_budget_usd,
+            )
+            raise ValueError("Tenant LLM budget exceeded")
+        else:
+            logger.warning(
+                "Tenant %s LLM budget exceeded: $%.2f >= $%.2f (soft limit, no cutoff)",
+                tenant_id, current_cost, settings.llm_monthly_budget_usd,
+            )
+
+
+async def _track_tenant_cost(
+    tenant_id: uuid.UUID | str | None,
+    cost_usd: float,
+    db: AsyncSession | None = None,
+) -> float:
+    """Track LLM cost in Redis and check for alert thresholds.
+
+    Increments ``cost:tenant:{tenant_id}:month:{YYYY-MM}`` by ``cost_usd``.
+    Returns the new total cost for the month.
+    """
+    if tenant_id is None or cost_usd <= 0:
+        return 0.0
+
+    try:
+        from app.core.auth import get_redis
+        r = get_redis()
+        key = _get_cost_key(tenant_id)
+        new_total = await r.incrbyfloat(key, cost_usd)
+        # Set TTL to 35 days so old months auto-expire
+        await r.expire(key, 35 * 24 * 3600)
+
+        # Check cost alert thresholds
+        await _check_cost_alerts(tenant_id, new_total, db)
+
+        return float(new_total)
+    except Exception:
+        logger.debug("Failed to track tenant cost in Redis")
+        return 0.0
+
+
+async def _check_cost_alerts(
+    tenant_id: uuid.UUID | str,
+    current_cost: float,
+    db: AsyncSession | None = None,
+) -> None:
+    """Send cost alerts at 50%, 80%, and 100% of the tenant budget.
+
+    Each threshold alert is sent only once per month (tracked via Redis flag
+    ``cost:alerted:{tenant_id}:{threshold}``).
+    """
+    from app.config import get_settings
+    settings = get_settings()
+
+    budget = settings.llm_monthly_budget_usd
+    if budget <= 0:
+        return
+
+    thresholds = [0.50, 0.80, 1.00]
+    now = datetime.now(timezone.utc)
+    month_str = now.strftime("%Y-%m")
+
+    try:
+        from app.core.auth import get_redis
+        r = get_redis()
+
+        for threshold in thresholds:
+            threshold_cost = budget * threshold
+            if current_cost >= threshold_cost:
+                alert_key = f"cost:alerted:{tenant_id}:{threshold}:{month_str}"
+                already_alerted = await r.set(alert_key, "1", nx=True, ex=35 * 24 * 3600)
+                if already_alerted:
+                    # This is a new alert — send notification
+                    pct = int(threshold * 100)
+                    logger.info(
+                        "Cost alert: tenant %s reached %d%% of budget ($%.2f / $%.2f)",
+                        tenant_id, pct, current_cost, budget,
+                    )
+                    # Best-effort system notification
+                    if db is not None:
+                        try:
+                            from app.core.notifications import post_system_message
+                            # Need a user_id — try to find an admin for this tenant
+                            from sqlalchemy import select as sa_select
+                            from app.models.user import User, UserTenant
+                            from app.models.role import Role
+
+                            async with db.begin_nested() if db.in_transaction() else _NoopCtx():
+                                result = await db.execute(
+                                    sa_select(User.id).join(UserTenant, UserTenant.user_id == User.id)
+                                    .where(UserTenant.tenant_id == tenant_id)
+                                    .where(User.is_active == True)  # noqa: E712
+                                    .limit(1)
+                                )
+                                admin_row = result.first()
+                                if admin_row:
+                                    await post_system_message(
+                                        db=db,
+                                        tenant_id=tenant_id if isinstance(tenant_id, uuid.UUID) else uuid.UUID(str(tenant_id)),
+                                        user_id=admin_row[0],
+                                        message_type="cost_alert",
+                                        title=f"LLM Cost Alert: {pct}% of budget reached",
+                                        body=f"Current monthly LLM cost: ${current_cost:.2f} / ${budget:.2f} ({pct}%).",
+                                        severity="warning" if pct < 100 else "error",
+                                    )
+                        except Exception:
+                            logger.debug("Failed to send cost alert notification")
+    except Exception:
+        logger.debug("Failed to check cost alerts")
+
+
+class _NoopCtx:
+    """No-op async context manager for optional transaction nesting."""
+    async def __aenter__(self):
+        return self
+    async def __aexit__(self, *args):
+        return False
+
+
 async def get_api_credentials(
    db: AsyncSession | None,
    tenant_id: uuid.UUID | None,
@@ -291,6 +461,9 @@ async def llm_complete(
    response_format: dict[str, Any] | None = None,
    timeout: int = DEFAULT_TIMEOUT,
    max_retries: int = DEFAULT_MAX_RETRIES,
+    trace_id: str | None = None,
+    tenant_id: uuid.UUID | str | None = None,
+    db: AsyncSession | None = None,
 ) -> dict[str, Any]:
    """Generic LLM chat completion via LiteLLM with retry and cost tracking.

@@ -310,17 +483,28 @@ async def llm_complete(
        response_format: Optional response format spec (e.g. JSON mode).
        timeout: Request timeout in seconds (default 30).
        max_retries: Max retry attempts for transient errors (default 2).
+        trace_id: Optional trace ID for request correlation/logging.
+        tenant_id: Optional tenant ID for cost tracking and budget enforcement.
+        db: Optional DB session for cost alert notifications.

    Returns:
        Dict with keys: ``content``, ``usage``, ``cost_usd``, ``model``,
        ``raw_response`` (the LiteLLM response object for advanced use).

    Raises:
+        ValueError: If tenant LLM budget is exceeded (hard cutoff).
        Exception: Permanent errors or after exhausting retries.
    """
+    # Check tenant budget before making the call
+    await _check_tenant_budget(tenant_id, db)
+
    # Build LiteLLM model string
    litellm_model = build_model(model, provider)

+    # Log trace_id correlation if provided
+    if trace_id:
+        logger.debug("llm_complete trace_id=%s model=%s", trace_id, litellm_model)
+
    # Build kwargs
    kwargs: dict[str, Any] = {
        "model": litellm_model,
@@ -348,12 +532,16 @@ async def llm_complete(
            cost_usd = _extract_cost_usd(response, litellm_model)

            logger.debug(
-                "llm_complete success: model=%s tokens=%d cost=$%.6f attempt=%d",
+                "llm_complete success: model=%s tokens=%d cost=$%.6f attempt=%d trace_id=%s",
                litellm_model,
                usage["total_tokens"],
                cost_usd,
                attempt + 1,
+                trace_id or "-",
            )
+            # Track cost in Redis for tenant budget enforcement
+            if tenant_id is not None and cost_usd > 0:
+                await _track_tenant_cost(tenant_id, cost_usd, db)
            return {
                "content": content,
                "usage": usage,
@@ -400,6 +588,7 @@ async def llm_embed(
    provider: str | None = None,
    dimensions: int | None = None,
    timeout: int = DEFAULT_TIMEOUT,
+    trace_id: str | None = None,
 ) -> list[list[float]]:
    """Generic text embedding via LiteLLM.

@@ -409,18 +598,25 @@ async def llm_embed(
    Args:
        texts: Single text string or list of texts to embed.
        model: Embedding model name (default: ``openai/text-embedding-3-small``).
-        db: Optional DB session for API key lookup.
-        tenant_id: Optional tenant ID for API key lookup.
+        db: Optional DB session for API key lookup and cost alert notifications.
+        tenant_id: Optional tenant ID for cost tracking and budget enforcement.
        api_key: Override API key. If ``None``, uses env/DB lookup.
        api_base: Override API base URL.
        provider: Provider prefix override.
        dimensions: Override embedding dimensions.
        timeout: Request timeout in seconds (default 30).
+        trace_id: Optional trace ID for request correlation/logging.

    Returns:
        List of embedding vectors (each a list of floats). For a single
        text input, returns a one-element list.
    """
+    # Check tenant budget before making the call
+    await _check_tenant_budget(tenant_id, db)
+
+    if trace_id:
+        logger.debug("llm_embed trace_id=%s", trace_id)
+
    # Normalise to list input
    single_input = isinstance(texts, str)
    text_list = [texts] if single_input else texts
@@ -467,11 +663,20 @@ async def llm_embed(
    try:
        response = await litellm.aembedding(**litellm_kwargs)
        embeddings = [d["embedding"] for d in response.data]
+        # Track embedding cost for tenant budget enforcement
+        if tenant_id is not None:
+            try:
+                emb_cost = _extract_cost_usd(response, litellm_model)
+                if emb_cost > 0:
+                    await _track_tenant_cost(tenant_id, emb_cost, db)
+            except Exception:
+                pass
        logger.debug(
-            "llm_embed success: model=%s count=%d dims=%d",
+            "llm_embed success: model=%s count=%d dims=%d trace_id=%s",
            litellm_model,
            len(embeddings),
            len(embeddings[0]) if embeddings else 0,
+            trace_id or "-",
        )
        return embeddings
    except Exception:
@@ -107,6 +107,10 @@ class Settings(BaseSettings):
    rate_limit_webhook_max: int = 100  # incoming webhooks
    rate_limit_webhook_window: int = 60  # 1 minute

+    # LLM Cost Overrun Protection (B.17)
+    llm_monthly_budget_usd: float = 100.0  # per-tenant monthly LLM budget
+    llm_hard_cutoff: bool = True  # block LLM calls when budget exceeded
+
    @property
    def cors_origin_list(self) -> list[str]:
        """Parse comma-separated CORS origins into a list."""
@@ -1,19 +1,177 @@
-"""Standardized error codes for consistent frontend handling."""
+"""Standardized error codes, categories, and unified error-response format.

-ERROR_CODES = {
-    'not_found': {'status': 404, 'message': 'Resource not found'},
-    'permission_denied': {'status': 403, 'message': 'Permission denied'},
-    'validation_error': {'status': 422, 'message': 'Validation failed'},
-    'rate_limited': {'status': 429, 'message': 'Too many requests'},
-    'internal_error': {'status': 500, 'message': 'Internal server error'},
-    'service_unavailable': {'status': 503, 'message': 'Service temporarily unavailable'},
+Every API error response follows the schema:
+
+    {
+        "code": "not_found",
+        "detail": "Resource not found",
+        "field": null,
+        "trace_id": "a1b2c3d4",
+        "retryable": false,
+        "category": "permanent"
+    }
+
+``ErrorCategory`` classifies errors so callers can decide retry strategy.
+"""
+
+from __future__ import annotations
+
+import enum
+from typing import Any
+
+
+class ErrorCategory(str, enum.Enum):
+    """Error classification for retry decisions."""
+
+    TRANSIENT = "transient"   # retryable: timeout, rate-limit, connection
+    PERMANENT = "permanent"    # non-retryable: validation, permission, not_found
+    PARTIAL = "partial"        # partly successful: batch, bulk operations
+
+
+ERROR_CODES: dict[str, dict[str, Any]] = {
+    # ── Original 6 codes ──
+    "not_found":          {"status": 404, "message": "Resource not found",          "category": ErrorCategory.PERMANENT,  "retryable": False},
+    "permission_denied":  {"status": 403, "message": "Permission denied",           "category": ErrorCategory.PERMANENT,  "retryable": False},
+    "validation_error":   {"status": 422, "message": "Validation failed",            "category": ErrorCategory.PERMANENT,  "retryable": False},
+    "rate_limited":       {"status": 429, "message": "Too many requests",            "category": ErrorCategory.TRANSIENT,  "retryable": True},
+    "internal_error":     {"status": 500, "message": "Internal server error",        "category": ErrorCategory.TRANSIENT,  "retryable": True},
+    "service_unavailable": {"status": 503, "message": "Service temporarily unavailable", "category": ErrorCategory.TRANSIENT, "retryable": True},
+    # ── New codes (B-ERR-FMT) ──
+    "forbidden":          {"status": 403, "message": "Forbidden",                    "category": ErrorCategory.PERMANENT,  "retryable": False},
+    "conflict":           {"status": 409, "message": "Conflict with current state",  "category": ErrorCategory.PERMANENT,  "retryable": False},
+    "unprocessable":      {"status": 422, "message": "Unprocessable entity",         "category": ErrorCategory.PERMANENT,  "retryable": False},
+    "not_implemented":    {"status": 501, "message": "Not implemented",               "category": ErrorCategory.PERMANENT,  "retryable": False},
+    "service_timeout":    {"status": 504, "message": "Service timed out",             "category": ErrorCategory.TRANSIENT,  "retryable": True},
+    "bad_gateway":        {"status": 502, "message": "Bad gateway",                   "category": ErrorCategory.TRANSIENT,  "retryable": True},
+    # ── Partial success ──
+    "partial_success":    {"status": 207, "message": "Partial success",               "category": ErrorCategory.PARTIAL,    "retryable": False},
 }


 class ApiError(Exception):
-    def __init__(self, code: str, detail: str = None, field: str = None, status: int = None):
+    """Application-level error with code, category, and retryable flag.
+
+    Attributes:
+        code: Error code key from ``ERROR_CODES``.
+        detail: Human-readable detail message.
+        field: Optional field name that caused the error.
+        status: HTTP status code.
+        category: ``ErrorCategory`` for retry decisions.
+        retryable: Whether the caller may retry the request.
+    """
+
+    def __init__(
+        self,
+        code: str,
+        detail: str | None = None,
+        field: str | None = None,
+        status: int | None = None,
+        category: ErrorCategory | None = None,
+        retryable: bool | None = None,
+    ):
+        meta = ERROR_CODES.get(code, {})
        self.code = code
-        self.detail = detail or ERROR_CODES.get(code, {}).get('message', 'Unknown error')
+        self.detail = detail or meta.get("message", "Unknown error")
        self.field = field
-        self.status = status or ERROR_CODES.get(code, {}).get('status', 500)
+        self.status = status or meta.get("status", 500)
+        self.category = category or meta.get("category", ErrorCategory.TRANSIENT)
+        self.retryable = retryable if retryable is not None else meta.get("retryable", True)
        super().__init__(self.detail)
+
+    def to_response(self, trace_id: str | None = None) -> dict[str, Any]:
+        """Build the unified error-response dict."""
+        resp: dict[str, Any] = {
+            "code": self.code,
+            "detail": self.detail,
+            "field": self.field,
+            "trace_id": trace_id,
+            "retryable": self.retryable,
+            "category": self.category.value if isinstance(self.category, ErrorCategory) else str(self.category),
+        }
+        return resp
+
+
+# ── Exception classification helper ──────────────────────────────────────────
+
+# Transient error indicators
+_TRANSIENT_KEYWORDS = frozenset({
+    "timeout", "timed out", "rate limit", "rate_limit", "429", "503", "502", "504",
+    "service unavailable", "overloaded", "connection reset", "connection aborted",
+    "temporary", "transient", "retry",
+})
+
+# Permanent error indicators
+_PERMANENT_KEYWORDS = frozenset({
+    "authentication", "auth", "401", "403", "unauthorized", "forbidden",
+    "invalid api key", "invalid_api_key", "validation", "invalid_request",
+    "400", "bad request", "model_not_found", "not found", "404", "409",
+    "conflict", "not implemented", "501", "permission",
+})
+
+# Partial error indicators
+_PARTIAL_KEYWORDS = frozenset({
+    "partial", "batch", "bulk", "some failed", "multi-status", "207",
+})
+
+
+def classify_exception(exc: Exception) -> ErrorCategory:
+    """Classify an exception into an ``ErrorCategory``.
+
+    Uses string matching on the exception message and type name.
+    Falls back to ``ErrorCategory.TRANSIENT`` for unknown errors (safer to retry).
+
+    Args:
+        exc: The exception to classify.
+
+    Returns:
+        ``ErrorCategory.TRANSIENT``, ``ErrorCategory.PERMANENT``, or
+        ``ErrorCategory.PARTIAL``.
+    """
+    # If it's already an ApiError, use its category
+    if isinstance(exc, ApiError):
+        return exc.category if isinstance(exc.category, ErrorCategory) else ErrorCategory(exc.category)
+
+    import asyncio as _asyncio
+
+    msg = str(exc).lower()
+    exc_type_name = type(exc).__name__.lower()
+
+    # Check partial first — batch/bulk errors
+    if any(kw in msg or kw in exc_type_name for kw in _PARTIAL_KEYWORDS):
+        return ErrorCategory.PARTIAL
+
+    # Check permanent — auth/validation/permission errors should never be retried
+    if any(kw in msg or kw in exc_type_name for kw in _PERMANENT_KEYWORDS):
+        return ErrorCategory.PERMANENT
+
+    # Check transient
+    if any(kw in msg or kw in exc_type_name for kw in _TRANSIENT_KEYWORDS):
+        return ErrorCategory.TRANSIENT
+
+    # asyncio.TimeoutError is always transient
+    if isinstance(exc, (_asyncio.TimeoutError, TimeoutError, ConnectionError)):
+        return ErrorCategory.TRANSIENT
+
+    # Default: treat as transient (safe to retry)
+    return ErrorCategory.TRANSIENT
+
+
+def build_error_response(
+    code: str,
+    detail: str | None = None,
+    field: str | None = None,
+    status: int | None = None,
+    trace_id: str | None = None,
+) -> dict[str, Any]:
+    """Build a unified error-response dict without raising an exception."""
+    meta = ERROR_CODES.get(code, {})
+    return {
+        "code": code,
+        "detail": detail or meta.get("message", "Unknown error"),
+        "field": field,
+        "trace_id": trace_id,
+        "retryable": meta.get("retryable", True),
+        "category": meta.get("category", ErrorCategory.TRANSIENT).value
+        if isinstance(meta.get("category"), ErrorCategory)
+        else str(meta.get("category", ErrorCategory.TRANSIENT)),
+    }
@@ -51,10 +51,50 @@ arq_jobs_total = Counter(

 # ─── Structured Logging (structlog) ───

+# Sensitive field names that should be redacted from log output
+_SENSITIVE_FIELDS = frozenset({
+    "password", "passwd", "secret", "api_key", "apikey", "token",
+    "authorization", "auth", "cookie", "session_id", "session",
+    "private_key", "privatekey", "credentials", "smtp_password",
+    "mail_password", "encryption_key", "secret_key",
+})
+
+
+def _sanitize_sensitive_fields(logger, method_name, event_dict):
+    """structlog processor that redacts sensitive field values."""
+    for key in list(event_dict.keys()):
+        if key.lower() in _SENSITIVE_FIELDS:
+            event_dict[key] = "***REDACTED***"
+    return event_dict
+
+
+def sanitize_dict(data: dict[str, Any], extra_sensitive: set[str] | None = None) -> dict[str, Any]:
+    """Sanitize a dictionary by redacting sensitive field values.
+
+    Args:
+        data: The dictionary to sanitize.
+        extra_sensitive: Additional field names to treat as sensitive.
+
+    Returns:
+        A copy of *data* with sensitive values replaced by ``"***REDACTED***"``.
+    """
+    sensitive = _SENSITIVE_FIELDS | (extra_sensitive or set())
+    result = {}
+    for key, value in data.items():
+        if key.lower() in sensitive:
+            result[key] = "***REDACTED***"
+        elif isinstance(value, dict):
+            result[key] = sanitize_dict(value, extra_sensitive)
+        else:
+            result[key] = value
+    return result
+
+
 structlog.configure(
    processors=[
        structlog.contextvars.merge_contextvars,
        structlog.processors.add_log_level,
+        _sanitize_sensitive_fields,
        structlog.processors.TimeStamper(fmt="iso"),
        structlog.processors.JSONRenderer(),
    ],
@@ -176,8 +176,35 @@ async def on_startup(ctx: dict[str, Any]) -> None:


 async def on_shutdown(ctx: dict[str, Any]) -> None:
-    """Called when worker shuts down."""
+    """Called when worker shuts down.
+
+    Pauses any running WorkflowRun instances so they can be resumed after
+    restart, then closes Redis.
+    """
    logger.info("ARQ worker shutting down...")
+
+    # Pause running workflow instances so they can be resumed after restart
+    try:
+        from app.core.db import get_worker_session_factory
+        from sqlalchemy import select as sa_select
+        from app.models.workflow import WorkflowInstance
+
+        session_factory = get_worker_session_factory()
+        async with session_factory() as db:
+            result = await db.execute(
+                sa_select(WorkflowInstance).where(
+                    WorkflowInstance.status == "running"
+                )
+            )
+            running = result.scalars().all()
+            if running:
+                for wf in running:
+                    wf.status = "paused"
+                await db.commit()
+                logger.info(f"Paused {len(running)} running workflow(s) for graceful shutdown")
+    except Exception as exc:
+        logger.warning(f"Failed to pause running workflows during shutdown: {exc}")
+
    from app.core.auth import close_redis
    await close_redis()

@@ -146,6 +146,57 @@ async def send_ws_error(
        logger.debug("Failed to send WS error to client")


+# ── Global connection registry for drain_all_connections ────────────────────
+# Plugin WS endpoints register their connection registries here so that
+# drain_all_connections() can close them all during graceful shutdown.
+_global_ws_registries: list[dict[str, list[WebSocket]]] = []
+
+
+def register_ws_registry(registry: dict[str, list[WebSocket]]) -> None:
+    """Register a WebSocket connection registry for graceful shutdown."""
+    if registry not in _global_ws_registries:
+        _global_ws_registries.append(registry)
+
+
+async def drain_all_connections(grace_period_seconds: float = 5.0) -> None:
+    """Notify all connected WS clients about reconnect and close connections.
+
+    Sends a ``reconnect`` hint message to every connected client, waits
+    for ``grace_period_seconds``, then forcefully closes all sockets.
+    Called during application graceful shutdown.
+    """
+    total_connections = 0
+    for registry in _global_ws_registries:
+        for user_id, conns in list(registry.items()):
+            for ws in list(conns):
+                try:
+                    await ws.send_text(json.dumps({
+                        "type": "reconnect",
+                        "reason": "server_shutdown",
+                        "message": "Server is shutting down. Please reconnect shortly.",
+                    }))
+                    total_connections += 1
+                except Exception:
+                    logger.debug("Failed to send reconnect hint to WS client")
+
+    logger.info(f"WS drain: notified {total_connections} connections, waiting {grace_period_seconds}s")
+
+    if grace_period_seconds > 0:
+        await asyncio.sleep(grace_period_seconds)
+
+    # Close all connections
+    for registry in _global_ws_registries:
+        for user_id, conns in list(registry.items()):
+            for ws in list(conns):
+                try:
+                    await ws.close(code=1001, reason="Server shutting down")
+                except Exception:
+                    logger.debug("WS already closed during drain")
+            registry.clear()
+    _global_ws_registries.clear()
+    logger.info("WS drain: all connections closed")
+
+
 async def handle_ws_message(
    websocket: WebSocket,
    message: str,
@@ -2,10 +2,13 @@

 from __future__ import annotations

+import asyncio
 import time
 import traceback
+import uuid as _uuid
 from contextlib import asynccontextmanager

+import structlog
 from fastapi import FastAPI, HTTPException, Request, Depends, APIRouter
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse
@@ -20,7 +23,7 @@ logger = logging.getLogger(__name__)

 from app.config import get_settings
 from app.core.db import close_engine, get_engine
-from app.core.error_codes import ApiError
+from app.core.error_codes import ApiError, ErrorCategory, classify_exception, build_error_response
 from app.core.middleware import CSRFMiddleware, SecurityHeadersMiddleware
 from app.core.rate_limit import GeneralRateLimitMiddleware
 from app.core.resilience import CircuitBreakerMiddleware
@@ -74,14 +77,58 @@ from app.routes import (
 )


+# ── Graceful shutdown signal ─────────────────────────────────────────────────
+# Set during lifespan shutdown so middleware and handlers can stop accepting work.
+_shutdown_event = asyncio.Event()
+
+# Track in-flight requests for graceful draining
+_inflight_requests: set[asyncio.Task] = set()
+
+
+async def _drain_inflight(timeout_per_request: float = 25.0) -> None:
+    """Wait for all in-flight request tasks to complete."""
+    if not _inflight_requests:
+        return
+    logger.info(f"Waiting for {len(_inflight_requests)} in-flight requests to complete")
+    # Give tasks a chance to finish; cancel remaining after timeout
+    done, pending = await asyncio.wait(
+        _inflight_requests,
+        timeout=timeout_per_request,
+    )
+    if pending:
+        logger.warning(f"Cancelling {len(pending)} in-flight requests that exceeded grace period")
+        for task in pending:
+            task.cancel()
+        await asyncio.gather(*pending, return_exceptions=True)
+
+
+def _generate_trace_id() -> str:
+    """Generate a short trace ID (first 8 chars of UUID4)."""
+    return _uuid.uuid4().hex[:8]
+
+
+def _get_trace_id() -> str | None:
+    """Get the current trace_id from structlog contextvars (best-effort)."""
+    try:
+        ctx = structlog.contextvars.get_contextvars()
+        return ctx.get("trace_id")
+    except Exception:
+        return None
+
+
 class RequestLoggingMiddleware(BaseHTTPMiddleware):
-    """Structured logging + Prometheus metrics for every HTTP request."""
+    """Structured logging + Prometheus metrics + trace_id for every HTTP request."""

    async def dispatch(self, request: Request, call_next):
        start_time = time.perf_counter()
        method = request.method
        path = request.url.path

+        # Generate trace_id and bind to structlog contextvars for this request
+        trace_id = _generate_trace_id()
+        structlog.contextvars.clear_contextvars()
+        structlog.contextvars.bind_contextvars(trace_id=trace_id)
+
        # Extract tenant_id from session cookie if available (best-effort)
        tenant_id = None

@@ -106,7 +153,7 @@ class RequestLoggingMiddleware(BaseHTTPMiddleware):
                    "message": f"[Backend] {method} {path}: {exc}",
                    "stack": tb_str,
                    "url": str(request.url),
-                    "context": {"method": method, "path": path, "source": "backend_middleware"},
+                    "context": {"method": method, "path": path, "source": "backend_middleware", "trace_id": trace_id},
                })
            except Exception:
                pass  # Never let error reporting break the request
@@ -115,6 +162,9 @@ class RequestLoggingMiddleware(BaseHTTPMiddleware):
        duration_ms = (time.perf_counter() - start_time) * 1000
        status_code = response.status_code

+        # Add trace_id to response header
+        response.headers["X-Trace-Id"] = trace_id
+
        # Report 4xx and 5xx errors to Forgejo (except 401/403 which are expected)
        if status_code >= 400 and status_code not in (401, 403):
            try:
@@ -122,7 +172,7 @@ class RequestLoggingMiddleware(BaseHTTPMiddleware):
                await report_error_to_forgejo({
                    "message": f"[Backend] {method} {path} → {status_code}",
                    "url": str(request.url),
-                    "context": {"method": method, "path": path, "status": status_code, "source": "backend_response"},
+                    "context": {"method": method, "path": path, "status": status_code, "source": "backend_response", "trace_id": trace_id},
                })
            except Exception:
                pass  # Never let error reporting break the response
@@ -137,6 +187,9 @@ class RequestLoggingMiddleware(BaseHTTPMiddleware):
            tenant_id=tenant_id,
        )

+        # Clear contextvars after request completes
+        structlog.contextvars.clear_contextvars()
+
        return response


@@ -308,10 +361,29 @@ async def lifespan(app: FastAPI):

    yield

-    # Shutdown: close global Redis and ARQ pool
+    # ── Graceful shutdown ────────────────────────────────────────────────────
+    # Signal that we're shutting down — no new requests should be accepted.
+    _shutdown_event.set()
+    logger.info("Graceful shutdown initiated — draining in-flight requests")
+
+    # Drain WebSocket connections (notify clients to reconnect)
+    try:
+        from app.core.ws_helpers import drain_all_connections
+        await drain_all_connections(grace_period_seconds=5)
+    except Exception as exc:
+        logger.warning(f"WS drain failed during shutdown: {exc}")
+
+    # Give in-flight requests time to complete (max 30s)
+    try:
+        await asyncio.wait_for(_drain_inflight(), timeout=30.0)
+    except asyncio.TimeoutError:
+        logger.warning("Graceful shutdown: 30s timeout reached, forcing shutdown")
+
+    # Close global Redis and ARQ pool
    await close_job_pool()
    await close_redis()
    await close_engine()
+    logger.info("Graceful shutdown complete")


 def create_app() -> FastAPI:
@@ -393,7 +465,8 @@ def create_app() -> FastAPI:
    # ── Global exception handler — catch ALL unhandled exceptions ──
    @app.exception_handler(Exception)
    async def global_exception_handler(request: Request, exc: Exception):
-        logger.error(f"Unhandled exception: {exc}", exc_info=True)
+        trace_id = _get_trace_id()
+        logger.error(f"Unhandled exception: {exc}", exc_info=True, extra={"trace_id": trace_id})
        record_error(
            event="unhandled_exception",
            method=request.method,
@@ -401,18 +474,51 @@ def create_app() -> FastAPI:
            status_code=500,
            error=str(exc),
        )
-        return JSONResponse(
-            status_code=500,
-            content={"detail": "Internal server error", "code": "internal_error"},
+        body = build_error_response(
+            code="internal_error",
+            detail="Internal server error",
+            trace_id=trace_id,
        )
+        resp = JSONResponse(status_code=500, content=body)
+        if trace_id:
+            resp.headers["X-Trace-Id"] = trace_id
+        return resp

-    # ── ApiError handler — structured error responses ──
+    # ── HTTPException handler — unified format ──
+    @app.exception_handler(HTTPException)
+    async def http_exception_handler(request: Request, exc: HTTPException):
+        trace_id = _get_trace_id()
+        # Map common HTTP status codes to error codes
+        status_to_code = {
+            404: "not_found",
+            403: "forbidden",
+            422: "unprocessable",
+            429: "rate_limited",
+            501: "not_implemented",
+            502: "bad_gateway",
+            503: "service_unavailable",
+            504: "service_timeout",
+        }
+        code = status_to_code.get(exc.status_code, "internal_error" if exc.status_code >= 500 else "validation_error")
+        body = build_error_response(
+            code=code,
+            detail=str(exc.detail) if exc.detail else None,
+            trace_id=trace_id,
+        )
+        resp = JSONResponse(status_code=exc.status_code, content=body)
+        if trace_id:
+            resp.headers["X-Trace-Id"] = trace_id
+        return resp
+
+    # ── ApiError handler — structured error responses with category/retryable ──
    @app.exception_handler(ApiError)
    async def api_error_handler(request: Request, exc: ApiError):
-        return JSONResponse(
-            status_code=exc.status,
-            content={'code': exc.code, 'detail': exc.detail, 'field': exc.field}
-        )
+        trace_id = _get_trace_id()
+        body = exc.to_response(trace_id=trace_id)
+        resp = JSONResponse(status_code=exc.status, content=body)
+        if trace_id:
+            resp.headers["X-Trace-Id"] = trace_id
+        return resp

    app.include_router(health.router)
    app.include_router(metrics.router)