fix(otel): export SERVER span on management-endpoint success without http_request (#28794)

Co-authored-by: Yassin Kortam <yassinkortam@Yassins-MacBook-Pro.local>

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
 

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}