refac

Author: taylorwilsdon

auth/auth_info_middleware.py

@@ -8,7 +8,7 @@
 from fastmcp.server.dependencies import get_access_token
 from fastmcp.server.dependencies import get_http_headers
 
-from auth.external_oauth_provider import SESSION_TIME
+from auth.external_oauth_provider import get_session_time
 from auth.oauth21_session_store import ensure_session_from_access_token
 from auth.oauth_types import WorkspaceAccessToken
 
@@ -134,7 +134,8 @@ async def _process_request_for_auth(self, context: MiddlewareContext):
                                                 session_id=f"google_oauth_{token_str[:8]}",
                                                 expires_at=verified_expires
                                                 if verified_expires is not None
-                                                else int(time.time()) + SESSION_TIME,
+                                                else int(time.time())
+                                                + get_session_time(),
                                                 claims=getattr(
                                                     verified_auth, "claims", {}
                                                 )

auth/external_oauth_provider.py

@@ -11,7 +11,6 @@
 import logging
 import os
 import time
-from datetime import datetime, timedelta, timezone
 from typing import Optional
 
 from starlette.routing import Route
@@ -45,7 +44,9 @@ def _get_session_time() -> int:
     return max(value, 1)
 
 
-SESSION_TIME = _get_session_time()
+def get_session_time() -> int:
+    """Return the configured session time in seconds."""
+    return _get_session_time()
 
 
 class ExternalOAuthProvider(GoogleProvider):
@@ -98,19 +99,18 @@ async def verify_token(self, token: str) -> Optional[AccessToken]:
                 from auth.google_auth import get_user_info
 
                 # Create minimal Credentials object for userinfo API call
-                # expiry must be set so credentials.valid returns True
                 credentials = Credentials(
                     token=token,
                     token_uri="https://oauth2.googleapis.com/token",
                     client_id=self._client_id,
                     client_secret=self._client_secret,
-                    expiry=datetime.now(timezone.utc) + timedelta(seconds=SESSION_TIME),
                 )
 
                 # Validate token by calling userinfo API
-                user_info = get_user_info(credentials)
+                user_info = get_user_info(credentials, skip_valid_check=True)
 
                 if user_info and user_info.get("email"):
+                    session_time = get_session_time()
                     # Token is valid - create AccessToken object
                     logger.info(
                         f"Validated external access token for: {user_info['email']}"
@@ -120,7 +120,7 @@ async def verify_token(self, token: str) -> Optional[AccessToken]:
                     access_token = WorkspaceAccessToken(
                         token=token,
                         scopes=scope_list,
-                        expires_at=int(time.time()) + SESSION_TIME,
+                        expires_at=int(time.time()) + session_time,
                         claims={
                             "email": user_info["email"],
                             "sub": user_info.get("id"),

auth/google_auth.py

@@ -759,11 +759,17 @@ def get_credentials(
         return None
 
 
-def get_user_info(credentials: Credentials) -> Optional[Dict[str, Any]]:
+def get_user_info(
+    credentials: Credentials, *, skip_valid_check: bool = False
+) -> Optional[Dict[str, Any]]:
     """Fetches basic user profile information (requires userinfo.email scope)."""
-    if not credentials or not credentials.valid:
-        logger.error("Cannot get user info: Invalid or missing credentials.")
+    if not credentials:
+        logger.error("Cannot get user info: Missing credentials.")
+        return None
+    if not skip_valid_check and not credentials.valid:
+        logger.error("Cannot get user info: Invalid credentials.")
         return None
+    service = None
     try:
         # Using googleapiclient discovery to get user info
         # Requires 'google-api-python-client' library