From 7b459ac1bc13912d7e10e0823e5ee42dcf4a8dc6 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 12 Sep 2026 23:40:18 +0000 Subject: [PATCH] fix(proxy): resolve team callbacks when key logging is empty and report GCS flush failures An empty key-level logging list is now treated as unset, so team logging, deprecated team callback_settings and default_team_settings apply. /key/health tests the effective callbacks, flushes the GCS logger explicitly and reports failed uploads. Failed GCS batches stay queued for the next flush instead of being dropped Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/gcs_bucket/gcs_bucket.py | 87 ++++++++++---- litellm/proxy/_types.py | 4 +- litellm/proxy/litellm_pre_call_utils.py | 11 +- .../key_management_endpoints.py | 107 +++++++++++------- litellm/types/integrations/gcs_bucket.py | 7 ++ .../gcs_bucket/test_gcs_bucket.py | 104 +++++++++++++++++ .../otel/test_otel_v2_destinations.py | 10 +- .../test_key_management_endpoints.py | 91 +++++++++++++++ .../proxy/test_litellm_pre_call_utils.py | 63 +++++++++++ 9 files changed, 413 insertions(+), 71 deletions(-) create mode 100644 tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 31ceb338dcd..e5f7173e3cc 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -3,6 +3,7 @@ import hashlib import json import os import time +from collections.abc import Sequence from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final @@ -68,10 +69,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") - # When queue is at maxsize, flush immediately to make room (no blocking, no data dropped) - if self.log_queue.full(): - await self.flush_queue() - await self.log_queue.put(GCSLogQueueItem(payload=logging_payload, kwargs=kwargs, response_obj=response_obj)) + await self._enqueue(GCSLogQueueItem(payload=logging_payload, kwargs=kwargs, response_obj=response_obj)) except Exception as e: verbose_logger.exception("GCS Bucket logging error: %s", e) @@ -87,14 +85,35 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") - # When queue is at maxsize, flush immediately to make room (no blocking, no data dropped) - if self.log_queue.full(): - await self.flush_queue() - await self.log_queue.put(GCSLogQueueItem(payload=logging_payload, kwargs=kwargs, response_obj=response_obj)) + await self._enqueue(GCSLogQueueItem(payload=logging_payload, kwargs=kwargs, response_obj=response_obj)) except Exception as e: verbose_logger.exception("GCS Bucket logging error: %s", e) + async def _enqueue(self, item: GCSLogQueueItem) -> None: + if self.log_queue.full(): + await self.flush_queue() + if self.log_queue.full(): + self.log_queue.get_nowait() + verbose_logger.error("GCS Bucket log queue still full after flush, dropped the oldest queued event") + self.log_queue.put_nowait(item) + + def _requeue(self, items: Sequence[GCSLogQueueItem]) -> None: + dropped: Final = sum(1 for item in items if not self._put_nowait_or_drop(item)) + verbose_logger.error( + "GCS Bucket upload failed for %s events, %s kept in queue for the next flush, %s dropped (queue full)", + len(items), + len(items) - dropped, + dropped, + ) + + def _put_nowait_or_drop(self, item: GCSLogQueueItem) -> bool: + try: + self.log_queue.put_nowait(item) + except asyncio.QueueFull: + return False + return True + def _drain_queue_batch(self) -> list[GCSLogQueueItem]: """ Drain items from the queue (non-blocking), respecting batch_size limit. @@ -219,17 +238,19 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): verbose_logger.exception("GCS Bucket error logging batch payload to GCS bucket: %s", e) return (success_count, error_count) - async def _send_individual_logs(self, items: list[GCSLogQueueItem]) -> None: + async def _send_individual_logs(self, items: list[GCSLogQueueItem]) -> GCSFlushResult: """ Send each log individually as separate GCS objects (legacy behavior). This is used when GCS_USE_BATCHED_LOGGING is disabled. """ - for item in items: - await self._send_single_log_item(item) + failed_items: Final = tuple([item for item in items if not await self._send_single_log_item(item)]) + if failed_items: + self._requeue(failed_items) + return GCSFlushResult(sent=len(items) - len(failed_items), failed=len(failed_items)) - async def _send_single_log_item(self, item: GCSLogQueueItem) -> None: + async def _send_single_log_item(self, item: GCSLogQueueItem) -> bool: """ - Send a single log item to GCS as an individual object. + Send a single log item to GCS as an individual object. Returns whether the upload succeeded. """ try: gcs_logging_config: Final[GCSLoggingConfig] = await self.get_gcs_logging_config(item["kwargs"]) @@ -254,8 +275,25 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): ) except Exception as e: verbose_logger.exception("GCS Bucket error logging individual payload to GCS bucket: %s", e) + return False + return True - async def async_send_batch(self): + async def _send_grouped_batches(self, items: list[GCSLogQueueItem]) -> GCSFlushResult: + results: Final = tuple( + [ + (group_items, await self._send_grouped_batch(group_items, config_key)) + for config_key, group_items in self._group_items_by_config(items).items() + ] + ) + for group_items, (_, group_failed) in results: + if group_failed: + self._requeue(group_items) + return GCSFlushResult( + sent=sum(group_sent for _, (group_sent, _) in results), + failed=sum(group_failed for _, (_, group_failed) in results), + ) + + async def async_send_batch(self) -> None: """ Process queued logs - sends logs to GCS Bucket. @@ -264,18 +302,17 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): If disabled, sends each log individually as separate GCS objects (legacy behavior). """ + await self._send_queued_events() + + async def _send_queued_events(self) -> GCSFlushResult: items_to_process: Final = self._drain_queue_batch() if not items_to_process: - return + return GCSFlushResult(sent=0, failed=0) if self.use_batched_logging: - grouped_items: Final = self._group_items_by_config(items_to_process) - - for config_key, group_items in grouped_items.items(): - await self._send_grouped_batch(group_items, config_key) - else: - await self._send_individual_logs(items_to_process) + return await self._send_grouped_batches(items_to_process) + return await self._send_individual_logs(items_to_process) def _get_object_name(self, kwargs: dict, logging_payload: StandardLoggingPayload, response_obj: Any) -> str: """ @@ -355,12 +392,16 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): def _get_object_date_from_datetime(self, datetime_obj: datetime) -> str: return datetime_obj.strftime("%Y-%m-%d") - async def flush_queue(self): + async def flush_queue(self) -> None: """ Override flush_queue to work with asyncio.Queue. """ - await self.async_send_batch() + await self.flush_queue_and_report() + + async def flush_queue_and_report(self) -> GCSFlushResult: + result: Final = await self._send_queued_events() self.last_flush_time = time.time() + return result async def periodic_flush(self): """ diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ae6c042ab3a..783102b2416 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,7 +1,7 @@ import enum import json import os -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple @@ -4441,7 +4441,7 @@ class CurrentItemRateLimit(TypedDict): class LoggingCallbackStatus(TypedDict, total=False): - callbacks: list[str] + callbacks: ReadOnly[Sequence[str]] status: Literal["healthy", "unhealthy"] details: str | None diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 59971e54e46..ba7b034b1d7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -904,18 +904,21 @@ def _get_validated_callback_metadata(item: dict, *, source: str) -> AddTeamCallb class KeyAndTeamLoggingSettings: """ Helper class to get the dynamic logging settings for the key and team + + An empty ``logging`` list is the same as no ``logging`` key: both return ``None`` so the + caller falls through to the next level. Disabling a callback is ``litellm_disabled_callbacks``. """ @staticmethod def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): if user_api_key_dict.metadata is not None and "logging" in user_api_key_dict.metadata: - return decrypt_callback_vars(user_api_key_dict.metadata).get("logging") + return decrypt_callback_vars(user_api_key_dict.metadata).get("logging") or None return None @staticmethod def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): if user_api_key_dict.team_metadata is not None and "logging" in user_api_key_dict.team_metadata: - return decrypt_callback_vars(user_api_key_dict.team_metadata).get("logging") + return decrypt_callback_vars(user_api_key_dict.team_metadata).get("logging") or None return None @@ -1029,8 +1032,8 @@ def resolve_tenant_otel_destinations( Key settings win over team settings outright, the same precedence ``_get_dynamic_logging_metadata`` applies, so one caller never exports the same - backend to two accounts. An empty key-level list counts as configured, since that - is what disabling a key's callbacks writes. Returns empty when OTEL V2 is off, when + backend to two accounts. An empty key-level list is unset and falls through to the + team, the way the runtime parser reads it. Returns empty when OTEL V2 is off, when neither level named a destination-capable backend, or when the config is incomplete, and the request then keeps the operator's own exporters. diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 1fe174763c4..3a97752e3fe 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -38,6 +38,7 @@ from litellm.constants import ( MINIMUM_CUSTOM_KEY_LENGTH, UI_SESSION_TOKEN_TEAM_ID, ) +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.models.credentials import CredentialItem @@ -6973,7 +6974,9 @@ async def key_health( Check the health of the key Checks: - - If key based logging is configured correctly - sends a test log + - If the logging that applies to this key (key metadata, team metadata, or + `default_team_settings` in the config) is configured correctly - sends a test log + and, for gcs_bucket, flushes the queue and reports the upload result Usage @@ -7015,29 +7018,40 @@ async def key_health( } ``` """ + from litellm.proxy.litellm_pre_call_utils import ( + _get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # the request-time resolver; the health check must report the same callbacks a request would use + ) + from litellm.proxy.proxy_server import proxy_config + try: - # Get the key's metadata key_metadata: Final = user_api_key_dict.metadata - - health_status: Final[KeyHealthResponse] = KeyHealthResponse( - key="healthy", - logging_callbacks=None, - ) - - # Check if logging is configured in metadata if key_metadata and "logging" in key_metadata: - logging_statuses: Final = await test_key_logging( - user_api_key_dict=user_api_key_dict, - request=request, - key_logging=decrypt_callback_vars(key_metadata)["logging"], + _raise_if_key_logging_missing_callback_name(decrypt_callback_vars(key_metadata)["logging"]) + + callback_settings: Final = _get_dynamic_logging_metadata( + user_api_key_dict=user_api_key_dict, proxy_config=proxy_config + ) + logging_callbacks: Final = ( + () + if callback_settings is None + else tuple( + dict.fromkeys( + (*(callback_settings.success_callback or ()), *(callback_settings.failure_callback or ())) + ) ) - health_status["logging_callbacks"] = logging_statuses + ) + if not logging_callbacks: + return KeyHealthResponse(key="healthy", logging_callbacks=None) - # Check if any logging callback is unhealthy - if logging_statuses.get("status") == "unhealthy": - health_status["key"] = "unhealthy" - - return KeyHealthResponse(**health_status) + logging_statuses: Final = await test_key_logging( + user_api_key_dict=user_api_key_dict, + request=request, + logging_callbacks=logging_callbacks, + ) + return KeyHealthResponse( + key="unhealthy" if logging_statuses.get("status") == "unhealthy" else "healthy", + logging_callbacks=logging_statuses, + ) except Exception as e: raise ProxyException( @@ -7072,31 +7086,41 @@ async def _can_user_query_key_info( return False +def _raise_if_key_logging_missing_callback_name(key_logging: Sequence[Mapping[str, str]]) -> None: + if any(callback.get("callback_name") is None for callback in key_logging): + raise ValueError("callback_name is required in key_logging") + + +async def flush_gcs_and_describe_failures(gcs_logger: CustomLogger | None) -> str | None: + from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger + + if not isinstance(gcs_logger, GCSBucketLogger): + return "gcs_bucket callback was selected but no GCS logger was initialized" + flush_result: Final = await gcs_logger.flush_queue_and_report() + if flush_result.failed == 0: + return None + return f"GCS upload failed for {flush_result.failed} event(s), {flush_result.sent} uploaded" + + async def test_key_logging( user_api_key_dict: UserAPIKeyAuth, request: Request, - key_logging: Sequence[Mapping[str, str]], + logging_callbacks: Sequence[str], ) -> LoggingCallbackStatus: """ - Test the key-based logging + Test the logging callbacks that apply to this key - - Test that key logging is correctly formatted and all args are passed correctly - Make a mock completion call -> user can check if it's correctly logged + - For gcs_bucket, flush the queue and report whether the upload succeeded - Check if any logger.exceptions were triggered -> if they were then returns it to the user client side """ import logging from io import StringIO + from litellm.litellm_core_utils.litellm_logging import get_custom_logger_compatible_class from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import general_settings, proxy_config - logging_callbacks: Final[list[str]] = [] - for callback in key_logging: - if callback.get("callback_name") is not None: - logging_callbacks.append(callback["callback_name"]) - else: - raise ValueError("callback_name is required in key_logging") - log_capture_string: Final = StringIO() ch: Final = logging.StreamHandler(log_capture_string) ch.setLevel(logging.ERROR) @@ -7131,21 +7155,28 @@ async def test_key_logging( await asyncio.sleep(2) # wait for callbacks to run, callbacks use batching so wait for the flush event - # Check if any logger exceptions were triggered + gcs_failure: Final = ( + await flush_gcs_and_describe_failures(get_custom_logger_compatible_class("gcs_bucket")) + if "gcs_bucket" in logging_callbacks + else None + ) + log_contents: Final = log_capture_string.getvalue() logger.removeHandler(ch) - if log_contents: + if gcs_failure is not None or log_contents: return LoggingCallbackStatus( callbacks=logging_callbacks, status="unhealthy", - details=f"Logger exceptions triggered, system is unhealthy: {log_contents}", - ) - else: - return LoggingCallbackStatus( - callbacks=logging_callbacks, - status="healthy", - details=f"No logger exceptions triggered, system is healthy. Manually check if logs were sent to {logging_callbacks} ", + details=f"Logger exceptions triggered, system is unhealthy: {gcs_failure or ''} {log_contents}".strip(), ) + return LoggingCallbackStatus( + callbacks=logging_callbacks, + status="healthy", + details=( + "No logger exceptions triggered, system is healthy. " + f"Manually check if logs were sent to {', '.join(logging_callbacks)}" + ), + ) _KEY_ALIAS_PATTERN: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_\-/\.@]{0,253}[a-zA-Z0-9]$") diff --git a/litellm/types/integrations/gcs_bucket.py b/litellm/types/integrations/gcs_bucket.py index 3840d7681a5..7f8ea7caee4 100644 --- a/litellm/types/integrations/gcs_bucket.py +++ b/litellm/types/integrations/gcs_bucket.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final from typing_extensions import TypedDict @@ -33,3 +34,9 @@ class GCSLogQueueItem(TypedDict): payload: StandardLoggingPayload kwargs: dict[str, Any] response_obj: Any | None + + +@dataclass(frozen=True, slots=True) +class GCSFlushResult: + sent: int + failed: int diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py new file mode 100644 index 00000000000..85306a10428 --- /dev/null +++ b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py @@ -0,0 +1,104 @@ +import asyncio +import json +from typing import Any, Final +from unittest.mock import patch + +import pytest + +from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.types.integrations.gcs_bucket import GCSFlushResult, GCSLoggingConfig, GCSLogQueueItem +from litellm.types.utils import StandardLoggingPayload + + +class _FakeUploadGCSLogger(GCSBucketLogger): + """Skips GCP auth; an upload raises when it carries any id in `failing_ids`, otherwise it is recorded""" + + def __init__(self, queue_maxsize: int = 0) -> None: + with patch("litellm.proxy.proxy_server.premium_user", True): # test-quality-ok: GCS logging is premium-gated + super().__init__(bucket_name="test-bucket") + self.log_queue = asyncio.Queue(maxsize=queue_maxsize) + self.failing_ids: frozenset[str] = frozenset() + self.uploaded: list[list[str]] = [] + + async def enqueue(self, request_id: str) -> None: + payload: Final = StandardLoggingPayload(id=request_id) # pyright: ignore[reportCallIssue] # partial payload is enough for queueing + await self._enqueue(GCSLogQueueItem(payload=payload, kwargs={}, response_obj={"id": request_id})) + + def queued_ids(self) -> list[str]: + return [self.log_queue.get_nowait()["payload"]["id"] for _ in range(self.log_queue.qsize())] + + async def get_gcs_logging_config(self, kwargs: dict[str, Any] | None = None) -> GCSLoggingConfig: + return GCSLoggingConfig(bucket_name="test-bucket", vertex_instance=None, path_service_account=None) + + async def construct_request_headers( + self, service_account_json: str | None, vertex_instance: VertexBase | None = None + ) -> dict[str, str]: + return {} + + async def _log_json_data_on_gcs( + self, headers: dict[str, str], bucket_name: str, object_name: str, logging_payload: StandardLoggingPayload | str + ) -> None: + ids: Final = ( + [json.loads(line)["id"] for line in logging_payload.splitlines()] + if isinstance(logging_payload, str) + else [logging_payload["id"]] + ) + if self.failing_ids.intersection(ids): + raise RuntimeError("storage.googleapis.com returned 404") + self.uploaded.append(ids) + + +@pytest.mark.asyncio +async def test_failed_batch_stays_queued_and_is_retried_on_the_next_flush(): + logger = _FakeUploadGCSLogger() + logger.failing_ids = frozenset({"req-1"}) + await logger.enqueue("req-1") + await logger.enqueue("req-2") + + failed_flush = await logger.flush_queue_and_report() + + assert failed_flush == GCSFlushResult(sent=0, failed=2) + assert logger.log_queue.qsize() == 2 + assert logger.uploaded == [] + + logger.failing_ids = frozenset() + retried_flush = await logger.flush_queue_and_report() + + assert retried_flush == GCSFlushResult(sent=2, failed=0) + assert logger.log_queue.qsize() == 0 + assert logger.uploaded == [["req-1", "req-2"]] + + +@pytest.mark.asyncio +async def test_individual_mode_requeues_only_the_failed_items(): + logger = _FakeUploadGCSLogger() + logger.use_batched_logging = False + logger.failing_ids = frozenset({"req-fail"}) + await logger.enqueue("req-ok") + await logger.enqueue("req-fail") + + result = await logger.flush_queue_and_report() + + assert result == GCSFlushResult(sent=1, failed=1) + assert logger.uploaded == [["req-ok"]] + assert logger.queued_ids() == ["req-fail"] + + +@pytest.mark.asyncio +async def test_enqueue_on_a_full_queue_whose_flush_failed_drops_the_oldest_event(): + logger = _FakeUploadGCSLogger(queue_maxsize=2) + logger.failing_ids = frozenset({"req-1", "req-2", "req-3"}) + await logger.enqueue("req-1") + await logger.enqueue("req-2") + + await logger.enqueue("req-3") + + assert logger.queued_ids() == ["req-2", "req-3"] + + +@pytest.mark.asyncio +async def test_empty_queue_flush_reports_nothing_sent_or_failed(): + logger = _FakeUploadGCSLogger() + + assert await logger.flush_queue_and_report() == GCSFlushResult(sent=0, failed=0) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py b/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py index 67695d5aed8..36c43396493 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py @@ -1803,15 +1803,17 @@ class TestTenantConfigAgreement: "callback_vars": {"langfuse_public_key": "pk", "langfuse_secret_key": "sk", "langfuse_host": host, **extra}, } - def test_a_key_that_disabled_its_callbacks_does_not_fall_back_to_the_team(self): - """Disabling a key's callbacks stores an empty list, which the sibling parser - reads as 'the key configured none'.""" + def test_a_key_with_an_empty_logging_list_falls_back_to_the_team(self): + """An empty key-level list is unset, the same way the sibling parser reads it, + so the team's destination applies.""" auth = UserAPIKeyAuth( metadata={"logging": []}, team_metadata={"logging": [self._entry("http://team.local")]}, ) - assert resolve_tenant_otel_destinations(auth) == () + destinations = resolve_tenant_otel_destinations(auth) + + assert [d.endpoint for d in destinations] == ["http://team.local/api/public/otel"] def test_two_entries_for_one_backend_merge_their_vars_last_wins(self): auth = UserAPIKeyAuth( diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 646ae43f37a..883b1f1a56e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15,6 +15,8 @@ import inspect from litellm.proxy._types import ( GenerateKeyRequest, + KeyHealthResponse, + LoggingCallbackStatus, NewUserRequest, LiteLLM_BudgetTable, LiteLLM_OrganizationTable, @@ -18186,3 +18188,92 @@ async def test_key_creator_cannot_detach_project_without_admin_access(): ) assert exc.value.status_code == 403 assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) + + +def _default_team_gcs_proxy_config(team_id: str): + from litellm.proxy.proxy_server import ProxyConfig + + pc: Final = ProxyConfig() + pc.config = { + "litellm_settings": { + "default_team_settings": [ + {"team_id": team_id, "success_callback": ["gcs_bucket"], "failure_callback": ["gcs_bucket"]} + ] + } + } + return pc + + +@pytest.mark.asyncio +async def test_key_health_tests_the_team_callbacks_an_empty_key_logging_list_falls_back_to(): + from litellm.proxy.management_endpoints.key_management_endpoints import key_health + + caller: Final = UserAPIKeyAuth(api_key="sk-1", team_id="team-gcs", metadata={"logging": []}, team_metadata={}) + logging_status: Final = LoggingCallbackStatus(callbacks=("gcs_bucket",), status="unhealthy", details="404") + with ( + patch("litellm.proxy.proxy_server.proxy_config", _default_team_gcs_proxy_config("team-gcs")), # test-quality-ok: key_health reads the module-level proxy config + patch( # test-quality-ok: the mock completion behind test_key_logging needs a running proxy + "litellm.proxy.management_endpoints.key_management_endpoints.test_key_logging", + AsyncMock(return_value=logging_status), + ) as test_logging, + ): + response = await key_health(request=MagicMock(), user_api_key_dict=caller) + + assert response == KeyHealthResponse(key="unhealthy", logging_callbacks=logging_status) + assert test_logging.await_args.kwargs["logging_callbacks"] == ("gcs_bucket",) + + +@pytest.mark.asyncio +async def test_key_health_without_any_effective_callbacks_reports_healthy_and_sends_no_test_log(): + from litellm.proxy.management_endpoints.key_management_endpoints import key_health + + caller: Final = UserAPIKeyAuth(api_key="sk-1", team_id="team-plain", metadata={"logging": []}, team_metadata={}) + with ( + patch("litellm.proxy.proxy_server.proxy_config", _default_team_gcs_proxy_config("team-gcs")), # test-quality-ok: key_health reads the module-level proxy config + patch( # test-quality-ok: the mock completion behind test_key_logging needs a running proxy + "litellm.proxy.management_endpoints.key_management_endpoints.test_key_logging", AsyncMock() + ) as test_logging, + ): + response = await key_health(request=MagicMock(), user_api_key_dict=caller) + + assert response == KeyHealthResponse(key="healthy", logging_callbacks=None) + test_logging.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_key_health_rejects_key_logging_entries_without_a_callback_name(): + from litellm.proxy.management_endpoints.key_management_endpoints import key_health + + caller: Final = UserAPIKeyAuth(api_key="sk-1", metadata={"logging": [{"callback_type": "success"}]}) + with patch("litellm.proxy.proxy_server.proxy_config", _default_team_gcs_proxy_config("team-gcs")): # test-quality-ok: key_health reads the module-level proxy config + with pytest.raises(ProxyException) as exc: + await key_health(request=MagicMock(), user_api_key_dict=caller) + + assert "callback_name is required" in exc.value.message + + +@pytest.mark.asyncio +async def test_flush_gcs_reports_the_failed_upload_count_from_the_registered_logger(): + from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger + from litellm.proxy.management_endpoints.key_management_endpoints import flush_gcs_and_describe_failures + from litellm.types.integrations.gcs_bucket import GCSFlushResult + + class _StuckUploadGCSLogger(GCSBucketLogger): + def __init__(self) -> None: + with patch("litellm.proxy.proxy_server.premium_user", True): # test-quality-ok: GCS logging is premium-gated + super().__init__(bucket_name="test-bucket") + + async def flush_queue_and_report(self) -> GCSFlushResult: + return GCSFlushResult(sent=0, failed=3) + + assert await flush_gcs_and_describe_failures(_StuckUploadGCSLogger()) == "GCS upload failed for 3 event(s), 0 uploaded" + + +@pytest.mark.asyncio +async def test_flush_gcs_names_a_missing_logger_when_the_callback_never_initialized(): + from litellm.proxy.management_endpoints.key_management_endpoints import flush_gcs_and_describe_failures + + assert ( + await flush_gcs_and_describe_failures(None) + == "gcs_bucket callback was selected but no GCS logger was initialized" + ) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 37b983d709a..fd786fcdeda 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2170,6 +2170,69 @@ def test_key_dynamic_logging_settings(): assert result is None +def test_empty_logging_list_on_key_and_team_is_unset(): + """A UI-generated `logging: []` is the same as no logging metadata, not an explicit override""" + auth = UserAPIKeyAuth(api_key="test-key", metadata={"logging": []}, team_metadata={"logging": []}) + + assert KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(auth) is None + assert KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(auth) is None + + +def test_empty_key_logging_falls_back_to_team_logging(): + auth = UserAPIKeyAuth( + api_key="test-key", + team_id="team-1", + metadata={"logging": []}, + team_metadata={ + "logging": [ + { + "callback_name": "gcs_bucket", + "callback_type": "success_and_failure", + "callback_vars": {"gcs_bucket_name": "team-bucket"}, + } + ] + }, + ) + + result = _get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=MagicMock()) + + assert result is not None + assert result.success_callback == ["gcs_bucket"] + assert result.failure_callback == ["gcs_bucket"] + assert result.callback_vars == {"gcs_bucket_name": "team-bucket"} + + +def test_empty_key_and_team_logging_falls_back_to_default_team_settings(): + from litellm.proxy.proxy_server import ProxyConfig + + pc = ProxyConfig() + pc.config = { + "litellm_settings": { + "default_team_settings": [ + { + "team_id": "team-gcs", + "success_callback": ["gcs_bucket"], + "failure_callback": ["gcs_bucket"], + "turn_off_message_logging": True, + } + ] + } + } + auth = UserAPIKeyAuth( + api_key="test-key", + team_id="team-gcs", + metadata={"logging": []}, + team_metadata={"logging": []}, + ) + + result = _get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=pc) + + assert result is not None + assert result.success_callback == ["gcs_bucket"] + assert result.failure_callback == ["gcs_bucket"] + assert result.callback_vars == {"turn_off_message_logging": "True"} + + def test_team_dynamic_logging_settings(): """ Test KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings method with arize and langfuse callbacks