""" infra.api.request_context — request-scoped correlation IDs (Stage 84). Production V1 polish. Before this stage, a customer reporting "500 error around 14:23" left the operator grep'ing logs by timestamp and hoping no one else was sharing that window. With correlation IDs every log line emitted during a request carries the same ``request_id``; the operator pastes it from the X-Request-ID response header and gets everything that happened on that request — middleware, auth, handler, DB queries, the exception trace. Design: * ``request_id`` and ``tenant_id`` live in ``contextvars.ContextVar`` so they're naturally scoped per request (asyncio task) without threading any explicit context object through every call. * The X-Request-ID middleware accepts an upstream header value if it matches a strict allowlist regex (length + charset), so a hostile peer can't inject 10KB into our log lines. Otherwise we generate a 16-char UUID hex. * The auth dependency calls ``set_tenant_context(tid)`` after a successful lookup, so handler logs include the tenant_id too. * ``JsonLogFormatter`` (in ``infra.deployment.observability``) reads these contextvars at format time and adds them to the payload — no code change needed at the call site of every ``logger.info(...)``. Stdlib only. """ from __future__ import annotations import re import uuid from contextvars import ContextVar from typing import Mapping, Optional # Per-request context. Default None when no request is in flight # (background tasks, scheduler ticks) — the formatter omits the # field rather than emitting "null". request_id_var: ContextVar[Optional[str]] = ContextVar( "orgstate_request_id", default=None, ) tenant_id_var: ContextVar[Optional[str]] = ContextVar( "orgstate_tenant_id", default=None, ) # Stage 154 — request cookies (read-only snapshot) stamped by # the request-context middleware so auth helpers that don't take # a FastAPI Request param (the legacy require_tenant_or_admin # pattern) can still see the SSO session cookie. Default is an # empty mapping so off-request callers (CLI, scheduler) just get # a no-op when they look up a cookie name. request_cookies_var: ContextVar[Mapping[str, str]] = ContextVar( "orgstate_request_cookies", default={}, ) # Allow upstream-supplied X-Request-ID when it looks sane: # alphanumerics + dash + underscore + dot, up to 128 chars. The # format covers UUIDs (with or without dashes), Cloudflare-style # CF-RAY values, and arbitrary trace IDs. Anything else is # discarded; we generate fresh. _VALID_REQUEST_ID = re.compile(r"^[\w.\-]{1,128}$") def generate_request_id() -> str: """16 hex chars — short enough to copy/paste in support tickets, big enough to be unique within a process lifetime. Full UUID would be 32 chars; we trim for ergonomics. Birthday collision after ~2^32 requests; we expect to roll well before that and a collision wouldn't corrupt anything, just confuse log grep.""" return uuid.uuid4().hex[:16] def sanitize_upstream_id(raw: Optional[str]) -> Optional[str]: """Return the upstream X-Request-ID iff it matches the allowlist. Otherwise None (caller generates fresh).""" if raw is None: return None raw = raw.strip() if not raw: return None if not _VALID_REQUEST_ID.match(raw): return None return raw def set_request_context(request_id: str, tenant_id: Optional[str] = None) -> tuple: """Set both contextvars. Returns the two tokens so caller can reset on response (FastAPI middleware does this via the contextvar's task-local nature anyway, but tests + non-async callers need the tokens).""" tok_req = request_id_var.set(request_id) tok_tid = tenant_id_var.set(tenant_id) return tok_req, tok_tid def set_tenant_context(tenant_id: Optional[str]) -> None: """Called from the auth dependency once the tenant is known. The request_id is already in context from the middleware.""" tenant_id_var.set(tenant_id) def clear_request_context(tokens: Optional[tuple] = None) -> None: """Reset both vars. Pass the tokens returned by ``set_request_context`` for a true reset, or call with no args to set both to None (good enough for sync teardown).""" if tokens is None: request_id_var.set(None) tenant_id_var.set(None) return tok_req, tok_tid = tokens request_id_var.reset(tok_req) tenant_id_var.reset(tok_tid) def current_request_id() -> Optional[str]: return request_id_var.get() def current_tenant_id() -> Optional[str]: return tenant_id_var.get() def set_request_cookies(cookies: Mapping[str, str]) -> None: """Stamp the incoming request's cookies on the contextvar so legacy auth helpers (require_key / require_tenant_or_admin) can read the SSO session cookie without a Request param.""" request_cookies_var.set(cookies) def current_request_cookies() -> Mapping[str, str]: return request_cookies_var.get()