feat(auth): AuthMiddleware injects workspace ContextVar from JWT wid/role

deps.get_current_user_from_request now stashes the decoded payload on
request.state.auth_payload so AuthMiddleware can populate the
workspace_context ContextVar without a second decode. Reset is paired
in the same try/finally as user_context to keep teardown atomic.

Also adds:
- auth.models.ActiveWorkspace — minimal proxy that satisfies the
  CurrentWorkspace protocol (id + role only).
- auth.models.User.default_workspace_id — surfaces the DB column added
  in T4.4 so the eventual /auth/me payload can reference it.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
1445043649
2026-05-12 22:25:43 +08:00
parent b4bef65079
commit 2145d36744
4 changed files with 173 additions and 2 deletions
+21
View File
@@ -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
+18 -2
View File
@@ -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)
+4
View File
@@ -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
@@ -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"}}