From a8c17ed63f343d6b85195ca35433cdb39f01c110 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:46:50 -0700 Subject: [PATCH] fix(proxy): keep request bodies off the model list span and tidy its tests The quiet span wrapper now skips the request body as well as the response, so no payload reaches the span on success or failure. The wrapper passes the endpoint's kwargs through instead of copying them, and reads the caller's key with an isinstance check. The model list test drives the real endpoint without swapping helpers, and the new payload test checks both wrappers against a real request body --- litellm/proxy/management_helpers/utils.py | 37 +++++---- .../test_otel_admin_endpoints.py | 78 +++++++++++-------- 2 files changed, 68 insertions(+), 47 deletions(-) diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 924bcbe496e..f514f233b0a 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -687,6 +687,7 @@ async def _emit_management_endpoint_otel_span( end_time: datetime, result: Any = None, exception: Exception | None = None, + log_payload: bool = True, ) -> None: """Stamp + end the parent OTEL SERVER span for a management endpoint. @@ -716,8 +717,8 @@ async def _emit_management_endpoint_otel_span( ) route = get_request_route(http_request) - request_body: dict = await read_request_body( # rebind-ok: pre-existing rebinding on a rename-only line - request=http_request + request_body: dict = ( # rebind-ok: pre-existing rebinding on a rename-only line + await read_request_body(request=http_request) if log_payload else {} ) else: route = func.__name__ @@ -738,7 +739,7 @@ async def _emit_management_endpoint_otel_span( ) _response: dict[str, object] | None = None - if exception is None and result is not None: + if log_payload and exception is None and result is not None: try: raw: Final[Mapping[str, object]] = dict(result) _response = {k: v for k, v in raw.items() if k not in _CREDENTIAL_FIELDS} @@ -782,26 +783,28 @@ def management_endpoint_wrapper(func: Callable[_P, Awaitable[_R]]) -> Callable[_ def management_endpoint_span_wrapper(func: Callable[_P, Awaitable[_R]]) -> Callable[_P, Awaitable[_R]]: - """``management_endpoint_wrapper`` without request/response payloads on the span, for - endpoints that serialize large or high-frequency responses (e.g. the model list).""" return _wrap_management_endpoint(func, log_payload=False) +def _caller_auth(endpoint_kwargs: Mapping[str, object]) -> UserAPIKeyAuth: + caller: Final = endpoint_kwargs.get("user_api_key_dict") + return caller if isinstance(caller, UserAPIKeyAuth) else UserAPIKeyAuth() + + def _wrap_management_endpoint( endpoint: Callable[_P, Awaitable[_R]], *, log_payload: bool ) -> Callable[_P, Awaitable[_R]]: @wraps(endpoint) async def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: - request_kwargs: Final[dict[str, Any]] = dict(kwargs) # mutable-ok: passed to dict-typed helpers start_time: Final = datetime.now() try: result: Final = await endpoint(*args, **kwargs) end_time = datetime.now() try: - user_api_key_dict: UserAPIKeyAuth = request_kwargs.get("user_api_key_dict") or UserAPIKeyAuth() + user_api_key_dict: UserAPIKeyAuth = _caller_auth(kwargs) await send_management_endpoint_alert( - request_kwargs=request_kwargs, + request_kwargs=kwargs, user_api_key_dict=user_api_key_dict, function_name=endpoint.__name__, ) @@ -809,18 +812,19 @@ def _wrap_management_endpoint( if parent_otel_span is not None: await _emit_management_endpoint_otel_span( func=endpoint, - kwargs=request_kwargs, + kwargs=kwargs, parent_otel_span=parent_otel_span, start_time=start_time, end_time=end_time, - result=result if log_payload else None, + result=result, + log_payload=log_payload, ) # Delete updated/deleted info from cache - _delete_api_key_from_cache(kwargs=request_kwargs) - _delete_user_id_from_cache(kwargs=request_kwargs) - _delete_team_id_from_cache(kwargs=request_kwargs) - _delete_customer_id_from_cache(kwargs=request_kwargs) + _delete_api_key_from_cache(kwargs=kwargs) + _delete_user_id_from_cache(kwargs=kwargs) + _delete_team_id_from_cache(kwargs=kwargs) + _delete_customer_id_from_cache(kwargs=kwargs) except Exception as e: # Non-Blocking Exception verbose_logger.debug("Error in management endpoint wrapper: %s", str(e)) @@ -829,17 +833,18 @@ def _wrap_management_endpoint( except Exception as e: end_time = datetime.now() - user_api_key_dict: UserAPIKeyAuth = request_kwargs.get("user_api_key_dict") or UserAPIKeyAuth() + user_api_key_dict: UserAPIKeyAuth = _caller_auth(kwargs) parent_otel_span = getattr(user_api_key_dict, "parent_otel_span", None) if parent_otel_span is not None: try: await _emit_management_endpoint_otel_span( func=endpoint, - kwargs=request_kwargs, + kwargs=kwargs, parent_otel_span=parent_otel_span, start_time=start_time, end_time=end_time, exception=e, + log_payload=log_payload, ) except Exception as otel_exc: # Non-Blocking Exception - never let OTEL failures swallow diff --git a/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py b/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py index 5ed512492c5..778845f6b62 100644 --- a/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py +++ b/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py @@ -3,9 +3,11 @@ async_management_endpoint_{success,failure}_hook integration points.""" import asyncio from datetime import datetime +from typing import Final from unittest.mock import MagicMock import pytest +from fastapi import Request from litellm.proxy._types import ( ManagementEndpointLoggingPayload, @@ -342,60 +344,74 @@ def test_management_wrapper_swallows_post_success_errors(server_span_factory, ot assert result == {"ok": True} -def test_span_wrapper_omits_payload(server_span_factory, otel_with_exporter, monkeypatch): - """The span-only wrapper keeps the endpoint span but drops the response - attributes — the models listing would otherwise serialize every model onto - the span.""" +TEAM_NEW_PATH: Final = "/team/new" + + +def _json_post(path: str, body: bytes) -> Request: + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": path, + "root_path": "", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + + +@pytest.mark.parametrize( + ("wrapper_name", "exports_payload"), + [("management_endpoint_wrapper", True), ("management_endpoint_span_wrapper", False)], +) +def test_span_wrapper_omits_payload(server_span_factory, otel_with_exporter, monkeypatch, wrapper_name, exports_payload): import litellm.proxy.proxy_server as proxy_server from litellm.proxy.management_helpers import utils as mgmt_utils otel, exporter = otel_with_exporter monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False) - monkeypatch.setattr(mgmt_utils, "send_management_endpoint_alert", _noop_alert) - server_span = server_span_factory(KEY_GENERATE_PATH) + server_span: Final = server_span_factory(TEAM_NEW_PATH) + wrap: Final = getattr(mgmt_utils, wrapper_name) - @mgmt_utils.management_endpoint_span_wrapper - async def fake_fn(data=None, user_api_key_dict=None): + @wrap + async def fake_fn(http_request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]: return {"data": [{"id": "gpt-4o"}]} - asyncio.run(fake_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span))) + asyncio.run( + fake_fn( + http_request=_json_post(TEAM_NEW_PATH, b'{"team_alias": "t"}'), + user_api_key_dict=_real_user_api_key_dict(server_span), + ) + ) - spans = [s for s in exporter.get_finished_spans() if s.name == "fake_fn"] + spans: Final = [s for s in exporter.get_finished_spans() if s.name == TEAM_NEW_PATH] assert len(spans) == 1 - assert not any(key.startswith("response.") for key in spans[0].attributes) + attribute_keys: Final = tuple(spans[0].attributes) + assert any(key.startswith("request.") for key in attribute_keys) is exports_payload + assert any(key.startswith("response.") for key in attribute_keys) is exports_payload -MODEL_LIST_PATH = "/v1/models" -INBOUND_TRACEPARENT = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" +MODEL_LIST_PATH: Final = "/v1/models" +INBOUND_TRACEPARENT: Final = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" def test_model_list_ends_server_span_in_inbound_trace(otel_with_exporter, monkeypatch): - """Regression: ``/v1/models`` created the OTEL SERVER span at the auth - boundary but never ended it, so the span was never exported and a caller - passing ``traceparent`` saw nothing for the models list. - - Drives the real endpoint and asserts the SERVER span reaches the exporter - under the caller's trace id.""" import litellm.proxy.proxy_server as proxy_server - import litellm.proxy.utils as proxy_utils - from litellm.proxy.management_helpers import utils as mgmt_utils otel, exporter = otel_with_exporter monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False) - monkeypatch.setattr(mgmt_utils, "send_management_endpoint_alert", _noop_alert) - async def _no_models_for_user(**kwargs): - return [] - - monkeypatch.setattr(proxy_utils, "get_available_models_for_user", _no_models_for_user) - - server_span = otel.create_litellm_proxy_request_started_span( + server_span: Final = otel.create_litellm_proxy_request_started_span( start_time=datetime.now(), headers={"traceparent": INBOUND_TRACEPARENT} ) otel.set_proxy_request_route_attributes(server_span, url_path=MODEL_LIST_PATH, http_route=MODEL_LIST_PATH) - result = asyncio.run( + result: Final = asyncio.run( proxy_server.model_list( user_api_key_dict=_real_user_api_key_dict(server_span), ) @@ -408,11 +424,11 @@ def test_model_list_ends_server_span_in_inbound_trace(otel_with_exporter, monkey expected_url_path=MODEL_LIST_PATH, where="model list success", ) - exported = get_server_span(exporter) + exported: Final = get_server_span(exporter) assert exported is not None assert exported.context.trace_id == int(INBOUND_TRACEPARENT.split("-")[1], 16) - endpoint_spans = [s for s in exporter.get_finished_spans() if s.name == "model_list"] + endpoint_spans: Final = [s for s in exporter.get_finished_spans() if s.name == "model_list"] assert len(endpoint_spans) == 1 assert endpoint_spans[0].parent is not None assert endpoint_spans[0].parent.span_id == exported.context.span_id