import asyncio, boto3, importlib, logging, os, pytest_asyncio import copy import json import re import sys import textwrap import threading import time import uuid from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine from contextlib import asynccontextmanager from contextvars import Context from datetime import datetime from pathlib import Path from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.s3_v2 import S3BatchUploadError, S3Logger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.integrations.s3_v2 import s3BatchLoggingElement from litellm.types.utils import StandardLoggingPayload from collections import defaultdict from litellm._logging import verbose_logger from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE, S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome _real_sleep: Final = asyncio.sleep _NOW: Final = 1_000_000.0 def _assert_jittered_backoff(sleep_mock: MagicMock, retries: int) -> None: """Each retry waits a jittered exponential backoff: attempt n sleeps somewhere in [2**n / 2, 2**n * 1.5).""" waits: Final = [sleep_call.args[0] for sleep_call in sleep_mock.call_args_list] assert len(waits) == retries, waits for attempt, wait in enumerate(waits): assert (1 << attempt) * 0.5 <= wait < (1 << attempt) * 1.5, (attempt, wait) class TestS3V2UnitTests: """Test that S3 v2 integration only uses safe_dumps and not json.dumps""" def test_s3_v2_source_code_analysis(self): """Test that S3 v2 source code only imports and uses safe_dumps""" import inspect from litellm.integrations import s3_v2 # Get the source code of the s3_v2 module source_code = inspect.getsource(s3_v2) # Verify that json.dumps is not used directly in the code assert "json.dumps(" not in source_code, "S3 v2 should not use json.dumps directly" @patch("asyncio.create_task") @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") def test_s3_v2_endpoint_url(self, mock_periodic_flush, mock_create_task): """testing s3 endpoint url""" from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement # Mock periodic_flush and create_task to prevent async task creation during init mock_periodic_flush.return_value = None mock_create_task.return_value = None # Mock response for all tests mock_response = MagicMock() mock_response.status_code = 200 mock_response.raise_for_status = MagicMock() # Create a test batch logging element test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-key.json", payload={"test": "data"}, s3_object_download_filename="test-file.json", ) # Test 1: Custom endpoint URL with bucket name 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 url = call_args[0][0] expected_url = "https://s3.amazonaws.com/test-bucket/2025-09-14/test-key.json" assert url == expected_url, f"Expected URL {expected_url}, got {url}" # Test 2: MinIO-compatible endpoint s3_logger_minio = S3Logger( s3_bucket_name="litellm-logs", s3_endpoint_url="https://minio.example.com:9000", s3_aws_access_key_id="minio-key", s3_aws_secret_access_key="minio-secret", s3_region_name="us-east-1", ) s3_logger_minio.async_httpx_client = AsyncMock() s3_logger_minio.async_httpx_client.put.return_value = mock_response asyncio.run(s3_logger_minio.async_upload_data_to_s3(test_element)) call_args_minio = s3_logger_minio.async_httpx_client.put.call_args assert call_args_minio is not None url_minio = call_args_minio[0][0] expected_minio_url = "https://minio.example.com:9000/litellm-logs/2025-09-14/test-key.json" assert url_minio == expected_minio_url, f"Expected MinIO URL {expected_minio_url}, got {url_minio}" # Test 3: Custom endpoint without bucket name (should fall back to default) s3_logger_no_bucket = S3Logger( 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_no_bucket.async_httpx_client = AsyncMock() s3_logger_no_bucket.async_httpx_client.put.return_value = mock_response asyncio.run(s3_logger_no_bucket.async_upload_data_to_s3(test_element)) call_args_no_bucket = s3_logger_no_bucket.async_httpx_client.put.call_args assert call_args_no_bucket is not None url_no_bucket = call_args_no_bucket[0][0] # Should use default S3 URL format when bucket is missing (bucket becomes None in URL) assert "s3.us-east-1.amazonaws.com" in url_no_bucket assert "https://" in url_no_bucket # Should not include the custom endpoint since bucket is missing assert "https://s3.amazonaws.com/" not in url_no_bucket # Test 4: Sync upload method with custom endpoint s3_logger_sync = S3Logger( s3_bucket_name="sync-bucket", s3_endpoint_url="https://custom.s3.endpoint.com", s3_aws_access_key_id="sync-key", s3_aws_secret_access_key="sync-secret", s3_region_name="us-east-1", ) mock_sync_client = MagicMock() mock_sync_client.put.return_value = mock_response with patch( "litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client, ): s3_logger_sync.upload_data_to_s3(test_element) call_args_sync = mock_sync_client.put.call_args assert call_args_sync is not None url_sync = call_args_sync[0][0] expected_sync_url = "https://custom.s3.endpoint.com/sync-bucket/2025-09-14/test-key.json" assert url_sync == expected_sync_url, f"Expected sync URL {expected_sync_url}, got {url_sync}" # Test 5: Download method with custom endpoint s3_logger_download = S3Logger( s3_bucket_name="download-bucket", s3_endpoint_url="https://download.s3.endpoint.com", s3_aws_access_key_id="download-key", s3_aws_secret_access_key="download-secret", s3_region_name="us-east-1", ) mock_download_response = MagicMock() mock_download_response.status_code = 200 mock_download_response.json = MagicMock(return_value={"downloaded": "data"}) s3_logger_download.async_httpx_client = AsyncMock() s3_logger_download.async_httpx_client.get.return_value = mock_download_response result = asyncio.run(s3_logger_download._download_object_from_s3("2025-09-14/download-test-key.json")) call_args_download = s3_logger_download.async_httpx_client.get.call_args assert call_args_download is not None url_download = call_args_download[0][0] expected_download_url = "https://download.s3.endpoint.com/download-bucket/2025-09-14/download-test-key.json" assert url_download == expected_download_url, ( f"Expected download URL {expected_download_url}, got {url_download}" ) assert result == {"downloaded": "data"} @patch("asyncio.create_task") @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") def test_s3_v2_virtual_hosted_style(self, mock_periodic_flush, mock_create_task): """Test s3_use_virtual_hosted_style parameter for virtual-hosted-style URLs""" from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement # Mock periodic_flush and create_task to prevent async task creation during init mock_periodic_flush.return_value = None mock_create_task.return_value = None # Mock response for all tests mock_response = MagicMock() mock_response.status_code = 200 mock_response.raise_for_status = MagicMock() # Create a test batch logging element test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-key.json", payload={"test": "data"}, s3_object_download_filename="test-file.json", ) # Test 1: Virtual-hosted-style with custom endpoint s3_logger_virtual = S3Logger( s3_bucket_name="test-bucket", s3_endpoint_url="https://s3.custom-endpoint.com", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_use_virtual_hosted_style=True, ) s3_logger_virtual.async_httpx_client = AsyncMock() s3_logger_virtual.async_httpx_client.put.return_value = mock_response asyncio.run(s3_logger_virtual.async_upload_data_to_s3(test_element)) call_args = s3_logger_virtual.async_httpx_client.put.call_args assert call_args is not None url = call_args[0][0] expected_url = "https://test-bucket.s3.custom-endpoint.com/2025-09-14/test-key.json" assert url == expected_url, f"Expected virtual-hosted-style URL {expected_url}, got {url}" # Test 2: Path-style (default behavior with s3_use_virtual_hosted_style=False) s3_logger_path = S3Logger( s3_bucket_name="test-bucket", s3_endpoint_url="https://s3.custom-endpoint.com", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_use_virtual_hosted_style=False, ) s3_logger_path.async_httpx_client = AsyncMock() s3_logger_path.async_httpx_client.put.return_value = mock_response asyncio.run(s3_logger_path.async_upload_data_to_s3(test_element)) call_args_path = s3_logger_path.async_httpx_client.put.call_args assert call_args_path is not None url_path = call_args_path[0][0] expected_path_url = "https://s3.custom-endpoint.com/test-bucket/2025-09-14/test-key.json" assert url_path == expected_path_url, f"Expected path-style URL {expected_path_url}, got {url_path}" # Test 3: Virtual-hosted-style with http protocol s3_logger_http = S3Logger( s3_bucket_name="http-bucket", s3_endpoint_url="http://minio.local:9000", s3_aws_access_key_id="minio-key", s3_aws_secret_access_key="minio-secret", s3_region_name="us-east-1", s3_use_virtual_hosted_style=True, ) s3_logger_http.async_httpx_client = AsyncMock() s3_logger_http.async_httpx_client.put.return_value = mock_response asyncio.run(s3_logger_http.async_upload_data_to_s3(test_element)) call_args_http = s3_logger_http.async_httpx_client.put.call_args assert call_args_http is not None url_http = call_args_http[0][0] expected_http_url = "http://http-bucket.minio.local:9000/2025-09-14/test-key.json" assert url_http == expected_http_url, ( f"Expected virtual-hosted-style URL with http {expected_http_url}, got {url_http}" ) # Test 4: Sync upload method with virtual-hosted-style s3_logger_sync_virtual = S3Logger( s3_bucket_name="sync-bucket", s3_endpoint_url="https://storage.example.com", s3_aws_access_key_id="sync-key", s3_aws_secret_access_key="sync-secret", s3_region_name="us-east-1", s3_use_virtual_hosted_style=True, ) mock_sync_client = MagicMock() mock_sync_client.put.return_value = mock_response with patch( "litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client, ): s3_logger_sync_virtual.upload_data_to_s3(test_element) call_args_sync = mock_sync_client.put.call_args assert call_args_sync is not None url_sync = call_args_sync[0][0] expected_sync_url = "https://sync-bucket.storage.example.com/2025-09-14/test-key.json" assert url_sync == expected_sync_url, ( f"Expected virtual-hosted-style sync URL {expected_sync_url}, got {url_sync}" ) # Test 5: Download method with virtual-hosted-style s3_logger_download_virtual = S3Logger( s3_bucket_name="download-bucket", s3_endpoint_url="https://download.endpoint.com", s3_aws_access_key_id="download-key", s3_aws_secret_access_key="download-secret", s3_region_name="us-east-1", s3_use_virtual_hosted_style=True, ) mock_download_response = MagicMock() mock_download_response.status_code = 200 mock_download_response.json = MagicMock(return_value={"downloaded": "data"}) s3_logger_download_virtual.async_httpx_client = AsyncMock() s3_logger_download_virtual.async_httpx_client.get.return_value = mock_download_response result = asyncio.run(s3_logger_download_virtual._download_object_from_s3("2025-09-14/download-test-key.json")) call_args_download = s3_logger_download_virtual.async_httpx_client.get.call_args assert call_args_download is not None url_download = call_args_download[0][0] expected_download_url = "https://download-bucket.download.endpoint.com/2025-09-14/download-test-key.json" assert url_download == expected_download_url, ( f"Expected virtual-hosted-style download URL {expected_download_url}, got {url_download}" ) 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): from unittest.mock import AsyncMock import requests 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(): """ Test that async_upload_data_to_s3 retries on transient S3 503 Slow Down and succeeds on the second attempt. """ from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-retry.json", payload={"test": "retry"}, s3_object_download_filename="test-retry.json", ) # First call returns 503, second call returns 200 response_503 = MagicMock() response_503.status_code = 503 response_503.text = "" response_200 = MagicMock() response_200.status_code = 200 response_200.text = "" response_200.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = AsyncMock(side_effect=[response_503, response_200]) with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: await logger.async_upload_data_to_s3(test_element) # Verify PUT was called twice (retry after 503) assert logger.async_httpx_client.put.call_count == 2 _assert_jittered_backoff(mock_sleep, 1) @pytest.mark.asyncio async def test_async_upload_retries_on_s3_500(): """ Test that async_upload_data_to_s3 retries on transient S3 500 errors. """ from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-retry-500.json", payload={"test": "retry-500"}, s3_object_download_filename="test-retry-500.json", ) response_500 = MagicMock() response_500.status_code = 500 response_500.text = "" response_200 = MagicMock() response_200.status_code = 200 response_200.text = "" response_200.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = AsyncMock(side_effect=[response_500, response_200]) with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: await logger.async_upload_data_to_s3(test_element) assert logger.async_httpx_client.put.call_count == 2 _assert_jittered_backoff(mock_sleep, 1) @pytest.mark.asyncio async def test_async_upload_exhausts_retries_on_persistent_503(): """ Test that async_upload_data_to_s3 raises after exhausting all retries on persistent S3 503. """ from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-exhaust.json", payload={"test": "exhaust"}, s3_object_download_filename="test-exhaust.json", ) # All 3 attempts return 503 response_503 = MagicMock() response_503.status_code = 503 response_503.text = "" response_503.raise_for_status = MagicMock(side_effect=Exception("503 Service Unavailable")) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = AsyncMock(return_value=response_503) with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: with patch.object(logger, "handle_callback_failure") as mock_failure: await logger.async_upload_data_to_s3(test_element) # 3 PUT attempts total assert logger.async_httpx_client.put.call_count == 3 # 2 sleeps (between attempts 1-2 and 2-3) assert mock_sleep.call_count == 2 # Callback failure handler called after exhausting retries mock_failure.assert_called_once_with(callback_name="S3Logger") @pytest.mark.asyncio async def test_async_upload_retries_400_with_an_unknown_error_code(): """ A 400 is outside the retry set, so an unknown gets a single PUT and the "retry" outcome for the flush-level requeue, never an in-call backoff. """ from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-retry-400.json", payload={"test": "retry-400"}, s3_object_download_filename="test-retry-400.json", ) response_400 = MagicMock() response_400.status_code = 400 response_400.text = "SomethingElse" response_400.raise_for_status = MagicMock(side_effect=Exception("400 Bad Request")) response_200 = MagicMock() response_200.status_code = 200 response_200.text = "" response_200.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = AsyncMock(side_effect=[response_400, response_200]) with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: outcome = await logger.async_upload_data_to_s3(test_element) assert logger.async_httpx_client.put.call_count == 1 mock_sleep.assert_not_awaited() assert outcome is False _SIGV4_ACCESS_KEY = re.compile(r"Credential=(AKIA\d+)/") @pytest.fixture def rotating_profile(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> str: """ A real botocore profile whose credential_process hands out a new key generation on every call and expires inside the advisory refresh window, so RefreshableCredentials re-runs it on every property read. """ counter = tmp_path / "generation" script = tmp_path / "rotate_credentials.py" script.write_text( textwrap.dedent( f""" import json, sys from datetime import datetime, timedelta, timezone from pathlib import Path counter = Path({str(counter)!r}) generation = int(counter.read_text()) if counter.exists() else 0 counter.write_text(str(generation + 1)) expiry = (datetime.now(timezone.utc) + timedelta(minutes=12)).strftime("%Y-%m-%dT%H:%M:%SZ") json.dump( {{ "Version": 1, "AccessKeyId": f"AKIA{{generation}}", "SecretAccessKey": f"secret-{{generation}}", "SessionToken": f"token-{{generation}}", "Expiration": expiry, }}, sys.stdout, ) """ ) ) profile = f"rotating-{uuid.uuid4().hex}" (tmp_path / "config").write_text(f"[profile {profile}]\ncredential_process = {sys.executable} {script}\n") monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) return profile def _generation(request: httpx.Request) -> tuple[str, str]: """(access key generation, session token generation) SigV4 baked into one request.""" access_key = _SIGV4_ACCESS_KEY.search(request.headers["Authorization"]) assert access_key is not None return access_key.group(1).removeprefix("AKIA"), request.headers["X-Amz-Security-Token"].removeprefix("token-") @asynccontextmanager async def _s3_logger_on_production_handler(profile: str, statuses: list[int | type[httpx.TransportError]]): """ S3Logger wired to the real AsyncHTTPHandler over an httpx MockTransport that answers with the given statuses (or raises the given transport errors) in order, so the handler's own error mapping is exercised. """ requests: list[httpx.Request] = [] replies = iter(statuses) def respond(request: httpx.Request) -> httpx.Response: requests.append(request) reply = next(replies) if not isinstance(reply, int): raise reply("S3 attempt failed", request=request) return httpx.Response(reply, request=request, text="SignatureDoesNotMatch") handler = AsyncHTTPHandler() handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) logger = S3Logger( s3_bucket_name="test-bucket", s3_region_name="us-east-1", s3_aws_profile_name=profile, s3_flush_interval=3600, ) logger.async_httpx_client = handler with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: yield logger, requests, mock_sleep await handler.client.aclose() @pytest.mark.asyncio async def test_async_upload_signs_with_one_frozen_credential_snapshot(rotating_profile: str, caplog): """ RefreshableCredentials refreshes on every property read once inside the advisory window, so signing off the live object would mix the access key of one generation with the token of the next. """ test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-frozen.json", payload={"test": "frozen"}, s3_object_download_filename="test-frozen.json", ) async with _s3_logger_on_production_handler(rotating_profile, [200]) as (logger, requests, _): await logger.async_upload_data_to_s3(test_element) assert len(requests) == 1 access_key_generation, token_generation = _generation(requests[0]) assert access_key_generation == token_generation assert "Error uploading to s3" not in caplog.text @pytest.mark.asyncio async def test_async_upload_retries_403_with_fresh_credentials_and_signature(rotating_profile: str, caplog): """ A 403 (SignatureDoesNotMatch after an IMDS rotation) must be retried, and the retry must fetch credentials again and carry a signature computed from that newer generation. """ test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-403.json", payload={"test": "403"}, s3_object_download_filename="test-403.json", ) async with _s3_logger_on_production_handler(rotating_profile, [403, 200]) as (logger, requests, mock_sleep): await logger.async_upload_data_to_s3(test_element) assert len(requests) == 2 first_key, first_token = _generation(requests[0]) second_key, second_token = _generation(requests[1]) assert first_key == first_token assert second_key == second_token assert int(second_key) > int(first_key) assert requests[1].headers["Authorization"] != requests[0].headers["Authorization"] _assert_jittered_backoff(mock_sleep, 1) assert "Error uploading to s3" not in caplog.text @pytest.mark.asyncio async def test_async_upload_exhausts_403_retries_through_production_http_handler(rotating_profile: str, caplog): test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-403-exhausted.json", payload={"test": "403-exhausted"}, s3_object_download_filename="test-403-exhausted.json", ) async with _s3_logger_on_production_handler(rotating_profile, [403, 403, 403]) as (logger, requests, mock_sleep): await logger.async_upload_data_to_s3(test_element) assert len(requests) == 3 _assert_jittered_backoff(mock_sleep, 2) assert "Error uploading to s3" in caplog.text @pytest.mark.asyncio @pytest.mark.parametrize("transient_status", [500, 503]) async def test_async_upload_recovers_from_transient_5xx_through_production_http_handler( transient_status: int, rotating_profile: str, caplog: pytest.LogCaptureFixture ): """ AsyncHTTPHandler.put raises MaskedHTTPStatusError on 5xx instead of returning the response, so a retry loop that only inspects returned status codes never runs (#42868). """ test_element = s3BatchLoggingElement( s3_object_key=f"2025-09-14/test-{transient_status}.json", payload={"test": str(transient_status)}, s3_object_download_filename=f"test-{transient_status}.json", ) async with _s3_logger_on_production_handler(rotating_profile, [transient_status, 200]) as ( logger, requests, mock_sleep, ): uploaded = await logger.async_upload_data_to_s3(test_element) assert uploaded is True assert len(requests) == 2 assert all(request.method == "PUT" for request in requests) assert requests[0].url == requests[1].url assert requests[0].content == requests[1].content _assert_jittered_backoff(mock_sleep, 1) assert "Error uploading to s3" not in caplog.text @pytest.mark.asyncio async def test_async_upload_is_single_attempted_on_404_through_production_http_handler(rotating_profile: str, caplog): test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-404.json", payload={"test": "404"}, s3_object_download_filename="test-404.json", ) async with _s3_logger_on_production_handler(rotating_profile, [404]) as (logger, requests, mock_sleep): outcome = await logger.async_upload_data_to_s3(test_element) assert len(requests) == 1 assert outcome is False mock_sleep.assert_not_awaited() assert "Error uploading to s3" in caplog.text @pytest.mark.asyncio async def test_async_upload_access_denied_403_is_retried_and_then_requeued(rotating_profile: str, caplog): test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-403-denied.json", payload={"test": "403-denied"}, s3_object_download_filename="test-403-denied.json", ) requests: list[httpx.Request] = [] def respond(request: httpx.Request) -> httpx.Response: requests.append(request) return httpx.Response(403, request=request, text="AccessDenied") handler = AsyncHTTPHandler() handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) logger = S3Logger( s3_bucket_name="test-bucket", s3_region_name="us-east-1", s3_aws_profile_name=rotating_profile, s3_flush_interval=3600, ) logger.async_httpx_client = handler with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: outcome = await logger.async_upload_data_to_s3(test_element) await handler.client.aclose() assert outcome is False assert len(requests) == 3 _assert_jittered_backoff(mock_sleep, 2) assert "Error uploading to s3" in caplog.text @pytest.mark.asyncio async def test_async_upload_retries_429_through_production_http_handler(rotating_profile: str, caplog): """S3 and S3-compatible stores answer throttling with 429 too, so it backs off and retries like a 503.""" async with _s3_logger_on_production_handler(rotating_profile, [429, 200]) as (logger, requests, mock_sleep): uploaded = await logger.async_upload_data_to_s3(_element({"test": "429"}, "429")) assert uploaded is True assert len(requests) == 2 assert requests[0].content == requests[1].content _assert_jittered_backoff(mock_sleep, 1) assert "Error uploading to s3" not in caplog.text @pytest.mark.asyncio async def test_async_upload_retries_a_timed_out_attempt_through_production_http_handler(rotating_profile: str, caplog): """AsyncHTTPHandler.put turns an httpx timeout into litellm.Timeout, which must count as a retryable attempt.""" async with _s3_logger_on_production_handler(rotating_profile, [httpx.ReadTimeout, 200]) as ( logger, requests, mock_sleep, ): uploaded = await logger.async_upload_data_to_s3(_element({"test": "timeout"}, "timeout")) assert uploaded is True assert len(requests) == 2 _assert_jittered_backoff(mock_sleep, 1) assert "Error uploading to s3" not in caplog.text @pytest.mark.asyncio @pytest.mark.parametrize( "error", [ pytest.param(httpx.ConnectError("connection refused"), id="connect-error"), pytest.param(httpx.RemoteProtocolError("server disconnected"), id="remote-protocol-error"), pytest.param(httpx.ReadTimeout("read timed out"), id="read-timeout"), ], ) async def test_async_upload_retries_transport_errors_then_succeeds(error: Exception) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = AsyncMock(side_effect=[error, _ok_response()]) with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: uploaded = await logger.async_upload_data_to_s3(_element({"test": "transport"}, "transport")) assert uploaded is True assert logger.async_httpx_client.put.call_count == 2 _assert_jittered_backoff(mock_sleep, 1) @pytest.mark.asyncio async def test_async_upload_exhausts_transport_error_retries_and_keeps_the_event_queued() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = AsyncMock(side_effect=httpx.ConnectError("connection refused")) logger.log_queue = [_element({"test": "down"}, "down")] with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: await logger.flush_queue() assert logger.async_httpx_client.put.call_count == 3 _assert_jittered_backoff(mock_sleep, 2) assert [element.s3_object_key for element in logger.log_queue] == ["2025-09-14/test-down.json"] @pytest.mark.asyncio async def test_async_upload_sends_every_attempt_with_a_finite_deadline(rotating_profile: str) -> None: """ The logger's handler is built with explicit params and no timeout, and the handler forwards that None to httpx as "no timeout", so without a per-attempt deadline a stalled S3 endpoint pins the upload forever. """ async with _s3_logger_on_production_handler(rotating_profile, [503, 200]) as (logger, requests, _): await logger.async_upload_data_to_s3(_element({"test": "deadline"}, "deadline")) assert len(requests) == 2 for request in requests: deadlines = request.extensions["timeout"] assert all(isinstance(limit, (int, float)) and limit > 0 for limit in deadlines.values()), deadlines @pytest.mark.asyncio async def test_backoff_is_jittered_so_throttled_uploads_do_not_retry_in_lockstep() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = AsyncMock( side_effect=[_transient_failure_response(), _ok_response(), _transient_failure_response(), _ok_response()] ) with ( patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, patch("random.random", side_effect=[0.0, 0.999]), ): await logger.async_upload_data_to_s3(_element({"i": 0}, "0")) await logger.async_upload_data_to_s3(_element({"i": 1}, "1")) first, second = (sleep_call.args[0] for sleep_call in mock_sleep.call_args_list) assert first != second assert 0.5 <= first < 1.5 and 0.5 <= second < 1.5 @pytest.mark.asyncio async def test_dropped_uploads_are_logged_once_per_flush_naming_each_reason(caplog) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) terminal = (_element({"i": "terminal-0"}, "terminal-0"), _element({"i": "terminal-1"}, "terminal-1")) expired = _element({"i": "expired"}, "expired").model_copy(update={"retrying_since": _NOW - 7200}) fresh = _element({"i": "fresh"}, "fresh") async def put( url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None ) -> MagicMock: if "test-terminal-" in url: return _terminal_failure_response() if url.endswith("test-expired.json"): return _transient_failure_response() return _ok_response() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [*terminal, expired, fresh] with ( patch("asyncio.sleep", new_callable=AsyncMock), patch("time.monotonic", return_value=_NOW), ): await logger.flush_queue() assert logger.log_queue == [] drop_lines = [record.getMessage() for record in caplog.records if "uploads dropped" in record.getMessage()] assert len(drop_lines) == 1, drop_lines assert "3 uploads dropped" in drop_lines[0] assert "2 rejected as terminal" in drop_lines[0] assert "1 retrying longer than s3_max_retry_age_seconds=3600" in drop_lines[0] def test_sync_upload_retries_403_with_fresh_signature(rotating_profile: str, monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("AWS_ACCESS_KEY_ID", raising=False) monkeypatch.delenv("AWS_SECRET_ACCESS_KEY", raising=False) monkeypatch.delenv("AWS_SESSION_TOKEN", raising=False) monkeypatch.setenv("AWS_PROFILE", rotating_profile) logger = S3Logger(s3_bucket_name="test-bucket", s3_region_name="us-east-1", s3_flush_interval=3600) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-sync-403.json", payload={"test": "sync-403"}, s3_object_download_filename="test-sync-403.json", ) requests: list[httpx.Request] = [] replies = iter([403, 200]) def respond(request: httpx.Request) -> httpx.Response: requests.append(request) return httpx.Response(next(replies), request=request, text="SignatureDoesNotMatch") handler = HTTPHandler() handler.client = httpx.Client(transport=httpx.MockTransport(respond)) with ( patch( # test-quality-ok: sync upload builds its HTTPHandler per call, there is no injection seam for it "litellm.integrations.s3_v2.get_httpx_client", return_value=handler ), patch("time.sleep") as mock_sleep, ): logger.upload_data_to_s3(test_element) assert len(requests) == 2 first_key, first_token = _generation(requests[0]) second_key, second_token = _generation(requests[1]) assert first_key == first_token assert second_key == second_token assert int(second_key) > int(first_key) _assert_jittered_backoff(mock_sleep, 1) def test_sync_upload_retries_on_s3_503(): """ Test that the sync upload_data_to_s3 retries on transient S3 503. """ from unittest.mock import MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-sync-retry.json", payload={"test": "sync-retry"}, s3_object_download_filename="test-sync-retry.json", ) response_503 = MagicMock() response_503.status_code = 503 response_200 = MagicMock() response_200.status_code = 200 response_200.raise_for_status = MagicMock() mock_sync_client = MagicMock() mock_sync_client.put = MagicMock(side_effect=[response_503, response_200]) with patch( "litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client, ): with patch("time.sleep") as mock_sleep: logger.upload_data_to_s3(test_element) assert mock_sync_client.put.call_count == 2 _assert_jittered_backoff(mock_sleep, 1) @pytest.mark.asyncio async def test_async_log_event_skips_when_standard_logging_object_missing(): """ Reproduces the bug where _async_log_event_base raises ValueError when kwargs has no standard_logging_object (e.g. call_type=afile_delete). The S3 logger should skip gracefully, not raise. """ logger = S3Logger( s3_bucket_name="test-bucket", s3_region_name="us-east-1", s3_aws_access_key_id="fake", s3_aws_secret_access_key="fake", ) kwargs_without_slo = { "call_type": "afile_delete", "model": None, "litellm_call_id": "test-call-id", } start_time = datetime.utcnow() end_time = datetime.utcnow() # Spy on handle_callback_failure — should NOT be called if we skip gracefully. # Without the fix, the ValueError is caught by the except block which calls # handle_callback_failure. With the fix, we return early and never hit except. with patch.object(logger, "handle_callback_failure") as mock_failure: await logger._async_log_event_base( kwargs=kwargs_without_slo, response_obj=None, start_time=start_time, end_time=end_time, ) assert not mock_failure.called, ( "handle_callback_failure should not be called — " "missing standard_logging_object should be a graceful skip, not an error" ) # Nothing should have been queued (catches the case where code falls # through without returning and appends None to the queue) assert len(logger.log_queue) == 0, "log_queue should be empty when standard_logging_object is missing" @pytest.mark.asyncio async def test_strip_base64_removes_file_and_nontext_entries(): logger = S3Logger(s3_strip_base64_files=True) payload = { "messages": [ { "role": "user", "content": [ {"type": "text", "text": "Hello world"}, { "type": "image", "file": {"file_data": "data:image/png;base64,AAAA"}, }, { "type": "file", "file": {"file_data": "data:application/pdf;base64,BBBB"}, }, ], }, { "role": "assistant", "content": [ {"type": "text", "text": "Response"}, { "type": "audio", "file": {"file_data": "data:audio/wav;base64,CCCC"}, }, ], }, ] } stripped = await logger._strip_base64_from_messages(payload) # 1️⃣ File/image/audio entries are removed assert len(stripped["messages"][0]["content"]) == 1 assert stripped["messages"][0]["content"][0]["text"] == "Hello world" assert len(stripped["messages"][1]["content"]) == 1 assert stripped["messages"][1]["content"][0]["text"] == "Response" # 2️⃣ No 'file' keys remain for msg in stripped["messages"]: for content in msg["content"]: assert "file" not in content assert content.get("type") == "text" @pytest.mark.asyncio async def test_strip_base64_keeps_non_file_content(): logger = S3Logger(s3_strip_base64_files=True) payload = { "messages": [ { "role": "user", "content": [ {"type": "text", "text": "Just text"}, {"type": "text", "text": "Another message"}, ], } ] } stripped = await logger._strip_base64_from_messages(payload) # Should not modify pure text messages assert stripped["messages"][0]["content"] == payload["messages"][0]["content"] @pytest.mark.asyncio async def test_strip_base64_handles_empty_or_missing_messages(): logger = S3Logger(s3_strip_base64_files=True) # Missing messages key payload_no_messages = {} stripped1 = await logger._strip_base64_from_messages(payload_no_messages) assert stripped1 == payload_no_messages # Empty messages list payload_empty = {"messages": []} stripped2 = await logger._strip_base64_from_messages(payload_empty) assert stripped2 == payload_empty @pytest.mark.asyncio async def test_strip_base64_mixed_nested_objects(): """ Handles weird/nested content structures gracefully. """ logger = S3Logger(s3_strip_base64_files=True) payload = { "messages": [ { "role": "system", "content": [ {"type": "text", "text": "Keep me"}, {"type": "custom", "metadata": "ignore but non-text"}, {"foo": "bar"}, {"file": {"file_data": "data:application/pdf;base64,XXX"}}, ], "extra": {"trace_id": "123"}, } ] } stripped = await logger._strip_base64_from_messages(payload) # Custom/non-text and file entries removed content = stripped["messages"][0]["content"] assert len(content) == 2 assert {"type": "text", "text": "Keep me"} in content assert {"foo": "bar"} in content # Extra metadata preserved assert stripped["messages"][0]["extra"]["trace_id"] == "123" @pytest.mark.asyncio async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch): """ Test that s3_verify=False is properly handled and not treated as None. This is a regression test for the bug where s3_verify=False was being ignored because 'False or s3_verify' would evaluate to s3_verify (None). """ from unittest.mock import AsyncMock, patch import litellm # Set up s3_callback_params with s3_verify=False monkeypatch.setattr( litellm, "s3_callback_params", { "s3_bucket_name": "test-bucket", "s3_endpoint_url": "https://localhost:443", "s3_aws_access_key_id": "minioadmin", "s3_aws_secret_access_key": "minioadmin", "s3_region_name": "us-east-1", "s3_verify": False, # This should NOT be ignored "s3_use_ssl": False, # This should also NOT be ignored }, ) with patch("asyncio.create_task"): with patch("litellm.integrations.s3_v2.get_async_httpx_client") as mock_get_client: mock_client = AsyncMock() mock_get_client.return_value = mock_client # Create logger logger = S3Logger() # Verify s3_verify is False, not None assert logger.s3_verify is False, f"Expected s3_verify=False, got {logger.s3_verify}" assert logger.s3_use_ssl is False, f"Expected s3_use_ssl=False, got {logger.s3_use_ssl}" # Verify that get_async_httpx_client was called with ssl_verify=False mock_get_client.assert_called_once() call_kwargs = mock_get_client.call_args.kwargs assert "params" in call_kwargs, "params should be passed to get_async_httpx_client" assert call_kwargs["params"] == {"ssl_verify": False}, ( f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}" ) @pytest.mark.asyncio async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch): """ Test that s3_verify=None uses default behavior. """ from unittest.mock import AsyncMock, patch import litellm # Set up s3_callback_params without s3_verify monkeypatch.setattr( litellm, "s3_callback_params", { "s3_bucket_name": "test-bucket", "s3_aws_access_key_id": "test-key", "s3_aws_secret_access_key": "test-secret", "s3_region_name": "us-east-1", }, ) with patch("asyncio.create_task"): with patch("litellm.integrations.s3_v2.get_async_httpx_client") as mock_get_client: mock_client = AsyncMock() mock_get_client.return_value = mock_client # Create logger without explicit s3_verify logger = S3Logger() # Verify s3_verify is None (default) assert logger.s3_verify is None, f"Expected s3_verify=None, got {logger.s3_verify}" # Verify that get_async_httpx_client was called mock_get_client.assert_called_once() call_kwargs = mock_get_client.call_args.kwargs # When s3_verify is None, params={'ssl_verify': None} which is fine - uses default behavior # The important thing is it's not False if "params" in call_kwargs and call_kwargs["params"] is not None: assert call_kwargs["params"].get("ssl_verify") is None # Either params is None or params={'ssl_verify': None} is acceptable @pytest.mark.asyncio async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatch: pytest.MonkeyPatch): """ Test that when s3_verify=False, the actual httpx client has verify=False. This validates that ssl_verify=False flows through to the httpx.AsyncClient. """ from unittest.mock import patch import litellm # Set up s3_callback_params with s3_verify=False monkeypatch.setattr( litellm, "s3_callback_params", { "s3_bucket_name": "test-bucket", "s3_endpoint_url": "https://localhost:443", "s3_aws_access_key_id": "minioadmin", "s3_aws_secret_access_key": "minioadmin", "s3_region_name": "us-east-1", "s3_verify": False, }, ) with patch("asyncio.create_task"): # Create logger - this creates the httpx client logger = S3Logger() # Verify the logger has s3_verify=False assert logger.s3_verify is False # Check the actual httpx client has verify=False # The async_httpx_client.client is the actual httpx.AsyncClient if hasattr(logger.async_httpx_client, "client"): httpx_client = logger.async_httpx_client.client # Check the _verify attribute (httpx internal) if hasattr(httpx_client, "_verify"): assert httpx_client._verify is False, f"Expected httpx client _verify=False, got {httpx_client._verify}" @pytest.mark.asyncio async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch): """ Test that the async httpx client respects s3_verify=False. """ from unittest.mock import AsyncMock, MagicMock, patch import litellm from litellm.types.integrations.s3_v2 import s3BatchLoggingElement # Set up s3_callback_params with s3_verify=False monkeypatch.setattr( litellm, "s3_callback_params", { "s3_bucket_name": "test-bucket", "s3_endpoint_url": "https://localhost:443", "s3_aws_access_key_id": "minioadmin", "s3_aws_secret_access_key": "minioadmin", "s3_region_name": "us-east-1", "s3_verify": False, }, ) with patch("asyncio.create_task"): logger = S3Logger() # Verify s3_verify is False assert logger.s3_verify is False # Create test element test_element = s3BatchLoggingElement( s3_object_key="2025-11-03/test-key.json", payload={"test": "data"}, s3_object_download_filename="test-file.json", ) # Mock the async httpx client's put method mock_response = MagicMock() mock_response.status_code = 200 mock_response.raise_for_status = MagicMock() logger.async_httpx_client.put = AsyncMock(return_value=mock_response) # Call async upload await logger.async_upload_data_to_s3(test_element) # Verify put was called assert logger.async_httpx_client.put.called # Check that the async httpx client was created with verify=False if hasattr(logger.async_httpx_client, "client"): httpx_client = logger.async_httpx_client.client if hasattr(httpx_client, "_verify"): assert httpx_client._verify is False, ( f"Expected async httpx client _verify=False, got {httpx_client._verify}" ) @pytest.mark.asyncio async def test_strip_base64_recursive_redaction(): logger = S3Logger(s3_strip_base64_files=True) payload = { "messages": [ { "content": [ {"type": "text", "text": "normal text"}, { "type": "text", "text": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUg", }, { "type": "text", "text": "Nested: {'data': 'data:application/pdf;base64,AAA...'}", }, {"file": {"file_data": "data:application/pdf;base64,AAAA"}}, {"metadata": {"preview": "data:audio/mp3;base64,AAAAA=="}}, ] } ] } result = await logger._strip_base64_from_messages(payload) content = result["messages"][0]["content"] # Dropped file-type entries assert not any("file" in c for c in content) # Base64 redacted globally import json for c in content: if isinstance(c, dict): s = json.dumps(c).lower() # "[base64_redacted]" is fine, but raw base64 is not assert "base64," not in s, f"Found real base64 blob in: {s}" # -------------------------------------------------------------- # Shared fixture that silences asyncio.create_task during tests # -------------------------------------------------------------- @pytest_asyncio.fixture(loop_scope="function") async def cancel_s3_periodic_flush_tasks(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[None]: periodic_flush_tasks: list[asyncio.Task[object]] = [] original_create_task = asyncio.create_task def track_create_task( coro: Coroutine[object, object, object], *, name: str | None = None, context: Context | None = None, ) -> asyncio.Task[object]: if context is None: task = original_create_task(coro, name=name) else: task = original_create_task(coro, name=name, context=context) if coro.__qualname__.endswith("periodic_flush"): periodic_flush_tasks.append(task) return task monkeypatch.setattr(asyncio, "create_task", track_create_task) yield for task in periodic_flush_tasks: if not task.done(): task.cancel() await asyncio.gather(*periodic_flush_tasks, return_exceptions=True) @pytest.fixture(autouse=True) def patch_asyncio_create_task(request): """Prevent 'no running event loop' errors when S3Logger calls asyncio.create_task().""" if request.node.originalname in { "test_basic_s3_logging", "test_basic_s3_v2_logging", "test_basic_s3_v2_logging_failure", }: yield return with patch("asyncio.create_task"): yield # -------------------------------------------------------------- # Parametrized prefix combination test # -------------------------------------------------------------- @pytest.mark.parametrize( "use_team_prefix,use_key_prefix,team_alias,key_alias,expected_prefix", [ (False, False, "teamA", "keyA", ""), (True, False, "teamA", "keyA", "teamA/"), (False, True, "teamA", "keyA", "keyA/"), (True, True, "teamA", "keyA", "teamA/keyA/"), (True, True, None, "keyA", "keyA/"), (True, True, "teamA", None, "teamA/"), (True, True, None, None, ""), ], ) def test_s3_object_key_prefix_combinations(use_team_prefix, use_key_prefix, team_alias, key_alias, expected_prefix): """ Validate correct S3 prefix composition for team alias + key alias combinations. """ with patch("litellm.integrations.s3_v2.get_s3_object_key") as mock_get_key: mock_get_key.return_value = "mocked/s3/object/key.json" logger = S3Logger( s3_bucket_name="test-bucket", s3_region_name="us-east-1", s3_use_team_prefix=use_team_prefix, s3_use_key_prefix=use_key_prefix, ) payload = StandardLoggingPayload( id="abc123", metadata={ "user_api_key_team_alias": team_alias, "user_api_key_alias": key_alias, }, messages=[{"role": "user", "content": [{"type": "text", "text": "hi"}]}], ) result = logger.create_s3_batch_logging_element(datetime.utcnow(), payload) assert result is not None mock_get_key.assert_called_once() prefix_arg = mock_get_key.call_args.kwargs.get("prefix") assert prefix_arg == expected_prefix, ( f"Expected prefix '{expected_prefix}', got '{prefix_arg}' " f"for team={team_alias}, key={key_alias}, " f"use_team_prefix={use_team_prefix}, use_key_prefix={use_key_prefix}" ) # -------------------------------------------------------------- # Test prefix priority and concatenation # -------------------------------------------------------------- def test_prefix_priority_and_path_construction(): """ Validate that prefix components are ordered and joined with '/' only once. """ with patch("litellm.integrations.s3_v2.get_s3_object_key") as mock_get_key: mock_get_key.return_value = "mocked/key" logger = S3Logger(s3_use_team_prefix=True, s3_use_key_prefix=True) payload = StandardLoggingPayload( id="xyz999", metadata={ "user_api_key_team_alias": "Team-Alpha", "user_api_key_alias": "API-12345", }, messages=[], ) logger.create_s3_batch_logging_element(datetime.utcnow(), payload) prefix_arg = mock_get_key.call_args.kwargs.get("prefix", "") assert prefix_arg == "Team-Alpha/API-12345/" assert "//" not in prefix_arg # -------------------------------------------------------------- # Test when prefixes are disabled # -------------------------------------------------------------- def test_prefix_absent_when_flags_disabled(): """ Verify prefix is omitted entirely when prefix flags are False. """ with patch("litellm.integrations.s3_v2.get_s3_object_key") as mock_get_key: mock_get_key.return_value = "mocked/key" logger = S3Logger(s3_use_team_prefix=False, s3_use_key_prefix=False) payload = StandardLoggingPayload( id="no-prefix", metadata={ "user_api_key_team_alias": "team-x", "user_api_key_alias": "key-x", }, messages=[], ) logger.create_s3_batch_logging_element(datetime.utcnow(), payload) prefix_arg = mock_get_key.call_args.kwargs.get("prefix", None) assert prefix_arg == "", f"Expected empty prefix, got {prefix_arg}" # -------------------------------------------------------------- # Integration-style test (asyncio fixture will patch create_task) # -------------------------------------------------------------- @pytest.mark.asyncio async def test_combined_prefix_reflects_in_s3_object_key(): """ Integration-style test ensuring final s3_object_key includes both prefixes correctly. """ logger = S3Logger(s3_use_team_prefix=True, s3_use_key_prefix=True) payload = StandardLoggingPayload( id="int-test", metadata={ "user_api_key_team_alias": "myteam", "user_api_key_alias": "apikey", }, messages=[], ) result = logger.create_s3_batch_logging_element(datetime.utcnow(), payload) key = result.s3_object_key assert "myteam/apikey/" in key, f"Expected both prefixes in key: {key}" def test_s3_object_key_sanitizes_slashes_and_colons_in_file_name(): from litellm.integrations.s3 import get_s3_object_key start_time = datetime(2026, 2, 11, 0, 35, 18, 391582) file_name = "time-00-35-18-391582_arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/gl18r6skk9yy" key = get_s3_object_key( s3_path="LiteLLMAPPLogs", prefix="myteam/", start_time=start_time, s3_file_name=file_name, ) assert key == ( "LiteLLMAPPLogs/myteam/2026-02-11/" "time-00-35-18-391582_arn_aws_bedrock_us-east-1_123456789012_model-invocation-job_gl18r6skk9yy.json" ) @pytest.mark.parametrize( "response_id", [ "s3://example-batch-bucket/litellm-bedrock-files/input.jsonl", "gs://example-batch-bucket/litellm-vertex-files/input.jsonl", ], ) def test_s3_object_key_has_no_colon_for_cloud_uri_file_ids(response_id: str): from litellm.integrations.s3 import get_s3_object_key key = get_s3_object_key( s3_path="", prefix="", start_time=datetime(2026, 9, 7, 4, 51, 6, 685889), s3_file_name=f"time-04-51-06-685889_{response_id}", ) filename = key.rsplit("/", 1)[-1] assert ":" not in filename assert filename.endswith("_input.jsonl.json") def test_create_s3_batch_logging_element_flat_key_for_arn_response_id(): """End-to-end through the s3_v2 element builder: an ARN response id must yield a flat file directly under the date segment.""" logger = S3Logger(s3_use_team_prefix=False, s3_use_key_prefix=False) payload = StandardLoggingPayload( id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/gl18r6skk9yy", metadata={}, messages=[], ) start_time = datetime(2026, 2, 11, 0, 35, 18, 391582) result = logger.create_s3_batch_logging_element(start_time, payload) assert result is not None date_segment = "2026-02-11/" file_segment = result.s3_object_key.split(date_segment, 1)[1] assert "/" not in file_segment, f"Expected flat file under date segment, got: {result.s3_object_key}" assert file_segment.endswith("model-invocation-job_gl18r6skk9yy.json") # -------------------------------------------------------------- # object keys bounded to S3 and filesystem-compatible limits # -------------------------------------------------------------- def _oversized_response_id() -> str: return "resp_" + "A" * 1100 def test_s3_object_key_at_the_segment_byte_limit_is_left_alone(): """A final key segment that still fits is left byte-identical.""" from litellm.integrations.s3 import get_s3_object_key start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) file_name = "x" * 250 key = get_s3_object_key(s3_path="input", prefix="", start_time=start_time, s3_file_name=file_name) assert key == f"input/2026-08-24/{file_name}.json" assert len(key.rsplit("/", 1)[1].encode("utf-8")) == 255 def test_s3_object_key_just_over_the_file_name_limit_is_bounded(): """A 256-byte final key segment is shortened below the local filesystem limit.""" from litellm.integrations.s3 import get_s3_object_key key = get_s3_object_key( s3_path="input", prefix="", start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), s3_file_name="x" * 251, ) assert len(key.rsplit("/", 1)[1].encode("utf-8")) <= 255 def test_s3_object_key_bounds_the_file_segment_when_the_full_key_fits(): """A locally stored object name is bounded even when its full S3 key is below 1024 bytes.""" import hashlib from litellm.constants import MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import get_s3_object_key file_name = "time-01-14-19-015434_resp_" + "A" * 300 key = get_s3_object_key( s3_path="requests", prefix="", start_time=datetime(2026, 9, 12, 1, 14, 19, 15434), s3_file_name=file_name, ) file_segment = key.rsplit("/", 1)[1] assert len(key.encode("utf-8")) < MAX_S3_OBJECT_KEY_BYTES assert len(file_segment.encode("utf-8")) <= 255 assert file_segment.startswith("time-01-14-19-015434_resp_") assert file_segment.endswith(f"_{hashlib.sha256(file_name.encode('utf-8')).hexdigest()}.json") def test_s3_object_key_file_segment_limit_counts_utf8_bytes(): """The final key segment is bounded by encoded bytes without splitting a character.""" from litellm.integrations.s3 import get_s3_object_key key = get_s3_object_key( s3_path="requests", prefix="", start_time=datetime(2026, 9, 12, 1, 14, 19, 15434), s3_file_name="time-01-14-19-015434_resp_" + "日" * 100, ) file_segment = key.rsplit("/", 1)[1] assert len(file_segment.encode("utf-8")) <= 255 assert "\ufffd" not in file_segment def test_s3_object_key_at_the_full_key_byte_limit_is_left_alone(): """A key of exactly 1024 bytes whose segments all fit is left byte-identical.""" from litellm.constants import MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import get_s3_object_key file_name = "x" * 200 fixed_len = len(f"/2026-08-24/{file_name}.json") s3_path = "/".join(["p" * 200] * 4) + "/" + "q" * (MAX_S3_OBJECT_KEY_BYTES - fixed_len - 804) key = get_s3_object_key( s3_path=s3_path, prefix="", start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), s3_file_name=file_name, ) assert key == f"{s3_path}/2026-08-24/{file_name}.json" assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES def test_s3_object_key_is_bounded_for_oversized_response_id(): """An oversized Responses API id is shortened to a readable head plus a digest.""" import hashlib from litellm.constants import MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import get_s3_object_key start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) file_name = f"time-06-18-41-948021_{_oversized_response_id()}" key = get_s3_object_key(s3_path="input", prefix="DefaultTeamProd/", start_time=start_time, s3_file_name=file_name) assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES assert key.startswith("input/DefaultTeamProd/2026-08-24/time-06-18-41-948021_resp_") assert key.endswith(f"_{hashlib.sha256(file_name.encode('utf-8')).hexdigest()}.json") @pytest.mark.parametrize( "s3_path,prefix", [ ("input", ""), ("a" * 900, ""), ("input", "team-" + "b" * 900 + "/"), ("c" * 600, "team-" + "d" * 600 + "/key-" + "e" * 600 + "/"), # many short segments, so the trim lands exactly on the budget edge ("", "ssss/" * 200), ], ) def test_s3_object_key_is_bounded_for_long_paths_and_aliases(s3_path: str, prefix: str): """Long paths, team aliases and key aliases stay within the cap.""" from litellm.constants import MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import get_s3_object_key key = get_s3_object_key( s3_path=s3_path, prefix=prefix, start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", ) assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES assert key.endswith(".json") assert "/2026-08-24/" in key or key.startswith("2026-08-24/") assert "/" not in key.rsplit("2026-08-24/", 1)[1] def test_s3_object_key_trimmed_prefixes_stay_distinct_per_operator(): """Prefixes that differ only past the trim point keep separate folders.""" from litellm.constants import MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import get_s3_object_key start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) keys = [ get_s3_object_key( s3_path="input", prefix="team-" + "b" * 1000 + suffix + "/", start_time=start_time, s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", ) for suffix in ("-one", "-two") ] assert keys[0] != keys[1] assert all(key.startswith("input/team-" + "b" * 900) for key in keys) assert all(len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES for key in keys) def test_s3_object_key_bounded_prefix_never_splits_a_multibyte_character(): """A multibyte prefix is trimmed on a character boundary.""" from litellm.constants import MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import get_s3_object_key s3_path = "\u65e5\u672c\u8a9e" * 200 key = get_s3_object_key( s3_path=s3_path, prefix="\u30c1\u30fc\u30e0" * 200 + "/", start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", ) assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES assert key.startswith(s3_path[:100]) assert "\ufffd" not in key def test_s3_object_key_stays_unique_for_ids_sharing_a_head(): """Ids sharing a visible head still get distinct keys.""" from litellm.integrations.s3 import get_s3_object_key start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) keys = { get_s3_object_key( s3_path="input", prefix="", start_time=start_time, s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}{suffix}", ) for suffix in ("first", "second", "third") } assert len(keys) == 3 def test_s3_object_key_bounding_matches_the_documented_layout(): """The bounded key is `//_.json`.""" import hashlib from litellm.integrations.s3 import get_s3_object_key file_name = f"time-06-18-41-948021_{_oversized_response_id()}" key = get_s3_object_key( s3_path="input", prefix="team/", start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), s3_file_name=file_name, ) digest = hashlib.sha256(file_name.encode("utf-8")).hexdigest() assert key == f"input/team/2026-08-24/{file_name[:64]}_{digest}.json" def test_s3_object_key_keeps_the_configured_prefix_when_only_the_id_overflows(): """A 940 byte configured prefix survives whole when only the id overflows.""" from litellm.constants import MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import get_s3_object_key prefix = "team-" + "b" * 934 + "/" key = get_s3_object_key( s3_path="", prefix=prefix, start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", ) assert key.startswith(prefix + "2026-08-24/") assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES def test_s3_object_key_spends_the_whole_budget_when_the_prefix_must_be_trimmed(): """A trimmed prefix keeps every byte the budget allows, not whole segments.""" from litellm.constants import MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import get_s3_object_key s3_path = "p" * 400 + "/" + "q" * 600 key = get_s3_object_key( s3_path=s3_path, prefix="", start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), s3_file_name="time-06-18-41-948021_abc", ) assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES assert key.startswith("p" * 400 + "/" + "q" * 500) def test_s3_object_key_keeps_a_single_segment_path_as_far_as_it_fits(): """A path with no separator is kept as far as it fits, never dropped to the bucket root.""" from litellm.constants import MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import get_s3_object_key key = get_s3_object_key( s3_path="a" * 1050, prefix="", start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), s3_file_name="time-06-18-41-948021_chatcmpl-xyz", ) assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES assert key.startswith("a" * 900) def test_create_s3_batch_logging_element_bounds_key_and_keeps_full_response_id(): """The batch element bounds the key and keeps the full response id in the payload.""" from litellm.constants import MAX_S3_OBJECT_KEY_BYTES logger = S3Logger(s3_use_team_prefix=True, s3_use_key_prefix=True) response_id = _oversized_response_id() payload = StandardLoggingPayload( id=response_id, metadata={"user_api_key_team_alias": "DefaultTeamProd", "user_api_key_alias": "prod-key"}, messages=[], ) result = logger.create_s3_batch_logging_element(datetime(2026, 8, 24, 6, 18, 41, 948021), payload) assert result is not None assert len(result.s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES assert result.s3_object_key.startswith("DefaultTeamProd/prod-key/2026-08-24/") assert result.payload["id"] == response_id def test_s3_object_download_filename_is_bounded_for_oversized_response_id(): """The Content-Disposition filename is bounded too, or the PUT fails with MetadataTooLarge.""" from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES from litellm.integrations.s3 import get_s3_object_download_filename file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), _oversized_response_id()) assert len(file_name.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES assert file_name.startswith("time-2026-08-24T06-18-41-948021_resp_") assert file_name.endswith(".json") def test_s3_object_download_filenames_stay_distinct_when_shortened(): """Shortened filenames stay distinct.""" from litellm.integrations.s3 import get_s3_object_download_filename start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) file_names = { get_s3_object_download_filename(start_time, _oversized_response_id() + suffix) for suffix in ("first", "second", "third") } assert len(file_names) == 3 def test_s3_object_download_filename_short_id_is_unchanged(): """An ordinary response id keeps the filename it had before.""" from litellm.integrations.s3 import get_s3_object_download_filename file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), "resp_abc123") assert file_name == "time-2026-08-24T06-18-41-948021_resp_abc123.json" def test_create_s3_batch_logging_element_bounds_the_download_filename(): """The batch element carries a bounded Content-Disposition filename.""" from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES logger = S3Logger() payload = StandardLoggingPayload(id=_oversized_response_id(), metadata={}, messages=[]) result = logger.create_s3_batch_logging_element(datetime(2026, 8, 24, 6, 18, 41, 948021), payload) assert result is not None assert len(result.s3_object_download_filename.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES @pytest.mark.asyncio async def test_audit_log_object_key_is_bounded_for_a_long_configured_path(): """Audit log keys are bounded by the same builder.""" from litellm.constants import MAX_S3_OBJECT_KEY_BYTES logger = S3Logger() logger.s3_path = "audit-archive/" + "z" * 1100 await logger.async_log_audit_log_event({"id": "1a4f7bd0-6f1e-4d0a-9b3c-9f2e1d5a7c88"}) assert len(logger.log_queue) == 1 assert len(logger.log_queue[0].s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES assert logger.log_queue[0].s3_object_key.startswith("audit-archive/" + "z" * 900) def test_s3_object_download_filename_drops_characters_that_break_the_header(): """A quote or separator in the response id cannot escape the quoted header value.""" from litellm.integrations.s3 import get_s3_object_download_filename file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), 'resp_a"b/c') assert file_name == "time-2026-08-24T06-18-41-948021_resp_a_b_c.json" # -------------------------------------------------------------- # params_source / s3_callback_params_override (audit-log decoupling) # -------------------------------------------------------------- def test_s3_callback_params_override_uses_alternate_dict(monkeypatch): """`s3_callback_params_override` makes the logger read its config from the override dict instead of `litellm.s3_callback_params`.""" import litellm monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-bucket"}) logger = S3Logger( s3_callback_params_override={ "s3_bucket_name": "audit-bucket", "s3_path": "audit-prefix", "s3_region_name": "us-west-2", } ) assert logger.s3_bucket_name == "audit-bucket" assert logger.s3_path == "audit-prefix" assert logger.s3_region_name == "us-west-2" def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch): """Resolving `os.environ/X` markers must not mutate the override dict or `litellm.s3_callback_params`.""" import litellm monkeypatch.setenv("MY_AUDIT_BUCKET", "resolved-bucket") override = {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"} monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"}) logger = S3Logger(s3_callback_params_override=override) assert logger.s3_bucket_name == "resolved-bucket" assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" assert litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" def test_s3_callback_params_override_none_falls_back_to_global(monkeypatch): """No override → behaves exactly as today (reads `litellm.s3_callback_params`).""" import litellm monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "from-global"}) logger = S3Logger() assert logger.s3_bucket_name == "from-global" def test_s3_callback_params_override_empty_dict_is_opt_in(monkeypatch): """An empty override dict skips the global entirely (env/IAM-only config).""" import litellm monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "from-global"}) logger = S3Logger(s3_callback_params_override={}) assert logger.s3_bucket_name is None def _expected_content_md5(payload: dict) -> str: import base64 import hashlib from litellm.litellm_core_utils.safe_json_dumps import safe_dumps json_string = safe_dumps(payload) return base64.b64encode(hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest()).decode() def _require_non_security_md5(monkeypatch): import hashlib original_md5 = hashlib.md5 def fips_md5(data=b"", *, usedforsecurity=True): if usedforsecurity: raise ValueError("MD5 blocked for security use") return original_md5(data, usedforsecurity=usedforsecurity) monkeypatch.setattr(hashlib, "md5", fips_md5) @pytest.mark.asyncio async def test_async_upload_sets_content_md5_header(monkeypatch): """ Object Lock buckets reject PUTs without a Content-MD5 header (AWS spec). The async upload must send a base64 md5 of the exact signed body. """ from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) payload = {"test": "content-md5"} test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-md5.json", payload=payload, s3_object_download_filename="test-md5.json", ) _require_non_security_md5(monkeypatch) response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put.return_value = response await logger.async_upload_data_to_s3(test_element) headers = logger.async_httpx_client.put.call_args.kwargs["headers"] assert headers["Content-MD5"] == _expected_content_md5(payload) assert "x-amz-server-side-encryption" not in headers def test_sync_upload_sets_content_md5_header(monkeypatch): """The sync upload path must also send Content-MD5 for Object Lock buckets.""" from unittest.mock import MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) payload = {"test": "sync-content-md5"} test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-sync-md5.json", payload=payload, s3_object_download_filename="test-sync-md5.json", ) _require_non_security_md5(monkeypatch) response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() mock_sync_client = MagicMock() mock_sync_client.put.return_value = response with patch( "litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client, ): logger.upload_data_to_s3(test_element) headers = mock_sync_client.put.call_args.kwargs["headers"] assert headers["Content-MD5"] == _expected_content_md5(payload) assert "x-amz-server-side-encryption" not in headers @pytest.mark.asyncio async def test_async_upload_sets_server_side_encryption_header_when_configured(): """ When s3_server_side_encryption is set (e.g. buckets with a KMS default encryption policy), the PUT must carry x-amz-server-side-encryption. """ from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_server_side_encryption="aws:kms", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-sse.json", payload={"test": "sse"}, s3_object_download_filename="test-sse.json", ) response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put.return_value = response await logger.async_upload_data_to_s3(test_element) headers = logger.async_httpx_client.put.call_args.kwargs["headers"] assert headers["x-amz-server-side-encryption"] == "aws:kms" def test_s3_server_side_encryption_read_from_callback_params(monkeypatch): """s3_server_side_encryption can be configured via s3_callback_params.""" import litellm monkeypatch.setattr( litellm, "s3_callback_params", { "s3_bucket_name": "from-global", "s3_server_side_encryption": "aws:kms", }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" @pytest.mark.asyncio async def test_async_upload_sets_sse_kms_key_id_header_when_configured(): """ When s3_sse_kms_key_id is set alongside aws:kms, the PUT must carry x-amz-server-side-encryption-aws-kms-key-id so objects are encrypted with the customer-managed KMS key instead of the bucket default. """ from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_server_side_encryption="aws:kms", s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-sse-kms.json", payload={"test": "sse-kms"}, s3_object_download_filename="test-sse-kms.json", ) response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put.return_value = response await logger.async_upload_data_to_s3(test_element) headers = logger.async_httpx_client.put.call_args.kwargs["headers"] assert headers["x-amz-server-side-encryption"] == "aws:kms" assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == ( "arn:aws:kms:us-east-1:111122223333:key/test-key-id" ) def test_sync_upload_sets_sse_kms_key_id_header_when_configured(): """The sync upload path must carry the same SSE-KMS headers.""" from unittest.mock import MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_server_side_encryption="aws:kms", s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-sync-sse-kms.json", payload={"test": "sync-sse-kms"}, s3_object_download_filename="test-sync-sse-kms.json", ) response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() mock_sync_client = MagicMock() mock_sync_client.put.return_value = response with patch( "litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client, ): logger.upload_data_to_s3(test_element) headers = mock_sync_client.put.call_args.kwargs["headers"] assert headers["x-amz-server-side-encryption"] == "aws:kms" assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == ( "arn:aws:kms:us-east-1:111122223333:key/test-key-id" ) @pytest.mark.asyncio async def test_async_upload_omits_kms_key_id_header_when_not_configured(): """SSE without a key id must not emit the KMS key id header.""" from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_server_side_encryption="AES256", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-aes256.json", payload={"test": "aes256"}, s3_object_download_filename="test-aes256.json", ) response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put.return_value = response await logger.async_upload_data_to_s3(test_element) headers = logger.async_httpx_client.put.call_args.kwargs["headers"] assert headers["x-amz-server-side-encryption"] == "AES256" assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers def test_s3_sse_kms_key_id_read_from_callback_params(monkeypatch): """s3_sse_kms_key_id can be configured via s3_callback_params.""" import litellm monkeypatch.setattr( litellm, "s3_callback_params", { "s3_bucket_name": "from-global", "s3_server_side_encryption": "aws:kms", "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", }, ) logger = S3Logger() assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") @pytest.mark.asyncio async def test_async_upload_infers_aws_kms_when_only_key_id_set(): """ Setting only s3_sse_kms_key_id must not produce an invalid request (S3 rejects a key id without an algorithm); aws:kms is inferred. """ from unittest.mock import AsyncMock, MagicMock from litellm.types.integrations.s3_v2 import s3BatchLoggingElement logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-kms-only.json", payload={"test": "kms-only"}, s3_object_download_filename="test-kms-only.json", ) response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put.return_value = response await logger.async_upload_data_to_s3(test_element) headers = logger.async_httpx_client.put.call_args.kwargs["headers"] assert headers["x-amz-server-side-encryption"] == "aws:kms" assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == ( "arn:aws:kms:us-east-1:111122223333:key/test-key-id" ) def test_s3_sse_kms_key_id_read_from_audit_override_params(monkeypatch): """The audit-log override path must honor s3_sse_kms_key_id too.""" import litellm monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-logs-bucket"}) logger = S3Logger( s3_callback_params_override={ "s3_bucket_name": "audit-logs-bucket", "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/audit-key-id", } ) assert logger.s3_bucket_name == "audit-logs-bucket" assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/audit-key-id") def test_kms_key_id_dropped_when_algorithm_is_not_kms(monkeypatch): """ AES256 plus a KMS key id is an invalid S3 combination; the key id must be dropped at init so uploads keep working instead of silently 400ing. """ import litellm monkeypatch.setattr( litellm, "s3_callback_params", { "s3_bucket_name": "from-global", "s3_server_side_encryption": "AES256", "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "AES256" assert logger.s3_sse_kms_key_id is None def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(monkeypatch): """ A YAML boolean in s3_server_side_encryption must not crash logger init and must not discard the valid key id; aws:kms is inferred from the key id. """ import litellm monkeypatch.setattr( litellm, "s3_callback_params", { "s3_bucket_name": "from-global", "s3_server_side_encryption": True, "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(monkeypatch): """A mistyped key id (unquoted YAML number) must not disable the valid algorithm.""" import litellm monkeypatch.setattr( litellm, "s3_callback_params", { "s3_bucket_name": "from-global", "s3_server_side_encryption": "aws:kms", "s3_sse_kms_key_id": 12345, }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" assert logger.s3_sse_kms_key_id is None _ACCESS_KEY = "AKIAIOSFODNN7EXAMPLE" _SECRET_KEY = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" _KEY_WITH_SPACE = "LOGS/LLM AI Projects/2026-08-04/time-13-01-00-abc.json" def _signature_for(signer_cls, url: str, method: str, body: bytes | None, headers: dict[str, str]) -> str: from botocore.awsrequest import AWSRequest from botocore.credentials import Credentials sent = {name.lower(): value for name, value in headers.items()} signed_header_names = sent["authorization"].split("SignedHeaders=")[1].split(", ")[0].split(";") request = AWSRequest( method=method, url=url, data=body, headers={name: sent[name] for name in signed_header_names if name in sent}, ) request.context["timestamp"] = sent["x-amz-date"] signer = signer_cls(Credentials(_ACCESS_KEY, _SECRET_KEY), "s3", "us-east-1") canonical_request = signer.canonical_request(request) return signer.signature(signer.string_to_sign(request, canonical_request), request) def _as_s3_canonicalizes(url: str) -> str: """ The path S3 rebuilds from the wire path: percent-encode everything outside the unreserved set, without normalizing or double-encoding. `=` becomes `%3D`, `%20` stays `%20`. """ from urllib.parse import quote, unquote, urlsplit, urlunsplit split = urlsplit(url) return urlunsplit(split._replace(path=quote(unquote(split.path), safe="/~"))) def _assert_signed_for_s3_canonicalization(url: str, method: str, body: bytes | None, headers: dict[str, str]) -> None: """ S3 rebuilds the canonical request from the wire path with single percent-encoding, which botocore models as S3SigV4Auth; plain SigV4Auth double-encodes it (%2520 for a space) and S3 answers 403 SignatureDoesNotMatch. Assert we sent an already-encoded path and signed it the way S3 reads it. """ from botocore.auth import S3SigV4Auth, SigV4Auth assert url == _as_s3_canonicalizes(url) sent_signature = headers["Authorization"].split("Signature=")[1].strip() assert sent_signature == _signature_for(S3SigV4Auth, url, method, body, headers) assert sent_signature != _signature_for(SigV4Auth, url, method, body, headers) def _logger_for_signing() -> S3Logger: return S3Logger( s3_bucket_name="logs-bucket", s3_aws_access_key_id=_ACCESS_KEY, s3_aws_secret_access_key=_SECRET_KEY, s3_region_name="us-east-1", ) def _element_with_space(): from litellm.types.integrations.s3_v2 import s3BatchLoggingElement return s3BatchLoggingElement( s3_object_key=_KEY_WITH_SPACE, payload={"test": "sigv4"}, s3_object_download_filename="log.json", ) @pytest.mark.asyncio async def test_async_upload_signs_object_key_with_space_the_way_s3_does(): from unittest.mock import AsyncMock, MagicMock logger = _logger_for_signing() response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put.return_value = response await logger.async_upload_data_to_s3(_element_with_space()) call = logger.async_httpx_client.put.call_args _assert_signed_for_s3_canonicalization( url=call[0][0], method="PUT", body=call.kwargs["data"].encode("utf-8"), headers=call.kwargs["headers"], ) def test_sync_upload_signs_object_key_with_space_the_way_s3_does(): from unittest.mock import MagicMock logger = _logger_for_signing() response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() mock_sync_client = MagicMock() mock_sync_client.put.return_value = response with patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client): logger.upload_data_to_s3(_element_with_space()) call = mock_sync_client.put.call_args _assert_signed_for_s3_canonicalization( url=call[0][0], method="PUT", body=call.kwargs["data"].encode("utf-8"), headers=call.kwargs["headers"], ) @pytest.mark.asyncio async def test_download_signs_object_key_with_space_the_way_s3_does(): from unittest.mock import AsyncMock, MagicMock logger = _logger_for_signing() response = MagicMock() response.status_code = 200 response.json = MagicMock(return_value={"downloaded": "data"}) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.get.return_value = response assert await logger._download_object_from_s3(_KEY_WITH_SPACE) == {"downloaded": "data"} call = logger.async_httpx_client.get.call_args _assert_signed_for_s3_canonicalization( url=call[0][0], method="GET", body=None, headers=call.kwargs["headers"], ) _RESERVED_CHAR_KEYS = ( "2026-08-21/time-05-29-36_resp_bGl0ZWxsbTpjdXN0b20=.json", "session=logs/2026-08-21/time-05-29-36_abc.json", "a+b/2026-08-21/time-05-29-36_abc.json", "a&b/2026-08-21/time-05-29-36_abc.json", "a#b/2026-08-21/time-05-29-36_abc.json", "a?b/2026-08-21/time-05-29-36_abc.json", "a%b/2026-08-21/time-05-29-36_abc.json", _KEY_WITH_SPACE, ) def _element_for(s3_object_key: str): from litellm.types.integrations.s3_v2 import s3BatchLoggingElement return s3BatchLoggingElement( s3_object_key=s3_object_key, payload={"test": "sigv4"}, s3_object_download_filename="log.json", ) def _expected_wire_url(s3_object_key: str) -> str: """The URL boto3 itself would put on the wire for this key.""" from urllib.parse import quote return f"https://logs-bucket.s3.us-east-1.amazonaws.com/{quote(s3_object_key, safe='/')}" @pytest.mark.parametrize("s3_object_key", _RESERVED_CHAR_KEYS) @pytest.mark.asyncio async def test_async_upload_percent_encodes_reserved_characters_in_object_key(s3_object_key): from unittest.mock import AsyncMock, MagicMock logger = _logger_for_signing() response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put.return_value = response await logger.async_upload_data_to_s3(_element_for(s3_object_key)) call = logger.async_httpx_client.put.call_args assert call[0][0] == _expected_wire_url(s3_object_key) _assert_signed_for_s3_canonicalization( url=call[0][0], method="PUT", body=call.kwargs["data"].encode("utf-8"), headers=call.kwargs["headers"], ) @pytest.mark.parametrize("s3_object_key", _RESERVED_CHAR_KEYS) def test_sync_upload_percent_encodes_reserved_characters_in_object_key(s3_object_key): from unittest.mock import MagicMock logger = _logger_for_signing() response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() mock_sync_client = MagicMock() mock_sync_client.put.return_value = response with patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client): logger.upload_data_to_s3(_element_for(s3_object_key)) call = mock_sync_client.put.call_args assert call[0][0] == _expected_wire_url(s3_object_key) _assert_signed_for_s3_canonicalization( url=call[0][0], method="PUT", body=call.kwargs["data"].encode("utf-8"), headers=call.kwargs["headers"], ) @pytest.mark.parametrize("s3_object_key", _RESERVED_CHAR_KEYS) @pytest.mark.asyncio async def test_download_percent_encodes_reserved_characters_in_object_key(s3_object_key): from unittest.mock import AsyncMock, MagicMock logger = _logger_for_signing() response = MagicMock() response.status_code = 200 response.json = MagicMock(return_value={"downloaded": "data"}) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.get.return_value = response assert await logger._download_object_from_s3(s3_object_key) == {"downloaded": "data"} call = logger.async_httpx_client.get.call_args assert call[0][0] == _expected_wire_url(s3_object_key) _assert_signed_for_s3_canonicalization( url=call[0][0], method="GET", body=None, headers=call.kwargs["headers"], ) def _s3_logger_for_region(region_name: str) -> S3Logger: logger = S3Logger.__new__(S3Logger) logger.s3_endpoint_url = None logger.s3_bucket_name = "my-litellm-audit" logger.s3_region_name = region_name return logger @pytest.mark.parametrize( "region_name,expected_url", [ ( "cn-northwest-1", "https://my-litellm-audit.s3.cn-northwest-1.amazonaws.com.cn/2025-01-01/key.json", ), ( "us-gov-west-1", "https://my-litellm-audit.s3.us-gov-west-1.amazonaws.com/2025-01-01/key.json", ), ( "us-east-1", "https://my-litellm-audit.s3.us-east-1.amazonaws.com/2025-01-01/key.json", ), ], ) def test_build_object_url_uses_partition_dns_suffix(region_name: str, expected_url: str) -> None: assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url def _prompts_only_logger(s3_log_prompts_only: bool | None = None) -> S3Logger: return S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_log_prompts_only=s3_log_prompts_only, ) def _chat_payload() -> StandardLoggingPayload: return StandardLoggingPayload( id="chatcmpl-prompts-only", messages=[{"role": "user", "content": "Reply with exactly the word PINEAPPLE."}], response={"choices": [{"message": {"role": "assistant", "content": "PINEAPPLE"}}]}, metadata={"user_api_key_team_alias": None}, ) async def _queued_body_via_async_upload( logger: S3Logger, log_event: Callable[..., Awaitable[None]] ) -> dict[str, object]: payload = _chat_payload() original = copy.deepcopy(payload) await log_event( kwargs={"standard_logging_object": payload}, response_obj=None, start_time=datetime(2026, 7, 30, 12, 0, 0), end_time=datetime(2026, 7, 30, 12, 0, 1), ) assert payload == original, "the caller's standard_logging_object must not be mutated" (element,) = logger.log_queue response = MagicMock() response.status_code = 200 response.raise_for_status = MagicMock() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put.return_value = response await logger.async_upload_data_to_s3(element) return json.loads(logger.async_httpx_client.put.call_args.kwargs["data"]) @pytest.mark.asyncio @pytest.mark.parametrize("event_name", ["async_log_success_event", "async_log_failure_event"]) async def test_prompts_only_drops_response_but_keeps_messages_in_uploaded_object( monkeypatch: pytest.MonkeyPatch, event_name: str ): import litellm monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": True}) logger = _prompts_only_logger() log_event: Callable[..., Awaitable[None]] = ( logger.async_log_success_event if event_name == "async_log_success_event" else logger.async_log_failure_event ) body = await _queued_body_via_async_upload(logger, log_event) assert body["messages"] == _chat_payload()["messages"] assert body["response"] is None assert body["id"] == "chatcmpl-prompts-only" @pytest.mark.asyncio async def test_prompts_only_default_off_keeps_response_in_uploaded_object(monkeypatch: pytest.MonkeyPatch): import litellm monkeypatch.setattr(litellm, "s3_callback_params", {}) monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) logger = _prompts_only_logger() body = await _queued_body_via_async_upload(logger, logger.async_log_success_event) assert body["response"] == _chat_payload()["response"] assert body["messages"] == _chat_payload()["messages"] @pytest.mark.asyncio async def test_prompts_only_explicit_false_in_params_beats_env_var(monkeypatch: pytest.MonkeyPatch): import litellm monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": False}) monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true") logger = _prompts_only_logger() body = await _queued_body_via_async_upload(logger, logger.async_log_success_event) assert body["response"] == _chat_payload()["response"] @pytest.mark.asyncio async def test_prompts_only_env_var_applies_when_param_unset(monkeypatch: pytest.MonkeyPatch): import litellm monkeypatch.setattr(litellm, "s3_callback_params", {}) logger = _prompts_only_logger() monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true") body = await _queued_body_via_async_upload(logger, logger.async_log_success_event) assert body["response"] is None assert body["messages"] == _chat_payload()["messages"] @respx.mock def test_prompts_only_constructor_kwarg_applies_to_sync_upload(monkeypatch: pytest.MonkeyPatch): import litellm monkeypatch.setattr(litellm, "s3_callback_params", {}) monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) logger = _prompts_only_logger(s3_log_prompts_only=True) payload = _chat_payload() element = logger.create_s3_batch_logging_element( start_time=datetime(2026, 7, 30, 12, 0, 0), standard_logging_payload=payload, ) assert element is not None assert payload["response"] == _chat_payload()["response"] put_route = respx.put(url__regex=r"https://test-bucket\.s3\..*").mock(return_value=httpx.Response(200)) logger.upload_data_to_s3(element) body = json.loads(put_route.calls.last.request.content) assert body["response"] is None assert body["messages"] == _chat_payload()["messages"] @pytest.mark.parametrize("callback_name", ["s3", "s3_v2"]) def test_prompts_only_toggle_is_exposed_to_admin_ui_for_both_s3_callbacks(callback_name: str): from litellm.integrations.custom_logger import CustomLogger assert "S3_LOG_PROMPTS_ONLY" in CustomLogger.get_callback_env_vars(callback_name) _PARTITION_START: Final = datetime(2026, 9, 29, 14, 5, 9, 123456) _PARTITION_ID: Final = "chatcmpl-partition" def _partition_payload(response_id: str = _PARTITION_ID) -> StandardLoggingPayload: return StandardLoggingPayload( id=response_id, metadata={"user_api_key_team_alias": "team-a", "user_api_key_alias": "key-a"}, messages=[], ) def _partition_logger( monkeypatch: pytest.MonkeyPatch, callback_params: dict[str, object], **kwargs: object ) -> S3Logger: import litellm monkeypatch.setattr( litellm, "s3_callback_params", {"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1", "s3_path": "logs", **callback_params}, ) return S3Logger( s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_use_team_prefix=True, s3_use_key_prefix=True, **kwargs, ) _DAILY_KEY: Final = f"logs/team-a/key-a/2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json" _HOURLY_KEY: Final = f"logs/team-a/key-a/2026-09-29/14/time-14-05-09-123456_{_PARTITION_ID}.json" @pytest.mark.parametrize( ("callback_params", "expected_key"), [ ({}, _DAILY_KEY), ({"s3_partition_granularity": None}, _DAILY_KEY), ({"s3_partition_granularity": "day"}, _DAILY_KEY), ({"s3_partition_granularity": "hour"}, _HOURLY_KEY), ], ) def test_partition_granularity_sets_request_log_folder( monkeypatch: pytest.MonkeyPatch, callback_params: dict[str, object], expected_key: str ) -> None: monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False) logger = _partition_logger(monkeypatch, callback_params) element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload()) assert element is not None assert element.s3_object_key == expected_key @pytest.mark.parametrize("invalid", ["hourly", "HOUR", "1", 1, True]) def test_invalid_partition_granularity_warns_and_keeps_daily_folder( monkeypatch: pytest.MonkeyPatch, invalid: object ) -> None: monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False) with patch("litellm.integrations.s3.verbose_logger") as mock_logger: logger = _partition_logger(monkeypatch, {"s3_partition_granularity": invalid}) element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload()) second = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload()) assert element is not None assert second is not None assert element.s3_object_key == second.s3_object_key == _DAILY_KEY mock_logger.warning.assert_called_once() assert mock_logger.warning.call_args.args[1:] == (invalid,) def test_partition_granularity_reads_admin_ui_env_var_below_callback_params(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("S3_PARTITION_GRANULARITY", "hour") from_env = _partition_logger(monkeypatch, {}).create_s3_batch_logging_element( _PARTITION_START, _partition_payload() ) from_params = _partition_logger(monkeypatch, {"s3_partition_granularity": "day"}).create_s3_batch_logging_element( _PARTITION_START, _partition_payload() ) assert from_env is not None and from_env.s3_object_key == _HOURLY_KEY assert from_params is not None and from_params.s3_object_key == _DAILY_KEY def test_partition_granularity_constructor_argument_and_os_environ_reference(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False) monkeypatch.setenv("MY_S3_PARTITION", "hour") from_ctor = _partition_logger(monkeypatch, {}, s3_partition_granularity="hour") from_secret = _partition_logger(monkeypatch, {"s3_partition_granularity": "os.environ/MY_S3_PARTITION"}) for logger in (from_ctor, from_secret): element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload()) assert element is not None and element.s3_object_key == _HOURLY_KEY def test_hourly_partition_long_key_keeps_hour_folder_within_s3_limit(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.constants import MAX_S3_OBJECT_KEY_BYTES monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False) logger = _partition_logger(monkeypatch, {"s3_partition_granularity": "hour", "s3_path": "p" * 1100}) element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload("r" * 600)) assert element is not None assert len(element.s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES assert re.search(r"/2026-09-29/14/[0-9a-f]{64}\.json$", element.s3_object_key) @pytest.mark.asyncio @pytest.mark.parametrize(("granularity", "hour_folder"), [("hour", True), ("day", False), (None, False)]) async def test_audit_log_key_follows_audit_callback_params_partition_granularity( monkeypatch: pytest.MonkeyPatch, granularity: str | None, hour_folder: bool ) -> None: monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False) logger = S3Logger( s3_callback_params_override={ "s3_bucket_name": "audit-bucket", "s3_path": "audit", "s3_partition_granularity": granularity, } ) await logger.async_log_audit_log_event({"id": "audit-1"}) (element,) = logger.log_queue match = re.fullmatch( r"audit/audit_logs/\d{4}-\d{2}-\d{2}/(?:(\d{2})/)?(\d{2})-\d{2}-\d{2}_audit-1\.json", element.s3_object_key ) assert match is not None, element.s3_object_key assert (match.group(1) is not None) is hour_folder if hour_folder: assert match.group(1) == match.group(2) @pytest.mark.asyncio async def test_hourly_batch_file_upload_writes_one_file_per_hour_folder(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False) logger = _partition_logger(monkeypatch, {"s3_partition_granularity": "hour"}, s3_batch_file_upload=True) put = _RecordingPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put before = logger.create_s3_batch_logging_element(datetime(2026, 9, 29, 13, 59, 59), _partition_payload("before")) after = logger.create_s3_batch_logging_element(datetime(2026, 9, 29, 14, 0, 1), _partition_payload("after")) assert before is not None and after is not None logger.log_queue = [before, after] await logger.async_send_batch() by_folder = { re.sub(r"/batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", "", url.split(".com/", 1)[-1]): data for url, data, _headers in put.calls } assert sorted(by_folder) == ["logs/team-a/key-a/2026-09-29/13", "logs/team-a/key-a/2026-09-29/14"] assert [json.loads(line)["id"] for line in (by_folder["logs/team-a/key-a/2026-09-29/13"] or "").splitlines()] == [ "before" ] assert [json.loads(line)["id"] for line in (by_folder["logs/team-a/key-a/2026-09-29/14"] or "").splitlines()] == [ "after" ] @pytest.mark.parametrize("granularity", [None, "day", "hour"]) def test_cold_storage_object_key_matches_the_uploaded_request_log_key( monkeypatch: pytest.MonkeyPatch, granularity: str | None ) -> None: import litellm from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False) monkeypatch.setattr( litellm, "s3_callback_params", {"s3_bucket_name": "test-bucket", "s3_path": "coldlogs", "s3_partition_granularity": granularity}, ) monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2") logger = S3Logger() uploaded = logger.create_s3_batch_logging_element( _PARTITION_START, StandardLoggingPayload(id=_PARTITION_ID, metadata={}, messages=[]) ) monkeypatch.setattr(litellm, "callbacks", [logger]) cold_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key( start_time=_PARTITION_START, response_id=_PARTITION_ID ) assert uploaded is not None assert cold_key == uploaded.s3_object_key assert ("/2026-09-29/14/" in cold_key) is (granularity == "hour") def test_cold_storage_key_matches_upload_when_env_var_changes_mid_request(monkeypatch: pytest.MonkeyPatch) -> None: import litellm from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False) monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "test-bucket", "s3_path": "coldlogs"}) monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2") logger = S3Logger() monkeypatch.setattr(litellm, "callbacks", [logger]) cold_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key( start_time=_PARTITION_START, response_id=_PARTITION_ID ) monkeypatch.setenv("S3_PARTITION_GRANULARITY", "hour") uploaded = logger.create_s3_batch_logging_element( _PARTITION_START, StandardLoggingPayload(id=_PARTITION_ID, metadata={"cold_storage_object_key": cold_key}, messages=[]), ) assert uploaded is not None assert cold_key == uploaded.s3_object_key == f"coldlogs/2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json" def test_hour_upload_ignores_a_cold_storage_key_owned_by_another_logger(monkeypatch: pytest.MonkeyPatch) -> None: import litellm from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False) monkeypatch.setattr( litellm, "s3_callback_params", {"s3_bucket_name": "test-bucket", "s3_partition_granularity": "hour"} ) monkeypatch.setattr(litellm, "cold_storage_custom_logger", "gcs_bucket") logger = S3Logger() cold_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key( start_time=_PARTITION_START, response_id=_PARTITION_ID ) uploaded = logger.create_s3_batch_logging_element( _PARTITION_START, StandardLoggingPayload(id=_PARTITION_ID, metadata={"cold_storage_object_key": cold_key}, messages=[]), ) assert cold_key == f"2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json" assert uploaded is not None assert uploaded.s3_object_key == f"2026-09-29/14/time-14-05-09-123456_{_PARTITION_ID}.json" @pytest.mark.parametrize("callback_name", ["s3", "s3_v2"]) def test_partition_granularity_is_exposed_to_admin_ui(callback_name: str) -> None: from litellm.integrations.custom_logger import CustomLogger assert "S3_PARTITION_GRANULARITY" in CustomLogger.get_callback_env_vars(callback_name) def _element(payload: dict[str, object], key_suffix: str) -> s3BatchLoggingElement: return s3BatchLoggingElement( s3_object_key=f"2025-09-14/test-{key_suffix}.json", payload=payload, s3_object_download_filename=f"test-{key_suffix}.json", ) def _ok_response() -> MagicMock: response = MagicMock() response.status_code = 200 response.text = "" response.raise_for_status = MagicMock() return response class _CountingPut: def __init__(self, width: int) -> None: self.width = width self.in_flight = 0 self.peak = 0 self.calls = 0 async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.in_flight += 1 self.peak = max(self.peak, self.in_flight) self.calls += 1 for _ in range(50): if self.in_flight >= self.width: break await _real_sleep(0) self.in_flight -= 1 return _ok_response() class _RecordingPut: def __init__(self) -> None: self.calls: tuple[tuple[str, str | None, dict[str, str] | None], ...] = () async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.calls = (*self.calls, (url, data, headers)) return _ok_response() class _LateAppendingPut: def __init__(self, logger: S3Logger, element: s3BatchLoggingElement, fail_first: bool = False) -> None: self.logger = logger self.element = element self.fail_first = fail_first self.appended = False self.failed_key: str | None = None async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: if not self.appended: self.appended = True self.logger.log_queue.append(self.element) if self.fail_first: self.failed_key = url if url == self.failed_key: return _transient_failure_response() return _ok_response() class _AppendingFailingPut: def __init__(self, logger: S3Logger, elements: tuple[s3BatchLoggingElement, ...]) -> None: self.logger = logger self.elements = elements self.appended = False async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: if not self.appended: self.appended = True for element in self.elements: self.logger.log_queue.append(element) return _transient_failure_response() class _FailOnSuffixPut: def __init__(self, suffixes: tuple[str, ...]) -> None: self.failing = True self.suffixes = suffixes self.calls: tuple[str, ...] = () async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.calls = (*self.calls, url) if self.failing and url.endswith(self.suffixes): return _transient_failure_response() return _ok_response() class _FailUntilClearedPut: def __init__(self, status: int = 503, code: str | None = "SlowDown", raw_body: str | None = None) -> None: self.failing = True self.response: Final = _coded_failure_response(status, code, raw_body) self.calls: tuple[tuple[str, str | None], ...] = () async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.calls = (*self.calls, (url, data)) if self.failing: return self.response return _ok_response() @pytest.mark.asyncio async def test_async_send_batch_bounds_concurrent_uploads() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_concurrent_uploads=4, ) put = _CountingPut(logger.s3_max_concurrent_uploads) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": i}, f"{i}") for i in range(40)] await logger.async_send_batch() assert put.peak == 4 assert put.calls == 40 @pytest.mark.asyncio async def test_async_send_batch_uploads_single_jsonl_file() -> None: import json logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_batch_file_upload=True, ) put = _RecordingPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put payloads = [{"id": "req-1"}, {"id": "req-2"}, {"id": "req-3"}] logger.log_queue = [_element(payload, f"{i}") for i, payload in enumerate(payloads)] await logger.async_send_batch() assert len(put.calls) == 1 url, data, headers = put.calls[0] assert url.endswith(".jsonl") assert data is not None assert headers is not None assert [json.loads(line) for line in data.splitlines()] == payloads assert headers["Content-Type"] == "application/x-ndjson" @pytest.mark.asyncio async def test_flush_queue_preserves_events_added_during_upload() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) late_element = _element({"id": "late"}, "late") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _LateAppendingPut(logger, late_element) logger.log_queue = [_element({"id": "first"}, "first")] await logger.flush_queue() assert logger.log_queue == [late_element] def _override_logger(**overrides: object) -> S3Logger: return S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_callback_params_override=overrides, ) def test_env_backed_false_string_keeps_per_request_uploads() -> None: assert _override_logger(s3_batch_file_upload="false").s3_batch_file_upload is False assert _override_logger(s3_batch_file_upload="true").s3_batch_file_upload is True logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_batch_file_upload=True, s3_callback_params_override={"s3_batch_file_upload": "false"}, ) assert logger.s3_batch_file_upload is True @pytest.mark.parametrize("bad", [0, -3, "0", "abc", ""]) def test_invalid_concurrency_falls_back_to_default(bad: object) -> None: from litellm.constants import DEFAULT_S3_MAX_CONCURRENT_UPLOADS logger = _override_logger(s3_max_concurrent_uploads=bad) assert logger.s3_max_concurrent_uploads == DEFAULT_S3_MAX_CONCURRENT_UPLOADS assert logger._upload_limiter._value == DEFAULT_S3_MAX_CONCURRENT_UPLOADS def test_env_backed_concurrency_string_is_parsed() -> None: logger = _override_logger(s3_max_concurrent_uploads="4") assert logger.s3_max_concurrent_uploads == 4 assert logger._upload_limiter._value == 4 @pytest.mark.parametrize("empty", [None, ""]) def test_empty_config_concurrency_falls_back_to_constructor_value(empty: object) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_concurrent_uploads=4, s3_callback_params_override={"s3_max_concurrent_uploads": empty}, ) assert logger.s3_max_concurrent_uploads == 4 assert logger._upload_limiter._value == 4 def _coded_failure_response(status: int, code: str | None, raw_body: str | None = None) -> MagicMock: body: Final = ( raw_body if raw_body is not None else (f"{code}" if code is not None else "") ) response = MagicMock() response.status_code = status response.text = body response.raise_for_status = MagicMock( side_effect=httpx.HTTPStatusError(str(status), request=MagicMock(), response=response) ) return response def _transient_failure_response(status: int = 503) -> MagicMock: return _coded_failure_response(status, "SlowDown") def _terminal_failure_response() -> MagicMock: return _coded_failure_response(400, "EntityTooLarge") @pytest.mark.asyncio async def test_failed_uploads_stay_queued_for_next_flush() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) elements = [_element({"i": i}, f"{i}") for i in range(5)] put = _FailOnSuffixPut(("test-2.json", "test-4.json")) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = list(elements) with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert [element.s3_object_key for element in logger.log_queue] == [ elements[2].s3_object_key, elements[4].s3_object_key, ] put.failing = False await logger.flush_queue() assert logger.log_queue == [] @pytest.mark.asyncio async def test_batch_file_upload_failure_keeps_whole_batch() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_batch_file_upload=True, ) put = _FailUntilClearedPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put elements = [_element({"i": i}, f"{i}") for i in range(3)] logger.log_queue = list(elements) with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert len(put.calls) == 3 assert len(logger.log_queue) == 1 assert logger.log_queue[0].body == "\n".join(json.dumps(element.payload) for element in elements) @pytest.mark.asyncio async def test_events_appended_during_failed_flush_survive() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) late = _element({"id": "late"}, "late") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _LateAppendingPut(logger, late, fail_first=True) first = _element({"id": "first"}, "first") logger.log_queue = [first] with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert [element.s3_object_key for element in logger.log_queue] == [first.s3_object_key, late.s3_object_key] assert logger.log_queue[0].retrying_since is None logger.async_httpx_client.put.failed_key = None await logger.flush_queue() assert logger.log_queue == [] @pytest.mark.asyncio async def test_batch_file_key_shape() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_path="logs", s3_batch_file_upload=True, ) put = _RecordingPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"id": "req-1"}, "0")] await logger.async_send_batch() ((url, _data, headers),) = put.calls assert headers is not None assert re.search(r".*/2025-09-14/batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", url) assert headers["Content-Disposition"].endswith('.jsonl"') @pytest.mark.asyncio async def test_batch_file_groups_raw_elements_by_key_parent() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_batch_file_upload=True, ) put = _RecordingPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put alpha = s3BatchLoggingElement( s3_object_key="logs/alpha/2026-01-01/a.json", payload={"id": "a"}, s3_object_download_filename="a.json" ) beta = s3BatchLoggingElement( s3_object_key="logs/beta/2026-01-01/b.json", payload={"id": "b"}, s3_object_download_filename="b.json" ) plain = s3BatchLoggingElement( s3_object_key="logs/2026-01-01/c.json", payload={"id": "c"}, s3_object_download_filename="c.json" ) root = s3BatchLoggingElement( s3_object_key="solo.json", payload={"id": "d"}, s3_object_download_filename="solo.json" ) logger.log_queue = [alpha, beta, plain, root] await logger.async_send_batch() assert len(put.calls) == 4 by_parent = { re.sub(r"(^|/)batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", "", url.split(".com/", 1)[-1]): (url, data) for url, data, _headers in put.calls } assert sorted(by_parent) == ["", "logs/2026-01-01", "logs/alpha/2026-01-01", "logs/beta/2026-01-01"] assert [line for line in by_parent[""][1].splitlines()] == [json.dumps({"id": "d"})] assert [line for line in by_parent["logs/alpha/2026-01-01"][1].splitlines()] == [json.dumps({"id": "a"})] assert [line for line in by_parent["logs/beta/2026-01-01"][1].splitlines()] == [json.dumps({"id": "b"})] assert [line for line in by_parent["logs/2026-01-01"][1].splitlines()] == [json.dumps({"id": "c"})] @pytest.mark.asyncio async def test_failed_batch_file_is_requeued_and_resent_unchanged() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_batch_file_upload=True, ) put = _FailUntilClearedPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": i}, f"{i}") for i in range(3)] with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert len(logger.log_queue) == 1 assert logger.log_queue[0].body is not None assert logger.log_queue[0].s3_object_key.endswith(".jsonl") put.failing = False await logger.flush_queue() assert logger.log_queue == [] assert len(put.calls) == 4 assert put.calls[0] == put.calls[1] @pytest.mark.asyncio async def test_elements_appended_after_failed_batch_file_get_their_own_file() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_batch_file_upload=True, ) put = _FailUntilClearedPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"id": "first"}, "first")] with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() late = _element({"id": "late"}, "late") logger.log_queue.append(late) put.failing = False await logger.flush_queue() assert logger.log_queue == [] assert len(put.calls) == 5 assert put.calls[0] == put.calls[1] second_flush: Final = put.calls[3:] assert put.calls[0] in second_flush late_call: Final = next(call for call in second_flush if call != put.calls[0]) assert late_call[0] != put.calls[0][0] assert late_call[1] == json.dumps({"id": "late"}) @pytest.mark.asyncio async def test_batch_file_mode_disabled_when_s3_v2_is_cold_storage_logger(monkeypatch: pytest.MonkeyPatch) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_batch_file_upload=True, ) put = _RecordingPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put import litellm monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2") logger.log_queue = [_element({"id": "req-1"}, "0")] await logger.async_send_batch() assert len(put.calls) == 1 assert put.calls[0][0].endswith("test-0.json") monkeypatch.setattr(litellm, "cold_storage_custom_logger", None) logger.log_queue = [_element({"id": "req-2"}, "1")] await logger.async_send_batch() assert len(put.calls) == 2 assert put.calls[1][0].endswith(".jsonl") class _FailOnSuffixCodedPut: def __init__( self, suffixes: tuple[str, ...], status: int, code: str | None = None, raw_body: str | None = None ) -> None: self.suffixes = suffixes self.response: Final = _coded_failure_response(status, code, raw_body) self.calls: tuple[str, ...] = () async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.calls = (*self.calls, url) if url.endswith(self.suffixes): return self.response return _ok_response() @pytest.mark.asyncio @pytest.mark.parametrize( ("status", "code", "raw_body", "puts_per_element"), [ pytest.param(403, "AccessDenied", None, 3, id="access-denied-403"), pytest.param(403, None, None, 3, id="empty-403"), pytest.param(403, None, "Forbidden", 3, id="html-403"), pytest.param(400, "KMS.DisabledException", None, 1, id="kms-disabled-400"), pytest.param(404, "NoSuchBucket", None, 1, id="no-such-bucket-404"), ], ) async def test_non_terminal_failure_is_requeued_and_delivered_on_recovery( status: int, code: str | None, raw_body: str | None, puts_per_element: int ) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) elements = [_element({"i": i}, f"{i}") for i in range(5)] put = _FailUntilClearedPut(status=status, code=code, raw_body=raw_body) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = list(elements) with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert len(logger.log_queue) == 5 assert len(put.calls) == 5 * puts_per_element put.failing = False await logger.flush_queue() assert logger.log_queue == [] assert len(put.calls) == 5 * puts_per_element + 5 landed: Final = frozenset( element.s3_object_key for element in elements if any(call[0].endswith(element.s3_object_key) for call in put.calls[-5:]) ) assert landed == frozenset(element.s3_object_key for element in elements) @pytest.mark.asyncio async def test_persistent_500_stays_queued_through_a_dozen_failed_flushes() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) put = _FailUntilClearedPut(status=500, code="InternalError") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": i}, f"{i}") for i in range(5)] with patch("asyncio.sleep", new_callable=AsyncMock): for _ in range(12): await logger.flush_queue() assert len(logger.log_queue) == 5 assert len(put.calls) == 12 * 15 put.failing = False await logger.flush_queue() assert logger.log_queue == [] assert len(put.calls) == 12 * 15 + 5 @pytest.mark.asyncio async def test_terminal_object_is_dropped_once_next_to_delivered_siblings_when_opted_in() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_drop_on_terminal_error=True, ) elements = [_element({"i": i}, f"{i}") for i in range(5)] put = _FailOnSuffixCodedPut(("test-1.json",), 400, "EntityTooLarge") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = list(elements) await logger.flush_queue() assert len(put.calls) == 5 assert logger.log_queue == [] @pytest.mark.asyncio async def test_terminal_object_is_requeued_when_opted_out() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_drop_on_terminal_error=False, ) put = _FailOnSuffixCodedPut(("test-1.json",), 400, "EntityTooLarge") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": i}, f"{i}") for i in range(5)] await logger.flush_queue() assert len(logger.log_queue) == 1 assert logger.log_queue[0].s3_object_key.endswith("test-1.json") assert sum(call.endswith("test-1.json") for call in put.calls) == 1 assert len(put.calls) == 5 @pytest.mark.asyncio async def test_terminal_objects_are_requeued_when_every_upload_in_the_flush_fails() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_drop_on_terminal_error=True, ) put = _FailUntilClearedPut(status=400, code="EntityTooLarge") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": i}, f"{i}") for i in range(5)] await logger.flush_queue() assert len(logger.log_queue) == 5 @pytest.mark.asyncio async def test_retrying_past_the_opted_in_budget_is_dropped_only_next_to_delivered_siblings() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_retry_age_seconds=60, ) aged = _element({"id": "aged"}, "aged").model_copy(update={"retrying_since": _NOW - 120}) fresh = _element({"id": "fresh"}, "fresh") put = _FailOnSuffixCodedPut(("test-aged.json",), 503, "SlowDown") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [aged, fresh] with ( patch("asyncio.sleep", new_callable=AsyncMock), patch("time.monotonic", return_value=_NOW), ): await logger.flush_queue() assert logger.log_queue == [] assert len(put.calls) == 4 @pytest.mark.asyncio async def test_retrying_past_the_budget_stays_queued_when_the_whole_flush_fails() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_retry_age_seconds=60, ) aged = _element({"id": "aged"}, "aged").model_copy(update={"retrying_since": _NOW - 120}) put = _FailUntilClearedPut(status=503, code="SlowDown") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [aged] with ( patch("asyncio.sleep", new_callable=AsyncMock), patch("time.monotonic", return_value=_NOW), ): await logger.flush_queue() assert len(logger.log_queue) == 1 @pytest.mark.asyncio async def test_overflow_after_a_failed_flush_trims_failed_first_and_counts_upload_failures_only( caplog, ) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_queue_size=4, ) late = tuple(_element({"id": f"late-{index}"}, f"late-{index}") for index in range(3)) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _AppendingFailingPut(logger, late) logger.log_queue = [_element({"id": "first"}, "first"), _element({"id": "second"}, "second")] with ( patch.object(logger, "handle_callback_failure") as mock_failure, patch("asyncio.sleep", new_callable=AsyncMock), pytest.raises(S3BatchUploadError), ): await logger.async_send_batch() assert [element.payload["id"] for element in logger.log_queue] == ["second", "late-0", "late-1", "late-2"] failed_uploads: Final = 2 assert mock_failure.call_count == failed_uploads mock_failure.assert_called_with(callback_name="S3Logger") assert "dropped 1 oldest events" in caplog.text @pytest.mark.asyncio async def test_default_logger_ages_out_elements_retrying_longer_than_an_hour(caplog) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) elements = [ _element({"i": index}, f"{index}").model_copy(update={"retrying_since": _NOW - 7200}) for index in range(3) ] put = _FailOnSuffixPut(("test-1.json", "test-2.json")) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = list(elements) with ( patch("asyncio.sleep", new_callable=AsyncMock), patch("time.monotonic", return_value=_NOW), ): await logger.flush_queue() assert logger.log_queue == [] assert "uploads dropped" in caplog.text @pytest.mark.asyncio async def test_opted_out_logger_never_ages_out_long_retrying_elements(caplog) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_retry_age_seconds=0, ) elements = [ _element({"i": index}, f"{index}").model_copy(update={"retrying_since": _NOW - 7200}) for index in range(3) ] put = _FailOnSuffixPut(("test-1.json", "test-2.json")) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = list(elements) with ( patch("asyncio.sleep", new_callable=AsyncMock), patch("time.monotonic", return_value=_NOW), ): await logger.flush_queue() assert [element.s3_object_key for element in logger.log_queue] == [ elements[1].s3_object_key, elements[2].s3_object_key, ] assert "uploads dropped" not in caplog.text put.failing = False await logger.flush_queue() assert logger.log_queue == [] landed: Final = frozenset(call_url.rsplit("/", 1)[-1] for call_url in put.calls) assert landed == frozenset(f"test-{index}.json" for index in range(3)) @pytest.mark.asyncio async def test_queue_grows_past_the_cap_while_the_sink_fails_and_everything_lands() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_queue_size=5, ) elements = [_element({"i": index}, f"{index}") for index in range(8)] put = _FailUntilClearedPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = list(elements[:5]) with ( patch.object(logger, "handle_callback_failure") as mock_failure, patch("asyncio.sleep", new_callable=AsyncMock), ): await logger.flush_queue() assert len(logger.log_queue) == 5 upload_failures: Final = 5 assert mock_failure.call_count == upload_failures for element in elements[5:]: logger.log_queue.append(element) put.failing = False await logger.flush_queue() assert logger.log_queue == [] landed: Final = frozenset(call[0].rsplit("/", 1)[-1] for call in put.calls[-8:]) assert landed == frozenset(f"test-{index}.json" for index in range(8)) # calls are (url, data) pairs @pytest.mark.asyncio @pytest.mark.parametrize( ("status", "code"), [ pytest.param(404, "NoSuchKey", id="404"), pytest.param(401, None, id="401"), pytest.param(400, None, id="uncoded-400"), ], ) async def test_unlisted_status_gets_one_put_and_stays_queued(status: int, code: str | None) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) put = _FailUntilClearedPut(status=status, code=code) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": index}, f"{index}") for index in range(4)] with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: await logger.flush_queue() assert len(put.calls) == 4 mock_sleep.assert_not_awaited() assert len(logger.log_queue) == 4 class _SyncRecordingClient: def __init__(self, response: httpx.Response) -> None: self.response: Final = response self.put_calls: list = [] # mutable-ok: call log appended once per PUT def put( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.put_calls.append(url) return self.response def test_sync_upload_404_is_single_attempt_without_sleep() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) sync_client: Final = _SyncRecordingClient(_coded_failure_response(404, "NoSuchKey")) with ( patch("litellm.integrations.s3_v2.get_httpx_client", return_value=sync_client), patch("time.sleep") as mock_sleep, ): logger.upload_data_to_s3(_element({"id": "sync-404"}, "sync-404")) assert len(sync_client.put_calls) == 1 mock_sleep.assert_not_called() @pytest.mark.parametrize( ("status", "expected_puts", "expected_retries"), [ pytest.param(408, 1, 0, id="408-single"), pytest.param(502, 1, 0, id="502-single"), pytest.param(504, 1, 0, id="504-single"), pytest.param(429, 3, 2, id="429-backoff"), pytest.param(503, 3, 2, id="503-backoff"), ], ) def test_sync_upload_retry_set(status: int, expected_puts: int, expected_retries: int) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) sync_client: Final = _SyncRecordingClient(_coded_failure_response(status, "SlowDown")) with ( patch("litellm.integrations.s3_v2.get_httpx_client", return_value=sync_client), patch("time.sleep") as mock_sleep, ): logger.upload_data_to_s3(_element({"id": "sync"}, "sync")) assert len(sync_client.put_calls) == expected_puts _assert_jittered_backoff(mock_sleep, expected_retries) @pytest.mark.parametrize( "failure", [pytest.param(httpx.ConnectError, id="connect-error"), pytest.param(httpx.ReadTimeout, id="read-timeout")], ) def test_sync_upload_retries_a_failed_attempt_through_production_http_handler( failure: type[httpx.TransportError], caplog: pytest.LogCaptureFixture ) -> None: """HTTPHandler.put turns an httpx timeout into litellm.Timeout, and each attempt carries its own deadline.""" requests: Final[list[httpx.Request]] = [] # mutable-ok: the transport records every attempt def respond(request: httpx.Request) -> httpx.Response: requests.append(request) if len(requests) == 1: raise failure("S3 attempt failed", request=request) return httpx.Response(200, request=request) handler: Final = HTTPHandler() handler.client = httpx.Client(transport=httpx.MockTransport(respond), timeout=None) logger: Final = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) with ( patch( # test-quality-ok: sync upload builds its HTTPHandler per call, there is no injection seam for it "litellm.integrations.s3_v2.get_httpx_client", return_value=handler ), patch("time.sleep") as mock_sleep, ): logger.upload_data_to_s3(_element({"id": "sync-transient"}, "sync-transient")) handler.client.close() assert len(requests) == 2 assert requests[0].content == requests[1].content assert all( request.extensions["timeout"] == httpx.Timeout(S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS).as_dict() for request in requests ) _assert_jittered_backoff(mock_sleep, 1) assert "Error uploading to s3" not in caplog.text @pytest.mark.asyncio @pytest.mark.parametrize( ("status", "code"), [ pytest.param(503, "SlowDown", id="503"), pytest.param(429, "SlowDown", id="429"), pytest.param(500, "InternalError", id="500"), pytest.param(403, "AccessDenied", id="access-denied-403"), ], ) async def test_retryable_statuses_back_off_three_attempts(status: int, code: str | None) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) put = _FailUntilClearedPut(status=status, code=code) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"id": "req"}, "req")] with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: await logger.flush_queue() assert len(put.calls) == 3 _assert_jittered_backoff(mock_sleep, 2) assert len(logger.log_queue) == 1 @pytest.mark.asyncio @pytest.mark.parametrize( ("status", "code"), [ pytest.param(408, None, id="408"), pytest.param(502, None, id="502"), pytest.param(504, None, id="504"), ], ) async def test_unlisted_statuses_are_not_retried_in_call(status: int, code: str | None) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) put = _FailUntilClearedPut(status=status, code=code) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"id": "req"}, "req")] with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: await logger.flush_queue() assert len(put.calls) == 1 assert mock_sleep.await_args_list == [] assert len(logger.log_queue) == 1 class _FirstFailThenOkPut: def __init__(self, fail_suffix: str) -> None: self.fail_suffix = fail_suffix self.failed_once = False self.calls: tuple[str, ...] = () async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.calls = (*self.calls, url) if url.endswith(self.fail_suffix) and not self.failed_once: self.failed_once = True return _transient_failure_response() return _ok_response() @pytest.mark.asyncio async def test_retry_finishes_before_the_next_first_attempt() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_concurrent_uploads=1, ) put = _FirstFailThenOkPut("test-a.json") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"id": "a"}, "a"), _element({"id": "b"}, "b")] with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))): await logger.flush_queue() assert [call_url.rsplit("/", 1)[-1] for call_url in put.calls] == ["test-a.json", "test-a.json", "test-b.json"] assert logger.log_queue == [] @pytest.mark.asyncio async def test_objects_in_backoff_are_bounded_by_the_slot_width() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_concurrent_uploads=2, ) put = _FailUntilClearedPut(status=503) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": i}, f"{i}") for i in range(20)] sleeping: Final[list[int]] = [0] peak: Final[list[int]] = [0] async def counting_sleep(delay: float) -> None: sleeping[0] += 1 peak[0] = max(peak[0], sleeping[0]) for _ in range(10): await _real_sleep(0) sleeping[0] -= 1 with patch("asyncio.sleep", new=counting_sleep): await logger.flush_queue() assert peak[0] <= 2, f"{peak[0]} objects slept at once, slot width is 2" assert len(put.calls) == 60 @pytest.mark.asyncio async def test_subclass_returning_true_drains_the_queue() -> None: class _TrueUploadLogger(S3Logger): async def async_upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement) -> bool: return True logger = _TrueUploadLogger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) logger.async_httpx_client = AsyncMock() logger.log_queue = [_element({"id": "a"}, "a")] await logger.flush_queue() assert logger.log_queue == [] logger.async_httpx_client.put.assert_not_called() @pytest.mark.asyncio async def test_failed_direct_upload_returns_false() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) put = _FailUntilClearedPut(status=500) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put test_element = _element({"id": "x"}, "x") with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))): outcome = await logger.async_upload_data_to_s3(test_element) assert outcome is False assert len(put.calls) == 3 @pytest.mark.asyncio async def test_terminal_drop_of_one_element_does_not_drop_a_sibling_with_the_same_key() -> None: class _TerminalForMarkerPut: def __init__(self) -> None: self.calls: tuple[str | None, ...] = () async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.calls = (*self.calls, data) if data is not None and "terminal-marker" in data: return _terminal_failure_response() return _transient_failure_response() logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_drop_on_terminal_error=True, ) put = _TerminalForMarkerPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put shared_key = "2025-09-14/shared.json" dropped = s3BatchLoggingElement( s3_object_key=shared_key, payload={"m": "terminal-marker"}, s3_object_download_filename="shared.json" ) sibling = s3BatchLoggingElement( s3_object_key=shared_key, payload={"m": "healthy"}, s3_object_download_filename="shared.json" ) with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))): assert await logger._upload_outcome(dropped) == "dropped" assert await logger._upload_outcome(sibling) == "retry" def test_upload_semaphore_alias_is_the_limiter() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) assert logger._upload_semaphore is logger._upload_limiter @pytest.mark.asyncio async def test_overridden_upload_stays_bounded_by_the_configured_width() -> None: class _InFlightUploadLogger(S3Logger): def __init__(self, **kwargs: object) -> None: super().__init__(**kwargs) self.in_flight = 0 self.peak = 0 async def async_upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement) -> bool: self.in_flight += 1 self.peak = max(self.peak, self.in_flight) for _ in range(10): await _real_sleep(0) self.in_flight -= 1 return True logger = _InFlightUploadLogger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_concurrent_uploads=4, ) logger.log_queue = [_element({"i": index}, f"{index}") for index in range(40)] await logger.flush_queue() assert logger.peak <= 4 assert logger.log_queue == [] @pytest.mark.asyncio async def test_holding_the_semaphore_during_a_direct_upload_does_not_deadlock() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_concurrent_uploads=1, ) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _RecordingPut() element = _element({"id": "x"}, "x") async def held_upload() -> bool: async with logger._upload_semaphore: return await logger.async_upload_data_to_s3(element) assert await asyncio.wait_for(held_upload(), timeout=5) is True @pytest.mark.asyncio async def test_assigning_a_semaphore_changes_the_upload_width() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) logger._upload_semaphore = asyncio.Semaphore(3) put = _CountingPut(width=3) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": index}, f"{index}") for index in range(30)] await logger.flush_queue() assert put.peak == 3 assert logger.log_queue == [] @pytest.mark.asyncio async def test_terminal_code_is_retried_like_base_when_the_drop_flag_is_off() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_drop_on_terminal_error=False, ) put = _StatusPut([_coded_failure_response(403, "InvalidRequest"), _ok_response()]) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put failures = AsyncMock() logger.handle_callback_failure = failures with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))): assert await logger.async_upload_data_to_s3(_element({"id": "x"}, "x")) is True assert put.calls == 2 failures.assert_not_called() dropping = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) put.calls = 0 dropping.async_httpx_client = AsyncMock() dropping.async_httpx_client.put = put with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))): assert await dropping.async_upload_data_to_s3(_element({"id": "x"}, "x")) is False assert put.calls == 1 def test_sync_terminal_code_is_retried_like_base_when_the_drop_flag_is_off() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_drop_on_terminal_error=False, ) mock_sync_client = MagicMock() mock_sync_client.put = MagicMock(side_effect=[_coded_failure_response(403, "InvalidRequest"), _ok_response()]) failures = MagicMock() logger.handle_callback_failure = failures with ( patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client), patch("time.sleep"), ): logger.upload_data_to_s3(_element({"id": "x"}, "x")) assert mock_sync_client.put.call_count == 2 failures.assert_not_called() dropping = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) mock_sync_client.put = MagicMock(side_effect=[_coded_failure_response(403, "InvalidRequest"), _ok_response()]) with ( patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client), patch("time.sleep"), ): dropping.upload_data_to_s3(_element({"id": "x"}, "x")) assert mock_sync_client.put.call_count == 1 def test_sync_retry_lines_stay_at_warning_level(caplog) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) mock_sync_client = MagicMock() mock_sync_client.put = MagicMock( side_effect=[_transient_failure_response(503), _transient_failure_response(503), _ok_response()] ) with ( caplog.at_level("WARNING"), patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client), patch("time.sleep"), ): logger.upload_data_to_s3(_element({"id": "x"}, "x")) assert mock_sync_client.put.call_count == 3 assert sum(1 for record in caplog.records if "retrying in" in record.getMessage()) == 2 @pytest.mark.asyncio async def test_direct_async_upload_logs_retry_lines_at_warning_level(caplog) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) put = _StatusPut([_transient_failure_response(503), _ok_response()]) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put with caplog.at_level("WARNING"), patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))): assert await logger.async_upload_data_to_s3(_element({"id": "x"}, "x")) is True assert put.calls == 2 assert sum(1 for record in caplog.records if "retrying in" in record.getMessage()) == 1 def _init_bypassed_logger() -> S3Logger: from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM logger = S3Logger.__new__(S3Logger) logger.iam_cache = BaseAWSLLM._shared_iam_cache logger.s3_endpoint_url = None logger.s3_bucket_name = "test-bucket" logger.s3_region_name = "us-east-1" logger.s3_use_virtual_hosted_style = False logger.s3_verify = None logger.s3_aws_access_key_id = "test-key" logger.s3_aws_secret_access_key = "test-secret" logger.s3_aws_session_token = None logger.s3_aws_session_name = None logger.s3_aws_profile_name = None logger.s3_aws_role_name = None logger.s3_aws_web_identity_token = None logger.s3_aws_sts_endpoint = None logger.s3_server_side_encryption = None logger.s3_sse_kms_key_id = None logger.s3_log_prompts_only = None return logger @pytest.mark.asyncio async def test_init_bypassed_logger_retries_a_503_and_reports_a_404() -> None: logger = _init_bypassed_logger() put = _StatusPut([_transient_failure_response(503), _ok_response()]) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))): assert await logger.async_upload_data_to_s3(_element({"id": "x"}, "x")) is True assert put.calls == 2 put.calls = 0 put.responses = [_coded_failure_response(404, None)] with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))): assert await logger.async_upload_data_to_s3(_element({"id": "y"}, "y")) is False assert put.calls == 1 def test_init_bypassed_sync_logger_retries_a_503_and_reports_a_404() -> None: logger = _init_bypassed_logger() mock_sync_client = MagicMock() mock_sync_client.put = MagicMock(side_effect=[_transient_failure_response(503), _ok_response()]) with ( patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client), patch("time.sleep"), ): logger.upload_data_to_s3(_element({"id": "x"}, "x")) assert mock_sync_client.put.call_count == 2 retried_headers: Final = dict(mock_sync_client.put.call_args.kwargs["headers"]) assert "X-Amz-Date" in retried_headers mock_sync_client.put = MagicMock(return_value=_coded_failure_response(404, None)) with ( patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client), patch("time.sleep"), ): logger.upload_data_to_s3(_element({"id": "y"}, "y")) assert mock_sync_client.put.call_count == 1 failed_url: Final = str(mock_sync_client.put.call_args[0][0]) assert "test-y.json" in failed_url @pytest.mark.asyncio async def test_subclass_with_base_style_upload_bounded_drains_the_queue() -> None: class _BaseStyleLogger(S3Logger): async def _upload_bounded(self, element: s3BatchLoggingElement) -> bool: return True logger = _BaseStyleLogger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) logger.async_httpx_client = AsyncMock() logger.log_queue = [_element({"id": "a"}, "a")] await logger.flush_queue() assert logger.log_queue == [] logger.async_httpx_client.put.assert_not_called() def test_bool_config_values_fall_back_to_the_default() -> None: from litellm.integrations.s3 import ( resolve_s3_max_concurrent_uploads, resolve_s3_max_queue_size, resolve_s3_max_retry_age_seconds, ) assert resolve_s3_max_concurrent_uploads(True, 16) == 1 assert resolve_s3_max_queue_size(True, 50000) == 50000 assert resolve_s3_max_retry_age_seconds(True, 3600) == 3600 def test_int_env_helper_falls_back_on_non_numeric(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.litellm_core_utils.env_utils import get_env_int monkeypatch.setenv("TEST_S3_INT_ENV", "abc") assert get_env_int("TEST_S3_INT_ENV", 3) == 3 monkeypatch.setenv("TEST_S3_INT_ENV", "7") assert get_env_int("TEST_S3_INT_ENV", 3) == 7 class _FailOncePerKeyPut: def __init__(self) -> None: self.failed: set[str] = set() self.calls: tuple[str, ...] = () async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.calls = (*self.calls, url) if url not in self.failed: self.failed.add(url) return _transient_failure_response() return _ok_response() class _SlowFailOncePerKeyPut: def __init__(self, dumps_count) -> None: self.failed: set[str] = set() self.dumps_count = dumps_count self.first_completed: int | None = None self.calls: tuple[str, ...] = () async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.calls = (*self.calls, url) await _real_sleep(0) if self.first_completed is None: self.first_completed = self.dumps_count() if url not in self.failed: self.failed.add(url) return _transient_failure_response() return _ok_response() @pytest.mark.asyncio async def test_peak_serialized_bodies_bounded_by_upload_width() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps as real_safe_dumps dumps_calls: list[object] = [] def counting_dumps(*args, **kwargs): dumps_calls.append(args) return real_safe_dumps(*args, **kwargs) put = _SlowFailOncePerKeyPut(lambda: len(dumps_calls)) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": index}, f"{index}") for index in range(64)] with ( patch("litellm.integrations.s3_v2.safe_dumps", side_effect=counting_dumps), patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))), ): await logger.flush_queue() assert put.first_completed is not None assert put.first_completed <= logger.s3_max_concurrent_uploads assert len(dumps_calls) == 64 assert len(put.calls) == 128 @pytest.mark.asyncio async def test_send_batch_calls_upload_with_one_positional_arg() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) uploaded: list[str] = [] # mutable-ok: appended once per upload by the double async def mock_upload(batch_logging_element) -> str: uploaded.append(batch_logging_element.s3_object_key) return "delivered" logger.async_upload_data_to_s3 = mock_upload logger.log_queue = [_element({"id": "a"}, "a"), _element({"id": "b"}, "b")] await logger.flush_queue() assert sorted(key.rsplit("/", 1)[-1] for key in uploaded) == ["test-a.json", "test-b.json"] assert logger.log_queue == [] @pytest.mark.asyncio async def test_retries_serialize_the_body_once_per_element_per_flush() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps as real_safe_dumps dumps_calls: list[object] = [] def counting_dumps(*args, **kwargs): dumps_calls.append(args) return real_safe_dumps(*args, **kwargs) put = _FailUntilClearedPut(status=503) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": index}, f"{index}") for index in range(8)] with ( patch("litellm.integrations.s3_v2.safe_dumps", side_effect=counting_dumps), patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))), ): await logger.flush_queue() assert len(dumps_calls) == 8 assert len(put.calls) == 24 assert len(logger.log_queue) == 8 @pytest.mark.asyncio async def test_async_flush_logs_one_retry_warning(caplog) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) put = _FailOncePerKeyPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"i": index}, f"{index}") for index in range(8)] with caplog.at_level("WARNING"), patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert logger.log_queue == [] assert sum(1 for record in caplog.records if "in-call retries" in record.getMessage()) == 1 assert all("retrying in" not in record.getMessage() for record in caplog.records) class _AppendingSuffixFailingPut: def __init__(self, logger: S3Logger, element: s3BatchLoggingElement, fail_suffixes: tuple[str, ...]) -> None: self.logger = logger self.element = element self.fail_suffixes = fail_suffixes self.appended = False async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: if not self.appended: self.appended = True self.logger.log_queue.append(self.element) if url.endswith(self.fail_suffixes): return _transient_failure_response() return _ok_response() @pytest.mark.asyncio async def test_failed_elements_stay_oldest_first_when_requeued() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_retry_age_seconds=3600, ) late = _element({"id": "late"}, "late") failed = _element({"id": "f2"}, "f2") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _AppendingSuffixFailingPut(logger, late, ("test-f2.json",)) logger.log_queue = [_element({"id": "f1"}, "f1"), failed] with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert [element.s3_object_key for element in logger.log_queue] == [failed.s3_object_key, late.s3_object_key] assert logger.log_queue[0].retrying_since is not None @pytest.mark.asyncio async def test_overflow_prefers_arrivals_over_failed_elements_without_counting_the_trim(caplog) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_queue_size=1, ) late = _element({"id": "late"}, "late") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _AppendingSuffixFailingPut(logger, late, ("test-f2.json",)) logger.log_queue = [_element({"id": "f1"}, "f1"), _element({"id": "f2"}, "f2")] with ( patch.object(logger, "handle_callback_failure") as mock_failure, patch("asyncio.sleep", new_callable=AsyncMock), ): await logger.flush_queue() assert [element.s3_object_key for element in logger.log_queue] == [late.s3_object_key] failed_uploads: Final = 1 assert mock_failure.call_count == failed_uploads assert "dropped 1 oldest events" in caplog.text @pytest.mark.asyncio async def test_fresh_elements_upload_before_stale_retries_after_a_failed_flush() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_concurrent_uploads=1, ) late = _element({"id": "late"}, "late") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _AppendingSuffixFailingPut(logger, late, ("test-f2.json",)) logger.log_queue = [_element({"id": "f1"}, "f1"), _element({"id": "f2"}, "f2")] with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() recovered = _FailOnSuffixPut(("never-matches",)) logger.async_httpx_client.put = recovered with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert [call_url.rsplit("/", 1)[-1] for call_url in recovered.calls] == ["test-late.json", "test-f2.json"] assert logger.log_queue == [] @pytest.mark.asyncio async def test_repeated_overflow_trims_oldest_across_failed_flushes(caplog) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_queue_size=3, ) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _AppendingFailingPut(logger, (_element({"id": "d"}, "d"),)) logger.log_queue = [_element({"id": name}, name) for name in ("a", "b", "c")] with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert [element.payload["id"] for element in logger.log_queue] == ["b", "c", "d"] logger.async_httpx_client.put = _AppendingFailingPut(logger, (_element({"id": "e"}, "e"),)) with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert [element.payload["id"] for element in logger.log_queue] == ["c", "d", "e"] assert caplog.text.count("dropped 1 oldest events") == 2 @pytest.mark.asyncio async def test_retry_age_budget_drops_after_the_clock_set_by_a_partial_failure(caplog) -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_retry_age_seconds=1, ) put = _FailOnSuffixCodedPut(("test-poison.json",), 503, "SlowDown") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"id": "poison"}, "poison"), _element({"id": "good"}, "good")] t0: Final = _NOW with patch.object(logger, "handle_callback_failure") as mock_failure: with ( patch("asyncio.sleep", new_callable=AsyncMock), patch("time.monotonic", return_value=t0), ): await logger.flush_queue() assert [element.s3_object_key for element in logger.log_queue] == ["2025-09-14/test-poison.json"] assert logger.log_queue[0].retrying_since == t0 logger.log_queue.append(_element({"id": "good-2"}, "good-2")) with ( patch("asyncio.sleep", new_callable=AsyncMock), patch("time.monotonic", return_value=t0 + 2), ): await logger.flush_queue() assert logger.log_queue == [] assert "retrying longer than s3_max_retry_age_seconds=1" in caplog.text poison_puts: Final = sum(1 for call_url in put.calls if call_url.endswith("test-poison.json")) assert poison_puts == 6 upload_failures: Final = 2 assert mock_failure.call_count == upload_failures @pytest.mark.asyncio async def test_the_retry_clock_starts_at_the_first_partial_failure_not_first_seen() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_retry_age_seconds=1, ) put = _FailUntilClearedPut(status=503, code="SlowDown") logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"id": "poison"}, "poison")] t0: Final = _NOW with ( patch("asyncio.sleep", new_callable=AsyncMock), patch("time.monotonic", return_value=t0), ): await logger.flush_queue() assert len(logger.log_queue) == 1 assert logger.log_queue[0].retrying_since is None logger.log_queue.append(_element({"id": "good"}, "good")) put.failing = False failing_poison: Final = _FailOnSuffixCodedPut(("test-poison.json",), 503, "SlowDown") logger.async_httpx_client.put = failing_poison with ( patch("asyncio.sleep", new_callable=AsyncMock), patch("time.monotonic", return_value=t0 + 500), ): await logger.flush_queue() assert [element.s3_object_key for element in logger.log_queue] == ["2025-09-14/test-poison.json"] assert logger.log_queue[0].retrying_since == t0 + 500 def test_sync_upload_retries_access_denied_403(caplog): logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-sync-403.json", payload={"test": "sync-403"}, s3_object_download_filename="test-sync-403.json", ) mock_sync_client = MagicMock() mock_sync_client.put = MagicMock(return_value=_coded_failure_response(403, "AccessDenied")) with ( patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client), patch("time.sleep") as mock_sleep, ): logger.upload_data_to_s3(test_element) assert mock_sync_client.put.call_count == 3 _assert_jittered_backoff(mock_sleep, 2) assert "dropping object" not in caplog.text def test_sync_upload_drops_terminal_object_once_and_logs_it_only_when_opted_in(caplog): logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_drop_on_terminal_error=True, ) test_element = s3BatchLoggingElement( s3_object_key="2025-09-14/test-sync-terminal.json", payload={"test": "sync-terminal"}, s3_object_download_filename="test-sync-terminal.json", ) mock_sync_client = MagicMock() mock_sync_client.put = MagicMock(return_value=_coded_failure_response(400, "EntityTooLarge")) with ( patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client), patch("time.sleep") as mock_sleep, ): logger.upload_data_to_s3(test_element) assert mock_sync_client.put.call_count == 1 mock_sleep.assert_not_called() assert "dropping object" in caplog.text @pytest.mark.asyncio async def test_requeued_batch_file_keeps_the_earliest_member_retrying_since() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_batch_file_upload=True, ) put = _FailUntilClearedPut() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put stale: Final = _NOW - 30 retried = _element({"id": "retried"}, "retried").model_copy(update={"retrying_since": stale}) fresh = _element({"id": "fresh"}, "fresh") logger.log_queue = [retried, fresh] with ( patch("asyncio.sleep", new_callable=AsyncMock), patch("time.monotonic", return_value=_NOW), ): await logger.flush_queue() assert len(logger.log_queue) == 1 assert logger.log_queue[0].s3_object_key.endswith(".jsonl") assert logger.log_queue[0].retrying_since == stale @pytest.mark.asyncio async def test_an_unlisted_5xx_is_requeued_without_an_extra_attempt() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) put = _FailUntilClearedPut(status=507) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = put logger.log_queue = [_element({"id": "req-507"}, "507")] with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert len(logger.log_queue) == 1 assert len(put.calls) == 1 @pytest.mark.parametrize("configured", [0, "0", None, ""]) def test_retry_age_resolution_disables_the_budget(configured: object) -> None: from litellm.integrations.s3 import resolve_s3_max_retry_age_seconds assert resolve_s3_max_retry_age_seconds(configured, 3600) is None @pytest.mark.parametrize("configured", ["abc", -5, True]) def test_invalid_retry_age_resolution_falls_back_with_a_warning(configured: object, caplog) -> None: from litellm.integrations.s3 import resolve_s3_max_retry_age_seconds assert resolve_s3_max_retry_age_seconds(configured, 3600) == 3600 assert "s3_max_retry_age_seconds" in caplog.text def test_retry_age_resolution_accepts_a_positive_int() -> None: from litellm.integrations.s3 import resolve_s3_max_retry_age_seconds assert resolve_s3_max_retry_age_seconds(30, 3600) == 30 def test_default_logger_sets_a_one_hour_retry_age_budget() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", ) assert logger.s3_max_retry_age_seconds == 3600 def test_constructor_zero_disables_the_retry_age_budget() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_retry_age_seconds=0, ) assert logger.s3_max_retry_age_seconds is None def test_invalid_callback_params_retry_age_falls_back_to_the_default() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_callback_params_override={"s3_max_retry_age_seconds": "abc"}, ) assert logger.s3_max_retry_age_seconds == 3600 def test_callback_params_retry_age_wins_over_constructor() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_retry_age_seconds=30, s3_callback_params_override={"s3_max_retry_age_seconds": 60}, ) assert logger.s3_max_retry_age_seconds == 60 def test_callback_params_drop_terminal_error_wins_over_constructor() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_drop_on_terminal_error=False, s3_callback_params_override={"s3_drop_on_terminal_error": True}, ) assert logger.s3_drop_on_terminal_error is True def test_invalid_callback_params_drop_terminal_error_falls_back_to_constructor_value() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_drop_on_terminal_error=True, s3_callback_params_override={"s3_drop_on_terminal_error": "banana"}, ) assert logger.s3_drop_on_terminal_error is True def test_callback_params_adaptive_concurrency_wins_over_constructor() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_adaptive_concurrency=False, s3_callback_params_override={"s3_adaptive_concurrency": "true"}, ) assert logger.s3_adaptive_concurrency is True assert logger._upload_limiter._ceiling > logger._upload_limiter.limit def test_invalid_callback_params_max_adaptive_concurrency_falls_back_to_default() -> None: from litellm.constants import DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY logger = _override_logger(s3_adaptive_concurrency=True, s3_max_adaptive_concurrency="abc") assert logger.s3_max_adaptive_concurrency == DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY assert logger._upload_limiter._ceiling == DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY def test_callback_params_queue_size_wins_over_constructor() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_queue_size=7, s3_callback_params_override={"s3_max_queue_size": 4}, ) assert logger.s3_max_queue_size == 4 assert logger.max_queue_size == 4 def test_invalid_callback_params_queue_size_falls_back_to_constructor_value() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_queue_size=7, s3_callback_params_override={"s3_max_queue_size": "abc"}, ) assert logger.s3_max_queue_size == 7 def test_invalid_constructor_queue_size_falls_back_to_default() -> None: from litellm.integrations.custom_batch_logger import CustomBatchLogger logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_queue_size="abc", ) assert logger.s3_max_queue_size == CustomBatchLogger.DEFAULT_MAX_QUEUE_SIZE assert logger.max_queue_size == CustomBatchLogger.DEFAULT_MAX_QUEUE_SIZE class _StatusPut: def __init__(self, responses: "list[MagicMock | Exception]") -> None: self.responses = responses self.calls = 0 async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.calls += 1 outcome = self.responses[min(self.calls - 1, len(self.responses) - 1)] if isinstance(outcome, Exception): raise outcome return outcome def _slow_down_response(status: int = 200) -> MagicMock: response = _ok_response() if status == 200 else _transient_failure_response(status) response.text = "SlowDown" return response @pytest.mark.asyncio async def test_503_response_lowers_the_adaptive_limit() -> None: logger = _override_logger(s3_adaptive_concurrency=True) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _StatusPut([_transient_failure_response(503), _ok_response()]) logger._upload_limiter._limit = 64 # mutable-ok: seed the AIMD state above the floor without replaying growth with patch("asyncio.sleep", new_callable=AsyncMock): logger.log_queue = [_element({"i": 0}, "0")] await logger.flush_queue() assert logger._upload_limiter.limit == 32 @pytest.mark.asyncio async def test_429_response_lowers_the_adaptive_limit() -> None: logger = _override_logger(s3_adaptive_concurrency=True) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _StatusPut([_transient_failure_response(429), _ok_response()]) logger._upload_limiter._limit = 64 # mutable-ok: seed the AIMD state above the floor without replaying growth with patch("asyncio.sleep", new_callable=AsyncMock): logger.log_queue = [_element({"i": 0}, "0")] await logger.flush_queue() assert logger._upload_limiter.limit == 32 @pytest.mark.asyncio async def test_slow_down_body_code_lowers_the_adaptive_limit() -> None: logger = _override_logger(s3_adaptive_concurrency=True) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _StatusPut([_slow_down_response()]) logger._upload_limiter._limit = 64 # mutable-ok: seed the AIMD state above the floor without replaying growth logger.log_queue = [_element({"i": 0}, "0")] await logger.async_send_batch() assert logger._upload_limiter.limit == 32 @pytest.mark.asyncio async def test_transport_error_lowers_the_adaptive_limit() -> None: logger = _override_logger(s3_adaptive_concurrency=True) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _StatusPut( [httpx.ConnectError("connect refused", request=MagicMock()), _ok_response()] ) logger._upload_limiter._limit = 64 # mutable-ok: seed the AIMD state above the floor without replaying growth with patch("asyncio.sleep", new_callable=AsyncMock): logger.log_queue = [_element({"i": 0}, "0")] await logger.flush_queue() assert logger._upload_limiter.limit == 32 @pytest.mark.asyncio async def test_fast_uploads_raise_the_adaptive_limit() -> None: logger = _override_logger(s3_adaptive_concurrency=True) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _RecordingPut() before: Final = logger._upload_limiter.limit logger.log_queue = [_element({"i": i}, f"{i}") for i in range(before)] await logger.async_send_batch() assert logger._upload_limiter.limit > before @pytest.mark.asyncio async def test_configured_concurrency_is_the_fixed_limit_when_adaptive_is_off() -> None: logger = S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", s3_max_concurrent_uploads=64, ) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _StatusPut([_transient_failure_response(503), _ok_response()]) assert logger._upload_limiter._value == 64 logger.log_queue = [_element({"i": 0}, "0")] with patch("asyncio.sleep", new_callable=AsyncMock): await logger.flush_queue() assert logger._upload_limiter._value == 64 @pytest.mark.asyncio async def test_the_limit_never_falls_below_the_configured_width() -> None: logger = _override_logger(s3_adaptive_concurrency=True, s3_max_concurrent_uploads=8) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _StatusPut( [_transient_failure_response(503), _transient_failure_response(503), _transient_failure_response(503)] ) with patch("asyncio.sleep", new_callable=AsyncMock): logger.log_queue = [_element({"i": 0}, "0")] await logger.flush_queue() assert logger._upload_limiter.limit == 8 def test_default_upload_width_is_16() -> None: from litellm.constants import DEFAULT_S3_MAX_CONCURRENT_UPLOADS logger = _override_logger() assert logger._upload_semaphore._value == DEFAULT_S3_MAX_CONCURRENT_UPLOADS assert DEFAULT_S3_MAX_CONCURRENT_UPLOADS == 16 @pytest.mark.asyncio async def test_a_slow_put_does_not_lower_the_adaptive_limit() -> None: logger = _override_logger(s3_adaptive_concurrency=True) logger.async_httpx_client = AsyncMock() async def slow_put( url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None ) -> MagicMock: await _real_sleep(0) return _ok_response() logger.async_httpx_client.put = slow_put before: Final = logger._upload_limiter.limit logger.log_queue = [_element({"i": 0}, "0")] await logger.async_send_batch() assert logger._upload_limiter.limit >= before class _FastOkPut: def __init__(self) -> None: self.calls = 0 async def __call__( self, url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, ) -> MagicMock: self.calls += 1 await _real_sleep(0) return _ok_response() async def _timed_send_batch(size: int) -> float: logger = _override_logger() logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = _FastOkPut() logger.log_queue = [_element({"i": i}, f"{i}") for i in range(size)] started = time.perf_counter() await logger.async_send_batch() return time.perf_counter() - started @pytest.mark.asyncio async def test_send_batch_time_grows_linearly_with_the_batch() -> None: baseline: Final = await _timed_send_batch(2_000) quadrupled: Final = await _timed_send_batch(8_000) assert quadrupled / baseline < 8, f"2k took {baseline:.3f}s, 8k took {quadrupled:.3f}s" class QueueOnlyS3LoggerWithoutCredentialsOrPeriodicFlush(S3Logger): def __init__(self) -> None: CustomLogger.__init__(self) self.log_queue: list[s3BatchLoggingElement] = [] self.batch_size = 100 self.s3_strip_base64_files = False self.s3_use_team_prefix = False self.s3_use_key_prefix = False self.s3_path = "" self.s3_log_prompts_only = None self.s3_batch_file_upload = False self._upload_semaphore = asyncio.Semaphore(1) @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("concurrent", [False, True]) @pytest.mark.asyncio async def test_repeated_failure_notifications_upload_once( monkeypatch: pytest.MonkeyPatch, stream: bool, concurrent: bool, ) -> None: sink: Final = QueueOnlyS3LoggerWithoutCredentialsOrPeriodicFlush() upload: Final = AsyncMock(return_value=True) monkeypatch.setattr(sink, "async_upload_data_to_s3", upload) monkeypatch.setattr(litellm, "callbacks", []) monkeypatch.setattr(litellm, "_async_failure_callback", [sink]) logging_obj: Final = Logging( model="openai/gpt-4o-mini", messages=[], stream=stream, call_type="acompletion", start_time=datetime.now(), litellm_call_id="synthetic-request", function_id="synthetic-function", ) error: Final = litellm.ServiceUnavailableError( message="synthetic deployment-selection failure", llm_provider="openai", model="openai/gpt-4o-mini", ) if concurrent: await asyncio.gather( *(logging_obj.async_failure_handler(error, "synthetic traceback") for _ in range(3)) ) else: for _ in range(3): await logging_obj.async_failure_handler(error, "synthetic traceback") assert len(sink.log_queue) == 1 assert all(entry.payload["status"] == "failure" for entry in sink.log_queue) await sink.async_send_batch() await asyncio.sleep(0) assert upload.await_count == 1 @pytest.fixture() def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) @pytest_asyncio.fixture(loop_scope="function") async def drain_logging_worker(isolate_litellm_state: None) -> AsyncIterator[None]: yield await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS: Final = LOGGING_WORKER_MAX_TIME_PER_COROUTINE + 5.0 @pytest.fixture(scope="function", autouse=True) def isolate_litellm_state(): """ Per-function isolation fixture. Resets litellm state to the true defaults captured at conftest import time, then restores after the test. This prevents module-level mutations (e.g. `litellm.num_retries = 3` at the top of test_langfuse_e2e_test.py) from leaking across tests within the same xdist worker. """ from litellm.litellm_core_utils import litellm_logging as ll_logging from litellm.proxy.management_helpers import audit_logs as ll_audit_logs if hasattr(litellm, "in_memory_llm_clients_cache"): litellm.in_memory_llm_clients_cache.flush_cache() ll_logging._in_memory_loggers.clear() ll_audit_logs._audit_log_callback_cache.clear() for attr in _LIST_ATTRS: if attr in _DEFAULTS: default = _DEFAULTS[attr] setattr(litellm, attr, default.copy() if isinstance(default, list) else default) for attr in _SCALAR_ATTRS: if attr in _DEFAULTS: setattr(litellm, attr, _DEFAULTS[attr]) yield if hasattr(litellm, "in_memory_llm_clients_cache"): litellm.in_memory_llm_clients_cache.flush_cache() ll_logging._in_memory_loggers.clear() ll_audit_logs._audit_log_callback_cache.clear() for attr in _LIST_ATTRS: if attr in _DEFAULTS: default = _DEFAULTS[attr] setattr(litellm, attr, default.copy() if isinstance(default, list) else default) for attr in _SCALAR_ATTRS: if attr in _DEFAULTS: setattr(litellm, attr, _DEFAULTS[attr]) _LIST_ATTRS = ( "callbacks", "success_callback", "failure_callback", "_async_success_callback", "_async_failure_callback", "service_callback", "pre_call_rules", "post_call_rules", ) _SCALAR_ATTRS = ( "set_verbose", "cache", "num_retries", "num_retries_per_request", "turn_off_message_logging", "redact_messages_in_exceptions", "redact_user_api_key_info", "s3_callback_params", "s3_audit_callback_params", "datadog_params", "vector_store_registry", ) _DEFAULTS: Final = { attr: getattr(litellm, attr).copy() if isinstance(getattr(litellm, attr), list) else getattr(litellm, attr) for attr in (*_LIST_ATTRS, *_SCALAR_ATTRS) if hasattr(litellm, attr) } @pytest.fixture(scope="module") def setup_and_teardown(): """ Module-scoped setup. Reloads litellm only in single-process mode (skipped under xdist to avoid cross-worker interference). """ import litellm worker_id = os.environ.get("PYTEST_XDIST_WORKER", None) if worker_id is None: importlib.reload(litellm) try: if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): import litellm.proxy.proxy_server importlib.reload(litellm.proxy.proxy_server) except Exception: pass if hasattr(litellm, "in_memory_llm_clients_cache"): litellm.in_memory_llm_clients_cache.flush_cache() yield @pytest.fixture def amazing_s3_retries(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "num_retries", 3) class _FakeS3Paginator: def __init__(self, objects): self.objects = objects def paginate(self, Bucket): keys = sorted(self.objects[Bucket]) if not keys: return [{}] return [{"Contents": [{"Key": key} for key in keys]}] class _FakeS3Client: def __init__(self): self.objects = defaultdict(dict) self.object_written = threading.Event() def clear(self): self.objects.clear() self.object_written.clear() def put_object(self, Bucket, Key, Body, **_kwargs): self.objects[Bucket][Key] = Body self.object_written.set() return {"ResponseMetadata": {"HTTPStatusCode": 200}} def delete_object(self, Bucket, Key): self.objects[Bucket].pop(Key, None) return {"ResponseMetadata": {"HTTPStatusCode": 204}} def get_paginator(self, name): assert name == "list_objects_v2" return _FakeS3Paginator(self.objects) def list_objects(self, Bucket): keys = sorted(self.objects[Bucket]) return {"Contents": [{"Key": key, "LastModified": 0} for key in keys]} _FAKE_S3_CLIENT = _FakeS3Client() @pytest.fixture def fake_s3_client(monkeypatch): _FAKE_S3_CLIENT.clear() def fake_boto3_client(service_name, *args, **kwargs): assert service_name == "s3" return _FAKE_S3_CLIENT monkeypatch.setattr(boto3, "client", fake_boto3_client) litellm.success_callback = [] litellm.callbacks = [] yield _FAKE_S3_CLIENT litellm.success_callback = [] litellm.callbacks = [] @pytest.mark.usefixtures( "cancel_s3_periodic_flush_tasks", "fake_s3_client", "_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown", "amazing_s3_retries", ) @pytest.mark.asyncio @pytest.mark.parametrize("sync_mode,streaming", [(True, True), (True, False), (False, True), (False, False)]) @pytest.mark.flaky(retries=3, delay=1) async def test_basic_s3_logging(sync_mode, streaming): verbose_logger.setLevel(level=logging.DEBUG) litellm.success_callback = ["s3"] litellm.s3_callback_params = { "s3_bucket_name": "load-testing-oct", "s3_aws_secret_access_key": "os.environ/AWS_SECRET_ACCESS_KEY", "s3_aws_access_key_id": "os.environ/AWS_ACCESS_KEY_ID", "s3_region_name": "us-west-2", } litellm.set_verbose = True response_id = None if sync_mode is True: response = litellm.completion( model="gpt-5-mini", messages=[{"role": "user", "content": "This is a test"}], mock_response="It's simple to use and easy to get started", stream=streaming, ) if streaming: for chunk in response: response_id = chunk.id else: response_id = response.id assert _FAKE_S3_CLIENT.object_written.wait(timeout=5) else: response = await litellm.acompletion( model="gpt-5-mini", messages=[{"role": "user", "content": "This is a test"}], mock_response="It's simple to use and easy to get started", stream=streaming, ) if streaming: async for chunk in response: response_id = chunk.id else: response_id = response.id assert await asyncio.to_thread(_FAKE_S3_CLIENT.object_written.wait, 5) total_objects, all_s3_keys = list_all_s3_objects("load-testing-oct") assert any(response_id in key for key in all_s3_keys) s3 = boto3.client("s3") for key in all_s3_keys: s3.delete_object(Bucket="load-testing-oct", Key=key) @pytest.mark.usefixtures( "cancel_s3_periodic_flush_tasks", "fake_s3_client", "_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown", "amazing_s3_retries", ) @pytest.mark.asyncio @pytest.mark.parametrize("streaming", [True]) @pytest.mark.flaky(retries=3, delay=1) async def test_basic_s3_v2_logging(streaming): from litellm.integrations.s3_v2 import S3Logger litellm.s3_callback_params = { "s3_bucket_name": "load-testing-oct", "s3_aws_secret_access_key": "test-secret", "s3_aws_access_key_id": "test-key", "s3_region_name": "us-west-2", } s3_v2_logger = S3Logger(s3_flush_interval=1) litellm.callbacks = [s3_v2_logger] uploaded_keys: list = [] uploaded = asyncio.Event() async def mock_upload(batch_logging_element): uploaded_keys.append(batch_logging_element.s3_object_key) uploaded.set() s3_v2_logger.async_upload_data_to_s3 = mock_upload litellm.set_verbose = True response_id = None response = await litellm.acompletion( model="gpt-5-mini", messages=[{"role": "user", "content": "This is a test"}], mock_response="It's simple to use and easy to get started", stream=streaming, ) if streaming: async for chunk in response: response_id = chunk.id else: response_id = response.id assert await asyncio.wait_for(uploaded.wait(), timeout=5) assert len(uploaded_keys) > 0, "S3 upload was never called" assert any(response_id in key for key in uploaded_keys), ( f"Expected response_id={response_id} in one of the uploaded S3 keys: {uploaded_keys}" ) @pytest.mark.usefixtures( "cancel_s3_periodic_flush_tasks", "fake_s3_client", "_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown", "amazing_s3_retries", ) @pytest.mark.asyncio @pytest.mark.flaky(retries=3, delay=1) async def test_basic_s3_v2_logging_failure(): """Test that S3 v2 logger makes httpx PUT request when logging failures""" from unittest.mock import AsyncMock, MagicMock from litellm.integrations.s3_v2 import S3Logger s3_v2_logger = S3Logger(s3_flush_interval=1) mock_response = MagicMock() mock_response.status_code = 200 mock_response.raise_for_status = MagicMock() s3_v2_logger.async_httpx_client = AsyncMock() s3_v2_logger.async_httpx_client.put.return_value = mock_response upload_called = False uploaded = asyncio.Event() async def mock_upload(batch_logging_element): nonlocal upload_called upload_called = True url = f"https://test-bucket.s3.us-west-2.amazonaws.com/{batch_logging_element.s3_object_key}" headers = {"Content-Type": "application/json"} data = '{"model": "gpt-5-mini"}' await s3_v2_logger.async_httpx_client.put(url=url, headers=headers, data=data) uploaded.set() s3_v2_logger.async_upload_data_to_s3 = mock_upload litellm.callbacks = [s3_v2_logger] litellm.s3_callback_params = { "s3_bucket_name": "test-bucket", "s3_aws_secret_access_key": "test-secret", "s3_aws_access_key_id": "test-key", "s3_region_name": "us-west-2", } litellm.set_verbose = True try: await litellm.acompletion( model="gpt-5-mini", api_key="invalid-api-key", messages=[{"role": "user", "content": "This is a test"}], mock_response=Exception("forced failure for S3 logging test"), ) except Exception: pass assert await asyncio.wait_for(uploaded.wait(), timeout=5) assert upload_called, "S3 upload method was not called" s3_v2_logger.async_httpx_client.put.assert_called() call_args = s3_v2_logger.async_httpx_client.put.call_args assert call_args is not None url = call_args[1]["url"] if "url" in call_args[1] else call_args[0][0] assert "test-bucket.s3.us-west-2.amazonaws.com" in url headers = call_args[1]["headers"] assert headers["Content-Type"] == "application/json" data = call_args[1]["data"] assert data is not None assert '"model": "gpt-5-mini"' in data def list_all_s3_objects(bucket_name): s3 = boto3.client("s3") all_s3_keys = [] paginator = s3.get_paginator("list_objects_v2") total_objects = 0 for page in paginator.paginate(Bucket=bucket_name): if "Contents" in page: total_objects += len(page["Contents"]) all_s3_keys.extend([obj["Key"] for obj in page["Contents"]]) return total_objects, all_s3_keys class TestS3Logger(S3Logger): def __init__(self, *args, **kwargs): self.recorded_requests = {} self.logged_standard_logging_payload = None super().__init__(*args, **kwargs) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self.recorded_requests[response_obj["id"]] = start_time self.logged_standard_logging_payload = kwargs["standard_logging_object"] return await super().async_log_success_event(kwargs, response_obj, start_time, end_time) def _batch_line_items_payload() -> dict: line = { "batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", "custom_id": "ok-1", "status": "success", "request": {"custom_id": "ok-1", "body": {"messages": [{"role": "user", "content": "hi"}]}}, "response": {"custom_id": "ok-1", "response": {"status_code": 200, "body": {"id": "chatcmpl-1"}}}, "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, "response_cost": 0.005, "model": "gpt-4o-mini", "custom_llm_provider": "openai", "metadata": {"user_api_key_hash": "hash-alice"}, } return { "batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", "finished_at": 1791494130, "part": 0, "lines": (line, {**line, "custom_id": "ok-2"}), } def _s3_logger_answering( status: int, s3_path: str | None, content: bytes = b"" ) -> tuple[S3Logger, list[httpx.Request]]: requests: list[httpx.Request] = [] def respond(request: httpx.Request) -> httpx.Response: requests.append(request) return httpx.Response(status, content=content, request=request) logger = S3Logger( s3_bucket_name="test-bucket", s3_region_name="us-east-1", s3_path=s3_path, s3_aws_access_key_id="AKIAIOSFODNN7EXAMPLE", s3_aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", s3_flush_interval=3600, ) logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) return logger, requests @pytest.mark.asyncio async def test_batch_line_items_are_one_ndjson_object_under_the_configured_path(): logger, requests = _s3_logger_answering(200, "cold/") payload = _batch_line_items_payload() await logger.async_log_batch_line_items(payload) (request,) = requests assert request.method == "PUT" assert request.url.path == "/cold/batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00000.jsonl" assert request.headers["Content-Type"] == "application/x-ndjson" assert request.headers["Content-Disposition"] == 'inline; filename="bGl0ZWxsbV9wcm94eTtiYXRjaA-part-00000.jsonl"' assert [json.loads(line) for line in request.content.decode().splitlines()] == list(payload["lines"]) assert logger.log_queue == [] @pytest.mark.asyncio async def test_batch_line_items_without_an_s3_path_sit_at_the_bucket_root(): logger, requests = _s3_logger_answering(200, None) await logger.async_log_batch_line_items(_batch_line_items_payload()) assert [request.url.path for request in requests] == [ "/batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00000.jsonl" ] @pytest.mark.asyncio async def test_a_rejected_batch_line_items_upload_is_logged_with_the_batch_id(caplog): logger, requests = _s3_logger_answering(400, "cold") with caplog.at_level(logging.ERROR, logger="LiteLLM"): await logger.async_log_batch_line_items(_batch_line_items_payload()) assert len(requests) == 1 assert any( "bGl0ZWxsbV9wcm94eTtiYXRjaA" in record.getMessage() and "were not stored" in record.getMessage() for record in caplog.records ) @pytest.mark.asyncio async def test_a_batch_line_items_object_the_bucket_refuses_for_good_is_not_kept_in_memory(): logger, requests = _s3_logger_answering(400, "cold", b"InvalidArgument") await logger.async_log_batch_line_items(_batch_line_items_payload()) await logger.flush_queue() assert len(requests) == 1 assert logger._flush_dropped == {} @pytest.mark.asyncio async def test_batch_line_items_are_serialized_and_hashed_off_the_event_loop(monkeypatch): import threading import litellm.integrations.s3_v2 as s3_v2_module logger, requests = _s3_logger_answering(200, "cold/") serialize = s3_v2_module.batch_line_items_ndjson prepare = logger._prepare_put threads = {} def _serialize(payload): threads["serialize"] = threading.current_thread() return serialize(payload) def _prepare(element): threads["hash"] = threading.current_thread() return prepare(element) monkeypatch.setattr(s3_v2_module, "batch_line_items_ndjson", _serialize) monkeypatch.setattr(logger, "_prepare_put", _prepare) await logger.async_log_batch_line_items(_batch_line_items_payload()) assert len(requests) == 1 assert set(threads) == {"serialize", "hash"} assert threading.main_thread() not in threads.values()