fix: Quota tracking was a process-local Python dict (app/routes/api.py

app/database.py

@@ -1,7 +1,7 @@
 """
 Database models and connection handling for Sugar-AI.
 """
-from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, Text
+from sqlalchemy import create_engine, Column, Integer, String, Boolean, Date, DateTime, Text
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, Session
 import datetime
@@ -39,6 +39,16 @@ def to_dict(self) -> Dict[str, Any]:
         }
 
 
+# per-key daily request quota, persisted so it survives restarts/reconnects
+class APIQuota(Base):
+    __tablename__ = "api_quotas"
+
+    api_key = Column(String, primary_key=True, index=True)
+    request_count = Column(Integer, default=0, nullable=False)
+    quota_date = Column(Date, nullable=False)
+    updated_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow)
+
+
 def create_tables() -> None:
     """Create database tables if they don't exist"""
     Base.metadata.create_all(bind=engine)

app/quota.py

@@ -0,0 +1,59 @@
+"""
+Persistent API quota tracking for Sugar-AI.
+"""
+import datetime
+from typing import Dict
+
+from fastapi import HTTPException
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+from app.config import settings
+
+
+def check_and_increment_quota(api_key: str, db: Session) -> Dict[str, int]:
+    """Check and increment a key's persisted daily request count.
+
+    The check and the increment are done via a single conditional UPDATE
+    statement so concurrent requests for the same key (FastAPI may run sync
+    dependencies like this one on separate threads/DB connections) can't
+    race on a read-modify-write and lose or double count updates.
+
+    Raises HTTPException(429) if the key has exhausted today's quota.
+    """
+    today = datetime.date.today()
+    max_requests = settings.MAX_DAILY_REQUESTS
+
+    # Ensure a row exists for this key, resetting it if it's from a
+    # previous day, before attempting the atomic increment below.
+    db.execute(
+        text(
+            "INSERT INTO api_quotas (api_key, request_count, quota_date) "
+            "VALUES (:api_key, 0, :today) "
+            "ON CONFLICT(api_key) DO UPDATE SET "
+            "request_count = CASE WHEN api_quotas.quota_date != :today THEN 0 "
+            "ELSE api_quotas.request_count END, "
+            "quota_date = :today"
+        ),
+        {"api_key": api_key, "today": today},
+    )
+
+    result = db.execute(
+        text(
+            "UPDATE api_quotas SET request_count = request_count + 1 "
+            "WHERE api_key = :api_key AND quota_date = :today "
+            "AND request_count < :max_requests"
+        ),
+        {"api_key": api_key, "today": today, "max_requests": max_requests},
+    )
+    db.commit()
+
+    if result.rowcount == 0:
+        raise HTTPException(status_code=429, detail="Daily request quota exceeded")
+
+    request_count = db.execute(
+        text("SELECT request_count FROM api_quotas WHERE api_key = :api_key"),
+        {"api_key": api_key},
+    ).scalar_one()
+
+    return {"remaining": max_requests - request_count, "total": max_requests}

app/routes/api.py

@@ -8,13 +8,13 @@
 import logging
 import os
 import json
-from datetime import datetime
-from typing import Dict, Optional, List
+from typing import Optional, List
 
 from app.database import get_db, APIKey
 from app.ai import RAGAgent
 from app.providers.base import GenerationParams
 from app.config import settings
+from app.quota import check_and_increment_quota
 
 # Pydantic models for chat completions
 class ChatMessage(BaseModel):
@@ -42,43 +42,30 @@ class PromptedLLMRequest(BaseModel):
 # Initialize the agent
 agent = None
 
-# user quotas tracking
-user_quotas: Dict[str, Dict] = {}
-
-def check_quota(api_key: str) -> bool:
-    """Check if a user has exceeded their daily quota"""
-    today = datetime.now().date()
-
-    if api_key not in user_quotas:
-        user_quotas[api_key] = {"count": 0, "date": today}
-        return True
-
-    # reset quota daily
-    if user_quotas[api_key]["date"] != today:
-        user_quotas[api_key]["count"] = 0
-        user_quotas[api_key]["date"] = today
-
-    if user_quotas[api_key]["count"] >= settings.MAX_DAILY_REQUESTS:
-        return False
-
-    user_quotas[api_key]["count"] += 1
-    return True
-
-def verify_api_key(api_key: Optional[str] = Header(None, alias="X-API-Key"), request: Request = None):
+def verify_api_key(
+    api_key: Optional[str] = Header(None, alias="X-API-Key"),
+    request: Request = None,
+    db: Session = Depends(get_db),
+):
     """Verify API key and check quota"""
     if not api_key:
         logger.warning(f"API key missing: {request.client.host if request else 'unknown'}")
         raise HTTPException(status_code=401, detail="API key is missing")
-    
+
     if api_key not in settings.API_KEYS:
         logger.warning(f"Invalid API key used: {api_key[:5]}... from {request.client.host if request else 'unknown'}")
         raise HTTPException(status_code=401, detail="Invalid API key")
-
-    if not check_quota(api_key):
+
+    try:
+        quota = check_and_increment_quota(api_key, db)
+    except HTTPException:
         logger.warning(f"Quota exceeded for user: {settings.API_KEYS[api_key]['name']}")
-        raise HTTPException(status_code=429, detail="Daily request quota exceeded")
-
-    return settings.API_KEYS[api_key]
+        raise
+
+    # copy so we don't mutate the shared settings.API_KEYS entry
+    user_info = dict(settings.API_KEYS[api_key])
+    user_info["quota"] = quota
+    return user_info
 
 @router.post("/ask")
 async def ask_question(
@@ -99,20 +86,10 @@ async def ask_question(
         process_time = time.time() - start_time
         logger.info(f"RESPONSE - User: {user_info['name']} - Success - Time: {process_time:.2f}s")
 
-        # check quota
-        api_key = next(
-            key for key, value in settings.API_KEYS.items()
-            if value['name'] == user_info['name']
-        )
-        remaining = (
-            settings.MAX_DAILY_REQUESTS
-            - user_quotas.get(api_key, {}).get("count", 0)
-        )
-
         return {
-            "answer": answer, 
+            "answer": answer,
             "user": user_info["name"],
-            "quota": {"remaining": remaining, "total": settings.MAX_DAILY_REQUESTS}
+            "quota": user_info["quota"]
         }
     except Exception as e:
         logger.error(f"ERROR - User: {user_info['name']} - Error: {str(e)}")
@@ -136,14 +113,10 @@ async def ask_llm(
         process_time = time.time() - start_time
         logger.info(f"RESPONSE - User: {user_info['name']} - Success - Time: {process_time:.2f}s")
 
-        # check quota
-        api_key = next(key for key, value in settings.API_KEYS.items() if value['name'] == user_info['name'])
-        remaining = settings.MAX_DAILY_REQUESTS - user_quotas.get(api_key, {}).get("count", 0)
-
         return {
-            "answer": answer, 
+            "answer": answer,
             "user": user_info["name"],
-            "quota": {"remaining": remaining, "total": settings.MAX_DAILY_REQUESTS}
+            "quota": user_info["quota"]
         }
     except Exception as e:
         logger.error(f"ERROR - User: {user_info['name']} - Error: {str(e)}")
@@ -160,11 +133,7 @@ async def ask_llm_prompted(
     """
     start_time = time.time()
     client_ip = request.client.host if request else "unknown"
-
-    # Check quota first
-    api_key = next(key for key, value in settings.API_KEYS.items() if value['name'] == user_info['name'])
-    remaining = settings.MAX_DAILY_REQUESTS - user_quotas.get(api_key, {}).get("count", 0)
-
+
     try:
         if request_data.chat:
             # Chat completions mode
@@ -213,7 +182,7 @@ async def ask_llm_prompted(
                     "finish_reason": "stop"
                 }],
                 "user": user_info["name"],
-                "quota": {"remaining": remaining, "total": settings.MAX_DAILY_REQUESTS},
+                "quota": user_info["quota"],
                 "generation_params": {
                     "max_length": request_data.max_length,
                     "truncation": request_data.truncation,
@@ -252,7 +221,7 @@ async def ask_llm_prompted(
             return {
                 "answer": answer, 
                 "user": user_info["name"],
-                "quota": {"remaining": remaining, "total": settings.MAX_DAILY_REQUESTS},
+                "quota": user_info["quota"],
                 "generation_params": {
                     "max_length": request_data.max_length,
                     "truncation": request_data.truncation,
@@ -289,16 +258,12 @@ async def debug(
         process_time = time.time() - start_time
         logger.info(f"RESPONSE - User: {user_info['name']} - Success - Time: {process_time:.2f}s")
 
-        # check quota
-        api_key = next(key for key, value in settings.API_KEYS.items() if value['name'] == user_info['name'])
-        remaining = settings.MAX_DAILY_REQUESTS - user_quotas.get(api_key, {}).get("count", 0)
-
         return {
-            "answer": answer, 
+            "answer": answer,
             "user": user_info["name"],
-            "quota": {"remaining": remaining, "total": settings.MAX_DAILY_REQUESTS}
+            "quota": user_info["quota"]
         }
-        
+
     except Exception as e:
         logger.error(f"ERROR - User: {user_info['name']} - Error: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")

tests/test_quota.py

@@ -0,0 +1,129 @@
+"""
+Tests for persistent, session-safe API quota tracking.
+"""
+from concurrent.futures import ThreadPoolExecutor
+from datetime import date, timedelta
+
+import pytest
+from fastapi import HTTPException
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.pool import StaticPool
+
+from app.config import settings
+from app.database import APIQuota, Base
+from app.quota import check_and_increment_quota
+
+
+def make_sessionmaker(db_path):
+    """Create a sessionmaker bound to a shared-cache SQLite file so multiple
+    engines/sessions (simulating separate connections/reconnects) can see
+    the same persisted data."""
+    engine = create_engine(
+        f"sqlite:///{db_path}",
+        connect_args={"check_same_thread": False},
+    )
+    Base.metadata.create_all(bind=engine)
+    return engine, sessionmaker(bind=engine)
+
+
+def test_quota_persists_across_recreated_sessions(tmp_path):
+    """Quota state must survive a brand new engine/session pointed at the
+    same database, simulating a process restart or reconnect."""
+    db_path = tmp_path / "quota.db"
+
+    engine1, SessionLocal1 = make_sessionmaker(db_path)
+    db1 = SessionLocal1()
+    check_and_increment_quota("key1", db1)
+    check_and_increment_quota("key1", db1)
+    db1.close()
+    engine1.dispose()
+
+    # Recreate the engine/session from scratch, as would happen on restart.
+    engine2, SessionLocal2 = make_sessionmaker(db_path)
+    db2 = SessionLocal2()
+    quota = check_and_increment_quota("key1", db2)
+    db2.close()
+    engine2.dispose()
+
+    assert quota["remaining"] == settings.MAX_DAILY_REQUESTS - 3
+
+
+def test_first_request_is_counted(tmp_path):
+    engine, SessionLocal = make_sessionmaker(tmp_path / "quota.db")
+    db = SessionLocal()
+
+    quota = check_and_increment_quota("key1", db)
+
+    assert quota["remaining"] == settings.MAX_DAILY_REQUESTS - 1
+    assert quota["total"] == settings.MAX_DAILY_REQUESTS
+
+    db.close()
+    engine.dispose()
+
+
+def test_quota_exceeded_raises_429(tmp_path):
+    engine, SessionLocal = make_sessionmaker(tmp_path / "quota.db")
+    db = SessionLocal()
+
+    for _ in range(settings.MAX_DAILY_REQUESTS):
+        check_and_increment_quota("key1", db)
+
+    with pytest.raises(HTTPException) as exc_info:
+        check_and_increment_quota("key1", db)
+
+    assert exc_info.value.status_code == 429
+    assert exc_info.value.detail == "Daily request quota exceeded"
+
+    db.close()
+    engine.dispose()
+
+
+def test_quota_resets_on_new_day(tmp_path):
+    engine, SessionLocal = make_sessionmaker(tmp_path / "quota.db")
+    db = SessionLocal()
+
+    check_and_increment_quota("key1", db)
+    row = db.query(APIQuota).filter(APIQuota.api_key == "key1").first()
+    row.quota_date = date.today() - timedelta(days=1)
+    row.request_count = settings.MAX_DAILY_REQUESTS
+    db.commit()
+
+    quota = check_and_increment_quota("key1", db)
+
+    assert quota["remaining"] == settings.MAX_DAILY_REQUESTS - 1
+
+    db.close()
+    engine.dispose()
+
+
+def test_concurrent_requests_never_exceed_quota(tmp_path):
+    """Concurrent requests for the same key (as FastAPI's threadpool would
+    dispatch for sync dependencies) must never be granted more than the
+    daily limit, even when they race on the read-modify-write."""
+    db_path = tmp_path / "quota.db"
+    engine, SessionLocal = make_sessionmaker(db_path)
+
+    attempts = settings.MAX_DAILY_REQUESTS * 3
+
+    def worker(_):
+        db = SessionLocal()
+        try:
+            check_and_increment_quota("key1", db)
+            return True
+        except HTTPException:
+            return False
+        finally:
+            db.close()
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        results = list(executor.map(worker, range(attempts)))
+
+    assert sum(results) == settings.MAX_DAILY_REQUESTS
+
+    db = SessionLocal()
+    row = db.query(APIQuota).filter(APIQuota.api_key == "key1").first()
+    db.close()
+    engine.dispose()
+
+    assert row.request_count == settings.MAX_DAILY_REQUESTS