diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 37b79e6d065..2d93bf3ece7 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,279 +1,297 @@ -import asyncio -import traceback -from datetime import datetime -from typing import Any, List, Optional, Union, cast - -import litellm -from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - get_litellm_metadata_from_kwargs, -) -from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import log_db_metrics -from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.utils import ProxyUpdateSpend -from litellm.types.utils import ( - StandardLoggingPayload, - StandardLoggingUserAPIKeyMetadata, -) -from litellm.utils import get_end_user_id_for_cost_tracking - - -class _ProxyDBLogger(CustomLogger): - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - await self._PROXY_track_cost_callback( - kwargs, response_obj, start_time, end_time - ) - - async def async_post_call_failure_hook( - self, - request_data: dict, - original_exception: Exception, - user_api_key_dict: UserAPIKeyAuth, - traceback_str: Optional[str] = None, - ): - request_route = user_api_key_dict.request_route - if _ProxyDBLogger._should_track_errors_in_db() is False: - return - elif request_route is not None and not RouteChecks.is_llm_api_route( - route=request_route - ): - return - - from litellm.proxy.proxy_server import proxy_logging_obj - - _metadata = dict( - StandardLoggingUserAPIKeyMetadata( - user_api_key_hash=user_api_key_dict.api_key, - user_api_key_alias=user_api_key_dict.key_alias, - user_api_key_spend=user_api_key_dict.spend, - user_api_key_max_budget=user_api_key_dict.max_budget, - user_api_key_budget_reset_at=( - user_api_key_dict.budget_reset_at.isoformat() - if user_api_key_dict.budget_reset_at - else None - ), - user_api_key_user_email=user_api_key_dict.user_email, - user_api_key_user_id=user_api_key_dict.user_id, - user_api_key_team_id=user_api_key_dict.team_id, - user_api_key_org_id=user_api_key_dict.org_id, - user_api_key_team_alias=user_api_key_dict.team_alias, - user_api_key_end_user_id=user_api_key_dict.end_user_id, - user_api_key_request_route=user_api_key_dict.request_route, - user_api_key_auth_metadata=user_api_key_dict.metadata, - ) - ) - _metadata["user_api_key"] = user_api_key_dict.api_key - _metadata["status"] = "failure" - _metadata["error_information"] = ( - StandardLoggingPayloadSetup.get_error_information( - original_exception=original_exception, - traceback_str=traceback_str, - ) - ) - - existing_metadata: dict = request_data.get("metadata", None) or {} - existing_metadata.update(_metadata) - - if "litellm_params" not in request_data: - request_data["litellm_params"] = {} - - existing_litellm_params = request_data.get("litellm_params", {}) - existing_litellm_metadata = existing_litellm_params.get("metadata", {}) or {} - - # Preserve tags from existing metadata - if existing_litellm_metadata.get("tags"): - existing_metadata["tags"] = existing_litellm_metadata.get("tags") - - request_data["litellm_params"]["proxy_server_request"] = ( - request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} - ) - request_data["litellm_params"]["metadata"] = existing_metadata - - # Preserve model name and custom_llm_provider - if "model" not in request_data: - request_data["model"] = existing_litellm_params.get("model") or request_data.get("model", "") - if "custom_llm_provider" not in request_data: - request_data["custom_llm_provider"] = existing_litellm_params.get("custom_llm_provider") or request_data.get("custom_llm_provider", "") - - await proxy_logging_obj.db_spend_update_writer.update_database( - token=user_api_key_dict.api_key, - response_cost=0.0, - user_id=user_api_key_dict.user_id, - end_user_id=user_api_key_dict.end_user_id, - team_id=user_api_key_dict.team_id, - kwargs=request_data, - completion_response=original_exception, - start_time=datetime.now(), - end_time=datetime.now(), - org_id=user_api_key_dict.org_id, - ) - - @log_db_metrics - async def _PROXY_track_cost_callback( - self, - kwargs, # kwargs to completion - completion_response: Optional[ - Union[litellm.ModelResponse, Any] - ], # response from completion - start_time=None, - end_time=None, # start/end time for completion - ): - from litellm.proxy.proxy_server import proxy_logging_obj, update_cache - - verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") - try: - verbose_proxy_logger.debug( - f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}" - ) - parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs=kwargs) - litellm_params = kwargs.get("litellm_params", {}) or {} - end_user_id = get_end_user_id_for_cost_tracking(litellm_params) - metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) - user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None)) - team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None)) - org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None)) - key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None)) - end_user_max_budget = metadata.get("user_api_end_user_max_budget", None) - sl_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) - response_cost = ( - sl_object.get("response_cost", None) - if sl_object is not None - else kwargs.get("response_cost", None) - ) - tags: Optional[List[str]] = ( - sl_object.get("request_tags", None) if sl_object is not None else None - ) - - if response_cost is not None: - user_api_key = metadata.get("user_api_key", None) - if kwargs.get("cache_hit", False) is True: - response_cost = 0.0 - verbose_proxy_logger.debug( - f"Cache Hit: response_cost {response_cost}, for user_id {user_id}" - ) - - verbose_proxy_logger.debug( - f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}" - ) - if _should_track_cost_callback( - user_api_key=user_api_key, - user_id=user_id, - team_id=team_id, - end_user_id=end_user_id, - ): - ## UPDATE DATABASE - await proxy_logging_obj.db_spend_update_writer.update_database( - token=user_api_key, - response_cost=response_cost, - user_id=user_id, - end_user_id=end_user_id, - team_id=team_id, - kwargs=kwargs, - completion_response=completion_response, - start_time=start_time, - end_time=end_time, - org_id=org_id, - ) - - # update cache - asyncio.create_task( - update_cache( - token=user_api_key, - user_id=user_id, - end_user_id=end_user_id, - response_cost=response_cost, - team_id=team_id, - parent_otel_span=parent_otel_span, - tags=tags, - ) - ) - - await proxy_logging_obj.slack_alerting_instance.customer_spend_alert( - token=user_api_key, - key_alias=key_alias, - end_user_id=end_user_id, - response_cost=response_cost, - max_budget=end_user_max_budget, - ) - else: - if kwargs.get("stream") is not True or ( - kwargs.get("stream") is True and "complete_streaming_response" in kwargs - ): - if sl_object is not None: - cost_tracking_failure_debug_info: Union[dict, str] = ( - sl_object["response_cost_failure_debug_info"] # type: ignore - or "response_cost_failure_debug_info is None in standard_logging_object" - ) - else: - cost_tracking_failure_debug_info = ( - "standard_logging_object not found" - ) - model = kwargs.get("model") - raise Exception( - f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing" - ) - except Exception as e: - error_msg = f"Error in tracking cost callback - {str(e)}\n Traceback:{traceback.format_exc()}" - model = kwargs.get("model", "") - metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) - litellm_metadata = kwargs.get("litellm_params", {}).get( - "litellm_metadata", {} - ) - old_metadata = kwargs.get("litellm_params", {}).get("metadata", {}) - call_type = kwargs.get("call_type", "") - error_msg += f"\n Args to _PROXY_track_cost_callback\n model: {model}\n chosen_metadata: {metadata}\n litellm_metadata: {litellm_metadata}\n old_metadata: {old_metadata}\n call_type: {call_type}\n" - asyncio.create_task( - proxy_logging_obj.failed_tracking_alert( - error_message=error_msg, - failing_model=model, - ) - ) - - verbose_proxy_logger.exception( - "Error in tracking cost callback - %s", str(e) - ) - - @staticmethod - def _should_track_errors_in_db(): - """ - Returns True if errors should be tracked in the database - - By default, errors are tracked in the database - - If users want to disable error tracking, they can set the disable_error_logs flag in the general_settings - """ - from litellm.proxy.proxy_server import general_settings - - if general_settings.get("disable_error_logs") is True: - return False - return - - -def _should_track_cost_callback( - user_api_key: Optional[str], - user_id: Optional[str], - team_id: Optional[str], - end_user_id: Optional[str], -) -> bool: - """ - Determine if the cost callback should be tracked based on the kwargs - """ - - # don't run track cost callback if user opted into disabling spend - if ProxyUpdateSpend.disable_spend_updates() is True: - return False - - if ( - user_api_key is not None - or user_id is not None - or team_id is not None - or end_user_id is not None - ): - return True - return False +import asyncio +import traceback +from datetime import datetime +from typing import Any, List, Optional, Union, cast + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + get_litellm_metadata_from_kwargs, +) +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import log_db_metrics +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.utils import ProxyUpdateSpend +from litellm.types.utils import ( + StandardLoggingPayload, + StandardLoggingUserAPIKeyMetadata, +) +from litellm.utils import get_end_user_id_for_cost_tracking + + +class _ProxyDBLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await self._PROXY_track_cost_callback( + kwargs, response_obj, start_time, end_time + ) + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: Optional[str] = None, + ): + request_route = user_api_key_dict.request_route + if _ProxyDBLogger._should_track_errors_in_db() is False: + return + elif request_route is not None and not RouteChecks.is_llm_api_route( + route=request_route + ): + return + + from litellm.proxy.proxy_server import proxy_logging_obj + + _metadata = dict( + StandardLoggingUserAPIKeyMetadata( + user_api_key_hash=user_api_key_dict.api_key, + user_api_key_alias=user_api_key_dict.key_alias, + user_api_key_spend=user_api_key_dict.spend, + user_api_key_max_budget=user_api_key_dict.max_budget, + user_api_key_budget_reset_at=( + user_api_key_dict.budget_reset_at.isoformat() + if user_api_key_dict.budget_reset_at + else None + ), + user_api_key_user_email=user_api_key_dict.user_email, + user_api_key_user_id=user_api_key_dict.user_id, + user_api_key_team_id=user_api_key_dict.team_id, + user_api_key_org_id=user_api_key_dict.org_id, + user_api_key_team_alias=user_api_key_dict.team_alias, + user_api_key_end_user_id=user_api_key_dict.end_user_id, + user_api_key_request_route=user_api_key_dict.request_route, + user_api_key_auth_metadata=user_api_key_dict.metadata, + ) + ) + _metadata["user_api_key"] = user_api_key_dict.api_key + _metadata["status"] = "failure" + _metadata["error_information"] = ( + StandardLoggingPayloadSetup.get_error_information( + original_exception=original_exception, + traceback_str=traceback_str, + ) + ) + + existing_metadata: dict = request_data.get("metadata", None) or {} + existing_metadata.update(_metadata) + + if "litellm_params" not in request_data: + request_data["litellm_params"] = {} + + existing_litellm_params = request_data.get("litellm_params", {}) + existing_litellm_metadata = existing_litellm_params.get("metadata", {}) or {} + + # Preserve tags from existing metadata + if existing_litellm_metadata.get("tags"): + existing_metadata["tags"] = existing_litellm_metadata.get("tags") + + request_data["litellm_params"]["proxy_server_request"] = ( + request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} + ) + request_data["litellm_params"]["metadata"] = existing_metadata + + # Preserve model name and custom_llm_provider + if "model" not in request_data: + request_data["model"] = existing_litellm_params.get("model") or request_data.get("model", "") + if "custom_llm_provider" not in request_data: + request_data["custom_llm_provider"] = existing_litellm_params.get( + "custom_llm_provider" + ) or request_data.get("custom_llm_provider", "") + + # Propagate standard_logging_object and litellm_trace_id from the + # Logging instance so that _get_session_id_for_spend_log uses the same + # trace_id that Langfuse received (via async_failure_handler). + # Without this, the DB session_id would be a random UUID that doesn't + # match the Langfuse trace_id, making failed requests unsearchable. + _litellm_logging_obj = request_data.get("litellm_logging_obj") + if _litellm_logging_obj is not None: + if "standard_logging_object" not in request_data: + request_data["standard_logging_object"] = getattr( + _litellm_logging_obj, "model_call_details", {} + ).get("standard_logging_object") + if request_data.get("litellm_trace_id") is None: + request_data["litellm_trace_id"] = getattr( + _litellm_logging_obj, "litellm_trace_id", None + ) + + await proxy_logging_obj.db_spend_update_writer.update_database( + token=user_api_key_dict.api_key, + response_cost=0.0, + user_id=user_api_key_dict.user_id, + end_user_id=user_api_key_dict.end_user_id, + team_id=user_api_key_dict.team_id, + kwargs=request_data, + completion_response=original_exception, + start_time=datetime.now(), + end_time=datetime.now(), + org_id=user_api_key_dict.org_id, + ) + + @log_db_metrics + async def _PROXY_track_cost_callback( + self, + kwargs, # kwargs to completion + completion_response: Optional[ + Union[litellm.ModelResponse, Any] + ], # response from completion + start_time=None, + end_time=None, # start/end time for completion + ): + from litellm.proxy.proxy_server import proxy_logging_obj, update_cache + + verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") + try: + verbose_proxy_logger.debug( + f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}" + ) + parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs=kwargs) + litellm_params = kwargs.get("litellm_params", {}) or {} + end_user_id = get_end_user_id_for_cost_tracking(litellm_params) + metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) + user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None)) + team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None)) + org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None)) + key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None)) + end_user_max_budget = metadata.get("user_api_end_user_max_budget", None) + sl_object: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object", None + ) + response_cost = ( + sl_object.get("response_cost", None) + if sl_object is not None + else kwargs.get("response_cost", None) + ) + tags: Optional[List[str]] = ( + sl_object.get("request_tags", None) if sl_object is not None else None + ) + + if response_cost is not None: + user_api_key = metadata.get("user_api_key", None) + if kwargs.get("cache_hit", False) is True: + response_cost = 0.0 + verbose_proxy_logger.debug( + f"Cache Hit: response_cost {response_cost}, for user_id {user_id}" + ) + + verbose_proxy_logger.debug( + f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}" + ) + if _should_track_cost_callback( + user_api_key=user_api_key, + user_id=user_id, + team_id=team_id, + end_user_id=end_user_id, + ): + ## UPDATE DATABASE + await proxy_logging_obj.db_spend_update_writer.update_database( + token=user_api_key, + response_cost=response_cost, + user_id=user_id, + end_user_id=end_user_id, + team_id=team_id, + kwargs=kwargs, + completion_response=completion_response, + start_time=start_time, + end_time=end_time, + org_id=org_id, + ) + + # update cache + asyncio.create_task( + update_cache( + token=user_api_key, + user_id=user_id, + end_user_id=end_user_id, + response_cost=response_cost, + team_id=team_id, + parent_otel_span=parent_otel_span, + tags=tags, + ) + ) + + await proxy_logging_obj.slack_alerting_instance.customer_spend_alert( + token=user_api_key, + key_alias=key_alias, + end_user_id=end_user_id, + response_cost=response_cost, + max_budget=end_user_max_budget, + ) + else: + if kwargs.get("stream") is not True or ( + kwargs.get("stream") is True and "complete_streaming_response" in kwargs + ): + if sl_object is not None: + cost_tracking_failure_debug_info: Union[dict, str] = ( + sl_object["response_cost_failure_debug_info"] # type: ignore + or "response_cost_failure_debug_info is None in standard_logging_object" + ) + else: + cost_tracking_failure_debug_info = ( + "standard_logging_object not found" + ) + model = kwargs.get("model") + raise Exception( + f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing" + ) + except Exception as e: + error_msg = f"Error in tracking cost callback - {str(e)}\n Traceback:{traceback.format_exc()}" + model = kwargs.get("model", "") + metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) + litellm_metadata = kwargs.get("litellm_params", {}).get( + "litellm_metadata", {} + ) + old_metadata = kwargs.get("litellm_params", {}).get("metadata", {}) + call_type = kwargs.get("call_type", "") + error_msg += f"\n Args to _PROXY_track_cost_callback\n model: {model}\n chosen_metadata: {metadata}\n litellm_metadata: {litellm_metadata}\n old_metadata: {old_metadata}\n call_type: {call_type}\n" + asyncio.create_task( + proxy_logging_obj.failed_tracking_alert( + error_message=error_msg, + failing_model=model, + ) + ) + + verbose_proxy_logger.exception( + "Error in tracking cost callback - %s", str(e) + ) + + @staticmethod + def _should_track_errors_in_db(): + """ + Returns True if errors should be tracked in the database + + By default, errors are tracked in the database + + If users want to disable error tracking, they can set the disable_error_logs flag in the general_settings + """ + from litellm.proxy.proxy_server import general_settings + + if general_settings.get("disable_error_logs") is True: + return False + return + + +def _should_track_cost_callback( + user_api_key: Optional[str], + user_id: Optional[str], + team_id: Optional[str], + end_user_id: Optional[str], +) -> bool: + """ + Determine if the cost callback should be tracked based on the kwargs + """ + + # don't run track cost callback if user opted into disabling spend + if ProxyUpdateSpend.disable_spend_updates() is True: + return False + + if ( + user_api_key is not None + or user_id is not None + or team_id is not None + or end_user_id is not None + ): + return True + return False diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index cb6d90103f7..d269a9531fd 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -126,3 +126,356 @@ async def test_async_post_call_failure_hook_non_llm_route(): # Assert that update_database was NOT called for non-LLM routes mock_update_database.assert_not_called() + + +@pytest.mark.asyncio +async def test_track_cost_callback_skips_when_no_standard_logging_object(): + """ + Reproduces the bug where _PROXY_track_cost_callback raises + 'Cost tracking failed for model=None' when kwargs has no + standard_logging_object (e.g. call_type=afile_delete). + + File operations have no model and no standard_logging_object. + The callback should skip gracefully instead of raising. + """ + logger = _ProxyDBLogger() + + kwargs = { + "call_type": "afile_delete", + "model": None, + "litellm_call_id": "test-call-id", + "litellm_params": {}, + "stream": False, + } + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # update_database should NOT be called — nothing to track + mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called() + + # failed_tracking_alert should NOT be called — this is not an error + mock_proxy_logging.failed_tracking_alert.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj(): + """ + When an LLM call fails, the proxy calls post_call_failure_hook with + request_data that doesn't contain standard_logging_object. But the + litellm_logging_obj (set by function_setup) is in request_data and + holds the standard_logging_object with the correct trace_id. + + The failure hook should propagate this so the DB spend log's session_id + matches the Langfuse trace_id. + """ + logger = _ProxyDBLogger() + + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", + user_id="test_user_id", + team_id="test_team_id", + ) + + # Simulate a litellm_logging_obj with model_call_details containing + # the standard_logging_object (as set by _failure_handler_helper_fn) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_trace_id = "trace-id-from-logging-obj" + mock_logging_obj.model_call_details = { + "standard_logging_object": { + "trace_id": "trace-id-from-logging-obj", + "error_str": "InternalServerError", + } + } + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + "litellm_logging_obj": mock_logging_obj, + # Note: no "standard_logging_object" and no "litellm_trace_id" + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Provider error"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + call_kwargs = mock_update_database.call_args[1]["kwargs"] + + # standard_logging_object should have been propagated from logging obj + assert call_kwargs.get("standard_logging_object") is not None + assert ( + call_kwargs["standard_logging_object"]["trace_id"] + == "trace-id-from-logging-obj" + ) + # litellm_trace_id should also be propagated as a fallback + assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_with_team_alias(): + """ + When team_id is set but team_alias is missing (and key_alias is present), + _enrich_failure_metadata_with_key_info should look up the team from cache + and populate user_api_key_team_alias. + """ + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "my-team-alias" + + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": "my-key-alias", # already set + "user_api_key_team_id": "test_team_id", + "user_api_key_team_alias": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + assert result["user_api_key_team_alias"] == "my-team-alias" + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_with_full_key_lookup(): + """ + When all key fields are null (auth error 401 scenario), _enrich_failure_metadata_with_key_info + should look up the key object from cache/DB and populate alias, user_id, team_id, + then look up the team to get team_alias. + """ + mock_key_obj = MagicMock() + mock_key_obj.key_alias = "fetched-key-alias" + mock_key_obj.user_id = "fetched-user-id" + mock_key_obj.team_id = "fetched-team-id" + mock_key_obj.org_id = "fetched-org-id" + + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "fetched-team-alias" + + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + return_value=mock_key_obj, + ), patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": None, # all null - simulates auth error path + "user_api_key_user_id": None, + "user_api_key_team_id": None, + "user_api_key_team_alias": None, + "user_api_key_org_id": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + assert result["user_api_key_alias"] == "fetched-key-alias" + assert result["user_api_key_user_id"] == "fetched-user-id" + assert result["user_api_key_team_id"] == "fetched-team-id" + assert result["user_api_key_org_id"] == "fetched-org-id" + assert result["user_api_key_team_alias"] == "fetched-team-alias" + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_skips_when_team_alias_present(): + """ + When team_alias is already populated, _enrich_failure_metadata_with_key_info + should not perform a team cache lookup. + """ + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + ) as mock_get_key, patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": "existing-alias", + "user_api_key_team_id": "test_team_id", + "user_api_key_team_alias": "already-set", + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + assert result["user_api_key_team_alias"] == "already-set" + mock_get_key.assert_not_called() + mock_get_team.assert_not_called() + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_skips_when_no_api_key(): + """ + When api_key hash is absent, _enrich_failure_metadata_with_key_info should + not perform any lookups. + """ + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + ) as mock_get_key: + metadata = { + "user_api_key": None, + "user_api_key_alias": None, + "user_api_key_team_id": None, + "user_api_key_team_alias": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + mock_get_key.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_enriches_auth_error_metadata(): + """ + Simulates a 401 ProxyException (e.g. can_key_call_model). In this case + UserAPIKeyAuth is created with only api_key set. The failure hook should + look up the key and team from cache/DB to populate all missing fields. + """ + logger = _ProxyDBLogger() + + # This is what auth_exception_handler creates for 401 errors + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed_key", + # key_alias, user_id, team_id, team_alias are all None + ) + + request_data = { + "model": "claude-haiku-4-5", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + } + + mock_key_obj = MagicMock() + mock_key_obj.key_alias = "my-key-alias" + mock_key_obj.user_id = "my-user-id" + mock_key_obj.team_id = "my-team-id" + mock_key_obj.org_id = None + + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "my-team-alias" + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + return_value=mock_key_obj, + ), patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("401 - model not allowed"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + call_args = mock_update_database.call_args[1] + metadata = call_args["kwargs"]["litellm_params"]["metadata"] + assert metadata["user_api_key_alias"] == "my-key-alias" + assert metadata["user_api_key_user_id"] == "my-user-id" + assert metadata["user_api_key_team_id"] == "my-team-id" + assert metadata["user_api_key_team_alias"] == "my-team-alias" + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_enriches_missing_team_alias(): + """ + When user_api_key_dict has a team_id but no team_alias, async_post_call_failure_hook + should look up the team from cache and populate user_api_key_team_alias in the + spend log metadata written to the DB. + """ + logger = _ProxyDBLogger() + + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", + key_alias="test_alias", + user_id="test_user_id", + team_id="test_team_id", + team_alias=None, # Missing - simulates regular key auth where SQL view omits team_alias + ) + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + } + + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "enriched-team-alias" + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Provider rate limit"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + call_args = mock_update_database.call_args[1] + metadata = call_args["kwargs"]["litellm_params"]["metadata"] + assert metadata["user_api_key_team_alias"] == "enriched-team-alias" + assert metadata["user_api_key_team_id"] == "test_team_id" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_value", [None, ""]) +async def test_track_cost_callback_skips_for_falsy_model_and_no_slo(model_value): + """ + Same bug as above but model can also be empty string (e.g. health check callbacks). + The guard should catch all falsy model values when sl_object is missing. + """ + logger = _ProxyDBLogger() + + kwargs = { + "call_type": "acompletion", + "model": model_value, + "litellm_params": {}, + "stream": False, + } + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_proxy_logging.failed_tracking_alert.assert_not_called()