diff --git a/tests/test_litellm_rust/chat/test_callback_mutation.py b/tests/test_litellm_rust/chat/test_callback_mutation.py index 5efe7036820..a51fbd7cfb3 100644 --- a/tests/test_litellm_rust/chat/test_callback_mutation.py +++ b/tests/test_litellm_rust/chat/test_callback_mutation.py @@ -265,25 +265,36 @@ async def test_native_chat_concurrent_calls_keep_callback_roots_isolated( recording_server: RecordingServer, ) -> None: import asyncio + import threading from litellm.integrations.custom_logger import CustomLogger + from tests.test_litellm_rust.support.callback_recorder import drain_logging from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE call_ids: Final = tuple(f"chat-{index}" for index in range(4)) recording_server.expected_requests = len(call_ids) - recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + accepted: Final = tuple(threading.Event() for _ in call_ids) + release: Final = threading.Event() + for signal in accepted: + recording_server.enqueue(ResponseSpec(body=MESSAGES_RESPONSE, accepted=signal, release_before_response=release)) + tokens: Final = {call_id: object() for call_id in call_ids} + terminal_state: Final = [] class Correlate(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): body: Final = kwargs["additional_args"]["complete_input_dict"] - body["messages"][0]["content"][0]["text"] = kwargs["litellm_call_id"] + body["messages"][0]["content"][0]["text"] = f"callback-{kwargs['litellm_call_id']}" + kwargs["correlation-token"] = tokens[kwargs["litellm_call_id"]] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + terminal_state.append((kwargs["litellm_call_id"], kwargs["correlation-token"])) logger: Final = Correlate() async def invoke(call_id: str) -> None: await litellm.acompletion( model="anthropic/claude-opus-5", - messages=[{"role": "user", "content": call_id}], + messages=[{"role": "user", "content": f"original-{call_id}"}], max_tokens=16, api_key="test-key", api_base=recording_server.base_url, @@ -292,10 +303,18 @@ async def test_native_chat_concurrent_calls_keep_callback_roots_isolated( num_retries=0, ) - await asyncio.gather(*(invoke(call_id) for call_id in call_ids)) + tasks: Final = tuple(asyncio.create_task(invoke(call_id)) for call_id in call_ids) + try: + assert all(await asyncio.gather(*(asyncio.to_thread(signal.wait, 3) for signal in accepted))) + finally: + release.set() + await asyncio.gather(*tasks) + await drain_logging() sent: Final = {request.body["messages"][0]["content"][0]["text"] for request in recording_server.requests} - assert sent == set(call_ids) + assert sent == {f"callback-{call_id}" for call_id in call_ids} + assert len(terminal_state) == len(call_ids) + assert all(token is tokens[call_id] for call_id, token in terminal_state) @pytest.mark.asyncio diff --git a/tests/test_litellm_rust/integrations/__init__.py b/tests/test_litellm_rust/integrations/__init__.py index 990fbd4fd23..c5f3569c9e0 100644 --- a/tests/test_litellm_rust/integrations/__init__.py +++ b/tests/test_litellm_rust/integrations/__init__.py @@ -5,6 +5,8 @@ from tests.test_litellm_rust.integrations.catalogue import ( GUARDRAIL_OBLIGATIONS, LOGGER_OBLIGATIONS, OSS_LOGGER_NAMES, + REQUIRED_GUARDRAIL_BEHAVIOR, + REQUIRED_LOGGER_BEHAVIOR, ) from tests.test_litellm_rust.integrations.harness import ( AsyncBoundaryLogger, @@ -50,6 +52,8 @@ __all__ = ( "OCR_ASYNC", "OCR_SYNC", "OSS_LOGGER_NAMES", + "REQUIRED_GUARDRAIL_BEHAVIOR", + "REQUIRED_LOGGER_BEHAVIOR", "AsyncBoundaryLogger", "GCSLiteralAIHarness", "GenericAPIExportHarness", diff --git a/tests/test_litellm_rust/integrations/catalogue.py b/tests/test_litellm_rust/integrations/catalogue.py index 344ff464015..3bcd0c9e7bf 100644 --- a/tests/test_litellm_rust/integrations/catalogue.py +++ b/tests/test_litellm_rust/integrations/catalogue.py @@ -127,6 +127,17 @@ class CoverageObligation: behavioral_cases: tuple[str, ...] +REQUIRED_LOGGER_BEHAVIOR: Final = frozenset({"generic_api", "gcs_bucket", "literalai", "prometheus", "opentelemetry"}) +REQUIRED_GUARDRAIL_BEHAVIOR: Final = frozenset( + { + "azure/text_moderations", + "crowdstrike_aidr", + "litellm_content_filter", + "microsoft_purview", + "rubrik", + } +) + LOGGER_BEHAVIORAL_CASES: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( { "generic_api": ("generic-api-success",), diff --git a/tests/test_litellm_rust/integrations/test_composition.py b/tests/test_litellm_rust/integrations/test_composition.py index 979db9f8b3b..5bc1b253786 100644 --- a/tests/test_litellm_rust/integrations/test_composition.py +++ b/tests/test_litellm_rust/integrations/test_composition.py @@ -1,8 +1,11 @@ import asyncio +import gc import json import threading +import weakref from collections.abc import Mapping from dataclasses import dataclass +from datetime import datetime from typing import Final import pytest @@ -13,18 +16,12 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger from litellm.integrations.prometheus import PrometheusLogger from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.guardrails.guardrail_registry import guardrail_initializer_registry from litellm.rust_bridge.provenance import has_rust_response_marker from litellm.types.guardrails import SupportedGuardrailIntegrations from litellm.types.utils import CallTypes -from tests.test_litellm_rust.support.callback_recorder import ( - LiveReferenceLogger, - RecordingLogger, - SecondaryLiveReferenceLogger, - drain_logging, -) from tests.test_litellm_rust.conftest import Backend, isolated_backend -from tests.test_litellm_rust.support.requests import MESSAGES_EVENTS from tests.test_litellm_rust.integrations import ( ASYNC_ROUTES, DISCOVERED_ONLY_GUARDRAIL_NAMES, @@ -37,6 +34,8 @@ from tests.test_litellm_rust.integrations import ( OCR_ASYNC, OCR_SYNC, OSS_LOGGER_NAMES, + REQUIRED_GUARDRAIL_BEHAVIOR, + REQUIRED_LOGGER_BEHAVIOR, AsyncBoundaryLogger, MutatingFailingLogger, OtelHarness, @@ -48,7 +47,20 @@ from tests.test_litellm_rust.integrations import ( route_id, wait_for_callback, ) +from tests.test_litellm_rust.support.callback_recorder import ( + LiveReferenceLogger, + RecordingLogger, + SecondaryLiveReferenceLogger, + drain_logging, +) from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service +from tests.test_litellm_rust.support.requests import ( + CHAT_MESSAGES, + CHAT_MODEL, + CHAT_RESPONSE, + MESSAGES, + MESSAGES_EVENTS, +) pytestmark = pytest.mark.requires_rust_extension @@ -104,6 +116,20 @@ def test_guardrail_catalogue_reconciles_enum_and_runtime_discovery() -> None: assert frozenset(GUARDRAIL_OBLIGATIONS) == initializer_names +def test_required_integrations_have_behavioral_cases() -> None: + logger_names: Final = frozenset( + name + for obligation in LOGGER_OBLIGATIONS.values() + if obligation.behavioral_cases + for name in obligation.registration_names + ) + guardrail_names: Final = frozenset( + name for name, obligation in GUARDRAIL_OBLIGATIONS.items() if obligation.behavioral_cases + ) + assert REQUIRED_LOGGER_BEHAVIOR <= logger_names + assert REQUIRED_GUARDRAIL_BEHAVIOR <= guardrail_names + + @pytest.mark.asyncio @pytest.mark.parametrize("route", ASYNC_ROUTES, ids=lambda route: f"otel-prometheus-generic-api-{route.name}-success") async def test_export_composition( @@ -162,6 +188,157 @@ async def test_terminal_callbacks_share_live_objects_within_one_run(backend: Bac second.release() +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_per_logger_redaction_preserves_the_shared_unredacted_payload(backend: Backend) -> None: + class Capture(CustomLogger): + def __init__(self, redact: bool) -> None: + super().__init__(turn_off_message_logging=redact) + self.kwargs: object = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.kwargs = kwargs + + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = provider_response(MESSAGES_ROUTE) + redacted: Final = Capture(True) + plain: Final = Capture(False) + await MESSAGES_ROUTE.invoke(provider, callbacks=[redacted, plain]) + await drain_logging() + + redacted_kwargs: Final = redacted.kwargs + plain_kwargs: Final = plain.kwargs + assert isinstance(redacted_kwargs, dict) + assert isinstance(plain_kwargs, dict) + redacted_payload: Final = redacted_kwargs["standard_logging_object"] + plain_payload: Final = plain_kwargs["standard_logging_object"] + assert redacted_kwargs is not plain_kwargs + assert redacted_payload is not plain_payload + assert redacted_payload["metadata"] is plain_payload["metadata"] + assert "redacted-by-litellm" in json.dumps(redacted_payload["messages"]) + assert MESSAGES_ROUTE.expected_text in json.dumps(plain_payload["response"]) + assert "redacted-by-litellm" not in json.dumps(plain_payload) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +@pytest.mark.parametrize("accepted", (True, False), ids=("accepted", "rejected")) +async def test_terminal_respects_proxy_deferred_completion_gate(backend: Backend, accepted: bool) -> None: + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = provider_response(MESSAGES_ROUTE) + recorder: Final = RecordingLogger() + logger: Final = Logging( + model=MESSAGES_ROUTE.provider_model, + messages=MESSAGES, + stream=False, + call_type=MESSAGES_ROUTE.call_type, + start_time=datetime.now(), + litellm_call_id="deferred-completion", + function_id="deferred-completion", + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + logger._defer_async_logging = True + response: Final = await MESSAGES_ROUTE.invoke(provider, litellm_logging_obj=logger) + await asyncio.sleep(0) + assert "async_log_success_event" not in recorder.names + + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( + logging_obj=logger, + exception_raised=not accepted, + ) + if accepted: + accepted_events: Final = await recorder.wait_for_async("async_log_success_event") + assert len(accepted_events) == 1 + assert "async_log_failure_event" not in recorder.names + else: + error: Final = RuntimeError("post-call guardrail rejected response") + await logger.async_failure_handler(error, str(error)) + rejected_events: Final = await recorder.wait_for_async("async_log_failure_event") + assert len(rejected_events) == 1 + assert "async_log_success_event" not in recorder.names + assert has_rust_response_marker(response) is (backend == "rust") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_suspended_terminal_callback_owns_payload_until_it_finishes(backend: Backend) -> None: + class Root: + pass + + class Suspended(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + assert kwargs["litellm_params"]["metadata"]["retained"] is not None + self.started.set() + await self.release.wait() + + async def invoke() -> weakref.ReferenceType[Root]: + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = provider_response(MESSAGES_ROUTE) + callback: Final = Suspended() + root = Root() + reference: Final = weakref.ref(root) + response: Final = await MESSAGES_ROUTE.invoke( + provider, + callbacks=[callback], + metadata={"retained": root}, + ) + del root + await asyncio.wait_for(callback.started.wait(), timeout=2) + gc.collect() + assert reference() is not None + callback.release.set() + await drain_logging() + del response + return reference + + reference: Final = await invoke() + await asyncio.sleep(0) + gc.collect() + assert reference() is None + + +async def observe_retry(backend: Backend) -> tuple[tuple[str, ...], bool]: + async with isolated_backend(backend): + with recording_service() as provider: + provider.expected_requests = 2 + provider.enqueue(ResponseSpec(body={"message": "retry this attempt"}, status=500)) + provider.default_response = ResponseSpec(body=CHAT_RESPONSE) + recorder: Final = RecordingLogger() + response: Final = await litellm.acompletion( + model=CHAT_MODEL, + messages=CHAT_MESSAGES, + api_key="test-key", + api_base=provider.base_url, + callbacks=[recorder], + num_retries=1, + ) + await drain_logging() + return recorder.names, has_rust_response_marker(response) + + +@pytest.mark.asyncio +async def test_retry_attempt_callback_sequence_matches_python() -> None: + python_names, python_rust_dispatch = await observe_retry("python") + rust_names, rust_rust_dispatch = await observe_retry("rust") + assert rust_names == python_names + assert rust_names.count("log_pre_api_call") == 2 + assert rust_names.count("async_log_failure_event") == 1 + assert rust_names.count("async_log_success_event") == 0 + assert python_rust_dispatch is False + assert rust_rust_dispatch is True + + @pytest.mark.asyncio async def test_failing_callback_preserves_prior_mutation_and_later_exporters() -> None: with recording_service() as provider, recording_service() as sink: