feat: implement refresh token handling and improve authentication flow; add email stats endpoint

Author: namanxdev

.claude/settings.local.json

@@ -19,7 +19,8 @@
       "WebFetch(domain:sarvam.ai)",
       "WebFetch(domain:relay.app)",
       "WebFetch(domain:stack-ai.com)",
-      "Bash(xargs basename:*)"
+      "Bash(xargs basename:*)",
+      "Bash(ls:*)"
     ]
   }
 }

backend/app/api/middleware/csrf.py

@@ -7,13 +7,11 @@
 
 from __future__ import annotations
 
-import hashlib
 import hmac
 import logging
 import secrets
-from datetime import UTC, datetime, timedelta
 
-from fastapi import Request, Response
+from fastapi import Request
 from starlette.middleware.base import BaseHTTPMiddleware
 
 logger = logging.getLogger(__name__)
@@ -27,6 +25,9 @@
 # Paths that are exempt from CSRF protection (e.g., webhook endpoints)
 EXEMPT_PATHS = frozenset({
     "/api/v1/auth/gmail/callback",  # OAuth callback doesn't need CSRF
+    "/api/v1/auth/login",
+    "/api/v1/auth/refresh",   # Uses httpOnly cookie — already secure
+    "/api/v1/auth/logout",    # Must work even when session is expired
     "/health",
     "/docs",
     "/redoc",
@@ -55,12 +56,12 @@ def _should_skip_csrf(request: Request) -> bool:
     # Skip for safe methods
     if request.method in SAFE_METHODS:
         return True
-    
+
     # Skip for exempt paths
     path = request.url.path
     if any(path.startswith(exempt) for exempt in EXEMPT_PATHS):
         return True
-    
+
     return False
 
 
@@ -80,19 +81,19 @@ class CSRFMiddleware(BaseHTTPMiddleware):
 
     async def dispatch(self, request: Request, call_next):
         """Process the request with CSRF protection."""
-        
+
         # Get existing CSRF token from cookie or generate new one
         csrf_cookie = request.cookies.get(CSRF_COOKIE_NAME)
         if not csrf_cookie:
             csrf_cookie = _generate_csrf_token()
-        
+
         # Store token in request state for response handler
         request.state.csrf_token = csrf_cookie
-        
+
         # Check if we need to validate CSRF token
         if not _should_skip_csrf(request):
             csrf_header = request.headers.get(CSRF_HEADER_NAME)
-            
+
             if not csrf_header:
                 logger.warning(
                     "CSRF token missing in header for %s %s",
@@ -107,7 +108,7 @@ async def dispatch(self, request: Request, call_next):
                         "code": "CSRF_TOKEN_MISSING"
                     }
                 )
-            
+
             if not hmac.compare_digest(csrf_cookie, csrf_header):
                 logger.warning(
                     "CSRF token mismatch for %s %s",
@@ -122,19 +123,19 @@ async def dispatch(self, request: Request, call_next):
                         "code": "CSRF_TOKEN_INVALID"
                     }
                 )
-        
+
         # Process the request
         response = await call_next(request)
-        
+
         # Set/update CSRF cookie
         # Use SameSite=Strict to prevent cross-site requests from sending the cookie
         # Use HttpOnly=False so JavaScript can read it (needed for double-submit pattern)
         # Use Secure in production
         from app.core.config import settings
-        
+
         cookie_value = request.state.csrf_token
         max_age = CSRF_TOKEN_EXPIRY_HOURS * 3600
-        
+
         response.set_cookie(
             key=CSRF_COOKIE_NAME,
             value=cookie_value,
@@ -144,7 +145,7 @@ async def dispatch(self, request: Request, call_next):
             samesite="strict",
             path="/",
         )
-        
+
         return response
 
 

backend/app/api/routes/auth.py

@@ -8,11 +8,14 @@
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from app.api.middleware.rate_limit import auth_rate_limit
+from app.core.config import settings
 from app.core.database import get_db
-from app.core.deps import get_current_user, oauth2_scheme
+from app.core.deps import oauth2_scheme
 from app.core.security import (
     create_access_token,
     create_password_reset_token,
+    create_refresh_token,
     decode_token,
     hash_password,
 )
@@ -28,9 +31,7 @@
     RegisterResponse,
     ResetPasswordRequest,
     TokenResponse,
-    TokenResponseWithRefresh,
 )
-from app.api.middleware.rate_limit import auth_rate_limit
 from app.services import auth_service
 
 logger = logging.getLogger(__name__)
@@ -78,7 +79,7 @@ async def login(
     Set use_cookies=false to receive the refresh token in the response body (less secure).
     """
     result = await auth_service.login(db, body.email, body.password)
-    
+
     # Set refresh token as httpOnly cookie for security (SEC-3 fix)
     if use_cookies and result.refresh_token:
         response.set_cookie(
@@ -90,7 +91,7 @@ async def login(
             max_age=7 * 24 * 60 * 60,  # 7 days
             path="/api/v1/auth/refresh",  # Only sent to refresh endpoint
         )
-    
+
     return TokenResponse(
         access_token=result.access_token,
         token_type=result.token_type,
@@ -116,16 +117,16 @@ async def refresh(
     """
     # Try to get refresh token from cookie first (more secure)
     refresh_token = request.cookies.get("refresh_token") or token
-    
+
     if not refresh_token:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Refresh token required",
             headers={"WWW-Authenticate": "Bearer"},
         )
-    
+
     result = await auth_service.refresh_token(db, refresh_token)
-    
+
     # Update refresh token cookie if using cookie-based auth
     if request.cookies.get("refresh_token") and result.refresh_token:
         response.set_cookie(
@@ -137,7 +138,7 @@ async def refresh(
             max_age=7 * 24 * 60 * 60,
             path="/api/v1/auth/refresh",
         )
-    
+
     return TokenResponse(
         access_token=result.access_token,
         token_type=result.token_type,
@@ -185,6 +186,7 @@ async def gmail_oauth_url() -> GmailUrlResponse:
 async def gmail_callback(
     body: GmailCallbackRequest,
     request: Request,
+    response: Response,
     db: AsyncSession = Depends(get_db),
 ) -> GmailCallbackResponse:
     """Exchange the Google OAuth authorization code.
@@ -218,6 +220,21 @@ async def gmail_callback(
 
     # Unauthenticated flow – login / register via Gmail.
     result = await auth_service.handle_gmail_login(db, body.code)
+
+    # Mirror the login endpoint: issue a refresh token and deliver it as an
+    # httpOnly cookie so the client can silently renew the access token.
+    if result.get("access_token"):
+        refresh_token = create_refresh_token({"sub": result["email"]})
+        response.set_cookie(
+            key="refresh_token",
+            value=refresh_token,
+            httponly=True,
+            secure=settings.is_production,
+            samesite="strict",
+            max_age=7 * 24 * 60 * 60,  # 7 days
+            path="/api/v1/auth/refresh",
+        )
+
     return GmailCallbackResponse(
         status=result["status"],
         email=result["email"],

backend/app/api/routes/emails.py

@@ -100,6 +100,34 @@ async def get_sync_status(
     }
 
 
+@router.get("/stats", response_model=EmailStatsResponse, summary="Get email counts by status")
+async def get_email_stats(
+    db: AsyncSession = Depends(get_db),
+    user: User = Depends(get_current_user),
+) -> EmailStatsResponse:
+    """Get email counts by status for the current user."""
+    result = await db.execute(
+        select(Email.status, func.count())
+        .where(Email.user_id == user.id)
+        .group_by(Email.status)
+    )
+
+    counts = {status: 0 for status in EmailStatus}
+    for status, count in result.all():
+        counts[status] = count
+
+    return EmailStatsResponse(
+        pending=counts.get(EmailStatus.PENDING, 0),
+        processing=counts.get(EmailStatus.PROCESSING, 0),
+        drafted=counts.get(EmailStatus.DRAFTED, 0),
+        needs_review=counts.get(EmailStatus.NEEDS_REVIEW, 0),
+        approved=counts.get(EmailStatus.APPROVED, 0),
+        sent=counts.get(EmailStatus.SENT, 0),
+        rejected=counts.get(EmailStatus.REJECTED, 0),
+        total=sum(counts.values()),
+    )
+
+
 @router.get(
     "/{email_id}",
     response_model=EmailDetailResponse,
@@ -133,31 +161,3 @@ async def process_email(
         trace_id=result["trace_id"],
         status=result["status"],
     )
-
-
-@router.get("/stats", response_model=EmailStatsResponse, summary="Get email counts by status")
-async def get_email_stats(
-    db: AsyncSession = Depends(get_db),
-    user: User = Depends(get_current_user),
-) -> EmailStatsResponse:
-    """Get email counts by status for the current user."""
-    result = await db.execute(
-        select(Email.status, func.count())
-        .where(Email.user_id == user.id)
-        .group_by(Email.status)
-    )
-
-    counts = {status: 0 for status in EmailStatus}
-    for status, count in result.all():
-        counts[status] = count
-
-    return EmailStatsResponse(
-        pending=counts.get(EmailStatus.PENDING, 0),
-        processing=counts.get(EmailStatus.PROCESSING, 0),
-        drafted=counts.get(EmailStatus.DRAFTED, 0),
-        needs_review=counts.get(EmailStatus.NEEDS_REVIEW, 0),
-        approved=counts.get(EmailStatus.APPROVED, 0),
-        sent=counts.get(EmailStatus.SENT, 0),
-        rejected=counts.get(EmailStatus.REJECTED, 0),
-        total=sum(counts.values()),
-    )

backend/app/api/routes/notifications.py

@@ -11,6 +11,7 @@
 from fastapi import APIRouter, WebSocket, WebSocketDisconnect
 from jose import JWTError, jwt
 from sqlalchemy import select
+from starlette.websockets import WebSocketState
 
 from app.core.config import settings
 from app.core.database import get_db
@@ -97,7 +98,7 @@ async def _heartbeat(websocket: WebSocket) -> None:
         while True:
             await asyncio.sleep(HEARTBEAT_INTERVAL)
             # Check if connection is still open before sending
-            if websocket.client_state.CONNECTED:
+            if websocket.client_state == WebSocketState.CONNECTED:
                 await websocket.send_json({"type": "ping"})
     except asyncio.CancelledError:
         # Task was cancelled, exit gracefully

backend/app/services/auth_service.py

@@ -22,7 +22,7 @@
 )
 from app.integrations.gmail.oauth import exchange_code
 from app.models.user import User
-from app.schemas.auth import TokenResponse
+from app.schemas.auth import TokenResponse, TokenResponseWithRefresh
 
 logger = logging.getLogger(__name__)
 
@@ -53,8 +53,8 @@ async def register(db: AsyncSession, email: str, password: str) -> User:
     return user
 
 
-async def login(db: AsyncSession, email: str, password: str) -> TokenResponse:
-    """Verify credentials and return a JWT access token.
+async def login(db: AsyncSession, email: str, password: str) -> TokenResponseWithRefresh:
+    """Verify credentials and return a JWT access token plus refresh token.
 
     Raises HTTP 401 if the email is not found or the password is wrong.
     """
@@ -72,7 +72,7 @@ async def login(db: AsyncSession, email: str, password: str) -> TokenResponse:
     access_token = create_access_token(token_data)
     refresh_token = create_refresh_token(token_data)
     logger.info("User logged in email=%s", email)
-    return TokenResponse(
+    return TokenResponseWithRefresh(
         access_token=access_token,
         refresh_token=refresh_token,
         token_type="bearer",
@@ -282,7 +282,6 @@ async def refresh_user_gmail_token(db: AsyncSession, user_id: uuid.UUID) -> bool
         True if token was refreshed (or still valid), False if refresh failed
         or user has no refresh token.
     """
-    import time
 
     from app.core.config import settings
 

backend/app/services/email_service.py

@@ -260,6 +260,8 @@ def on_token_refresh(updated_creds: dict[str, Any]) -> None:
     except Exception as crm_exc:
         logger.warning("CRM auto-populate failed: %s", crm_exc)
 
+    return {"fetched": len(raw_emails), "created": created}
+
 
 def _enrich_and_create_contact(
     crm,

backend/app/services/scheduler.py

@@ -40,7 +40,7 @@ async def _sync_all_users() -> None:
     for user in users:
         try:
             async with async_session_factory() as db:
-                sync_result = await _sync_emails_core(db, user)
+                sync_result = await _sync_emails_core(db, user) or {}
                 await db.commit()
 
                 # Auto-classify if new emails were found