diff --git a/backend/app/gateway/auth/models.py b/backend/app/gateway/auth/models.py index d8f9b954..1e9f91a7 100644 --- a/backend/app/gateway/auth/models.py +++ b/backend/app/gateway/auth/models.py @@ -31,6 +31,12 @@ class User(BaseModel): needs_setup: bool = Field(default=False, description="True for auto-created admin until setup completes") token_version: int = Field(default=0, description="Incremented on password change to invalidate old JWTs") + # Workspace linkage (Stage 0 PR4) + default_workspace_id: str | None = Field( + default=None, + description="The workspace the user lands in by default after login. NULL → /select-workspace.", + ) + class UserResponse(BaseModel): """Response model for user info endpoint.""" @@ -39,3 +45,18 @@ class UserResponse(BaseModel): email: str system_role: Literal["admin", "user"] needs_setup: bool = False + + +class ActiveWorkspace(BaseModel): + """Lightweight workspace proxy injected into the request-scoped contextvar. + + Implements the structural ``CurrentWorkspace`` protocol expected by + ``deerflow.runtime.workspace_context``: only ``.id`` (str) and + ``.role`` (str) are required. We intentionally do *not* embed the + full ``WorkspaceRow`` here — the middleware needs to set the + contextvar on every request and an extra DB lookup just to populate + a name/slug we don't use yet would be wasted work. + """ + + id: str + role: str diff --git a/backend/app/gateway/auth_middleware.py b/backend/app/gateway/auth_middleware.py index 6b645226..7895efa7 100644 --- a/backend/app/gateway/auth_middleware.py +++ b/backend/app/gateway/auth_middleware.py @@ -17,9 +17,11 @@ from starlette.responses import JSONResponse from starlette.types import ASGIApp from app.gateway.auth.errors import AuthErrorCode, AuthErrorResponse +from app.gateway.auth.models import ActiveWorkspace from app.gateway.authz import _ALL_PERMISSIONS, AuthContext from app.gateway.internal_auth import INTERNAL_AUTH_HEADER_NAME, get_internal_user, is_valid_internal_auth_token from deerflow.runtime.user_context import reset_current_user, set_current_user +from deerflow.runtime.workspace_context import reset_current_workspace, set_current_workspace # Paths that never require authentication. _PUBLIC_PATH_PREFIXES: tuple[str, ...] = ( @@ -119,8 +121,22 @@ class AuthMiddleware(BaseHTTPMiddleware): # JWT-decode + DB-lookup pipeline a second time per request). request.state.user = user request.state.auth = AuthContext(user=user, permissions=_ALL_PERMISSIONS) - token = set_current_user(user) + user_token = set_current_user(user) + + # Inject workspace contextvar from the JWT's wid/role claims. + # decode_token has already rejected legacy no-wid tokens upstream, + # so by the time we get here payload.wid is guaranteed non-None + # for cookie-authenticated requests. Internal-auth requests skip + # the workspace contextvar (they don't have a workspace scope — + # the internal user is a system actor). + ws_token = None + payload = getattr(request.state, "auth_payload", None) + if payload is not None and payload.wid is not None: + ws_token = set_current_workspace(ActiveWorkspace(id=payload.wid, role=payload.role or "owner")) + try: return await call_next(request) finally: - reset_current_user(token) + if ws_token is not None: + reset_current_workspace(ws_token) + reset_current_user(user_token) diff --git a/backend/app/gateway/deps.py b/backend/app/gateway/deps.py index 96ea7c5e..96ddecd1 100644 --- a/backend/app/gateway/deps.py +++ b/backend/app/gateway/deps.py @@ -220,6 +220,10 @@ async def get_current_user_from_request(request: Request): detail=AuthErrorResponse(code=AuthErrorCode.TOKEN_INVALID, message="Token revoked (password changed)").model_dump(), ) + # Stash decoded payload on request.state so AuthMiddleware can read + # wid/role for the workspace contextvar without a second decode. + request.state.auth_payload = payload + return user diff --git a/backend/tests/test_auth_middleware_workspace.py b/backend/tests/test_auth_middleware_workspace.py new file mode 100644 index 00000000..9b208537 --- /dev/null +++ b/backend/tests/test_auth_middleware_workspace.py @@ -0,0 +1,130 @@ +"""AuthMiddleware injects the workspace ContextVar (Stage 0 PR4 T4.7). + +After PR4 every authenticated request has a workspace bound on +``deerflow.runtime.workspace_context._current_workspace``. The middleware +populates it from the JWT's ``wid`` / ``role`` claims, mirrors what it +already does for ``user_context``, and tears both down in a single +``try/finally`` so leaks don't cross requests. + +Legacy 4-field tokens (no ``wid``) are rejected upstream by +``decode_token`` (T4.6) — those should never reach the workspace +injection branch; this file pins that the 401 they trigger carries +``AuthErrorCode.WORKSPACE_REQUIRED``. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from unittest.mock import AsyncMock, patch +from uuid import uuid4 + +import jwt +import pytest +from fastapi import FastAPI +from starlette.testclient import TestClient + +from app.gateway.auth import create_access_token +from app.gateway.auth.config import get_auth_config +from app.gateway.auth.models import User +from app.gateway.auth_middleware import AuthMiddleware +from deerflow.runtime.workspace_context import get_current_workspace + + +@pytest.fixture(autouse=True) +def _stable_jwt_secret(monkeypatch): + monkeypatch.setenv("AUTH_JWT_SECRET", "test-secret-key-for-jwt-testing-minimum-32-chars") + yield + + +def _make_app() -> FastAPI: + """App with AuthMiddleware + an inspect route that surfaces the contextvar.""" + app = FastAPI() + app.add_middleware(AuthMiddleware) + + @app.get("/api/v1/auth/setup-status") # public — never gates on wid + async def setup_status(): + return {"needs_setup": False} + + @app.get("/api/models") # protected — exercises wid injection + async def inspect_workspace(): + ws = get_current_workspace() + if ws is None: + return {"workspace": None} + return {"workspace": {"id": ws.id, "role": ws.role}} + + return app + + +def _make_user(uid: str) -> User: + return User(id=uid, email="t@example.com", password_hash="hash", token_version=0) + + +def _make_legacy_token() -> str: + """Encode a pre-PR4 JWT (no wid/role) directly.""" + now = datetime.now(UTC) + payload = { + "sub": str(uuid4()), + "exp": now + timedelta(hours=1), + "iat": now, + "ver": 0, + } + return jwt.encode(payload, get_auth_config().jwt_secret, algorithm="HS256") + + +def test_new_jwt_injects_workspace_into_contextvar() -> None: + """Cookie with wid+role → route observes the workspace via the contextvar.""" + uid = str(uuid4()) + token = create_access_token(uid, workspace_id="ws-abc", role="owner") + + with patch("app.gateway.deps.get_local_provider") as fn: + fn.return_value.get_user = AsyncMock(return_value=_make_user(uid)) + client = TestClient(_make_app()) + res = client.get("/api/models", cookies={"access_token": token}) + + assert res.status_code == 200, res.text + assert res.json() == {"workspace": {"id": "ws-abc", "role": "owner"}} + + +def test_legacy_jwt_rejected_with_workspace_required() -> None: + """No-wid tokens get 401 with AuthErrorCode.WORKSPACE_REQUIRED, not generic token_invalid.""" + client = TestClient(_make_app()) + res = client.get("/api/models", cookies={"access_token": _make_legacy_token()}) + + assert res.status_code == 401 + assert res.json()["detail"]["code"] == "workspace_required" + + +def test_public_path_skips_workspace_check() -> None: + """Public whitelist (e.g. /api/v1/auth/setup-status) does not require wid.""" + client = TestClient(_make_app()) + res = client.get("/api/v1/auth/setup-status") # no cookie at all + assert res.status_code == 200 + + +def test_workspace_contextvar_resets_between_requests() -> None: + """After dispatch returns the contextvar must be clear (no leak across requests). + + Why we test this: if the try/finally is wired only for user_context but + not workspace_context, two back-to-back requests can see each other's + workspace under asyncio task switching. + """ + uid = str(uuid4()) + token = create_access_token(uid, workspace_id="ws-first", role="owner") + + # First request resolves to ws-first + with patch("app.gateway.deps.get_local_provider") as fn: + fn.return_value.get_user = AsyncMock(return_value=_make_user(uid)) + client = TestClient(_make_app()) + res1 = client.get("/api/models", cookies={"access_token": token}) + assert res1.json() == {"workspace": {"id": "ws-first", "role": "owner"}} + + # Outside the request scope the contextvar must be empty again. + assert get_current_workspace() is None + + # Second request with a different workspace must not see ws-first. + token2 = create_access_token(uid, workspace_id="ws-second", role="owner") + with patch("app.gateway.deps.get_local_provider") as fn: + fn.return_value.get_user = AsyncMock(return_value=_make_user(uid)) + client = TestClient(_make_app()) + res2 = client.get("/api/models", cookies={"access_token": token2}) + assert res2.json() == {"workspace": {"id": "ws-second", "role": "owner"}}