33import atexit
44import inspect
55import json
6+ import operator
67import tempfile
78import threading
89import time
@@ -463,71 +464,78 @@ def _get_token(self, convert_timestamps: bool = True) -> dict[str, Any]:
463464 return token
464465
465466 def _refresh_access_token (self ) -> tuple [str , float ]:
466- # First check if a token cache exists on disk. If yes, find and use:
467- # - A valid access token.
468- # - A valid refresh token, and if so, use it automatically to redeem a new access token.
467+ # Token resolution order (cheapest option first):
468+ # 1. Valid access token in cache → use directly, no network call
469+ # 2. Refresh token in cache → exchange for new AT, one network call
470+ # 3. Device code flow → interactive, requires user action
469471 credentials = None
470- for token in self .__app .token_cache .search (self .__app .token_cache .CredentialType .REFRESH_TOKEN ):
471- if "expires_on" in token and token ["expires_on" ] > time .time ():
472- credentials = token
472+
473+ # 1. Check for a still-valid access token. search() does NOT filter by expiry,
474+ # so we check manually and respect the leeway to avoid handing out a near-expired token.
475+ for token in self .__app .token_cache .search (
476+ self .__app .token_cache .CredentialType .ACCESS_TOKEN ,
477+ query = {"client_id" : self .client_id },
478+ ):
479+ expiry = int (token .get ("expires_on" , 0 )) - time .time () - self .token_expiry_leeway_seconds
480+ if expiry > 0 :
481+ credentials = {"access_token" : token ["secret" ], "expires_in" : expiry }
473482 break
483+
484+ # 2. No valid AT — try to silently redeem a refresh token.
485+ if credentials is None :
486+ rt_entry = None
487+ for token in self .__app .token_cache .search (
488+ self .__app .token_cache .CredentialType .REFRESH_TOKEN ,
489+ query = {"client_id" : self .client_id },
490+ ):
491+ rt_entry = token
492+ break # MSAL RTs have no 'expires_on'; use the first found
493+
494+ if rt_entry is not None :
495+ # Pass the full RT cache entry (not just the secret string) so MSAL's
496+ # on_removing_rt callback can properly remove it on invalid_grant.
497+ # Exclude OIDC meta-scopes that are not valid in token-endpoint requests.
498+ oidc_scopes = frozenset ({"openid" , "profile" , "email" , "offline_access" })
499+ resp = self .__app .client .obtain_token_by_refresh_token (
500+ rt_entry ,
501+ rt_getter = operator .itemgetter ("secret" ),
502+ scope = " " .join (s for s in self .__scopes if s not in oidc_scopes ),
503+ )
504+ if isinstance (resp , dict ) and "error" not in resp :
505+ credentials = resp
506+ # else: RT rejected by server, fall through to device code flow
507+
474508 if credentials is not None :
475- credentials = self .__app .client .obtain_token_by_refresh_token (credentials .get ("secret" , "" ))
476- else :
477- for token in self .__app .token_cache .search (self .__app .token_cache .CredentialType .ACCESS_TOKEN ):
478- if expiry := int (token .get ("expires_on" , 0 )) - time .time () > 0 :
479- credentials = {
480- "access_token" : token .get ("secret" ),
481- "expires_in" : expiry ,
482- }
483- break
484- # If we're unable to find (or acquire a new) access token, we initiate the device code auth flow.
509+ self ._verify_credentials (credentials )
510+ return credentials ["access_token" ], time .time () + float (credentials ["expires_in" ])
511+
512+ # 3. If we're unable to find (or acquire a new) access token, we initiate the device code auth flow.
485513 # The msal device_code flow does not support setting the audience, so we need to handle it manually.
486514 # We use the http client instantiated as part of the msal client, as well as the details found
487515 # in oauth discovery.
488- if credentials is None :
489- data = {
490- "scope" : self .scope_string (),
491- "client_id" : self .client_id ,
492- }
493- for key , value in self .__token_custom_args .items ():
494- data [key ] = value
495-
496- device_flow_endpoint = self ._get_device_authorization_endpoint ()
497- device_flow_response = self ._get_device_code_response (device_flow_endpoint , data )
498- if "verification_uri" in device_flow_response :
499- print ( # noqa: T201
500- f"Visit { device_flow_response ['verification_uri' ]} and enter the code: { device_flow_response .get ('user_code' , 'ERROR' )} "
501- )
502- elif "message" in device_flow_response :
503- print ( # noqa: T201
504- f"Device code: { device_flow_response .get ('message' , device_flow_response .get ('user_code' , 'ERROR' ))} "
505- )
506- else :
507- raise CogniteAuthError (
508- f"Error initiating device flow: { device_flow_response .get ('error' )} - { device_flow_response .get ('error_description' )} "
509- )
510- if "interval" not in device_flow_response :
511- # Set default interval according to standard
512- device_flow_response ["interval" ] = 5
513- if "expires_in" in device_flow_response :
514- # msal library uses expires_at instead of the standard expires_in
515- device_flow_response ["expires_at" ] = device_flow_response ["expires_in" ] + time .time ()
516- # Poll for token
517- credentials = self .__app .client .obtain_token_by_device_flow (
518- flow = device_flow_response ,
519- data = dict (
520- data ,
521- code = device_flow_response .get (
522- "device_code"
523- ), # Hack from msal library to get the code from the device flow, not standard
524- ),
516+ data = {"scope" : self .scope_string (), "client_id" : self .client_id , ** self .__token_custom_args }
517+ device_flow_endpoint = self ._get_device_authorization_endpoint ()
518+ response = self ._get_device_code_response (device_flow_endpoint , data )
519+ if "verification_uri" in response :
520+ print (f"Visit { response ['verification_uri' ]} and enter the code: { response .get ('user_code' , 'ERROR' )} " ) # noqa: T201
521+ elif "message" in response :
522+ print (f"Device code: { response .get ('message' , response .get ('user_code' , 'ERROR' ))} " ) # noqa: T201
523+ else :
524+ raise CogniteAuthError (
525+ f"Error initiating device flow: { response .get ('error' )} - { response .get ('error_description' )} "
525526 )
526-
527- self ._verify_credentials (credentials )
528- self .__app .token_cache .add (
529- dict (credentials , environment = self .__app .authority .instance ),
527+ if "interval" not in response :
528+ response ["interval" ] = 5 # Set default interval according to standard
529+ if "expires_in" in response :
530+ # msal library uses expires_at instead of the standard expires_in
531+ response ["expires_at" ] = float (response ["expires_in" ]) + time .time ()
532+
533+ credentials = self .__app .client .obtain_token_by_device_flow (
534+ flow = response ,
535+ # Hack from msal library to get the code from the device flow, not standard:
536+ data = dict (data , code = response .get ("device_code" )),
530537 )
538+ self ._verify_credentials (credentials )
531539 return credentials ["access_token" ], time .time () + float (credentials ["expires_in" ])
532540
533541 @classmethod
@@ -579,18 +587,16 @@ def default_for_azure_ad(
579587 mem_cache_only : bool = False ,
580588 ) -> OAuthDeviceCode :
581589 """
582- Create an OAuthDeviceCode instance for Azure with default URLs and scopes. It uses the pre-configured Cognite
583- app registration for device code flow. If you need device code flow with another app registration, instantiate
584- OAuthDeviceCode directly.
590+ Create an OAuthDeviceCode instance for Azure with default URLs and scopes.
585591
586592 The default configuration creates the URLs based on the tenant id and cluster:
587593
588594 * Authority URL: "https://login.microsoftonline.com/{tenant_id}"
589- * Scopes: [f"https://{cdf_cluster}.cognitedata.com/.default "]
595+ * Scopes: [f"https://{cdf_cluster}.cognitedata.com/IDENTITY", f"https://{cdf_cluster}.cognitedata.com/user_impersonation", "profile", "openid", "offline_access "]
590596
591597 Args:
592598 tenant_id (str): The Azure tenant id
593- client_id (str): An app registration that allows device code flow.
599+ client_id (str): Your app registration client id. Must have device code flow enabled .
594600 cdf_cluster (str): The CDF cluster where the CDF project is located.
595601 token_cache_path (Path | None): Location to store token cache, defaults to os temp directory/cognitetokencache.{client_id}.bin.
596602 token_expiry_leeway_seconds (int): The token is refreshed at the earliest when this number of seconds is left before expiry. Default: 30 sec
@@ -601,18 +607,18 @@ def default_for_azure_ad(
601607 """
602608 return cls (
603609 authority_url = f"https://login.microsoftonline.com/{ tenant_id } " ,
604- client_id = client_id , # Default application for CDF API for device code flow
610+ client_id = client_id ,
605611 scopes = [
606612 f"https://{ cdf_cluster } .cognitedata.com/IDENTITY" ,
607613 f"https://{ cdf_cluster } .cognitedata.com/user_impersonation" ,
608614 "profile" ,
609615 "openid" ,
616+ "offline_access" , # required for Azure to issue a refresh token
610617 ],
611618 token_cache_path = token_cache_path ,
612619 token_expiry_leeway_seconds = token_expiry_leeway_seconds ,
613620 clear_cache = clear_cache ,
614621 mem_cache_only = mem_cache_only ,
615- audience = f"https://{ cdf_cluster } .cognitedata.com" ,
616622 )
617623
618624
@@ -681,9 +687,8 @@ def scopes(self) -> list[str]:
681687 return self .__scopes
682688
683689 def _refresh_access_token (self ) -> tuple [str , float ]:
684- # First check if a token cache exists on disk. If yes, find and use:
685- # - A valid access token.
686- # - A valid refresh token, and if so, use it automatically to redeem a new access token.
690+ # Try the in-memory token cache silently (MSAL checks AT first, then RT automatically).
691+ # Falls through to interactive flow if nothing usable is found.
687692 credentials = None
688693 if accounts := self .__app .get_accounts ():
689694 credentials = self .__app .acquire_token_silent (scopes = self .__scopes , account = accounts [0 ])
0 commit comments