authentication refactor - added setup token and token refresh mechanism

.gitignore

@@ -35,3 +35,4 @@ Thumbs.db
 # Other
 resources/
 tmp/
+auth.db

studio/backend/auth/__init__.py

@@ -1,24 +1,46 @@
 """
 Authentication module for JWT-based auth with SQLite storage.
 """
-from .authentication import create_access_token, get_current_subject, reload_secret
+from .authentication import (
+    create_access_token,
+    create_refresh_token,
+    refresh_access_token,
+    get_current_subject,
+    reload_secret,
+)
 from .storage import (
     is_initialized,
     create_initial_user,
     get_user_and_secret,
     load_jwt_secret,
+    save_setup_token,
+    consume_setup_token,
+    has_pending_setup_token,
+    save_refresh_token,
+    verify_refresh_token,
+    revoke_user_refresh_tokens,
 )
 from .hashing import hash_password, verify_password
 
 __all__ = [
     "create_access_token",
+    "create_refresh_token",
+    "refresh_access_token",
     "get_current_subject",
     "reload_secret",
     "is_initialized",
     "create_initial_user",
     "get_user_and_secret",
     "load_jwt_secret",
+    "save_setup_token",
+    "consume_setup_token",
+    "has_pending_setup_token",
+    "save_refresh_token",
+    "verify_refresh_token",
+    "revoke_user_refresh_tokens",
     "hash_password",
     "verify_password",
 ]
 
+
+

studio/backend/auth/authentication.py

@@ -6,10 +6,11 @@ from fastapi import Depends, HTTPException, status
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 import jwt
 
-from .storage import load_jwt_secret
+from .storage import load_jwt_secret, save_refresh_token, verify_refresh_token
 
 ALGORITHM = "HS256"
 ACCESS_TOKEN_EXPIRE_MINUTES = 60
+REFRESH_TOKEN_EXPIRE_DAYS = 7
 
 # Load stable secret from SQLite (set during first-time setup)
 # This will raise RuntimeError if auth hasn't been initialized yet
@@ -40,6 +41,31 @@ def create_access_token(
     return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
 
 
+def create_refresh_token(subject: str) -> str:
+    """
+    Create a random refresh token, store its hash in SQLite, and return it.
+
+    Refresh tokens are opaque (not JWTs) and expire after REFRESH_TOKEN_EXPIRE_DAYS.
+    """
+    token = secrets.token_urlsafe(48)
+    expires_at = datetime.now(UTC) + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
+    save_refresh_token(token, subject, expires_at.isoformat())
+    return token
+
+
+def refresh_access_token(refresh_token: str) -> Optional[str]:
+    """
+    Validate a refresh token and issue a new access token.
+
+    The refresh token itself is NOT consumed — it stays valid until expiry.
+    Returns a new access_token or None if the refresh token is invalid/expired.
+    """
+    username = verify_refresh_token(refresh_token)
+    if username is None:
+        return None
+    return create_access_token(subject=username)
+
+
 def reload_secret() -> None:
     """
     Reload the JWT secret from SQLite.
@@ -77,7 +103,3 @@ async def get_current_subject(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Invalid or expired token",
         )
-# token = create_access_token("local-user")
-# print(token)
-
-

studio/backend/auth/storage.py

@@ -1,13 +1,20 @@
 """
 SQLite storage for authentication data (user credentials + JWT secret).
 """
+import hashlib
 import sqlite3
+from datetime import UTC, datetime
 from pathlib import Path
 from typing import Optional, Tuple
 
 DB_PATH = Path(__file__).parent / "auth.db"
 
 
+def _hash_token(token: str) -> str:
+    """SHA-256 hash a setup token for safe storage."""
+    return hashlib.sha256(token.encode("utf-8")).hexdigest()
+
+
 def get_connection() -> sqlite3.Connection:
     """Get a connection to the auth database, creating tables if needed."""
     conn = sqlite3.connect(DB_PATH)
@@ -23,6 +30,24 @@ def get_connection() -> sqlite3.Connection:
         );
         """
     )
+    conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS setup_tokens (
+            id INTEGER PRIMARY KEY,
+            token_hash TEXT NOT NULL
+        );
+        """
+    )
+    conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS refresh_tokens (
+            id INTEGER PRIMARY KEY,
+            token_hash TEXT NOT NULL,
+            username TEXT NOT NULL,
+            expires_at TEXT NOT NULL
+        );
+        """
+    )
     conn.commit()
     return conn
 
@@ -99,3 +124,119 @@ def load_jwt_secret() -> str:
     finally:
         conn.close()
 
+
+def save_setup_token(token: str) -> None:
+    """
+    Store a hashed setup token, replacing any existing one.
+    """
+    token_hash = _hash_token(token)
+    conn = get_connection()
+    try:
+        conn.execute("DELETE FROM setup_tokens")
+        conn.execute("INSERT INTO setup_tokens (token_hash) VALUES (?)", (token_hash,))
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def consume_setup_token(token: str) -> bool:
+    """
+    Verify a setup token and delete it if valid.
+    
+    Returns True if the token was valid (and is now consumed), False otherwise.
+    """
+    token_hash = _hash_token(token)
+    conn = get_connection()
+    try:
+        cur = conn.execute(
+            "SELECT id FROM setup_tokens WHERE token_hash = ?", (token_hash,)
+        )
+        row = cur.fetchone()
+        if row is None:
+            return False
+        conn.execute("DELETE FROM setup_tokens WHERE id = ?", (row["id"],))
+        conn.commit()
+        return True
+    finally:
+        conn.close()
+
+
+def has_pending_setup_token() -> bool:
+    """Check if a setup token is waiting to be consumed."""
+    conn = get_connection()
+    try:
+        cur = conn.execute("SELECT COUNT(*) AS c FROM setup_tokens")
+        row = cur.fetchone()
+        return bool(row["c"])
+    finally:
+        conn.close()
+
+
+def save_refresh_token(token: str, username: str, expires_at: str) -> None:
+    """
+    Store a hashed refresh token with its associated username and expiry.
+    """
+    token_hash = _hash_token(token)
+    conn = get_connection()
+    try:
+        conn.execute(
+            """
+            INSERT INTO refresh_tokens (token_hash, username, expires_at)
+            VALUES (?, ?, ?)
+            """,
+            (token_hash, username, expires_at),
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def verify_refresh_token(token: str) -> Optional[str]:
+    """
+    Verify a refresh token and return the username.
+
+    Returns the username if valid and not expired, None otherwise.
+    The token is NOT consumed — it stays valid until it expires.
+    """
+    token_hash = _hash_token(token)
+    conn = get_connection()
+    try:
+        # Clean up any expired tokens while we're here
+        conn.execute(
+            "DELETE FROM refresh_tokens WHERE expires_at < ?",
+            (datetime.now(UTC).isoformat(),),
+        )
+        conn.commit()
+
+        cur = conn.execute(
+            """
+            SELECT id, username, expires_at FROM refresh_tokens
+            WHERE token_hash = ?
+            """,
+            (token_hash,),
+        )
+        row = cur.fetchone()
+        if row is None:
+            return None
+
+        # Check expiry
+        expires_at = datetime.fromisoformat(row["expires_at"])
+        if datetime.now(UTC) > expires_at:
+            conn.execute("DELETE FROM refresh_tokens WHERE id = ?", (row["id"],))
+            conn.commit()
+            return None
+
+        return row["username"]
+    finally:
+        conn.close()
+
+
+def revoke_user_refresh_tokens(username: str) -> None:
+    """Revoke all refresh tokens for a user (e.g. on logout)."""
+    conn = get_connection()
+    try:
+        conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,))
+        conn.commit()
+    finally:
+        conn.close()
+

studio/backend/main.py

@@ -1,6 +1,9 @@
 """
 Main FastAPI application for Unsloth UI Backend
 """
+import secrets
+from contextlib import asynccontextmanager
+
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
@@ -10,12 +13,30 @@ from datetime import datetime
 
 # Import routers
 from routes import training_router, models_router, inference_router, datasets_router, auth_router
+from auth import storage
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Generate and print a setup token on startup if auth is not yet initialized."""
+    if not storage.is_initialized():
+        setup_token = secrets.token_urlsafe(32)
+        storage.save_setup_token(setup_token)
+        print("\n" + "=" * 60)
+        print("FIRST-TIME SETUP REQUIRED")
+        print("Use this one-time setup token to create your admin account:\n")
+        print(f"    {setup_token}\n")
+        print("This token can only be used once.")
+        print("=" * 60 + "\n")
+    yield
+
 
 # Create FastAPI app
 app = FastAPI(
     title="Unsloth UI Backend",
     version="1.0.0",
-    description="Backend API for Unsloth UI - Training and Model Management"
+    description="Backend API for Unsloth UI - Training and Model Management",
+    lifespan=lifespan,
 )
 
 # CORS middleware

studio/backend/models/auth.py

@@ -6,6 +6,7 @@ from pydantic import BaseModel, Field
 
 class AuthSetupRequest(BaseModel):
     """First-time setup: create the initial admin user + password."""
+    setup_token: str = Field(..., description="One-time setup token printed to the server console")
     username: str = Field(..., description="Admin username")
     password: str = Field(..., min_length=8, description="Admin password (minimum 8 characters)")
 
@@ -16,6 +17,11 @@ class AuthLoginRequest(BaseModel):
     password: str = Field(..., description="Password")
 
 
+class RefreshTokenRequest(BaseModel):
+    """Refresh token payload to obtain new access + refresh tokens."""
+    refresh_token: str = Field(..., description="Refresh token from a previous login or refresh")
+
+
 class AuthStatusResponse(BaseModel):
     """Indicate whether auth has been initialized."""
     initialized: bool = Field(..., description="True if auth setup has been completed")

studio/backend/models/users.py

@@ -4,7 +4,7 @@ This module defines the data models used for user authentication and management
 in the FastAPI application.
 """
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class User(BaseModel):
@@ -20,9 +20,10 @@ class UserInDB(BaseModel):
 
 
 class Token(BaseModel):
-    """Authentication token model with access token and type."""
+    """Authentication token model with access and refresh tokens."""
 
     access_token: str
+    refresh_token: str
     token_type: str
 
 
@@ -30,3 +31,4 @@ class TokenData(BaseModel):
     """Token payload model containing username."""
 
     username: str | None = None
+

studio/backend/routes/auth.py

@@ -7,11 +7,17 @@ import secrets
 from models.auth import (
     AuthSetupRequest,
     AuthLoginRequest,
+    RefreshTokenRequest,
     AuthStatusResponse,
 )
 from models.users import Token
 from auth import storage, hashing
-from auth.authentication import create_access_token, reload_secret
+from auth.authentication import (
+    create_access_token,
+    create_refresh_token,
+    refresh_access_token,
+    reload_secret,
+)
 
 router = APIRouter()
 
@@ -32,6 +38,7 @@ async def setup_auth(payload: AuthSetupRequest) -> Token:
     """
     First-time setup: create the admin user and a JWT secret.
 
+    Requires a valid setup token (printed to the server console on startup).
     Can only be called once. Subsequent calls will return 400.
     """
     if storage.is_initialized():
@@ -40,6 +47,13 @@ async def setup_auth(payload: AuthSetupRequest) -> Token:
             detail="Auth is already initialized.",
         )
 
+    # Validate the one-time setup token
+    if not storage.consume_setup_token(payload.setup_token):
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Invalid or expired setup token.",
+        )
+
     # Generate a strong random JWT secret for this installation
     jwt_secret = secrets.token_urlsafe(64)
 
@@ -59,15 +73,20 @@ async def setup_auth(payload: AuthSetupRequest) -> Token:
     # Reload JWT secret from DB (so authentication.py picks it up)
     reload_secret()
 
-    # Issue a token for the new user
+    # Issue access + refresh tokens for the new user
     access_token = create_access_token(subject=payload.username)
-    return Token(access_token=access_token, token_type="bearer")
+    refresh_token = create_refresh_token(subject=payload.username)
+    return Token(
+        access_token=access_token,
+        refresh_token=refresh_token,
+        token_type="bearer",
+    )
 
 
 @router.post("/login", response_model=Token)
 async def login(payload: AuthLoginRequest) -> Token:
     """
-    Login with username/password and receive a JWT.
+    Login with username/password and receive access + refresh tokens.
     """
     record = storage.get_user_and_secret(payload.username)
     if record is None:
@@ -84,5 +103,31 @@ async def login(payload: AuthLoginRequest) -> Token:
         )
 
     access_token = create_access_token(subject=payload.username)
-    return Token(access_token=access_token, token_type="bearer")
+    refresh_token = create_refresh_token(subject=payload.username)
+    return Token(
+        access_token=access_token,
+        refresh_token=refresh_token,
+        token_type="bearer",
+    )
+
+
+@router.post("/refresh", response_model=Token)
+async def refresh(payload: RefreshTokenRequest) -> Token:
+    """
+    Exchange a valid refresh token for a new access token.
+
+    The refresh token itself is reusable until it expires (7 days).
+    """
+    new_access_token = refresh_access_token(payload.refresh_token)
+    if new_access_token is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or expired refresh token",
+        )
+
+    return Token(
+        access_token=new_access_token,
+        refresh_token=payload.refresh_token,
+        token_type="bearer",
+    )
 

studio/backend/utils/datasets/dataset_utils.py

@@ -12,7 +12,7 @@ All internal utilities have been moved to separate modules:
 - chat_templates: apply_chat_template_to_dataset, get_tokenizer_chat_template, etc.
 - vlm_processing: generate_smart_vlm_instruction
 - data_collators: DeepSeekOCRDataCollator, VLMDataCollator
-- model_mappings: TEMPLATE_TO_MODEL_MAPPER, RESPONSE_MARKERS
+- model_mappings: TEMPLATE_TO_MODEL_MAPPER
 """
 
 # Import from modular files
@@ -37,7 +37,7 @@ from .chat_templates import (
 )
 from .vlm_processing import generate_smart_vlm_instruction
 from .data_collators import DeepSeekOCRDataCollator, VLMDataCollator
-from .model_mappings import TEMPLATE_TO_MODEL_MAPPER, RESPONSE_MARKERS
+from .model_mappings import TEMPLATE_TO_MODEL_MAPPER
 
 
 def check_dataset_format(dataset, is_vlm: bool = False) -> dict: