diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 406a4f8c98a..6a68ba8c4d1 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -312,8 +312,11 @@ class Cache: verbose_logger.debug("\nCreated cache key: %s", cache_key) hashed_cache_key = Cache._get_hashed_cache_key(cache_key) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) + # Remove preset_cache_key from kwargs to avoid "got multiple values" TypeError + # when kwargs already contains preset_cache_key from upstream callers + kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"} self._set_preset_cache_key_in_kwargs( - preset_cache_key=hashed_cache_key, **kwargs + preset_cache_key=hashed_cache_key, **kwargs_for_preset ) return hashed_cache_key diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index f764b07941b..7f7d47b3150 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -404,11 +404,14 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Prepare the signed headers signed_headers = dict(aws_request.headers.items()) + # Use prepared URL so path segments match SigV4 canonical request (e.g. %20 for spaces). + request_url = prepped.url or url + # Make the request with retry for transient S3 errors (500/503) max_retries = 3 for attempt in range(max_retries): response = await self.async_httpx_client.put( - url, data=json_string, headers=signed_headers + request_url, data=json_string, headers=signed_headers ) if response.status_code in (500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s @@ -590,6 +593,9 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Prepare the signed headers signed_headers = dict(aws_request.headers.items()) + # Use prepared URL so path segments match SigV4 canonical request (e.g. %20 for spaces). + request_url = prepped.url or url + httpx_client = _get_httpx_client( params={"ssl_verify": self.s3_verify} if self.s3_verify is not None @@ -599,7 +605,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): max_retries = 3 for attempt in range(max_retries): response = httpx_client.put( - url, data=json_string, headers=signed_headers + request_url, data=json_string, headers=signed_headers ) if response.status_code in (500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s @@ -701,8 +707,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Prepare the signed headers signed_headers = dict(aws_request.headers.items()) - # Make the request - response = await self.async_httpx_client.get(url, headers=signed_headers) + request_url = prepped.url or url + response = await self.async_httpx_client.get( + request_url, headers=signed_headers + ) if response.status_code != 200: verbose_logger.exception( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8d3831e75fb..c50bfe3ab1a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2596,7 +2596,12 @@ class MCPServerManager: return server # If not found and tool name is prefixed, try extracting server name from prefix - if is_tool_name_prefixed(tool_name): + known_prefixes = { + normalize_server_name(get_server_prefix(s)) + for s in self.get_registry().values() + if get_server_prefix(s) + } + if is_tool_name_prefixed(tool_name, known_server_prefixes=known_prefixes): ( original_tool_name, server_name_from_prefix, diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 8189f212bcb..79942eda54e 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -100,17 +100,39 @@ def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: return prefixed_name, "" -def is_tool_name_prefixed(tool_name: str) -> bool: +def is_tool_name_prefixed( + tool_name: str, + known_server_prefixes: Optional[set] = None, +) -> bool: """ - Check if tool name has server prefix + Check if tool name has a known MCP server prefix. + + When ``known_server_prefixes`` is provided the function verifies that the + substring before the first separator is an actual registered server + prefix. Without it the check falls back to the legacy heuristic + (separator present anywhere in the name), which can produce false + positives for non-MCP tools whose names contain hyphens + (e.g. ``text-to-speech``, ``code-review``). Args: - tool_name: Tool name to check + tool_name: Tool name to check. + known_server_prefixes: Optional set of normalised server prefixes + currently registered in the MCP manager. Pass this whenever + the caller has access to the server registry so that the check + is accurate. Returns: - True if tool name is prefixed, False otherwise + True if tool name is prefixed, False otherwise. """ - return MCP_TOOL_PREFIX_SEPARATOR in tool_name + if MCP_TOOL_PREFIX_SEPARATOR not in tool_name: + return False + + if known_server_prefixes is not None: + candidate_prefix = tool_name.split(MCP_TOOL_PREFIX_SEPARATOR, 1)[0] + return normalize_server_name(candidate_prefix) in known_server_prefixes + + # Legacy fallback – separator present somewhere in the name. + return True def validate_mcp_server_name( diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 0f4ebbd4880..67cb281029c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -433,6 +433,109 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # contain API keys or other secrets) in error responses. raise Exception(f"Presidio PII analysis failed: {type(e).__name__}") from e + async def _post_presidio_anonymize( + self, text: str, analyze_results: Any + ) -> Any: + """POST to Presidio anonymize; returns parsed JSON body.""" + # Use shared session to prevent memory leak (issue #14540) + async with self._get_session_iterator() as session: + 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, + headers={"Accept": "application/json"}, + ) as response: + if response.status >= 400: + error_body = await response.text() + raise Exception( + f"Presidio anonymizer returned HTTP {response.status}: {error_body[:200]}" + ) + content_type = getattr( + response, + "content_type", + response.headers.get("Content-Type", ""), + ) + if "application/json" not in content_type: + error_body = await response.text() + raise Exception( + f"Presidio anonymizer returned non-JSON Content-Type '{content_type}'; body: '{error_body[:200]}'" + ) + return await response.json() + + def _finalize_presidio_anonymize_simple( + self, + redacted_text: Dict[str, Any], + masked_entity_count: Dict[str, int], + ) -> str: + # No need to build numbered tokens — just use Presidio's + # already-anonymized text directly. The old code incorrectly + # applied anonymizer item positions (which reference the + # *output* text) to the *original* text, causing offset errors. + for item in redacted_text.get("items", []): + entity_type = item.get("entity_type", None) + if entity_type is not None: + masked_entity_count[entity_type] = ( + masked_entity_count.get(entity_type, 0) + 1 + ) + return redacted_text["text"] + + def _finalize_presidio_anonymize_numbered_tokens( + self, + text: str, + analyze_results: Any, + request_data: Optional[Dict], + masked_entity_count: Dict[str, int], + ) -> str: + # output_parse_pii is True — we need sequentially numbered + # tokens and a pii_tokens mapping for later unmasking. + # Use analyze_results positions (which reference the ORIGINAL + # text) instead of anonymizer items (which reference the output). + new_text = text + if request_data is None: + verbose_proxy_logger.warning( + "Presidio anonymize_text called without request_data — " + "PII tokens cannot be stored per-request. " + "This may indicate a missing caller update." + ) + request_data = {} + if not request_data.get("metadata"): + request_data["metadata"] = {} + if "pii_tokens" not in request_data["metadata"]: + request_data["metadata"]["pii_tokens"] = {} + pii_tokens = request_data["metadata"]["pii_tokens"] + + # Assign sequence numbers in forward (left-to-right) order so + # that is the first entity in the text, etc. + sorted_forward = sorted(analyze_results, key=lambda x: x["start"]) + seq_map = {} + for idx, ar in enumerate(sorted_forward, start=1): + seq_map[(ar["start"], ar["end"])] = idx + + # Apply replacements in reverse order by start position so + # that replacing later spans first does not shift earlier + # coordinates in the original text. + for ar in reversed(sorted_forward): + start = ar["start"] + end = ar["end"] + entity_type = ar["entity_type"] + replacement = f"<{entity_type}>" + seq = seq_map[(start, end)] + if replacement.endswith(">"): + replacement = f"{replacement[:-1]}_{seq}>" + else: + replacement = f"{replacement}_{seq}" + pii_tokens[replacement] = text[start:end] + new_text = new_text[:start] + replacement + new_text[end:] + masked_entity_count[entity_type] = ( + masked_entity_count.get(entity_type, 0) + 1 + ) + return new_text + async def anonymize_text( self, text: str, @@ -449,100 +552,20 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if isinstance(analyze_results, list) and len(analyze_results) == 0: return text - # Use shared session to prevent memory leak (issue #14540) - 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, - } - - async with session.post( - anonymize_url, - json=anonymize_payload, - headers={"Accept": "application/json"}, - ) as response: - # Validate HTTP status - if response.status >= 400: - error_body = await response.text() - raise Exception( - f"Presidio anonymizer returned HTTP {response.status}: {error_body[:200]}" - ) - - # Validate Content-Type is JSON - content_type = getattr( - response, - "content_type", - response.headers.get("Content-Type", ""), - ) - if "application/json" not in content_type: - error_body = await response.text() - raise Exception( - f"Presidio anonymizer returned non-JSON Content-Type '{content_type}'; body: '{error_body[:200]}'" - ) - - redacted_text = await response.json() - - new_text = text - if redacted_text is not None: - verbose_proxy_logger.debug("redacted_text: %s", redacted_text) - # Process items in reverse order by start position so that - # replacing later spans first does not shift earlier coordinates. - for item in sorted( - redacted_text["items"], key=lambda x: x["start"], reverse=True - ): - start = item["start"] - end = item["end"] - replacement = item["text"] # replacement token - if item["operator"] == "replace" and output_parse_pii is True: - if request_data is None: - verbose_proxy_logger.warning( - "Presidio anonymize_text called without request_data — " - "PII tokens cannot be stored per-request. " - "This may indicate a missing caller update." - ) - request_data = {} - # Store pii_tokens in metadata to avoid leaking to LLM providers. - # Providers like Anthropic reject unknown top-level fields. - if not request_data.get("metadata"): - request_data["metadata"] = {} - if "pii_tokens" not in request_data["metadata"]: - request_data["metadata"]["pii_tokens"] = {} - pii_tokens = request_data["metadata"]["pii_tokens"] - - # Append a sequential number to make each token unique - # per request, so unmasking maps back to the correct - # original value. Format: , - # This is LLM-friendly and degrades gracefully if the - # LLM doesn't echo the token verbatim. - seq = len(pii_tokens) + 1 - if replacement.endswith(">"): - replacement = f"{replacement[:-1]}_{seq}>" - else: - replacement = f"{replacement}_{seq}" - - # Use ORIGINAL text (not new_text) since start/end - # reference the original text's coordinates. - pii_tokens[replacement] = text[start:end] - - new_text = new_text[:start] + replacement + new_text[end:] - entity_type = item.get("entity_type", None) - if entity_type is not None: - masked_entity_count[entity_type] = ( - masked_entity_count.get(entity_type, 0) + 1 - ) - # When output_parse_pii is True, new_text contains sequentially - # numbered tokens (e.g. ) that match the keys - # in pii_tokens. Returning redacted_text["text"] (Presidio's - # original output) would send un-numbered tokens to the LLM, - # making unmasking impossible. - # When output_parse_pii is False, new_text == redacted_text["text"] - # because no suffix is appended. - return new_text - else: + redacted_text = await self._post_presidio_anonymize(text, analyze_results) + if redacted_text is None: raise Exception("Invalid anonymizer response: received None") + + verbose_proxy_logger.debug("redacted_text: %s", redacted_text) + + if not output_parse_pii: + return self._finalize_presidio_anonymize_simple( + redacted_text, masked_entity_count + ) + + return self._finalize_presidio_anonymize_numbered_tokens( + text, analyze_results, request_data, masked_entity_count + ) except Exception as e: # Sanitize exception to avoid leaking the original text (which may # contain API keys or other secrets) in error responses. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 85a12f70f58..db9468356dd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11527,8 +11527,11 @@ async def login_v2(request: Request): # noqa: PLR0915 litellm_dashboard_ui += "/ui/" litellm_dashboard_ui += "?login=success" + # Token is included in the response body so the UI can set a JS-accessible + # cookie even when a reverse proxy (e.g. nginx-ingress) adds HttpOnly to the + # server-set cookie, which would otherwise cause an infinite login redirect. json_response = JSONResponse( - content={"redirect_url": litellm_dashboard_ui}, + content={"redirect_url": litellm_dashboard_ui, "token": jwt_token}, status_code=status.HTTP_200_OK, ) json_response.set_cookie(key="token", value=jwt_token) diff --git a/tests/local_testing/test_cache_preset_key.py b/tests/local_testing/test_cache_preset_key.py new file mode 100644 index 00000000000..de0ec05603c --- /dev/null +++ b/tests/local_testing/test_cache_preset_key.py @@ -0,0 +1,87 @@ +""" +Test for preset_cache_key multiple values bug fix. + +This test verifies that get_cache_key doesn't raise TypeError when kwargs +already contains preset_cache_key. + +Issue: When get_cache_key(**kwargs) is called with kwargs containing +preset_cache_key, the call to _set_preset_cache_key_in_kwargs() would fail with: + TypeError: got multiple values for keyword argument 'preset_cache_key' +""" + +import pytest +from unittest.mock import MagicMock, patch + + +class TestPresetCacheKeyFix: + """Tests for the preset_cache_key multiple values fix.""" + + def test_get_cache_key_with_preset_cache_key_in_kwargs(self): + """ + Test that get_cache_key handles kwargs that already contain preset_cache_key. + + This was causing: + TypeError: _set_preset_cache_key_in_kwargs() got multiple values + for keyword argument 'preset_cache_key' + """ + from litellm.caching.caching import Cache + + cache = Cache() + + # Simulate kwargs that already has preset_cache_key (as can happen + # when the cache key is recomputed in certain code paths) + kwargs_with_preset = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "preset_cache_key": "existing_key_12345", # This caused the bug + "litellm_params": {}, + } + + # This should NOT raise TypeError + try: + result = cache.get_cache_key(**kwargs_with_preset) + assert result is not None + assert isinstance(result, str) + except TypeError as e: + if "multiple values for keyword argument" in str(e): + pytest.fail(f"Bug not fixed: {e}") + raise + + def test_get_cache_key_without_preset_cache_key(self): + """Test normal case without preset_cache_key in kwargs still works.""" + from litellm.caching.caching import Cache + + cache = Cache() + + kwargs_normal = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "litellm_params": {}, + } + + result = cache.get_cache_key(**kwargs_normal) + assert result is not None + assert isinstance(result, str) + + def test_preset_cache_key_is_set_in_litellm_params(self): + """Verify that preset_cache_key is correctly set in litellm_params.""" + from litellm.caching.caching import Cache + + cache = Cache() + + litellm_params = {} + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "litellm_params": litellm_params, + } + + result = cache.get_cache_key(**kwargs) + + # The method should set preset_cache_key in litellm_params + assert "preset_cache_key" in litellm_params + assert litellm_params["preset_cache_key"] == result + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index ab4ac1aa68a..2ad8358cc94 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -292,6 +292,50 @@ class TestS3V2UnitTests: assert result == {"downloaded": "data"} + @patch("asyncio.create_task") + @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") + def test_s3_v2_put_url_encodes_spaces_in_object_key( + self, mock_periodic_flush, mock_create_task + ): + import requests + from unittest.mock import AsyncMock + + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement + + mock_periodic_flush.return_value = None + mock_create_task.return_value = None + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + + s3_object_key = "My Team/2025-09-14/test-key.json" + test_element = s3BatchLoggingElement( + s3_object_key=s3_object_key, + payload={"test": "data"}, + s3_object_download_filename="test-file.json", + ) + + s3_logger = S3Logger( + s3_bucket_name="test-bucket", + s3_endpoint_url="https://s3.amazonaws.com", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + ) + s3_logger.async_httpx_client = AsyncMock() + s3_logger.async_httpx_client.put.return_value = mock_response + + asyncio.run(s3_logger.async_upload_data_to_s3(test_element)) + + call_args = s3_logger.async_httpx_client.put.call_args + assert call_args is not None + actual_url = call_args[0][0] + raw_url = f"https://s3.amazonaws.com/test-bucket/{s3_object_key}" + expected_url = requests.Request("PUT", raw_url).prepare().url + assert actual_url == expected_url + assert " " not in actual_url + @pytest.mark.asyncio async def test_async_upload_retries_on_s3_503(): """ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py new file mode 100644 index 00000000000..d761d9c54cc --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py @@ -0,0 +1,90 @@ +""" +Tests for is_tool_name_prefixed with known_server_prefixes parameter. + +Verifies fix for https://github.com/BerriAI/litellm/issues/25081 +""" + +import pytest + +from litellm.proxy._experimental.mcp_server.utils import is_tool_name_prefixed + + +# --------------------------------------------------------------------------- +# Legacy behaviour (no known_server_prefixes passed) +# --------------------------------------------------------------------------- + + +class TestLegacyBehaviour: + """Without known_server_prefixes the function falls back to heuristic.""" + + def test_plain_name_returns_false(self): + assert is_tool_name_prefixed("get_weather") is False + + def test_hyphenated_name_returns_true_legacy(self): + """Legacy heuristic: any hyphen → True (the bug this issue reports).""" + assert is_tool_name_prefixed("text-to-speech") is True + + def test_prefixed_name_returns_true_legacy(self): + assert is_tool_name_prefixed("myserver-get_weather") is True + + +# --------------------------------------------------------------------------- +# New behaviour (known_server_prefixes supplied) +# --------------------------------------------------------------------------- + + +class TestWithKnownPrefixes: + """When known_server_prefixes is supplied, only real prefixes match.""" + + PREFIXES = {"myserver", "weather_api", "code_tools"} + + def test_known_prefix_returns_true(self): + assert ( + is_tool_name_prefixed( + "myserver-get_weather", known_server_prefixes=self.PREFIXES + ) + is True + ) + + def test_hyphenated_non_mcp_tool_returns_false(self): + """This is the core fix: 'text-to-speech' is NOT an MCP-prefixed tool.""" + assert ( + is_tool_name_prefixed( + "text-to-speech", known_server_prefixes=self.PREFIXES + ) + is False + ) + + def test_code_review_not_misclassified(self): + assert ( + is_tool_name_prefixed( + "code-review", known_server_prefixes=self.PREFIXES + ) + is False + ) + + def test_no_separator_returns_false(self): + assert ( + is_tool_name_prefixed( + "simple_tool", known_server_prefixes=self.PREFIXES + ) + is False + ) + + def test_empty_prefixes_set_rejects_all(self): + """With an empty registry, nothing can be prefixed.""" + assert ( + is_tool_name_prefixed("myserver-get_weather", known_server_prefixes=set()) + is False + ) + + def test_prefix_normalisation(self): + """Server names with spaces are normalised to underscores.""" + prefixes = {"my_server"} + # add_server_prefix_to_name normalises spaces → underscores + assert ( + is_tool_name_prefixed( + "my_server-list_files", known_server_prefixes=prefixes + ) + is True + ) 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 32a8c1b1070..38ea42285c1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2230,3 +2230,171 @@ async def test_apply_to_output_streaming_bytes_only_logs_warning(): mock_logger.warning.assert_called_once() warning_msg = mock_logger.warning.call_args[0][0] assert "Output PII masking was skipped" in warning_msg + + +@pytest.mark.asyncio +async def test_anonymize_text_uses_correct_positions_no_parse_pii(): + """ + Regression test for anonymizer offset bug (fixes #24160). + + The Presidio anonymizer returns items with start/end positions that + reference the *anonymized output* text, not the original input text. + When output_parse_pii is False, anonymize_text must return + redacted_text["text"] directly instead of manually splicing the + original text using those positions, which produces garbled output + with remnants of original PII data. + """ + original_text = ( + "My name is John Smith, my email is john@example.com, phone 555-867-5309" + ) + # Positions as returned by the analyzer (reference original text) + analyze_results = [ + {"end": 51, "entity_type": "EMAIL_ADDRESS", "score": 1.0, "start": 35}, + {"end": 21, "entity_type": "PERSON", "score": 0.85, "start": 11}, + {"end": 71, "entity_type": "PHONE_NUMBER", "score": 0.75, "start": 59}, + ] + # Anonymizer response — positions reference the *anonymized* text + anonymizer_response = { + "text": "My name is , my email is , phone ", + "items": [ + { + "start": 56, + "end": 70, + "entity_type": "PHONE_NUMBER", + "text": "", + "operator": "replace", + }, + { + "start": 33, + "end": 48, + "entity_type": "EMAIL_ADDRESS", + "text": "", + "operator": "replace", + }, + { + "start": 11, + "end": 19, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + }, + ], + } + + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=anonymizer_response, + ) + + masked_entity_count = {} + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + result = await guardrail.anonymize_text( + text=original_text, + analyze_results=analyze_results, + output_parse_pii=False, + masked_entity_count=masked_entity_count, + ) + + expected = "My name is , my email is , phone " + assert result == expected, ( + f"anonymize_text produced garbled output with PII remnants.\n" + f"Expected: {expected!r}\n" + f"Got: {result!r}" + ) + assert masked_entity_count == { + "PERSON": 1, + "EMAIL_ADDRESS": 1, + "PHONE_NUMBER": 1, + } + + +@pytest.mark.asyncio +async def test_anonymize_text_uses_correct_positions_with_parse_pii(): + """ + Regression test for anonymizer offset bug with output_parse_pii=True + (fixes #24160). + + When output_parse_pii is True, anonymize_text must use positions from + analyze_results (which reference the original text) to build numbered + tokens and the pii_tokens mapping, not positions from anonymizer items + (which reference the anonymized output text). + """ + original_text = ( + "My name is John Smith, my email is john@example.com, phone 555-867-5309" + ) + analyze_results = [ + {"end": 51, "entity_type": "EMAIL_ADDRESS", "score": 1.0, "start": 35}, + {"end": 21, "entity_type": "PERSON", "score": 0.85, "start": 11}, + {"end": 71, "entity_type": "PHONE_NUMBER", "score": 0.75, "start": 59}, + ] + anonymizer_response = { + "text": "My name is , my email is , phone ", + "items": [ + { + "start": 56, + "end": 70, + "entity_type": "PHONE_NUMBER", + "text": "", + "operator": "replace", + }, + { + "start": 33, + "end": 48, + "entity_type": "EMAIL_ADDRESS", + "text": "", + "operator": "replace", + }, + { + "start": 11, + "end": 19, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + }, + ], + } + + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + output_parse_pii=True, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=anonymizer_response, + ) + + masked_entity_count = {} + request_data = {"metadata": {}} + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + result = await guardrail.anonymize_text( + text=original_text, + analyze_results=analyze_results, + output_parse_pii=True, + masked_entity_count=masked_entity_count, + request_data=request_data, + ) + + # Result must not contain any remnants of original PII + assert "John" not in result + assert "john@example.com" not in result + assert "555-867-5309" not in result + + # pii_tokens must map numbered tokens back to correct original values + pii_tokens = request_data["metadata"]["pii_tokens"] + token_values = set(pii_tokens.values()) + assert "John Smith" in token_values + assert "john@example.com" in token_values + assert "555-867-5309" in token_values + + # Tokens must be numbered in left-to-right order of appearance: + # PERSON (pos 11) → _1, EMAIL_ADDRESS (pos 35) → _2, PHONE_NUMBER (pos 59) → _3 + assert pii_tokens.get("") == "John Smith" + assert pii_tokens.get("") == "john@example.com" + assert pii_tokens.get("") == "555-867-5309" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index daabed0def1..c32a1bdd463 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -104,7 +104,10 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): ) assert response.status_code == 200 - assert response.json() == {"redirect_url": "http://testserver/ui/?login=success"} + assert response.json() == { + "redirect_url": "http://testserver/ui/?login=success", + "token": "signed-token", + } assert response.cookies.get("token") == "signed-token" mock_authenticate_user.assert_awaited_once_with( diff --git a/tests/test_litellm/types/test_prometheus_latency_buckets.py b/tests/test_litellm/types/test_prometheus_latency_buckets.py new file mode 100644 index 00000000000..85670bb0b74 --- /dev/null +++ b/tests/test_litellm/types/test_prometheus_latency_buckets.py @@ -0,0 +1,17 @@ +"""LATENCY_BUCKETS covers long-running LLM calls (histograms are in seconds).""" + +import math + +from litellm.types.integrations.prometheus import LATENCY_BUCKETS + + +def test_latency_buckets_include_seven_and_ten_minutes(): + """Buckets beyond 5 min so histograms resolve requests up to default LLM timeouts.""" + assert 300.0 in LATENCY_BUCKETS + assert 420.0 in LATENCY_BUCKETS # 7 min + assert 600.0 in LATENCY_BUCKETS # 10 min + assert math.isinf(LATENCY_BUCKETS[-1]) + idx_300 = LATENCY_BUCKETS.index(300.0) + idx_420 = LATENCY_BUCKETS.index(420.0) + idx_600 = LATENCY_BUCKETS.index(600.0) + assert idx_300 < idx_420 < idx_600 diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index d3fab5cf5bb..5b3ec1d1fb1 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -42,6 +42,7 @@ import ToolPoliciesView from "@/components/ToolPoliciesView"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; +import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { isJwtExpired } from "@/utils/jwtUtils"; import { buildLoginUrlWithReturn, consumeReturnUrl, isValidReturnUrl, normalizeUrlForCompare, storeReturnUrl } from "@/utils/returnUrlUtils"; import { formatUserRole, isAdminRole } from "@/utils/roles"; @@ -51,21 +52,12 @@ import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useMemo, useRef, useState } from "react"; import { ConfigProvider, theme } from "antd"; -function getCookie(name: string) { - // Safer cookie read + decoding; handles '=' inside values - const match = document.cookie.split("; ").find((row) => row.startsWith(name + "=")); - if (!match) return null; - const value = match.slice(name.length + 1); - try { - return decodeURIComponent(value); - } catch { - return value; - } -} - function deleteCookie(name: string, path = "/") { // Best-effort client-side clear (works for non-HttpOnly cookies without Domain) document.cookie = `${name}=; Max-Age=0; Path=${path}`; + if (name === "token") { + clearTokenCookies(); + } } interface ProxySettings { diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index c57dcb97eb9..3c107fa586f 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -5,6 +5,7 @@ import * as Networking from "./networking"; vi.mock("@/utils/cookieUtils", () => ({ clearTokenCookies: vi.fn(), getCookie: vi.fn(), + storeLoginToken: vi.fn(), })); vi.mock("./molecules/notifications_manager", () => ({ @@ -79,6 +80,38 @@ describe("networking - expired session handling", () => { }); }); +describe("loginCall - storeLoginToken integration", () => { + const originalFetch = global.fetch; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + global.fetch = originalFetch; + }); + + it("calls storeLoginToken when response includes token", async () => { + global.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ redirect_url: "/ui/?login=success", token: "my-jwt" }), + }) as any; + const { storeLoginToken } = await import("@/utils/cookieUtils"); + await Networking.loginCall("admin", "pass"); + expect(storeLoginToken).toHaveBeenCalledWith("my-jwt"); + }); + + it("does not call storeLoginToken when response has no token", async () => { + global.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ redirect_url: "/ui/?login=success" }), + }) as any; + const { storeLoginToken } = await import("@/utils/cookieUtils"); + await Networking.loginCall("admin", "pass"); + expect(storeLoginToken).not.toHaveBeenCalled(); + }); +}); + describe("daily activity helpers", () => { const startTime = new Date("2025-02-12T00:00:00.000Z"); const endTime = new Date("2025-02-19T00:00:00.000Z"); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 28f8d308de7..16f35605877 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -69,7 +69,7 @@ export const getInProductNudgesCall = async (accessToken: string) => { * Helper file for calls being made to proxy */ import MessageManager from "@/components/molecules/message_manager"; -import { clearTokenCookies } from "@/utils/cookieUtils"; +import { clearTokenCookies, storeLoginToken } from "@/utils/cookieUtils"; import { TagNewRequest, TagUpdateRequest, TagListResponse, TagInfoResponse } from "./tag_management/types"; import { Team } from "./key_team_helpers/key_list"; import { UserInfo } from "./view_users/types"; @@ -9255,14 +9255,14 @@ export const loginCall = async (username: string, password: string, useV3?: bool const exchangeData: LoginResponse = await exchangeResponse.json(); if (exchangeData.token) { - document.cookie = `token=${exchangeData.token}; path=/; SameSite=Lax`; + storeLoginToken(exchangeData.token); } return exchangeData; } // Backwards compatibility: v2 or old v3 returns token directly if (data.token) { - document.cookie = `token=${data.token}; path=/; SameSite=Lax`; + storeLoginToken(data.token); } return data; diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index f97d8ffab04..90eac56540d 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -1,5 +1,5 @@ "use client"; -import { clearTokenCookies } from "@/utils/cookieUtils"; +import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { Col, Grid } from "@tremor/react"; import { Typography } from "antd"; import { jwtDecode } from "jwt-decode"; @@ -35,12 +35,6 @@ export type UserInfo = { spend: number; }; -function getCookie(name: string) { - console.log("COOKIES", document.cookie); - const cookieValue = document.cookie.split("; ").find((row) => row.startsWith(name + "=")); - return cookieValue ? cookieValue.split("=")[1] : null; -} - interface UserDashboardProps { userID: string | null; userRole: string | null; @@ -103,7 +97,11 @@ const UserDashboard: React.FC = ({ // They are only cleared on logout useEffect(() => { const handleBeforeUnload = () => { + const token = sessionStorage.getItem("token"); sessionStorage.clear(); + if (token) { + sessionStorage.setItem("token", token); + } }; window.addEventListener("beforeunload", handleBeforeUnload); return () => window.removeEventListener("beforeunload", handleBeforeUnload); diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts index 8b066e6a8ea..c7bd27a6a85 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts @@ -1,11 +1,12 @@ import { describe, it, expect, beforeEach, vi } from "vitest"; -import { clearTokenCookies, getCookie } from "./cookieUtils"; +import { clearTokenCookies, getCookie, storeLoginToken } from "./cookieUtils"; describe("cookieUtils", () => { beforeEach(() => { document.cookie.split(";").forEach((c) => { document.cookie = c.replace(/^ +/, "").replace(/=.*/, "=;expires=" + new Date().toUTCString() + ";path=/"); }); + sessionStorage.clear(); vi.spyOn(console, "log").mockImplementation(() => {}); }); @@ -116,6 +117,55 @@ describe("cookieUtils", () => { vi.restoreAllMocks(); }); + + it("should clear sessionStorage token", () => { + sessionStorage.setItem("token", "stored-token"); + clearTokenCookies(); + expect(sessionStorage.getItem("token")).toBeNull(); + }); + }); + + describe("storeLoginToken", () => { + it("should store the token in sessionStorage", () => { + storeLoginToken("my-jwt-token"); + expect(sessionStorage.getItem("token")).toBe("my-jwt-token"); + }); + + it("should overwrite an existing token in sessionStorage", () => { + storeLoginToken("old-token"); + expect(sessionStorage.getItem("token")).toBe("old-token"); + + storeLoginToken("new-token"); + expect(sessionStorage.getItem("token")).toBe("new-token"); + }); + + it("should not throw when window is undefined (server-side rendering)", () => { + const originalWindow = global.window; + delete (global as any).window; + + expect(() => storeLoginToken("token")).not.toThrow(); + + global.window = originalWindow; + }); + + it("should not store empty string token", () => { + storeLoginToken(""); + expect(sessionStorage.getItem("token")).toBeNull(); + }); + + it("should not store whitespace-only token", () => { + storeLoginToken(" "); + expect(sessionStorage.getItem("token")).toBeNull(); + }); + + it("should set a JS-accessible cookie at /ui path", () => { + const cookieSpy = vi.spyOn(document, "cookie", "set"); + storeLoginToken("my-jwt-token"); + expect(cookieSpy).toHaveBeenCalledWith( + expect.stringContaining("path=/ui") + ); + vi.restoreAllMocks(); + }); }); describe("getCookie", () => { @@ -141,5 +191,26 @@ describe("cookieUtils", () => { expect(getCookie("token")).toBe("token-value"); expect(getCookie("other")).toBe("other-value"); }); + + it("should handle values containing '=' characters", () => { + document.cookie = "token=abc=def=ghi; path=/"; + expect(getCookie("token")).toBe("abc=def=ghi"); + }); + + it("should fall back to sessionStorage when cookie is not found", () => { + sessionStorage.setItem("token", "session-stored-jwt"); + expect(getCookie("token")).toBe("session-stored-jwt"); + }); + + it("should prefer cookie over sessionStorage", () => { + document.cookie = "token=cookie-value; path=/"; + sessionStorage.setItem("token", "session-value"); + expect(getCookie("token")).toBe("cookie-value"); + }); + + it("should not fall back to sessionStorage for non-token keys", () => { + sessionStorage.setItem("other", "other-value"); + expect(getCookie("other")).toBeNull(); + }); }); }); diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.ts b/ui/litellm-dashboard/src/utils/cookieUtils.ts index 01add36542c..b4493744ad4 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.ts @@ -2,6 +2,23 @@ * Utility functions for managing cookies */ +/** + * Returns the cookie path for the UI. + * Derives the path from window.location.pathname so it works when + * LiteLLM is deployed behind a subpath (e.g. /myapp/ui instead of /ui). + * No imports from networking.tsx to avoid circular dependencies. + */ +function getUiCookiePath(): string { + if (typeof window === "undefined") return "/ui"; + // Match "/ui" only as a full path segment (followed by "/" or end of string) + // to avoid false matches like "/my-ui-tool/login" → "/my-ui". + const match = window.location.pathname.match(/\/ui(?=\/|$)/); + if (match && match.index !== undefined) { + return window.location.pathname.substring(0, match.index + 3); + } + return "/ui"; +} + /** * Clears the token cookie from both root and /ui paths */ @@ -16,7 +33,8 @@ export function clearTokenCookies() { // Clear with various combinations of path and SameSite // Include current path in case of custom server root path const currentPath = window.location.pathname; - const paths = ["/", "/ui"]; + const uiCookiePath = getUiCookiePath(); + const paths = ["/", uiCookiePath]; // Add the current path directory if it's different from root and /ui if (currentPath && currentPath !== "/" && !currentPath.startsWith("/ui")) { @@ -43,7 +61,45 @@ export function clearTokenCookies() { }); }); - console.log("After clearing cookies:", document.cookie); + try { + sessionStorage.removeItem("token"); + } catch { + // sessionStorage may be unavailable + } + +} + +/** + * Stores the login token so the UI can read it even when a reverse proxy + * (e.g. nginx-ingress) adds HttpOnly to the server-set cookie. + * + * Strategy: + * 1. Set a JS-accessible cookie at path "/ui". Because nginx only modifies + * server-set Set-Cookie headers, a cookie created via document.cookie will + * never carry HttpOnly. Using path "/ui" avoids colliding with the + * server-set HttpOnly cookie at path "/". + * 2. Also store in sessionStorage as a secondary fallback. + */ +export function storeLoginToken(token: string) { + if (typeof window === "undefined") return; + if (!token || !token.trim()) return; + + // 1. JS-accessible cookie at /ui — survives same-tab navigations and + // is readable by getCookie() via document.cookie. + try { + const secure = window.location.protocol === "https:" ? "; Secure" : ""; + const cookiePath = getUiCookiePath(); + document.cookie = `token=${encodeURIComponent(token)}; path=${cookiePath}; SameSite=Lax${secure}`; + } catch { + // cookie setting may fail in restrictive environments + } + + // 2. sessionStorage backup + try { + sessionStorage.setItem("token", token); + } catch { + // sessionStorage may be unavailable (e.g. private browsing quota exceeded) + } } /** @@ -53,6 +109,23 @@ export function clearTokenCookies() { */ export function getCookie(name: string) { if (typeof document === "undefined") return null; - const cookieValue = document.cookie.split("; ").find((row) => row.startsWith(name + "=")); - return cookieValue ? cookieValue.split("=")[1] : null; + const row = document.cookie.split("; ").find((r) => r.startsWith(name + "=")); + if (row) { + const raw = row.split("=").slice(1).join("="); + try { + return decodeURIComponent(raw); + } catch { + return raw; + } + } + // Fallback to sessionStorage — covers the case where a reverse proxy + // added HttpOnly to the server-set cookie, making it invisible to JS. + if (name === "token" && typeof window !== "undefined") { + try { + return sessionStorage.getItem(name); + } catch { + return null; + } + } + return null; }