refactor: replace daily usage increment with atomic reservation

main.py

@@ -50,23 +50,23 @@ async def lifespan(app: FastAPI):
         logger.warning(f"Embedder failed to load: {e}")
 
     # Eagerly load ASR models so warmup runs at startup, not on first request
-    try:
-        from src.asr.models import get_audio_model_en
-        get_audio_model_en()
-    except Exception as e:
-        logger.warning(f"English ASR model failed to load: {e}")
-
-    try:
-        from src.asr.models import get_audio_model_ar
-        get_audio_model_ar()
-    except Exception as e:
-        logger.warning(f"Arabic ASR model failed to load: {e}")
-
-    from src.rag.batch_workers import start_workers
-    start_workers()
-
-    from src.asr.batch_workers import start_asr_workers
-    start_asr_workers()
+    # try:
+    #     from src.asr.models import get_audio_model_en
+    #     get_audio_model_en()
+    # except Exception as e:
+    #     logger.warning(f"English ASR model failed to load: {e}")
+
+    # try:
+    #     from src.asr.models import get_audio_model_ar
+    #     get_audio_model_ar()
+    # except Exception as e:
+    #     logger.warning(f"Arabic ASR model failed to load: {e}")
+
+    # from src.rag.batch_workers import start_workers
+    # start_workers()
+
+    # from src.asr.batch_workers import start_asr_workers
+    # start_asr_workers()
 
     yield
 

quiz_generator/routes.py

@@ -3,7 +3,7 @@ from loguru import logger
 from fastapi import APIRouter, HTTPException, Depends
 from typing import Optional
 from src.quiz_generator.quiz import smart_quiz_generator
-from src.store import get_material, get_chunks, save_chunks, get_summary, save_quiz, get_quizzes, save_quiz_result, get_quiz_results, check_daily_limit, increment_daily_usage, ADMIN_EMAILS
+from src.store import get_material, get_chunks, save_chunks, get_summary, save_quiz, get_quizzes, save_quiz_result, get_quiz_results, atomic_check_and_increment_daily_limit, decrement_daily_usage, ADMIN_EMAILS
 from src.dependencies import get_current_user_id, get_current_user
 from src.config import settings
 from .schemas import QuizRequest, QuizResponse, SaveQuizResultRequest
@@ -31,9 +31,10 @@ async def generate_quiz(
     user_id: str = Depends(get_current_user_id),
     current_user=Depends(get_current_user),
 ):
-    # Rate limit check
+    # Atomically reserve a daily request slot before starting expensive generation.
     user_email = current_user.get("email") if isinstance(current_user, dict) else getattr(current_user, "email", None)
-    if not check_daily_limit(user_id, email=user_email, limit=20):
+    is_admin = bool(user_email and user_email in ADMIN_EMAILS)
+    if not atomic_check_and_increment_daily_limit(user_id, email=user_email, limit=20):
         raise HTTPException(429, "Daily limit of 20 requests reached. Come back tomorrow!")
 
     body.difficulty = body.difficulty.capitalize()
@@ -158,16 +159,23 @@ async def generate_quiz(
             model_name=settings.model_name,
         )
 
-        # Only increment limit if the quiz was generated successfully and saved without error
-        if not (user_email and user_email in ADMIN_EMAILS):
-            increment_daily_usage(user_id)
-
         return QuizResponse(quiz=quiz, quiz_id=saved["id"])
     except ValueError as e:
         logger.warning(f"Validation error in generate_quiz: {str(e)}")
+        # Refund the reserved slot — validation errors are our fault, not the user's.
+        if not is_admin:
+            decrement_daily_usage(user_id)
         raise HTTPException(400, str(e))
+    except HTTPException:
+        # Re-raise HTTP exceptions as-is (403, 400, etc.).
+        if not is_admin:
+            decrement_daily_usage(user_id)
+        raise
     except Exception as e:
         logger.error(f"Quiz generation failed: {str(e)}", exc_info=True)
+        # Refund: generation/save failed through no fault of the user.
+        if not is_admin:
+            decrement_daily_usage(user_id)
         raise HTTPException(500, f"Quiz generation failed: {e}")
 
 

rag/rag.py

@@ -75,40 +75,60 @@ async def store_embeddings_async(material_id: str, chunk_ids: list[str], chunks:
     """
     Async variant of store_embeddings that routes embedding inference through
     the batch worker queue for batching across concurrent requests.
+    Falls back to the synchronous path if the batch workers are not running.
     """
-    from src.rag.batch_workers import embedding_queue, job_store
-    from .schemas import EmbeddingJob
+    try:
+        from src.rag.batch_workers import embedding_queue, job_store
+        from .schemas import EmbeddingJob
 
-    job = EmbeddingJob(job_id=str(uuid.uuid4()), texts=chunks)
-    job_store[job.job_id] = {"status": "pending", "result": None, "error": None}
-    await embedding_queue.put(job)
-    await job.done.wait()
+        job = EmbeddingJob(job_id=str(uuid.uuid4()), texts=chunks)
+        job_store[job.job_id] = {"status": "pending", "result": None, "error": None}
+        await embedding_queue.put(job)
 
-    entry = job_store.pop(job.job_id)
-    if entry["status"] == "error":
-        raise RuntimeError(f"Embedding failed: {entry['error']}")
+        # Wait for the worker to process the job, but with a timeout.
+        # If batch workers are not running (e.g. commented out in main.py),
+        # this would hang forever — the timeout triggers a fallback to the sync path.
+        try:
+            await asyncio.wait_for(job.done.wait(), timeout=30.0)
+        except asyncio.TimeoutError:
+            # Workers not running — clean up and fall back to sync embedding
+            job_store.pop(job.job_id, None)
+            logger.warning(
+                f"store_embeddings_async timed out waiting for batch worker "
+                f"(material={material_id}). Falling back to synchronous embedding."
+            )
+            await asyncio.to_thread(store_embeddings, material_id, chunk_ids, chunks)
+            return
+
+        entry = job_store.pop(job.job_id)
+        if entry["status"] == "error":
+            raise RuntimeError(f"Embedding failed: {entry['error']}")
+
+        embeddings = entry["result"]
+
+        records = [
+            {"chunk_id": cid, "material_id": material_id, "embedding": emb}
+            for cid, emb in zip(chunk_ids, embeddings)
+        ]
 
-    embeddings = entry["result"]
+        db = get_supabase()
+        if db is None:
+            logger.warning("Supabase not connected — embeddings computed but NOT stored (no DB).")
+            return
 
-    records = [
-        {"chunk_id": cid, "material_id": material_id, "embedding": emb}
-        for cid, emb in zip(chunk_ids, embeddings)
-    ]
+        logger.info(f"Storing {len(records)} embeddings in Supabase for material {material_id}...")
 
-    db = get_supabase()
-    if db is None:
-        logger.warning("Supabase not connected — embeddings computed but NOT stored (no DB).")
-        return
+        def _insert_records():
+            for i in range(0, len(records), 50):
+                db.table("material_embeddings").insert(records[i:i + 50]).execute()
 
-    logger.info(f"Storing {len(records)} embeddings in Supabase for material {material_id}...")
-    
-    def _insert_records():
-        for i in range(0, len(records), 50):
-            db.table("material_embeddings").insert(records[i:i + 50]).execute()
-            
-    loop = asyncio.get_event_loop()
-    await loop.run_in_executor(None, _insert_records)
-    logger.info(f"Embeddings stored successfully for material {material_id}.")
+        await asyncio.to_thread(_insert_records)
+        logger.info(f"Embeddings stored successfully for material {material_id}.")
+
+    except Exception as e:
+        # If anything unexpected fails, fall back to sync to avoid blocking the caller.
+        logger.warning(f"store_embeddings_async failed ({e}); falling back to sync.")
+        await asyncio.to_thread(store_embeddings, material_id, chunk_ids, chunks)
 
 
 def similarity_search(query: str, material_id: str, k: int = 5) -> list[dict]:

store.py

@@ -816,14 +816,154 @@ def increment_daily_usage(user_id: str) -> None:
         logger.error(f"Failed to increment daily usage: {e}")
 
 
+def decrement_daily_usage(user_id: str) -> None:
+    """
+    Rollback helper — decrements the daily counter by 1 (clamped to 0).
+
+    Called when a generation request was reserved (incremented) but failed
+    before producing a usable result, so the user is not penalised for a
+    server-side error.
+    """
+    today = _get_today_date_str()
+
+    # ── Redis path — atomic DECR clamped to 0 ────────────────────────────────
+    r = get_redis()
+    if r is not None:
+        try:
+            rkey = f"{DAILY_RATE_LIMIT_KEY_PREFIX}{user_id}:{today}"
+            # DECR is atomic; clamp to 0 so we never go negative
+            current = r.decr(rkey)
+            if current < 0:
+                r.set(rkey, 0)
+                r.expire(rkey, DAILY_RATE_LIMIT_REDIS_TTL)
+        except Exception as e:
+            logger.warning("Redis decrement_daily_usage failed: %s", e)
+
+    # ── Supabase — keep DB in sync ────────────────────────────────────────────
+    try:
+        result = _robust_execute(
+            _table_supabase("profiles")
+            .select("daily_requests, last_request_date")
+            .eq("id", user_id)
+        )
+        if not result.data:
+            return
+        profile = result.data[0] if result.data else None
+        if not profile:
+            return
+
+        last_date = profile.get("last_request_date")
+        count = profile.get("daily_requests", 0) or 0
+        if last_date != today:
+            count = 0
+        new_count = max(0, count - 1)
+
+        _robust_execute(
+            _table_supabase("profiles")
+            .update({"daily_requests": new_count, "last_request_date": today})
+            .eq("id", user_id)
+        )
+    except Exception as e:
+        logger.error(f"Failed to decrement daily usage: {e}")
+
+
+# Lua script: atomically check count < limit, then INCR + EXPIRE if allowed.
+# Returns 1 (allowed and incremented) or 0 (limit exceeded, no change).
+_ATOMIC_RATE_LIMIT_LUA = """
+local key   = KEYS[1]
+local limit = tonumber(ARGV[1])
+local ttl   = tonumber(ARGV[2])
+local count = tonumber(redis.call('GET', key) or 0)
+if count >= limit then
+    return 0
+end
+redis.call('INCR', key)
+redis.call('EXPIRE', key, ttl)
+return 1
+"""
+
+
+def atomic_check_and_increment_daily_limit(
+    user_id: str, email: Optional[str] = None, limit: int = 20
+) -> bool:
+    """
+    Atomically check the daily limit AND increment in one operation.
+
+    This eliminates the TOCTOU race condition that existed when `check_daily_limit`
+    and `increment_daily_usage` were called as two separate steps: concurrent
+    requests could both pass the check before either had incremented the counter.
+
+    Redis path: executes a Lua script so the GET + conditional INCR is one
+    indivisible command — Redis serialises all commands within a script.
+
+    Supabase fallback: increment first, then verify; rollback if over limit.
+    This is safe because Supabase operations are individually atomic (but not
+    as tight as the Lua approach under extreme concurrency).
+
+    Returns True if the request is allowed (counter was incremented),
+    False if the daily limit is already reached (counter unchanged).
+    """
+    # Admins are always allowed and never counted
+    if email and email in ADMIN_EMAILS:
+        return True
+
+    today = _get_today_date_str()
+
+    # ── Redis path: true atomic check-and-increment via Lua ───────────────────
+    r = get_redis()
+    if r is not None:
+        try:
+            rkey = f"{DAILY_RATE_LIMIT_KEY_PREFIX}{user_id}:{today}"
+            result = r.eval(_ATOMIC_RATE_LIMIT_LUA, 1, rkey, limit, DAILY_RATE_LIMIT_REDIS_TTL)
+            return bool(result)   # 1 → allowed, 0 → exceeded
+        except Exception as e:
+            logger.warning(
+                "Redis atomic_check_and_increment failed: %s — falling back to Supabase", e
+            )
+
+    # ── Supabase fallback: increment-first strategy ───────────────────────────
+    try:
+        result = _robust_execute(
+            _table_supabase("profiles")
+            .select("daily_requests, last_request_date")
+            .eq("id", user_id)
+        )
+        if not result.data:
+            # No profile row yet — treat as first request (allowed)
+            _robust_execute(
+                _table_supabase("profiles")
+                .update({"daily_requests": 1, "last_request_date": today})
+                .eq("id", user_id)
+            )
+            return True
+
+        profile = result.data[0]
+        last_date = profile.get("last_request_date")
+        count = profile.get("daily_requests", 0) or 0
+        if last_date != today:
+            count = 0   # day rolled over — reset
+
+        if count >= limit:
+            return False  # already at limit — do not increment
+
+        # Increment in Supabase
+        _robust_execute(
+            _table_supabase("profiles")
+            .update({"daily_requests": count + 1, "last_request_date": today})
+            .eq("id", user_id)
+        )
+        return True
+    except Exception as e:
+        logger.error(f"Supabase atomic_check_and_increment failed: {e}")
+        return True   # fail open rather than block the user on a DB error
+
+
 def check_and_increment_daily_limit(user_id: str, email: Optional[str] = None, limit: int = 20) -> bool:
     """
-    Check limit and increment if allowed. Deprecated/kept for compatibility.
+    Legacy wrapper kept for backward compatibility (used in rag/routes.py).
+    Delegates to the new atomic implementation.
     """
-    allowed = check_daily_limit(user_id, email, limit)
-    if allowed and not (email and email in ADMIN_EMAILS):
-        increment_daily_usage(user_id)
-    return allowed
+    return atomic_check_and_increment_daily_limit(user_id, email, limit)
 
 
 def get_usage(user_id: str) -> dict:

summary_generator/routes.py

@@ -9,7 +9,7 @@ from src.rag.rag import store_embeddings_async
 from src.store import (
     get_material, get_chunks, save_chunks, save_summary,
     get_summary as get_stored_summary, update_material_status,
-    check_daily_limit, increment_daily_usage,
+    atomic_check_and_increment_daily_limit, decrement_daily_usage, ADMIN_EMAILS,
 )
 from src.dependencies import get_current_user_id, get_current_user
 from src.config import settings
@@ -25,15 +25,20 @@ async def generate_summary(
     user_id: str = Depends(get_current_user_id),
     current_user=Depends(get_current_user),
 ):
-    # Rate limit check
+    # Atomically reserve a daily request slot before starting expensive generation.
     user_email = current_user.get("email") if isinstance(current_user, dict) else getattr(current_user, "email", None)
-    if not check_daily_limit(user_id, email=user_email, limit=20):
+    is_admin = bool(user_email and user_email in ADMIN_EMAILS)
+    if not atomic_check_and_increment_daily_limit(user_id, email=user_email, limit=20):
         raise HTTPException(429, "Daily limit of 20 requests reached. Come back tomorrow!")
 
     mat = get_material(body.material_id)
     if not mat:
+        if not is_admin:
+            decrement_daily_usage(user_id)
         raise HTTPException(404, "Material not found")
     if mat.get("user_id") != user_id:
+        if not is_admin:
+            decrement_daily_usage(user_id)
         raise HTTPException(403, "Access denied")
 
     try:
@@ -106,12 +111,15 @@ async def generate_summary(
             model_name=settings.model_name,
         )
 
-        # Only increment limit if the summary was generated successfully and saved without error
-        if not (user_email and user_email in settings.admin_emails if hasattr(settings, "admin_emails") else False):
-            increment_daily_usage(user_id)
-
         return SummarizeResponse(summary=summary, time_taken=elapsed)
+    except HTTPException:
+        # Refund on any HTTP exception raised inside the try block (e.g. 400).
+        if not is_admin:
+            decrement_daily_usage(user_id)
+        raise
     except Exception as e:
+        if not is_admin:
+            decrement_daily_usage(user_id)
         raise HTTPException(500, f"Summarization failed: {e}")