From 02199b670bbbda106e6a35afe9263ca6bfb91f6a Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Tue, 20 Oct 2026 08:21:49 -0700 Subject: [PATCH] fix(proxy): share the billing logger across callback registries (#47510) * fix(logging): deduplicate equivalent completed success callbacks * fix(proxy): share the billing logger across callback registries * test(proxy): use fixed timestamps for billing regressions --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 5 +- .../test_litellm_logging.py | 72 +++++++++++++++---- .../unit/proxy/proxy_server/test_lifecycle.py | 72 +++++++++++++------ 3 files changed, 112 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f87b1edc2e5..2c6d2e2736a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3118,8 +3118,9 @@ def cost_tracking(): from litellm.integrations.shadow_eval_logger import ShadowEvalLogger spend_event_producer = build_spend_event_producer(CollectorSettings(), fallback=run_spend_event) - litellm.logging_callback_manager.add_litellm_callback(ProxyDBLogger(spend_event_producer)) - litellm.logging_callback_manager.add_litellm_async_success_callback(ProxyDBLogger(spend_event_producer)) + db_logger: Final = ProxyDBLogger(spend_event_producer) + litellm.logging_callback_manager.add_litellm_callback(db_logger) + litellm.logging_callback_manager.add_litellm_async_success_callback(db_logger) litellm.logging_callback_manager.add_litellm_callback(ShadowEvalLogger()) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 7aac43298d4..981c43a9792 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -2462,7 +2462,41 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests( @pytest.mark.asyncio -@pytest.mark.parametrize("registration", ["sync", "async", "overlapping"]) +@pytest.mark.parametrize("different_headers", [False, True]) +async def test_completed_success_preserves_generic_logger_configuration( + monkeypatch: pytest.MonkeyPatch, different_headers: bool +) -> None: + from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger + + first: Final = GenericAPILogger( + endpoint="https://logs.example.test/events", + headers={"Authorization": "Bearer first"}, + event_types=["llm_api_success"], + ) + second: Final = GenericAPILogger( + endpoint="https://logs.example.test/events", + headers={"Authorization": "Bearer second" if different_headers else "Bearer first"}, + event_types=["llm_api_success"] if different_headers else ["llm_api_failure"], + ) + monkeypatch.setattr(litellm, "success_callback", [first]) + monkeypatch.setattr(litellm, "_async_success_callback", [second]) + monkeypatch.setattr(litellm, "generic_api_use_v1", False) + logging_obj: Final = _async_request_logging_obj("distinct-generic-loggers", []) + response: Final = ModelResponse(id="chatcmpl-distinct-generic-loggers") + now: Final = datetime.datetime(2026, 1, 1) + + await logging_obj.async_success_handler(response, now, now) + + assert len(first.log_queue) == 1 + assert first.log_queue[0]["id"] == response.id + if different_headers: + assert second.log_queue == first.log_queue + else: + assert second.log_queue == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("registration", ["sync", "async", "overlapping", "equivalent", "distinct"]) @pytest.mark.parametrize("kind", ["sync", "async", "dual"]) @pytest.mark.parametrize("compatibility_first", [False, True]) async def test_completed_success_selects_one_implemented_hook( @@ -2517,9 +2551,15 @@ async def test_completed_success_selects_one_implemented_hook( ) -> None: record("async", kwargs, response_obj) - callback: Final = {"sync": SyncRecorder, "async": AsyncRecorder, "dual": DualRecorder}[kind]() + recorder: Final = {"sync": SyncRecorder, "async": AsyncRecorder, "dual": DualRecorder}[kind] + callback: Final = recorder() + async_callback: Final = ( + recorder(turn_off_message_logging=registration == "distinct") + if registration in ("equivalent", "distinct") + else callback + ) monkeypatch.setattr(litellm, "success_callback", [callback] if registration != "async" else []) - monkeypatch.setattr(litellm, "_async_success_callback", [callback] if registration != "sync" else []) + monkeypatch.setattr(litellm, "_async_success_callback", [async_callback] if registration != "sync" else []) logging_obj: Final = _async_request_logging_obj("completed-hook", []) response: Final = ModelResponse(id="chatcmpl-completed-hook") now: Final = datetime.datetime.now() @@ -2531,20 +2571,26 @@ async def test_completed_success_selects_one_implemented_hook( await logging_obj.async_success_handler(response, now, now) logging_obj.handle_sync_success_callbacks_for_async_calls(response, now, now) - assert events.qsize() == 1 - hook, response_id, cost, thread_id = events.get_nowait() - assert hook == ("sync" if kind == "sync" else "async") - assert response_id == response.id - assert cost == logging_obj.model_call_details["response_cost"] - assert (thread_id == loop_thread) is (kind != "sync") + expected_count: Final = 2 if registration in ("equivalent", "distinct") else 1 + assert events.qsize() == expected_count + for _ in range(expected_count): + hook, response_id, cost, thread_id = events.get_nowait() + assert hook == ("sync" if kind == "sync" else "async") + assert response_id == response.id + assert cost == logging_obj.model_call_details["response_cost"] + assert (thread_id == loop_thread) is (kind != "sync") -@pytest.mark.parametrize("for_async_request", [False, True]) +@pytest.mark.parametrize("for_async_request,split_registries", [(False, False), (True, False), (True, True)]) @pytest.mark.parametrize("private_first", [False, True]) @pytest.mark.parametrize("response_kind", ["responses", "completion"]) def test_sync_success_redacts_only_opted_out_loggers( - monkeypatch: pytest.MonkeyPatch, for_async_request: bool, private_first: bool, response_kind: str + monkeypatch: pytest.MonkeyPatch, + for_async_request: bool, + split_registries: bool, + private_first: bool, + response_kind: str, ) -> None: from queue import SimpleQueue @@ -2563,8 +2609,8 @@ def test_sync_success_redacts_only_opted_out_loggers( private_logger: Final = PayloadRecorder(turn_off_message_logging=True) public_logger: Final = PayloadRecorder() callbacks: Final = [private_logger, public_logger] if private_first else [public_logger, private_logger] - monkeypatch.setattr(litellm, "success_callback", callbacks) - monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", callbacks[:1] if split_registries else callbacks) + monkeypatch.setattr(litellm, "_async_success_callback", callbacks[1:] if split_registries else []) monkeypatch.setattr(litellm, "turn_off_message_logging", False) logging_obj: Final = LitellmLogging( model="openai/gpt-6.1-sol", diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 3914c337f59..11a44a0ec16 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -697,33 +697,61 @@ def test_load_from_azure_key_vault_missing_uri_failure_is_swallowed(monkeypatch) # --------------------------------------------------------------------------- -def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set(monkeypatch): +@pytest.mark.asyncio +@pytest.mark.parametrize("setup_count", [1, 2]) +async def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set( + monkeypatch: pytest.MonkeyPatch, setup_count: int +) -> None: + from datetime import datetime + import litellm from litellm.integrations.shadow_eval_logger import ShadowEvalLogger + from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, decode_spend_event + from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer + from litellm.types.utils import ModelResponse, Usage + from litellm.utils import Rules, function_setup - fake_prisma = MagicMock() - monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) - monkeypatch.setattr(litellm, "callbacks", [], raising=False) - monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False) + producer: Final = AsyncMock(spec=SpendEventProducer) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(ps, "spend_event_producer", None) + monkeypatch.setattr(ps, "build_spend_event_producer", lambda settings, fallback: producer) + for name in ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, name, []) - before_callbacks = len(litellm.callbacks) - before_async = len(litellm._async_success_callback) + for _ in range(setup_count): + cost_tracking() - cost_tracking() - cost_tracking() - - observed = { - "added_to_callbacks": len(litellm.callbacks) - before_callbacks, - "added_to_async_success": len(litellm._async_success_callback) - before_async, - "shadow_eval_loggers": sum(isinstance(cb, ShadowEvalLogger) for cb in litellm.callbacks), - "prisma_was_set": True, - } - assert normalize(observed) == { - "added_to_callbacks": 2, - "added_to_async_success": 1, - "shadow_eval_loggers": 1, - "prisma_was_set": True, - } + assert len(litellm.callbacks) == 2 + assert sum(isinstance(cb, ShadowEvalLogger) for cb in litellm.callbacks) == 1 + assert len(litellm._async_success_callback) == 1 + now: Final = datetime(2026, 1, 1) + logging_obj, _ = function_setup( + "acompletion", + Rules(), + now, + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + litellm_call_id="single-accounting-event", + ) + logging_obj.update_environment_variables(litellm_params={"acompletion": True, "metadata": {}}, optional_params={}) + response: Final = ModelResponse( + id="chatcmpl-single-accounting-event", + model="gpt-4o-mini", + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + await logging_obj.async_success_handler(response, now, now) + producer.publish.assert_awaited_once() + event: Final = decode_spend_event(producer.publish.call_args.args[0]) + assert not isinstance(event, SpendEventDecodeError) + assert event.response_cost == logging_obj.model_call_details["response_cost"] + assert event.response_cost > 0 def test_cost_tracking_no_op_when_prisma_missing(monkeypatch):