timeout.py

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.responses import JSONResponse, Response
from starlette.types import ASGIApp
from fastapi import Request
import asyncio
import os
import time
from typing import Dict, Mapping, Optional


class TimeoutMiddleware(BaseHTTPMiddleware):
    def __init__(
        self,
        app: ASGIApp,
        methods_timeout: Optional[Mapping[str, object]] = None,
        paths_timeout: Optional[Mapping[str, object]] = None,
    ) -> None:
        super().__init__(app)

        self.default_timeout = self._get_timeout_from_env("HTTP_DEFAULT_TIMEOUT", default=2 * 60)
        self.maximum_age_seconds = self._get_timeout_from_env("HTTP_MAXIMUM_AGE_SECONDS", default=5 * 60)
        self.clock_skew_allowance = self._get_timeout_from_env("HTTP_CLOCK_SKEW_ALLOWANCE", default=5 * 60)

        self.methods_timeout = self._parse_methods_timeout(methods_timeout or {})
        self.paths_timeout = self._parse_paths_timeout(paths_timeout or {})

    @staticmethod
    def _get_timeout_from_env(env_var: str, default: float) -> float:
        timeout = os.environ.get(env_var, default)
        try:
            return float(timeout)
        except (TypeError, ValueError):
            raise ValueError(f"Invalid timeout value in env {env_var}: {timeout}")

    @staticmethod
    def _parse_methods_timeout(methods_timeout: Mapping[str, object]) -> Dict[str, float]:
        result: Dict[str, float] = {}
        for method, timeout in methods_timeout.items():
            if timeout is None:
                continue
            try:
                result[method.upper()] = float(timeout)  # type: ignore[arg-type]  # guarded by try/except
            except (TypeError, ValueError):
                raise ValueError(f"Invalid timeout value for method {method}: {timeout}")
        return result

    @staticmethod
    def _parse_paths_timeout(paths_timeout: Mapping[str, object]) -> Dict[str, float]:
        result: Dict[str, float] = {}
        for path, timeout in paths_timeout.items():
            if timeout is None:
                continue
            try:
                result[path] = float(timeout)  # type: ignore[arg-type]  # guarded by try/except
            except (TypeError, ValueError):
                raise ValueError(f"Invalid timeout value for path {path}: {timeout}")
        return result

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        # Check for stale request header first
        # Uses clock_skew_allowance to tolerate client/server clock drift (#5929)
        request_start_header = request.headers.get("x-request-start-time")
        if request_start_header:
            try:
                request_start_time = float(request_start_header)
                current_time = time.time()
                request_age = current_time - request_start_time

                if request_age > self.maximum_age_seconds + self.clock_skew_allowance:
                    return JSONResponse(
                        status_code=408,
                        content={
                            "error": "clock_skew",
                            "message": "Request rejected — your device clock may be out of sync",
                            "server_time": current_time,
                            "client_time": request_start_time,
                            "skew_seconds": round(request_age, 1),
                            "hint": "Check your device date/time settings and enable automatic time",
                        },
                    )
            except (ValueError, TypeError):
                pass

        path_timeout = self.paths_timeout.get(request.url.path)
        timeout = (
            path_timeout if path_timeout is not None else self.methods_timeout.get(request.method, self.default_timeout)
        )
        # Stamp the monotonic request start so request-scoped read budgets
        # (utils.other.list_budget) can derive an internal deadline that leaves
        # serialization headroom under this middleware's hard cutoff.
        request.state.omi_request_started_monotonic = time.monotonic()
        try:
            return await asyncio.wait_for(call_next(request), timeout=timeout)
        except asyncio.TimeoutError:
            return Response(status_code=504)