fix(proxy): trace model list requests with the caller's traceparent (#49141)

* fix(proxy): trace model list requests with the caller's traceparent

The proxy creates an OTEL SERVER span at the auth boundary whenever a request
carries traceparent, but /v1/models never closed it. Unended spans are not
exported, so callers saw nothing for the model list while inference requests
showed up in the same trace

Route the endpoint through the management span wrapper so the SERVER span is
closed and parented to the caller's trace. The wrapper gains a variant that
skips request and response payloads, since attaching the full model list would
bloat every poll

* refactor(proxy): type the span wrapper with ParamSpec

The strict ruff rules ban new typing.cast imports, so the wrapper keeps the
wrapped endpoint's exact signature via ParamSpec instead of casting the
wrapped callable back to its type. No behavior change

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

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

* fix(proxy): keep the 5xx status on the model list span and trace model retrieve

The span wrapper ended the server span on every failure, so a 500 or 503
from /v1/models lost its status code. It now leaves that span to the
proxy's exception handler, which knows the answered status, and only ends
it itself for an HTTPException, which no handler closes on these routes.

GET /v1/models/{model_id} (the OpenAI SDK's models.retrieve) gets the same
wrapper, so it joins the caller's trace like the listing does.

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Wiley Kestner 2026-10-31 09:22:45 +01:00 • committed by GitHub
parent c4b01b2ac8
commit da71a03f6e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 248 additions and 15 deletions

View file

@ -1,14 +1,15 @@
# 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, ParamSpec, Protocol, TypeVar
from fastapi import HTTPException, Request
from pydantic import BaseModel, TypeAdapter
from starlette.exceptions import HTTPException as StarletteHTTPException
import litellm
from litellm._internal_context import with_service_target
@ -687,6 +688,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 +718,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 +740,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}
@ -767,37 +769,60 @@ async def _emit_management_endpoint_otel_span(
)
def management_endpoint_wrapper(func):
_P: Final = ParamSpec("_P")
_R: Final = TypeVar("_R")
def management_endpoint_wrapper(func: Callable[_P, Awaitable[_R]]) -> Callable[_P, Awaitable[_R]]:
"""
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)
async def wrapper(*args, **kwargs):
def management_endpoint_span_wrapper(func: Callable[_P, Awaitable[_R]]) -> Callable[_P, Awaitable[_R]]:
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 _proxy_exception_handler_ends_server_span(exception: Exception) -> bool:
return not isinstance(exception, StarletteHTTPException)
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:
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()
user_api_key_dict: UserAPIKeyAuth = _caller_auth(kwargs)
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,
log_payload=log_payload,
)
# Delete updated/deleted info from cache
@ -813,17 +838,18 @@ def management_endpoint_wrapper(func):
except Exception as e:
end_time = datetime.now()
user_api_key_dict: UserAPIKeyAuth = 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:
if parent_otel_span is not None and (log_payload or not _proxy_exception_handler_ends_server_span(e)):
try:
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,
exception=e,
log_payload=log_payload,
)
except Exception as otel_exc:
# Non-Blocking Exception - never let OTEL failures swallow

View file

@ -181,6 +181,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,
@ -12339,6 +12340,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),
@ -12649,6 +12651,7 @@ async def model_list(
dependencies=[Depends(user_api_key_auth)],
tags=["model management"],
)
@management_endpoint_span_wrapper
async def model_info(
model_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),

View file

@ -3,10 +3,14 @@ async_management_endpoint_{success,failure}_hook integration points."""
import asyncio
from datetime import datetime
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
from fastapi import Request
from litellm import Router
from litellm.proxy._types import (
ManagementEndpointLoggingPayload,
UserAPIKeyAuth,
@ -340,3 +344,203 @@ 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}
TEAM_NEW_PATH: Final = "/team/new"
MODEL_LIST_PATH: Final = "/v1/models"
MODEL_RETRIEVE_PATH: Final = "/v1/models/gpt-4o"
INBOUND_TRACEPARENT: Final = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
def _request(method: str, path: str, body: bytes = b"") -> Request:
async def receive() -> dict[str, object]:
return {"type": "http.request", "body": body, "more_body": False}
return Request(
{
"type": "http",
"method": method,
"path": path,
"root_path": "",
"query_string": b"",
"headers": [(b"content-type", b"application/json")],
},
receive,
)
async def _answer_like_the_proxy(request: Request, exception: Exception) -> None:
from litellm.proxy.proxy_server import app
handler: Final = next(
app.exception_handlers[cls] for cls in type(exception).__mro__ if cls in app.exception_handlers
)
await handler(request, exception)
@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)
server_span: Final = server_span_factory(TEAM_NEW_PATH)
wrap: Final = getattr(mgmt_utils, wrapper_name)
@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(
http_request=_request("POST", TEAM_NEW_PATH, b'{"team_alias": "t"}'),
user_api_key_dict=_real_user_api_key_dict(server_span),
)
)
spans: Final = [s for s in exporter.get_finished_spans() if s.name == TEAM_NEW_PATH]
assert len(spans) == 1
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
@pytest.mark.parametrize(
("wrapper_name", "exports_request_body"),
[("management_endpoint_wrapper", True), ("management_endpoint_span_wrapper", False)],
)
def test_span_wrapper_omits_request_body_on_failure(
server_span_factory, otel_with_exporter, monkeypatch, wrapper_name, exports_request_body
):
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: Final = server_span_factory(TEAM_NEW_PATH)
wrap: Final = getattr(mgmt_utils, wrapper_name)
@wrap
async def fake_fn(http_request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]:
raise make_fastapi_http_exception(400, "team_alias taken")
with pytest.raises(Exception, match="team_alias taken"):
asyncio.run(
fake_fn(
http_request=_request("POST", TEAM_NEW_PATH, b'{"team_alias": "t"}'),
user_api_key_dict=_real_user_api_key_dict(server_span),
)
)
spans: Final = [s for s in exporter.get_finished_spans() if s.name == TEAM_NEW_PATH]
assert len(spans) == 1
assert any(key.startswith("request.") for key in spans[0].attributes) is exports_request_body
@pytest.mark.parametrize(
("raised", "status_code"),
[
(RuntimeError("router settings lookup failed"), 500),
(httpx.ConnectError("All connection attempts failed"), 503),
(make_fastapi_http_exception(400, "Invalid scope parameter"), 400),
],
)
def test_span_wrapper_failure_keeps_the_answered_status(otel_with_exporter, monkeypatch, raised, status_code):
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: Final = _server_span(otel, MODEL_LIST_PATH, headers={})
request: Final = _request("GET", MODEL_LIST_PATH)
request.state.parent_otel_span = server_span
@mgmt_utils.management_endpoint_span_wrapper
async def failing_model_list(request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]:
raise raised
async def call_and_answer() -> None:
try:
await failing_model_list(request=request, user_api_key_dict=_real_user_api_key_dict(server_span))
except Exception as exc:
await _answer_like_the_proxy(request, exc)
asyncio.run(call_and_answer())
assert_server_span_attrs(
exporter,
expected_status=status_code,
expected_url_path=MODEL_LIST_PATH,
where=f"model list {type(raised).__name__}",
)
exported: Final = get_server_span(exporter)
assert exported is not None
assert exported.attributes.get("error.code") == str(status_code)
def _server_span(otel, path: str, headers: dict[str, str]):
server_span: Final = otel.create_litellm_proxy_request_started_span(start_time=datetime.now(), headers=headers)
otel.set_proxy_request_route_attributes(server_span, url_path=path, http_route=path)
return server_span
def _assert_ended_in_inbound_trace(exporter, path: str, handler_name: str) -> None:
assert_server_span_attrs(exporter, expected_status=200, expected_url_path=path, where=f"{path} success")
exported: Final = get_server_span(exporter)
assert exported is not None
assert exported.context.trace_id == int(INBOUND_TRACEPARENT.split("-")[1], 16)
endpoint_spans: Final = [s for s in exporter.get_finished_spans() if s.name == handler_name]
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)
def _install_proxy_state(monkeypatch, router: Router | None) -> None:
import litellm.proxy.proxy_server as proxy_server
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "user_model", None)
def test_model_list_ends_server_span_in_inbound_trace(otel_with_exporter, monkeypatch):
import litellm.proxy.proxy_server as proxy_server
otel, exporter = otel_with_exporter
monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False)
_install_proxy_state(monkeypatch, None)
server_span: Final = _server_span(otel, MODEL_LIST_PATH, headers={"traceparent": INBOUND_TRACEPARENT})
result: Final = asyncio.run(proxy_server.model_list(user_api_key_dict=_real_user_api_key_dict(server_span)))
assert result == {"data": [], "object": "list"}
_assert_ended_in_inbound_trace(exporter, MODEL_LIST_PATH, "model_list")
def test_model_retrieve_ends_server_span_in_inbound_trace(otel_with_exporter, monkeypatch):
import litellm.proxy.proxy_server as proxy_server
otel, exporter = otel_with_exporter
monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False)
_install_proxy_state(
monkeypatch,
Router(model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-x"}}]),
)
server_span: Final = _server_span(otel, MODEL_RETRIEVE_PATH, headers={"traceparent": INBOUND_TRACEPARENT})
result: Final = asyncio.run(
proxy_server.model_info(model_id="gpt-4o", user_api_key_dict=_real_user_api_key_dict(server_span))
)
assert result["id"] == "gpt-4o"
_assert_ended_in_inbound_trace(exporter, MODEL_RETRIEVE_PATH, "model_info")