diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 477edf47c5f..fb9c49fc898 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 = [ diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 60939fc622e..7c62e4f4147 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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", ()) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index a9385c637e8..6a70a0137a4 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -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, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 69a5e2e233a..c78434eda83 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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"] = {} diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index c515108ab11..ed5843f4bce 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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() diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 4da62c29e86..6eb60f06ed0 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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, diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index a4a2a70ed51..0da1148ac84 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -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, diff --git a/tests/unit/proxy/management_endpoints/test_common_utils.py b/tests/unit/proxy/management_endpoints/test_common_utils.py index 43c74286d75..d783f4c64f6 100644 --- a/tests/unit/proxy/management_endpoints/test_common_utils.py +++ b/tests/unit/proxy/management_endpoints/test_common_utils.py @@ -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 ( diff --git a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py index 51a067f250f..a593431523b 100644 --- a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py +++ b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -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"}) diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 5e9c350431c..b81c1c9f3d5 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -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 diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index add6aabb224..b2d76af7f2d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -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(): """ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/page.tsx index 2d99041c1ce..c074401e6bc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/page.tsx @@ -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 ; + const { accessToken, userId, userRole, isViewOnly, premiumUser } = useAuthorized(); + return ( + + ); } diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 85c42b2d09b..7ab6ec2679e 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -552,6 +552,17 @@ describe("Teams - Create Team CTA is grouped with the tabs on the left", () => { renderWithQueryClient(); expect(screen.queryByTestId("create-team-button")).not.toBeInTheDocument(); }); + + it("should hide the Require Trace ID switch from a view-only Admin session", async () => { + renderWithQueryClient(); + + 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(); + 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, + }), + ); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 9e7f622bdab..0b773fd0f04 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -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 = ({ accessToken, userID, userRole, premiumUser = false }) => { +const Teams: React.FC = ({ 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 = ({ accessToken, userID, userRole, premiumUser const [editTeam, setEditTeam] = useState(false); const [isTeamModalVisible, setIsTeamModalVisible] = useState(false); + const canEditAsProxyAdmin = !isViewOnly && isProxyAdminRole(userRole || ""); const [userModels, setUserModels] = useState([]); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [teamToDelete, setTeamToDelete] = useState(null); @@ -545,7 +550,12 @@ const Teams: React.FC = ({ 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 = ({ accessToken, userID, userRole, premiumUser )} )} + {canEditAsProxyAdmin && ( + + {({ id, value, onChange }) => ( + + )} + + )} {canViewPolicies && ( ({ 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) => { - vi.mocked(networking.teamInfoCall).mockResolvedValue(storedTeam()); + const openEditor = async ( + user: ReturnType, + 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(); + renderWithProviders(); 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) => { 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(); + + 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(); + + 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(); + + expect(await screen.findByText("Require Trace ID: Disabled")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 3f3c3a9ecc0..3f933c4a4f6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -146,6 +146,7 @@ const UI_MANAGED_METADATA_KEYS: ReadonlySet = new Set([ "guardrails", "opted_out_global_guardrails", "disable_global_guardrails", + "require_trace_id", ]); const TEAM_MODEL_BADGE_TONES: Record = { @@ -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 = ({ const [teamModelMaxBudget, setTeamModelMaxBudget] = useState({}); const routerSettingsRef = React.useRef(null); const [organization, setOrganization] = useState(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 = ({ 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 = ({ 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 = ({ + +

Request Settings

+
+

Require Trace ID: {info.metadata?.require_trace_id === true ? "Enabled" : "Disabled"}

+
+ +

Models

@@ -1964,6 +1985,21 @@ const TeamInfoView: React.FC = ({ )} + {canEditAsProxyAdmin && ( + + {({ id, value, onChange }) => ( + + )} + + )} + {canViewPolicies && (