Accept access tokens on protected routes; persist OIDC state

- Add transitional @jwt_access_required (accepts access tokens, and for now
  also refresh tokens) and use it on every protected route and inside
  admin_required. /api/user/refresh stays refresh-only.
- Move pending OIDC state/PKCE verifiers from an in-process dict to an
  oidc_states table (10 min TTL, single-use) so logins survive a blue/green
  flip. Drop the unused redirect_after_login.
- Unit tests for the decorator and the state store.

Deploy this commit on its own, before the access-token frontend: the old
backend answers access tokens with 422, so the new frontend only works once
this backend is live on both colors.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

blueprints/beli/__init__.py

@@ -3,7 +3,9 @@ import logging
 import os
 
 from quart import Blueprint, jsonify, request
-from quart_jwt_extended import get_jwt_identity, jwt_refresh_token_required
+from quart_jwt_extended import get_jwt_identity
+
+from blueprints.users.decorators import jwt_access_required
 
 from .logic import (
     MAX_PROFILES_PER_USER,
@@ -60,7 +62,7 @@ def _serialize_restaurant(r: BeliRestaurant) -> dict:
 
 
 @beli_blueprint.route("/profiles", methods=["GET"])
-@jwt_refresh_token_required
+@jwt_access_required
 async def list_profiles():
     user_id = get_jwt_identity()
     profiles = await BeliProfile.filter(user_id=user_id).order_by("created_at").all()
@@ -68,7 +70,7 @@ async def list_profiles():
 
 
 @beli_blueprint.route("/profiles", methods=["POST"])
-@jwt_refresh_token_required
+@jwt_access_required
 async def create_profile():
     user_id = get_jwt_identity()
 
@@ -112,7 +114,7 @@ async def create_profile():
 
 
 @beli_blueprint.route("/profiles/<profile_id>", methods=["DELETE"])
-@jwt_refresh_token_required
+@jwt_access_required
 async def delete_profile(profile_id: str):
     user_id = get_jwt_identity()
     profile = await BeliProfile.get_or_none(id=profile_id, user_id=user_id)
@@ -125,7 +127,7 @@ async def delete_profile(profile_id: str):
 
 
 @beli_blueprint.route("/profiles/<profile_id>/refresh", methods=["POST"])
-@jwt_refresh_token_required
+@jwt_access_required
 async def refresh_profile_endpoint(profile_id: str):
     user_id = get_jwt_identity()
     profile = await BeliProfile.get_or_none(id=profile_id, user_id=user_id)
@@ -137,7 +139,7 @@ async def refresh_profile_endpoint(profile_id: str):
 
 
 @beli_blueprint.route("/restaurants", methods=["GET"])
-@jwt_refresh_token_required
+@jwt_access_required
 async def list_restaurants():
     user_id = get_jwt_identity()
     profile_ids = [
@@ -161,7 +163,7 @@ async def list_restaurants():
 
 
 @beli_blueprint.route("/profiles/<profile_id>/refresh-friends", methods=["POST"])
-@jwt_refresh_token_required
+@jwt_access_required
 async def refresh_friends_endpoint(profile_id: str):
     user_id = get_jwt_identity()
     profile = await BeliProfile.get_or_none(id=profile_id, user_id=user_id)

blueprints/conversation/__init__.py

@@ -7,12 +7,10 @@ import time
 import uuid
 
 from quart import Blueprint, jsonify, make_response, request
-from quart_jwt_extended import (
-    get_jwt_identity,
-    jwt_refresh_token_required,
-)
+from quart_jwt_extended import get_jwt_identity
 
 import blueprints.users.models
+from blueprints.users.decorators import jwt_access_required
 from metrics import count_result, observe_response, time_response
 from utils.image_process import image_content_block
 from utils.strip_markdown import extract_agent_text
@@ -126,7 +124,7 @@ def _build_user_message(
 
 
 @conversation_blueprint.post("/query")
-@jwt_refresh_token_required
+@jwt_access_required
 async def query():
     current_user_uuid = get_jwt_identity()
     user = await blueprints.users.models.User.get(id=current_user_uuid)
@@ -170,7 +168,7 @@ async def query():
 
 
 @conversation_blueprint.post("/upload-image")
-@jwt_refresh_token_required
+@jwt_access_required
 async def upload_image():
     current_user_uuid = get_jwt_identity()
     await blueprints.users.models.User.get(id=current_user_uuid)
@@ -202,14 +200,14 @@ async def upload_image():
 
 
 @conversation_blueprint.get("/image/<path:image_key>")
-@jwt_refresh_token_required
+@jwt_access_required
 async def serve_image(image_key: str):
     url = await s3_presigned_url(image_key)
     return jsonify({"url": url})
 
 
 @conversation_blueprint.post("/stream-query")
-@jwt_refresh_token_required
+@jwt_access_required
 async def stream_query():
     current_user_uuid = get_jwt_identity()
     user = await blueprints.users.models.User.get(id=current_user_uuid)
@@ -448,7 +446,7 @@ async def stream_query():
 
 
 @conversation_blueprint.route("/<conversation_id>")
-@jwt_refresh_token_required
+@jwt_access_required
 async def get_conversation(conversation_id: str):
     user_uuid = get_jwt_identity()
     conversation = await _get_owned_conversation(conversation_id, user_uuid)
@@ -483,7 +481,7 @@ async def get_conversation(conversation_id: str):
 
 
 @conversation_blueprint.post("/")
-@jwt_refresh_token_required
+@jwt_access_required
 async def create_conversation():
     user_uuid = get_jwt_identity()
     user = await blueprints.users.models.User.get(id=user_uuid)
@@ -497,7 +495,7 @@ async def create_conversation():
 
 
 @conversation_blueprint.get("/")
-@jwt_refresh_token_required
+@jwt_access_required
 async def get_all_conversations():
     user_uuid = get_jwt_identity()
     user = await blueprints.users.models.User.get(id=user_uuid)
@@ -523,7 +521,7 @@ async def get_all_conversations():
 
 
 @conversation_blueprint.patch("/<conversation_id>")
-@jwt_refresh_token_required
+@jwt_access_required
 async def rename_conversation(conversation_id: str):
     user_uuid = get_jwt_identity()
     body = await request.get_json()
@@ -543,7 +541,7 @@ async def rename_conversation(conversation_id: str):
 
 
 @conversation_blueprint.delete("/<conversation_id>")
-@jwt_refresh_token_required
+@jwt_access_required
 async def delete_conversation(conversation_id: str):
     user_uuid = get_jwt_identity()
     conversation = await _get_owned_conversation(conversation_id, user_uuid)

blueprints/rag/__init__.py

@@ -1,5 +1,4 @@
 from quart import Blueprint, jsonify
-from quart_jwt_extended import jwt_refresh_token_required
 
 from .logic import (
     delete_all_documents,
@@ -7,13 +6,13 @@ from .logic import (
     index_documents,
     sync_obsidian_documents,
 )
-from blueprints.users.decorators import admin_required
+from blueprints.users.decorators import admin_required, jwt_access_required
 
 rag_blueprint = Blueprint("rag_api", __name__, url_prefix="/api/rag")
 
 
 @rag_blueprint.get("/stats")
-@jwt_refresh_token_required
+@jwt_access_required
 async def get_stats():
     """Get vector store statistics."""
     stats = get_vector_store_stats()

blueprints/users/__init__.py

@@ -7,7 +7,8 @@ from quart_jwt_extended import (
 )
 from .models import User
 from .oidc_service import OIDCUserService
-from .decorators import admin_required
+from .oidc_state import consume_oidc_state, save_oidc_state
+from .decorators import admin_required, jwt_access_required
 from config.oidc_config import oidc_config
 import os
 import secrets
@@ -19,10 +20,6 @@ import base64
 
 user_blueprint = Blueprint("user_api", __name__, url_prefix="/api/user")
 
-# In-memory storage for OIDC state/PKCE (production: use Redis or database)
-# Format: {state: {"pkce_verifier": str, "redirect_after_login": str}}
-_oidc_sessions = {}
-
 
 @user_blueprint.route("/oidc/login", methods=["GET"])
 async def oidc_login():
@@ -47,11 +44,8 @@ async def oidc_login():
         # Generate state for CSRF protection
         state = secrets.token_urlsafe(32)
 
-        # Store PKCE verifier and state for callback validation
-        _oidc_sessions[state] = {
-            "pkce_verifier": code_verifier,
-            "redirect_after_login": request.args.get("redirect", "/"),
-        }
+        # Persist PKCE verifier keyed by state for callback validation
+        await save_oidc_state(state, code_verifier)
 
         # Get authorization endpoint from discovery
         discovery = await oidc_config.get_discovery_document()
@@ -93,12 +87,10 @@ async def oidc_callback():
         return jsonify({"error": "Missing code or state"}), 400
 
     # Validate state and retrieve PKCE verifier
-    session = _oidc_sessions.pop(state, None)
-    if not session:
+    pkce_verifier = await consume_oidc_state(state)
+    if not pkce_verifier:
         return jsonify({"error": "Invalid or expired state"}), 400
 
-    pkce_verifier = session["pkce_verifier"]
-
     # Exchange authorization code for tokens
     discovery = await oidc_config.get_discovery_document()
     token_endpoint = discovery.get("token_endpoint")
@@ -206,7 +198,7 @@ async def login():
 
 
 @user_blueprint.route("/me", methods=["GET"])
-@jwt_refresh_token_required
+@jwt_access_required
 async def me():
     user_id = get_jwt_identity()
     user = await User.get_or_none(id=user_id)

blueprints/users/decorators.py

@@ -4,10 +4,47 @@ Authentication decorators for role-based access control.
 
 from functools import wraps
 from quart import jsonify
-from quart_jwt_extended import jwt_refresh_token_required, get_jwt_identity
+from quart_jwt_extended import get_jwt_identity
+from quart_jwt_extended.exceptions import WrongTokenError
+from quart_jwt_extended.view_decorators import (
+    verify_jwt_in_request,
+    verify_jwt_refresh_token_in_request,
+)
 from .models import User
 
 
+def jwt_access_required(fn):
+    """
+    Protect an endpoint with the (1h) access token.
+
+    TRANSITIONAL: this also accepts the refresh token, so an old frontend
+    (which sends the refresh token on every call) keeps working against this
+    backend, including tabs that stay open across deploys.
+
+    It does NOT make a new frontend work against an older backend: routes
+    still on `@jwt_refresh_token_required` answer an access token with 422,
+    which the frontend does not retry. So this backend must be live on both
+    colors before the access-token frontend ships (deploy in two releases).
+
+    TODO: once the access-token frontend has been live for a while, replace
+    the body with `await verify_jwt_in_request()` (i.e. plain `@jwt_required`)
+    and map WrongTokenError to 401, so a stale tab still sending the refresh
+    token refreshes and retries instead of failing with 422.
+    """
+
+    @wraps(fn)
+    async def wrapper(*args, **kwargs):
+        try:
+            await verify_jwt_in_request()
+        except WrongTokenError:
+            # Legacy client sent the refresh token. Expired/invalid tokens
+            # still raise from verify_jwt_in_request and are not retried here.
+            await verify_jwt_refresh_token_in_request()
+        return await fn(*args, **kwargs)
+
+    return wrapper
+
+
 def admin_required(fn):
     """
     Decorator that requires the user to be an admin (member of lldap_admin group).
@@ -15,7 +52,7 @@ def admin_required(fn):
     """
 
     @wraps(fn)
-    @jwt_refresh_token_required
+    @jwt_access_required
     async def wrapper(*args, **kwargs):
         user_id = get_jwt_identity()
         user = await User.get_or_none(id=user_id)

blueprints/users/models.py

@@ -56,3 +56,19 @@ class User(Model):
         if not self.password:
             return False
         return bcrypt.checkpw(plain_password.encode("utf-8"), self.password)
+
+
+class OIDCState(Model):
+    """Pending OIDC login: maps the `state` param to its PKCE verifier.
+
+    Stored in Postgres (not process memory) so a login that starts on one
+    blue/green color and finishes on the other still validates.
+    """
+
+    state = fields.CharField(max_length=128, primary_key=True)
+    pkce_verifier = fields.CharField(max_length=128)
+    expires_at = fields.DatetimeField(index=True)
+    created_at = fields.DatetimeField(auto_now_add=True)
+
+    class Meta:
+        table = "oidc_states"

blueprints/users/oidc_state.py

@@ -0,0 +1,36 @@
+"""
+Persistent, single-use store for pending OIDC logins (state -> PKCE verifier).
+
+Lives in Postgres so a login started on one blue/green color can finish on
+the other.
+"""
+
+from datetime import datetime, timedelta, timezone
+
+from .models import OIDCState
+
+OIDC_STATE_TTL = timedelta(minutes=10)
+
+
+async def save_oidc_state(
+    state: str, pkce_verifier: str, ttl: timedelta = OIDC_STATE_TTL
+) -> None:
+    now = datetime.now(timezone.utc)
+    # Opportunistic cleanup of abandoned logins
+    await OIDCState.filter(expires_at__lte=now).delete()
+    await OIDCState.create(
+        state=state, pkce_verifier=pkce_verifier, expires_at=now + ttl
+    )
+
+
+async def consume_oidc_state(state: str) -> str | None:
+    """Return the PKCE verifier for `state` and delete it, or None if the
+    state is unknown, expired, or was already consumed."""
+    row = await OIDCState.get_or_none(state=state)
+    if row is None:
+        return None
+    # The delete is the single-use guard: only one concurrent caller wins.
+    deleted = await OIDCState.filter(state=state).delete()
+    if not deleted or row.expires_at <= datetime.now(timezone.utc):
+        return None
+    return row.pkce_verifier