From 02572d815fc54e9feb40e554deb5c0a37bbf5cfd Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 8 Sep 2026 07:51:52 -0700 Subject: [PATCH] fix(rust): isolate callback tests and preserve messages logging lifecycle --- .../python-bridge/src/routes/messages.rs | 18 +- tests/_prometheus_helpers.py | 24 ++ tests/test_litellm_rust/callback_recorder.py | 8 + tests/test_litellm_rust/conftest.py | 123 +++++--- tests/test_litellm_rust/integrations.py | 201 +++++++++++++ tests/test_litellm_rust/test_conftest.py | 29 ++ tests/test_litellm_rust/test_integrations.py | 272 ++++++++++++++++++ .../test_messages_callbacks.py | 45 ++- 8 files changed, 662 insertions(+), 58 deletions(-) create mode 100644 tests/_prometheus_helpers.py create mode 100644 tests/test_litellm_rust/integrations.py create mode 100644 tests/test_litellm_rust/test_conftest.py create mode 100644 tests/test_litellm_rust/test_integrations.py diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index 07f15f1cebd..a6adfd66f9f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -199,11 +199,7 @@ fn prepare(py: Python<'_>, arguments: Py) -> PyResult> let message = PyDict::new(py); message.set_item("role", "user")?; message.set_item("content", serialized)?; - let messages = vec![message]; - logging - .bind(py) - .call_method1("update_messages", (&messages,))?; - kwargs.set_item("input", messages)?; + kwargs.set_item("input", vec![message])?; kwargs.set_item("api_key", "")?; kwargs.set_item("additional_args", additional)?; logging @@ -340,7 +336,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.lifecycle_owned = self.logger is None self.state = None self.response = None self.error = None @@ -354,7 +350,7 @@ class Host: self.streaming = self.logger.stream is True async def deployment_pre(self): - if not self.deployment_hooks_owned: + if not self.lifecycle_owned: return modified = await utils.async_pre_call_deployment_hook(self.current, 'anthropic_messages') if modified is not None: @@ -373,7 +369,7 @@ class Host: self.end = datetime.now() async def deployment_success(self): - if self.deployment_hooks_owned: + if self.lifecycle_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( @@ -384,11 +380,11 @@ class Host: ) async def deployment_failure(self): - if self.deployment_hooks_owned: + if self.lifecycle_owned: await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'anthropic_messages') def terminal(self, action, value): - if self.streaming: + if self.streaming or not self.lifecycle_owned: return None return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, None, value, self.start, self.end) @@ -398,7 +394,7 @@ class Host: def sync_failure(self): return self.terminal('sync_failure', self.error) def async_failure(self): return self.terminal('async_failure', self.error) def restore(self): - if not self.streaming: + if not self.streaming and self.lifecycle_owned: utils._restore_correlation_context_if_supported(self.logger) def advance(self, outcome, error=None): diff --git a/tests/_prometheus_helpers.py b/tests/_prometheus_helpers.py new file mode 100644 index 00000000000..f537253237f --- /dev/null +++ b/tests/_prometheus_helpers.py @@ -0,0 +1,24 @@ +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Final + +from prometheus_client import REGISTRY, CollectorRegistry + + +def clear_prometheus_registry() -> None: + for collector in tuple(REGISTRY._collector_to_names): # pyright: ignore[reportPrivateUsage] # prometheus_client has no public collector enumeration + REGISTRY.unregister(collector) + + +@contextmanager +def isolated_prometheus_registry(registry: CollectorRegistry = REGISTRY) -> Iterator[None]: + original: Final = tuple(registry._collector_to_names) # pyright: ignore[reportPrivateUsage] # prometheus_client has no public collector enumeration + for collector in original: + registry.unregister(collector) + try: + yield + finally: + for collector in tuple(registry._collector_to_names): # pyright: ignore[reportPrivateUsage] # remove collectors created inside the isolation scope + registry.unregister(collector) + for collector in original: + registry.register(collector) diff --git a/tests/test_litellm_rust/callback_recorder.py b/tests/test_litellm_rust/callback_recorder.py index d94a579864f..02ba541b015 100644 --- a/tests/test_litellm_rust/callback_recorder.py +++ b/tests/test_litellm_rust/callback_recorder.py @@ -6,6 +6,14 @@ from dataclasses import dataclass from typing import Final from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker + + +async def drain_logging(worker: LoggingWorker = GLOBAL_LOGGING_WORKER) -> None: + # The wrapper's scheduled helper enqueues without awaiting; let it run before joining the queue. + await asyncio.sleep(0) + worker.start() + await asyncio.wait_for(worker.flush(), timeout=10) @dataclass(frozen=True, slots=True) diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 42946dd64c4..08b3572b250 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -1,49 +1,106 @@ +import asyncio import os -from collections.abc import Iterator -from typing import Final +from collections.abc import AsyncIterator, Generator +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, contextmanager +from types import ModuleType +from typing import Final, cast import pytest +import pytest_asyncio import litellm -from litellm.rust_bridge.configuration import reset_rust_configuration +from litellm import utils +from litellm.litellm_core_utils import litellm_logging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivateUsage] # share the canonical env parsing with the module under test + _parse_env_bool, + reset_rust_configuration, +) +from tests._prometheus_helpers import isolated_prometheus_registry +from tests.test_litellm_rust.callback_recorder import drain_logging +from tests.test_litellm_rust.integrations import ( + otel, # noqa: F401 # pytest fixture export + prometheus, # noqa: F401 # pytest fixture export + provider, # noqa: F401 # pytest fixture export +) from tests.test_litellm_rust.recording_server import recording_server # noqa: F401 # pytest fixture export +CALLBACK_ATTRIBUTES: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", +) -@pytest.fixture(autouse=True) -def isolate_rust_state() -> Iterator[None]: - callback_attributes: Final = ( - "callbacks", - "input_callback", - "success_callback", - "failure_callback", - "_async_input_callback", - "_async_success_callback", - "_async_failure_callback", - ) - original_callbacks: Final = {attribute: list(getattr(litellm, attribute)) for attribute in callback_attributes} - original_cache: Final = litellm.cache - for attribute in callback_attributes: - getattr(litellm, attribute).clear() - litellm.cache = None # test-quality-ok: isolate the process-global cache from native extension tests + +def _list_attribute(container: ModuleType, attribute: str) -> list[object]: + value: Final = getattr(container, attribute) + if not isinstance(value, list): + raise AssertionError(f"{container.__name__}.{attribute} is not a list") + return cast(list[object], value) + + +@contextmanager +def _isolated_list(container: ModuleType, attribute: str) -> Generator[None]: + source: Final = _list_attribute(container, attribute) + original: Final = list(source) + source.clear() # mutable-ok: test isolation mutates global registries by design + try: + yield + finally: + source.clear() + source.extend(original) + setattr(container, attribute, source) + + +@contextmanager +def _rebound(container: ModuleType, attribute: str, value: object) -> Generator[None]: + original: Final[object] = getattr(container, attribute) + setattr(container, attribute, value) + try: + yield + finally: + setattr(container, attribute, original) + + +@contextmanager +def _rust_mode() -> Generator[None]: reset_rust_configuration() litellm.rust(True) - yield - for attribute, callbacks in original_callbacks.items(): - target = getattr(litellm, attribute) - target.clear() - target.extend(callbacks) - litellm.cache = original_cache # test-quality-ok: restore the process-global cache after native extension tests - reset_rust_configuration() + try: + yield + finally: + reset_rust_configuration() + + +@pytest_asyncio.fixture(autouse=True, loop_scope="function") +async def isolate_rust_state() -> AsyncIterator[ExitStack]: + with ExitStack() as stack: + for attribute in CALLBACK_ATTRIBUTES: + stack.enter_context(_isolated_list(litellm, attribute)) + stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # string-name callback cache has no public accessor + stack.enter_context(_rebound(utils, "callback_list", [])) # rebind-ok: legacy global registry mutated by set_callbacks + stack.enter_context(_rebound(litellm, "cache", None)) # test-quality-ok: isolate the process-global cache from native extension tests + stack.enter_context(isolated_prometheus_registry()) + stack.enter_context(_rust_mode()) + executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-test-logging") + stack.enter_context(_rebound(utils, "executor", executor)) + try: + yield stack + finally: + try: + await drain_logging() + finally: + await asyncio.to_thread(executor.shutdown, wait=True) + await GLOBAL_LOGGING_WORKER.stop() def pytest_collection_modifyitems(items): - rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - if not rust_enabled: + if not _parse_env_bool(os.environ.get("LITELLM_RUST")): skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") for item in items: item.add_marker(skip) diff --git a/tests/test_litellm_rust/integrations.py b/tests/test_litellm_rust/integrations.py new file mode 100644 index 00000000000..f8ef0b93205 --- /dev/null +++ b/tests/test_litellm_rust/integrations.py @@ -0,0 +1,201 @@ +import asyncio +import time +from collections.abc import Awaitable, Callable +from contextlib import ExitStack +from dataclasses import dataclass +from typing import Final, Literal + +import pytest +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from prometheus_client import REGISTRY + +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.opentelemetry import LITELLM_REQUEST_SPAN_NAME, OpenTelemetry, OpenTelemetryConfig +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs +from tests.test_litellm_rust.contracts import ( + MESSAGES, + MESSAGES_MODEL, + MESSAGES_RESPONSE, + OCR_RESPONSE, + call_native_aocr, + call_native_ocr, +) +from tests.test_litellm_rust.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.callback_recorder import drain_logging + +RouteName = Literal["ocr-sync", "ocr-async", "messages", "messages-stream"] +GuardrailObservation = tuple[Literal["request", "response"], tuple[str, ...]] + + +@dataclass(frozen=True, slots=True) +class Route: + name: RouteName + call_type: str + provider_model: str + provider_response: dict[str, object] + response_text: str + provider: str + expected_cost: float + logging_only_scan: GuardrailObservation + fires_async_hooks: bool + + async def invoke(self, server: RecordingServer, **kwargs: object) -> object: + match self.name: + case "ocr-sync": + return await asyncio.to_thread(call_native_ocr, server, **kwargs) + case "ocr-async": + return await call_native_aocr(server, **kwargs) + case "messages": + return await _call_messages(server, **kwargs) + case "messages-stream": + stream: Final = await _call_messages(server, stream=True, **kwargs) + return [chunk async for chunk in stream] + + +async def _call_messages(server: RecordingServer, **kwargs: object) -> object: + return await litellm.anthropic.messages.acreate( + model=MESSAGES_MODEL, + messages=MESSAGES, + max_tokens=64, + api_key="test-key", + api_base=server.base_url, + **kwargs, + ) + + +MESSAGES_COST: Final = 5 * 3e-06 + 4 * 1.5e-05 +OCR_COST: Final = 0.004 + +OCR_SYNC: Final = Route( + name="ocr-sync", + call_type="ocr", + provider_model="mistral-ocr-latest", + provider_response=OCR_RESPONSE, + response_text="native OCR response", + provider="mistral", + expected_cost=OCR_COST, + logging_only_scan=("response", ("native OCR response",)), + fires_async_hooks=False, +) +OCR_ASYNC: Final = Route( + name="ocr-async", + call_type="aocr", + provider_model="mistral-ocr-latest", + provider_response=OCR_RESPONSE, + response_text="native OCR response", + provider="mistral", + expected_cost=OCR_COST, + logging_only_scan=("response", ("native OCR response",)), + fires_async_hooks=True, +) +MESSAGES_ROUTE: Final = Route( + name="messages", + call_type="anthropic_messages", + provider_model="claude-sonnet-4-5-20250929", + provider_response=MESSAGES_RESPONSE, + response_text="Hello from native Messages", + provider="anthropic", + expected_cost=MESSAGES_COST, + logging_only_scan=("request", ("Hello",)), + fires_async_hooks=True, +) +MESSAGES_STREAM: Final = Route( + name="messages-stream", + call_type="anthropic_messages", + provider_model="claude-sonnet-4-5-20250929", + provider_response=MESSAGES_RESPONSE, + response_text="Hello from native Messages", + provider="anthropic", + expected_cost=MESSAGES_COST, + logging_only_scan=("request", ("Hello",)), + fires_async_hooks=True, +) +ALL_ROUTES: Final = (OCR_SYNC, OCR_ASYNC, MESSAGES_ROUTE, MESSAGES_STREAM) +ASYNC_ROUTES: Final = tuple(route for route in ALL_ROUTES if route.fires_async_hooks) +NON_STREAM_ASYNC_ROUTES: Final = (OCR_ASYNC, MESSAGES_ROUTE) + + +def route_id(route: Route) -> str: + return route.name + + +@pytest.fixture +def provider(recording_server: RecordingServer, route: Route) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=route.provider_response) + return recording_server + + +@dataclass(frozen=True, slots=True) +class OtelHarness: + logger: OpenTelemetry + exporter: InMemorySpanExporter + + def spans(self, name: str = LITELLM_REQUEST_SPAN_NAME) -> tuple[ReadableSpan, ...]: + return tuple(span for span in self.exporter.get_finished_spans() if span.name == name) + + async def wait_for_spans( + self, name: str = LITELLM_REQUEST_SPAN_NAME, count: int = 1, timeout: float = 10 + ) -> tuple[ReadableSpan, ...]: + deadline: Final = time.monotonic() + timeout + while len(self.spans(name)) < count: + if time.monotonic() >= deadline: + raise TimeoutError(f"Timed out waiting for {count} {name} spans; saw {self.spans(name)}") + await asyncio.sleep(0.01) + await drain_logging() + return self.spans(name) + + +@pytest.fixture +def otel(isolate_rust_state: ExitStack) -> OtelHarness: + exporter: Final = InMemorySpanExporter() + tracer_provider: Final = TracerProvider() + isolate_rust_state.callback(tracer_provider.shutdown) + tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + logger: Final = OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter), tracer_provider=tracer_provider) + return OtelHarness(logger=logger, exporter=exporter) + + +@pytest.fixture +def prometheus() -> PrometheusLogger: + return PrometheusLogger() + + +def metric_value(name: str, **labels: str) -> float: + for metric in REGISTRY.collect(): + for sample in metric.samples: + if sample.name == name and all(sample.labels.get(key) == value for key, value in labels.items()): + return sample.value + return 0.0 + + +class RecordingGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str = "rust-review", fail_with: Exception | None = None) -> None: + super().__init__( + guardrail_name=guardrail_name, + event_hook=GuardrailEventHooks.logging_only, + default_on=True, + ) + self.observations: list[GuardrailObservation] = [] + self._fail_with = fail_with + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None) -> GenericGuardrailAPIInputs: + self.observations.append((input_type, tuple(inputs.get("texts") or ()))) + if self._fail_with is not None: + raise self._fail_with + return inputs + + +class ReviewGuardrail(CustomGuardrail): + def __init__(self, review: Callable[[object], Awaitable[object]]) -> None: + super().__init__(guardrail_name="rust-review", event_hook=GuardrailEventHooks.post_call, default_on=True) + self._review = review + self.call_types: list[object] = [] + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self.call_types.append(call_type) + return await self._review(response) diff --git a/tests/test_litellm_rust/test_conftest.py b/tests/test_litellm_rust/test_conftest.py new file mode 100644 index 00000000000..1ea9a929c7b --- /dev/null +++ b/tests/test_litellm_rust/test_conftest.py @@ -0,0 +1,29 @@ +from types import ModuleType +from typing import Final + +import pytest + +from tests.test_litellm_rust.conftest import _isolated_list # pyright: ignore[reportPrivateUsage] # exercise fixture restoration without mutating SDK state + +pytestmark = pytest.mark.requires_rust_extension + + +def test_callback_list_restores_identity_after_rebinding_and_failure() -> None: + container: Final = ModuleType("callback_registry") + callback: Final = object() + original: Final = [callback] + setattr(container, "callbacks", original) + + def failing_test() -> None: + with _isolated_list(container, "callbacks"): + assert getattr(container, "callbacks") is original + assert original == [] + original.append(object()) + setattr(container, "callbacks", [object()]) + raise RuntimeError("test failed") + + with pytest.raises(RuntimeError, match="test failed"): + failing_test() + + assert getattr(container, "callbacks") is original + assert original == [callback] diff --git a/tests/test_litellm_rust/test_integrations.py b/tests/test_litellm_rust/test_integrations.py new file mode 100644 index 00000000000..e2afcb57d36 --- /dev/null +++ b/tests/test_litellm_rust/test_integrations.py @@ -0,0 +1,272 @@ +import json +from typing import Final + +import pytest +from opentelemetry.trace import StatusCode +from prometheus_client import CollectorRegistry, Counter + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils import litellm_logging +from litellm.types.utils import CallTypes +from tests._prometheus_helpers import isolated_prometheus_registry +from tests.test_litellm_rust.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.integrations import ( + ALL_ROUTES, + ASYNC_ROUTES, + MESSAGES_ROUTE, + NON_STREAM_ASYNC_ROUTES, + OCR_ASYNC, + OCR_SYNC, + OtelHarness, + RecordingGuardrail, + ReviewGuardrail, + Route, + metric_value, + route_id, +) +from tests.test_litellm_rust.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + +FAILURE_RESPONSE: Final = ResponseSpec(body={"message": "provider unavailable"}, status=500) + + +def test_prometheus_registry_restores_collectors_after_failure() -> None: + registry: Final = CollectorRegistry() + original: Final = Counter("original", "Original collector", registry=registry) + original.inc(2) + + def failing_test() -> None: + with isolated_prometheus_registry(registry): + assert registry.get_sample_value("original_total") is None + Counter("original", "Temporary replacement", registry=registry).inc(7) + Counter("temporary", "Temporary collector", registry=registry).inc() + raise RuntimeError("test failed") + + with pytest.raises(RuntimeError, match="test failed"): + failing_test() + + assert registry.get_sample_value("original_total") == 2 + assert registry.get_sample_value("temporary_total") is None + registry.unregister(original) + assert registry.get_sample_value("original_total") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ALL_ROUTES, ids=route_id) +async def test_otel_emits_one_request_span_on_success( + route: Route, provider: RecordingServer, otel: OtelHarness +) -> None: + await route.invoke(provider, callbacks=[otel.logger]) + + spans: Final = await otel.wait_for_spans() + assert len(spans) == 1 + span: Final = spans[0] + assert span.status.status_code is StatusCode.OK + assert span.attributes["llm.request.type"] == route.call_type + assert span.attributes["gen_ai.request.model"] == route.provider_model + assert json.loads(span.attributes["hidden_params"])["response_cost"] == pytest.approx(route.expected_cost) + + +@pytest.mark.asyncio +async def test_otel_stream_span_appears_only_after_exhaustion( + otel: OtelHarness, recording_server: RecordingServer +) -> None: + recording_server.default_response = ResponseSpec(body=MESSAGES_ROUTE.provider_response) + + stream: Final = await MESSAGES_ROUTE.invoke(recording_server, callbacks=[otel.logger], stream=True) + await drain_logging() + assert otel.spans() == () + + chunks: Final = [chunk async for chunk in stream] + + assert chunks + assert len(await otel.wait_for_spans()) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", (OCR_SYNC, OCR_ASYNC), ids=route_id) +async def test_otel_emits_one_error_span_on_provider_failure( + route: Route, provider: RecordingServer, otel: OtelHarness +) -> None: + provider.enqueue(FAILURE_RESPONSE) + + with pytest.raises(litellm.InternalServerError): + await route.invoke(provider, callbacks=[otel.logger]) + + spans: Final = await otel.wait_for_spans() + assert len(spans) == 1 + assert spans[0].status.status_code is StatusCode.ERROR + exception_events: Final = [event for event in spans[0].events if event.name == "exception"] + assert len(exception_events) == 1 + assert "InternalServerError" in exception_events[0].attributes["exception.type"] + assert spans[0].attributes["llm.request.type"] == route.call_type + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ALL_ROUTES, ids=route_id) +async def test_otel_and_custom_logger_observe_same_standard_logging_object( + route: Route, provider: RecordingServer, otel: OtelHarness +) -> None: + recorder: Final = RecordingLogger() + + await route.invoke(provider, callbacks=[otel.logger, recorder]) + + events: Final = ( + await recorder.wait_for_async("async_log_success_event") + if route.fires_async_hooks + else recorder.wait_for("log_success_event") + ) + payload: Final = events[0].kwargs["standard_logging_object"] + span: Final = (await otel.wait_for_spans())[0] + assert payload["call_type"] == route.call_type + assert payload["model"] == route.provider_model + assert json.loads(span.attributes["hidden_params"]) == payload["hidden_params"] + assert span.attributes["llm.request.type"] == payload["call_type"] + assert span.attributes["litellm.provider.model"] == payload["model"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=route_id) +async def test_prometheus_counts_one_successful_request( + route: Route, provider: RecordingServer, prometheus: PrometheusLogger +) -> None: + before: Final = metric_value("litellm_requests_metric_total", model=route.provider_model) + recorder: Final = RecordingLogger() + + await route.invoke(provider, callbacks=[prometheus, recorder]) + await recorder.wait_for_async("async_log_success_event") + + assert metric_value("litellm_requests_metric_total", model=route.provider_model) == before + 1 + assert metric_value("litellm_llm_api_failed_requests_metric_total", model=route.provider_model) == 0 + + +@pytest.mark.asyncio +async def test_prometheus_counts_tokens_from_messages_usage( + recording_server: RecordingServer, prometheus: PrometheusLogger +) -> None: + recording_server.default_response = ResponseSpec(body=MESSAGES_ROUTE.provider_response) + recorder: Final = RecordingLogger() + + await MESSAGES_ROUTE.invoke(recording_server, callbacks=[prometheus, recorder]) + await recorder.wait_for_async("async_log_success_event") + + assert metric_value("litellm_input_tokens_metric_total", model=MESSAGES_ROUTE.provider_model) == 5 + assert metric_value("litellm_output_tokens_metric_total", model=MESSAGES_ROUTE.provider_model) == 4 + + +@pytest.mark.asyncio +async def test_prometheus_counts_one_failed_request( + otel: OtelHarness, prometheus: PrometheusLogger, recording_server: RecordingServer +) -> None: + recording_server.enqueue(FAILURE_RESPONSE) + + with pytest.raises(litellm.InternalServerError): + await OCR_ASYNC.invoke(recording_server, callbacks=[prometheus, otel.logger]) + + assert len(await otel.wait_for_spans()) == 1 + assert metric_value("litellm_llm_api_failed_requests_metric_total", model=OCR_ASYNC.provider_model) == 1 + assert metric_value("litellm_requests_metric_total", model=OCR_ASYNC.provider_model) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id) +async def test_prometheus_by_string_name_is_initialized_once(route: Route, provider: RecordingServer) -> None: + provider.expected_requests = 2 + litellm.success_callback = ["prometheus"] # test-quality-ok: tests public string-name registration; isolate_rust_state restores this registry + recorder: Final = RecordingLogger() + + await route.invoke(provider, callbacks=[recorder]) + await route.invoke(provider, callbacks=[recorder]) + await recorder.wait_for_async("async_log_success_event", count=2) + + instances: Final = [cb for cb in litellm_logging._in_memory_loggers if isinstance(cb, PrometheusLogger)] # pyright: ignore[reportPrivateUsage] # string-name cache has no public accessor + assert len(instances) == 1 + assert metric_value("litellm_requests_metric_total", model=route.provider_model) == 2 + assert "prometheus" not in litellm.success_callback + assert instances[0] in litellm._async_success_callback # pyright: ignore[reportPrivateUsage] # callback registry has no public accessor + + +@pytest.mark.asyncio +async def test_sync_ocr_reaches_sync_hooks_only( + recording_server: RecordingServer, otel: OtelHarness, prometheus: PrometheusLogger +) -> None: + recording_server.default_response = ResponseSpec(body=OCR_SYNC.provider_response) + + await OCR_SYNC.invoke(recording_server, callbacks=[otel.logger, prometheus]) + + assert len(await otel.wait_for_spans()) == 1 + assert metric_value("litellm_requests_metric_total", model=OCR_SYNC.provider_model) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id) +async def test_logging_only_guardrail_verdict_reaches_otel_and_custom_logger( + route: Route, provider: RecordingServer, otel: OtelHarness +) -> None: + guardrail: Final = RecordingGuardrail() + recorder: Final = RecordingLogger() + + await route.invoke(provider, callbacks=[guardrail, otel.logger, recorder]) + + payload: Final = (await recorder.wait_for_async("async_log_success_event"))[0].kwargs["standard_logging_object"] + assert guardrail.observations == [route.logging_only_scan] + verdicts: Final = payload["guardrail_information"] + assert [verdict["guardrail_name"] for verdict in verdicts] == ["rust-review"] + assert verdicts[0]["guardrail_status"] == "success" + guardrail_spans: Final = await otel.wait_for_spans("guardrail") + assert len(guardrail_spans) == 1 + assert guardrail_spans[0].attributes["guardrail_name"] == "rust-review" + assert guardrail_spans[0].attributes["guardrail_status"] == "success" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id) +async def test_logging_only_guardrail_failure_does_not_block_loggers( + route: Route, provider: RecordingServer, otel: OtelHarness, prometheus: PrometheusLogger +) -> None: + guardrail: Final = RecordingGuardrail(fail_with=RuntimeError("review service unavailable")) + recorder: Final = RecordingLogger() + + response: Final = await route.invoke(provider, callbacks=[guardrail, otel.logger, prometheus, recorder]) + + assert response is not None + assert len(await otel.wait_for_spans()) == 1 + assert metric_value("litellm_requests_metric_total", model=route.provider_model) == 1 + payload: Final = (await recorder.wait_for_async("async_log_success_event"))[0].kwargs["standard_logging_object"] + assert payload["guardrail_information"][0]["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id) +async def test_post_call_guardrail_replacement_is_what_loggers_see( + route: Route, provider: RecordingServer, otel: OtelHarness +) -> None: + async def review(response: object) -> object: + match route.name: + case "ocr-async": + return response.model_copy( + update={"pages": [response.pages[0].model_copy(update={"markdown": "Reviewed OCR"})]} + ) + case _: + return {**response, "content": [{"type": "text", "text": "Reviewed Messages"}]} + + guardrail: Final = ReviewGuardrail(review) + recorder: Final = RecordingLogger() + litellm.callbacks.append(guardrail) + + response: Final = await route.invoke(provider, callbacks=[otel.logger, recorder], guardrails=["rust-review"]) + + assert guardrail.call_types == [CallTypes(route.call_type)] + event: Final = (await recorder.wait_for_async("async_log_success_event"))[0] + span: Final = (await otel.wait_for_spans())[0] + match route.name: + case "ocr-async": + assert response.pages[0].markdown == "Reviewed OCR" + assert event.response.pages[0].markdown == "Reviewed OCR" + case _: + assert response["content"][0]["text"] == "Reviewed Messages" + assert event.response.choices[0].message.content == "Reviewed Messages" + assert "Reviewed Messages" in json.dumps(dict(span.attributes)) + assert "guardrails" not in provider.requests[0].body diff --git a/tests/test_litellm_rust/test_messages_callbacks.py b/tests/test_litellm_rust/test_messages_callbacks.py index e21ad4661db..d6644e27fdf 100644 --- a/tests/test_litellm_rust/test_messages_callbacks.py +++ b/tests/test_litellm_rust/test_messages_callbacks.py @@ -10,7 +10,7 @@ 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.callback_recorder import RecordingLogger, drain_logging from tests.test_litellm_rust.contracts import ( MESSAGES, MESSAGES_MODEL, @@ -54,18 +54,7 @@ async def test_messages_pre_call_receives_expected_provider_request(messages_ser assert len(observations) == 1 model, messages, additional_args = observations[0] assert model == "claude-sonnet-4-5-20250929" - assert messages == [ - { - "role": "user", - "content": json.dumps( - { - "model": "claude-sonnet-4-5-20250929", - "messages": MESSAGES, - "max_tokens": 64, - } - ), - } - ] + assert messages == MESSAGES assert additional_args["api_base"] == f"{messages_server.base_url}/v1/messages" assert additional_args["complete_input_dict"] == { "model": "claude-sonnet-4-5-20250929", @@ -155,6 +144,32 @@ async def test_messages_callbacks_run_once(messages_server: RecordingServer) -> assert "async_log_failure_event" not in recorder.names +@pytest.mark.asyncio +async def test_messages_logging_drain_waits_for_suspended_callback(messages_server: RecordingServer) -> None: + started: Final = asyncio.Event() + release: Final = asyncio.Event() + finished: Final = asyncio.Event() + + class SuspendedLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + started.set() + await release.wait() + finished.set() + + await call_messages(messages_server, [SuspendedLogger()]) + draining: Final = asyncio.create_task(drain_logging()) + try: + await asyncio.wait_for(started.wait(), timeout=10) + await asyncio.sleep(0) + assert not draining.done() + assert not finished.is_set() + finally: + release.set() + await asyncio.wait_for(draining, timeout=10) + + assert finished.is_set() + + @pytest.mark.asyncio @pytest.mark.xfail(strict=True, reason="accepted native Messages errors are replayed through the Python transport") async def test_messages_failure_callbacks_receive_original_provider_error(messages_server: RecordingServer) -> None: @@ -268,6 +283,7 @@ async def test_messages_post_call_guardrail_replacement_reaches_caller_and_loggi @pytest.mark.asyncio async def test_messages_logging_hook_replacement_reaches_later_loggers_only(messages_server: RecordingServer) -> None: observations: Final = [] + exported: Final = asyncio.Event() class RecordGuardrailVerdict(CustomLogger): async def async_logging_hook(self, kwargs, result, call_type): @@ -280,9 +296,10 @@ async def test_messages_logging_hook_replacement_reaches_later_loggers_only(mess async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): observations.append(("export", kwargs["guardrail-verdict"])) + exported.set() response: Final = await call_messages(messages_server, [RecordGuardrailVerdict(), ExportLog()]) - await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + await asyncio.wait_for(exported.wait(), timeout=10) assert observations == ["guardrail", ("export", "allowed")] assert response["content"][0]["text"] == "Hello from native Messages"