From 30551de371fe37b1ee8de1e5fe949ddcdad272e8 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 25 May 2026 12:13:17 -0700 Subject: [PATCH] fix(otel): export SERVER span on management-endpoint success without http_request (#28794) Co-authored-by: Yassin Kortam --- litellm/proxy/management_helpers/utils.py | 124 ++++++------ .../test_otel_admin_endpoints.py | 176 ++++++++++++++++++ 2 files changed, 245 insertions(+), 55 deletions(-) diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index b7d5cc30c4..ab100615b4 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -2,7 +2,7 @@ ## Helper utils for the management endpoints (keys/users/teams) from datetime import datetime from functools import wraps -from typing import List, Optional, Tuple +from typing import Any, Callable, List, Optional, Tuple from fastapi import HTTPException, Request @@ -435,6 +435,58 @@ async def send_management_endpoint_alert( ) +async def _emit_management_endpoint_otel_span( + func: Callable, + kwargs: dict, + parent_otel_span: Any, + start_time: datetime, + end_time: datetime, + result: Any = None, + exception: Optional[Exception] = None, +) -> None: + """Stamp + end the parent OTEL SERVER span for a management endpoint. + + Routes the request/response (or exception) through the OTEL success/failure + hook. Falls back to ``func.__name__`` for the route when the handler has no + ``http_request`` param — endpoints like ``/key/generate`` never receive one, + and gating the hook on it leaked their SERVER span (created in auth, never + ended → never exported). Always emitting keeps both success and failure + paths consistent. + """ + from litellm.proxy.proxy_server import open_telemetry_logger + + if open_telemetry_logger is None: + return + + http_request: Optional[Request] = kwargs.get("http_request") + if http_request is not None: + route = http_request.url.path + request_body: dict = await _read_request_body(request=http_request) + else: + route = func.__name__ + request_body = {} + + logging_payload = ManagementEndpointLoggingPayload( + route=route, + request_data=request_body, + response=None, + start_time=start_time, + end_time=end_time, + exception=exception, + ) + + if exception is None: + await open_telemetry_logger.async_management_endpoint_success_hook( + logging_payload=logging_payload, + parent_otel_span=parent_otel_span, + ) + else: + await open_telemetry_logger.async_management_endpoint_failure_hook( + logging_payload=logging_payload, + parent_otel_span=parent_otel_span, + ) + + def management_endpoint_wrapper(func): """ This wrapper does the following: @@ -446,13 +498,10 @@ def management_endpoint_wrapper(func): @wraps(func) async def wrapper(*args, **kwargs): start_time = datetime.now() - _http_request: Optional[Request] = None try: result = await func(*args, **kwargs) end_time = datetime.now() try: - if kwargs is None: - kwargs = {} user_api_key_dict: UserAPIKeyAuth = ( kwargs.get("user_api_key_dict") or UserAPIKeyAuth() ) @@ -462,31 +511,16 @@ def management_endpoint_wrapper(func): user_api_key_dict=user_api_key_dict, function_name=func.__name__, ) - _http_request = kwargs.get("http_request", None) parent_otel_span = getattr(user_api_key_dict, "parent_otel_span", None) if parent_otel_span is not None: - from litellm.proxy.proxy_server import open_telemetry_logger - - if open_telemetry_logger is not None: - if _http_request: - _route = _http_request.url.path - _request_body: dict = await _read_request_body( - request=_http_request - ) - _response = dict(result) if result is not None else None - - logging_payload = ManagementEndpointLoggingPayload( - route=_route, - request_data=_request_body, - response=_response, - start_time=start_time, - end_time=end_time, - ) - - await open_telemetry_logger.async_management_endpoint_success_hook( # type: ignore - logging_payload=logging_payload, - parent_otel_span=parent_otel_span, - ) + await _emit_management_endpoint_otel_span( + func=func, + kwargs=kwargs, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, + result=result, + ) # Delete updated/deleted info from cache _delete_api_key_from_cache(kwargs=kwargs) @@ -502,39 +536,19 @@ def management_endpoint_wrapper(func): except Exception as e: end_time = datetime.now() - if kwargs is None: - kwargs = {} user_api_key_dict: UserAPIKeyAuth = ( kwargs.get("user_api_key_dict") or UserAPIKeyAuth() ) parent_otel_span = getattr(user_api_key_dict, "parent_otel_span", None) if parent_otel_span is not None: - from litellm.proxy.proxy_server import open_telemetry_logger - - if open_telemetry_logger is not None: - _http_request = kwargs.get("http_request") - if _http_request: - _route = _http_request.url.path - _request_body: dict = await _read_request_body( - request=_http_request - ) - else: - _route = func.__name__ - _request_body = {} - - logging_payload = ManagementEndpointLoggingPayload( - route=_route, - request_data=_request_body, - response=None, - start_time=start_time, - end_time=end_time, - exception=e, - ) - - await open_telemetry_logger.async_management_endpoint_failure_hook( # type: ignore - logging_payload=logging_payload, - parent_otel_span=parent_otel_span, - ) + await _emit_management_endpoint_otel_span( + func=func, + kwargs=kwargs, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, + exception=e, + ) raise e diff --git a/tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py b/tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py index b1a3b834c3..34103449da 100644 --- a/tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py +++ b/tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py @@ -3,6 +3,7 @@ async_management_endpoint_{success,failure}_hook integration points.""" import asyncio from datetime import datetime +from unittest.mock import MagicMock import pytest @@ -14,6 +15,7 @@ from litellm.proxy._types import ( from ._helpers import ( HttpStatusException, assert_server_span_attrs, + get_server_span, make_fastapi_http_exception, make_httpx_status_error, ) @@ -28,6 +30,10 @@ def _real_user_api_key_dict(parent_span): ) +async def _noop_alert(*args, **kwargs): + return None + + async def _drive_admin_failure(*, otel, exception, parent_span, route): payload = ManagementEndpointLoggingPayload( route=route, @@ -180,3 +186,173 @@ def test_admin_endpoint_failure_stamps_server_span( expected_url_path=path, where=f"{path} {expected_status}", ) + + +def test_management_wrapper_success_ends_server_span_without_http_request( + server_span_factory, otel_with_exporter, monkeypatch +): + """Regression: management endpoints whose handler does not declare an + ``http_request`` parameter (``/key/generate``, ``/user/new``, ``/mcp/*``, + ...) must still get their parent SERVER span stamped + ended on success. + + The success hook itself stamps 200 and ``end()``s the parent, but the + wrapper only invoked it when ``http_request`` was present — so on success + the span (created in auth) was never ended and never exported. This drives + the real wrapper around an ``http_request``-less handler and asserts the + SERVER span reaches the exporter with status 200. + """ + 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_wrapper + async def fake_generate_key_fn(data=None, user_api_key_dict=None): + # No ``http_request`` parameter — mirrors generate_key_fn et al. + return {"key": "sk-xyz", "key_name": "k"} + + asyncio.run( + fake_generate_key_fn( + data={}, + user_api_key_dict=_real_user_api_key_dict(server_span), + ) + ) + + assert_server_span_attrs( + exporter, + expected_status=200, + expected_url_path=KEY_GENERATE_PATH, + where="management wrapper success without http_request", + ) + + +def test_management_wrapper_failure_ends_server_span( + server_span_factory, otel_with_exporter, monkeypatch +): + """When the handler raises, the wrapper must route through the failure hook + and stamp + end the parent SERVER span with the error status — even for an + ``http_request``-less handler (route falls back to ``func.__name__``).""" + 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) + + server_span = server_span_factory(KEY_GENERATE_PATH) + + @mgmt_utils.management_endpoint_wrapper + async def failing_fn(data=None, user_api_key_dict=None): + raise HttpStatusException(500, "boom") + + with pytest.raises(HttpStatusException): + asyncio.run( + failing_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span)) + ) + + assert_server_span_attrs( + exporter, + expected_status=500, + expected_url_path=KEY_GENERATE_PATH, + where="management wrapper failure", + ) + + +def test_management_wrapper_success_with_http_request( + server_span_factory, otel_with_exporter, monkeypatch +): + """Cover the branch where the handler DOES declare ``http_request``: the + route comes from ``http_request.url.path`` and the body is read from it.""" + 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) + + async def _fake_body(request=None): + return {"team_alias": "t"} + + monkeypatch.setattr(mgmt_utils, "_read_request_body", _fake_body) + + server_span = server_span_factory("/team/new") + http_request = MagicMock() + http_request.url.path = "/team/new" + + @mgmt_utils.management_endpoint_wrapper + async def fake_new_team(data=None, http_request=None, user_api_key_dict=None): + return {"team_id": "t-1"} + + asyncio.run( + fake_new_team( + data={}, + http_request=http_request, + user_api_key_dict=_real_user_api_key_dict(server_span), + ) + ) + + assert_server_span_attrs( + exporter, + expected_status=200, + expected_url_path="/team/new", + where="management wrapper success with http_request", + ) + + +def test_management_wrapper_noop_when_otel_logger_absent( + server_span_factory, otel_with_exporter, monkeypatch +): + """When no OTEL logger is registered, the helper early-returns and no SERVER + span is exported — and the handler result is still returned unchanged.""" + 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", None, raising=False) + monkeypatch.setattr(mgmt_utils, "send_management_endpoint_alert", _noop_alert) + + server_span = server_span_factory(KEY_GENERATE_PATH) + + @mgmt_utils.management_endpoint_wrapper + async def fake_fn(data=None, user_api_key_dict=None): + return {"ok": True} + + result = asyncio.run( + fake_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span)) + ) + + assert result == {"ok": True} + assert get_server_span(exporter) is None + + +def test_management_wrapper_swallows_post_success_errors( + server_span_factory, otel_with_exporter, monkeypatch +): + """A failure in post-success bookkeeping (cache invalidation, alerting) must + not propagate — the handler result is returned regardless (non-blocking).""" + 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) + + def _boom(*args, **kwargs): + raise RuntimeError("cache backend down") + + monkeypatch.setattr(mgmt_utils, "_delete_api_key_from_cache", _boom) + + server_span = server_span_factory(KEY_GENERATE_PATH) + + @mgmt_utils.management_endpoint_wrapper + async def fake_fn(data=None, user_api_key_dict=None): + return {"ok": True} + + result = asyncio.run( + fake_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span)) + ) + + assert result == {"ok": True}