scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

fix(proxy): dispatch team callbacks for passthrough failures

Browse source
This commit is contained in:
Yucheng Zhu 2026-08-28 09:59:26 -07:00
parent 5bcd494e88
commit c8ae1cf4ed
9 changed files with 601 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
basedpyright-code-budget.json
View file

@ -84,7 +84,7 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 1808
"limit": 1806
},
"reportRedeclaration": {
"limit": 8

16
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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:

4
litellm/proxy/management_endpoints/team_callback_endpoints.py
View file

@ -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)

181
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -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")

4
litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py
View file

@ -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
)

8
tests/proxy_unit_tests/test_proxy_utils.py
View file

@ -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
)

18
tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py
View file

@ -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": {}}),
)

410
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -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 (

24
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -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():