Add login/register: JWT auth, SQLite users, protected generate, Login/Register pages and AuthContext

.env.example

@@ -11,5 +11,8 @@ GROQ_API_KEY=
 # CLAUDE_MODEL=claude-3-5-haiku-20241022
 # GROQ_MODEL=llama-3.3-70b-versatile
 
+# Auth: secret for JWT signing (set a long random string in production)
+SECRET_KEY=change-me-in-production
+
 # Optional: CORS origins (comma-separated). Use * for development.
 # CORS_ORIGINS=https://your-frontend.fly.dev

app.py

@@ -11,7 +11,7 @@ import os
 from pathlib import Path
 
 from dotenv import load_dotenv
-from fastapi import FastAPI, HTTPException, Request, APIRouter
+from fastapi import Depends, FastAPI, HTTPException, Request, APIRouter
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
 from slowapi import Limiter, _rate_limit_exceeded_handler
@@ -22,6 +22,9 @@ import logging
 
 from ai_prompts import TEMPLATES
 from utils import call_llm
+from database import init_db
+from auth import get_current_user
+from routers.auth import router as auth_router
 
 # -----------------------------------------------------------------------------
 # Environment setup
@@ -105,35 +108,43 @@ class GenerateRequest(BaseModel):
 # API Routes
 # -----------------------------------------------------------------------------
 
+@app.on_event("startup")
+def on_startup():
+    init_db()
+
+
+app.include_router(auth_router)
+
+
 @app.get("/")
 async def root():
     """Health/root endpoint so the backend URL doesn't return 404 when visited."""
-    return {"service": "Pseudogen API", "docs": "/docs", "generate": "POST /generate-pseudocode", "v1": "POST /v1/generate-pseudocode"}
+    return {"service": "Pseudogen API", "docs": "/docs", "auth": "/auth/login", "generate": "POST /generate-pseudocode", "v1": "POST /v1/generate-pseudocode"}
 
 # v1 API router (versioned endpoint; same behavior as legacy path)
 v1_router = APIRouter(prefix="/v1", tags=["v1"])
 
 @v1_router.post("/generate-pseudocode")
 @limiter.limit("30/minute")
-async def generate_v1(request: Request, req: GenerateRequest):
-    """Same as generate; versioned under /v1."""
-    return await generate_impl(request, req)
+async def generate_v1(request: Request, req: GenerateRequest, user: dict = Depends(get_current_user)):
+    """Same as generate; versioned under /v1. Requires auth."""
+    return await generate_impl(request, req, user)
 
 app.include_router(v1_router)
 
 @app.post("/generate-pseudocode")
 @limiter.limit("30/minute")
-async def generate(request: Request, req: GenerateRequest):
+async def generate(request: Request, req: GenerateRequest, user: dict = Depends(get_current_user)):
     """
-    Generate pseudocode from a given problem description.
-    Free plan: up to 4000 chars. Send X-Plan: premium for up to 12000 chars.
+    Generate pseudocode from a given problem description. Requires Bearer token.
+    Plan (free/premium) is taken from the logged-in user.
     """
-    return await generate_impl(request, req)
+    return await generate_impl(request, req, user)
 
 
-async def generate_impl(request: Request, req: GenerateRequest):
+async def generate_impl(request: Request, req: GenerateRequest, user: dict):
     """Shared implementation for generate and generate_v1."""
-    plan = (request.headers.get("X-Plan") or "free").strip().lower()
+    plan = (user.get("plan") or "free").strip().lower()
     if plan != "premium" and len(req.problem_description) > FREE_MAX_INPUT_LEN:
         raise HTTPException(
             status_code=400,

auth.py

@@ -0,0 +1,72 @@
+# backend/auth.py
+"""
+JWT and password hashing for Pseudogen. get_current_user dependency for protected routes.
+"""
+import os
+from datetime import datetime, timedelta, timezone
+
+import jwt
+from fastapi import Depends, HTTPException, status
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from passlib.context import CryptContext
+from dotenv import load_dotenv
+from pathlib import Path
+
+from database import get_user_by_email, get_user_by_id
+
+load_dotenv(dotenv_path=Path(__file__).resolve().parent / ".env")
+
+SECRET_KEY = os.getenv("SECRET_KEY", "change-me-in-production-use-env")
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 60 * 24 * 7  # 7 days
+
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+security = HTTPBearer(auto_error=False)
+
+
+def hash_password(password: str) -> str:
+    return pwd_context.hash(password)
+
+
+def verify_password(plain: str, hashed: str) -> bool:
+    return pwd_context.verify(plain, hashed)
+
+
+def create_access_token(data: dict) -> str:
+    to_encode = data.copy()
+    expire = datetime.now(timezone.utc) + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    to_encode["exp"] = expire
+    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+
+
+def decode_token(token: str) -> dict | None:
+    try:
+        return jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+    except jwt.PyJWTError:
+        return None
+
+
+async def get_current_user(
+    credentials: HTTPAuthorizationCredentials | None = Depends(security),
+) -> dict:
+    """Dependency: require valid Bearer token and return user dict (id, email, plan)."""
+    if credentials is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Not authenticated",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    payload = decode_token(credentials.credentials)
+    if payload is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or expired token",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    user_id = payload.get("sub")
+    if not user_id:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token")
+    user = get_user_by_id(int(user_id))
+    if user is None:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="User not found")
+    return user

database.py

@@ -0,0 +1,75 @@
+# backend/database.py
+"""
+SQLite database and user CRUD for Pseudogen auth.
+"""
+import sqlite3
+from pathlib import Path
+
+_BACKEND_DIR = Path(__file__).resolve().parent
+DB_PATH = _BACKEND_DIR / "pseudogen.db"
+
+
+def get_connection():
+    conn = sqlite3.connect(DB_PATH)
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def init_db():
+    """Create users table if it does not exist."""
+    conn = get_connection()
+    try:
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS users (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                email TEXT UNIQUE NOT NULL,
+                hashed_password TEXT NOT NULL,
+                plan TEXT NOT NULL DEFAULT 'free',
+                created_at TEXT NOT NULL DEFAULT (datetime('now'))
+            )
+        """)
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def get_user_by_email(email: str) -> dict | None:
+    conn = get_connection()
+    try:
+        row = conn.execute(
+            "SELECT id, email, hashed_password, plan, created_at FROM users WHERE email = ?",
+            (email.strip().lower(),),
+        ).fetchone()
+        if row is None:
+            return None
+        return dict(row)
+    finally:
+        conn.close()
+
+
+def get_user_by_id(user_id: int) -> dict | None:
+    conn = get_connection()
+    try:
+        row = conn.execute(
+            "SELECT id, email, plan, created_at FROM users WHERE id = ?",
+            (user_id,),
+        ).fetchone()
+        if row is None:
+            return None
+        return dict(row)
+    finally:
+        conn.close()
+
+
+def create_user(email: str, hashed_password: str, plan: str = "free") -> dict:
+    conn = get_connection()
+    try:
+        cursor = conn.execute(
+            "INSERT INTO users (email, hashed_password, plan) VALUES (?, ?, ?)",
+            (email.strip().lower(), hashed_password, plan),
+        )
+        conn.commit()
+        user_id = cursor.lastrowid
+        return {"id": user_id, "email": email.strip().lower(), "plan": plan}
+    finally:
+        conn.close()

requirements.txt

@@ -13,6 +13,11 @@ PyYAML==6.0.3
 # Rate limiting
 slowapi==0.1.9
 
+# Auth
+passlib[bcrypt]==1.7.4
+PyJWT==2.10.1
+email-validator==2.2.0
+
 # Testing
 pytest==8.3.4
 pytest-asyncio==0.24.0

routers/__init__.py

@@ -0,0 +1 @@
+# backend/routers

routers/auth.py

@@ -0,0 +1,66 @@
+# backend/routers/auth.py
+from fastapi import APIRouter, Depends, HTTPException, status
+from pydantic import BaseModel, EmailStr
+
+from database import init_db, get_user_by_email, create_user
+from auth import hash_password, verify_password, create_access_token, get_current_user
+
+router = APIRouter(prefix="/auth", tags=["auth"])
+
+
+class RegisterRequest(BaseModel):
+    email: EmailStr
+    password: str
+
+
+class LoginRequest(BaseModel):
+    email: EmailStr
+    password: str
+
+
+class TokenResponse(BaseModel):
+    access_token: str
+    token_type: str = "bearer"
+
+
+class UserResponse(BaseModel):
+    id: int
+    email: str
+    plan: str
+
+
+@router.post("/register", response_model=TokenResponse)
+def register(req: RegisterRequest):
+    if len(req.password) < 8:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Password must be at least 8 characters",
+        )
+    if get_user_by_email(req.email):
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Email already registered",
+        )
+    init_db()
+    hashed = hash_password(req.password)
+    user = create_user(req.email, hashed)
+    token = create_access_token({"sub": str(user["id"])})
+    return TokenResponse(access_token=token)
+
+
+@router.post("/login", response_model=TokenResponse)
+def login(req: LoginRequest):
+    init_db()
+    user = get_user_by_email(req.email)
+    if not user or not verify_password(req.password, user["hashed_password"]):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid email or password",
+        )
+    token = create_access_token({"sub": str(user["id"])})
+    return TokenResponse(access_token=token)
+
+
+@router.get("/me", response_model=UserResponse)
+def me(user: dict = Depends(get_current_user)):
+    return UserResponse(id=user["id"], email=user["email"], plan=user["plan"])

tests/test_app.py

@@ -1,14 +1,29 @@
 """
 Unit tests for Pseudogen FastAPI app.
 Tests GenerateRequest validation, style/detail handling, and generate endpoint with mocked LLM.
+Auth is overridden so generate endpoints receive a fake user.
 """
 import pytest
 from fastapi.testclient import TestClient
 from unittest.mock import patch
 
-# Import app after setting up mocks so we can patch utils.call_llm
 from app import app, GenerateRequest
+from auth import get_current_user
 
+# Override auth so generate endpoints see a fake user (plan from user, not header)
+FAKE_USER_FREE = {"id": 1, "email": "test@test.com", "plan": "free"}
+FAKE_USER_PREMIUM = {"id": 2, "email": "premium@test.com", "plan": "premium"}
+
+
+async def _fake_user_free():
+    return FAKE_USER_FREE
+
+
+async def _fake_user_premium():
+    return FAKE_USER_PREMIUM
+
+
+app.dependency_overrides[get_current_user] = _fake_user_free
 client = TestClient(app)
 
 
@@ -132,17 +147,20 @@ def test_free_plan_rejects_input_over_4000_chars():
 
 
 def test_premium_plan_accepts_input_up_to_12000_chars():
-    """With X-Plan: premium, input up to 12000 chars is accepted."""
-    with patch("app.call_llm") as mock_llm:
-        mock_llm.return_value = "BEGIN\nEND"
-        r = client.post(
-            "/generate-pseudocode",
-            headers={"X-Plan": "premium"},
-            json={
-                "problem_description": "y" * 10000,
-                "style": "Academic",
-                "detail": "Concise",
-            },
-        )
-    assert r.status_code == 200
-    assert r.json().get("markdown") == "BEGIN\nEND"
+    """With user plan premium, input up to 12000 chars is accepted."""
+    app.dependency_overrides[get_current_user] = _fake_user_premium
+    try:
+        with patch("app.call_llm") as mock_llm:
+            mock_llm.return_value = "BEGIN\nEND"
+            r = client.post(
+                "/generate-pseudocode",
+                json={
+                    "problem_description": "y" * 10000,
+                    "style": "Academic",
+                    "detail": "Concise",
+                },
+            )
+        assert r.status_code == 200
+        assert r.json().get("markdown") == "BEGIN\nEND"
+    finally:
+        app.dependency_overrides[get_current_user] = _fake_user_free

tests/test_auth.py

@@ -0,0 +1,69 @@
+"""
+Tests for auth endpoints: register, login, /me.
+Uses real SQLite DB (backend/pseudogen.db). Register uses a unique email per run to avoid conflicts.
+"""
+import time
+import pytest
+from fastapi.testclient import TestClient
+
+from app import app
+
+client = TestClient(app)
+
+
+def _unique_email():
+    return f"test-{int(time.time() * 1000)}@example.com"
+
+
+def test_register_returns_token():
+    email = _unique_email()
+    r = client.post("/auth/register", json={"email": email, "password": "password123"})
+    assert r.status_code == 200
+    data = r.json()
+    assert "access_token" in data
+    assert data.get("token_type") == "bearer"
+
+
+def test_register_rejects_short_password():
+    r = client.post("/auth/register", json={"email": "a@b.com", "password": "short"})
+    assert r.status_code == 400
+    assert "8" in r.json().get("detail", "")
+
+
+def test_register_rejects_duplicate_email():
+    email = _unique_email()
+    client.post("/auth/register", json={"email": email, "password": "password123"})
+    r = client.post("/auth/register", json={"email": email, "password": "password123"})
+    assert r.status_code == 400
+    assert "already" in r.json().get("detail", "").lower()
+
+
+def test_login_returns_token():
+    email = _unique_email()
+    client.post("/auth/register", json={"email": email, "password": "password123"})
+    r = client.post("/auth/login", json={"email": email, "password": "password123"})
+    assert r.status_code == 200
+    assert "access_token" in r.json()
+
+
+def test_login_rejects_wrong_password():
+    email = _unique_email()
+    client.post("/auth/register", json={"email": email, "password": "password123"})
+    r = client.post("/auth/login", json={"email": email, "password": "wrong"})
+    assert r.status_code == 401
+
+
+def test_me_returns_user_with_valid_token():
+    email = _unique_email()
+    reg = client.post("/auth/register", json={"email": email, "password": "password123"})
+    token = reg.json()["access_token"]
+    r = client.get("/auth/me", headers={"Authorization": f"Bearer {token}"})
+    assert r.status_code == 200
+    data = r.json()
+    assert data["email"] == email
+    assert data["plan"] == "free"
+
+
+def test_me_rejects_invalid_token():
+    r = client.get("/auth/me", headers={"Authorization": "Bearer invalid"})
+    assert r.status_code == 401