From 05fdd099ba8e99fc4b0f5238a28a41c03c385d08 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 25 Jan 2026 05:06:49 +0530 Subject: [PATCH] =?UTF-8?q?fix(presidio):=20resolve=20runtime=20error=20by?= =?UTF-8?q?=20handling=20asyncio=20loops=20in=20bac=E2=80=A6=20(#19714)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(presidio): resolve runtime error by handling asyncio loops in background threads * add test case for thread safety --- .../guardrails/guardrail_hooks/presidio.py | 132 ++++++++-------- .../guardrail_hooks/test_presidio.py | 141 +++++++++++++++--- 2 files changed, 188 insertions(+), 85 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 20df54b62ce..fbe6d753679 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -9,8 +9,10 @@ import asyncio +import threading import json from datetime import datetime +from contextlib import asynccontextmanager from typing import ( TYPE_CHECKING, Any, @@ -107,6 +109,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # Lock to prevent race conditions when creating session under concurrent load # Note: asyncio.Lock() can be created without an event loop; it only needs one when awaited self._session_lock: asyncio.Lock = asyncio.Lock() + + # Track main thread ID to safely identity when we are running in main loop vs background thread + + self._main_thread_id = threading.get_ident() + if mock_testing is True: # for testing purposes only return @@ -172,46 +179,47 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): "http://" + self.presidio_anonymizer_api_base ) - async def _get_http_session(self) -> aiohttp.ClientSession: + @asynccontextmanager + async def _get_session_iterator( + self, + ) -> AsyncGenerator[aiohttp.ClientSession, None]: """ - Get or create the shared HTTP session for Presidio API calls. + Async context manager for yielding an HTTP session. - Fixes memory leak (issue #14540) where every guardrail check created - a new aiohttp.ClientSession that was never properly closed. - - Thread-safe: Uses asyncio.Lock to prevent race conditions when - multiple concurrent requests try to create the session simultaneously. + Logic: + 1. If running in the main thread (where the object was initialized/destined to live normally), + use the shared `self._http_session` (protected by a lock). + 2. If running in a background thread (e.g. logging hook), yield a NEW ephemeral session + and ensure it is closed after use. """ - async with self._session_lock: - if self._http_session is None or self._http_session.closed: - self._http_session = aiohttp.ClientSession() - return self._http_session + + # Check if we are in the stored main thread + if threading.get_ident() == self._main_thread_id: + # Main thread -> use shared session + async with self._session_lock: + if self._http_session is None or self._http_session.closed: + self._http_session = aiohttp.ClientSession() + yield self._http_session + else: + # Background thread -> create ephemeral session + # This avoids "attached to a different loop" or "no running event loop" errors + # when accessing the shared session created in the main loop + session = aiohttp.ClientSession() + try: + yield session + finally: + if not session.closed: + await session.close() async def _close_http_session(self) -> None: - """Close the HTTP session if it exists.""" + """Close the shared HTTP session if it exists.""" if self._http_session is not None and not self._http_session.closed: await self._http_session.close() self._http_session = None def __del__(self): - """Cleanup: close HTTP session on instance destruction.""" - if self._http_session is not None and not self._http_session.closed: - try: - # Try to close the session, but don't fail if event loop is gone - import asyncio - try: - loop = asyncio.get_event_loop() - if loop.is_running(): - # Schedule cleanup, don't block __del__ - asyncio.create_task(self._close_http_session()) - else: - loop.run_until_complete(self._close_http_session()) - except RuntimeError: - # Event loop is closed, can't clean up - not ideal but better than crashing - pass - except Exception: - # Suppress all exceptions in __del__ to avoid issues during shutdown - pass + """Cleanup: we try to close, but doing async cleanup in __del__ is risky.""" + pass def _get_presidio_analyze_request_payload( self, @@ -273,28 +281,27 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return self.mock_redacted_text # Use shared session to prevent memory leak (issue #14540) - session = await self._get_http_session() + async with self._get_session_iterator() as session: + # Make the request to /analyze + analyze_url = f"{self.presidio_analyzer_api_base}analyze" - # Make the request to /analyze - analyze_url = f"{self.presidio_analyzer_api_base}analyze" - - analyze_payload: PresidioAnalyzeRequest = ( - self._get_presidio_analyze_request_payload( - text=text, - presidio_config=presidio_config, - request_data=request_data, + analyze_payload: PresidioAnalyzeRequest = ( + self._get_presidio_analyze_request_payload( + text=text, + presidio_config=presidio_config, + request_data=request_data, + ) ) - ) - verbose_proxy_logger.debug( - "Making request to: %s with payload: %s", - analyze_url, - analyze_payload, - ) + verbose_proxy_logger.debug( + "Making request to: %s with payload: %s", + analyze_url, + analyze_payload, + ) - async with session.post(analyze_url, json=analyze_payload) as response: - analyze_results = await response.json() - verbose_proxy_logger.debug("analyze_results: %s", analyze_results) + async with session.post(analyze_url, json=analyze_payload) as response: + analyze_results = await response.json() + verbose_proxy_logger.debug("analyze_results: %s", analyze_results) # Handle error responses from Presidio (e.g., {'error': 'No text provided'}) # Presidio may return a dict instead of a list when errors occur @@ -302,7 +309,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if "error" in analyze_results: verbose_proxy_logger.warning( "Presidio analyzer returned error: %s, returning empty list", - analyze_results.get("error") + analyze_results.get("error"), ) return [] # If it's a dict but not an error, try to process it as a single item @@ -314,7 +321,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): except Exception as e: verbose_proxy_logger.warning( "Failed to parse Presidio dict response: %s, returning empty list", - e + e, ) return [] @@ -351,20 +358,19 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return text # Use shared session to prevent memory leak (issue #14540) - session = await self._get_http_session() + async with self._get_session_iterator() as session: + # Make the request to /anonymize + anonymize_url = f"{self.presidio_anonymizer_api_base}anonymize" + verbose_proxy_logger.debug("Making request to: %s", anonymize_url) + anonymize_payload = { + "text": text, + "analyzer_results": analyze_results, + } - # Make the request to /anonymize - anonymize_url = f"{self.presidio_anonymizer_api_base}anonymize" - verbose_proxy_logger.debug("Making request to: %s", anonymize_url) - anonymize_payload = { - "text": text, - "analyzer_results": analyze_results, - } - - async with session.post( - anonymize_url, json=anonymize_payload - ) as response: - redacted_text = await response.json() + async with session.post( + anonymize_url, json=anonymize_payload + ) as response: + redacted_text = await response.json() new_text = text if redacted_text is not None: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 42af3942f1a..fc4ff28c774 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -4,10 +4,9 @@ Tests PII detection and masking for different message formats """ import asyncio -import json import os import sys -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest @@ -625,7 +624,7 @@ async def test_request_data_flows_to_apply_guardrail(): return text with patch.object(presidio, "check_pii", mock_check_pii): - result = await presidio.apply_guardrail( + await presidio.apply_guardrail( inputs={"texts": ["Test message"]}, request_data=request_data, input_type="request", @@ -709,16 +708,21 @@ async def test_presidio_filter_scope_initializer(monkeypatch): monkeypatch.setattr(litellm, "logging_callback_manager", mgr, raising=False) import litellm.proxy.guardrails.guardrail_initializers as gi import litellm.proxy.guardrails.guardrail_hooks.presidio as presidio_mod + monkeypatch.setattr( presidio_mod, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False ) - monkeypatch.setattr(gi, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False) + monkeypatch.setattr( + gi, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False + ) # input-only created.clear() from litellm.proxy.guardrails.guardrail_initializers import initialize_presidio - params_input = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="input") + params_input = LitellmParams( + guardrail="presidio", mode="pre_call", presidio_filter_scope="input" + ) guardrail_dict = {"guardrail_name": "g1"} cb = initialize_presidio(params_input, guardrail_dict) assert cb is created[0] @@ -726,14 +730,18 @@ async def test_presidio_filter_scope_initializer(monkeypatch): # output-only created.clear() - params_output = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="output") + params_output = LitellmParams( + guardrail="presidio", mode="pre_call", presidio_filter_scope="output" + ) cb = initialize_presidio(params_output, guardrail_dict) assert len(created) == 1 assert created[0].apply_to_output is True # both -> expect two callbacks (input + output) created.clear() - params_both = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="both") + params_both = LitellmParams( + guardrail="presidio", mode="pre_call", presidio_filter_scope="both" + ) cb = initialize_presidio(params_both, guardrail_dict) assert len(created) == 2 assert any(not c.apply_to_output for c in created) @@ -741,13 +749,15 @@ async def test_presidio_filter_scope_initializer(monkeypatch): @pytest.mark.asyncio -async def test_empty_content_handling(presidio_guardrail, mock_user_api_key, mock_cache): +async def test_empty_content_handling( + presidio_guardrail, mock_user_api_key, mock_cache +): """ Test that Presidio handles empty content gracefully. - + This is common in tool/function calling where assistant messages have empty content but include tool_calls. - + Bug fix: Previously crashed with: TypeError: argument after ** must be a mapping, not str """ @@ -761,7 +771,10 @@ async def test_empty_content_handling(presidio_guardrail, mock_user_api_key, moc { "id": "call_123", "type": "function", - "function": {"name": "calculator", "arguments": '{"a":2,"b":2}'}, + "function": { + "name": "calculator", + "arguments": '{"a":2,"b":2}', + }, } ], }, @@ -794,10 +807,12 @@ async def test_empty_content_handling(presidio_guardrail, mock_user_api_key, moc @pytest.mark.asyncio -async def test_whitespace_only_content(presidio_guardrail, mock_user_api_key, mock_cache): +async def test_whitespace_only_content( + presidio_guardrail, mock_user_api_key, mock_cache +): """ Test that Presidio handles whitespace-only content gracefully. - + Whitespace-only content should be treated the same as empty content. """ test_data = { @@ -832,7 +847,7 @@ async def test_whitespace_only_content(presidio_guardrail, mock_user_api_key, mo async def test_analyze_text_with_empty_string(): """ Test analyze_text method directly with empty string. - + Should return empty list without making API call to Presidio. """ presidio = _OPTIONAL_PresidioPIIMasking( @@ -864,7 +879,7 @@ async def test_analyze_text_with_empty_string(): async def test_analyze_text_error_dict_handling(): """ Test that analyze_text handles error dict responses from Presidio API. - + When Presidio returns {'error': 'No text provided'}, should handle gracefully instead of crashing with TypeError. """ @@ -878,16 +893,20 @@ async def test_analyze_text_error_dict_handling(): class MockResponse: async def json(self): return {"error": "No text provided"} + async def __aenter__(self): return self + async def __aexit__(self, *args): pass - + class MockSession: def post(self, *args, **kwargs): return MockResponse() + async def __aenter__(self): return self + async def __aexit__(self, *args): pass @@ -904,10 +923,12 @@ async def test_analyze_text_error_dict_handling(): @pytest.mark.asyncio -async def test_tool_calling_complete_scenario(presidio_guardrail, mock_user_api_key, mock_cache): +async def test_tool_calling_complete_scenario( + presidio_guardrail, mock_user_api_key, mock_cache +): """ Test complete tool calling scenario with PII in user message. - + This tests the real-world scenario where: 1. User provides a query with PII 2. Assistant responds with empty content + tool_calls @@ -1002,7 +1023,12 @@ def test_no_thresholds_returns_all(): guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) analyze_results = [ {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.1, "start": 0, "end": 4}, - {"entity_type": PiiEntityType.EMAIL_ADDRESS, "score": 0.2, "start": 5, "end": 9}, + { + "entity_type": PiiEntityType.EMAIL_ADDRESS, + "score": 0.2, + "start": 5, + "end": 9, + }, ] filtered = guardrail.filter_analyze_results_by_score(analyze_results) @@ -1019,7 +1045,12 @@ def test_entity_specific_threshold_only_applies_to_that_entity(): ) analyze_results = [ {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4}, - {"entity_type": PiiEntityType.EMAIL_ADDRESS, "score": 0.1, "start": 5, "end": 9}, + { + "entity_type": PiiEntityType.EMAIL_ADDRESS, + "score": 0.1, + "start": 5, + "end": 9, + }, ] filtered = guardrail.filter_analyze_results_by_score(analyze_results) @@ -1038,7 +1069,12 @@ def test_filter_uses_default_all_threshold(): ) analyze_results = [ {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4}, - {"entity_type": PiiEntityType.EMAIL_ADDRESS, "score": 0.8, "start": 5, "end": 9}, + { + "entity_type": PiiEntityType.EMAIL_ADDRESS, + "score": 0.8, + "start": 5, + "end": 9, + }, ] filtered = guardrail.filter_analyze_results_by_score(analyze_results) @@ -1059,7 +1095,12 @@ def test_entity_specific_overrides_default_threshold(): ) analyze_results = [ {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.65, "start": 0, "end": 4}, - {"entity_type": PiiEntityType.EMAIL_ADDRESS, "score": 0.75, "start": 5, "end": 9}, + { + "entity_type": PiiEntityType.EMAIL_ADDRESS, + "score": 0.75, + "start": 5, + "end": 9, + }, ] filtered = guardrail.filter_analyze_results_by_score(analyze_results) @@ -1134,3 +1175,59 @@ def test_update_in_memory_applies_score_thresholds(): guardrail.update_in_memory_litellm_params(params) assert guardrail.presidio_score_thresholds == {PiiEntityType.CREDIT_CARD: 0.85} + + +@pytest.mark.asyncio +async def test_get_session_iterator_thread_safety(presidio_guardrail): + """ + Test that _get_session_iterator yields: + 1. The shared session when in the main thread. + 2. A new session when in a background thread. + """ + import threading + import aiohttp + + # 1. Main Thread Case + # We are in the "main thread" relative to the guardrail initialization + async with presidio_guardrail._get_session_iterator() as session: + assert isinstance(session, aiohttp.ClientSession) + assert session is presidio_guardrail._http_session + shared_session_id = id(session) + + # 2. Background Thread Case + # Define a helper function to run in a thread + def thread_target(loop, result_future): + async def run_in_loop(): + # This runs in the thread's loop + async with presidio_guardrail._get_session_iterator() as session: + return session, id(session) + + try: + # Create a new loop for this thread to run async code + new_loop = asyncio.new_event_loop() + asyncio.set_event_loop(new_loop) + session_obj, session_id = new_loop.run_until_complete(run_in_loop()) + result_future.set_result((session_obj, session_id)) + new_loop.close() + except Exception as e: + result_future.set_exception(e) + + # Run the background thread test + bg_future = asyncio.Future() + t = threading.Thread( + target=thread_target, args=(asyncio.get_running_loop(), bg_future) + ) + t.start() + t.join() + + bg_session, bg_session_id = await bg_future + + # Assertions + # The background session should be DIFFERENT from the shared session + assert bg_session_id != shared_session_id + # The shared session should still be open (not closed by the background thread) + assert not presidio_guardrail._http_session.closed + # The background session should be closed (handled by the context manager in the thread) + assert bg_session.closed + + print("✓ Session iterator thread safety test passed")