test(rust): expand callback lifecycle coverage

litellm-rust/crates/python-bridge/src/routes/messages.rs

@@ -340,6 +340,7 @@ class Host:
         self.current = arguments
         self.asynchronous = asynchronous
         self.logger = arguments.get('litellm_logging_obj')
+        self.deployment_hooks_owned = self.logger is None
         self.state = None
         self.response = None
         self.error = None
@@ -353,7 +354,9 @@ class Host:
         self.streaming = self.logger.stream is True
 
     async def deployment_pre(self):
-        modified = await utils.async_pre_call_deployment_hook(self.current, 'amessages')
+        if not self.deployment_hooks_owned:
+            return
+        modified = await utils.async_pre_call_deployment_hook(self.current, 'anthropic_messages')
         if modified is not None:
             self.current = modified
         self.current['litellm_logging_obj'] = self.logger
@@ -370,7 +373,8 @@ class Host:
         self.end = datetime.now()
 
     async def deployment_success(self):
-        self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aanthropic_messages)
+        if self.deployment_hooks_owned:
+            self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aanthropic_messages)
         if self.streaming:
             self.response = retain_stream_response(
                 self.response,
@@ -380,7 +384,8 @@ class Host:
             )
 
     async def deployment_failure(self):
-        await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'amessages')
+        if self.deployment_hooks_owned:
+            await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'anthropic_messages')
 
     def terminal(self, action, value):
         if self.streaming:

tests/test_litellm_rust/test_messages_callbacks.py

@@ -8,6 +8,8 @@ import pytest
 
 import litellm
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+from litellm.types.utils import CallTypes
 from tests.test_litellm_rust.callback_recorder import RecordingLogger
 from tests.test_litellm_rust.contracts import (
     MESSAGES,
@@ -219,3 +221,132 @@ async def test_messages_stream_logs_success_after_exhaustion(messages_server: Re
     assert len(events) == 1
     assert "async_log_stream_event" not in recorder.names
     assert events[0].kwargs["complete_streaming_response"] is not None
+
+
+@pytest.mark.asyncio
+async def test_messages_compression_hook_replaces_messages_sent_to_provider(messages_server: RecordingServer) -> None:
+    compressed_messages: Final = [{"role": "user", "content": "Compressed context"}]
+    call_types: Final = []
+
+    class CompressMessages(CustomLogger):
+        async def async_pre_call_deployment_hook(self, kwargs, call_type):
+            call_types.append(call_type)
+            kwargs["messages"] = compressed_messages
+            return kwargs
+
+    litellm.callbacks.append(CompressMessages())
+
+    await call_messages(messages_server, [])
+
+    assert call_types == [CallTypes.anthropic_messages]
+    assert messages_server.requests[0].body["messages"] == compressed_messages
+
+
+@pytest.mark.asyncio
+async def test_messages_post_call_guardrail_replacement_reaches_caller_and_logging(
+    messages_server: RecordingServer,
+) -> None:
+    recorder: Final = RecordingLogger()
+    call_types: Final = []
+
+    class ReviewResponse(CustomLogger):
+        async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
+            call_types.append(call_type)
+            response["content"][0]["text"] = "Reviewed response"
+            return response
+
+    litellm.callbacks.append(ReviewResponse())
+
+    response: Final = await call_messages(messages_server, [recorder])
+    events: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert call_types == [CallTypes.anthropic_messages]
+    assert response["content"][0]["text"] == "Reviewed response"
+    assert events[0].response.choices[0].message.content == "Reviewed response"
+
+
+@pytest.mark.asyncio
+async def test_messages_logging_hook_replacement_reaches_later_loggers_only(messages_server: RecordingServer) -> None:
+    observations: Final = []
+
+    class RecordGuardrailVerdict(CustomLogger):
+        async def async_logging_hook(self, kwargs, result, call_type):
+            observations.append("guardrail")
+            return {**kwargs, "guardrail-verdict": "allowed"}, result
+
+    class ExportLog(CustomLogger):
+        async def async_logging_hook(self, kwargs, result, call_type):
+            return kwargs, result
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            observations.append(("export", kwargs["guardrail-verdict"]))
+
+    response: Final = await call_messages(messages_server, [RecordGuardrailVerdict(), ExportLog()])
+    await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
+
+    assert observations == ["guardrail", ("export", "allowed")]
+    assert response["content"][0]["text"] == "Hello from native Messages"
+
+
+@pytest.mark.asyncio
+async def test_messages_success_callback_failure_does_not_skip_later_loggers(
+    messages_server: RecordingServer,
+) -> None:
+    recorder: Final = RecordingLogger()
+
+    class UnavailableExporter(CustomLogger):
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            raise RuntimeError("exporter unavailable")
+
+    response: Final = await call_messages(messages_server, [UnavailableExporter(), recorder])
+    events: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert response["content"][0]["text"] == "Hello from native Messages"
+    assert len(events) == 1
+    assert "async_log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+async def test_messages_concurrent_calls_keep_callback_state_isolated(messages_server: RecordingServer) -> None:
+    messages_server.expected_requests = 4
+    tokens: Final = {f"messages-{index}": object() for index in range(4)}
+    terminal_state: Final = []
+
+    class CorrelateCallState(RecordingLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            kwargs["correlation-token"] = tokens[kwargs["litellm_call_id"]]
+            super().log_pre_api_call(model, messages, kwargs)
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            terminal_state.append((kwargs["litellm_call_id"], kwargs["correlation-token"]))
+            await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
+
+    correlate: Final = CorrelateCallState()
+
+    await asyncio.gather(*(call_messages(messages_server, [correlate], litellm_call_id=call_id) for call_id in tokens))
+    await correlate.wait_for_async("async_log_success_event", count=4)
+
+    assert len(terminal_state) == 4
+    assert all(token is tokens[call_id] for call_id, token in terminal_state)
+
+
+@pytest.mark.asyncio
+async def test_messages_cancelled_call_runs_no_terminal_callbacks(messages_server: RecordingServer) -> None:
+    messages_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE, delay=0.5)
+    recorder: Final = RecordingLogger()
+    task: Final = asyncio.create_task(call_messages(messages_server, [recorder]))
+
+    async with asyncio.timeout(10):
+        while not messages_server.requests:
+            await asyncio.sleep(0.01)
+    task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
+
+    assert recorder.names.count("log_pre_api_call") == 1
+    assert "log_success_event" not in recorder.names
+    assert "async_log_success_event" not in recorder.names
+    assert "log_failure_event" not in recorder.names
+    assert "async_log_failure_event" not in recorder.names

tests/test_litellm_rust/test_ocr_callbacks.py

@@ -11,7 +11,6 @@ from litellm.integrations.custom_logger import CustomLogger
 from tests.test_litellm_rust.callback_recorder import RecordingLogger
 from tests.test_litellm_rust.contracts import (
     OCR_DOCUMENT,
-    OCR_MODEL,
     OCR_RESPONSE,
     call_native_aocr,
     call_native_ocr,
@@ -247,3 +246,69 @@ async def test_pre_call_runs_in_callers_execution_context(ocr_server: RecordingS
     assert len(events) == 1
     assert events[0].loop is caller_loop
     assert events[0].thread is caller_thread
+
+
+@pytest.mark.asyncio
+async def test_ocr_failure_callbacks_receive_pre_call_state(ocr_server: RecordingServer) -> None:
+    ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
+    token: Final = object()
+    observed: Final = []
+
+    class TrackInFlightRequest(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            kwargs["request-token"] = token
+
+        def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            observed.append(("sync", kwargs["request-token"]))
+
+        async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            observed.append(("async", kwargs["request-token"]))
+
+    with pytest.raises(litellm.InternalServerError):
+        await call_aocr(ocr_server, [TrackInFlightRequest()])
+
+    assert [event for event, _ in observed] == ["sync", "async"]
+    assert all(observed_token is token for _, observed_token in observed)
+
+
+@pytest.mark.asyncio
+async def test_ocr_failure_callback_error_does_not_mask_provider_error_or_later_callbacks(
+    ocr_server: RecordingServer,
+) -> None:
+    ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
+    recorder: Final = RecordingLogger()
+
+    class UnavailableExporter(CustomLogger):
+        def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            raise RuntimeError("exporter unavailable")
+
+        async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            raise RuntimeError("exporter unavailable")
+
+    with pytest.raises(litellm.InternalServerError) as caught:
+        await call_aocr(ocr_server, [UnavailableExporter(), recorder])
+
+    sync_events: Final = tuple(event for event in recorder.events if event.name == "log_failure_event")
+    async_events: Final = tuple(event for event in recorder.events if event.name == "async_log_failure_event")
+    assert len(sync_events) == 1
+    assert len(async_events) == 1
+    assert sync_events[0].kwargs["exception"] is caught.value
+    assert async_events[0].kwargs["exception"] is caught.value
+    assert "async_log_success_event" not in recorder.names
+
+
+def test_ocr_duplicate_callback_registration_dispatches_once(ocr_server: RecordingServer) -> None:
+    recorder: Final = RecordingLogger()
+
+    call_ocr(
+        ocr_server,
+        [recorder, recorder],
+        success_callback=[recorder],
+        failure_callback=[recorder],
+    )
+    recorder.wait_for("log_success_event")
+
+    assert recorder.names.count("log_pre_api_call") == 1
+    assert recorder.names.count("logging_hook") == 1
+    assert recorder.names.count("log_success_event") == 1
+    assert "log_failure_event" not in recorder.names