From 7dd28d2e37a889e6489538ae393838003e39e7e7 Mon Sep 17 00:00:00 2001 From: nate Date: Sat, 10 Oct 2026 01:13:31 +0000 Subject: [PATCH] fix(proxy): keep callbacks registered before config load and warn when telemetry cannot register Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/callback_utils.py | 17 ++-- litellm/proxy/proxy_server.py | 4 +- litellm/proxy/telemetry/runtime.py | 11 +++ .../proxy/common_utils/test_callback_utils.py | 80 ++++++++----------- tests/unit/proxy/telemetry/test_runtime.py | 33 +++++++- 5 files changed, 88 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 4fa7709e5c7..5d89c73933d 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -404,15 +404,14 @@ def initialize_callbacks_on_proxy( PrometheusLogger._mount_metrics_endpoint() else: - litellm.callbacks = [ - _loaded_callback_or_raise( - entry=value, - loaded=get_instance_fn( - value=value, - config_file_path=config_file_path, - ), - ) - ] + loaded_callback: Final = _loaded_callback_or_raise( + entry=value, + loaded=get_instance_fn( + value=value, + config_file_path=config_file_path, + ), + ) + litellm.callbacks = [*litellm.callbacks, loaded_callback] verbose_proxy_logger.debug("%s Initialized Callbacks - %s %s", blue_color_code, litellm.callbacks, reset_color_code) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5030b1512a6..fd38de03660 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -800,7 +800,7 @@ from litellm.proxy.spend_tracking.spend_event_producer import ( ) from litellm.proxy.telemetry.endpoints import router as telemetry_router from litellm.proxy.telemetry.middleware import TelemetryMiddleware -from litellm.proxy.telemetry.runtime import TelemetryRuntime +from litellm.proxy.telemetry.runtime import TelemetryRuntime, register_attempt_logger from litellm.proxy.telemetry.settings import describe_errors as describe_telemetry_errors from litellm.proxy.telemetry.settings import load_settings as load_telemetry_settings @@ -1789,7 +1789,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState db=lambda: ( # pyright: ignore[reportArgumentType] # PrismaWrapper forwards raw queries via __getattr__ prisma_client.writer_db if prisma_client is not None else None ), - register=litellm.logging_callback_manager.add_litellm_callback, + register=register_attempt_logger, ) ) ) diff --git a/litellm/proxy/telemetry/runtime.py b/litellm/proxy/telemetry/runtime.py index c9b49ee284f..49fa52b87a3 100644 --- a/litellm/proxy/telemetry/runtime.py +++ b/litellm/proxy/telemetry/runtime.py @@ -11,7 +11,9 @@ from typing import Final import httpx +import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import MAX_CALLBACKS from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # legacy params: dict signature ) @@ -57,6 +59,15 @@ async def deployment_hash_secret(store: TelemetryStore | None, fallback: Callabl return fallback() +def register_attempt_logger(logger: TelemetryAttemptLogger) -> None: + litellm.logging_callback_manager.add_litellm_callback(logger) + if logger not in litellm.callbacks: + verbose_proxy_logger.warning( + "telemetry: litellm.callbacks already holds %s callbacks, provider attempts will not be recorded", + MAX_CALLBACKS, + ) + + def deployment_hasher(secret: bytes) -> Callable[[str], str]: key: Final = hashlib.sha256(b"litellm-telemetry-deployment:" + secret).digest() return lambda model_id: hmac.new(key, model_id.encode(), hashlib.sha256).hexdigest()[:16] diff --git a/tests/unit/proxy/common_utils/test_callback_utils.py b/tests/unit/proxy/common_utils/test_callback_utils.py index ecb2375d495..7d0b3cbd9df 100644 --- a/tests/unit/proxy/common_utils/test_callback_utils.py +++ b/tests/unit/proxy/common_utils/test_callback_utils.py @@ -154,9 +154,7 @@ def test_add_policy_to_applied_policies_header_uses_litellm_metadata_bucket(): "litellm_metadata": {}, } - add_policy_to_applied_policies_header( - request_data=request_data, policy_name="global-baseline" - ) + add_policy_to_applied_policies_header(request_data=request_data, policy_name="global-baseline") assert request_data["litellm_metadata"]["applied_policies"] == ["global-baseline"] assert "applied_policies" not in request_data.get("metadata", {}) @@ -311,9 +309,7 @@ def test_initialize_callbacks_on_proxy_instantiates_compression_interception( lambda litellm_settings, callback_specific_params: dummy_callback, ) - original_callbacks = ( - list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] - ) + original_callbacks = list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] litellm.callbacks = [] try: initialize_callbacks_on_proxy( @@ -366,20 +362,11 @@ def test_encrypt_callback_vars_round_trip(monkeypatch): enc_vars = encrypted["logging"][0]["callback_vars"] assert enc_vars["langfuse_secret_key"] != "sk-lf-secret" assert enc_vars["langfuse_public_key"] != "pk-lf-public" - assert ( - encrypted["callback_settings"]["callback_vars"]["langsmith_api_key"] - != "ls-api-key" - ) + assert encrypted["callback_settings"]["callback_vars"]["langsmith_api_key"] != "ls-api-key" decrypted = decrypt_callback_vars(encrypted) - assert ( - decrypted["logging"][0]["callback_vars"] - == original["logging"][0]["callback_vars"] - ) - assert ( - decrypted["callback_settings"]["callback_vars"] - == original["callback_settings"]["callback_vars"] - ) + assert decrypted["logging"][0]["callback_vars"] == original["logging"][0]["callback_vars"] + assert decrypted["callback_settings"]["callback_vars"] == original["callback_settings"]["callback_vars"] def test_encrypt_callback_vars_is_idempotent(monkeypatch): @@ -402,10 +389,7 @@ def test_decrypt_callback_vars_passes_through_legacy_plaintext(monkeypatch): plaintext = _sample_metadata() decrypted = decrypt_callback_vars(plaintext) # legacy rows decrypt-fail and fall through unchanged - assert ( - decrypted["logging"][0]["callback_vars"]["langfuse_secret_key"] - == "sk-lf-secret" - ) + assert decrypted["logging"][0]["callback_vars"]["langfuse_secret_key"] == "sk-lf-secret" def test_callback_vars_helpers_handle_edge_shapes(monkeypatch): @@ -506,9 +490,7 @@ def test_initialize_callbacks_on_proxy_lakera_ignores_non_dict_callback_settings SimpleNamespace(prisma_client=None), ) - original_callbacks = ( - list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] - ) + original_callbacks = list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] litellm.callbacks = [] try: # A non-dict value must be ignored (init_params stays {}), not **-unpacked. @@ -526,9 +508,7 @@ def test_initialize_callbacks_on_proxy_lakera_ignores_non_dict_callback_settings @pytest.mark.parametrize("bad_root", [None, True]) -def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root( - monkeypatch, bad_root -): +def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root(monkeypatch, bad_root): """Regression: a blank `callback_settings:` key in YAML loads as None (and `callback_settings: true` as a bool); load_config forwards that value verbatim as callback_specific_params. Membership tests like @@ -545,9 +525,7 @@ def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root( CompressionInterceptionLogger, ) - original_callbacks = ( - list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] - ) + original_callbacks = list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] litellm.callbacks = [] try: initialize_callbacks_on_proxy( @@ -557,9 +535,7 @@ def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root( litellm_settings={}, callback_specific_params=bad_root, ) - assert any( - isinstance(c, CompressionInterceptionLogger) for c in litellm.callbacks - ) + assert any(isinstance(c, CompressionInterceptionLogger) for c in litellm.callbacks) finally: litellm.callbacks = original_callbacks @@ -611,7 +587,7 @@ def test_strip_callback_config_passes_through_non_dicts(value): _PROBE_MODULE_NAME = "custom_callback_probe" -_PROBE_MODULE_SOURCE = ''' +_PROBE_MODULE_SOURCE = """ from litellm.integrations.custom_logger import CustomLogger @@ -632,7 +608,7 @@ def log_event_fn(kwargs, response_obj, start_time, end_time): NOT_A_CALLBACK = "some-plain-string" proxy_handler_instance = FloorMaxTokens() -''' +""" @pytest.fixture @@ -646,9 +622,7 @@ def probe_config_path(tmp_path): """ (tmp_path / f"{_PROBE_MODULE_NAME}.py").write_text(_PROBE_MODULE_SOURCE) - original_callbacks = ( - list(litellm.callbacks) if isinstance(litellm.callbacks, list) else litellm.callbacks - ) + original_callbacks = list(litellm.callbacks) if isinstance(litellm.callbacks, list) else litellm.callbacks litellm.callbacks = [] ProxyLogging._callback_capabilities_cache.clear() try: @@ -674,7 +648,7 @@ def test_initialize_callbacks_on_proxy_rejects_class_valued_entry(probe_config_p silently never run the hook. Config load must fail instead.""" entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" - with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: + with pytest.raises(ValueError, match="litellm_settings\\.callbacks entry") as exc_info: _load_callbacks([entry], probe_config_path) message = str(exc_info.value) @@ -692,12 +666,10 @@ def test_initialize_callbacks_on_proxy_rejects_class_valued_entry(probe_config_p ("NOT_A_CALLBACK", "str 'some-plain-string'"), ], ) -def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values( - probe_config_path, attribute, expected_fragment -): +def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values(probe_config_path, attribute, expected_fragment): entry = f"{_PROBE_MODULE_NAME}.{attribute}" - with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: + with pytest.raises(ValueError, match="litellm_settings\\.callbacks entry") as exc_info: _load_callbacks([entry], probe_config_path) message = str(exc_info.value) @@ -709,7 +681,7 @@ def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values( def test_initialize_callbacks_on_proxy_rejects_class_valued_non_list_value(probe_config_path): entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" - with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: + with pytest.raises(ValueError, match="litellm_settings\\.callbacks entry") as exc_info: _load_callbacks(entry, probe_config_path) assert entry in str(exc_info.value) @@ -762,3 +734,21 @@ def test_initialize_callbacks_on_proxy_accepts_instance_non_list_value(probe_con assert len(litellm.callbacks) == 1 assert isinstance(litellm.callbacks[0], CustomLogger) + + +@pytest.mark.parametrize( + "value", + [f"{_PROBE_MODULE_NAME}.proxy_handler_instance", [f"{_PROBE_MODULE_NAME}.proxy_handler_instance"]], + ids=["single_string", "list"], +) +def test_initialize_callbacks_on_proxy_keeps_callbacks_registered_before_it_across_reloads( + probe_config_path, value, monkeypatch: pytest.MonkeyPatch +): + internal = CustomLogger() + monkeypatch.setattr(litellm, "callbacks", [internal]) + + _load_callbacks(value, probe_config_path) + _load_callbacks(value, probe_config_path) + + assert litellm.callbacks[0] is internal + assert len(litellm.callbacks) == 3 diff --git a/tests/unit/proxy/telemetry/test_runtime.py b/tests/unit/proxy/telemetry/test_runtime.py index 253b0997403..588dcc8105b 100644 --- a/tests/unit/proxy/telemetry/test_runtime.py +++ b/tests/unit/proxy/telemetry/test_runtime.py @@ -3,10 +3,13 @@ from typing import Final import httpx import pytest + +import litellm from pydantic import HttpUrl from litellm.proxy.telemetry.attempt_logger import TelemetryAttemptLogger -from litellm.proxy.telemetry.runtime import TelemetryRuntime, deployment_hasher +from litellm.constants import MAX_CALLBACKS +from litellm.proxy.telemetry.runtime import TelemetryRuntime, deployment_hasher, register_attempt_logger from litellm.proxy.telemetry.settings import TelemetryDestination, TelemetrySettings from litellm.telemetry.consent import ConsentGatedSink, TelemetryConsent from litellm.telemetry.records import AttemptRecord, InstanceInfo, RequestRecord, TelemetryGroup, UIEvent @@ -262,3 +265,31 @@ async def test_settle_timeout_reads_the_configured_value_after_start_and_the_def ) assert runtime.settings.settle_timeout_seconds == 7.5 await runtime.stop() + + +@pytest.fixture +def full_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", [f"callback_{i}" for i in range(MAX_CALLBACKS)]) + + +@pytest.mark.usefixtures("full_callbacks") +def test_a_full_callback_list_warns_that_attempts_will_not_be_recorded(caplog: pytest.LogCaptureFixture) -> None: + logger: Final = TelemetryAttemptLogger(lambda: None, deployment_hasher(b"secret")) + + register_attempt_logger(logger) + + assert logger not in litellm.callbacks + assert "provider attempts will not be recorded" in caplog.text + + +def test_registering_the_attempt_logger_adds_it_once_without_a_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + logger: Final = TelemetryAttemptLogger(lambda: None, deployment_hasher(b"secret")) + + register_attempt_logger(logger) + register_attempt_logger(logger) + + assert litellm.callbacks == [logger] + assert "will not be recorded" not in caplog.text