diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 30aaaacad04..5fb26f89db0 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -1,11 +1,11 @@ # What is this? ## Helper utils for the management endpoints (keys/users/teams) -from collections.abc import Callable, Mapping, MutableMapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence from collections.abc import Set as AbstractSet from datetime import datetime from functools import wraps from types import MappingProxyType -from typing import Any, Final, Protocol +from typing import Any, Final, Protocol, TypeVar, cast from fastapi import HTTPException, Request from pydantic import BaseModel, TypeAdapter @@ -767,19 +767,31 @@ async def _emit_management_endpoint_otel_span( ) -def management_endpoint_wrapper(func): +_Endpoint: Final = TypeVar("_Endpoint", bound=Callable[..., Awaitable[Any]]) + + +def management_endpoint_wrapper(func: _Endpoint) -> _Endpoint: """ This wrapper does the following: 1. Log I/O, Exceptions to OTEL 2. Create an Audit log for success calls """ + return _wrap_management_endpoint(func, log_payload=True) - @wraps(func) + +def management_endpoint_span_wrapper(func: _Endpoint) -> _Endpoint: + """``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 _wrap_management_endpoint(endpoint: _Endpoint, *, log_payload: bool) -> _Endpoint: + @wraps(endpoint) async def wrapper(*args, **kwargs): start_time: Final = datetime.now() try: - result: Final = await func(*args, **kwargs) + result: Final = await endpoint(*args, **kwargs) end_time = datetime.now() try: user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict") or UserAPIKeyAuth() @@ -787,17 +799,17 @@ def management_endpoint_wrapper(func): await send_management_endpoint_alert( request_kwargs=kwargs, user_api_key_dict=user_api_key_dict, - function_name=func.__name__, + function_name=endpoint.__name__, ) parent_otel_span: Span | None = getattr(user_api_key_dict, "parent_otel_span", None) if parent_otel_span is not None: await _emit_management_endpoint_otel_span( - func=func, + func=endpoint, kwargs=kwargs, parent_otel_span=parent_otel_span, start_time=start_time, end_time=end_time, - result=result, + result=result if log_payload else None, ) # Delete updated/deleted info from cache @@ -818,7 +830,7 @@ def management_endpoint_wrapper(func): if parent_otel_span is not None: try: await _emit_management_endpoint_otel_span( - func=func, + func=endpoint, kwargs=kwargs, parent_otel_span=parent_otel_span, start_time=start_time, @@ -835,4 +847,4 @@ def management_endpoint_wrapper(func): raise e - return wrapper + return cast(_Endpoint, wrapper) # cast-ok: wraps() preserves the endpoint's signature diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 257740b4c7e..4d3a569bfdd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -174,6 +174,7 @@ from litellm.proxy.common_utils.token_count_settings import ( ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields +from litellm.proxy.management_helpers.utils import management_endpoint_span_wrapper from litellm.repositories import database as database_runtime from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, @@ -12343,6 +12344,7 @@ async def _deployment_hidden_by_listing_callbacks(deployment: Deployment, user_a @router.get( "/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"] ) # if project requires model list +@management_endpoint_span_wrapper async def model_list( request: Request = None, # pyright: ignore[reportArgumentType] # FastAPI always injects the Request; the None default only serves direct in-process callers user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 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 ebeb514ffff..5ed512492c5 100644 --- a/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py +++ b/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py @@ -340,3 +340,80 @@ def test_management_wrapper_swallows_post_success_errors(server_span_factory, ot result = asyncio.run(fake_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span))) 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.""" + 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) + + @mgmt_utils.management_endpoint_span_wrapper + async def fake_fn(data=None, user_api_key_dict=None): + return {"data": [{"id": "gpt-4o"}]} + + asyncio.run(fake_fn(data={}, 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"] + assert len(spans) == 1 + assert not any(key.startswith("response.") for key in spans[0].attributes) + + +MODEL_LIST_PATH = "/v1/models" +INBOUND_TRACEPARENT = "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( + 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( + proxy_server.model_list( + user_api_key_dict=_real_user_api_key_dict(server_span), + ) + ) + + assert result == {"data": [], "object": "list"} + assert_server_span_attrs( + exporter, + expected_status=200, + expected_url_path=MODEL_LIST_PATH, + where="model list success", + ) + exported = 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"] + assert len(endpoint_spans) == 1 + assert endpoint_spans[0].parent is not None + assert endpoint_spans[0].parent.span_id == exported.context.span_id + assert not any(key.startswith("response.") for key in endpoint_spans[0].attributes)