scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 118

feat(proxy): require trace IDs per team on LLM, MCP and agent requests (#44933)

Browse source 
* feat(proxy): require trace IDs per team on LLM, MCP and agent requests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): assert trace ID authorization outcomes

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): clear trace ID type gate errors

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style(proxy): format trace ID type check

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): validate trace ID metadata

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): use read-only trace metadata types

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): skip trace ID metadata validation for teams without dict metadata

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* perf(proxy): check team trace ID setting before route matching

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-10 11:50:28 -07:00 • committed by GitHub
parent b9c49aeb35
commit 6fcf9ef62d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 1078 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/_types.py
View file

@ -540,11 +540,7 @@ class LiteLLMRoutes(enum.Enum):
#########################################################
passthrough_routes_wildcard = [f"{route}/*" for route in mapped_pass_through_routes]
litellm_native_routes = [
"/rag/ingest",
"/v1/rag/ingest",
"/rag/query",
"/v1/rag/query",
trace_telemetry_routes = (
"/v1/traces",
"/v1/logs",
"/v1/traces/query",
@ -552,6 +548,14 @@ class LiteLLMRoutes(enum.Enum):
"/v1/traces/{trace_id}",
"/v1/traces/{trace_id}/spans/{span_id}",
"/v1/traces/{trace_id}/spans/{span_id}/error",
)
litellm_native_routes = [
"/rag/ingest",
"/v1/rag/ingest",
"/rag/query",
"/v1/rag/query",
*trace_telemetry_routes,
]
anthropic_routes = [
@ -2215,6 +2219,7 @@ class NewTeamRequest(TeamBase):
allowed_passthrough_routes: list | None = None
denied_passthrough_routes: list[str] | None = None
disable_global_guardrails: bool | None = None
require_trace_id: bool | None = None
secret_manager_settings: dict | None = None
model_rpm_limit: dict[str, int] | None = None
rpm_limit_type: Literal["guaranteed_throughput", "best_effort_throughput"] | None = (
@ -2289,6 +2294,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
policies: list[str] | None = None
object_permission: LiteLLM_ObjectPermissionBase | None = None
disable_global_guardrails: bool | None = None
require_trace_id: bool | None = None
team_member_budget: float | None = None
team_member_budget_duration: str | None = None
team_member_rpm_limit: int | None = None
@ -5156,6 +5162,7 @@ LiteLLM_ManagementEndpoint_MetadataFields: Final = [
"throttle_on_budget_exceeded",
"enable_prompt_caching",
"end_user_budget_id",
"require_trace_id",
]
LiteLLM_ManagementEndpoint_MetadataFields_Premium: Final = [

116
litellm/proxy/auth/auth_checks.py
View file

@ -162,7 +162,13 @@ from .auth_checks_organization import (
add_team_org_context_to_request_body,
organization_role_based_access_check,
)
from .auth_utils import get_model_from_request, get_request_route_template, request_fallback_model_names
from .auth_utils import (
get_model_from_request,
get_request_route,
get_request_route_template,
request_dispatched_to_pass_through_endpoint,
request_fallback_model_names,
)
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -395,6 +401,61 @@ def _typed_request_body(request_body: dict) -> Mapping[str, object]:
typed_general_settings: Final = _typed_request_body
_TRACE_ID_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def _metadata_mapping_has_trace_id(metadata: Mapping[str, object]) -> bool:
trace_id: Final[object] = metadata.get("trace_id")
return isinstance(trace_id, str) and bool(trace_id)
def _metadata_has_trace_id(metadata: object) -> bool:
if not isinstance(metadata, dict):
return False
metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(metadata)
return _metadata_mapping_has_trace_id(metadata_dict)
def _team_metadata_requires_trace_id(team_object: LiteLLM_TeamTable) -> bool:
if not isinstance(
team_object.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # narrow the untyped field
dict,
):
return False
return (
team_object.metadata.get("require_trace_id") # pyright: ignore[reportUnknownMemberType] # narrowed to dict
is True
)
def _request_has_trace_id(request_body: Mapping[str, object], request: Request, *, headers_only: bool = False) -> bool:
from litellm.proxy.litellm_pre_call_utils import (
_trace_id_from_traceparent, # pyright: ignore[reportPrivateUsage] # reuse W3C traceparent parser
get_chain_id_from_headers,
metadata_variable_name_for_route,
)
if get_chain_id_from_headers(dict(request.headers)):
return True
if headers_only:
return False
litellm_trace_id: Final[object] = request_body.get("litellm_trace_id")
if isinstance(litellm_trace_id, str) and litellm_trace_id:
return True
if "litellm_trace_id" not in request_body:
traceparent: Final[str | None] = request.headers.get("traceparent")
if isinstance(traceparent, str) and _trace_id_from_traceparent(traceparent):
return True
metadata_variable_name: Final = metadata_variable_name_for_route(get_request_route(request))
selected_metadata: Final[object] = request_body.get(metadata_variable_name)
if isinstance(selected_metadata, dict):
selected_metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(selected_metadata)
if "trace_id" in selected_metadata_dict:
return _metadata_mapping_has_trace_id(selected_metadata_dict)
other_metadata_variable_name: Final = "litellm_metadata" if metadata_variable_name == "metadata" else "metadata"
return _metadata_has_trace_id(request_body.get(other_metadata_variable_name))
class _JsonLoadsObj(Protocol):
def __call__(self, data: str) -> object: ...
@ -1141,6 +1202,59 @@ async def common_checks(
code=status.HTTP_400_BAD_REQUEST,
)
pass_through_route: Final = RouteChecks.check_route_access(
route=route,
allowed_routes=LiteLLMRoutes.passthrough_routes_wildcard.value,
) or request_dispatched_to_pass_through_endpoint(request)
headers_only: Final = (
RouteChecks.check_route_access(
route=route,
allowed_routes=LiteLLMRoutes.mcp_inference_routes.value,
)
or RouteChecks.check_route_access(
route=route,
allowed_routes=LiteLLMRoutes.agent_inference_routes.value,
)
or pass_through_route
)
team_trace_id_gate_applies: Final = (
team_object is not None
and _team_metadata_requires_trace_id(team_object)
and request.method not in ("GET", "HEAD", "OPTIONS")
and (RouteChecks.is_llm_api_route(route=route) or pass_through_route)
and not RouteChecks.check_route_access(
route=route,
allowed_routes=LiteLLMRoutes.trace_telemetry_routes.value,
)
)
team_metadata: Final[Mapping[str, object] | None] = (
_TRACE_ID_METADATA_ADAPTER.validate_python(
team_object.metadata # pyright: ignore[reportUnknownMemberType] # validate the narrowed untyped field
)
if team_trace_id_gate_applies and team_object is not None
else None
)
if (
team_trace_id_gate_applies
and team_object is not None
and team_metadata is not None
and not _request_has_trace_id(
request_body=_typed_request_body(request_body),
request=request,
headers_only=headers_only,
)
):
raise ProxyException(
message=(
f"Team '{team_object.team_id}' requires a trace ID on every LLM, MCP and agent request. "
"Send an x-litellm-trace-id header. "
"LLM requests can also send a W3C traceparent header or set metadata.trace_id in the body."
),
type=ProxyErrorTypes.bad_request_error,
param="trace_id",
code=status.HTTP_400_BAD_REQUEST,
)
managed_policy: Final = managed_agent_policy(valid_token)
if _model and valid_token is not None and managed_policy is not None:
managed_models: Final = (managed_policy.object_permission or MappingProxyType({})).get("models", ())

23
litellm/proxy/management_endpoints/common_utils.py
View file

@ -272,6 +272,29 @@ def check_disable_global_guardrails_caller_permission(
_check_disable_global_guardrails_caller_permission: Final = check_disable_global_guardrails_caller_permission
def check_require_trace_id_caller_permission(
require_trace_id: bool | None,
metadata: Mapping[str, object] | None,
user_api_key_dict: UserAPIKeyAuth,
*,
metadata_sent: bool,
existing_metadata: Mapping[str, object] | None = None,
) -> None:
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
stored: Final = existing_metadata is not None and existing_metadata.get("require_trace_id") is True
metadata_value: Final = metadata is not None and metadata.get("require_trace_id") is True
requested: Final = (
require_trace_id if require_trace_id is not None else (metadata_value if metadata_sent else stored)
)
if requested is stored:
return
raise HTTPException(
status_code=403,
detail={"error": "Only proxy admins can set `require_trace_id` on a team."},
)
def team_member_has_permission(
user_api_key_dict: UserAPIKeyAuth,
team_obj: LiteLLM_TeamTable,

18
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -140,6 +140,7 @@ from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # l
_user_has_admin_view, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
check_disable_global_guardrails_caller_permission,
check_passthrough_routes_caller_permission,
check_require_trace_id_caller_permission,
member_budget_patch,
set_object_metadata_field,
team_member_has_permission,
@ -1415,6 +1416,7 @@ async def new_team(
- guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
- policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
- require_trace_id: Optional[bool] - Reject LLM, MCP and agent requests from this team that carry no trace ID.
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
- team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
- team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
@ -1643,6 +1645,12 @@ async def new_team(
user_api_key_dict,
entity="team",
)
check_require_trace_id_caller_permission(
data.require_trace_id,
data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict
user_api_key_dict,
metadata_sent="metadata" in data.model_fields_set,
)
if isinstance(data.metadata, dict):
TeamMemberBudgetHandler.strip_system_managed_metadata_keys(data.metadata)
@ -2178,6 +2186,7 @@ async def update_team(
- guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
- policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
- require_trace_id: Optional[bool] - Reject LLM, MCP and agent requests from this team that carry no trace ID.
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
- team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
- team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
@ -2331,6 +2340,13 @@ async def update_team(
entity="team",
existing_metadata=_existing_team_metadata if isinstance(_existing_team_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # existing_team_row.metadata is a bare dict
)
check_require_trace_id_caller_permission(
data.require_trace_id,
data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict
user_api_key_dict,
metadata_sent="metadata" in data.model_fields_set,
existing_metadata=_existing_team_metadata if isinstance(_existing_team_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # existing_team_row.metadata is a bare dict
)
if data.soft_budget is not None:
max_budget_to_check = data.max_budget if data.max_budget is not None else existing_team_row.max_budget
@ -2440,6 +2456,8 @@ async def update_team(
)
updated_kv = data.json(exclude_unset=True)
if updated_kv.get("require_trace_id") is None:
updated_kv.pop("require_trace_id", None)
if "model_max_budget" in updated_kv and updated_kv["model_max_budget"] is None:
updated_kv["model_max_budget"] = {}

2
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -3448,7 +3448,7 @@ class SSOAuthenticationHandler:
_team_request: Final = deepcopy(default_team_params)
_team_request["team_id"] = litellm_team_id
_team_request["team_alias"] = litellm_team_name
team_request = NewTeamRequest(**_team_request)
team_request = NewTeamRequest.model_validate(_team_request)
elif isinstance(litellm.default_team_params, DefaultTeamSSOParams):
_default_team_params: Final = deepcopy(litellm.default_team_params)
_new_team_request: Final = team_request.model_dump()

3
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -1118,6 +1118,7 @@ async def pass_through_request(
"""
from litellm.exceptions import ModifyResponseException
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers
from litellm.proxy.pass_through_endpoints.passthrough_guardrails import (
PassthroughGuardrailHandler,
)
@ -1223,6 +1224,7 @@ async def pass_through_request(
proxy_config=proxy_config,
route_description="pass_through_endpoint",
)
litellm_trace_id: Final[str | None] = get_chain_id_from_headers(dict(request.headers))
logging_obj = Logging(
model=passthrough_model,
messages=[{"role": "user", "content": safe_dumps(_parsed_body)}],
@ -1230,6 +1232,7 @@ async def pass_through_request(
call_type="pass_through_endpoint",
start_time=start_time,
litellm_call_id=litellm_call_id,
litellm_trace_id=litellm_trace_id,
function_id="1245",
dynamic_success_callbacks=team_callbacks.success_callbacks,
dynamic_failure_callbacks=team_callbacks.failure_callbacks,

415
tests/unit/proxy/auth/test_auth_checks.py
View file

@ -7,13 +7,15 @@ from dotenv import load_dotenv
load_dotenv()
from collections.abc import Iterator
from collections.abc import Iterator, Mapping
from copy import deepcopy
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest, litellm
import httpx
from fastapi import Request, status
from prisma import Prisma
from litellm._service_logger import ServiceTypes
from litellm.proxy._types import LiteLLM_OrganizationTable, UserAPIKeyAuth
@ -28,6 +30,8 @@ from litellm.proxy._types import (
LiteLLM_UserTable,
LiteLLM_TeamTable,
Litellm_EntityType,
LitellmUserRoles,
ProxyException,
)
from litellm.proxy.utils import PrismaClient
from litellm.proxy.auth.auth_checks import (
@ -35,9 +39,406 @@ from litellm.proxy.auth.auth_checks import (
is_model_cost_zero,
virtual_key_soft_budget_check,
_team_soft_budget_check,
common_checks,
)
from litellm.proxy.utils import ProxyLogging
from litellm.proxy.utils import CallInfo
from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER
async def _run_team_trace_id_check(
*,
route: str,
method: str = "POST",
request_body: dict[str, object] | None = None,
headers: Mapping[str, str] | None = None,
team_metadata: Mapping[str, object] | None = None,
team_object: LiteLLM_TeamTable | None = None,
registered_pass_through_endpoint: bool = False,
) -> bool:
request_payload: Final[dict[str, object]] = {} if request_body is None else request_body
raw_headers: Final = tuple(
(name.lower().encode("latin-1"), value.encode("latin-1"))
for name, value in (() if headers is None else headers.items())
)
request: Final = Request(
{
"type": "http",
"method": method,
"path": route,
"headers": raw_headers,
"query_string": b"",
}
)
if registered_pass_through_endpoint:
endpoint: Final = SimpleNamespace()
setattr(endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True)
request.scope["endpoint"] = endpoint
team: Final = (
team_object
if team_object is not None
else LiteLLM_TeamTable(
team_id="trace-id-team",
metadata=None if team_metadata is None else dict(team_metadata),
)
)
valid_token: Final = UserAPIKeyAuth(token="test-token", team_id=team.team_id)
user: Final = LiteLLM_UserTable(user_id="trace-id-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
return await common_checks(
request_body=request_payload,
team_object=team,
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(spec=ProxyLogging),
valid_token=valid_token,
request=request,
skip_budget_checks=True,
)
@pytest.mark.parametrize(
"route",
["/chat/completions", "/v1/messages", "/mcp/", "/a2a/my-agent"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_llm_mcp_and_agent_requests(route: str) -> None:
request_body: Final[dict[str, object]] = {"metadata": {"source": "unit-test"}}
original_body: Final = deepcopy(request_body)
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route=route,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
assert exc_info.value.message == (
"Team 'trace-id-team' requires a trace ID on every LLM, MCP and agent request. "
"Send an x-litellm-trace-id header. "
"LLM requests can also send a W3C traceparent header or set metadata.trace_id in the body."
)
assert request_body == original_body
@pytest.mark.parametrize(
("headers", "request_body"),
[
({"x-litellm-trace-id": "trace-123"}, {}),
({"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}, {}),
({}, {"metadata": {"trace_id": "trace-123"}}),
({}, {"litellm_metadata": {"trace_id": "trace-123"}}),
({}, {"litellm_trace_id": "trace-123"}),
({}, {"litellm_trace_id": "", "metadata": {"trace_id": "trace-123"}}),
({}, {"litellm_trace_id": None, "metadata": {"trace_id": "trace-123"}}),
],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_supported_trace_carriers(
headers: Mapping[str, str],
request_body: dict[str, object],
) -> None:
original_body: Final = deepcopy(request_body)
await _run_team_trace_id_check(
route="/chat/completions",
headers=headers,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert request_body == original_body
@pytest.mark.parametrize("route", ["/v1/traces", "/v1/logs", "/v1/traces/query"])
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_trace_telemetry_routes(route: str) -> None:
assert await _run_team_trace_id_check(
route=route,
method="POST",
request_body={},
team_metadata={"require_trace_id": True},
) is True
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_overridden_body_metadata_on_chat() -> None:
request_body: Final[dict[str, object]] = {
"metadata": {"trace_id": ""},
"litellm_metadata": {"trace_id": "trace-123"},
}
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route="/chat/completions",
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.parametrize(
"request_body",
[
{"litellm_metadata": {"trace_id": "trace-123"}},
{"metadata": {"trace_id": "trace-123"}},
],
ids=["selected-litellm-metadata", "promoted-requester-metadata"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_effective_metadata_on_responses_route(
request_body: dict[str, object],
) -> None:
assert await _run_team_trace_id_check(
route="/v1/responses",
request_body=request_body,
team_metadata={"require_trace_id": True},
) is True
@pytest.mark.parametrize(
("headers", "request_body"),
[
(
{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
{"litellm_trace_id": ""},
),
(
{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
{"litellm_trace_id": None},
),
],
ids=[
"traceparent-with-empty-top-level",
"traceparent-with-null-top-level",
],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_fallback_carriers_when_litellm_trace_id_is_explicit(
headers: Mapping[str, str],
request_body: dict[str, object],
) -> None:
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route="/chat/completions",
headers=headers,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_header_on_mcp_with_empty_body() -> None:
assert await _run_team_trace_id_check(
route="/mcp/",
request_body={},
headers={"x-litellm-trace-id": "trace-mcp-123"},
team_metadata={"require_trace_id": True},
) is True
@pytest.mark.parametrize(
("headers", "request_body"),
[
(
{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
{},
),
({}, {"metadata": {"trace_id": "trace-body-123"}}),
],
ids=["traceparent", "body-metadata"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_nonheader_carriers_on_pass_through_routes(
headers: Mapping[str, str],
request_body: dict[str, object],
) -> None:
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route="/anthropic/v1/messages",
headers=headers,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_header_on_pass_through_route() -> None:
assert await _run_team_trace_id_check(
route="/anthropic/v1/messages",
request_body={},
headers={"x-litellm-trace-id": "trace-provider-123"},
team_metadata={"require_trace_id": True},
) is True
@pytest.mark.parametrize(
("headers", "request_body"),
[
(
{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
{},
),
({}, {"metadata": {"trace_id": "trace-body-123"}}),
],
ids=["traceparent", "body-metadata"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_nonheader_carriers_on_registered_pass_through_routes(
headers: Mapping[str, str],
request_body: dict[str, object],
) -> None:
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route="/configured-passthrough/infer",
headers=headers,
request_body=request_body,
team_metadata={"require_trace_id": True},
registered_pass_through_endpoint=True,
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.parametrize(
("route", "request_body"),
[
("/mcp/", {}),
(
"/mcp-rest/tools/call",
{"server_id": "math", "name": "add", "arguments": {"left": 2, "right": 3}},
),
(
"/a2a/my-agent",
{
"jsonrpc": "2.0",
"id": "1",
"method": "message/send",
"params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
},
),
],
ids=["streamable-mcp", "mcp-rest", "a2a"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_traceparent_on_mcp_and_agent_routes(
route: str,
request_body: dict[str, object],
) -> None:
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route=route,
headers={"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.parametrize(
("route", "request_body"),
[
(
"/mcp-rest/tools/call",
{
"server_id": "math",
"name": "add",
"arguments": {"left": 2, "right": 3},
"metadata": {"trace_id": "trace-body-123"},
},
),
(
"/a2a/my-agent",
{
"jsonrpc": "2.0",
"id": "1",
"method": "message/send",
"params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
"metadata": {"trace_id": "trace-body-123"},
},
),
],
ids=["mcp-rest", "a2a"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_body_metadata_on_mcp_and_agent_routes(
route: str,
request_body: dict[str, object],
) -> None:
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route=route,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.parametrize(
("headers", "request_body"),
[
({"traceparent": "invalid"}, {}),
({}, {"metadata": {"trace_id": ""}}),
],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_invalid_trace_carriers(
headers: Mapping[str, str],
request_body: dict[str, object],
) -> None:
original_body: Final = deepcopy(request_body)
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route="/chat/completions",
headers=headers,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
assert request_body == original_body
@pytest.mark.parametrize(
("route", "method", "team_metadata"),
[
("/v1/models", "GET", {"require_trace_id": True}),
("/chat/completions", "POST", None),
("/chat/completions", "POST", {"require_trace_id": False}),
("/key/generate", "POST", {"require_trace_id": True}),
],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_exempts_methods_routes_and_disabled_teams(
route: str,
method: str,
team_metadata: Mapping[str, object] | None,
) -> None:
assert await _run_team_trace_id_check(route=route, method=method, team_metadata=team_metadata) is True
@pytest.mark.asyncio
async def test_team_require_trace_id_skips_non_dict_team_metadata() -> None:
team: Final[LiteLLM_TeamTable] = LiteLLM_TeamTable.model_construct(
team_id="trace-id-team",
metadata="not a mapping",
)
assert await _run_team_trace_id_check(route="/v1/chat/completions", team_object=team) is True
@pytest.mark.parametrize("route", ("/lens/datasets", "/lens/tracing/keys"))
@ -59,7 +460,7 @@ async def test_get_end_user_object(customer_spend, customer_budget):
Scenario 1: normal - get_end_user_object returns the cached user
Scenario 2: user over budget - NOTE: budget enforcement now happens in
common_checks() via _check_end_user_budget(), not in get_end_user_object()
This test verifies that get_end_user_object correctly retrieves the end user
from cache. Budget enforcement is tested separately in test_check_end_user_budget().
"""
@ -99,12 +500,12 @@ async def test_check_end_user_budget(customer_spend, customer_budget):
Test _check_end_user_budget enforcement:
- Scenario 1: customer_spend=0, customer_budget=10 - should pass (under budget)
- Scenario 2: customer_spend=10, customer_budget=0 - should fail (over budget)
Note: Budget enforcement for end users happens in common_checks() via
Note: Budget enforcement for end users happens in common_checks() via
_check_end_user_budget(), not in get_end_user_object().
"""
from litellm.proxy.auth.auth_checks import check_end_user_budget
_budget = LiteLLM_BudgetTable(max_budget=customer_budget)
end_user_obj = LiteLLM_EndUserTable(
user_id="my-test-customer",
@ -112,9 +513,9 @@ async def test_check_end_user_budget(customer_spend, customer_budget):
litellm_budget_table=_budget,
blocked=False,
)
should_exceed = customer_spend > customer_budget
if not should_exceed:
await check_end_user_budget(
end_user_obj=end_user_obj,

89
tests/unit/proxy/management_endpoints/test_common_utils.py
View file

@ -30,6 +30,7 @@ from litellm.proxy.management_endpoints.common_utils import (
_org_admin_can_invite_user,
_team_admin_can_invite_user,
admin_can_invite_user,
check_require_trace_id_caller_permission,
set_object_metadata_field,
update_metadata_fields,
user_api_key_has_admin_view,
@ -960,6 +961,94 @@ class TestCheckDisableGlobalGuardrailsCallerPermission:
)
@pytest.mark.parametrize(
("require_trace_id", "metadata", "metadata_sent", "existing_metadata"),
[
pytest.param(True, None, False, None, id="create-with-top-level-true"),
pytest.param(False, None, False, {"require_trace_id": True}, id="clear-stored-value-top-level"),
pytest.param(None, {}, True, {"require_trace_id": True}, id="replace-metadata-omitting-stored-value"),
pytest.param(
None,
None,
True,
{"require_trace_id": True},
id="explicit-null-metadata-clears-stored-value",
),
pytest.param(
False,
{"require_trace_id": True},
True,
{"require_trace_id": True},
id="top-level-false-wins-over-metadata",
),
],
)
def test_non_proxy_admin_cannot_change_require_trace_id(
require_trace_id: bool | None,
metadata: dict[str, object] | None,
metadata_sent: bool,
existing_metadata: dict[str, object] | None,
) -> None:
exc_info: Final[pytest.ExceptionInfo[HTTPException]]
with pytest.raises(HTTPException) as exc_info:
check_require_trace_id_caller_permission(
require_trace_id,
metadata,
UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER),
metadata_sent=metadata_sent,
existing_metadata=existing_metadata,
)
assert exc_info.value.status_code == 403
assert exc_info.value.detail == {"error": "Only proxy admins can set `require_trace_id` on a team."}
@pytest.mark.parametrize(
("require_trace_id", "metadata", "metadata_sent", "existing_metadata"),
[
pytest.param(True, {"require_trace_id": True}, True, {"require_trace_id": True}, id="resend-stored-value"),
pytest.param(None, {}, True, None, id="metadata-without-stored-value"),
pytest.param(
True,
{},
True,
{"require_trace_id": True},
id="top-level-true-with-empty-metadata-keeps-stored-value",
),
pytest.param(None, None, False, {"require_trace_id": True}, id="metadata-not-sent"),
],
)
def test_non_proxy_admin_can_preserve_or_omit_require_trace_id(
require_trace_id: bool | None,
metadata: dict[str, object] | None,
metadata_sent: bool,
existing_metadata: dict[str, object] | None,
) -> None:
assert (
check_require_trace_id_caller_permission(
require_trace_id,
metadata,
UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER),
metadata_sent=metadata_sent,
existing_metadata=existing_metadata,
)
is None
)
def test_proxy_admin_can_flip_stored_require_trace_id() -> None:
assert (
check_require_trace_id_caller_permission(
False,
None,
UserAPIKeyAuth(user_id="u2", user_role=LitellmUserRoles.PROXY_ADMIN),
metadata_sent=False,
existing_metadata={"require_trace_id": True},
)
is None
)
class TestTeamMemberHasPermission:
def test_requires_caller_to_be_a_team_member(self):
from litellm.proxy.management_endpoints.common_utils import (

26
tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py
View file

@ -1,3 +1,5 @@
from typing import Final
import pytest
from fastapi import HTTPException
@ -89,6 +91,30 @@ class TestChangedTeamFields:
data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["b"]})
assert changed_team_fields(data, _team(metadata={"guardrails": ["a"]})) == frozenset({"guardrails"})
def test_top_level_require_trace_id_change_is_attributed_to_the_folded_field(self):
existing: Final = _team(metadata={"require_trace_id": True})
data: Final = UpdateTeamRequest(team_id="team-1", require_trace_id=False)
assert changed_team_fields(data, existing) == frozenset({"require_trace_id"})
assert team_admin_edit_verdict(data, existing, frozenset({"tpm_limit"})) == TeamAdminFieldNotPermitted(
field="require_trace_id"
)
assert changed_team_fields(UpdateTeamRequest(team_id="team-1", require_trace_id=True), existing) == frozenset()
def test_metadata_replacement_cannot_drop_stored_require_trace_id(self):
data: Final = UpdateTeamRequest(team_id="team-1", metadata={})
existing: Final = _team(metadata={"require_trace_id": True})
assert changed_team_fields(data, existing) == frozenset({"require_trace_id"})
assert team_admin_edit_verdict(data, existing, frozenset({"tpm_limit"})) == TeamAdminFieldNotPermitted(
field="require_trace_id"
)
def test_resending_stored_require_trace_id_is_not_a_change(self):
data: Final = UpdateTeamRequest(team_id="team-1", metadata={"require_trace_id": True})
assert changed_team_fields(data, _team(metadata={"require_trace_id": True})) == frozenset()
def test_custom_metadata_key_change_is_attributed_to_metadata(self):
data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["a"], "cost_center": "b"})
existing = _team(metadata={"guardrails": ["a"], "cost_center": "a"})

166
tests/unit/proxy/management_endpoints/test_team_endpoints.py
View file

@ -9,8 +9,9 @@ from typing import Final, Optional, cast
from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch
import httpx
import litellm
import pytest
from fastapi import HTTPException
from fastapi import HTTPException, Request
from fastapi.testclient import TestClient
from pydantic import ValidationError
@ -30,6 +31,7 @@ from litellm.proxy._types import (
LitellmTableNames,
LitellmUserRoles,
Member,
NewTeamRequest,
ProxyErrorTypes,
ProxyException,
ResetSpendRequest,
@ -60,6 +62,7 @@ from litellm.proxy.management_endpoints.team_endpoints import (
aggregated_date_range_error,
delete_team,
list_available_teams,
new_team,
reset_team_member_budget_fn,
reset_team_member_spend_fn,
router,
@ -11706,6 +11709,147 @@ async def test_update_team_blocks_non_admin_disable_global_guardrails(mock_db_cl
assert "disable_global_guardrails" in str(exc.value.message)
@pytest.mark.asyncio
async def test_new_team_blocks_org_admin_require_trace_id(mock_db_client: MagicMock) -> None:
mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
with (
_org_admins(("u-team-admin", "org-1")),
patch(
"litellm.proxy.management_endpoints.team_endpoints._should_auto_add_team_creator",
return_value=False,
),
patch(
"litellm.proxy.management_endpoints.team_endpoints.get_org_object",
new=AsyncMock(return_value=MagicMock()),
),
patch(
"litellm.proxy.management_endpoints.team_endpoints._check_org_team_limits",
new=AsyncMock(return_value=None),
),
):
exc: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc:
await new_team(
data=NewTeamRequest(
team_alias="t",
organization_id="org-1",
require_trace_id=True,
),
http_request=MagicMock(spec=Request),
user_api_key_dict=_non_admin_auth(),
)
assert str(exc.value.code) == "403"
assert "Only proxy admins can set `require_trace_id` on a team." in str(exc.value.message)
@pytest.mark.asyncio
async def test_update_team_blocks_org_admin_require_trace_id_change(mock_db_client: MagicMock) -> None:
existing: Final[MagicMock] = MagicMock()
existing.metadata = {"require_trace_id": True}
existing.model_dump.return_value = {
"team_id": "t1",
"organization_id": "org-1",
"metadata": {"require_trace_id": True},
}
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing)
with _org_admins(("u-team-admin", "org-1")):
exc: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc:
await update_team(
data=UpdateTeamRequest(team_id="t1", require_trace_id=False),
http_request=MagicMock(spec=Request),
user_api_key_dict=_non_admin_auth(),
)
assert str(exc.value.code) == "403"
assert "Only proxy admins can set `require_trace_id` on a team." in str(exc.value.message)
@pytest.mark.asyncio
async def test_update_team_blocks_org_admin_null_metadata_clear_require_trace_id(
mock_db_client: MagicMock,
) -> None:
existing: Final[MagicMock] = MagicMock()
existing.metadata = {"require_trace_id": True}
existing.model_dump.return_value = {
"team_id": "t1",
"organization_id": "org-1",
"metadata": {"require_trace_id": True},
}
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing)
with _org_admins(("u-team-admin", "org-1")):
exc: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc:
await update_team(
data=UpdateTeamRequest(team_id="t1", metadata=None),
http_request=MagicMock(spec=Request),
user_api_key_dict=_non_admin_auth(),
)
assert str(exc.value.code) == "403"
assert "Only proxy admins can set `require_trace_id` on a team." in str(exc.value.message)
@pytest.mark.asyncio
async def test_update_team_ignores_explicit_null_require_trace_id_and_preserves_metadata(
mock_db_client: MagicMock,
mock_admin_auth: UserAPIKeyAuth,
disable_audit_logging_for_mocked_team: None,
) -> None:
mock_cache: Final[MagicMock]
with (
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()),
):
existing_team: Final[MagicMock] = MagicMock(
team_id="trace-null-team",
organization_id=None,
model_id=None,
metadata={"require_trace_id": True},
members_with_roles=[],
models=[],
)
existing_team.model_dump.return_value = {
"team_id": "trace-null-team",
"organization_id": None,
"metadata": {"require_trace_id": True},
"members_with_roles": [],
"models": [],
}
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team)
mock_db_client.jsonify_team_object = lambda db_data: db_data
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.async_set_cache = AsyncMock()
updated_team: Final[MagicMock] = MagicMock(
team_id="trace-null-team",
organization_id=None,
litellm_model_table=None,
metadata={"require_trace_id": True},
)
updated_team.model_dump.return_value = {
"team_id": "trace-null-team",
"organization_id": None,
"metadata": {"require_trace_id": True},
}
mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=updated_team)
await update_team(
data=UpdateTeamRequest(team_id="trace-null-team", require_trace_id=None),
http_request=MagicMock(spec=Request),
user_api_key_dict=mock_admin_auth,
)
written: Final = mock_db_client.db.litellm_teamtable.update.call_args.kwargs["data"]
assert "require_trace_id" not in written
assert written.get("metadata", existing_team.metadata)["require_trace_id"] is True
assert updated_team.metadata["require_trace_id"] is True
def test_set_budget_reset_at_clears_when_budget_duration_null():
"""
When budget_duration is explicitly set to null, _set_budget_reset_at
@ -14637,6 +14781,26 @@ def _wire_new_team_prisma(mock_db_client):
return mock_db_client.db.litellm_teamtable.create
@pytest.mark.asyncio
async def test_new_team_persists_require_trace_id_in_metadata(
mock_db_client: MagicMock,
mock_admin_auth: UserAPIKeyAuth,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "default_team_settings", None)
monkeypatch.setattr(litellm, "default_team_params", {})
mock_team_create: Final = _wire_new_team_prisma(mock_db_client)
await new_team(
data=NewTeamRequest(team_alias="trace-required-team", require_trace_id=True),
http_request=MagicMock(spec=Request),
user_api_key_dict=mock_admin_auth,
)
team_data: Final = mock_team_create.call_args.kwargs["data"]
assert team_data["metadata"]["require_trace_id"] is True
@pytest.mark.asyncio
async def test_new_team_explicit_null_budget_duration_beats_configured_default(
mock_db_client, mock_admin_auth, monkeypatch

36
tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -13,8 +13,9 @@ from dataclasses import dataclass
from datetime import datetime
from io import BytesIO
from types import MappingProxyType, ModuleType, SimpleNamespace
from typing import Final, Optional
from typing import Final, Optional, cast
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import UUID
import httpx
import pytest
@ -5618,14 +5619,14 @@ def _enter_relay_logging_mocks(stack, parsed_body):
return mock_proxy_logging, mock_success_handler
def _relay_client_request(method="GET"):
mock_request = MagicMock(spec=Request)
def _relay_client_request(method: str = "GET", headers: Mapping[str, str] | None = None) -> Request:
mock_request: Final = MagicMock(spec=Request)
mock_request.method = method
mock_request.url = httpx.URL("http://localhost:4000/passthrough-relay/results")
mock_request.body = AsyncMock(return_value=b"")
mock_request.headers = Headers({})
mock_request.headers = Headers({} if headers is None else headers)
mock_request.query_params = QueryParams({})
return mock_request
return cast(Request, mock_request)
@pytest.mark.asyncio
@ -6140,7 +6141,12 @@ def _enter_upstream_usage_mocks(stack, parsed_body):
return mock_proxy_logging, enqueued
async def _run_upstream_reporting_passthrough(upstream_headers, status_code=200, cost_per_request=None):
async def _run_upstream_reporting_passthrough(
upstream_headers: Mapping[str, str],
status_code: int = 200,
cost_per_request: float | None = None,
request_headers: Mapping[str, str] | None = None,
) -> tuple[list[object], MagicMock]:
"""Drive a generic pass-through against an upstream that reports its own
cost/usage. Returns (recorded standard logging payloads, proxy logging mock)."""
from litellm.proxy._types import UserAPIKeyAuth
@ -6158,7 +6164,7 @@ async def _run_upstream_reporting_passthrough(upstream_headers, status_code=200,
mock_proxy_logging, enqueued = _enter_upstream_usage_mocks(stack, {})
with _recording_success_callback() as recorder:
await pass_through_request(
request=_relay_client_request(method="POST"),
request=_relay_client_request(method="POST", headers=request_headers),
target="http://internal-api.test/v1/summarize",
custom_headers={},
user_api_key_dict=UserAPIKeyAuth(api_key="sk-upstream-usage", team_id="team-fil"),
@ -6172,6 +6178,22 @@ async def _run_upstream_reporting_passthrough(upstream_headers, status_code=200,
await fake_client.aclose()
@pytest.mark.asyncio
async def test_pass_through_logging_uses_trace_header_and_generates_without_one() -> None:
trace_id: Final[str] = "4bf92f3577b34da6a3ce929d0e0e4736"
trace_payloads, _ = await _run_upstream_reporting_passthrough(
upstream_headers=MappingProxyType({}),
request_headers=MappingProxyType({"x-litellm-trace-id": trace_id}),
)
logged_trace_id: Final[str] = cast(types_utils.StandardLoggingPayload, trace_payloads[0])["trace_id"]
assert logged_trace_id == trace_id
generated_payloads, _ = await _run_upstream_reporting_passthrough(upstream_headers=MappingProxyType({}))
generated_trace_id: Final[str] = cast(types_utils.StandardLoggingPayload, generated_payloads[0])["trace_id"]
assert str(UUID(generated_trace_id)) == generated_trace_id
assert generated_trace_id != trace_id
@pytest.mark.asyncio
async def test_passthrough_records_cost_and_tokens_reported_by_upstream():
"""

12
ui/litellm-dashboard/src/app/(dashboard)/teams/page.tsx
View file

@ -4,6 +4,14 @@ import Teams from "@/components/Teams";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
export default function TeamsPage() {
const { accessToken, userId, userRole, premiumUser } = useAuthorized();
return <Teams accessToken={accessToken} userID={userId} userRole={userRole} premiumUser={premiumUser ?? false} />;
const { accessToken, userId, userRole, isViewOnly, premiumUser } = useAuthorized();
return (
<Teams
accessToken={accessToken}
userID={userId}
userRole={userRole}
isViewOnly={isViewOnly}
premiumUser={premiumUser ?? false}
/>
);
}

49
ui/litellm-dashboard/src/components/Teams.test.tsx
View file

@ -552,6 +552,17 @@ describe("Teams - Create Team CTA is grouped with the tabs on the left", () => {
renderWithQueryClient(<Teams accessToken="test-token" userID="user-123" userRole="Admin Viewer" />);
expect(screen.queryByTestId("create-team-button")).not.toBeInTheDocument();
});
it("should hide the Require Trace ID switch from a view-only Admin session", async () => {
renderWithQueryClient(<Teams accessToken="test-token" userID="user-123" userRole="Admin" isViewOnly />);
fireEvent.click(screen.getByTestId("create-team-button"));
await screen.findByLabelText(/team name/i);
fireEvent.click(screen.getByText("Additional Settings"));
expect(screen.getByRole("switch", { name: /Disable Global Guardrails/i })).toBeInTheDocument();
expect(screen.queryByRole("switch", { name: /Require Trace ID/i })).not.toBeInTheDocument();
});
});
describe("Teams - Default Team Settings tab visibility", () => {
@ -1447,8 +1458,7 @@ describe("Teams - the exact bytes the create call sends", () => {
await openCreateModal({ premiumUser: true });
await openSection("Additional Settings", /Team Member Key Duration/);
const switches = screen.getAllByRole("switch");
fireEvent.click(switches[switches.length - 1]);
fireEvent.click(screen.getByRole("switch", { name: /Disable Global Guardrails/i }));
const payload = await submit();
@ -1459,8 +1469,7 @@ describe("Teams - the exact bytes the create call sends", () => {
await openCreateModal();
await openSection("Additional Settings", /Team Member Key Duration/);
const switches = screen.getAllByRole("switch");
fireEvent.click(switches[switches.length - 1]);
fireEvent.click(screen.getByRole("switch", { name: /Disable Global Guardrails/i }));
const payload = await submit();
@ -1797,6 +1806,7 @@ describe("Teams - disable_global_guardrails switch gating", () => {
fireEvent.click(screen.getByText("Additional Settings"));
expect(screen.queryByRole("switch", { name: /Disable Global Guardrails/i })).not.toBeInTheDocument();
expect(screen.queryByRole("switch", { name: /Require Trace ID/i })).not.toBeInTheDocument();
});
it("shows the Disable Global Guardrails switch to a proxy admin", async () => {
@ -1807,4 +1817,35 @@ describe("Teams - disable_global_guardrails switch gating", () => {
expect(await screen.findByRole("switch", { name: /Disable Global Guardrails/i })).toBeInTheDocument();
});
it("sends require_trace_id when a proxy admin enables it during team creation", async () => {
const createdTeamResponse = {
team_id: "new-team-1",
team_alias: "Trace Required",
models: [],
organization_id: null,
keys: [],
members_with_roles: [],
spend: 0,
};
vi.mocked(teamCreateCall).mockResolvedValue(createdTeamResponse);
renderWithQueryClient(<Teams accessToken="test-token" userID="user-123" userRole="Admin" />);
await openCreateModal();
fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Trace Required" } });
fireEvent.click(screen.getByText("Additional Settings"));
fireEvent.click(await screen.findByRole("switch", { name: /Require Trace ID/i }));
const submitButtons = screen.getAllByRole("button", { name: /create team/i });
fireEvent.click(submitButtons[submitButtons.length - 1]);
await waitFor(() => {
expect(teamCreateCall).toHaveBeenCalledWith(
"test-token",
expect.objectContaining({
team_alias: "Trace Required",
require_trace_id: true,
}),
);
});
});
});

29
ui/litellm-dashboard/src/components/Teams.tsx
View file

@ -60,6 +60,7 @@ interface TeamProps {
accessToken: string | null;
userID: string | null;
userRole: string | null;
isViewOnly?: boolean;
premiumUser?: boolean;
}
@ -90,6 +91,7 @@ const teamCreateFieldsSchema = z.object({
secret_manager_settings: z.string().optional(),
guardrails: z.array(z.string()).optional(),
disable_global_guardrails: z.boolean().optional(),
require_trace_id: z.boolean().optional(),
policies: z.array(z.string()).optional(),
access_group_ids: z.array(z.string()).optional(),
allowed_vector_store_ids: z.array(z.string()).optional(),
@ -127,6 +129,7 @@ const EMPTY_TEAM_CREATE_VALUES: TeamCreateFormValues = {
secret_manager_settings: undefined,
guardrails: undefined,
disable_global_guardrails: undefined,
require_trace_id: undefined,
policies: undefined,
access_group_ids: undefined,
allowed_vector_store_ids: undefined,
@ -147,6 +150,7 @@ const ADDITIONAL_SETTINGS_FIELDS = [
"secret_manager_settings",
"guardrails",
"disable_global_guardrails",
"require_trace_id",
"policies",
"access_group_ids",
"allowed_vector_store_ids",
@ -210,7 +214,7 @@ const getAdminOrganizations = (
};
// @deprecated
const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser = false }) => {
const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, isViewOnly = false, premiumUser = false }) => {
const { data: organizationsData } = useOrganizations();
const organizations = organizationsData ?? null;
const { data: teamMetadataSchemaFields = [], isLoading: isTeamMetadataSchemaLoading } = useTeamMetadataSchema();
@ -264,6 +268,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
const [editTeam, setEditTeam] = useState<boolean>(false);
const [isTeamModalVisible, setIsTeamModalVisible] = useState(false);
const canEditAsProxyAdmin = !isViewOnly && isProxyAdminRole(userRole || "");
const [userModels, setUserModels] = useState<string[]>([]);
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
const [teamToDelete, setTeamToDelete] = useState<Team | null>(null);
@ -545,7 +550,12 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
}
}
await teamCreateCall(accessToken, { ...formValues, models: normalizeTeamModelSelection(formValues.models) });
const { require_trace_id: requireTraceId, ...teamCreateValues } = formValues;
await teamCreateCall(accessToken, {
...teamCreateValues,
...(requireTraceId === undefined ? {} : { require_trace_id: requireTraceId }),
models: normalizeTeamModelSelection(formValues.models),
});
toast.success("Team created");
await refreshTeams();
resetCreateForm();
@ -1005,6 +1015,21 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
)}
</FormField>
)}
{canEditAsProxyAdmin && (
<FormField
control={form.control}
name="require_trace_id"
className="mt-4"
label={labelWithHint(
"Require Trace ID",
"Reject LLM, MCP and agent requests from this team that do not send a trace ID (x-litellm-trace-id header, LLM requests can also use traceparent or metadata.trace_id)",
)}
>
{({ id, value, onChange }) => (
<Switch id={id} checked={value === true} onCheckedChange={onChange} />
)}
</FormField>
)}
{canViewPolicies && (
<FormField
control={form.control}

67
ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
View file

@ -15,7 +15,7 @@ import TeamInfoView, {
type TeamData,
} from "./TeamInfo";
const authState = vi.hoisted(() => ({ userRole: "Admin" }));
const authState = vi.hoisted(() => ({ userRole: "Admin", isViewOnly: false }));
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
default: () => ({
@ -24,6 +24,7 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
userId: "user-1",
userEmail: "user@example.com",
userRole: authState.userRole,
isViewOnly: authState.isViewOnly,
premiumUser: false,
disabledPersonalKeyCreation: null,
showSSOBanner: false,
@ -361,6 +362,7 @@ describe("TeamInfoView", () => {
afterEach(() => {
vi.clearAllMocks();
authState.userRole = "Admin";
authState.isViewOnly = false;
});
describe("display and rendering", () => {
@ -2226,9 +2228,11 @@ describe("TeamInfoView - the exact bytes the update call sends", () => {
afterEach(() => {
vi.clearAllMocks();
authState.userRole = "Admin";
authState.isViewOnly = false;
});
const storedTeam = () =>
const storedTeam = (requireTraceId = false) =>
createMockTeamData({
models: ["gpt-4"],
max_budget: 100,
@ -2238,13 +2242,17 @@ describe("TeamInfoView - the exact bytes the update call sends", () => {
team_member_budget_table: { max_budget: 42, budget_duration: "30d", tpm_limit: 11, rpm_limit: 22 },
default_team_member_models: ["gpt-4"],
object_permission: { search_tools: ["tool-a"], vector_stores: ["vs-1"] },
metadata: requireTraceId ? { require_trace_id: true } : {},
});
const openEditor = async (user: ReturnType<typeof userEvent.setup>) => {
vi.mocked(networking.teamInfoCall).mockResolvedValue(storedTeam());
const openEditor = async (
user: ReturnType<typeof userEvent.setup>,
options: { requireTraceId?: boolean; isProxyAdmin?: boolean } = {},
) => {
vi.mocked(networking.teamInfoCall).mockResolvedValue(storedTeam(options.requireTraceId === true));
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any);
renderWithProviders(<TeamInfoView {...props} />);
renderWithProviders(<TeamInfoView {...props} is_proxy_admin={options.isProxyAdmin ?? true} />);
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
await user.click(screen.getByRole("tab", { name: "Settings" }));
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
@ -2298,6 +2306,7 @@ describe("TeamInfoView - the exact bytes the update call sends", () => {
const payload = await save(user);
expect(payload.metadata).not.toHaveProperty("require_trace_id");
expect(payload).toStrictEqual({
...alwaysSent,
team_member_budget_duration: undefined,
@ -2309,6 +2318,27 @@ describe("TeamInfoView - the exact bytes the update call sends", () => {
});
});
it("sends require_trace_id when a proxy admin enables it", async () => {
const user = userEvent.setup({ delay: null });
await openEditor(user);
fireEvent.click(screen.getByRole("switch", { name: /Require Trace ID/i }));
const payload = await save(user);
expect(payload.metadata).toMatchObject({ require_trace_id: true });
});
it("hides the switch from non-admins and preserves the stored value on an unrelated save", async () => {
const user = userEvent.setup({ delay: null });
authState.userRole = "Internal User";
await openEditor(user, { requireTraceId: true, isProxyAdmin: false });
expect(screen.queryByRole("switch", { name: /Require Trace ID/i })).not.toBeInTheDocument();
const payload = await save(user);
expect(payload.metadata).toMatchObject({ require_trace_id: true });
});
const openEditorWithAgents = async (user: ReturnType<typeof userEvent.setup>) => {
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({
@ -3191,6 +3221,7 @@ describe("TeamInfoView - disable_global_guardrails switch gating", () => {
await openEditForm();
expect(screen.queryByRole("switch", { name: /Disable all global guardrails/i })).not.toBeInTheDocument();
expect(screen.queryByRole("switch", { name: /Require Trace ID/i })).not.toBeInTheDocument();
});
it("shows the Disable all global guardrails switch to a proxy admin", async () => {
@ -3198,5 +3229,31 @@ describe("TeamInfoView - disable_global_guardrails switch gating", () => {
await openEditForm();
expect(await screen.findByRole("switch", { name: /Disable all global guardrails/i })).toBeInTheDocument();
expect(screen.getByRole("switch", { name: /Require Trace ID/i })).toBeInTheDocument();
});
it("shows the stored trace requirement to a non-admin on the team overview", async () => {
authState.userRole = "Internal User";
vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ metadata: { require_trace_id: true } }));
renderWithProviders(<TeamInfoView {...props} is_proxy_admin={false} />);
expect(await screen.findByText("Require Trace ID: Enabled")).toBeInTheDocument();
});
it("keeps the Require Trace ID value read-only for a proxy admin viewer", async () => {
authState.isViewOnly = true;
vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ metadata: { require_trace_id: true } }));
renderWithProviders(<TeamInfoView {...props} />);
expect(await screen.findByText("Require Trace ID: Enabled")).toBeInTheDocument();
expect(screen.queryByRole("switch", { name: /Require Trace ID/i })).not.toBeInTheDocument();
});
it("shows a disabled trace requirement to a non-admin on the team overview", async () => {
authState.userRole = "Internal User";
vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ metadata: {} }));
renderWithProviders(<TeamInfoView {...props} is_proxy_admin={false} />);
expect(await screen.findByText("Require Trace ID: Disabled")).toBeInTheDocument();
});
});

38
ui/litellm-dashboard/src/components/team/TeamInfo.tsx
View file

@ -146,6 +146,7 @@ const UI_MANAGED_METADATA_KEYS: ReadonlySet<string> = new Set([
"guardrails",
"opted_out_global_guardrails",
"disable_global_guardrails",
"require_trace_id",
]);
const TEAM_MODEL_BADGE_TONES: Record<TeamModelBadgeKind, StatusTone> = {
@ -418,6 +419,7 @@ const teamUpdateFieldsSchema = z.object({
.refine(estimateChecks.perModel.isValid, estimateChecks.perModel.message),
guardrails: z.array(z.string()).optional(),
disable_global_guardrails: z.boolean().optional(),
require_trace_id: z.boolean().optional(),
policies: z.array(z.string()).optional(),
access_group_ids: z.array(z.string()).optional(),
vector_stores: z.array(z.string()).optional(),
@ -476,6 +478,7 @@ const EMPTY_TEAM_UPDATE_VALUES: TeamUpdateFormValues = {
default_estimated_output_tokens_per_model: "",
guardrails: [],
disable_global_guardrails: false,
require_trace_id: false,
policies: [],
access_group_ids: [],
vector_stores: [],
@ -538,6 +541,7 @@ const toTeamFormValues = (info: TeamInfoRecord, effectiveGuardrails: string[]):
: "",
guardrails: effectiveGuardrails,
disable_global_guardrails: info.metadata?.disable_global_guardrails || false,
require_trace_id: info.metadata?.require_trace_id === true,
policies: info.policies || [],
access_group_ids: info.access_group_ids || [],
vector_stores: info.object_permission?.vector_stores || [],
@ -628,11 +632,12 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
const [teamModelMaxBudget, setTeamModelMaxBudget] = useState<ModelMaxBudget>({});
const routerSettingsRef = React.useRef<RouterSettingsAccordionRef>(null);
const [organization, setOrganization] = useState<Organization | null>(null);
const { userRole } = useAuthorized();
const { userRole, isViewOnly } = useAuthorized();
const accessGroupNames = useModelAccessGroupNames();
const { data: allMcpServers = [], isError: mcpServersFailed, isLoading: mcpServersLoading } = useMCPServers();
const { data: allMcpToolsets = [], isError: mcpToolsetsFailed, isLoading: mcpToolsetsLoading } = useMCPToolsets();
const { data: allAccessGroups = [], isError: accessGroupsFailed, isLoading: accessGroupsLoading } = useAccessGroups();
const canEditAsProxyAdmin = is_proxy_admin && !isViewOnly;
const canEditTeamEstimates = isProxyAdminRole(userRole);
const teamEstimateTooltip = estimateTooltips(canEditTeamEstimates, "team");
const { data: userOrganizations = [] } = useOrganizations();
@ -1033,6 +1038,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
memberBudgetAlertEmails !== undefined && Object.keys(memberBudgetAlertEmails).length > 0
? { [TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY]: memberBudgetAlertEmails }
: {};
const storedRequireTraceIdMetadata =
info.metadata && Object.prototype.hasOwnProperty.call(info.metadata, "require_trace_id")
? { require_trace_id: info.metadata.require_trace_id }
: {};
const requireTraceIdMetadata =
canEditAsProxyAdmin && values.require_trace_id !== (info.metadata?.require_trace_id === true)
? { require_trace_id: values.require_trace_id === true }
: storedRequireTraceIdMetadata;
const updateData: any = {
team_id: teamId,
@ -1053,6 +1066,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
opted_out_global_guardrails: optedOutGlobalGuardrails,
...(values.logging_settings?.length > 0 ? { logging: values.logging_settings } : {}),
disable_global_guardrails: killSwitchOnAtSave,
...requireTraceIdMetadata,
...(estimatedOutputTokens !== null ? { default_estimated_output_tokens: estimatedOutputTokens } : {}),
...(estimatedOutputTokensPerModel !== undefined
? { default_estimated_output_tokens_per_model: estimatedOutputTokensPerModel }
@ -1354,6 +1368,13 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
</div>
</Card>
<Card className="block p-6">
<p>Request Settings</p>
<div className="mt-2">
<p>Require Trace ID: {info.metadata?.require_trace_id === true ? "Enabled" : "Disabled"}</p>
</div>
</Card>
<Card className="block p-6">
<p>Models</p>
<div className="mt-2 flex flex-wrap gap-2">
@ -1964,6 +1985,21 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
</FormField>
)}
{canEditAsProxyAdmin && (
<FormField
control={form.control}
name="require_trace_id"
label={labelWithHint(
"Require Trace ID",
"Reject LLM, MCP and agent requests from this team that do not send a trace ID (x-litellm-trace-id header, LLM requests can also use traceparent or metadata.trace_id)",
)}
>
{({ id, value, onChange }) => (
<Switch id={id} checked={value === true} onCheckedChange={onChange} />
)}
</FormField>
)}
{canViewPolicies && (
<FormField
control={form.control}

8
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -16878,6 +16878,7 @@ export interface paths {
* - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
* - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
* - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
* - require_trace_id: Optional[bool] - Reject LLM, MCP and agent requests from this team that carry no trace ID.
* - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
* - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
* - team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
@ -17106,6 +17107,7 @@ export interface paths {
* - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
* - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
* - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
* - require_trace_id: Optional[bool] - Reject LLM, MCP and agent requests from this team that carry no trace ID.
* - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
* - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
* - team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
@ -38443,6 +38445,8 @@ export interface components {
policies?: string[] | null;
/** Prompts */
prompts?: string[] | null;
/** Require Trace Id */
require_trace_id?: boolean | null;
/** Router Settings */
router_settings?: components["schemas"]["UpdateRouterConfig"] | null;
/** Rpm Limit */
@ -39528,6 +39532,8 @@ export interface components {
policies?: string[] | null;
/** Prompts */
prompts?: string[] | null;
/** Require Trace Id */
require_trace_id?: boolean | null;
/** Router Settings */
router_settings?: components["schemas"]["UpdateRouterConfig"] | null;
/** Rpm Limit */
@ -47501,6 +47507,8 @@ export interface components {
policies?: string[] | null;
/** Prompts */
prompts?: string[] | null;
/** Require Trace Id */
require_trace_id?: boolean | null;
/** Router Settings */
router_settings?: components["schemas"]["UpdateRouterConfig"] | null;
/** Rpm Limit */