svc-infra 0.1.589__py3-none-any.whl → 0.1.706__py3-none-any.whl

svc_infra/api/fastapi/dx.py

@@ -59,19 +59,36 @@ from svc_infra.api.fastapi.auth.security import (
     RequireUser,
 )
 from svc_infra.api.fastapi.auth.settings import AuthSettings, get_auth_settings
-from svc_infra.api.fastapi.dual.protected import (
+
+# ----------------
+# WebSocket identity primitives (lightweight JWT, no DB required)
+# ----------------
+from svc_infra.api.fastapi.auth.ws_security import (
+    AllowWSIdentity,
+    OptionalWSIdentity,
+    RequireWSAnyScope,
+    RequireWSIdentity,
+    RequireWSScopes,
+    WSIdentity,
+    WSPrincipal,
+)
+from svc_infra.api.fastapi.dual.protected import (  # WebSocket routers (DualAPIRouter with JWT auth, no DB required)
     optional_identity_router,
     protected_router,
     roles_router,
     scopes_router,
     service_router,
     user_router,
+    ws_optional_router,
+    ws_protected_router,
+    ws_scopes_router,
+    ws_user_router,
 )
 
 # ----------------
 # Pre-wired routers (OpenAPI security auto-injected)
 # ----------------
-from svc_infra.api.fastapi.dual.public import public_router
+from svc_infra.api.fastapi.dual.public import public_router, ws_public_router
 
 # ----------------
 # App bootstrap
@@ -127,6 +144,20 @@ __all__ = [
     "service_router",
     "scopes_router",
     "roles_router",
+    # WebSocket identity
+    "WSPrincipal",
+    "WSIdentity",
+    "OptionalWSIdentity",
+    "RequireWSIdentity",
+    "AllowWSIdentity",
+    "RequireWSScopes",
+    "RequireWSAnyScope",
+    # WebSocket routers
+    "ws_public_router",
+    "ws_protected_router",
+    "ws_user_router",
+    "ws_scopes_router",
+    "ws_optional_router",
     # Feature routers
     "apikey_router",
     "mfa_router",

svc_infra/api/fastapi/ease.py

@@ -104,9 +104,13 @@ class EasyAppOptions(BaseModel):
                 else self.logging.enable
             ),
             level=(
-                override.logging.level if override.logging.level is not None else self.logging.level
+                override.logging.level
+                if override.logging.level is not None
+                else self.logging.level
             ),
-            fmt=override.logging.fmt if override.logging.fmt is not None else self.logging.fmt,
+            fmt=override.logging.fmt
+            if override.logging.fmt is not None
+            else self.logging.fmt,
         )
 
         # observability
@@ -148,6 +152,7 @@ def easy_service_api(
     public_cors_origins: list[str] | str | None = None,
     root_public_base_url: str | None = None,
     root_include_api_key: bool | None = None,
+    **fastapi_kwargs,  # Forward all other FastAPI kwargs
 ) -> FastAPI:
     service = ServiceInfo(name=name, release=release)
     specs = [
@@ -161,6 +166,7 @@ def easy_service_api(
         public_cors_origins=public_cors_origins,
         root_public_base_url=root_public_base_url,
         root_include_api_key=root_include_api_key,
+        **fastapi_kwargs,  # Forward to setup_service_api
     )
 
 
@@ -176,6 +182,7 @@ def easy_service_app(
     options: EasyAppOptions | None = None,
     enable_logging: bool | None = None,
     enable_observability: bool | None = None,
+    **fastapi_kwargs,  # Forward all other FastAPI kwargs (lifespan, etc.)
 ) -> FastAPI:
     """
     One-call bootstrap with env + options + flags:
@@ -226,6 +233,7 @@ def easy_service_app(
         public_cors_origins=public_cors_origins,
         root_public_base_url=root_public_base_url,
         root_include_api_key=root_include_api_key,
+        **fastapi_kwargs,  # Forward FastAPI kwargs (lifespan, etc.)
     )
 
     # 5) Observability

svc_infra/api/fastapi/http/concurrency.py

@@ -10,5 +10,6 @@ def require_if_match(request: Request, current_etag: str):
         )
     if current_etag not in [t.strip() for t in val.split(",")]:
         raise HTTPException(
-            status_code=status.HTTP_412_PRECONDITION_FAILED, detail="ETag precondition failed."
+            status_code=status.HTTP_412_PRECONDITION_FAILED,
+            detail="ETag precondition failed.",
         )

svc_infra/api/fastapi/http/conditional.py

@@ -20,7 +20,9 @@ def set_conditional_headers(
         resp.headers["Last-Modified"] = format_datetime(last_modified)
 
 
-def maybe_not_modified(request: Request, etag: str | None, last_modified: datetime | None) -> bool:
+def maybe_not_modified(
+    request: Request, etag: str | None, last_modified: datetime | None
+) -> bool:
     inm = request.headers.get("If-None-Match")
     ims = request.headers.get("If-Modified-Since")
     etag_ok = etag and inm and etag in [t.strip() for t in inm.split(",")]

svc_infra/api/fastapi/middleware/debug.py

@@ -15,6 +15,9 @@ class RouteDebugMiddleware(BaseHTTPMiddleware):
         route = request.scope.get("route")
         ep = getattr(route, "endpoint", None) if route else None
         log.info(
-            "MATCHED %s %s -> %s", request.method, request.url.path, getattr(ep, "__name__", ep)
+            "MATCHED %s %s -> %s",
+            request.method,
+            request.url.path,
+            getattr(ep, "__name__", ep),
         )
         return response

svc_infra/api/fastapi/middleware/errors/catchall.py

@@ -31,7 +31,9 @@ class CatchAllExceptionMiddleware:
 
             if response_started:
                 try:
-                    await send({"type": "http.response.body", "body": b"", "more_body": False})
+                    await send(
+                        {"type": "http.response.body", "body": b"", "more_body": False}
+                    )
                 except Exception:
                     pass
             else:
@@ -52,4 +54,6 @@ class CatchAllExceptionMiddleware:
                         "headers": [(b"content-type", PROBLEM_MT.encode("ascii"))],
                     }
                 )
-                await send({"type": "http.response.body", "body": body, "more_body": False})
+                await send(
+                    {"type": "http.response.body", "body": body, "more_body": False}
+                )

svc_infra/api/fastapi/middleware/errors/exceptions.py

@@ -12,7 +12,7 @@ class FastApiException(Exception):
         detail: Optional[str] = None,
         status_code: int = 400,
         *,
-        code: str | None = None
+        code: str | None = None,
     ):
         self.title = title
         self.detail = detail

svc_infra/api/fastapi/middleware/errors/handlers.py

@@ -2,6 +2,7 @@ import logging
 import traceback
 from typing import Any, Dict, Optional
 
+import httpx
 from fastapi import Request
 from fastapi.exceptions import HTTPException, RequestValidationError
 from fastapi.responses import JSONResponse, Response
@@ -46,6 +47,7 @@ def problem_response(
     code: str | None = None,
     errors: list[dict] | None = None,
     trace_id: str | None = None,
+    headers: dict[str, str] | None = None,
 ) -> Response:
     body: Dict[str, Any] = {
         "type": type_uri,
@@ -62,10 +64,30 @@ def problem_response(
         body["errors"] = errors
     if trace_id:
         body["trace_id"] = trace_id
-    return JSONResponse(status_code=status, content=body, media_type=PROBLEM_MT)
+    return JSONResponse(
+        status_code=status, content=body, media_type=PROBLEM_MT, headers=headers
+    )
 
 
 def register_error_handlers(app):
+    @app.exception_handler(httpx.TimeoutException)
+    async def handle_httpx_timeout(request: Request, exc: httpx.TimeoutException):
+        trace_id = _trace_id_from_request(request)
+        # Map outbound HTTP client timeouts to 504 Gateway Timeout
+        # Keep details generic in prod
+        return problem_response(
+            status=504,
+            title="Gateway Timeout",
+            detail=(
+                "Upstream request timed out."
+                if IS_PROD
+                else (str(exc) or "httpx timeout")
+            ),
+            code="GATEWAY_TIMEOUT",
+            instance=str(request.url),
+            trace_id=trace_id,
+        )
+
     @app.exception_handler(FastApiException)
     async def handle_app_exception(request: Request, exc: FastApiException):
         trace_id = _trace_id_from_request(request)
@@ -104,14 +126,25 @@ def register_error_handlers(app):
     @app.exception_handler(HTTPException)
     async def handle_http_exception(request: Request, exc: HTTPException):
         trace_id = _trace_id_from_request(request)
-        title = {401: "Unauthorized", 403: "Forbidden", 404: "Not Found"}.get(
-            exc.status_code, "Error"
-        )
+        title = {
+            401: "Unauthorized",
+            403: "Forbidden",
+            404: "Not Found",
+            429: "Too Many Requests",
+        }.get(exc.status_code, "Error")
         detail = (
             exc.detail
             if not IS_PROD or exc.status_code < 500
             else "Something went wrong. Please contact support."
         )
+        # Preserve headers set on the exception (e.g., Retry-After for rate limits)
+        hdrs: dict[str, str] | None = None
+        try:
+            if getattr(exc, "headers", None):
+                # FastAPI/Starlette exceptions store headers as a dict[str, str]
+                hdrs = dict(getattr(exc, "headers"))
+        except Exception:
+            hdrs = None
         return problem_response(
             status=exc.status_code,
             title=title,
@@ -119,19 +152,31 @@ def register_error_handlers(app):
             code=title.replace(" ", "_").upper(),
             instance=str(request.url),
             trace_id=trace_id,
+            headers=hdrs,
         )
 
     @app.exception_handler(StarletteHTTPException)
-    async def handle_starlette_http_exception(request: Request, exc: StarletteHTTPException):
+    async def handle_starlette_http_exception(
+        request: Request, exc: StarletteHTTPException
+    ):
         trace_id = _trace_id_from_request(request)
-        title = {401: "Unauthorized", 403: "Forbidden", 404: "Not Found"}.get(
-            exc.status_code, "Error"
-        )
+        title = {
+            401: "Unauthorized",
+            403: "Forbidden",
+            404: "Not Found",
+            429: "Too Many Requests",
+        }.get(exc.status_code, "Error")
         detail = (
             exc.detail
             if not IS_PROD or exc.status_code < 500
             else "Something went wrong. Please contact support."
         )
+        hdrs: dict[str, str] | None = None
+        try:
+            if getattr(exc, "headers", None):
+                hdrs = dict(getattr(exc, "headers"))
+        except Exception:
+            hdrs = None
         return problem_response(
             status=exc.status_code,
             title=title,
@@ -139,6 +184,7 @@ def register_error_handlers(app):
             code=title.replace(" ", "_").upper(),
             instance=str(request.url),
             trace_id=trace_id,
+            headers=hdrs,
         )
 
     @app.exception_handler(IntegrityError)

svc_infra/api/fastapi/middleware/graceful_shutdown.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from contextlib import asynccontextmanager
+from typing import Optional
+
+from fastapi import FastAPI
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from svc_infra.app.env import pick
+
+logger = logging.getLogger(__name__)
+
+
+def _get_grace_period_seconds() -> float:
+    default = pick(prod=20.0, nonprod=5.0)
+    raw = os.getenv("SHUTDOWN_GRACE_PERIOD_SECONDS")
+    if raw is None or raw == "":
+        return float(default)
+    try:
+        return float(raw)
+    except ValueError:
+        return float(default)
+
+
+class InflightTrackerMiddleware:
+    """Tracks number of in-flight requests to support graceful shutdown drains."""
+
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        if scope.get("type") != "http":
+            await self.app(scope, receive, send)
+            return
+        app = scope.get("app")
+        if app is None:
+            await self.app(scope, receive, send)
+            return
+        state = getattr(app, "state", None)
+        if state is None:
+            await self.app(scope, receive, send)
+            return
+        state._inflight_requests = getattr(state, "_inflight_requests", 0) + 1
+        try:
+            await self.app(scope, receive, send)
+        finally:
+            state._inflight_requests = max(
+                0, getattr(state, "_inflight_requests", 1) - 1
+            )
+
+
+async def _wait_for_drain(app: FastAPI, grace: float) -> None:
+    interval = 0.1
+    waited = 0.0
+    while waited < grace:
+        inflight = int(getattr(app.state, "_inflight_requests", 0))
+        if inflight <= 0:
+            return
+        await asyncio.sleep(interval)
+        waited += interval
+    inflight = int(getattr(app.state, "_inflight_requests", 0))
+    if inflight > 0:
+        logger.warning(
+            "Graceful shutdown timeout: %s in-flight request(s) after %.2fs",
+            inflight,
+            waited,
+        )
+
+
+def install_graceful_shutdown(
+    app: FastAPI, *, grace_seconds: Optional[float] = None
+) -> None:
+    """Install inflight tracking and lifespan hooks to wait for requests to drain.
+
+    - Adds InflightTrackerMiddleware
+    - Registers a lifespan handler that initializes state and waits up to grace_seconds on shutdown
+    """
+    app.add_middleware(InflightTrackerMiddleware)
+
+    g = (
+        float(grace_seconds)
+        if grace_seconds is not None
+        else _get_grace_period_seconds()
+    )
+
+    # Preserve any existing lifespan and wrap it so our drain runs on shutdown.
+    previous_lifespan = getattr(app.router, "lifespan_context", None)
+
+    @asynccontextmanager
+    async def _lifespan(a: FastAPI):  # noqa: ANN202
+        # Startup: initialize inflight counter
+        a.state._inflight_requests = 0
+        if previous_lifespan is not None:
+            async with previous_lifespan(a):
+                yield
+        else:
+            yield
+        # Shutdown: wait for in-flight requests to drain (up to grace period)
+        await _wait_for_drain(a, g)
+
+    app.router.lifespan_context = _lifespan

svc_infra/api/fastapi/middleware/idempotency.py

@@ -1,81 +1,206 @@
+import base64
 import hashlib
+import json
 import time
-from typing import Annotated
+from typing import Annotated, Optional
 
 from fastapi import Header, HTTPException, Request
-from starlette.middleware.base import BaseHTTPMiddleware
-from starlette.responses import Response
+from starlette.types import ASGIApp, Receive, Scope, Send
 
+from .idempotency_store import IdempotencyStore, InMemoryIdempotencyStore
 
-class IdempotencyMiddleware(BaseHTTPMiddleware):
-    def __init__(self, app, ttl_seconds: int = 24 * 3600, store=None):
-        super().__init__(app)
+
+class IdempotencyMiddleware:
+    """
+    Pure ASGI idempotency middleware.
+
+    Caches responses for requests with Idempotency-Key header to ensure
+    duplicate requests return the same response. Use skip_paths for endpoints
+    where idempotency caching is not appropriate (e.g., streaming responses).
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        ttl_seconds: int = 24 * 3600,
+        store: Optional[IdempotencyStore] = None,
+        header_name: str = "Idempotency-Key",
+        skip_paths: Optional[list[str]] = None,
+    ):
+        self.app = app
         self.ttl = ttl_seconds
-        self.store = store or {}  # replace with Redis
-
-    def _cache_key(self, request, idkey: str):
-        body = getattr(request, "_body", None)
-        if body is None:
-            body = b""
-
-            async def _read():
-                data = await request.body()
-                request._body = data  # stash for downstream
-                return data
-
-            # read once
-            # note: starlette Request is awaitable; we read in dispatch below
-
-        sig = hashlib.sha256(
-            (
-                request.method + "|" + request.url.path + "|" + idkey + "|" + (request._body or b"")
-            ).encode()
-            if isinstance(request._body, str)
-            else (request.method + "|" + request.url.path + "|" + idkey).encode()
-            + (request._body or b"")
-        ).hexdigest()
+        self.store: IdempotencyStore = store or InMemoryIdempotencyStore()
+        self.header_name = header_name.lower()
+        self.skip_paths = skip_paths or []
+
+    def _cache_key(self, method: str, path: str, idkey: str) -> str:
+        sig = hashlib.sha256((method + "|" + path + "|" + idkey).encode()).hexdigest()
         return f"idmp:{sig}"
 
-    async def dispatch(self, request, call_next):
-        if request.method in {"POST", "PATCH", "DELETE"}:
-            # read & buffer body once
-            body = await request.body()
-            request._body = body
-            idkey = request.headers.get("Idempotency-Key")
-            if idkey:
-                k = self._cache_key(request, idkey)
-                entry = self.store.get(k)
-                now = time.time()
-                if entry and entry["exp"] > now:
-                    cached = entry["resp"]
-                    return Response(
-                        content=cached["body"],
-                        status_code=cached["status"],
-                        headers=cached["headers"],
-                        media_type=cached.get("media_type"),
-                    )
-                resp = await call_next(request)
-                # cache only 2xx/201 responses
-                if 200 <= resp.status_code < 300:
-                    body_bytes = b"".join([section async for section in resp.body_iterator])
-                    headers = dict(resp.headers)
-                    self.store[k] = {
-                        "resp": {
-                            "status": resp.status_code,
-                            "body": body_bytes,
-                            "headers": headers,
-                            "media_type": resp.media_type,
-                        },
-                        "exp": now + self.ttl,
-                    }
-                    return Response(
-                        content=body_bytes,
-                        status_code=resp.status_code,
-                        headers=headers,
-                        media_type=resp.media_type,
-                    )
-                return resp
-        return await call_next(request)
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope.get("type") != "http":
+            await self.app(scope, receive, send)
+            return
+
+        path = scope.get("path", "")
+        method = scope.get("method", "GET")
+
+        # Skip specified paths
+        if any(skip in path for skip in self.skip_paths):
+            await self.app(scope, receive, send)
+            return
+
+        # Only apply to mutating methods
+        if method not in {"POST", "PATCH", "DELETE"}:
+            await self.app(scope, receive, send)
+            return
+
+        # Get idempotency key from headers
+        headers = {k.decode().lower(): v.decode() for k, v in scope.get("headers", [])}
+        idkey = headers.get(self.header_name)
+
+        if not idkey:
+            # No idempotency key - pass through
+            await self.app(scope, receive, send)
+            return
+
+        # Buffer the request body
+        body_parts = []
+        while True:
+            message = await receive()
+            if message["type"] == "http.request":
+                body_parts.append(message.get("body", b"") or b"")
+                if not message.get("more_body", False):
+                    break
+            elif message["type"] == "http.disconnect":
+                break
+        body = b"".join(body_parts)
+
+        k = self._cache_key(method, path, idkey)
+        now = time.time()
+        req_hash = hashlib.sha256(body).hexdigest()
+
+        existing = self.store.get(k)
+        if existing and existing.exp > now:
+            # If payload mismatches, return conflict
+            if existing.req_hash and existing.req_hash != req_hash:
+                await self._send_json_response(
+                    send,
+                    409,
+                    {
+                        "type": "about:blank",
+                        "title": "Conflict",
+                        "detail": "Idempotency-Key re-used with different request payload.",
+                    },
+                )
+                return
+            # If response cached and payload matches, replay it
+            if existing.status is not None and existing.body_b64 is not None:
+                await self._send_cached_response(send, existing)
+                return
+
+        # Claim the key
+        exp = now + self.ttl
+        created = self.store.set_initial(k, req_hash, exp)
+        if not created:
+            existing = self.store.get(k)
+            if existing and existing.req_hash and existing.req_hash != req_hash:
+                await self._send_json_response(
+                    send,
+                    409,
+                    {
+                        "type": "about:blank",
+                        "title": "Conflict",
+                        "detail": "Idempotency-Key re-used with different request payload.",
+                    },
+                )
+                return
+            if (
+                existing
+                and existing.status is not None
+                and existing.body_b64 is not None
+            ):
+                await self._send_cached_response(send, existing)
+                return
+
+        # Create a replay receive that returns buffered body
+        # IMPORTANT: After replaying the body, we must forward to original receive()
+        # so that Starlette's listen_for_disconnect can properly detect client disconnects.
+        # This is required for streaming responses on ASGI spec < 2.4.
+        body_sent = False
+
+        async def replay_receive():
+            nonlocal body_sent
+            if not body_sent:
+                body_sent = True
+                return {"type": "http.request", "body": body, "more_body": False}
+            # After body is sent, forward to original receive for disconnect detection
+            return await receive()
+
+        # Capture response for caching
+        response_started = False
+        response_status = 0
+        response_headers: list = []
+        response_body_parts = []
+
+        async def capture_send(message):
+            nonlocal response_started, response_status, response_headers
+            if message["type"] == "http.response.start":
+                response_started = True
+                response_status = message.get("status", 200)
+                response_headers = list(message.get("headers", []))
+            elif message["type"] == "http.response.body":
+                body_chunk = message.get("body", b"")
+                if body_chunk:
+                    response_body_parts.append(body_chunk)
+            await send(message)
+
+        await self.app(scope, replay_receive, capture_send)
+
+        # Cache successful responses
+        if 200 <= response_status < 300:
+            response_body = b"".join(response_body_parts)
+            headers_dict = {k.decode(): v.decode() for k, v in response_headers}
+            media_type = headers_dict.get("content-type", "application/octet-stream")
+            self.store.set_response(
+                k,
+                status=response_status,
+                body=response_body,
+                headers=headers_dict,
+                media_type=media_type,
+            )
+
+    async def _send_json_response(self, send, status: int, content: dict) -> None:
+        body = json.dumps(content).encode("utf-8")
+        await send(
+            {
+                "type": "http.response.start",
+                "status": status,
+                "headers": [(b"content-type", b"application/json")],
+            }
+        )
+        await send({"type": "http.response.body", "body": body, "more_body": False})
+
+    async def _send_cached_response(self, send, existing) -> None:
+        headers = [
+            (k.encode(), v.encode()) for k, v in (existing.headers or {}).items()
+        ]
+        if existing.media_type:
+            headers.append((b"content-type", existing.media_type.encode()))
+        await send(
+            {
+                "type": "http.response.start",
+                "status": existing.status,
+                "headers": headers,
+            }
+        )
+        await send(
+            {
+                "type": "http.response.body",
+                "body": base64.b64decode(existing.body_b64),
+                "more_body": False,
+            }
+        )
 
 
 async def require_idempotency_key(
@@ -83,4 +208,6 @@ async def require_idempotency_key(
     request: Request,
 ) -> None:
     if not idempotency_key.strip():
-        raise HTTPException(status_code=400, detail="Idempotency-Key must not be empty.")
+        raise HTTPException(
+            status_code=400, detail="Idempotency-Key must not be empty."
+        )