Implement stateless JWT authentication for Hugging Face Spaces

- Add PyJWT dependency for JWT token handling
- Create JWT token functions: create_auth_token, verify_auth_token
- Implement stateless authentication using JWT cookies
- Update hf_login_required to use JWT tokens instead of sessions
- Set auth_token cookie on login with 48-hour expiry
- Clear auth_token cookie on logout
- Use very permissive cookie settings for HF Spaces iframe compatibility

This solves the load balancing issue where sessions are not shared
between different servers in Hugging Face Spaces by using stateless
JWT tokens stored in cookies instead of server-side sessions.

app.py

@@ -3,7 +3,7 @@ Chatty - Ethics Chat Application
 A Flask web application that provides a chat interface for the findEthics-Atlas API.
 """
 
-from flask import Flask, render_template, request, jsonify, redirect, url_for, flash, session
+from flask import Flask, render_template, request, jsonify, redirect, url_for, flash, session, make_response
 from flask_login import login_user, logout_user, login_required, current_user
 from flask_wtf.csrf import CSRFProtect
 import requests
@@ -13,6 +13,8 @@ import logging
 from datetime import datetime, timedelta
 from collections import defaultdict
 import time
+import jwt
+import hashlib
 
 from config import get_config, validate_environment
 from auth import init_login_manager, authenticate_user, create_user_account
@@ -69,34 +71,72 @@ logging.basicConfig(
 # Initialize Flask-Login
 init_login_manager(app)
 
+# Stateless authentication functions for Hugging Face Spaces
+def create_auth_token(user_id, email):
+    """Create a JWT token for stateless authentication"""
+    payload = {
+        'user_id': str(user_id),
+        'email': email,
+        'exp': datetime.utcnow() + timedelta(hours=48),  # 48 hour expiry
+        'iat': datetime.utcnow()
+    }
+    return jwt.encode(payload, app.config['SECRET_KEY'], algorithm='HS256')
+
+def verify_auth_token(token):
+    """Verify and decode JWT token"""
+    try:
+        payload = jwt.decode(token, app.config['SECRET_KEY'], algorithms=['HS256'])
+        return payload
+    except jwt.ExpiredSignatureError:
+        app.logger.info("Auth token expired")
+        return None
+    except jwt.InvalidTokenError:
+        app.logger.info("Invalid auth token")
+        return None
+
+def get_current_user_from_token():
+    """Get current user from JWT token in cookie"""
+    token = request.cookies.get('auth_token')
+    if not token:
+        return None
+    
+    payload = verify_auth_token(token)
+    if not payload:
+        return None
+    
+    try:
+        user = User.get_by_id(payload['user_id'])
+        if user and user.email == payload['email']:
+            return user
+    except Exception as e:
+        app.logger.warning(f"Failed to load user from token: {e}")
+    
+    return None
+
 # Custom login_required decorator for Hugging Face Spaces load balancing
 def hf_login_required(f):
     """
-    Custom login_required decorator that handles Hugging Face Spaces load balancing
-    by attempting to restore user session from stored session data
+    Custom login_required decorator that uses stateless JWT authentication
+    for Hugging Face Spaces load balancing compatibility
     """
     @wraps(f)
     def decorated_function(*args, **kwargs):
-        # First check if user is already authenticated
+        # First check if user is already authenticated via Flask-Login
         if current_user.is_authenticated:
             return f(*args, **kwargs)
         
-        # If not authenticated, try to restore from session data (for load balancing)
-        if 'user_id' in session and 'user_email' in session:
-            try:
-                # Try to load user from database
-                user = User.get_by_id(session['user_id'])
-                if user:
-                    # Restore the user session
-                    login_user(user, remember=True)
-                    app.logger.info(f"Restored user session for load balancing: {user.email}")
-                    return f(*args, **kwargs)
-            except Exception as e:
-                app.logger.warning(f"Failed to restore user session: {e}")
+        # Try to authenticate using JWT token from cookie
+        user = get_current_user_from_token()
+        if user:
+            # Restore the user session for this request
+            login_user(user, remember=True)
+            app.logger.info(f"Restored user from JWT token: {user.email}")
+            return f(*args, **kwargs)
         
         # If all else fails, redirect to login
         app.logger.info(f"Authentication required - redirecting to login")
         app.logger.info(f"Session data: {dict(session)}")
+        app.logger.info(f"Auth token present: {'auth_token' in request.cookies}")
         app.logger.info(f"Request IP: {request.environ.get('HTTP_X_FORWARDED_FOR', request.remote_addr)}")
         
         return redirect(url_for('login', next=request.url))
@@ -273,29 +313,45 @@ def login():
             success, user, error_message = authenticate_user(email, password)
             
             if success and user:
-                # Login successful - create session
-                # For Hugging Face Spaces, we need to make sessions permanent to work in iframe
+                # Login successful - create session and JWT token
                 session.permanent = True
-                login_user(user, remember=True)  # Use remember=True for iframe compatibility
+                login_user(user, remember=True)
+                
+                # Create JWT token for stateless authentication across load-balanced servers
+                auth_token = create_auth_token(user.id, user.email)
                 
-                # Store user info in session as backup for load balancing issues
+                # Store user info in session as backup
                 session['user_id'] = str(user.id)
                 session['user_email'] = user.email
                 session['login_time'] = datetime.utcnow().isoformat()
                 
-                # Debug session information
+                # Debug information
                 app.logger.info(f"Session after login_user: {dict(session)}")
                 app.logger.info(f"Current user authenticated: {current_user.is_authenticated}")
                 app.logger.info(f"Current user ID: {getattr(current_user, 'id', 'None')}")
+                app.logger.info(f"JWT token created: {bool(auth_token)}")
                 app.logger.info(f"Request IP: {request.environ.get('HTTP_X_FORWARDED_FOR', request.remote_addr)}")
                 
                 flash('Login successful!', 'success')
                 
-                # Redirect to next page or chat interface
+                # Create response with JWT token cookie
                 next_page = request.args.get('next')
                 if next_page and next_page.startswith('/'):
-                    return redirect(next_page)
-                return redirect(url_for('chat_interface'))
+                    response = make_response(redirect(next_page))
+                else:
+                    response = make_response(redirect(url_for('chat_interface')))
+                
+                # Set JWT token cookie with very permissive settings for HF Spaces
+                response.set_cookie(
+                    'auth_token',
+                    auth_token,
+                    max_age=48*60*60,  # 48 hours
+                    secure=False,      # HF handles HTTPS at proxy level
+                    httponly=False,    # Allow JS access for iframe compatibility
+                    samesite=None      # Most permissive for cross-origin
+                )
+                
+                return response
             else:
                 # Login failed - record attempt and display error
                 record_login_attempt(client_ip)
@@ -322,15 +378,23 @@ def logout():
         # Flash success message
         flash('You have been logged out successfully.', 'info')
         
-        # Redirect to login page
-        return redirect(url_for('login'))
+        # Create response and clear JWT token cookie
+        response = make_response(redirect(url_for('login')))
+        response.set_cookie('auth_token', '', expires=0)  # Clear the cookie
+        
+        return response
         
     except Exception as e:
         # Even if there's an error, still try to log out
         logout_user()
         session.clear()
         flash('Logout completed.', 'info')
-        return redirect(url_for('login'))
+        
+        # Clear JWT token cookie even on error
+        response = make_response(redirect(url_for('login')))
+        response.set_cookie('auth_token', '', expires=0)
+        
+        return response
 
 @app.route('/chat')
 @hf_login_required

requirements.txt

@@ -21,4 +21,5 @@ gradio==5.0.1
 
 # Additional dependencies for production
 gunicorn==21.2.0  # WSGI server for production deployment
-python-dateutil==2.8.2  # Enhanced date handling
+python-dateutil==2.8.2  # Enhanced date handling
+PyJWT==2.8.0  # JWT tokens for stateless authentication