diff --git a/msal/application.py b/msal/application.py index 085914c4..dd65ac6c 100644 --- a/msal/application.py +++ b/msal/application.py @@ -246,7 +246,8 @@ def _build_client(self, client_credential, authority): default_body=default_body, client_assertion=client_assertion, client_assertion_type=client_assertion_type, - on_obtaining_tokens=self.token_cache.add, + on_obtaining_tokens=lambda event: self.token_cache.add(dict( + event, environment=authority.instance)), on_removing_rt=self.token_cache.remove_rt, on_updating_rt=self.token_cache.update_rt) diff --git a/msal/token_cache.py b/msal/token_cache.py index 6884075d..83fc1891 100644 --- a/msal/token_cache.py +++ b/msal/token_cache.py @@ -126,6 +126,8 @@ def __add(self, event, now=None): environment = realm = None if "token_endpoint" in event: _, environment, realm = canonicalize(event["token_endpoint"]) + if "environment" in event: # Always available unless in legacy test cases + environment = event["environment"] # Set by application.py response = event.get("response", {}) data = event.get("data", {}) access_token = response.get("access_token")