fix(guardrails): avoid blocking the event loop during HiddenLayer JWT refresh (#45437)

* fix(guardrails): avoid blocking the event loop during HiddenLayer JWT refresh

* test(guardrails): mock HiddenLayer detection requests with respx

litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import os
 from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Final, Literal, Protocol
@@ -402,7 +403,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
                 verbose_proxy_logger.debug(
                     "Unable to authenticate to Hiddenlayer, JWT token is invalid or expired, trying to refresh the token."
                 )
-                self.jwt_token = self.refresh_jwt_func()
+                self.jwt_token = await asyncio.to_thread(self.refresh_jwt_func)
                 headers["Authorization"] = f"Bearer {self.jwt_token}"
                 response = await self._http_client.post(
                     f"{self.api_base}/detection/v1/interactions",
@@ -609,7 +610,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
                 verbose_proxy_logger.debug(
                     "Unable to authenticate to Hiddenlayer, JWT token is invalid or expired, trying to refresh the token."
                 )
-                self.jwt_token = self.refresh_jwt_func()
+                self.jwt_token = await asyncio.to_thread(self.refresh_jwt_func)
                 headers["Authorization"] = f"Bearer {self.jwt_token}"
                 response = await self._http_client.post(
                     f"{self.api_base}/{path}",

tests/unit/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py

@@ -1,15 +1,17 @@
+import asyncio
 import os
 import threading
 import time
 import uuid
-from typing import List, cast
+from typing import Final, List, Literal, cast
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from fastapi import HTTPException
-from httpx import Request, Response
 import requests
-
+import responses
+import respx
+from fastapi import HTTPException
+from httpx import HTTPStatusError, Request, Response
 
 import litellm
 from litellm import ModelResponse
@@ -1153,3 +1155,130 @@ def test_get_jwt_gives_up_at_the_timeout_instead_of_blocking_the_event_loop(hang
         guardrail.refresh_jwt_func()
 
     assert [call.kwargs["timeout"] for call in get_jwt.call_args_list] == [2, 2]
+
+
+@pytest.fixture(params=("v1", "v2"))
+def hiddenlayer_jwt_guardrail(
+    request: pytest.FixtureRequest, httpx_transport: None
+) -> HiddenlayerGuardrail | HiddenlayerGuardrailV2:
+    guardrail_cls: Final = HiddenlayerGuardrail if request.param == "v1" else HiddenlayerGuardrailV2
+    with responses.RequestsMock() as auth:
+        auth.post(
+            "https://auth.example.test/oauth2/token?grant_type=client_credentials", json={"access_token": "expired"}
+        )
+        return guardrail_cls(
+            guardrail_name="hiddenlayer",
+            api_id="id",
+            api_key="secret",
+            api_base="https://api.hiddenlayer.ai",
+            auth_url="https://auth.example.test",
+            timeout=2,
+        )
+
+
+@pytest.mark.parametrize("input_type", ("request", "response"))
+async def test_jwt_refresh_allows_other_coroutines_to_run_before_it_finishes(
+    hiddenlayer_jwt_guardrail: HiddenlayerGuardrail | HiddenlayerGuardrailV2,
+    input_type: Literal["request", "response"],
+    respx_mock: respx.MockRouter,
+) -> None:
+    guardrail: Final = hiddenlayer_jwt_guardrail
+    loop: Final = asyncio.get_running_loop()
+    refresh_started: Final = asyncio.Event()
+    release_refresh: Final = threading.Event()
+    successful_body: Final = (
+        {"evaluation": {"action": "Allow"}}
+        if isinstance(guardrail, HiddenlayerGuardrail)
+        else {"messages": [{"role": "user", "content": "hi"}]}
+        if input_type == "request"
+        else {"choices": [{"message": {"role": "assistant", "content": "hi"}}]}
+    )
+
+    def refresh_token(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
+        assert request.headers["Authorization"] == "Basic aWQ6c2VjcmV0"
+        loop.call_soon_threadsafe(refresh_started.set)
+        assert release_refresh.wait(timeout=5), "event loop could not run while JWT refresh was pending"
+        return 200, {}, '{"access_token": "fresh"}'
+
+    def evaluate(request: Request) -> Response:
+        return Response(
+            401 if request.headers["Authorization"] == "Bearer expired" else 200,
+            json=successful_body,
+        )
+
+    async def other_request() -> None:
+        await refresh_started.wait()
+        release_refresh.set()
+
+    route: Final = respx_mock.route(method="POST", host="api.hiddenlayer.ai").mock(side_effect=evaluate)
+    peer: Final = asyncio.create_task(other_request())
+    with responses.RequestsMock() as auth:
+        auth.add_callback(
+            responses.POST,
+            "https://auth.example.test/oauth2/token?grant_type=client_credentials",
+            callback=refresh_token,
+            content_type="application/json",
+        )
+        try:
+            result, _ = await asyncio.wait_for(
+                asyncio.gather(
+                    guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type=input_type),
+                    peer,
+                ),
+                timeout=10,
+            )
+        finally:
+            release_refresh.set()
+            peer.cancel()
+            await asyncio.gather(peer, return_exceptions=True)
+        assert len(auth.calls) == 1
+
+    assert result["texts"] == ["hi"]
+    assert guardrail.jwt_token == "fresh"
+    assert route.call_count == 2
+    assert route.calls[0].request.headers["Authorization"] == "Bearer expired"
+    assert route.calls[1].request.headers["Authorization"] == "Bearer fresh"
+
+
+@pytest.mark.parametrize("input_type", ("request", "response"))
+async def test_jwt_refresh_exception_propagates_without_replacing_the_token(
+    hiddenlayer_jwt_guardrail: HiddenlayerGuardrail | HiddenlayerGuardrailV2,
+    input_type: Literal["request", "response"],
+    respx_mock: respx.MockRouter,
+) -> None:
+    guardrail: Final = hiddenlayer_jwt_guardrail
+    refresh_error: Final = requests.exceptions.Timeout("token endpoint timed out")
+    route: Final = respx_mock.route(method="POST", host="api.hiddenlayer.ai").respond(401)
+    with responses.RequestsMock() as auth:
+        auth.post("https://auth.example.test/oauth2/token?grant_type=client_credentials", body=refresh_error)
+        with pytest.raises(requests.exceptions.Timeout, match="token endpoint timed out") as caught:
+            await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type=input_type)
+
+    assert caught.value is refresh_error
+    assert guardrail.jwt_token == "expired"
+    assert route.call_count == 1
+    assert route.calls[0].request.headers["Authorization"] == "Bearer expired"
+
+
+@pytest.mark.parametrize("input_type", ("request", "response"))
+async def test_second_detection_401_propagates_without_another_jwt_refresh(
+    hiddenlayer_jwt_guardrail: HiddenlayerGuardrail | HiddenlayerGuardrailV2,
+    input_type: Literal["request", "response"],
+    respx_mock: respx.MockRouter,
+) -> None:
+    guardrail: Final = hiddenlayer_jwt_guardrail
+    route: Final = respx_mock.route(method="POST", host="api.hiddenlayer.ai").respond(401)
+    with responses.RequestsMock() as auth:
+        auth.post(
+            "https://auth.example.test/oauth2/token?grant_type=client_credentials", json={"access_token": "fresh"}
+        )
+        with pytest.raises(HTTPStatusError) as caught:
+            await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type=input_type)
+        assert len(auth.calls) == 1
+
+    assert caught.value.response.status_code == 401
+    assert caught.value.request.headers["Authorization"] == "Bearer fresh"
+    assert guardrail.jwt_token == "fresh"
+    assert route.call_count == 2
+    assert route.calls[0].request.headers["Authorization"] == "Bearer expired"
+    assert route.calls[1].request.headers["Authorization"] == "Bearer fresh"