scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(proxy): keep request bodies off the model list span and tidy its tests

Browse source 
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
This commit is contained in:
mateo-berri 2026-10-30 07:46:50 -07:00 • committed by devin-ai-integration[bot]
parent 38e6292a85
commit a8c17ed63f
2 changed files with 68 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
litellm/proxy/management_helpers/utils.py
View file

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

78
tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py
View file

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