diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py index e94145b3efe..e7b608e162e 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py @@ -10,17 +10,18 @@ from collections.abc import Mapping, Sequence from typing import Final from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.safe_json_dumps import strip_null_bytes def optional_str(value: object) -> str | None: return value if isinstance(value, str) else None -def _optional_str_tuple(value: object) -> tuple[str, ...] | None: +def _sanitized_str_tuple(value: object) -> tuple[str, ...] | None: if not isinstance(value, list): return None items: Final[Sequence[object]] = value - return tuple(tag for tag in items if isinstance(tag, str)) + return tuple(strip_null_bytes(tag) for tag in items if isinstance(tag, str)) def is_collection_route(url_route: str, collection_suffix: str) -> bool: @@ -37,12 +38,12 @@ def request_tags_from_metadata(request_metadata: Mapping[str, object]) -> tuple[ tagged key does not put its tags in the top-level metadata "tags" on the passthrough path) """ - tags: Final = _optional_str_tuple(request_metadata.get("tags")) + tags: Final = _sanitized_str_tuple(request_metadata.get("tags")) if tags: return tags key_auth_metadata: Final = request_metadata.get("user_api_key_auth_metadata") if isinstance(key_auth_metadata, dict): - return _optional_str_tuple(key_auth_metadata.get("tags")) + return _sanitized_str_tuple(key_auth_metadata.get("tags")) return None diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 9b96387675b..7985faa9e4b 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -954,6 +954,16 @@ class TestAnthropicBatchPassthroughCostTracking: mock_managed_files_hook.store_unified_object_id.assert_awaited_once() return mock_managed_files_hook.store_unified_object_id.call_args[1] + @pytest.mark.asyncio + async def test_persisted_tags_are_db_safe(self, mock_logging_obj): + """Regression for PostgreSQL 22P05, asserted on the value that actually reaches + store_unified_object_id so it stays pinned if the sanitation moves.""" + call_kwargs = await self._store_with_metadata( + mock_logging_obj, {"user_api_key": "hashed-key-a", "tags": ["clean", "bad\x00tag"]} + ) + + assert call_kwargs["request_tags"] == ("clean", "badtag") + @pytest.mark.asyncio async def test_create_persists_key_hash_and_tags(self, mock_logging_obj): """Regression (LIT-5288): the batch create must persist the creating key's hashed diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_batch_attribution.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_batch_attribution.py index 5109cbb5991..1f7acd0723f 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_batch_attribution.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_batch_attribution.py @@ -70,6 +70,28 @@ class TestRequestTagsFromMetadata: def test_malformed_key_auth_metadata_is_ignored(self): assert request_tags_from_metadata({"user_api_key_auth_metadata": "nope"}) is None + @pytest.mark.parametrize( + "raw, expected", + [ + (["bad\x00tag"], ("badtag",)), + (["\x00leading"], ("leading",)), + (["trailing\x00"], ("trailing",)), + (["a\x00b\x00c"], ("abc",)), + (["\x00"], ("",)), + # every element, not just the first + (["clean", "bad\x00tag"], ("clean", "badtag")), + (["one\x00", "two\x00", "three\x00"], ("one", "two", "three")), + ], + ) + def test_nul_bytes_are_stripped_from_request_tags(self, raw, expected): + """Regression for PostgreSQL 22P05: an unstripped NUL aborts the managed object row + insert, so the batch is never cost tracked.""" + assert request_tags_from_metadata({"tags": raw}) == expected + + def test_nul_bytes_are_stripped_from_key_tags_fallback(self): + """Regression for PostgreSQL 22P05: the key-tags fallback shares the same helper.""" + assert request_tags_from_metadata({"user_api_key_auth_metadata": {"tags": ["key\x00tag"]}}) == ("keytag",) + @pytest.mark.parametrize( "url_route, suffix, expected", diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index d53e6dedf0b..ac79c183ca3 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -336,6 +336,17 @@ class TestVertexAIBatchPassthroughHandler: mock_managed_files_hook.store_unified_object_id.assert_called_once() return mock_managed_files_hook.store_unified_object_id.call_args[1] + def test_persisted_tags_are_db_safe(self, mock_logging_obj, mock_managed_files_hook): + """Regression for PostgreSQL 22P05, asserted for Vertex too so moving the + sanitation somewhere that only covers Anthropic fails loudly.""" + call_kwargs = self._store_with_metadata( + mock_logging_obj, + mock_managed_files_hook, + {"user_api_key": "hashed-key-a", "tags": ["clean", "bad\x00tag"]}, + ) + + assert call_kwargs["request_tags"] == ("clean", "badtag") + def test_create_persists_key_hash_and_tags( self, mock_logging_obj, mock_managed_files_hook ):