wip

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

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",

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",),

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: