diff --git a/src/main.py b/src/main.py index 4b7b06e..ca8a46c 100644 --- a/src/main.py +++ b/src/main.py @@ -41,8 +41,7 @@ from src.services.otel_config import setup_otel from src.services.permission_service import initialize_permission_service from src.middleware.evo_auth import EvoAuthMiddleware -from src.evo_extension_points import runtime_context -from starlette.requests import Request as StarletteRequest +from src.middleware.tenant_context import TenantContextMiddleware from fastapi.exceptions import HTTPException from src.core.exceptions import BaseAPIException from src.core.exception_handlers import ( @@ -105,73 +104,7 @@ app.add_exception_handler(BaseAPIException, base_api_exception_handler) -class TenantContextMiddleware: - """Bind the per-request operational context (tenant) via the runtime_context - extension point, so downstream lazy DB transactions see the correct tenant GUC. - - Why a PURE ASGI middleware (NOT BaseHTTPMiddleware): - BaseHTTPMiddleware runs the downstream app in a separate anyio task via a - memory stream. That task boundary means any ContextVar set here (which is - how the enterprise runtime_context impl propagates the tenant, and how - SQLAlchemy's lazy transaction later reads it) would NOT be visible where the - DB session actually opens its connection — the bind would silently no-op. - A pure ASGI middleware calls ``self.app(...)`` in the SAME task, so the - ContextVar set by ``bind_context`` stays live all the way down to the lazy - transaction. This is the same reason the SQLAlchemy/tenant docs warn against - BaseHTTPMiddleware for context propagation. - - Why the ordering (this must run AFTER EvoAuth on the inbound path): - ``current_context_id`` resolves the tenant from ``request.state.user_context``, - which EvoAuthMiddleware populates. In Starlette the LAST ``add_middleware`` is - the OUTERMOST (runs first inbound). So to be MORE INTERNAL than EvoAuth (i.e. - run after it inbound), this middleware is added BEFORE EvoAuth's - ``add_middleware`` call below. Flow: EvoAuth populates user_context first, - then this middleware reads it and binds the tenant. - """ - - def __init__(self, app): - self.app = app - - async def __call__(self, scope, receive, send): - if scope["type"] != "http": - await self.app(scope, receive, send) - return - - try: - request = StarletteRequest(scope, receive) - cid = runtime_context.current_context_id(request) - except Exception as e: - # Fail-open: never drop a request because context resolution blew up. - # Community default has no tenant anyway (single-scope); a failure here - # must degrade to "no tenant bound", not a 500. - logger.warning(f"TenantContextMiddleware: context resolution failed: {e}") - cid = None - - if not cid: - await self.app(scope, receive, send) - return - - # Fail-open on the BIND SETUP only, never on the downstream app. If entering - # ``bind_context`` blows up we fall back to running the app unbound - # (community-default = no tenant) instead of 500-ing. But once the app has - # started, its exceptions must propagate untouched: catching them here would - # both swallow real errors AND risk running the app twice (a second response - # on the same send channel). ``app_started`` guards that boundary. - app_started = False - try: - async with runtime_context.bind_context(cid): - app_started = True - await self.app(scope, receive, send) - except Exception as e: - if app_started: - raise - logger.warning(f"TenantContextMiddleware: bind_context failed: {e}") - await self.app(scope, receive, send) - - -# Tenant-context middleware — added BEFORE EvoAuth so it is MORE INTERNAL -# (runs AFTER EvoAuth inbound, once request.state.user_context exists). See the -# class docstring for the ordering + pure-ASGI rationale. +# Before EvoAuth on purpose: it must run after EvoAuth inbound. app.add_middleware(TenantContextMiddleware) # EvoAuth middleware for authentication and user context diff --git a/src/middleware/tenant_context.py b/src/middleware/tenant_context.py new file mode 100644 index 0000000..efacd49 --- /dev/null +++ b/src/middleware/tenant_context.py @@ -0,0 +1,75 @@ +"""Binds the per-request context resolved by the ``runtime_context`` extension +point around the downstream app.""" + +from http import HTTPStatus + +from starlette.exceptions import HTTPException +from starlette.requests import Request as StarletteRequest +from starlette.responses import JSONResponse + +from src.evo_extension_points import runtime_context +from src.utils.logger import setup_logger + +logger = setup_logger(__name__) + + +class TenantContextMiddleware: + """Pure ASGI, not BaseHTTPMiddleware: that one runs the app in another task, + where a ContextVar set here is invisible to the lazy DB transaction. + Added before EvoAuthMiddleware in main.py (Starlette wraps in reverse) so it + can read ``request.state.user_context``. + """ + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + try: + request = StarletteRequest(scope, receive) + cid = runtime_context.current_context_id(request) + except HTTPException as e: + # A refusal from the extension point is a response, not a failure. + await self._error_response(e)(scope, receive, send) + return + except Exception as e: + # Resolution failure degrades to "no context bound". + logger.warning(f"TenantContextMiddleware: context resolution failed: {e}") + cid = None + + if not cid: + await self.app(scope, receive, send) + return + + # Fail-open only on the bind setup: once the app has started, its + # exceptions propagate (retrying here would run it twice). + app_started = False + try: + async with runtime_context.bind_context(cid): + app_started = True + await self.app(scope, receive, send) + except Exception as e: + if app_started: + raise + logger.warning(f"TenantContextMiddleware: bind_context failed: {e}") + await self.app(scope, receive, send) + + @staticmethod + def _error_response(exc: HTTPException) -> JSONResponse: + try: + phrase = HTTPStatus(exc.status_code).phrase + except ValueError: # non-standard code such as 499 + phrase = "Error" + return JSONResponse( + status_code=exc.status_code, + content={ + "error": phrase, + "code": "ERR_" + phrase.upper().replace(" ", "_"), + "message": str(exc.detail), + }, + headers=getattr(exc, "headers", None), + ) + diff --git a/tests/unit/test_tenant_context_middleware.py b/tests/unit/test_tenant_context_middleware.py new file mode 100644 index 0000000..fc9a953 --- /dev/null +++ b/tests/unit/test_tenant_context_middleware.py @@ -0,0 +1,97 @@ +"""TenantContextMiddleware: a refusal raised by the runtime_context extension +point reaches the client; a plain resolution failure still degrades to "no +tenant bound" instead of dropping the request.""" + +from __future__ import annotations + +from contextlib import asynccontextmanager + +import pytest +from fastapi import FastAPI, HTTPException +from fastapi.testclient import TestClient + +from src import evo_extension_points +from src.middleware.tenant_context import TenantContextMiddleware + + +class _Context: + def __init__(self, resolve): + self._resolve = resolve + self.bound: list[str] = [] + + def current_context_id(self, source): + return self._resolve() + + def with_context(self, context_id, fn): + return fn() + + @asynccontextmanager + async def bind_context(self, context_id): + self.bound.append(context_id) + yield + + +@pytest.fixture(autouse=True) +def _reset_runtime_context(): + evo_extension_points.reset("runtime_context") + yield + evo_extension_points.reset("runtime_context") + + +def build_client(resolve): + app = FastAPI() + + @app.get("/ping") + async def ping(): + return {"ok": True} + + app.add_middleware(TenantContextMiddleware) + ctx = _Context(resolve) + evo_extension_points.replace("runtime_context", ctx) + return TestClient(app), ctx + + +def test_http_exception_from_the_extension_point_is_returned_as_is(): + def refuse(): + raise HTTPException(status_code=403, detail="user is not a member of tenant") + + client, ctx = build_client(refuse) + response = client.get("/ping") + + assert response.status_code == 403 + assert response.json() == { + "error": "Forbidden", + "code": "ERR_FORBIDDEN", + "message": "user is not a member of tenant", + } + assert ctx.bound == [], "a refused request must not reach bind_context" + + +def test_resolution_failure_still_runs_the_request_unbound(): + def blow_up(): + raise RuntimeError("membership lookup unavailable") + + client, ctx = build_client(blow_up) + response = client.get("/ping") + + assert response.status_code == 200 + assert ctx.bound == [] + + +def test_resolved_context_is_bound_around_the_request(): + client, ctx = build_client(lambda: "ctx-a") + response = client.get("/ping") + + assert response.status_code == 200 + assert ctx.bound == ["ctx-a"] + + +def test_non_standard_status_code_is_still_returned(): + def refuse(): + raise HTTPException(status_code=499, detail="client closed request") + + client, _ = build_client(refuse) + response = client.get("/ping") + + assert response.status_code == 499 + assert response.json()["code"] == "ERR_ERROR"