Added refresh token logic

Author: lucaseduoli

src/auth/ibm_auth.py

@@ -6,10 +6,13 @@
                   IBM_JWT_PUBLIC_KEY_URL is configured.
 fetch_ibm_public_key — fetch and cache IBM's public key PEM.
 """
+import asyncio
+
 import httpx
 import jwt
 from cryptography.hazmat.primitives.serialization import load_pem_public_key
 
+from config.settings import PLATFORM_REFRESH_URL
 from utils.logging_config import get_logger
 
 logger = get_logger(__name__)
@@ -88,3 +91,30 @@ def validate_ibm_jwt(token: str, public_key) -> dict | None:
     except jwt.InvalidTokenError as exc:
         logger.warning("IBM JWT validation failed", error=str(exc))
         return None
+
+async def refresh_ibm_jwt(token: str) -> str | None:
+    """Refresh the IBM JWT token using the configured PLATFORM_REFRESH_URL."""
+    if not PLATFORM_REFRESH_URL:
+        return None
+
+    headers = {
+        "Authorization": f"Bearer {token}",
+        "User-Agent": "curl/7.64.1"
+    }
+
+    async with httpx.AsyncClient() as client:
+        for attempt in range(10):
+            try:
+                resp = await client.post(PLATFORM_REFRESH_URL, headers=headers, timeout=30.0)
+                if resp.status_code == 200:
+                    data = resp.json()
+                    new_token = data.get("token") or data.get("access_token") or data.get("refresh_token") or resp.text.strip().strip('"')
+                    if '.' in new_token:
+                        logger.info("Successfully refreshed IBM JWT token.")
+                        return new_token
+            except Exception as e:
+                logger.warning(f"Failed to refresh IBM token (attempt {attempt + 1}): {e}")
+            
+            await asyncio.sleep(10)
+    
+    return None

src/config/settings.py

@@ -63,6 +63,8 @@
 IBM_AUTH_ENABLED = os.getenv("IBM_AUTH_ENABLED", "false").lower() in ("true", "1", "yes")
 PLATFORM_USERNAME = os.getenv("PLATFORM_USERNAME")
 PLATFORM_PASSWORD = os.getenv("PLATFORM_PASSWORD")
+PLATFORM_REFRESH_URL = os.getenv("PLATFORM_REFRESH_URL")
+PLATFORM_REFRESH_PERIOD = get_env_int("PLATFORM_REFRESH_PERIOD", 60)
 IBM_JWT_PUBLIC_KEY_URL = os.getenv("IBM_JWT_PUBLIC_KEY_URL", "")
 IBM_SESSION_COOKIE_NAME = os.getenv("IBM_SESSION_COOKIE_NAME", "ibm-openrag-session")
 IBM_CREDENTIALS_HEADER = os.getenv("IBM_CREDENTIALS_HEADER", "X-IBM-LH-Credentials")

src/dependencies.py

@@ -21,7 +21,7 @@ async def my_endpoint(
 from typing import Optional
 
 from cachetools import TTLCache
-from fastapi import Depends, HTTPException, Request
+from fastapi import Depends, HTTPException, Request, Response
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from session_manager import User
@@ -342,7 +342,7 @@ async def _dep(
 # ─────────────────────────────────────────────
 
 
-async def _get_ibm_user(request: Request, required: bool) -> Optional["User"]:
+async def _get_ibm_user(request: Request, response: Response, required: bool) -> Optional["User"]:
     """Authenticate via IBM AMS.
 
     0. X-IBM-LH-Credentials header (configurable via IBM_CREDENTIALS_HEADER) —
@@ -407,6 +407,32 @@ async def _get_ibm_user(request: Request, required: bool) -> Optional["User"]:
                 user_id = claims.get("username", sub)
                 email = claims.get("username", sub)
                 name = claims.get("display_name", claims.get("username", sub))
+                
+                exp = claims.get("exp")
+                if exp:
+                    import datetime
+
+                    from config.settings import PLATFORM_REFRESH_PERIOD
+                    
+                    now = datetime.datetime.now(datetime.UTC).timestamp()
+                    if exp - now < PLATFORM_REFRESH_PERIOD:
+                        logger.info("IBM JWT token is expiring soon, attempting refresh...")
+                        new_token = await ibm_auth.refresh_ibm_jwt(ibm_token)
+                        if new_token:
+                            ibm_token = new_token
+                            response.set_cookie(
+                                key=IBM_SESSION_COOKIE_NAME,
+                                value=new_token,
+                                httponly=True,
+                                secure=True,
+                                samesite="lax",
+                            )
+                            logger.info("Successfully refreshed IBM JWT token and updated cookie.")
+                        else:
+                            if now >= exp:
+                                if required:
+                                    raise HTTPException(status_code=401, detail="IBM session expired. Please log in again.")
+                                return None
 
     if lh_credentials and lh_credentials.strip() != "":
         logger.debug("[AUTH] IBM LH credentials found in request headers")
@@ -518,6 +544,7 @@ async def _get_ibm_user(request: Request, required: bool) -> Optional["User"]:
 
 async def get_current_user(
     request: Request,
+    response: Response,
     session_manager=Depends(get_session_manager),
 ) -> User:
     """
@@ -532,7 +559,7 @@ async def get_current_user(
     # IBM AMS cookie auth takes priority when enabled
     if IBM_AUTH_ENABLED:
         logger.debug("[AUTH] IBM auth mode enabled, getting current user")
-        user = await _get_ibm_user(request, required=True)
+        user = await _get_ibm_user(request, response, required=True)
         if user and user.user_id and user.user_id not in session_manager.users:
             session_manager.users[user.user_id] = user
         return await _attach_db_user_id(request, user)
@@ -560,6 +587,7 @@ async def get_current_user(
 
 async def get_optional_user(
     request: Request,
+    response: Response,
     session_manager=Depends(get_session_manager),
 ) -> User | None:
     """
@@ -574,7 +602,7 @@ async def get_optional_user(
     # IBM AMS cookie auth takes priority when enabled
     if IBM_AUTH_ENABLED:
         logger.debug("[AUTH] IBM auth mode enabled, getting optional user")
-        user = await _get_ibm_user(request, required=False)
+        user = await _get_ibm_user(request, response, required=False)
         if user and user.user_id and user.user_id not in session_manager.users:
             session_manager.users[user.user_id] = user
         if user: