From 3cb5121163705e7f730d844fa8893291a6f43c99 Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Fri, 6 Mar 2026 14:29:50 -0800 Subject: [PATCH 01/10] Adding PyJWKClient caching --- .../core/authorization/jwt_token_validator.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py index e64e0987..1aa02f36 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py @@ -12,11 +12,19 @@ logger = logging.getLogger(__name__) - class JwtTokenValidator: + + _jwk_clients_cache: dict[str, PyJWKClient] = {} + def __init__(self, configuration: AgentAuthConfiguration): self.configuration = configuration + @staticmethod + def _get_jwk_client(jwks_uri: str) -> PyJWKClient: + if jwks_uri not in JwtTokenValidator._jwk_clients_cache: + JwtTokenValidator._jwk_clients_cache[jwks_uri] = PyJWKClient(jwks_uri) + return JwtTokenValidator._jwk_clients_cache[jwks_uri] + async def validate_token(self, token: str) -> ClaimsIdentity: logger.debug("Validating JWT token.") @@ -49,8 +57,7 @@ async def _get_public_key_or_secret(self, token: str) -> PyJWK: if unverified_payload.get("iss") == "https://api.botframework.com" else f"https://login.microsoftonline.com/{self.configuration.TENANT_ID}/discovery/v2.0/keys" ) - jwks_client = PyJWKClient(jwksUri) - + jwks_client = JwtTokenValidator._get_jwk_client(jwksUri) key = await asyncio.to_thread(jwks_client.get_signing_key, header["kid"]) return key From 3155b54527086723535ec084d2beb40b35f25b80 Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Fri, 6 Mar 2026 14:33:26 -0800 Subject: [PATCH 02/10] Formatting and documentation --- .../core/authorization/jwt_token_validator.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py index 1aa02f36..1a5fb75f 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py @@ -12,20 +12,33 @@ logger = logging.getLogger(__name__) + class JwtTokenValidator: + """Utility class for validating JWT tokens using the PyJWT library and JWKs from a specified URI.""" _jwk_clients_cache: dict[str, PyJWKClient] = {} def __init__(self, configuration: AgentAuthConfiguration): + """Initializes the JwtTokenValidator with the given configuration. + + :param configuration: An instance of AgentAuthConfiguration containing the necessary settings for token validation. + """ self.configuration = configuration @staticmethod def _get_jwk_client(jwks_uri: str) -> PyJWKClient: + """Retrieves a PyJWKClient for the specified JWKS URI, using a cache to avoid redundant clients.""" if jwks_uri not in JwtTokenValidator._jwk_clients_cache: JwtTokenValidator._jwk_clients_cache[jwks_uri] = PyJWKClient(jwks_uri) return JwtTokenValidator._jwk_clients_cache[jwks_uri] async def validate_token(self, token: str) -> ClaimsIdentity: + """Validates a JWT token. + + :param token: The JWT token to validate. + :return: A ClaimsIdentity object containing the token's claims if validation is successful. + :raises ValueError: If the token is invalid or if the audience claim is not valid + """ logger.debug("Validating JWT token.") key = await self._get_public_key_or_secret(token) @@ -45,10 +58,12 @@ async def validate_token(self, token: str) -> ClaimsIdentity: return ClaimsIdentity(decoded_token, True, security_token=token) def get_anonymous_claims(self) -> ClaimsIdentity: + """Returns a ClaimsIdentity for an anonymous user.""" logger.debug("Returning anonymous claims identity.") return ClaimsIdentity({}, False, authentication_type="Anonymous") async def _get_public_key_or_secret(self, token: str) -> PyJWK: + """Retrieves the public key or secret for validating the JWT token.""" header = get_unverified_header(token) unverified_payload: dict = decode(token, options={"verify_signature": False}) @@ -57,7 +72,7 @@ async def _get_public_key_or_secret(self, token: str) -> PyJWK: if unverified_payload.get("iss") == "https://api.botframework.com" else f"https://login.microsoftonline.com/{self.configuration.TENANT_ID}/discovery/v2.0/keys" ) - jwks_client = JwtTokenValidator._get_jwk_client(jwksUri) + jwks_client = JwtTokenValidator._get_jwk_client(jwksUri) key = await asyncio.to_thread(jwks_client.get_signing_key, header["kid"]) return key From 504f38e7fec9394ef916fb7f7e2489127e92819b Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Fri, 6 Mar 2026 14:34:15 -0800 Subject: [PATCH 03/10] Removing unnecessary import --- .../hosting/core/authorization/jwt_token_validator.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py index 1a5fb75f..a3d63964 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py @@ -3,9 +3,13 @@ import asyncio import logging -import jwt -from jwt import PyJWKClient, PyJWK, decode, get_unverified_header +from jwt import ( + PyJWKClient, + PyJWK, + decode, + get_unverified_header +) from .agent_auth_configuration import AgentAuthConfiguration from .claims_identity import ClaimsIdentity @@ -42,7 +46,7 @@ async def validate_token(self, token: str) -> ClaimsIdentity: logger.debug("Validating JWT token.") key = await self._get_public_key_or_secret(token) - decoded_token = jwt.decode( + decoded_token = decode( token, key=key, algorithms=["RS256"], From d94b4bf5ac80edaf43ef79f15baf53a1a88713a4 Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Fri, 6 Mar 2026 14:35:04 -0800 Subject: [PATCH 04/10] More formatting --- .../hosting/core/authorization/jwt_token_validator.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py index a3d63964..178e3c6c 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py @@ -4,12 +4,7 @@ import asyncio import logging -from jwt import ( - PyJWKClient, - PyJWK, - decode, - get_unverified_header -) +from jwt import PyJWKClient, PyJWK, decode, get_unverified_header from .agent_auth_configuration import AgentAuthConfiguration from .claims_identity import ClaimsIdentity From ade8d587417c33a07da6374a0aad5203838208f8 Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Fri, 6 Mar 2026 15:30:27 -0800 Subject: [PATCH 05/10] Locking _get_jwk_client --- .../core/authorization/jwt_token_validator.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py index 178e3c6c..663ff277 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py @@ -3,6 +3,7 @@ import asyncio import logging +from typing import Any from jwt import PyJWKClient, PyJWK, decode, get_unverified_header @@ -16,6 +17,7 @@ class JwtTokenValidator: """Utility class for validating JWT tokens using the PyJWT library and JWKs from a specified URI.""" _jwk_clients_cache: dict[str, PyJWKClient] = {} + _get_signing_key_lock = asyncio.Lock() def __init__(self, configuration: AgentAuthConfiguration): """Initializes the JwtTokenValidator with the given configuration. @@ -25,11 +27,18 @@ def __init__(self, configuration: AgentAuthConfiguration): self.configuration = configuration @staticmethod - def _get_jwk_client(jwks_uri: str) -> PyJWKClient: + def _get_jwk_client(jwks_uri: str, header: dict[str, Any]) -> PyJWKClient: """Retrieves a PyJWKClient for the specified JWKS URI, using a cache to avoid redundant clients.""" if jwks_uri not in JwtTokenValidator._jwk_clients_cache: JwtTokenValidator._jwk_clients_cache[jwks_uri] = PyJWKClient(jwks_uri) return JwtTokenValidator._jwk_clients_cache[jwks_uri] + + @staticmethod + async def _get_signing_key(jwks_client: PyJWKClient, token: str) -> PyJWK: + async with JwtTokenValidator._get_signing_key_lock: + # get_signing_key is not guaranteed to be thread-safe, so we run it in a thread to avoid blocking the event loop + key = await asyncio.to_thread(jwks_client.get_signing_key, header["kid"]) + return key async def validate_token(self, token: str) -> ClaimsIdentity: """Validates a JWT token. @@ -66,12 +75,13 @@ async def _get_public_key_or_secret(self, token: str) -> PyJWK: header = get_unverified_header(token) unverified_payload: dict = decode(token, options={"verify_signature": False}) - jwksUri = ( + jwks_uri = ( "https://login.botframework.com/v1/.well-known/keys" if unverified_payload.get("iss") == "https://api.botframework.com" else f"https://login.microsoftonline.com/{self.configuration.TENANT_ID}/discovery/v2.0/keys" ) - jwks_client = JwtTokenValidator._get_jwk_client(jwksUri) - key = await asyncio.to_thread(jwks_client.get_signing_key, header["kid"]) + + jwks_client = JwtTokenValidator._get_jwk_client(jwks_uri, header) + key = await JwtTokenValidator._get_signing_key(jwks_client) return key From ceb4103672b43dc696a3e04ab8e30bcce8f2f31d Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Sun, 8 Mar 2026 19:44:14 -0700 Subject: [PATCH 06/10] Moving jwk client caching to its own class --- .../core/authorization/jwt_token_validator.py | 41 ++++++++++--------- 1 file changed, 22 insertions(+), 19 deletions(-) diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py index 663ff277..1a6d858d 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py @@ -12,12 +12,30 @@ logger = logging.getLogger(__name__) +class _JwkClientManager: + """Helper class to manage PyJWKClient instances for different JWKS URIs, with caching and thread safety.""" + + _cache: dict[str, PyJWKClient] + + def __init__(self): + self._cache = {} + + def _get_jwk_client(self, jwks_uri: str) -> PyJWKClient: + """Retrieves a PyJWKClient for the given JWKS URI, using a cache to avoid creating multiple clients for the same URI.""" + if jwks_uri not in self._cache: + self._cache[jwks_uri] = PyJWKClient(jwks_uri) + return self._cache[jwks_uri] + + async def get_signing_key(self, jwks_uri: str, header: dict[str, Any]) -> PyJWK: + """Retrieves the signing key from the JWK client for the given token header.""" + jwks_client = self._get_jwk_client(jwks_uri) + key = await asyncio.to_thread(jwks_client.get_signing_key, header["kid"]) + return key class JwtTokenValidator: """Utility class for validating JWT tokens using the PyJWT library and JWKs from a specified URI.""" - _jwk_clients_cache: dict[str, PyJWKClient] = {} - _get_signing_key_lock = asyncio.Lock() + _jwk_client_manager = _JwkClientManager() def __init__(self, configuration: AgentAuthConfiguration): """Initializes the JwtTokenValidator with the given configuration. @@ -26,20 +44,6 @@ def __init__(self, configuration: AgentAuthConfiguration): """ self.configuration = configuration - @staticmethod - def _get_jwk_client(jwks_uri: str, header: dict[str, Any]) -> PyJWKClient: - """Retrieves a PyJWKClient for the specified JWKS URI, using a cache to avoid redundant clients.""" - if jwks_uri not in JwtTokenValidator._jwk_clients_cache: - JwtTokenValidator._jwk_clients_cache[jwks_uri] = PyJWKClient(jwks_uri) - return JwtTokenValidator._jwk_clients_cache[jwks_uri] - - @staticmethod - async def _get_signing_key(jwks_client: PyJWKClient, token: str) -> PyJWK: - async with JwtTokenValidator._get_signing_key_lock: - # get_signing_key is not guaranteed to be thread-safe, so we run it in a thread to avoid blocking the event loop - key = await asyncio.to_thread(jwks_client.get_signing_key, header["kid"]) - return key - async def validate_token(self, token: str) -> ClaimsIdentity: """Validates a JWT token. @@ -80,8 +84,7 @@ async def _get_public_key_or_secret(self, token: str) -> PyJWK: if unverified_payload.get("iss") == "https://api.botframework.com" else f"https://login.microsoftonline.com/{self.configuration.TENANT_ID}/discovery/v2.0/keys" ) - - jwks_client = JwtTokenValidator._get_jwk_client(jwks_uri, header) - key = await JwtTokenValidator._get_signing_key(jwks_client) + + key = await self._jwk_client_manager.get_signing_key(jwks_uri, header) return key From 4713ac489e8678737d06f9109a3792879794f7eb Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Mon, 9 Mar 2026 03:17:24 -0700 Subject: [PATCH 07/10] Adding tests for token validation helper class --- .../core/authorization/jwt_token_validator.py | 29 +++- tests/hosting_core/authorization/__init__.py | 0 .../authorization/test_jwk_client_manager.py | 130 ++++++++++++++++++ 3 files changed, 152 insertions(+), 7 deletions(-) create mode 100644 tests/hosting_core/authorization/__init__.py create mode 100644 tests/hosting_core/authorization/test_jwk_client_manager.py diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py index 1a6d858d..21f6e8ea 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py @@ -4,6 +4,7 @@ import asyncio import logging from typing import Any +from dataclasses import dataclass from jwt import PyJWKClient, PyJWK, decode, get_unverified_header @@ -12,25 +13,39 @@ logger = logging.getLogger(__name__) +@dataclass +class _JwkClientCacheEntry: + + jwk_client: PyJWKClient + lock: asyncio.Lock + class _JwkClientManager: - """Helper class to manage PyJWKClient instances for different JWKS URIs, with caching and thread safety.""" + """Helper class to manage PyJWKClient instances for different JWKS URIs, with caching and async-safety""" - _cache: dict[str, PyJWKClient] + _cache: dict[str, _JwkClientCacheEntry] def __init__(self): self._cache = {} - def _get_jwk_client(self, jwks_uri: str) -> PyJWKClient: + def _get_jwk_client(self, jwks_uri: str) -> _JwkClientCacheEntry: """Retrieves a PyJWKClient for the given JWKS URI, using a cache to avoid creating multiple clients for the same URI.""" if jwks_uri not in self._cache: - self._cache[jwks_uri] = PyJWKClient(jwks_uri) + self._cache[jwks_uri] = _JwkClientCacheEntry(PyJWKClient(jwks_uri), asyncio.Lock()) return self._cache[jwks_uri] async def get_signing_key(self, jwks_uri: str, header: dict[str, Any]) -> PyJWK: """Retrieves the signing key from the JWK client for the given token header.""" - jwks_client = self._get_jwk_client(jwks_uri) - key = await asyncio.to_thread(jwks_client.get_signing_key, header["kid"]) - return key + + jwk_cache_entry = self._get_jwk_client(jwks_uri) + async with jwk_cache_entry.lock: + # locking and creating a new thread seems strange, + # but PyJWKClient.get_signing_key is synchronous, so we spawn another thread + # to make the call non-blocking, allowing other queued coroutines to run in the meantime. + # Meanwhile, the lock ensures safety for the PyJWKClient's underlying cache and + # prevents duplicate calls to the JWKS endpoint for the same URI when multiple + # coroutines are trying to get signing keys concurrently. + key = await asyncio.to_thread(jwk_cache_entry.jwk_client.get_signing_key, header["kid"]) + return key class JwtTokenValidator: """Utility class for validating JWT tokens using the PyJWT library and JWKs from a specified URI.""" diff --git a/tests/hosting_core/authorization/__init__.py b/tests/hosting_core/authorization/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/hosting_core/authorization/test_jwk_client_manager.py b/tests/hosting_core/authorization/test_jwk_client_manager.py new file mode 100644 index 00000000..0db24f9c --- /dev/null +++ b/tests/hosting_core/authorization/test_jwk_client_manager.py @@ -0,0 +1,130 @@ +import asyncio +import threading +import time + +import pytest +from jwt import PyJWKClient + +from microsoft_agents.hosting.core.authorization.jwt_token_validator import ( + _JwkClientManager, +) + + +async def _wait_until_set(event: threading.Event, timeout: float = 1.0) -> None: + start = time.monotonic() + while not event.is_set(): + if time.monotonic() - start > timeout: + raise AssertionError("Timed out waiting for threading event.") + await asyncio.sleep(0.01) + + +class TestJwkClientManager: + def test_get_jwk_client_reuses_cache_for_same_uri(self): + manager = _JwkClientManager() + jwks_uri = "https://issuer.example.com/keys" + + first = manager._get_jwk_client(jwks_uri) + second = manager._get_jwk_client(jwks_uri) + + assert first is second + assert len(manager._cache) == 1 + + def test_get_jwk_client_creates_distinct_entries_for_distinct_uris(self): + manager = _JwkClientManager() + + first = manager._get_jwk_client("https://issuer-a.example.com/keys") + second = manager._get_jwk_client("https://issuer-b.example.com/keys") + + assert first is not second + assert first.lock is not second.lock + assert len(manager._cache) == 2 + + @pytest.mark.asyncio + async def test_get_signing_key_calls_pyjwkclient_with_header_kid( + self, monkeypatch + ): + manager = _JwkClientManager() + jwks_uri = "https://issuer.example.com/keys" + seen_kids = [] + expected_key = object() + + def fake_get_signing_key(self, kid): + seen_kids.append(kid) + return expected_key + + # Only mocked member: PyJWKClient.get_signing_key + monkeypatch.setattr(PyJWKClient, "get_signing_key", fake_get_signing_key) + + key = await manager.get_signing_key(jwks_uri, {"kid": "kid-123"}) + + assert key is expected_key + assert seen_kids == ["kid-123"] + + @pytest.mark.asyncio + async def test_get_signing_key_reuses_same_client_for_same_uri(self, monkeypatch): + manager = _JwkClientManager() + jwks_uri = "https://issuer.example.com/keys" + client_ids = [] + + def fake_get_signing_key(self, kid): + client_ids.append(id(self)) + return {"kid": kid} + + # Only mocked member: PyJWKClient.get_signing_key + monkeypatch.setattr(PyJWKClient, "get_signing_key", fake_get_signing_key) + + await manager.get_signing_key(jwks_uri, {"kid": "kid-a"}) + await manager.get_signing_key(jwks_uri, {"kid": "kid-b"}) + + assert client_ids[0] == client_ids[1] + assert len(manager._cache) == 1 + + @pytest.mark.asyncio + async def test_get_signing_key_serializes_concurrent_calls_per_uri( + self, monkeypatch + ): + manager = _JwkClientManager() + jwks_uri = "https://issuer.example.com/keys" + + first_entered = threading.Event() + second_entered = threading.Event() + release_first = threading.Event() + + def fake_get_signing_key(self, kid): + if kid == "kid-1": + first_entered.set() + if not release_first.wait(timeout=2): + raise TimeoutError("First call was not released in time.") + elif kid == "kid-2": + second_entered.set() + return {"kid": kid} + + # Only mocked member: PyJWKClient.get_signing_key + monkeypatch.setattr(PyJWKClient, "get_signing_key", fake_get_signing_key) + + first_task = asyncio.create_task( + manager.get_signing_key(jwks_uri, {"kid": "kid-1"}) + ) + await _wait_until_set(first_entered) + + second_task = asyncio.create_task( + manager.get_signing_key(jwks_uri, {"kid": "kid-2"}) + ) + + # If per-URI lock works, second call must not enter get_signing_key yet. + await asyncio.sleep(0.05) + assert not second_entered.is_set() + + release_first.set() + results = await asyncio.gather(first_task, second_task) + + assert results[0]["kid"] == "kid-1" + assert results[1]["kid"] == "kid-2" + assert second_entered.is_set() + + @pytest.mark.asyncio + async def test_get_signing_key_raises_key_error_when_header_has_no_kid(self): + manager = _JwkClientManager() + + with pytest.raises(KeyError): + await manager.get_signing_key("https://issuer.example.com/keys", {}) \ No newline at end of file From 830825b53da19011a2b6e43846c5e71ce5c49b78 Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Mon, 9 Mar 2026 03:18:26 -0700 Subject: [PATCH 08/10] formatting --- .../core/authorization/jwt_token_validator.py | 13 ++++++++++--- tests/copilotstudio_client/test_copilot_client.py | 1 + .../authorization/test_jwk_client_manager.py | 6 ++---- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py index 21f6e8ea..7c4d3f0c 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py @@ -13,12 +13,14 @@ logger = logging.getLogger(__name__) + @dataclass class _JwkClientCacheEntry: jwk_client: PyJWKClient lock: asyncio.Lock + class _JwkClientManager: """Helper class to manage PyJWKClient instances for different JWKS URIs, with caching and async-safety""" @@ -30,9 +32,11 @@ def __init__(self): def _get_jwk_client(self, jwks_uri: str) -> _JwkClientCacheEntry: """Retrieves a PyJWKClient for the given JWKS URI, using a cache to avoid creating multiple clients for the same URI.""" if jwks_uri not in self._cache: - self._cache[jwks_uri] = _JwkClientCacheEntry(PyJWKClient(jwks_uri), asyncio.Lock()) + self._cache[jwks_uri] = _JwkClientCacheEntry( + PyJWKClient(jwks_uri), asyncio.Lock() + ) return self._cache[jwks_uri] - + async def get_signing_key(self, jwks_uri: str, header: dict[str, Any]) -> PyJWK: """Retrieves the signing key from the JWK client for the given token header.""" @@ -44,9 +48,12 @@ async def get_signing_key(self, jwks_uri: str, header: dict[str, Any]) -> PyJWK: # Meanwhile, the lock ensures safety for the PyJWKClient's underlying cache and # prevents duplicate calls to the JWKS endpoint for the same URI when multiple # coroutines are trying to get signing keys concurrently. - key = await asyncio.to_thread(jwk_cache_entry.jwk_client.get_signing_key, header["kid"]) + key = await asyncio.to_thread( + jwk_cache_entry.jwk_client.get_signing_key, header["kid"] + ) return key + class JwtTokenValidator: """Utility class for validating JWT tokens using the PyJWT library and JWKs from a specified URI.""" diff --git a/tests/copilotstudio_client/test_copilot_client.py b/tests/copilotstudio_client/test_copilot_client.py index e82e2869..cffc82f6 100644 --- a/tests/copilotstudio_client/test_copilot_client.py +++ b/tests/copilotstudio_client/test_copilot_client.py @@ -16,6 +16,7 @@ from aiohttp import ClientSession, ClientError from urllib.parse import urlparse + @pytest.mark.asyncio async def test_copilot_client_error(mocker): # Define the connection settings diff --git a/tests/hosting_core/authorization/test_jwk_client_manager.py b/tests/hosting_core/authorization/test_jwk_client_manager.py index 0db24f9c..89aea317 100644 --- a/tests/hosting_core/authorization/test_jwk_client_manager.py +++ b/tests/hosting_core/authorization/test_jwk_client_manager.py @@ -40,9 +40,7 @@ def test_get_jwk_client_creates_distinct_entries_for_distinct_uris(self): assert len(manager._cache) == 2 @pytest.mark.asyncio - async def test_get_signing_key_calls_pyjwkclient_with_header_kid( - self, monkeypatch - ): + async def test_get_signing_key_calls_pyjwkclient_with_header_kid(self, monkeypatch): manager = _JwkClientManager() jwks_uri = "https://issuer.example.com/keys" seen_kids = [] @@ -127,4 +125,4 @@ async def test_get_signing_key_raises_key_error_when_header_has_no_kid(self): manager = _JwkClientManager() with pytest.raises(KeyError): - await manager.get_signing_key("https://issuer.example.com/keys", {}) \ No newline at end of file + await manager.get_signing_key("https://issuer.example.com/keys", {}) From 6d4a5bbe24e3f56226d5ca393d66266bda463b5a Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Mon, 9 Mar 2026 11:15:07 -0700 Subject: [PATCH 09/10] Using a threading.Lock instead --- .../aiohttp/jwt_authorization_middleware.py | 1 + .../core/authorization/jwt_token_validator.py | 31 +++++++++++-------- 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/libraries/microsoft-agents-hosting-aiohttp/microsoft_agents/hosting/aiohttp/jwt_authorization_middleware.py b/libraries/microsoft-agents-hosting-aiohttp/microsoft_agents/hosting/aiohttp/jwt_authorization_middleware.py index f09604ba..02eae158 100644 --- a/libraries/microsoft-agents-hosting-aiohttp/microsoft_agents/hosting/aiohttp/jwt_authorization_middleware.py +++ b/libraries/microsoft-agents-hosting-aiohttp/microsoft_agents/hosting/aiohttp/jwt_authorization_middleware.py @@ -7,6 +7,7 @@ ) + @middleware async def jwt_authorization_middleware(request: Request, handler): diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py index 7c4d3f0c..53e235db 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py @@ -3,6 +3,7 @@ import asyncio import logging +import threading from typing import Any from dataclasses import dataclass @@ -18,7 +19,7 @@ class _JwkClientCacheEntry: jwk_client: PyJWKClient - lock: asyncio.Lock + lock: threading.Lock class _JwkClientManager: @@ -33,7 +34,7 @@ def _get_jwk_client(self, jwks_uri: str) -> _JwkClientCacheEntry: """Retrieves a PyJWKClient for the given JWKS URI, using a cache to avoid creating multiple clients for the same URI.""" if jwks_uri not in self._cache: self._cache[jwks_uri] = _JwkClientCacheEntry( - PyJWKClient(jwks_uri), asyncio.Lock() + PyJWKClient(jwks_uri), threading.Lock() ) return self._cache[jwks_uri] @@ -41,17 +42,21 @@ async def get_signing_key(self, jwks_uri: str, header: dict[str, Any]) -> PyJWK: """Retrieves the signing key from the JWK client for the given token header.""" jwk_cache_entry = self._get_jwk_client(jwks_uri) - async with jwk_cache_entry.lock: - # locking and creating a new thread seems strange, - # but PyJWKClient.get_signing_key is synchronous, so we spawn another thread - # to make the call non-blocking, allowing other queued coroutines to run in the meantime. - # Meanwhile, the lock ensures safety for the PyJWKClient's underlying cache and - # prevents duplicate calls to the JWKS endpoint for the same URI when multiple - # coroutines are trying to get signing keys concurrently. - key = await asyncio.to_thread( - jwk_cache_entry.jwk_client.get_signing_key, header["kid"] - ) - return key + + # locking and creating a new thread seems strange, + # but PyJWKClient.get_signing_key is synchronous, so we spawn another thread + # to make the call non-blocking, allowing other queued coroutines to run in the meantime. + # Meanwhile, the lock ensures safety for the PyJWKClient's underlying cache and + # prevents duplicate calls to the JWKS endpoint for the same URI when multiple + # coroutines are trying to get signing keys concurrently. + + def _helper(): + with jwk_cache_entry.lock: + return jwk_cache_entry.jwk_client.get_signing_key(header["kid"]) + + + key = await asyncio.to_thread(_helper) + return key class JwtTokenValidator: From 96258a2c99e8c9a61a00a2b1305b7fcfccb2a3c6 Mon Sep 17 00:00:00 2001 From: Rodrigo Brandao Date: Mon, 9 Mar 2026 11:16:03 -0700 Subject: [PATCH 10/10] Formatting --- .../hosting/aiohttp/jwt_authorization_middleware.py | 1 - .../hosting/core/authorization/jwt_token_validator.py | 1 - 2 files changed, 2 deletions(-) diff --git a/libraries/microsoft-agents-hosting-aiohttp/microsoft_agents/hosting/aiohttp/jwt_authorization_middleware.py b/libraries/microsoft-agents-hosting-aiohttp/microsoft_agents/hosting/aiohttp/jwt_authorization_middleware.py index 02eae158..f09604ba 100644 --- a/libraries/microsoft-agents-hosting-aiohttp/microsoft_agents/hosting/aiohttp/jwt_authorization_middleware.py +++ b/libraries/microsoft-agents-hosting-aiohttp/microsoft_agents/hosting/aiohttp/jwt_authorization_middleware.py @@ -7,7 +7,6 @@ ) - @middleware async def jwt_authorization_middleware(request: Request, handler): diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py index 53e235db..2c09c8c4 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/authorization/jwt_token_validator.py @@ -54,7 +54,6 @@ def _helper(): with jwk_cache_entry.lock: return jwk_cache_entry.jwk_client.get_signing_key(header["kid"]) - key = await asyncio.to_thread(_helper) return key