svc-infra 0.1.595__py3-none-any.whl → 1.1.0__py3-none-any.whl

svc_infra/api/fastapi/middleware/idempotency.py

@@ -1,81 +1,203 @@
+import base64
 import hashlib
+import json
 import time
 from typing import Annotated
 
 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).
+
+    Matching uses prefix matching: "/v1/chat" matches "/v1/chat", "/v1/chat/stream",
+    but not "/api/v1/chat" or "/v1/chatter".
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        ttl_seconds: int = 24 * 3600,
+        store: IdempotencyStore | None = None,
+        header_name: str = "Idempotency-Key",
+        skip_paths: list[str] | None = 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 using prefix matching
+        if any(path.startswith(skip) 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(

svc_infra/api/fastapi/middleware/idempotency_store.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+import base64
+import json
+import time
+from dataclasses import dataclass
+from typing import Protocol
+
+
+@dataclass
+class IdempotencyEntry:
+    req_hash: str
+    exp: float
+    # Optional response fields when available
+    status: int | None = None
+    body_b64: str | None = None
+    headers: dict[str, str] | None = None
+    media_type: str | None = None
+
+
+class IdempotencyStore(Protocol):
+    def get(self, key: str) -> IdempotencyEntry | None:
+        pass
+
+    def set_initial(self, key: str, req_hash: str, exp: float) -> bool:
+        """Atomically create an entry if absent. Returns True if created, False if already exists."""
+        pass
+
+    def set_response(
+        self,
+        key: str,
+        *,
+        status: int,
+        body: bytes,
+        headers: dict[str, str],
+        media_type: str | None,
+    ) -> None:
+        pass
+
+    def delete(self, key: str) -> None:
+        pass
+
+
+class InMemoryIdempotencyStore:
+    def __init__(self):
+        self._store: dict[str, IdempotencyEntry] = {}
+
+    def get(self, key: str) -> IdempotencyEntry | None:
+        entry = self._store.get(key)
+        if not entry:
+            return None
+        # expire lazily
+        if entry.exp <= time.time():
+            self._store.pop(key, None)
+            return None
+        return entry
+
+    def set_initial(self, key: str, req_hash: str, exp: float) -> bool:
+        now = time.time()
+        existing = self._store.get(key)
+        if existing and existing.exp > now:
+            return False
+        self._store[key] = IdempotencyEntry(req_hash=req_hash, exp=exp)
+        return True
+
+    def set_response(
+        self,
+        key: str,
+        *,
+        status: int,
+        body: bytes,
+        headers: dict[str, str],
+        media_type: str | None,
+    ) -> None:
+        entry = self._store.get(key)
+        if not entry:
+            # Create if missing to ensure replay works until exp
+            entry = IdempotencyEntry(req_hash="", exp=time.time() + 60)
+            self._store[key] = entry
+        entry.status = status
+        entry.body_b64 = base64.b64encode(body).decode()
+        entry.headers = dict(headers)
+        entry.media_type = media_type
+
+    def delete(self, key: str) -> None:
+        self._store.pop(key, None)
+
+
+class RedisIdempotencyStore:
+    """A simple Redis-backed store.
+
+    Notes:
+        - Uses GET/SET with JSON payload; initial claim uses SETNX semantics.
+        - Not fully atomic for response update; sufficient for basic dedupe.
+        - For strict guarantees, replace with a Lua script (future improvement).
+    """
+
+    def __init__(self, redis_client, *, prefix: str = "idmp"):
+        self.r = redis_client
+        self.prefix = prefix
+
+    def _k(self, key: str) -> str:
+        return f"{self.prefix}:{key}"
+
+    def get(self, key: str) -> IdempotencyEntry | None:
+        raw = self.r.get(self._k(key))
+        if not raw:
+            return None
+        try:
+            data = json.loads(raw)
+        except Exception:
+            return None
+        entry = IdempotencyEntry(
+            req_hash=data.get("req_hash", ""),
+            exp=float(data.get("exp", 0)),
+            status=data.get("status"),
+            body_b64=data.get("body_b64"),
+            headers=data.get("headers"),
+            media_type=data.get("media_type"),
+        )
+        if entry.exp <= time.time():
+            try:
+                self.r.delete(self._k(key))
+            except Exception:
+                pass
+            return None
+        return entry
+
+    def set_initial(self, key: str, req_hash: str, exp: float) -> bool:
+        payload = json.dumps({"req_hash": req_hash, "exp": exp})
+        # Attempt NX set
+        ok = self.r.set(self._k(key), payload, nx=True)
+        # If set, also set TTL (expire at exp)
+        if ok:
+            ttl = max(1, int(exp - time.time()))
+            try:
+                self.r.expire(self._k(key), ttl)
+            except Exception:
+                pass
+            return True
+        # If exists but expired, overwrite
+        entry = self.get(key)
+        if not entry:
+            self.r.set(self._k(key), payload)
+            ttl = max(1, int(exp - time.time()))
+            try:
+                self.r.expire(self._k(key), ttl)
+            except Exception:
+                pass
+            return True
+        return False
+
+    def set_response(
+        self,
+        key: str,
+        *,
+        status: int,
+        body: bytes,
+        headers: dict[str, str],
+        media_type: str | None,
+    ) -> None:
+        entry = self.get(key)
+        if not entry:
+            # default short ttl if missing; caller should have set initial
+            entry = IdempotencyEntry(req_hash="", exp=time.time() + 60)
+        entry.status = status
+        entry.body_b64 = base64.b64encode(body).decode()
+        entry.headers = dict(headers)
+        entry.media_type = media_type
+        ttl = max(1, int(entry.exp - time.time()))
+        payload = json.dumps(
+            {
+                "req_hash": entry.req_hash,
+                "exp": entry.exp,
+                "status": entry.status,
+                "body_b64": entry.body_b64,
+                "headers": entry.headers,
+                "media_type": entry.media_type,
+            }
+        )
+        self.r.set(self._k(key), payload, ex=ttl)
+
+    def delete(self, key: str) -> None:
+        try:
+            self.r.delete(self._k(key))
+        except Exception:
+            pass

svc_infra/api/fastapi/middleware/optimistic_lock.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Annotated, Any
+
+from fastapi import Header, HTTPException
+
+
+async def require_if_match(
+    version: Annotated[str | None, Header(alias="If-Match")] = None,
+) -> str:
+    """Require If-Match header for optimistic locking on mutating operations.
+
+    Returns the header value. Raises 428 if missing.
+    """
+    if not version:
+        raise HTTPException(
+            status_code=428, detail="Missing If-Match header for optimistic locking."
+        )
+    return version
+
+
+def check_version_or_409(get_current_version: Callable[[], Any], provided: str) -> None:
+    """Compare provided version with current version; raise 409 on mismatch.
+
+    - get_current_version: callable returning the resource's current version (int/str)
+    - provided: header value; attempts to coerce to int if current is int
+    """
+    current = get_current_version()
+    p: int | str
+    if isinstance(current, int):
+        try:
+            p = int(provided)
+        except Exception:
+            raise HTTPException(status_code=400, detail="Invalid If-Match value; expected integer.")
+    else:
+        p = provided
+    if p != current:
+        raise HTTPException(status_code=409, detail="Version mismatch (optimistic locking).")

svc_infra/api/fastapi/middleware/ratelimit.py

@@ -1,61 +1,158 @@
+import json
 import time
 
-from starlette.middleware.base import BaseHTTPMiddleware
-from starlette.responses import JSONResponse
+from fastapi import Request
+from starlette.types import ASGIApp, Receive, Scope, Send
 
 from svc_infra.obs.metrics import emit_rate_limited
 
 from .ratelimit_store import InMemoryRateLimitStore, RateLimitStore
 
+try:
+    # Optional import: tenancy may not be enabled in all apps
+    from svc_infra.api.fastapi.tenancy.context import (
+        resolve_tenant_id as _resolve_tenant_id,
+    )
+except Exception:  # pragma: no cover - fallback for minimal builds
+    _resolve_tenant_id = None  # type: ignore[assignment]
+
+
+class SimpleRateLimitMiddleware:
+    """
+    Pure ASGI rate limiting middleware.
+
+    Applies per-key rate limits with configurable windows. Use skip_paths for
+    endpoints that should bypass rate limiting (e.g., health checks, webhooks).
+
+    Matching uses prefix matching: "/v1/chat" matches "/v1/chat", "/v1/chat/stream",
+    but not "/api/v1/chat" or "/v1/chatter".
+    """
 
-class SimpleRateLimitMiddleware(BaseHTTPMiddleware):
     def __init__(
         self,
-        app,
+        app: ASGIApp,
         limit: int = 120,
         window: int = 60,
         key_fn=None,
+        *,
+        # When provided, dynamically computes a limit for the current request (e.g. per-tenant quotas)
+        # Signature: (request: Request, tenant_id: Optional[str]) -> int | None
+        limit_resolver=None,
+        # If True, automatically scopes the bucket key by tenant id when available
+        scope_by_tenant: bool = False,
+        # When True, allows unresolved tenant IDs to fall back to an "X-Tenant-Id" header value.
+        # Disabled by default to avoid trusting arbitrary client-provided headers which could
+        # otherwise be used to evade per-tenant limits when authentication fails.
+        allow_untrusted_tenant_header: bool = False,
         store: RateLimitStore | None = None,
+        skip_paths: list[str] | None = None,
     ):
-        super().__init__(app)
+        self.app = app
         self.limit, self.window = limit, window
-        self.key_fn = key_fn or (lambda r: r.headers.get("X-API-Key") or r.client.host)
+        self.key_fn = key_fn
+        self._limit_resolver = limit_resolver
+        self.scope_by_tenant = scope_by_tenant
+        self._allow_untrusted_tenant_header = allow_untrusted_tenant_header
         self.store = store or InMemoryRateLimitStore(limit=limit)
+        self.skip_paths = skip_paths or []
+
+    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", "")
+
+        # Skip specified paths using prefix matching
+        if any(path.startswith(skip) for skip in self.skip_paths):
+            await self.app(scope, receive, send)
+            return
+
+        # Create a Request object for key extraction and tenant resolution
+        request = Request(scope, receive)
+
+        # Default key function
+        key_fn = self.key_fn or (
+            lambda r: r.headers.get("X-API-Key") or (r.client.host if r.client else "unknown")
+        )
+
+        # Resolve tenant when possible
+        tenant_id = None
+        if self.scope_by_tenant or self._limit_resolver:
+            try:
+                if _resolve_tenant_id is not None:
+                    tenant_id = await _resolve_tenant_id(request)
+            except Exception:
+                tenant_id = None
+            # Fallback header behavior - ONLY if explicitly allowed
+            # Never trust untrusted headers by default to prevent rate limit evasion
+            if not tenant_id and self._allow_untrusted_tenant_header:
+                tenant_id = request.headers.get("X-Tenant-Id") or request.headers.get("X-Tenant-ID")
+
+        key = key_fn(request)
+        if self.scope_by_tenant and tenant_id:
+            key = f"{key}:tenant:{tenant_id}"
+
+        # Allow dynamic limit overrides
+        eff_limit = self.limit
+        if self._limit_resolver:
+            try:
+                v = self._limit_resolver(request, tenant_id)
+                eff_limit = int(v) if v is not None else self.limit
+            except Exception:
+                eff_limit = self.limit
 
-    async def dispatch(self, request, call_next):
-        key = self.key_fn(request)
         now = int(time.time())
-        # Increment counter in store
-        count, limit, reset = self.store.incr(str(key), self.window)
+        count, _store_limit, reset = self.store.incr(str(key), self.window)
+        limit = eff_limit
         remaining = max(0, limit - count)
 
-        if remaining < 0:  # defensive clamp
-            remaining = 0
-
         if count > limit:
+            # Rate limited - return 429
             retry = max(0, reset - now)
             try:
                 emit_rate_limited(str(key), limit, retry)
             except Exception:
                 pass
-            return JSONResponse(
-                status_code=429,
-                content={
+
+            body = json.dumps(
+                {
                     "title": "Too Many Requests",
                     "status": 429,
                     "detail": "Rate limit exceeded.",
                     "code": "RATE_LIMITED",
-                },
-                headers={
-                    "X-RateLimit-Limit": str(limit),
-                    "X-RateLimit-Remaining": "0",
-                    "X-RateLimit-Reset": str(reset),
-                    "Retry-After": str(retry),
-                },
+                }
+            ).encode("utf-8")
+
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": 429,
+                    "headers": [
+                        (b"content-type", b"application/json"),
+                        (b"x-ratelimit-limit", str(limit).encode()),
+                        (b"x-ratelimit-remaining", b"0"),
+                        (b"x-ratelimit-reset", str(reset).encode()),
+                        (b"retry-after", str(retry).encode()),
+                    ],
+                }
             )
+            await send({"type": "http.response.body", "body": body, "more_body": False})
+            return
+
+        # Not rate limited - add headers to response
+        async def send_with_headers(message):
+            if message["type"] == "http.response.start":
+                headers = list(message.get("headers", []))
+                # Add rate limit headers if not already present
+                header_names = {h[0].lower() for h in headers}
+                if b"x-ratelimit-limit" not in header_names:
+                    headers.append((b"x-ratelimit-limit", str(limit).encode()))
+                if b"x-ratelimit-remaining" not in header_names:
+                    headers.append((b"x-ratelimit-remaining", str(remaining).encode()))
+                if b"x-ratelimit-reset" not in header_names:
+                    headers.append((b"x-ratelimit-reset", str(reset).encode()))
+                message = {**message, "headers": headers}
+            await send(message)
 
-        resp = await call_next(request)
-        resp.headers.setdefault("X-RateLimit-Limit", str(limit))
-        resp.headers.setdefault("X-RateLimit-Remaining", str(remaining))
-        resp.headers.setdefault("X-RateLimit-Reset", str(reset))
-        return resp
+        await self.app(scope, receive, send_with_headers)