From c8ae1cf4ed7e531fd0b73ee795cafcf8181725be Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Fri, 28 Aug 2026 09:59:26 -0700 Subject: [PATCH] fix(proxy): dispatch team callbacks for passthrough failures --- basedpyright-code-budget.json | 2 +- litellm/proxy/litellm_pre_call_utils.py | 16 +- .../team_callback_endpoints.py | 4 +- .../pass_through_endpoints.py | 181 ++++++-- .../vertex_ai_endpoints/langfuse_endpoints.py | 4 +- tests/proxy_unit_tests/test_proxy_utils.py | 8 +- .../test_team_callback_endpoints.py | 18 +- .../test_pass_through_endpoints.py | 410 +++++++++++++++++- .../proxy/test_litellm_pre_call_utils.py | 24 +- 9 files changed, 601 insertions(+), 66 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index cd39aa3931f..9757b7b6f77 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1808 + "limit": 1806 }, "reportRedeclaration": { "limit": 8 diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index fb269e846fc..5490004742a 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -801,7 +801,7 @@ class KeyAndTeamLoggingSettings: return None -def _get_dynamic_logging_metadata( +def get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> TeamCallbackMetadata | None: callback_settings_obj: TeamCallbackMetadata | None = None @@ -915,12 +915,16 @@ def clean_headers( return clean_headers -def _is_credential_header(header: str) -> bool: +def _is_credential_header(header: str, litellm_key_header_name: str | None = None) -> bool: """Whether `header` carries a caller credential rather than request context.""" - return header.lower() in _CREDENTIAL_HEADER_NAMES + return header.lower() in _CREDENTIAL_HEADER_NAMES or ( + litellm_key_header_name is not None and header.lower() == litellm_key_header_name.lower() + ) -def redact_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: +def redact_credential_headers( + headers: Mapping[str, str], litellm_key_header_name: str | None = None +) -> Mapping[str, str]: """Return a copy of `headers` with credential-bearing values masked. `clean_headers` deliberately preserves some credential headers so they can be @@ -934,7 +938,7 @@ def redact_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: the stored copy and the logging callbacks JSON-serialize it. """ return { - header: (_REDACTED_HEADER_VALUE if _is_credential_header(header) else value) + header: (_REDACTED_HEADER_VALUE if _is_credential_header(header, litellm_key_header_name) else value) for header, value in headers.items() } @@ -2135,7 +2139,7 @@ async def add_litellm_data_to_request( ) # Team Callbacks controls - callback_settings_obj: Final = _get_dynamic_logging_metadata( + callback_settings_obj: Final = get_dynamic_logging_metadata( user_api_key_dict=user_api_key_dict, proxy_config=proxy_config ) if callback_settings_obj is not None: diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 08346983f32..5d5a01a8ed5 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -122,7 +122,7 @@ def _resolve_team_callbacks(team_metadata: object) -> TeamCallbackMetadata: ``metadata["logging"]`` holds the ``AddTeamCallback`` entries written by ``POST /team/{team_id}/callback`` and by the Admin UI, while ``metadata["callback_settings"]`` holds the older ``TeamCallbackMetadata`` - shape. Request-time resolution in ``_get_dynamic_logging_metadata`` treats + shape. Request-time resolution in ``get_dynamic_logging_metadata`` treats the two as mutually exclusive: a populated ``logging`` slot wins outright and ``callback_settings`` is consulted only as the deprecated fallback. This reader applies the same precedence so it reports what a request would @@ -605,7 +605,7 @@ async def disable_team_logging( # Update metadata team_metadata["callback_settings"] = team_callback_settings_obj.model_dump() - # _get_dynamic_logging_metadata stops at metadata["logging"], where the API + # get_dynamic_logging_metadata stops at metadata["logging"], where the API # and Admin UI register callbacks, without ever reading callback_settings. team_metadata["logging"] = [] # mutable-ok: the disabled state is persisted as an empty JSON array team_metadata = encrypt_callback_vars(team_metadata) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3d60f4f5f3a..8d2668fed52 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -77,7 +77,11 @@ from litellm.proxy.common_utils.http_parsing_utils import ( from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + get_dynamic_logging_metadata, + redact_credential_headers, +) from litellm.proxy.utils import normalize_route_for_root_path from litellm.repositories.team_repository import TeamRepository from litellm.secret_managers.main import get_secret_str @@ -89,7 +93,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( EndpointType, PassthroughStandardLoggingPayload, ) -from litellm.types.utils import Usage +from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage from .streaming_handler import PassThroughStreamingHandler from .success_handler import PassThroughEndpointLogging @@ -605,7 +609,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): "url": str(request.url), "method": request.method, "body": copy.copy(_parsed_body), # use copy instead of deepcopy - "headers": request.headers, + "headers": redact_credential_headers( + _safe_get_request_headers(request), + litellm_key_header_name=_get_custom_litellm_key_header_name(), + ), }, }, "call_type": "pass_through_endpoint", @@ -750,10 +757,33 @@ def _build_passthrough_failure_request_payload( return request_payload +async def _dispatch_passthrough_dynamic_failure( + logging_obj: LiteLLMLoggingObj, + exception: Exception, + traceback_str: str, +) -> None: + if logging_obj.model_call_details.get("has_logged_async_failure", False): + return + try: + await logging_obj.dispatch_failure_handlers( + exception=exception, + traceback_exception=traceback_str, + prefer_async_handlers=True, + ) + except Exception: + verbose_proxy_logger.warning( + "pass_through_endpoint: dynamic failure callback raised", + exc_info=True, + ) + finally: + logging_obj.model_call_details["has_logged_async_failure"] = True + + async def _log_passthrough_upstream_failure( response: httpx.Response, user_api_key_dict: UserAPIKeyAuth, request_payload: dict, + logging_obj: LiteLLMLoggingObj, ) -> None: """Fire LiteLLM-side failure hooks (spend tracking, alerting callbacks) for an upstream 4xx/5xx passthrough response. @@ -781,12 +811,18 @@ async def _log_passthrough_upstream_failure( status_code=response.status_code, detail=f"Upstream passthrough request failed with status {response.status_code}", ) + traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + await _dispatch_passthrough_dynamic_failure( + logging_obj=logging_obj, + exception=synthetic_exception, + traceback_str=traceback_str, + ) try: await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=synthetic_exception, request_data=request_payload, - traceback_str=traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG), + traceback_str=traceback_str, ) except Exception: # noqa: BLE001 - a failing logging callback must never break the passthrough response verbose_proxy_logger.warning( @@ -795,6 +831,50 @@ async def _log_passthrough_upstream_failure( ) +def _get_custom_litellm_key_header_name() -> str | None: + from litellm.proxy.proxy_server import general_settings + + configured_header: Final = general_settings.get("litellm_key_header_name") if general_settings else None + return configured_header if isinstance(configured_header, str) and configured_header else None + + +def _get_passthrough_logging_init_params( + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[ + list[str | Callable | CustomLogger] | None, + list[str | Callable | CustomLogger] | None, + dict[str, object] | None, +]: + from litellm.proxy.proxy_server import proxy_config + + callback_settings: Final = get_dynamic_logging_metadata( + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + ) + if callback_settings is None: + return None, None, None + + success_callback_names: Final = callback_settings.success_callback + dynamic_success_callbacks: Final[list[str | Callable | CustomLogger] | None] = ( + [*success_callback_names] if success_callback_names is not None else None + ) + failure_callback_names: Final = callback_settings.failure_callback + dynamic_failure_callbacks: Final[list[str | Callable | CustomLogger] | None] = ( + [*failure_callback_names] if failure_callback_names is not None else None + ) + callback_vars: Final = callback_settings.callback_vars + if not callback_vars: + return dynamic_success_callbacks, dynamic_failure_callbacks, None + + return ( + dynamic_success_callbacks, + dynamic_failure_callbacks, + dict( + (*callback_vars.items(), (TRUSTED_CALLBACK_VARS_FIELD, callback_vars)), + ), + ) + + from litellm.passthrough.timeout_utils import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, # noqa: F401 - re-exported for backward compat resolve_llm_passthrough_timeout, # noqa: F401 - re-exported for backward compat @@ -902,7 +982,10 @@ async def pass_through_request( verbose_proxy_logger.debug( "Pass through endpoint sending request to \nURL %s\nheaders: %s\nbody: %s\n", url, - headers, + redact_credential_headers( + headers, + litellm_key_header_name=_get_custom_litellm_key_header_name(), + ), _parsed_body, ) @@ -928,6 +1011,9 @@ async def pass_through_request( # read e.g. ``chat gpt-4o`` instead of ``chat unknown``. passthrough_model: Final = (_parsed_body.get("model") if isinstance(_parsed_body, dict) else None) or "unknown" start_time: Final = datetime.now() + dynamic_success_callbacks, dynamic_failure_callbacks, callback_kwargs = _get_passthrough_logging_init_params( + user_api_key_dict=user_api_key_dict + ) logging_obj = Logging( model=passthrough_model, messages=[{"role": "user", "content": safe_dumps(_parsed_body)}], @@ -936,6 +1022,9 @@ async def pass_through_request( start_time=start_time, litellm_call_id=litellm_call_id, function_id="1245", + dynamic_success_callbacks=dynamic_success_callbacks, + dynamic_failure_callbacks=dynamic_failure_callbacks, + kwargs=callback_kwargs, ) # Store passthrough guardrails config on logging_obj for field targeting @@ -1132,7 +1221,10 @@ async def pass_through_request( additional_args={ "complete_input_dict": _parsed_body, "api_base": str(logging_url), - "headers": headers, + "headers": redact_credential_headers( + headers, + litellm_key_header_name=_get_custom_litellm_key_header_name(), + ), }, ) stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( @@ -1191,6 +1283,7 @@ async def pass_through_request( custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, ), + logging_obj=logging_obj, ) # Call response headers hook for streaming pass-through @@ -1272,6 +1365,7 @@ async def pass_through_request( custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, ), + logging_obj=logging_obj, ) # Call response headers hook for detected streaming pass-through @@ -1366,6 +1460,7 @@ async def pass_through_request( response=response, user_api_key_dict=user_api_key_dict, request_payload=failure_request_payload, + logging_obj=logging_obj, ) if response.status_code < 400 and response_body is not None and guardrails_to_run: @@ -1580,13 +1675,18 @@ async def pass_through_request( _carry_guardrail_logging_info(request_payload, post_call_guardrail_data) + traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + if logging_obj is not None: + await _dispatch_passthrough_dynamic_failure( + logging_obj=logging_obj, + exception=e, + traceback_str=traceback_str, + ) await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=request_payload, - traceback_str=traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, - ), + traceback_str=traceback_str, ) ######################################################### @@ -2032,6 +2132,10 @@ async def websocket_passthrough_request( verbose_proxy_logger.info("WebSocket passthrough (%s): Starting WebSocket connection to %s", endpoint, target) + dynamic_success_callbacks, dynamic_failure_callbacks, callback_kwargs = _get_passthrough_logging_init_params( + user_api_key_dict=user_api_key_dict + ) + # Only accept the WebSocket if requested (for generic usage) if accept_websocket: await websocket.accept() @@ -2052,7 +2156,6 @@ async def websocket_passthrough_request( ]: upstream_headers[header_name] = header_value - # Initialize logging object similar to HTTP passthrough logging_obj: Final = Logging( model="unknown", messages=[{"role": "user", "content": "WebSocket connection"}], @@ -2061,6 +2164,9 @@ async def websocket_passthrough_request( start_time=start_time, litellm_call_id=litellm_call_id, function_id="websocket_passthrough", + dynamic_success_callbacks=dynamic_success_callbacks, + dynamic_failure_callbacks=dynamic_failure_callbacks, + kwargs=callback_kwargs, ) # Create passthrough logging payload @@ -2115,7 +2221,10 @@ async def websocket_passthrough_request( additional_args={ "complete_input_dict": {}, "api_base": target, - "headers": upstream_headers, + "headers": redact_credential_headers( + upstream_headers, + litellm_key_header_name=_get_custom_litellm_key_header_name(), + ), }, ) @@ -2399,15 +2508,24 @@ async def websocket_passthrough_request( if logging_obj is not None: request_payload["litellm_logging_obj"] = logging_obj - # Log the connection failure using the same pattern as HTTP - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=exc, - request_data=request_payload, - traceback_str=traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, - ), + traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + await _dispatch_passthrough_dynamic_failure( + logging_obj=logging_obj, + exception=exc, + traceback_str=traceback_str, ) + try: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=exc, + request_data=request_payload, + traceback_str=traceback_str, + ) + except Exception: # noqa: BLE001 - a failing logging callback must never change the WebSocket close behavior + verbose_proxy_logger.warning( + "pass_through_endpoint: post_call_failure_hook raised for WebSocket connection failure", + exc_info=True, + ) if _client_socket_is_open(websocket): await websocket.close( @@ -2427,15 +2545,24 @@ async def websocket_passthrough_request( if logging_obj is not None: request_payload["litellm_logging_obj"] = logging_obj - # Log the unexpected error using the same pattern as HTTP - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=request_payload, - traceback_str=traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, - ), + traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + await _dispatch_passthrough_dynamic_failure( + logging_obj=logging_obj, + exception=e, + traceback_str=traceback_str, ) + try: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=request_payload, + traceback_str=traceback_str, + ) + except Exception: # noqa: BLE001 - a failing logging callback must never change the WebSocket close behavior + verbose_proxy_logger.warning( + "pass_through_endpoint: post_call_failure_hook raised for WebSocket failure", + exc_info=True, + ) if _client_socket_is_open(websocket): await websocket.close(code=1011, reason="WebSocket passthrough error") diff --git a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py index 523b669280e..d62942c1d73 100644 --- a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py @@ -22,7 +22,7 @@ from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers -from litellm.proxy.litellm_pre_call_utils import _get_dynamic_logging_metadata +from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( create_pass_through_route, ) @@ -175,7 +175,7 @@ async def langfuse_proxy_route( user_api_key_dict: Final = await user_api_key_auth(request=request, api_key=f"Bearer {api_key}") - callback_settings_obj: Final[TeamCallbackMetadata | None] = _get_dynamic_logging_metadata( + callback_settings_obj: Final[TeamCallbackMetadata | None] = get_dynamic_logging_metadata( user_api_key_dict=user_api_key_dict, proxy_config=proxy_config ) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 3bde72ccd49..7e9a062ba18 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -22,7 +22,7 @@ from litellm.proxy.auth.auth_utils import ( is_request_body_safe, ) from litellm.proxy.litellm_pre_call_utils import ( - _get_dynamic_logging_metadata, + get_dynamic_logging_metadata, add_litellm_data_to_request, ) from pydantic import ValidationError @@ -294,7 +294,7 @@ def test_dynamic_logging_metadata_key_and_team_metadata(callback_vars): rpm_limit_per_model=None, tpm_limit_per_model=None, ) - callbacks = _get_dynamic_logging_metadata( + callbacks = get_dynamic_logging_metadata( user_api_key_dict=user_api_key_dict, proxy_config=proxy_config ) @@ -332,7 +332,7 @@ def test_dynamic_logging_metadata_ignores_env_references_from_key_metadata( team_metadata={}, ) - callbacks = _get_dynamic_logging_metadata( + callbacks = get_dynamic_logging_metadata( user_api_key_dict=user_api_key_dict, proxy_config=proxy_config ) @@ -410,7 +410,7 @@ def test_dynamic_turn_off_message_logging(callback_vars): rpm_limit_per_model=None, tpm_limit_per_model=None, ) - callbacks = _get_dynamic_logging_metadata( + callbacks = get_dynamic_logging_metadata( user_api_key_dict=user_api_key_dict, proxy_config=proxy_config ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py index 08e931e6405..821a68d386c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py @@ -539,7 +539,7 @@ async def test_get_team_callbacks_returns_callbacks_registered_via_post(monkeypa async def test_get_team_callbacks_prefers_logging_over_deprecated_callback_settings(): """A team carrying both shapes must report only the one that actually fires. - Request-time resolution in _get_dynamic_logging_metadata stops at the + Request-time resolution in get_dynamic_logging_metadata stops at the first populated slot: metadata["logging"] wins and callback_settings is never consulted. Reporting the union here would tell an operator that gcs_bucket is active on a team whose requests never send to it. @@ -752,7 +752,7 @@ async def test_disable_team_logging_stops_callbacks_registered_via_api(): the endpoint and then asks the real request-time resolver what the written row would do. """ - from litellm.proxy.litellm_pre_call_utils import _get_dynamic_logging_metadata + from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata metadata = { "logging": [ @@ -780,7 +780,7 @@ async def test_disable_team_logging_stops_callbacks_registered_via_api(): written = json.loads(mock_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"]) assert written["logging"] == [] - resolved = _get_dynamic_logging_metadata( + resolved = get_dynamic_logging_metadata( UserAPIKeyAuth(api_key="hashed", team_id="team-1", team_metadata=written), proxy_config=MagicMock(**{"load_team_config.return_value": {}}), ) @@ -859,7 +859,7 @@ async def test_add_team_callbacks_refreshes_cached_team(stub_team_cache_refresh) @pytest.mark.asyncio async def test_disable_team_logging_clears_both_metadata_shapes(): """A team carrying both shapes ends up with neither active.""" - from litellm.proxy.litellm_pre_call_utils import _get_dynamic_logging_metadata + from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata metadata = { "logging": [ @@ -893,7 +893,7 @@ async def test_disable_team_logging_clears_both_metadata_shapes(): assert written["callback_settings"]["success_callback"] == [] assert written["callback_settings"]["failure_callback"] == [] - resolved = _get_dynamic_logging_metadata( + resolved = get_dynamic_logging_metadata( UserAPIKeyAuth(api_key="hashed", team_id="team-1", team_metadata=written), proxy_config=MagicMock(**{"load_team_config.return_value": {}}), ) @@ -1023,7 +1023,7 @@ async def test_delete_team_callback_leaves_the_other_callback_firing(): Asks the real request-time resolver what the written row would do, the same way the disable_logging regression test does. """ - from litellm.proxy.litellm_pre_call_utils import _get_dynamic_logging_metadata + from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata mock_prisma = _patch_prisma(_team_row(team_id="team-1", metadata=_two_callback_metadata())) @@ -1040,7 +1040,7 @@ async def test_delete_team_callback_leaves_the_other_callback_firing(): ) written = json.loads(mock_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"]) - resolved = _get_dynamic_logging_metadata( + resolved = get_dynamic_logging_metadata( UserAPIKeyAuth(api_key="hashed", team_id="team-1", team_metadata=written), proxy_config=MagicMock(**{"load_team_config.return_value": {}}), ) @@ -1215,7 +1215,7 @@ async def test_delete_team_callback_keeps_last_removal_from_reviving_legacy_shap dropping the key would fall through to a legacy callback_settings block and silently re-enable a destination the caller just removed. """ - from litellm.proxy.litellm_pre_call_utils import _get_dynamic_logging_metadata + from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata metadata = { "logging": [ @@ -1249,7 +1249,7 @@ async def test_delete_team_callback_keeps_last_removal_from_reviving_legacy_shap assert written["logging"] == [] assert response.data.success_callbacks == () - resolved = _get_dynamic_logging_metadata( + resolved = get_dynamic_logging_metadata( UserAPIKeyAuth(api_key="hashed", team_id="team-1", team_metadata=written), proxy_config=MagicMock(**{"load_team_config.return_value": {}}), ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index a3f56adb86f..4514784cc16 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -14,30 +14,28 @@ from fastapi import Request, UploadFile from starlette.datastructures import FormData, Headers, QueryParams from starlette.datastructures import UploadFile as StarletteUploadFile - +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, + LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, - LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, _registered_pass_through_routes, create_pass_through_route, initialize_pass_through_endpoints, pass_through_request, - resolve_pass_through_request_timeout, resolve_llm_passthrough_timeout, + resolve_pass_through_request_timeout, websocket_passthrough_request, ) -from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import ProxyException, UserAPIKeyAuth -from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, -) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) - -import litellm +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, +) MESSAGE_START_SSE_FRAME = b'event: message_start\ndata: {"type": "message_start"}\n\n' @@ -1308,7 +1306,14 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): mock_request.method = "POST" mock_request.url = "http://test-proxy.com/api/endpoint" mock_request.body = AsyncMock(return_value=b'{"message": "test request"}') - mock_request.headers = Headers({}) + mock_request.headers = Headers( + { + "authorization": "Bearer caller-secret", + "x-api-key": "caller-secret", + "x-goog-api-key": "caller-secret", + "langfuse_trace_id": "preserved-context", + } + ) mock_request.query_params = QueryParams({}) # Create mock user API key dict @@ -1355,6 +1360,12 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): assert proxy_server_request["method"] == mock_request.method # The body should be the value returned by pre_call_hook, not the original request body assert proxy_server_request["body"] == {"test": "data"} + assert proxy_server_request["headers"] == { + "authorization": "***REDACTED***", + "x-api-key": "***REDACTED***", + "x-goog-api-key": "***REDACTED***", + "langfuse_trace_id": "preserved-context", + } # Verify other required kwargs are present assert "call_type" in call_kwargs @@ -4096,8 +4107,10 @@ class _FakeUpstreamTransport(httpx.AsyncBaseTransport): self._status_code = status_code self._headers = headers self._stream = stream + self.request = None async def handle_async_request(self, request): + self.request = request return httpx.Response( status_code=self._status_code, headers=self._headers, @@ -5369,6 +5382,381 @@ def test_passthrough_budget_metadata_cannot_be_forged_by_the_request_body(): assert metadata["user_api_key_model_max_budget"] == key_budget +def _team_scoped_langfuse_otel_key(callback_type: str = "success") -> UserAPIKeyAuth: + return UserAPIKeyAuth( + team_id="team-id", + team_metadata={ + "logging": [ + { + "callback_name": "langfuse_otel", + "callback_type": callback_type, + "callback_vars": { + "langfuse_public_key": "team-public-key", + "langfuse_secret_key": "team-secret-key", + "langfuse_host": "https://team.langfuse.example", + }, + }, + ] + }, + ) + + +@pytest.mark.asyncio +async def test_pass_through_request_initializes_team_logging_before_dispatch(): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + fake_client, cleanup = _inject_fake_passthrough_client( + _FakeUpstreamTransport( + status_code=200, + headers={"content-type": "application/json"}, + stream=_RecordingUpstreamByteStream((b'{"answer": "ok"}',)), + ), + timeout=None, + ) + + try: + with ExitStack() as stack: + mock_proxy_logging, mock_success_handler = _enter_relay_logging_mocks(stack, {}) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_proxy_logging.get_proxy_hook = MagicMock(return_value=None) + enqueued = [] + stack.enter_context( + patch.object( + GLOBAL_LOGGING_WORKER, + "ensure_initialized_and_enqueue", + new=MagicMock(side_effect=lambda async_coroutine: enqueued.append(async_coroutine)), + ) + ) + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", MagicMock())) + + response = await pass_through_request( + request=_relay_client_request(method="POST"), + target="http://upstream.test/v1/chat/completions", + custom_headers={}, + user_api_key_dict=_team_scoped_langfuse_otel_key(), + ) + + assert response.status_code == 200 + pre_call_logging_obj = mock_proxy_logging.pre_call_hook.await_args.kwargs["data"]["litellm_logging_obj"] + success_logging_obj = mock_success_handler.call_args.kwargs["logging_obj"] + assert pre_call_logging_obj is success_logging_obj + assert success_logging_obj.standard_callback_dynamic_params == { + "langfuse_public_key": "team-public-key", + "langfuse_secret_key": "team-secret-key", + "langfuse_host": "https://team.langfuse.example", + } + assert success_logging_obj._trusted_callback_vars == ( + ("langfuse_public_key", "team-public-key"), + ("langfuse_secret_key", "team-secret-key"), + ("langfuse_host", "https://team.langfuse.example"), + ) + assert [callback.callback_name for callback in success_logging_obj.dynamic_success_callbacks] == [ + "langfuse_otel" + ] + assert [callback.callback_name for callback in success_logging_obj.dynamic_async_success_callbacks] == [ + "langfuse_otel" + ] + for async_coroutine in enqueued: + async_coroutine.close() + finally: + cleanup() + await fake_client.aclose() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_initializes_team_logging_before_dispatch(): + upstream_ws = FakeUpstreamWebSocket(b'{"type": "session.created"}') + websocket = _client_websocket(AsyncMock(return_value={"type": "websocket.disconnect"})) + + with ExitStack() as stack: + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", MagicMock())) + mock_success_handler = stack.enter_context( + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "pass_through_endpoint_logging.pass_through_async_success_handler", + new=AsyncMock(), + ) + ) + with _patched_websocket_passthrough_environment(upstream_ws): + await websocket_passthrough_request( + websocket=websocket, + target="wss://upstream.test/v1/realtime", + custom_headers={}, + user_api_key_dict=_team_scoped_langfuse_otel_key(), + endpoint="/openai/v1/realtime", + ) + + success_logging_obj = mock_success_handler.call_args.kwargs["logging_obj"] + assert success_logging_obj.standard_callback_dynamic_params == { + "langfuse_public_key": "team-public-key", + "langfuse_secret_key": "team-secret-key", + "langfuse_host": "https://team.langfuse.example", + } + assert success_logging_obj._trusted_callback_vars == ( + ("langfuse_public_key", "team-public-key"), + ("langfuse_secret_key", "team-secret-key"), + ("langfuse_host", "https://team.langfuse.example"), + ) + assert [callback.callback_name for callback in success_logging_obj.dynamic_success_callbacks] == ["langfuse_otel"] + assert [callback.callback_name for callback in success_logging_obj.dynamic_async_success_callbacks] == [ + "langfuse_otel" + ] + + +class _FailureCallbackRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.failure_event_kwargs: list[dict] = [] + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self.failure_event_kwargs.append(kwargs) + + +@pytest.mark.asyncio +async def test_pass_through_request_dispatches_team_failure_callback_once_for_upstream_error(): + failure_callback_recorder = _FailureCallbackRecorder() + upstream_transport = _FakeUpstreamTransport( + status_code=403, + headers={"content-type": "application/json"}, + stream=_RecordingUpstreamByteStream((b'{"error": "upstream denied"}',)), + ) + fake_client, cleanup = _inject_fake_passthrough_client(upstream_transport, timeout=None) + request = _relay_client_request(method="POST") + request.headers = Headers( + { + "Authorization": "Bearer provider-secret", + "x-api-key": "provider-api-key", + "X-Custom-LiteLLM-Key": "virtual-key-secret", + "x-request-id": "request-123", + } + ) + + try: + with ExitStack() as stack: + mock_proxy_logging, mock_success_handler = _enter_relay_logging_mocks(stack, {}) + mock_proxy_logging.get_proxy_hook = MagicMock(return_value=None) + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", MagicMock())) + stack.enter_context( + patch("litellm.proxy.proxy_server.general_settings", {"litellm_key_header_name": "x-custom-litellm-key"}) + ) + stack.enter_context( + patch.object( + litellm, + "_known_custom_logger_compatible_callbacks", + ["langfuse_otel"], + ) + ) + stack.enter_context( + patch( + "litellm.litellm_core_utils.litellm_logging." + "_init_custom_logger_compatible_class", + return_value=failure_callback_recorder, + ) + ) + + response = await pass_through_request( + request=request, + target="http://upstream.test/v1/chat/completions", + custom_headers={}, + user_api_key_dict=_team_scoped_langfuse_otel_key(callback_type="failure"), + forward_headers=True, + ) + + assert response.status_code == 403 + assert json.loads(response.body) == {"error": "upstream denied"} + assert len(failure_callback_recorder.failure_event_kwargs) == 1 + callback_kwargs = failure_callback_recorder.failure_event_kwargs[0] + assert callback_kwargs["has_logged_async_failure"] is True + callback_headers = callback_kwargs["additional_args"]["headers"] + assert callback_headers == { + "authorization": "***REDACTED***", + "x-api-key": "***REDACTED***", + "x-custom-litellm-key": "***REDACTED***", + "x-request-id": "request-123", + } + serialized_callback_kwargs = json.dumps(callback_kwargs, default=str) + assert "provider-secret" not in serialized_callback_kwargs + assert "provider-api-key" not in serialized_callback_kwargs + assert "virtual-key-secret" not in serialized_callback_kwargs + assert upstream_transport.request is not None + assert upstream_transport.request.headers["authorization"] == "Bearer provider-secret" + assert upstream_transport.request.headers["x-api-key"] == "provider-api-key" + assert upstream_transport.request.headers["X-Custom-LiteLLM-Key"] == "virtual-key-secret" + assert request.headers["Authorization"] == "Bearer provider-secret" + mock_success_handler.assert_not_called() + mock_proxy_logging.post_call_failure_hook.assert_awaited_once() + finally: + cleanup() + await fake_client.aclose() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_dispatches_team_failure_callback_once(): + failure_callback_recorder = _FailureCallbackRecorder() + websocket = _client_websocket(AsyncMock(return_value={"type": "websocket.disconnect"})) + + with ExitStack() as stack: + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", MagicMock())) + stack.enter_context( + patch.object( + litellm, + "_known_custom_logger_compatible_callbacks", + ["langfuse_otel"], + ) + ) + stack.enter_context( + patch( + "litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class", + return_value=failure_callback_recorder, + ) + ) + with _patched_websocket_passthrough_environment(ClosingUpstreamWebSocket(RuntimeError("upstream failed"))): + await websocket_passthrough_request( + websocket=websocket, + target="wss://upstream.test/v1/realtime", + custom_headers={}, + user_api_key_dict=_team_scoped_langfuse_otel_key(callback_type="failure"), + endpoint="/openai/v1/realtime", + ) + + assert len(failure_callback_recorder.failure_event_kwargs) == 1 + assert failure_callback_recorder.failure_event_kwargs[0]["additional_args"]["headers"] == {} + + +@pytest.mark.asyncio +async def test_websocket_passthrough_dispatches_team_failure_callback_when_upstream_rejects_handshake(): + from websockets.datastructures import Headers as WebSocketHeaders + from websockets.exceptions import InvalidStatus + from websockets.http11 import Response as WebSocketResponse + + failure_callback_recorder = _FailureCallbackRecorder() + websocket = _client_websocket(AsyncMock(return_value={"type": "websocket.disconnect"})) + upstream_response = WebSocketResponse( + status_code=403, + reason_phrase="Forbidden", + headers=WebSocketHeaders(), + ) + + with ExitStack() as stack: + mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj")) + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + stack.enter_context( + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", + side_effect=InvalidStatus(upstream_response), + ) + ) + mock_worker = stack.enter_context( + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") + ) + mock_worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: async_coroutine.close() + ) + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", MagicMock())) + stack.enter_context( + patch.object( + litellm, + "_known_custom_logger_compatible_callbacks", + ["langfuse_otel"], + ) + ) + stack.enter_context( + patch( + "litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class", + return_value=failure_callback_recorder, + ) + ) + + await websocket_passthrough_request( + websocket=websocket, + target="wss://upstream.test/v1/realtime", + custom_headers={}, + user_api_key_dict=_team_scoped_langfuse_otel_key(callback_type="failure"), + endpoint="/openai/v1/realtime", + ) + + assert len(failure_callback_recorder.failure_event_kwargs) == 1 + mock_proxy_logging.post_call_failure_hook.assert_awaited_once() + websocket.close.assert_awaited_once_with(code=1011, reason="Upstream connection rejected") + + +@pytest.mark.asyncio +async def test_websocket_passthrough_forwards_credentials_without_exposing_them_to_failure_callback(): + failure_callback_recorder = _FailureCallbackRecorder() + websocket = _client_websocket(AsyncMock(return_value={"type": "websocket.disconnect"})) + websocket.headers = Headers( + { + "Authorization": "Bearer provider-secret", + "x-api-key": "provider-api-key", + "x-request-id": "request-123", + } + ) + upstream_ws = ClosingUpstreamWebSocket(RuntimeError("upstream failed")) + + with ExitStack() as stack: + mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj")) + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_connect = stack.enter_context( + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", + return_value=FakeUpstreamConnect(upstream_ws), + ) + ) + mock_worker = stack.enter_context( + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") + ) + mock_worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: async_coroutine.close() + ) + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", MagicMock())) + stack.enter_context( + patch("litellm.proxy.proxy_server.general_settings", {"litellm_key_header_name": "x-custom-litellm-key"}) + ) + stack.enter_context( + patch.object( + litellm, + "_known_custom_logger_compatible_callbacks", + ["langfuse_otel"], + ) + ) + stack.enter_context( + patch( + "litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class", + return_value=failure_callback_recorder, + ) + ) + + await websocket_passthrough_request( + websocket=websocket, + target="wss://upstream.test/v1/realtime", + custom_headers={"X-Custom-LiteLLM-Key": "virtual-key-secret"}, + user_api_key_dict=_team_scoped_langfuse_otel_key(callback_type="failure"), + endpoint="/openai/v1/realtime", + forward_headers=True, + ) + + assert len(failure_callback_recorder.failure_event_kwargs) == 1 + callback_kwargs = failure_callback_recorder.failure_event_kwargs[0] + callback_headers = callback_kwargs["additional_args"]["headers"] + assert callback_headers == { + "X-Custom-LiteLLM-Key": "***REDACTED***", + "authorization": "***REDACTED***", + "x-api-key": "***REDACTED***", + } + serialized_callback_kwargs = json.dumps(callback_kwargs, default=str) + assert "provider-secret" not in serialized_callback_kwargs + assert "provider-api-key" not in serialized_callback_kwargs + assert "virtual-key-secret" not in serialized_callback_kwargs + assert mock_connect.call_args.kwargs["additional_headers"] == { + "X-Custom-LiteLLM-Key": "virtual-key-secret", + "authorization": "Bearer provider-secret", + "x-api-key": "provider-api-key", + } + assert websocket.headers["Authorization"] == "Bearer provider-secret" + + def _marked_pass_through_endpoint(): """An endpoint carrying the marker ``create_pass_through_route`` sets.""" from litellm.types.passthrough_endpoints.pass_through_endpoints import ( 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 a10f3672dff..6a98b589f16 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -20,7 +20,7 @@ from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, _apply_credential_overrides_from_model_config, _extract_credential_from_entry, - _get_dynamic_logging_metadata, + get_dynamic_logging_metadata, _get_enforced_params, _get_metadata_variable_name, _promoted_trace_control_fields, @@ -2311,7 +2311,7 @@ def test_team_dynamic_logging_settings_decrypts_callback_vars(monkeypatch): def test_get_dynamic_logging_metadata_with_arize_team_logging(): """ - Test _get_dynamic_logging_metadata function with arize team logging and dynamic parameters + Test get_dynamic_logging_metadata function with arize team logging and dynamic parameters """ # Setup user with arize team logging including callback_vars user_api_key_dict = UserAPIKeyAuth( @@ -2335,7 +2335,7 @@ def test_get_dynamic_logging_metadata_with_arize_team_logging(): mock_proxy_config = MagicMock() # Call the function - result = _get_dynamic_logging_metadata( + result = get_dynamic_logging_metadata( user_api_key_dict=user_api_key_dict, proxy_config=mock_proxy_config ) @@ -2386,7 +2386,7 @@ def test_get_dynamic_logging_metadata_ignores_env_reference_from_key_metadata( team_metadata={}, ) - result = _get_dynamic_logging_metadata( + result = get_dynamic_logging_metadata( user_api_key_dict=user_api_key_dict, proxy_config=MagicMock() ) @@ -6569,6 +6569,22 @@ def test_redact_credential_headers_classifies_each_header(header, expected_redac assert headers[header] == "secret-value" +def test_redact_credential_headers_redacts_configured_litellm_key_header_without_mutating_input(): + from litellm.proxy.litellm_pre_call_utils import redact_credential_headers + + headers = { + "X-Custom-LiteLLM-Key": "virtual-key-secret", + "x-request-id": "request-123", + } + + redacted = redact_credential_headers(headers, litellm_key_header_name="x-custom-litellm-key") + + assert redacted == { + "X-Custom-LiteLLM-Key": "***REDACTED***", + "x-request-id": "request-123", + } + assert headers["X-Custom-LiteLLM-Key"] == "virtual-key-secret" + @pytest.mark.asyncio async def test_add_litellm_data_to_request_debug_log_does_not_print_credentials():