ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Yassin Kortam <yassinkortam@Yassins-MacBook-Pro.local>
This commit is contained in:
Yassin Kortam 2026-05-25 12:13:17 -07:00 committed by GitHub
parent f9407bc036
commit 30551de371
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 245 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

124
litellm/proxy/management_helpers/utils.py
View File

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

176
tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py
View File

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