Authentication Flow Improvements: fix token refresh handling

- Implement automatic token refresh when tokens expire
- Tokens now refresh transparently without user re-authentication
- Add comprehensive refresh token tests

Author: eranco74

assisted_service_mcp/src/mcp.py

@@ -65,7 +65,9 @@ def _create_oauth_token(self) -> Callable[[Any], Optional[str]]:
             Function that can extract OAuth tokens from MCP context
         """
 
-        def get_oauth_token(mcp: Any) -> Optional[str]:
+        def get_oauth_token(
+            mcp: Any,
+        ) -> Optional[str]:  # pylint: disable=too-many-locals
             if not settings.OAUTH_ENABLED:
                 return None
 
@@ -86,11 +88,24 @@ def get_oauth_token(mcp: Any) -> Optional[str]:
                 # Get client identifier for OAuth flow
                 client_id = self._get_mcp_client_identifier(mcp)
 
-                # Check if we have a completed OAuth token for this client
-                token = oauth_manager.token_store.get_access_token_by_client(client_id)
-                if token:
-                    log.info("Using cached OAuth token for MCP client %s", client_id)
-                    return token
+                # Try to get token (with automatic refresh if expired)
+                # Use async method which handles refresh, wrapped in asyncio.run for sync context
+                try:
+                    token = asyncio.run(
+                        oauth_manager.get_access_token_by_client(client_id)
+                    )
+                    if token:
+                        log.info(
+                            "Using cached OAuth token for MCP client %s", client_id
+                        )
+                        return token
+                except Exception as e:
+                    log.debug(
+                        "Could not get/refresh token for client %s: %s", client_id, e
+                    )
+                    # Fall through to start new OAuth flow
+
+                # No valid token found - need to start OAuth flow
 
                 # Check if OAuth flow is already in progress for this client
                 for (

assisted_service_mcp/src/oauth/manager.py

@@ -23,6 +23,10 @@
 )
 from assisted_service_mcp.src.settings import settings
 
+# Token expiration constants
+DEFAULT_TOKEN_EXPIRES_IN = 900  # 15 minutes in seconds
+TOKEN_EXPIRY_BUFFER = 300  # 5 minutes safety margin before expiration
+
 
 class OAuthManager:
     """Manages OAuth authentication flow for the MCP server.
@@ -214,14 +218,14 @@ async def exchange_code_for_token(
 
             # Create token object
             token_id = secrets.token_hex(16)
-            expires_in = token_data.get("expires_in", 3600)
+            expires_in = token_data.get("expires_in", DEFAULT_TOKEN_EXPIRES_IN)
 
             token = OAuthToken(
                 token_id=token_id,
                 client_id=state.client_id,
                 access_token=token_data["access_token"],
                 refresh_token=token_data.get("refresh_token"),
-                expires_at=time.time() + expires_in - 300,  # 5 min safety margin
+                expires_at=time.time() + expires_in - TOKEN_EXPIRY_BUFFER,
             )
 
             # Store token
@@ -248,7 +252,8 @@ async def get_access_token_by_id(self, token_id: str) -> Optional[str]:
         Returns:
             Access token if found and valid, None otherwise
         """
-        token = self.token_store.get_token_by_id(token_id)
+        # Get token including expired ones (for refresh purposes)
+        token = self.token_store.get_token_by_id(token_id, include_expired=True)
         if not token:
             return None
 
@@ -260,6 +265,8 @@ async def get_access_token_by_id(self, token_id: str) -> Optional[str]:
                 token = self.token_store.get_token_by_id(token_id)
                 return token.access_token if token else None
             log.warning("Failed to refresh token %s", token_id)
+            # Clean up expired token if refresh failed
+            self.token_store.remove_token(token_id)
             return None
 
         return token.access_token
@@ -287,7 +294,8 @@ async def get_access_token_by_client(self, client_id: str) -> Optional[str]:
         Returns:
             Access token if found and valid, None otherwise
         """
-        token = self.token_store.get_token_by_client(client_id)
+        # Get token including expired ones (for refresh purposes)
+        token = self.token_store.get_token_by_client(client_id, include_expired=True)
         if not token:
             return None
 
@@ -299,6 +307,8 @@ async def get_access_token_by_client(self, client_id: str) -> Optional[str]:
                 token = self.token_store.get_token_by_client(client_id)
                 return token.access_token if token else None
             log.warning("Failed to refresh token for client %s", client_id)
+            # Clean up expired token if refresh failed
+            self.token_store.remove_client_token(client_id)
             return None
 
         return token.access_token
@@ -333,8 +343,8 @@ async def _refresh_token(self, token: OAuthToken) -> bool:
             # Update token in store
             new_access_token = token_data["access_token"]
             new_refresh_token = token_data.get("refresh_token", token.refresh_token)
-            expires_in = token_data.get("expires_in", 3600)
-            new_expires_at = time.time() + expires_in - 300
+            expires_in = token_data.get("expires_in", DEFAULT_TOKEN_EXPIRES_IN)
+            new_expires_at = time.time() + expires_in - TOKEN_EXPIRY_BUFFER
 
             self.token_store.update_token(
                 token.token_id, new_access_token, new_refresh_token, new_expires_at
@@ -536,7 +546,7 @@ async def oauth_token_handler(request: Request) -> Dict[str, Any]:
             return {
                 "access_token": token.get("access_token"),
                 "token_type": token.get("token_type", "Bearer"),
-                "expires_in": token.get("expires_in", 3600),
+                "expires_in": token.get("expires_in", DEFAULT_TOKEN_EXPIRES_IN),
                 "refresh_token": token.get("refresh_token"),
                 "scope": token.get("scope", "openid profile email"),
             }

assisted_service_mcp/src/oauth/middleware.py

@@ -42,8 +42,8 @@ async def handle_mcp_request(self, request: Request, call_next: Any) -> Response
 
         client_id = self._get_client_identifier(request)
 
-        # Try to use existing token
-        token = oauth_manager.token_store.get_access_token_by_client(client_id)
+        # Try to use existing token (with automatic refresh if expired)
+        token = await oauth_manager.get_access_token_by_client(client_id)
         if token:
             log.info("Using cached token for client %s", client_id)
             return await self._create_authenticated_request(request, call_next, token)
@@ -176,7 +176,7 @@ async def _wait_for_oauth_completion(
             waited_time += poll_interval
 
             # Check if OAuth completed (token available for client)
-            token = oauth_manager.token_store.get_access_token_by_client(client_id)
+            token = await oauth_manager.get_access_token_by_client(client_id)
             if token:
                 log.info(
                     "OAuth completed for client %s, proceeding with request", client_id

assisted_service_mcp/src/oauth/store.py

@@ -47,43 +47,49 @@ def store_token(self, token: OAuthToken) -> None:
                 token.expires_at,
             )
 
-    def get_token_by_id(self, token_id: str) -> Optional[OAuthToken]:
+    def get_token_by_id(
+        self, token_id: str, include_expired: bool = False
+    ) -> Optional[OAuthToken]:
         """Get token by token ID.
 
         Args:
             token_id: Token identifier
+            include_expired: If True, return expired tokens (for refresh purposes)
 
         Returns:
-            OAuthToken if found and valid, None otherwise
+            OAuthToken if found and valid (or expired if include_expired=True), None otherwise
         """
         with self._lock:
             token = self._tokens.get(token_id)
             if not token:
                 return None
 
-            if token.is_expired():
+            if token.is_expired() and not include_expired:
                 log.debug("Token %s is expired, removing from store", token_id)
                 self._remove_token_unsafe(token_id)
                 return None
 
             return token
 
-    def get_token_by_client(self, client_id: str) -> Optional[OAuthToken]:
+    def get_token_by_client(
+        self, client_id: str, include_expired: bool = False
+    ) -> Optional[OAuthToken]:
         """Get token for a client.
 
         Args:
             client_id: Client identifier
+            include_expired: If True, return expired tokens (for refresh purposes)
 
         Returns:
-            OAuthToken if found and valid, None otherwise
+            OAuthToken if found and valid (or expired if include_expired=True), None otherwise
         """
         with self._lock:
             token_id = self._client_tokens.get(client_id)
             if not token_id:
                 return None
 
             # Re-use get_token_by_id which also acquires lock (RLock allows re-entrance)
-            return self.get_token_by_id(token_id)
+            return self.get_token_by_id(token_id, include_expired=include_expired)
 
     def get_access_token_by_id(self, token_id: str) -> Optional[str]:
         """Get access token string by token ID.