From cd87fee9c5c1557069920b6385a357c27bdd88e7 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 3 Aug 2026 18:37:45 -0700 Subject: [PATCH] feat(team): custom metadata validation hook for team create and update (#33353) * feat(team): custom metadata validation hook for team create and update Operators can point general_settings.custom_team_metadata_validate at an async Python function that validates team metadata before /team/new, POST /team/update, and PATCH /team/{team_id} commit their writes. The hook receives the metadata that will actually be written (the merged result on PATCH) plus the stored metadata and requester context, and fails closed: a rejected value returns the function's own message as a 400 while any exception or timeout blocks the write with a configurable generic message as a 503. Premium-gated like enforced_params. * fix(team): validate metadata before model alias writes and strip system keys from validator input Review follow-ups on the team metadata validation hook: run the validator before the model_aliases table insert so a rejected create leaves no orphaned model rows, strip system-managed keys from existing_metadata so the validator sees symmetric input on both fields, and accept class instances exposing an async __call__ as validators. Adds a three-way validator implementation matrix (allowlist function, HTTP-service-backed function, immutability-enforcing class instance) driven through the real create, update, and patch endpoints, including an HTTP stub service and outage coverage. * test(team): run the metadata validation matrix against the DB-backed proxy in CI Adds the validator matrix to the proxy_store_model_in_db_tests CircleCI job so every scenario runs full e2e against a Postgres-backed proxy. The proxy config registers a dispatching validator that routes each request to one of the three implementations via a metadata key and accepts anything that does not opt in, keeping the rest of the suite unaffected. CI starts a stand-in cost center service on the host for the HTTP-backed implementation, reached from the container via host.docker.internal, and the outage path targets a closed port to prove the fail-closed 503 without stopping services. * feat(ui): edit team metadata as key-value pairs in team create and edit forms The team create and edit forms asked for metadata as a raw JSON blob in a textarea buried under Additional Settings. Both forms now render a key-value pair editor directly under the TPM/RPM limit fields, backed by a shared MetadataKeyValueFields component. Values round-trip losslessly: non-string values display as JSON and parse back to their typed form on save, and JSON-ambiguous strings are quoted so their type survives the trip. The edit form hides UI-managed keys (logging, guardrails, model rate limits, etc.) that dedicated controls already own and re-add on save. * fix(ui): explain typed JSON parsing in the team metadata help text * feat(team): schema-driven metadata fields from team_metadata_schema config * refactor(team): render schema metadata fields as locked key-value rows, drop allowed_values * refactor(team): schema fields reduce to key and label, tag-rendered keys, clean rejection toasts * refactor(ui): prepopulate declared metadata keys as ordinary key-value rows * fix(team): let non-admin dashboard users read the team metadata schema * test(proxy): pin timeout wiring, boundary, and error-message contracts for team metadata validation * fix(proxy): use pooled async httpx client in the e2e team metadata validator example * refactor(team): satisfy staging lint ratchets inherited by the merge --- .circleci/config.yml | 23 + litellm/proxy/_types.py | 2 + .../custom_team_metadata_validate.py | 40 + .../store_model_db_config.yaml | 3 + .../team_metadata_validator_e2e.py | 107 +++ .../management_endpoints/team_endpoints.py | 57 +- .../team_metadata_validation.py | 193 +++++ litellm/proxy/proxy_server.py | 16 + .../management_endpoints/team_endpoints.py | 22 +- .../cost_center_service.py | 54 ++ .../test_team_metadata_validation_e2e.py | 171 +++++ .../test_team_endpoints.py | 318 ++++++++ .../team_metadata_validator_impls.py | 92 +++ .../test_team_metadata_validation.py | 704 ++++++++++++++++++ ui/litellm-dashboard/eslint-suppressions.json | 10 + .../hooks/teams/useTeamMetadataSchema.test.ts | 39 + .../hooks/teams/useTeamMetadataSchema.ts | 37 + .../src/components/Teams.test.tsx | 210 ++++++ ui/litellm-dashboard/src/components/Teams.tsx | 48 +- .../MetadataKeyValueFields.test.tsx | 282 +++++++ .../MetadataKeyValueFields.tsx | 135 ++++ .../src/components/networking.tsx | 4 +- .../src/components/team/TeamInfo.test.tsx | 137 ++++ .../src/components/team/TeamInfo.tsx | 60 +- .../src/lib/http/client.test.ts | 46 +- ui/litellm-dashboard/src/lib/http/client.ts | 32 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 66 ++ 27 files changed, 2844 insertions(+), 64 deletions(-) create mode 100644 litellm/proxy/example_config_yaml/custom_team_metadata_validate.py create mode 100644 litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py create mode 100644 litellm/proxy/management_helpers/team_metadata_validation.py create mode 100644 tests/store_model_in_db_tests/cost_center_service.py create mode 100644 tests/store_model_in_db_tests/test_team_metadata_validation_e2e.py create mode 100644 tests/test_litellm/proxy/management_helpers/team_metadata_validator_impls.py create mode 100644 tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMetadataSchema.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMetadataSchema.ts create mode 100644 ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.test.tsx create mode 100644 ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.tsx diff --git a/.circleci/config.yml b/.circleci/config.yml index c3a34a97b3b..9063ba1d4a6 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -163,6 +163,26 @@ commands: done echo "fake OpenAI endpoint did not become ready" >&2 exit 1 + start_cost_center_service: + description: "Start the stand-in cost center validation service (tests/store_model_in_db_tests/cost_center_service.py) on host port 9414 and wait until healthy. The proxy's team-metadata validator (team_metadata_validator_e2e.py, impl 'http') reaches it via TEAM_METADATA_VALIDATION_SERVICE_URL=http://host.docker.internal:9414/validate. Run after uv deps are synced." + steps: + - run: + name: Start cost center validation service + background: true + command: | + uv run --no-sync python tests/store_model_in_db_tests/cost_center_service.py --host 0.0.0.0 --port 9414 + - run: + name: Wait for cost center validation service + command: | + for i in $(seq 1 30); do + if curl -sf http://localhost:9414/health >/dev/null 2>&1; then + echo "cost center validation service is up" + exit 0 + fi + sleep 1 + done + echo "cost center validation service did not become ready" >&2 + exit 1 setup_litellm_enterprise_pip: steps: - run: @@ -2264,6 +2284,7 @@ jobs: uv sync --frozen --all-groups --all-extras --python 3.12 - start_postgres - start_fake_openai_endpoint + - start_cost_center_service - attach_workspace: at: ~/project - run: @@ -2283,11 +2304,13 @@ jobs: -e STORE_MODEL_IN_DB="True" \ -e LITELLM_MASTER_KEY="sk-1234" \ -e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \ + -e TEAM_METADATA_VALIDATION_SERVICE_URL=http://host.docker.internal:9414/validate \ -e LITELLM_LICENSE=$LITELLM_LICENSE \ -e LITELLM_LOG=ERROR \ --add-host host.docker.internal:host-gateway \ --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \ + -v $(pwd)/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py:/app/team_metadata_validator_e2e.py \ litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3ccf3ea9952..268b08d0f2e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -543,6 +543,7 @@ class LiteLLMRoutes(enum.Enum): "/v2/team/list", "/organization/list", "/team/available", + "/team/metadata_schema", "/user/info", "/v2/user/info", "/model/info", @@ -609,6 +610,7 @@ class LiteLLMRoutes(enum.Enum): "/team/block", "/team/unblock", "/team/available", + "/team/metadata_schema", "/team/permissions_list", "/team/permissions_update", "/team/permissions_bulk_update", diff --git a/litellm/proxy/example_config_yaml/custom_team_metadata_validate.py b/litellm/proxy/example_config_yaml/custom_team_metadata_validate.py new file mode 100644 index 00000000000..1ca21bc3104 --- /dev/null +++ b/litellm/proxy/example_config_yaml/custom_team_metadata_validate.py @@ -0,0 +1,40 @@ +"""Example validator for `general_settings.custom_team_metadata_validate`. + +Wire it up in the proxy config: + +```yaml +general_settings: + custom_team_metadata_validate: custom_team_metadata_validate.validate_team_metadata + team_metadata_validation_timeout: 5 + team_metadata_validation_error_message: "Validation service unavailable, contact your admin." +``` + +Return `valid=False` with an `error_message` to reject the write with that +message (HTTP 400). Raise any exception (for example, when the upstream +validation service is unreachable) to fail closed with the generic +`team_metadata_validation_error_message` (HTTP 503). +""" + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataValidationPayload, + TeamMetadataValidationResult, +) + +VALID_COST_CENTERS = frozenset({"CC-1001", "CC-1002", "CC-2001"}) + + +async def validate_team_metadata( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + cost_center = payload.metadata.get("cost_center") + if cost_center is None: + return TeamMetadataValidationResult( + valid=False, + error_message="Team metadata must include a cost_center. Contact the FinOps team.", + ) + if cost_center not in VALID_COST_CENTERS: + return TeamMetadataValidationResult( + valid=False, + error_message=f"Cost center {cost_center} is not recognized. Contact the FinOps team.", + ) + return TeamMetadataValidationResult(valid=True) diff --git a/litellm/proxy/example_config_yaml/store_model_db_config.yaml b/litellm/proxy/example_config_yaml/store_model_db_config.yaml index 5b77a53b4b7..3f19ab39a02 100644 --- a/litellm/proxy/example_config_yaml/store_model_db_config.yaml +++ b/litellm/proxy/example_config_yaml/store_model_db_config.yaml @@ -7,4 +7,7 @@ model_list: general_settings: store_model_in_db: true + custom_team_metadata_validate: team_metadata_validator_e2e.validate_team_metadata + team_metadata_validation_timeout: 5 + team_metadata_validation_error_message: "Cost center validation is unavailable right now; the team was not saved. Contact FinOps." diff --git a/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py b/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py new file mode 100644 index 00000000000..23390ffec1c --- /dev/null +++ b/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py @@ -0,0 +1,107 @@ +"""Dispatching team metadata validator for the store_model_in_db e2e suite. + +The suite runs one proxy with one config, so a single registered validator +dispatches to one of three independent implementations chosen per request via +the `_e2e_validator_impl` metadata key: + +- `allowlist`: requires `cost_center` and checks it against a static set +- `http`: POSTs the metadata to the cost center service at + `TEAM_METADATA_VALIDATION_SERVICE_URL`; transport errors raise (fail closed) +- `http_down`: like `http` but targets a closed port, proving the 503 path +- `immutable`: requires `cost_center` and forbids changing it once set + +A request whose metadata carries no `_e2e_validator_impl` key is accepted +untouched, so the rest of the suite's team operations are unaffected. An +unknown impl value raises, which the proxy converts to the fail-closed 503. +""" + +import os + +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataValidationPayload, + TeamMetadataValidationResult, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + +ALLOWED_COST_CENTERS = frozenset({"CC-1001", "CC-1002"}) +CLOSED_PORT_URL = "http://127.0.0.1:9/validate" + + +async def _validate_allowlist(payload: TeamMetadataValidationPayload) -> TeamMetadataValidationResult: + cost_center = payload.metadata.get("cost_center") + if cost_center is None: + return TeamMetadataValidationResult( + valid=False, + error_message="cost_center is required in team metadata. Contact the FinOps team.", + ) + if cost_center not in ALLOWED_COST_CENTERS: + return TeamMetadataValidationResult( + valid=False, + error_message=f"Cost center {cost_center} is not recognized. Contact the FinOps team.", + ) + return TeamMetadataValidationResult(valid=True) + + +async def _validate_via_http(payload: TeamMetadataValidationPayload, service_url: str) -> TeamMetadataValidationResult: + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + response = await client.post( + service_url, + json={ # mutable-ok: httpx serializes the request body from a plain dict + "operation": payload.operation, + "metadata": payload.metadata, + }, + timeout=2.0, + ) + body = response.json() + if body.get("ok") is True: + return TeamMetadataValidationResult(valid=True) + return TeamMetadataValidationResult( + valid=False, + error_message=body.get("reason", "Rejected by the cost center service."), + ) + + +class _ImmutableCostCenterValidator: + def __init__(self, immutable_key: str = "cost_center") -> None: + self.immutable_key = immutable_key + + async def __call__(self, payload: TeamMetadataValidationPayload) -> TeamMetadataValidationResult: + current = payload.metadata.get(self.immutable_key) + if current is None: + return TeamMetadataValidationResult( + valid=False, + error_message=f"{self.immutable_key} is required in team metadata. Contact the FinOps team.", + ) + if payload.operation == "update" and payload.existing_metadata is not None: + prior = payload.existing_metadata.get(self.immutable_key) + if prior is not None and prior != current: + return TeamMetadataValidationResult( + valid=False, + error_message=( + f"{self.immutable_key} is immutable once set " + f"(stored: {prior}, requested: {current}). Contact the FinOps team." + ), + ) + return TeamMetadataValidationResult(valid=True) + + +_IMMUTABLE_VALIDATOR = _ImmutableCostCenterValidator() + + +async def validate_team_metadata( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + impl = payload.metadata.get("_e2e_validator_impl") + if impl is None: + return TeamMetadataValidationResult(valid=True) + if impl == "allowlist": + return await _validate_allowlist(payload) + if impl == "http": + service_url = os.environ.get("TEAM_METADATA_VALIDATION_SERVICE_URL", "http://localhost:9414/validate") + return await _validate_via_http(payload, service_url) + if impl == "http_down": + return await _validate_via_http(payload, CLOSED_PORT_URL) + if impl == "immutable": + return await _IMMUTABLE_VALIDATOR(payload) + raise ValueError(f"unknown _e2e_validator_impl: {impl}") diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 4dd86e5769d..bdd7e18a850 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -113,6 +113,10 @@ from litellm.proxy.management_helpers.object_permission_utils import ( from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + validate_team_metadata_if_configured, +) from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, @@ -147,6 +151,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( TeamListResponse, TeamMemberAddResult, TeamMemberInfoResponse, + TeamMetadataSchemaResponse, UpdateTeamMemberPermissionsRequest, ) @@ -1287,6 +1292,18 @@ async def new_team( _check_passthrough_routes_caller_permission(data, user_api_key_dict, entity="team") + if isinstance(data.metadata, dict): + TeamMemberBudgetHandler.strip_system_managed_metadata_keys(data.metadata) + + await validate_team_metadata_if_configured( + operation="create", + metadata=data.metadata, + existing_metadata=None, + team_id=data.team_id, + team_alias=data.team_alias, + user_api_key_dict=user_api_key_dict, + ) + ## ADD TO MODEL TABLE _model_id = None if data.model_aliases is not None and isinstance(data.model_aliases, dict): @@ -1301,9 +1318,6 @@ async def new_team( _model_id = model_dict.id - ## Create Team Member Budget Table - if isinstance(data.metadata, dict): - TeamMemberBudgetHandler.strip_system_managed_metadata_keys(data.metadata) data_json = data.json() ## Handle Object Permission - MCP, Vector Stores etc. @@ -1965,6 +1979,25 @@ async def update_team( if isinstance(updated_kv.get("metadata"), dict): TeamMemberBudgetHandler.strip_system_managed_metadata_keys(updated_kv["metadata"]) + if "metadata" in updated_kv: + stored_metadata = ( + { # mutable-ok: the validator payload's isinstance guard requires a plain dict + key: value + for key, value in existing_team_row.metadata.items() + if key not in TeamMemberBudgetHandler.SYSTEM_MANAGED_METADATA_KEYS + } + if isinstance(existing_team_row.metadata, dict) + else None + ) + await validate_team_metadata_if_configured( + operation="update", + metadata=updated_kv.get("metadata"), + existing_metadata=stored_metadata, + team_id=data.team_id, + team_alias=data.team_alias if data.team_alias is not None else existing_team_row.team_alias, + user_api_key_dict=user_api_key_dict, + ) + # Check budget_duration and budget_reset_at _set_budget_reset_at(data, updated_kv) @@ -4179,6 +4212,24 @@ async def unblock_team( return record +@router.get( + "/team/metadata_schema", + tags=["team management"], # mutable-ok: fastapi's decorator signature types tags as a list + dependencies=(Depends(user_api_key_auth),), + response_model=TeamMetadataSchemaResponse, +) +async def get_team_metadata_schema(): + """ + Get the team metadata fields declared in ``general_settings.team_metadata_schema``. + + The UI uses this to prepopulate the team metadata form with the declared + keys. Returns an empty ``fields`` list when no schema is configured. This + schema is advisory; server-side enforcement stays with + ``custom_team_metadata_validate``. + """ + return TeamMetadataSchemaResponse(fields=TEAM_METADATA_SCHEMA_REGISTRY.get()) + + @router.get("/team/available") async def list_available_teams( http_request: Request, diff --git a/litellm/proxy/management_helpers/team_metadata_validation.py b/litellm/proxy/management_helpers/team_metadata_validation.py new file mode 100644 index 00000000000..2b8ba05c602 --- /dev/null +++ b/litellm/proxy/management_helpers/team_metadata_validation.py @@ -0,0 +1,193 @@ +"""Custom validation of team metadata on team create/update. + +Operators point `general_settings.custom_team_metadata_validate` at an async +Python function (loaded via `get_instance_fn`, like `custom_key_generate`). +The function receives a `TeamMetadataValidationPayload` and returns a +`TeamMetadataValidationResult`. The proxy awaits it before committing a team +write and fails closed: a rejected value surfaces the function's own message +(HTTP 400), while any raised exception or timeout blocks the write with a +generic message (HTTP 503). +""" + +import asyncio +import inspect +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import Literal, Protocol + +from fastapi import HTTPException, status +from pydantic import BaseModel, JsonValue, TypeAdapter + +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.types.proxy.management_endpoints.team_endpoints import ( + TeamMetadataFieldSchema, +) + +DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS = 5.0 +DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE = ( + "Team metadata validation is currently unavailable, so the team was not saved. Contact your proxy admin." +) +DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE = "Team metadata failed validation." + + +class TeamMetadataRequester(BaseModel): + user_id: str | None = None + user_email: str | None = None + user_role: str | None = None + + +class TeamMetadataValidationPayload(BaseModel): + operation: Literal["create", "update"] + metadata: Mapping[str, JsonValue] + existing_metadata: Mapping[str, JsonValue] | None = None + team_id: str | None = None + team_alias: str | None = None + requester: TeamMetadataRequester + + +class TeamMetadataValidationResult(BaseModel): + valid: bool + error_message: str | None = None + + +_EMPTY_METADATA: Mapping[str, JsonValue] = MappingProxyType({}) + + +class TeamMetadataValidator(Protocol): + def __call__(self, payload: TeamMetadataValidationPayload, /) -> Awaitable[TeamMetadataValidationResult]: ... + + +class TeamMetadataValidatorRegistry: + def __init__(self) -> None: + self._validator: TeamMetadataValidator | None = None + + def set(self, validator: TeamMetadataValidator | None) -> None: + self._validator = validator + + def get(self) -> TeamMetadataValidator | None: + return self._validator + + +TEAM_METADATA_VALIDATOR_REGISTRY = TeamMetadataValidatorRegistry() + +_TEAM_METADATA_SCHEMA_ADAPTER: TypeAdapter[tuple[TeamMetadataFieldSchema, ...]] = TypeAdapter( + tuple[TeamMetadataFieldSchema, ...] +) + + +def parse_team_metadata_schema(raw_schema: object) -> tuple[TeamMetadataFieldSchema, ...]: + """Parse ``general_settings.team_metadata_schema``; raises on a malformed schema so config load fails fast.""" + if raw_schema is None: + return () + fields = _TEAM_METADATA_SCHEMA_ADAPTER.validate_python(raw_schema) + keys = tuple(field.key for field in fields) + duplicate_keys = sorted(frozenset(key for key in keys if keys.count(key) > 1)) + if duplicate_keys: + raise ValueError(f"team_metadata_schema contains duplicate keys: {', '.join(duplicate_keys)}") + return fields + + +class TeamMetadataSchemaRegistry: + def __init__(self) -> None: + self._fields: tuple[TeamMetadataFieldSchema, ...] = () + + def set(self, fields: tuple[TeamMetadataFieldSchema, ...]) -> None: + self._fields = fields + + def get(self) -> tuple[TeamMetadataFieldSchema, ...]: + return self._fields + + +TEAM_METADATA_SCHEMA_REGISTRY = TeamMetadataSchemaRegistry() + + +async def run_team_metadata_validation( + validator: TeamMetadataValidator, + payload: TeamMetadataValidationPayload, + premium_user: bool, + timeout_seconds: float, + unavailable_message: str, +) -> None: + if premium_user is not True: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": f"custom_team_metadata_validate is an Enterprise feature. {CommonProxyErrors.not_premium_user.value}" + }, + ) + if not ( + inspect.iscoroutinefunction(validator) or inspect.iscoroutinefunction(getattr(validator, "__call__", None)) + ): + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": "custom_team_metadata_validate must be an async function" + }, + ) + + try: + raw_result = await asyncio.wait_for(validator(payload), timeout=timeout_seconds) + result = TeamMetadataValidationResult.model_validate(raw_result) + except Exception: # noqa: BLE001 # fail closed: any validator failure must block the team write + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={"error": unavailable_message}, # mutable-ok: HTTPException.detail has no immutable form + ) + + if not result.valid: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": result.error_message or DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE + }, + ) + + +def _read_timeout_seconds(general_settings: Mapping[str, object]) -> float: + raw_timeout = general_settings.get("team_metadata_validation_timeout") + if isinstance(raw_timeout, (int, float)) and not isinstance(raw_timeout, bool) and raw_timeout > 0: + return float(raw_timeout) + return DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS + + +def _read_unavailable_message(general_settings: Mapping[str, object]) -> str: + raw_message = general_settings.get("team_metadata_validation_error_message") + if isinstance(raw_message, str) and raw_message.strip(): + return raw_message + return DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE + + +async def validate_team_metadata_if_configured( + operation: Literal["create", "update"], + metadata: Mapping[str, JsonValue] | None, + existing_metadata: Mapping[str, JsonValue] | None, + team_id: str | None, + team_alias: str | None, + user_api_key_dict: UserAPIKeyAuth, + registry: TeamMetadataValidatorRegistry = TEAM_METADATA_VALIDATOR_REGISTRY, +) -> None: + from litellm.proxy.proxy_server import general_settings, premium_user + + validator = registry.get() + if validator is None: + return + + payload = TeamMetadataValidationPayload( + operation=operation, + metadata=metadata if isinstance(metadata, dict) else _EMPTY_METADATA, + existing_metadata=existing_metadata if isinstance(existing_metadata, dict) else None, + team_id=team_id, + team_alias=team_alias, + requester=TeamMetadataRequester( + user_id=user_api_key_dict.user_id, + user_email=user_api_key_dict.user_email, + user_role=user_api_key_dict.user_role.value if user_api_key_dict.user_role is not None else None, + ), + ) + await run_team_metadata_validation( + validator=validator, + payload=payload, + premium_user=premium_user, + timeout_seconds=_read_timeout_seconds(general_settings), + unavailable_message=_read_unavailable_message(general_settings), + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ac45898ce0b..2ff39164d80 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -461,6 +461,11 @@ from litellm.proxy.management_helpers.audit_logs import ( create_audit_log_for_update, create_object_audit_log, ) +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + TEAM_METADATA_VALIDATOR_REGISTRY, + parse_team_metadata_schema, +) from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.middleware.billable_request_metrics_middleware import ( BillableRequestMetricsMiddleware, @@ -783,6 +788,8 @@ def cleanup_router_config_variables(): user_custom_auth_path = None user_custom_key_generate = None user_custom_key_update = None + TEAM_METADATA_VALIDATOR_REGISTRY.set(None) + TEAM_METADATA_SCHEMA_REGISTRY.set(()) user_custom_sso = None user_custom_ui_sso_sign_in_handler = None use_background_health_checks = None @@ -3499,6 +3506,7 @@ _DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: dict[str, tuple[str, ...]] = { "custom_auth", "custom_key_generate", "custom_key_update", + "custom_team_metadata_validate", "custom_sso", "custom_ui_sso_sign_in_handler", ), @@ -4829,6 +4837,14 @@ class ProxyConfig: if custom_key_update is not None: user_custom_key_update = get_instance_fn(value=custom_key_update, config_file_path=config_file_path) + custom_team_metadata_validate = general_settings.get("custom_team_metadata_validate", None) + TEAM_METADATA_VALIDATOR_REGISTRY.set( + get_instance_fn(value=custom_team_metadata_validate, config_file_path=config_file_path) + if custom_team_metadata_validate is not None + else None + ) + TEAM_METADATA_SCHEMA_REGISTRY.set(parse_team_metadata_schema(general_settings.get("team_metadata_schema"))) + custom_sso = general_settings.get("custom_sso", None) if custom_sso is not None: user_custom_sso = get_instance_fn(value=custom_sso, config_file_path=config_file_path) diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 2d9387da956..f3d0529a7ea 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,13 +1,12 @@ from typing import Any, Dict, List, Literal, Optional, Union -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, Field from litellm.proxy._types import ( KeyManagementRoutes, LiteLLM_DeletedTeamTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, - LiteLLM_UserTable, Member, ) @@ -125,3 +124,22 @@ class TeamMemberInfoResponse(LiteLLM_TeamMembership): role: Optional[str] = None user_email: Optional[str] = None team_alias: Optional[str] = None + + +class TeamMetadataFieldSchema(BaseModel): + """One declared team metadata field from ``general_settings.team_metadata_schema``. + + Advisory only: the UI uses it to prepopulate the team metadata form. + Enforcement stays with ``custom_team_metadata_validate``. + """ + + model_config = ConfigDict(extra="forbid") + + key: str = Field(min_length=1) + label: Optional[str] = None + + +class TeamMetadataSchemaResponse(BaseModel): + """Response for GET /team/metadata_schema; ``fields`` is empty when no schema is configured.""" + + fields: tuple[TeamMetadataFieldSchema, ...] diff --git a/tests/store_model_in_db_tests/cost_center_service.py b/tests/store_model_in_db_tests/cost_center_service.py new file mode 100644 index 00000000000..c2820dfbe7e --- /dev/null +++ b/tests/store_model_in_db_tests/cost_center_service.py @@ -0,0 +1,54 @@ +"""Stand-in cost center validation service for the team metadata e2e tests. + +Accepts POST /validate with {"operation": ..., "metadata": {...}} and answers +{"ok": true} or {"ok": false, "reason": ...} based on a static allowlist. +GET /health answers 200 for the CI wait loop. +""" + +import argparse +import json +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +ALLOWED_COST_CENTERS = {"CC-1001", "CC-1002"} + + +class CostCenterHandler(BaseHTTPRequestHandler): + def _respond(self, status: int, body: dict) -> None: + payload = json.dumps(body).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def do_GET(self): + if self.path == "/health": + self._respond(200, {"status": "healthy"}) + return + self._respond(404, {"error": "not found"}) + + def do_POST(self): + if self.path != "/validate": + self._respond(404, {"error": "not found"}) + return + length = int(self.headers.get("Content-Length", 0)) + body = json.loads(self.rfile.read(length) or b"{}") + cost_center = (body.get("metadata") or {}).get("cost_center") + if cost_center is None: + self._respond(200, {"ok": False, "reason": "cost_center missing per cost center service"}) + elif cost_center not in ALLOWED_COST_CENTERS: + self._respond(200, {"ok": False, "reason": f"cost center {cost_center} rejected by cost center service"}) + else: + self._respond(200, {"ok": True}) + + def log_message(self, format, *args): + pass + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=9414) + args = parser.parse_args() + print(f"cost center service listening on {args.host}:{args.port}") + ThreadingHTTPServer((args.host, args.port), CostCenterHandler).serve_forever() diff --git a/tests/store_model_in_db_tests/test_team_metadata_validation_e2e.py b/tests/store_model_in_db_tests/test_team_metadata_validation_e2e.py new file mode 100644 index 00000000000..050b59ef426 --- /dev/null +++ b/tests/store_model_in_db_tests/test_team_metadata_validation_e2e.py @@ -0,0 +1,171 @@ +"""E2E matrix for custom team metadata validation against the DB-backed proxy. + +The proxy (see store_model_db_config.yaml) registers +team_metadata_validator_e2e.validate_team_metadata, which dispatches per +request to one of three independent implementations via the +`_e2e_validator_impl` metadata key: a static allowlist function, an +HTTP-backed function calling the cost center service started by CI, and an +immutability-enforcing class instance. Metadata without the dispatch key is +accepted untouched, so the rest of this suite is unaffected. +""" + +import os +import uuid + +import httpx +import pytest + +PROXY_BASE_URL = os.getenv("PROXY_BASE_URL", "http://localhost:4000") +MASTER_KEY = os.getenv("LITELLM_MASTER_KEY", "sk-1234") +HEADERS = {"Authorization": f"Bearer {MASTER_KEY}", "Content-Type": "application/json"} + +UNAVAILABLE_MESSAGE = "Cost center validation is unavailable right now; the team was not saved. Contact FinOps." + +IMPLS = ["allowlist", "http", "immutable"] + +REQUIRED_MESSAGES = { + "allowlist": "cost_center is required in team metadata", + "http": "cost_center missing per cost center service", + "immutable": "cost_center is required in team metadata", +} +UNKNOWN_MESSAGES = { + "allowlist": "is not recognized", + "http": "rejected by cost center service", +} + + +def _meta(impl, **fields): + return {"_e2e_validator_impl": impl, **fields} + + +def _create_team(metadata, team_id=None): + body = {"team_alias": f"meta-validate-{uuid.uuid4().hex[:8]}"} + if team_id is not None: + body["team_id"] = team_id + if metadata is not None: + body["metadata"] = metadata + return httpx.post(f"{PROXY_BASE_URL}/team/new", headers=HEADERS, json=body, timeout=30) + + +def _patch_team(team_id, body): + return httpx.patch(f"{PROXY_BASE_URL}/team/{team_id}", headers=HEADERS, json=body, timeout=30) + + +def _post_update(team_id, body): + return httpx.post(f"{PROXY_BASE_URL}/team/update", headers=HEADERS, json={"team_id": team_id, **body}, timeout=30) + + +def _team_info(team_id): + return httpx.get(f"{PROXY_BASE_URL}/team/info", headers=HEADERS, params={"team_id": team_id}, timeout=30) + + +def _delete_team(team_id): + httpx.post(f"{PROXY_BASE_URL}/team/delete", headers=HEADERS, json={"team_ids": [team_id]}, timeout=30) + + +@pytest.fixture +def team_with_cost_center(request): + impl = request.param + team_id = f"meta-validate-{impl}-{uuid.uuid4().hex[:8]}" + response = _create_team(metadata=_meta(impl, cost_center="CC-1001"), team_id=team_id) + assert response.status_code == 200, response.text + yield impl, team_id + _delete_team(team_id) + + +@pytest.mark.parametrize("impl", IMPLS) +def test_create_with_valid_cost_center_succeeds(impl): + response = _create_team(metadata=_meta(impl, cost_center="CC-1001")) + assert response.status_code == 200, response.text + team_id = response.json()["team_id"] + try: + assert response.json()["metadata"]["cost_center"] == "CC-1001" + finally: + _delete_team(team_id) + + +@pytest.mark.parametrize("impl", IMPLS) +def test_create_without_cost_center_is_rejected(impl): + team_id = f"meta-validate-reject-{impl}-{uuid.uuid4().hex[:8]}" + response = _create_team(metadata=_meta(impl), team_id=team_id) + assert response.status_code == 400, response.text + assert REQUIRED_MESSAGES[impl] in response.text + info = _team_info(team_id) + assert info.status_code == 404, "rejected create must not leave a team row behind" + + +@pytest.mark.parametrize("impl", IMPLS) +def test_create_with_unknown_cost_center(impl): + response = _create_team(metadata=_meta(impl, cost_center="CC-9999")) + if impl == "immutable": + assert response.status_code == 200, response.text + _delete_team(response.json()["team_id"]) + return + assert response.status_code == 400, response.text + assert UNKNOWN_MESSAGES[impl] in response.text + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_patch_changing_cost_center(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _patch_team(team_id, {"metadata": {"cost_center": "CC-1002"}}) + if impl == "immutable": + assert response.status_code == 400, response.text + assert "immutable once set" in response.text + info = _team_info(team_id).json()["team_info"]["metadata"] + assert info["cost_center"] == "CC-1001", "blocked update must leave stored metadata intact" + return + assert response.status_code == 200, response.text + assert response.json()["metadata"]["cost_center"] == "CC-1002" + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_patch_unrelated_key_validates_merged_result(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _patch_team(team_id, {"metadata": {"team_notes": "hello"}}) + assert response.status_code == 200, response.text + merged = response.json()["metadata"] + assert merged["cost_center"] == "CC-1001" + assert merged["team_notes"] == "hello" + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_patch_null_deleting_cost_center_is_rejected(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _patch_team(team_id, {"metadata": {"cost_center": None}}) + assert response.status_code == 400, response.text + assert REQUIRED_MESSAGES[impl] in response.text + info = _team_info(team_id).json()["team_info"]["metadata"] + assert info["cost_center"] == "CC-1001" + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_post_update_dropping_cost_center_is_rejected(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _post_update(team_id, {"metadata": _meta(impl, team_notes="only-notes")}) + assert response.status_code == 400, response.text + assert REQUIRED_MESSAGES[impl] in response.text + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_update_without_metadata_skips_validation(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _post_update(team_id, {"tpm_limit": 55}) + assert response.status_code == 200, response.text + + +def test_http_service_outage_fails_closed_with_configured_message(): + response = _create_team(metadata=_meta("http_down", cost_center="CC-1001")) + assert response.status_code == 503, response.text + assert UNAVAILABLE_MESSAGE in response.text + + +def test_metadata_without_dispatch_key_is_untouched(): + response = _create_team(metadata={"any_key": "any_value"}) + assert response.status_code == 200, response.text + team_id = response.json()["team_id"] + try: + update = _post_update(team_id, {"metadata": {"any_key": "changed"}}) + assert update.status_code == 200, update.text + finally: + _delete_team(team_id) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 0de2c1ac71d..fc018a5333e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -9863,11 +9863,14 @@ async def _drive_team_write( raw_body=None, user=None, find_returns_none=False, + mock_sink=None, ): """Drive POST ``update_team`` or PATCH ``patch_team`` against a mocked team. Returns ``(endpoint_result, update_mock)``; propagates whatever the endpoint raises. Inspect ``update_mock.call_args.kwargs["data"]`` for the DB write. + Pass a dict as ``mock_sink`` to receive the update mock even when the + endpoint raises. """ from unittest.mock import AsyncMock, MagicMock, Mock from unittest.mock import patch as _patch @@ -9913,6 +9916,8 @@ async def _drive_team_write( return_value=LiteLLM_TeamTable(team_id=_PATCH_TEAM_ID, team_alias="t") ) pc.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + if mock_sink is not None: + mock_sink["update"] = pc.db.litellm_teamtable.update req = Mock(spec=Request) if kind == "post": @@ -10175,6 +10180,249 @@ async def test_patch_returns_full_team_object_not_wrapper(): assert result.team_id == _PATCH_TEAM_ID +# --------------------------------------------------------------------------- +# custom_team_metadata_validate wiring: the configured validator must gate +# every team write path (POST /team/new, POST /team/update, PATCH /team/{id}) +# and must see the metadata that will actually be written. +# --------------------------------------------------------------------------- + +from contextlib import contextmanager + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_VALIDATOR_REGISTRY, + TeamMetadataValidationResult, +) + + +@contextmanager +def _configured_team_metadata_validator(validator): + TEAM_METADATA_VALIDATOR_REGISTRY.set(validator) + try: + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + yield + finally: + TEAM_METADATA_VALIDATOR_REGISTRY.set(None) + + +def _recording_validator(recorded, valid=True, error_message=None): + async def validator(payload): + recorded.append(payload) + return TeamMetadataValidationResult(valid=valid, error_message=error_message) + + return validator + + +@pytest.mark.asyncio +async def test_update_validator_sees_replacement_on_post_and_merged_on_patch(): + """POST hands the validator the wholesale replacement; PATCH hands it the + RFC 7386 merged result including preserved keys.""" + existing = {"cost_center": "OLD", "keep": 1} + body = {"metadata": {"cost_center": "NEW"}} + + recorded_post = [] + with _configured_team_metadata_validator(_recording_validator(recorded_post)): + await _drive_team_write("post", existing_metadata=existing, payload=body) + + recorded_patch = [] + with _configured_team_metadata_validator(_recording_validator(recorded_patch)): + await _drive_team_write("patch", existing_metadata=existing, payload=body) + + assert len(recorded_post) == 1 + assert recorded_post[0].operation == "update" + assert recorded_post[0].metadata == {"cost_center": "NEW"} + assert recorded_post[0].existing_metadata == existing + + assert len(recorded_patch) == 1 + assert recorded_patch[0].operation == "update" + assert recorded_patch[0].metadata == {"cost_center": "NEW", "keep": 1} + assert recorded_patch[0].existing_metadata == existing + + +@pytest.mark.asyncio +async def test_patch_null_delete_removes_key_from_validated_metadata(): + """Deleting a key via PATCH null must be visible to the validator as the + key's absence in the resulting metadata, so a required key cannot be + silently dropped.""" + recorded = [] + with _configured_team_metadata_validator(_recording_validator(recorded)): + await _drive_team_write( + "patch", + existing_metadata={"cost_center": "OLD", "keep": 1}, + payload={"metadata": {"cost_center": None}}, + ) + + assert len(recorded) == 1 + assert recorded[0].metadata == {"keep": 1} + assert recorded[0].existing_metadata == {"cost_center": "OLD", "keep": 1} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_update_without_metadata_skips_validator(kind): + recorded = [] + with _configured_team_metadata_validator(_recording_validator(recorded)): + await _drive_team_write(kind, existing_metadata={"k": "v"}, payload={"tpm_limit": 5}) + + assert recorded == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_update_validator_rejection_blocks_db_write(kind): + recorded = [] + validator = _recording_validator(recorded, valid=False, error_message="cost center rejected, contact FinOps") + sink = {} + + with _configured_team_metadata_validator(validator): + with pytest.raises(ProxyException) as exc_info: + await _drive_team_write( + kind, + existing_metadata={"cost_center": "OLD"}, + payload={"metadata": {"cost_center": "BAD"}}, + mock_sink=sink, + ) + + assert str(exc_info.value.code) == "400" + assert "cost center rejected, contact FinOps" in str(exc_info.value.message) + assert len(recorded) == 1 + sink["update"].assert_not_awaited() + + +@pytest.mark.asyncio +async def test_new_team_validator_runs_without_metadata_and_rejection_blocks_create(): + """Create always validates, even when the request carries no metadata, so a + required-key policy can reject a team created without one.""" + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + recorded = [] + validator = _recording_validator(recorded, valid=False, error_message="cost_center is required") + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + _configured_team_metadata_validator(validator), + ): + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.db.litellm_teamtable.create = AsyncMock() + mock_license.is_team_count_over_limit.return_value = False + + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=NewTeamRequest(team_alias="no-metadata-team"), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"), + ) + + assert str(exc_info.value.code) == "400" + assert "cost_center is required" in str(exc_info.value.message) + assert len(recorded) == 1 + assert recorded[0].operation == "create" + assert recorded[0].metadata == {} + assert recorded[0].existing_metadata is None + mock_prisma.db.litellm_teamtable.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_new_team_validator_accept_proceeds_to_create(mock_db_client, mock_admin_auth): + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + mock_db_client.jsonify_team_object = lambda db_data: db_data + mock_db_client.get_data = AsyncMock(return_value=None) + mock_db_client.update_data = AsyncMock(return_value=MagicMock()) + mock_db_client.db = MagicMock() + + team_create_result = MagicMock(team_id="team-accept-1") + team_create_result.model_dump.return_value = {"team_id": "team-accept-1"} + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + recorded = [] + with _configured_team_metadata_validator(_recording_validator(recorded)): + await new_team( + data=NewTeamRequest(team_alias="accepted-team", metadata={"cost_center": "CC-1001"}), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + assert len(recorded) == 1 + assert recorded[0].operation == "create" + assert recorded[0].metadata == {"cost_center": "CC-1001"} + mock_db_client.db.litellm_teamtable.create.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_new_team_rejection_precedes_model_alias_write(): + """A rejected create must not leave an orphaned LiteLLM_ModelTable row: + validation runs before the model_aliases insert.""" + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + validator = _recording_validator([], valid=False, error_message="cost_center is required") + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + _configured_team_metadata_validator(validator), + ): + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.db.litellm_teamtable.create = AsyncMock() + mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1")) + mock_license.is_team_count_over_limit.return_value = False + + with pytest.raises(ProxyException): + await new_team( + data=NewTeamRequest( + team_alias="alias-orphan-check", + model_aliases={"alias-model": "gpt-4o"}, + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"), + ) + + mock_prisma.db.litellm_modeltable.create.assert_not_awaited() + mock_prisma.db.litellm_teamtable.create.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_update_existing_metadata_excludes_system_managed_keys(kind): + """The validator's existing_metadata must be symmetric with metadata: + server-owned keys (team_member_budget_id) are stripped from both, so a + key-preservation validator never sees them 'disappear'.""" + recorded = [] + stored = {"cost_center": "CC-1001", "team_member_budget_id": "budget-abc"} + + with _configured_team_metadata_validator(_recording_validator(recorded)): + await _drive_team_write( + kind, + existing_metadata=dict(stored), + payload={"metadata": {"notes": "x"}}, + ) + + assert len(recorded) == 1 + assert recorded[0].existing_metadata == {"cost_center": "CC-1001"} + assert "team_member_budget_id" not in recorded[0].metadata + + # --------------------------------------------------------------------------- # PATCH body is validated through PatchTeamRequest before it is handed to # update_team. The write below must stay byte-identical to what the untyped @@ -10340,6 +10588,76 @@ async def test_list_available_teams_filters_joined_and_validates_rows(monkeypatc assert find_many_kwargs["where"] == {"team_id": {"in": ["team-open"]}} +@pytest.mark.asyncio +async def test_get_team_metadata_schema_returns_configured_fields(): + from litellm.proxy.management_endpoints.team_endpoints import get_team_metadata_schema + from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + parse_team_metadata_schema, + ) + + TEAM_METADATA_SCHEMA_REGISTRY.set( + parse_team_metadata_schema( + [ + {"key": "cost_center", "label": "Cost Center"}, + {"key": "app_name", "label": "Application Name"}, + ] + ) + ) + try: + result = await get_team_metadata_schema() + finally: + TEAM_METADATA_SCHEMA_REGISTRY.set(()) + + assert [field.key for field in result.fields] == ["cost_center", "app_name"] + assert result.fields[0].label == "Cost Center" + assert result.fields[1].label == "Application Name" + + +@pytest.mark.asyncio +async def test_get_team_metadata_schema_empty_when_unconfigured(): + from litellm.proxy.management_endpoints.team_endpoints import get_team_metadata_schema + from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + ) + + TEAM_METADATA_SCHEMA_REGISTRY.set(()) + result = await get_team_metadata_schema() + + assert result.fields == () + + +def test_get_team_metadata_schema_route_requires_auth(): + from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + parse_team_metadata_schema, + ) + + with patch("litellm.proxy.proxy_server.master_key", "sk-1234"): + response = client.get("/team/metadata_schema") + assert response.status_code == 401 + + TEAM_METADATA_SCHEMA_REGISTRY.set(parse_team_metadata_schema([{"key": "cost_center", "label": "Cost Center"}])) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1" + ) + try: + authed = client.get("/team/metadata_schema") + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + TEAM_METADATA_SCHEMA_REGISTRY.set(()) + + assert authed.status_code == 200 + assert authed.json() == {"fields": [{"key": "cost_center", "label": "Cost Center"}]} + + +def test_team_metadata_schema_route_is_readable_by_non_admins(): + from litellm.proxy._types import LiteLLMRoutes + + assert "/team/metadata_schema" in LiteLLMRoutes.info_routes.value + assert "/team/metadata_schema" in LiteLLMRoutes.management_routes.value + + def _provisioning_caller(role: LitellmUserRoles) -> UserAPIKeyAuth: return UserAPIKeyAuth(user_id="caller-1", user_role=role) diff --git a/tests/test_litellm/proxy/management_helpers/team_metadata_validator_impls.py b/tests/test_litellm/proxy/management_helpers/team_metadata_validator_impls.py new file mode 100644 index 00000000000..39ebeb97ec1 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/team_metadata_validator_impls.py @@ -0,0 +1,92 @@ +"""Three independent `custom_team_metadata_validate` implementations. + +Used by the matrix tests in `test_team_metadata_validation.py` and loadable +directly from a proxy config via `get_instance_fn` for live verification: + +- `validate_allowlist`: plain async function; requires `cost_center` and + checks it against a static allowlist. +- `validate_via_http`: async function that POSTs the metadata to an external + validation service (`TEAM_METADATA_VALIDATION_SERVICE_URL`); any transport + error or non-2xx response raises, exercising the fail-closed path. +- `IMMUTABLE_COST_CENTER_VALIDATOR`: class instance with an async + `__call__`; requires `cost_center` and forbids changing it once set, using + `existing_metadata` and `operation`. +""" + +import os + +import httpx + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataValidationPayload, + TeamMetadataValidationResult, +) + +ALLOWED_COST_CENTERS = frozenset({"CC-1001", "CC-1002"}) +DEFAULT_SERVICE_URL = "http://localhost:9414/validate" + + +async def validate_allowlist( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + cost_center = payload.metadata.get("cost_center") + if cost_center is None: + return TeamMetadataValidationResult( + valid=False, + error_message="cost_center is required in team metadata. Contact the FinOps team.", + ) + if cost_center not in ALLOWED_COST_CENTERS: + return TeamMetadataValidationResult( + valid=False, + error_message=f"Cost center {cost_center} is not recognized. Contact the FinOps team.", + ) + return TeamMetadataValidationResult(valid=True) + + +async def validate_via_http( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + service_url = os.environ.get("TEAM_METADATA_VALIDATION_SERVICE_URL", DEFAULT_SERVICE_URL) + async with httpx.AsyncClient(timeout=2.0) as client: + response = await client.post( + service_url, + json={"operation": payload.operation, "metadata": payload.metadata}, + ) + response.raise_for_status() + body = response.json() + if body.get("ok") is True: + return TeamMetadataValidationResult(valid=True) + return TeamMetadataValidationResult( + valid=False, + error_message=body.get("reason", "Rejected by the cost center service."), + ) + + +class ImmutableCostCenterValidator: + def __init__(self, immutable_key: str = "cost_center") -> None: + self.immutable_key = immutable_key + + async def __call__( + self, + payload: TeamMetadataValidationPayload, + ) -> TeamMetadataValidationResult: + current = payload.metadata.get(self.immutable_key) + if current is None: + return TeamMetadataValidationResult( + valid=False, + error_message=f"{self.immutable_key} is required in team metadata. Contact the FinOps team.", + ) + if payload.operation == "update" and payload.existing_metadata is not None: + prior = payload.existing_metadata.get(self.immutable_key) + if prior is not None and prior != current: + return TeamMetadataValidationResult( + valid=False, + error_message=( + f"{self.immutable_key} is immutable once set " + f"(stored: {prior}, requested: {current}). Contact the FinOps team." + ), + ) + return TeamMetadataValidationResult(valid=True) + + +IMMUTABLE_COST_CENTER_VALIDATOR = ImmutableCostCenterValidator() diff --git a/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py new file mode 100644 index 00000000000..d22c3db0f7e --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py @@ -0,0 +1,704 @@ +import asyncio +import os +import sys +from unittest.mock import patch + +import pytest +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.management_helpers.team_metadata_validation import ( + DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE, + DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS, + DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE, + TeamMetadataRequester, + TeamMetadataValidationPayload, + TeamMetadataValidationResult, + TeamMetadataValidatorRegistry, + _read_timeout_seconds, + _read_unavailable_message, + run_team_metadata_validation, + validate_team_metadata_if_configured, +) + + +def _registry_with(validator): + registry = TeamMetadataValidatorRegistry() + registry.set(validator) + return registry + +UNAVAILABLE_MESSAGE = "validation system down, contact ops" + + +def _payload(**overrides): + values = { + "operation": "create", + "metadata": {"cost_center": "CC-1001"}, + "existing_metadata": None, + "team_id": "team-1", + "team_alias": "alias-1", + "requester": TeamMetadataRequester(user_id="u1"), + } + values.update(overrides) + return TeamMetadataValidationPayload(**values) + + +async def _run(validator, payload=None, premium_user=True, timeout_seconds=1.0): + await run_team_metadata_validation( + validator=validator, + payload=payload or _payload(), + premium_user=premium_user, + timeout_seconds=timeout_seconds, + unavailable_message=UNAVAILABLE_MESSAGE, + ) + + +@pytest.mark.asyncio +async def test_valid_result_passes(): + async def validator(payload): + return TeamMetadataValidationResult(valid=True) + + await _run(validator) + + +@pytest.mark.asyncio +async def test_rejection_raises_400_with_validator_message(): + async def validator(payload): + return TeamMetadataValidationResult(valid=False, error_message="cost center rejected, contact FinOps") + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "cost center rejected, contact FinOps"} + + +@pytest.mark.asyncio +async def test_rejection_without_message_uses_default(): + async def validator(payload): + return TeamMetadataValidationResult(valid=False) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE} + + +@pytest.mark.asyncio +async def test_dict_shaped_return_is_accepted(): + async def validator(payload): + return {"valid": False, "error_message": "rejected via dict"} + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "rejected via dict"} + + +@pytest.mark.asyncio +async def test_validator_exception_fails_closed_with_generic_message(): + async def validator(payload): + raise RuntimeError("internal validation service is down") + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": UNAVAILABLE_MESSAGE} + + +@pytest.mark.asyncio +async def test_validator_timeout_fails_closed(): + async def validator(payload): + await asyncio.sleep(1.0) + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator, timeout_seconds=0.01) + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": UNAVAILABLE_MESSAGE} + + +@pytest.mark.asyncio +async def test_malformed_return_shape_fails_closed(): + async def validator(payload): + return "not-a-validation-result" + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": UNAVAILABLE_MESSAGE} + + +@pytest.mark.asyncio +async def test_non_premium_user_is_rejected(): + async def validator(payload): + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator, premium_user=False) + assert exc_info.value.status_code == 400 + assert CommonProxyErrors.not_premium_user.value in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_non_coroutine_validator_is_rejected(): + def validator(payload): + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == {"error": "custom_team_metadata_validate must be an async function"} + + +@pytest.mark.parametrize( + "general_settings, expected", + [ + ({}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": 2}, 2.0), + ({"team_metadata_validation_timeout": 0.5}, 0.5), + ({"team_metadata_validation_timeout": True}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": -1}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": 0}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": "3"}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ], +) +def test_read_timeout_seconds(general_settings, expected): + assert _read_timeout_seconds(general_settings) == expected + + +@pytest.mark.parametrize( + "general_settings, expected", + [ + ({}, DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE), + ( + {"team_metadata_validation_error_message": "call the help desk"}, + "call the help desk", + ), + ({"team_metadata_validation_error_message": " "}, DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE), + ({"team_metadata_validation_error_message": None}, DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE), + ], +) +def test_read_unavailable_message(general_settings, expected): + assert _read_unavailable_message(general_settings) == expected + + +@pytest.mark.asyncio +async def test_adapter_is_noop_when_unconfigured(): + calls = [] + + async def validator(payload): + calls.append(payload) + return TeamMetadataValidationResult(valid=True) + + await validate_team_metadata_if_configured( + operation="create", + metadata={"cost_center": "CC-1001"}, + existing_metadata=None, + team_id="team-1", + team_alias="alias-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u1"), + registry=TeamMetadataValidatorRegistry(), + ) + assert calls == [] + + +@pytest.mark.asyncio +async def test_adapter_builds_payload_and_reads_settings(): + recorded = [] + + async def validator(payload): + recorded.append(payload) + return TeamMetadataValidationResult(valid=True) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.proxy_server.general_settings", + {"team_metadata_validation_timeout": 3, "team_metadata_validation_error_message": "ops msg"}, + ), + ): + await validate_team_metadata_if_configured( + operation="update", + metadata={"cost_center": "CC-2001"}, + existing_metadata={"cost_center": "CC-1001", "keep": 1}, + team_id="team-9", + team_alias="alias-9", + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user-9", + user_email="user-9@example.com", + ), + registry=_registry_with(validator), + ) + + assert len(recorded) == 1 + payload = recorded[0] + assert payload.operation == "update" + assert payload.metadata == {"cost_center": "CC-2001"} + assert payload.existing_metadata == {"cost_center": "CC-1001", "keep": 1} + assert payload.team_id == "team-9" + assert payload.team_alias == "alias-9" + assert payload.requester.user_id == "user-9" + assert payload.requester.user_email == "user-9@example.com" + assert payload.requester.user_role == LitellmUserRoles.INTERNAL_USER.value + + +@pytest.mark.asyncio +async def test_adapter_normalizes_non_dict_metadata_to_empty_dict(): + recorded = [] + + async def validator(payload): + recorded.append(payload) + return TeamMetadataValidationResult(valid=True) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + await validate_team_metadata_if_configured( + operation="create", + metadata=None, + existing_metadata=None, + team_id="team-1", + team_alias=None, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u1"), + registry=_registry_with(validator), + ) + + assert len(recorded) == 1 + assert recorded[0].metadata == {} + assert recorded[0].existing_metadata is None + + +# --------------------------------------------------------------------------- +# Validator implementation matrix: three independent implementations +# (allowlist function, HTTP-service-backed function, immutability class +# instance) driven through the real team write endpoints. +# --------------------------------------------------------------------------- + +import json as _json +import socket +import threading +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest.mock import AsyncMock, MagicMock, Mock + +import team_metadata_validator_impls as impls + +from litellm.proxy._types import ProxyException +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_VALIDATOR_REGISTRY, +) + + +class _CostCenterServiceHandler(BaseHTTPRequestHandler): + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + body = _json.loads(self.rfile.read(length) or b"{}") + cost_center = (body.get("metadata") or {}).get("cost_center") + if cost_center is None: + resp = {"ok": False, "reason": "cost_center missing per cost center service"} + elif cost_center not in impls.ALLOWED_COST_CENTERS: + resp = {"ok": False, "reason": f"cost center {cost_center} rejected by cost center service"} + else: + resp = {"ok": True} + payload = _json.dumps(resp).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format, *args): + pass + + +@pytest.fixture(scope="module") +def cost_center_service_url(): + server = ThreadingHTTPServer(("127.0.0.1", 0), _CostCenterServiceHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}/validate" + finally: + server.shutdown() + thread.join(timeout=5) + + +def _closed_port_url(): + probe = socket.socket() + probe.bind(("127.0.0.1", 0)) + port = probe.getsockname()[1] + probe.close() + return f"http://127.0.0.1:{port}/validate" + + +@contextmanager +def _configured(validator): + TEAM_METADATA_VALIDATOR_REGISTRY.set(validator) + try: + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + yield + finally: + TEAM_METADATA_VALIDATOR_REGISTRY.set(None) + + +async def _drive_create(metadata, mock_sink=None): + from fastapi import Request + + from litellm.proxy._types import LiteLLM_TeamTable, NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as pc, + patch("litellm.proxy.proxy_server._license_check") as lic, + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + ): + team_row = MagicMock(team_id="matrix-team-1") + team_row.model_dump.return_value = {"team_id": "matrix-team-1"} + pc.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + pc.get_data = AsyncMock(return_value=None) + pc.update_data = AsyncMock(return_value=MagicMock()) + pc.db.litellm_teamtable.create = AsyncMock(return_value=team_row) + pc.db.litellm_teamtable.count = AsyncMock(return_value=0) + pc.db.litellm_teamtable.update = AsyncMock(return_value=team_row) + pc.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + pc.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1")) + lic.is_team_count_over_limit.return_value = False + if mock_sink is not None: + mock_sink["team_create"] = pc.db.litellm_teamtable.create + mock_sink["model_create"] = pc.db.litellm_modeltable.create + + request_kwargs = {"team_alias": "matrix-team"} + if metadata is not None: + request_kwargs["metadata"] = metadata + return await new_team( + data=NewTeamRequest(**request_kwargs), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + +async def _drive_update(kind, existing_metadata, payload): + from fastapi import Request + + from litellm.proxy._types import LiteLLM_TeamTable, PatchTeamRequest, UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import patch_team, update_team + + team_id = "matrix-team-upd" + existing = LiteLLM_TeamTable( + team_id=team_id, + team_alias="matrix", + metadata=existing_metadata, + organization_id=None, + ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as pc, + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.team_endpoints._refresh_cached_team", + new=AsyncMock(), + ), + ): + pc.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing) + pc.db.litellm_teamtable.update = AsyncMock( + return_value=LiteLLM_TeamTable(team_id=team_id, team_alias="matrix") + ) + pc.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + + req = Mock(spec=Request) + if kind == "post": + return await update_team( + data=UpdateTeamRequest(team_id=team_id, **payload), + http_request=req, + user_api_key_dict=auth, + litellm_changed_by=None, + ) + return await patch_team( + team_id=team_id, + data=PatchTeamRequest.model_validate(dict(payload)), + http_request=req, + user_api_key_dict=auth, + litellm_changed_by=None, + ) + + +OK = ("ok", None) + +_MATRIX_IMPLS = { + "allowlist": lambda: impls.validate_allowlist, + "http": lambda: impls.validate_via_http, + "immutable_class": lambda: impls.IMMUTABLE_COST_CENTER_VALIDATOR, +} + +# (scenario, kind, existing_metadata, request payload, {impl: expected}) +# expected is ("ok", None) or ("reject", ) +_MATRIX_SCENARIOS = [ + ( + "create-valid-cost-center", + "create", + None, + {"cost_center": "CC-1001"}, + {"allowlist": OK, "http": OK, "immutable_class": OK}, + ), + ( + "create-missing-cost-center", + "create", + None, + None, + { + "allowlist": ("reject", "cost_center is required"), + "http": ("reject", "cost_center missing per cost center service"), + "immutable_class": ("reject", "cost_center is required"), + }, + ), + ( + "create-unknown-cost-center", + "create", + None, + {"cost_center": "CC-9999"}, + { + "allowlist": ("reject", "is not recognized"), + "http": ("reject", "rejected by cost center service"), + "immutable_class": OK, + }, + ), + ( + "patch-change-cost-center", + "patch", + {"cost_center": "CC-1001"}, + {"metadata": {"cost_center": "CC-1002"}}, + { + "allowlist": OK, + "http": OK, + "immutable_class": ("reject", "immutable once set"), + }, + ), + ( + "patch-unrelated-key-preserves-cost-center", + "patch", + {"cost_center": "CC-1001"}, + {"metadata": {"notes": "hello"}}, + {"allowlist": OK, "http": OK, "immutable_class": OK}, + ), + ( + "patch-null-deletes-cost-center", + "patch", + {"cost_center": "CC-1001"}, + {"metadata": {"cost_center": None}}, + { + "allowlist": ("reject", "cost_center is required"), + "http": ("reject", "cost_center missing per cost center service"), + "immutable_class": ("reject", "cost_center is required"), + }, + ), + ( + "post-replace-drops-cost-center", + "post", + {"cost_center": "CC-1001"}, + {"metadata": {"notes": "only-notes"}}, + { + "allowlist": ("reject", "cost_center is required"), + "http": ("reject", "cost_center missing per cost center service"), + "immutable_class": ("reject", "cost_center is required"), + }, + ), + ( + "update-without-metadata-skips-validation", + "post", + {"cost_center": "CC-1001"}, + {"tpm_limit": 5}, + {"allowlist": OK, "http": OK, "immutable_class": OK}, + ), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("impl_name", sorted(_MATRIX_IMPLS)) +@pytest.mark.parametrize( + "scenario, kind, existing_metadata, request_payload, expectations", + _MATRIX_SCENARIOS, + ids=[row[0] for row in _MATRIX_SCENARIOS], +) +async def test_validator_implementation_matrix( + monkeypatch, + cost_center_service_url, + impl_name, + scenario, + kind, + existing_metadata, + request_payload, + expectations, +): + monkeypatch.setenv("TEAM_METADATA_VALIDATION_SERVICE_URL", cost_center_service_url) + validator = _MATRIX_IMPLS[impl_name]() + outcome, message_part = expectations[impl_name] + + async def drive(): + if kind == "create": + return await _drive_create(metadata=request_payload) + return await _drive_update(kind, existing_metadata, request_payload) + + with _configured(validator): + if outcome == "ok": + await drive() + else: + with pytest.raises(ProxyException) as exc_info: + await drive() + assert str(exc_info.value.code) == "400", f"{scenario} x {impl_name}" + assert message_part in str(exc_info.value.message), f"{scenario} x {impl_name}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "kind, existing_metadata, request_payload", + [ + ("create", None, {"cost_center": "CC-1001"}), + ("patch", {"cost_center": "CC-1001"}, {"metadata": {"notes": "x"}}), + ], +) +async def test_http_validator_service_outage_fails_closed(monkeypatch, kind, existing_metadata, request_payload): + monkeypatch.setenv("TEAM_METADATA_VALIDATION_SERVICE_URL", _closed_port_url()) + + with _configured(impls.validate_via_http): + with pytest.raises(ProxyException) as exc_info: + if kind == "create": + await _drive_create(metadata=request_payload) + else: + await _drive_update(kind, existing_metadata, request_payload) + + assert str(exc_info.value.code) == "503" + assert DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_class_instance_with_async_call_is_accepted(): + await _run(impls.ImmutableCostCenterValidator(), payload=_payload(metadata={"cost_center": "CC-1"})) + + +@pytest.mark.asyncio +async def test_class_instance_with_sync_call_is_rejected(): + class SyncValidator: + def __call__(self, payload): + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(SyncValidator()) + assert exc_info.value.status_code == 500 + + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataSchemaRegistry, + parse_team_metadata_schema, +) + + +def test_parse_schema_none_returns_empty(): + assert parse_team_metadata_schema(None) == () + + +def test_parse_schema_round_trips_fields_in_order(): + raw = [ + {"key": "cost_center", "label": "Cost Center"}, + {"key": "app_name"}, + ] + + fields = parse_team_metadata_schema(raw) + + assert [field.key for field in fields] == ["cost_center", "app_name"] + assert fields[0].label == "Cost Center" + assert fields[1].label is None + + +@pytest.mark.parametrize( + "raw", + [ + "cost_center", + {"key": "cost_center"}, + [{"label": "missing key"}], + [{"key": ""}], + [{"key": "cost_center", "required": True}], + [{"key": "cost_center", "description": "Cost center code"}], + [{"key": "cost_center", "allowed_values": ["CC-1001"]}], + ], +) +def test_parse_schema_malformed_raises(raw): + with pytest.raises(Exception): + parse_team_metadata_schema(raw) + + +def test_parse_schema_duplicate_keys_raise(): + with pytest.raises(ValueError, match="duplicate"): + parse_team_metadata_schema([{"key": "cost_center"}, {"key": "app_name"}, {"key": "cost_center"}]) + + +def test_schema_registry_defaults_empty_and_round_trips(): + registry = TeamMetadataSchemaRegistry() + assert registry.get() == () + + fields = parse_team_metadata_schema([{"key": "cost_center"}]) + registry.set(fields) + assert registry.get() == fields + + registry.set(()) + assert registry.get() == () + + +@pytest.mark.asyncio +async def test_non_callable_validator_is_rejected_with_clean_500(): + class NotCallable: + pass + + with pytest.raises(HTTPException) as exc_info: + await _run(NotCallable()) + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == {"error": "custom_team_metadata_validate must be an async function"} + + +def test_parse_schema_duplicate_error_lists_offending_keys(): + with pytest.raises(ValueError) as exc_info: + parse_team_metadata_schema( + [{"key": "cost_center"}, {"key": "app_name"}, {"key": "cost_center"}, {"key": "app_name"}] + ) + assert str(exc_info.value) == "team_metadata_schema contains duplicate keys: app_name, cost_center" + + +@pytest.mark.asyncio +async def test_adapter_applies_configured_timeout_to_slow_validator(): + async def slow_validator(payload): + await asyncio.sleep(0.2) + return TeamMetadataValidationResult(valid=True) + + registry = TeamMetadataValidatorRegistry() + registry.set(slow_validator) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.proxy_server.general_settings", + {"team_metadata_validation_timeout": 0.01}, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await validate_team_metadata_if_configured( + operation="create", + metadata={"cost_center": "CC-1001"}, + existing_metadata=None, + team_id="team-1", + team_alias="alias-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u1"), + registry=registry, + ) + + assert exc_info.value.status_code == 503 diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 2602b8fa7ad..50b652ff57e 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2862,6 +2862,16 @@ "count": 2 } }, + "src/components/common_components/MetadataKeyValueFields.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/common_components/MetadataKeyValueFields.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/common_components/ModelAliasManager.tsx": { "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMetadataSchema.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMetadataSchema.test.ts new file mode 100644 index 00000000000..0d8dc9fdee4 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMetadataSchema.test.ts @@ -0,0 +1,39 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { fetchTeamMetadataSchema } from "./useTeamMetadataSchema"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => ""), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), +})); + +describe("fetchTeamMetadataSchema", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("should return the declared fields from the response", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue({ + ok: true, + text: async () => JSON.stringify({ fields: [{ key: "cost_center", label: "Cost Center", required: true }] }), + }), + ); + + await expect(fetchTeamMetadataSchema("sk-test")).resolves.toEqual([ + { key: "cost_center", label: "Cost Center", required: true }, + ]); + }); + + it("should return an empty list when the response has no fields array", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: true, text: async () => "{}" })); + + await expect(fetchTeamMetadataSchema("sk-test")).resolves.toEqual([]); + }); + + it("should throw on a non-ok response so the query can retry and fail open", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: false, status: 404, text: async () => "" })); + + await expect(fetchTeamMetadataSchema("sk-test")).rejects.toThrow("404"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMetadataSchema.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMetadataSchema.ts new file mode 100644 index 00000000000..8dc0157d4ba --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMetadataSchema.ts @@ -0,0 +1,37 @@ +import { useQuery } from "@tanstack/react-query"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { createQueryKeys } from "@/app/(dashboard)/hooks/common/queryKeysFactory"; +import { getGlobalLitellmHeaderName, getProxyBaseUrl } from "@/components/networking"; +import { createApiClient } from "@/lib/http/client"; + +export interface TeamMetadataField { + key: string; + label?: string | null; +} + +const TWENTY_FOUR_HOURS_MS = 24 * 60 * 60 * 1000; + +const silentApiClient = createApiClient({ + getBaseUrl: getProxyBaseUrl, + getAuthHeaderName: getGlobalLitellmHeaderName, +}); + +export const fetchTeamMetadataSchema = async (accessToken: string): Promise => { + const data = await silentApiClient.get<{ fields?: TeamMetadataField[] }>("/team/metadata_schema", { accessToken }); + return Array.isArray(data?.fields) ? data.fields : []; +}; + +export const teamMetadataSchemaKeys = createQueryKeys("teamMetadataSchema"); + +export const useTeamMetadataSchema = () => { + const { accessToken } = useAuthorized(); + + return useQuery({ + queryKey: teamMetadataSchemaKeys.list({}), + queryFn: async () => await fetchTeamMetadataSchema(accessToken!), + enabled: Boolean(accessToken), + staleTime: TWENTY_FOUR_HOURS_MS, + gcTime: TWENTY_FOUR_HOURS_MS, + retry: 1, + }); +}; diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 742a88d864c..2afc4bbc586 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -2,6 +2,8 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; import React from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; +import NotificationsManager from "./molecules/notifications_manager"; import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; import { fetchMCPAccessGroups, getGuardrailsList, teamCreateCall } from "./networking"; import Teams from "./Teams"; @@ -33,6 +35,10 @@ vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ teamsTableKeys: { all: ["teamsTable"] }, })); +vi.mock("@/app/(dashboard)/hooks/teams/useTeamMetadataSchema", () => ({ + useTeamMetadataSchema: vi.fn(() => ({ data: [], isLoading: false })), +})); + vi.mock("./molecules/notifications_manager", () => ({ default: { info: vi.fn(), @@ -617,6 +623,210 @@ describe("Teams - access_group_ids in team create", () => { }); }); +describe("Teams - metadata key-value pairs in team create", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(teamCreateCall).mockResolvedValue({ + team_id: "new-team-1", + team_alias: "Test Team", + models: ["gpt-4"], + organization_id: null, + keys: [], + members_with_roles: [], + spend: 0, + }); + mockUseOrganizations.mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }], + }); + }); + + const openCreateModal = async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + }; + + it("renders the metadata editor in the main form without opening Additional Settings", async () => { + await openCreateModal(); + + expect(screen.getByRole("button", { name: /add key-value pair/i })).toBeInTheDocument(); + }); + + it("submits metadata built from key-value pairs as a typed JSON object", async () => { + await openCreateModal(); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Test Team" } }); + fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); + + fireEvent.click(screen.getByRole("button", { name: /add key-value pair/i })); + await waitFor(() => { + expect(screen.getByPlaceholderText("Key")).toBeInTheDocument(); + }); + fireEvent.change(screen.getByPlaceholderText("Key"), { target: { value: "cost_center" } }); + fireEvent.change(screen.getByPlaceholderText("Value"), { target: { value: "eng-42" } }); + + fireEvent.click(screen.getByRole("button", { name: /add key-value pair/i })); + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key")).toHaveLength(2); + }); + fireEvent.change(screen.getAllByPlaceholderText("Key")[1], { target: { value: "tier" } }); + fireEvent.change(screen.getAllByPlaceholderText("Value")[1], { target: { value: "3" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalled(); + }); + + const submittedValues = vi.mocked(teamCreateCall).mock.calls[0][1]; + expect(JSON.parse(submittedValues.metadata)).toEqual({ cost_center: "eng-42", tier: 3 }); + }); + + it("omits metadata entirely when no pairs are added", async () => { + await openCreateModal(); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Test Team" } }); + fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalled(); + }); + + expect(vi.mocked(teamCreateCall).mock.calls[0][1].metadata).toBeUndefined(); + }); +}); + +describe("Teams - schema-declared metadata fields in team create", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(teamCreateCall).mockResolvedValue({ + team_id: "new-team-1", + team_alias: "Test Team", + models: ["gpt-4"], + organization_id: null, + keys: [], + members_with_roles: [], + spend: 0, + }); + mockUseOrganizations.mockReturnValue({ data: null }); + vi.mocked(useTeamMetadataSchema).mockReturnValue({ + data: [{ key: "cost_center", label: "Cost Center" }], + isLoading: false, + } as any); + }); + + const openCreateModal = async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + }; + + it("should prepopulate the declared key as an ordinary pair row and submit its value", async () => { + await openCreateModal(); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Test Team" } }); + fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); + + await waitFor(() => { + expect((screen.getByPlaceholderText("Key") as HTMLInputElement).value).toBe("cost_center"); + }); + fireEvent.change(screen.getByPlaceholderText("Value"), { target: { value: "CC-1001" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalled(); + }); + + const submittedValues = vi.mocked(teamCreateCall).mock.calls[0][1]; + expect(JSON.parse(submittedValues.metadata)).toEqual({ cost_center: "CC-1001" }); + }); + + it("should toast only the validator's own message when the backend rejects the create", async () => { + vi.mocked(teamCreateCall).mockRejectedValue( + new Error("{'error': 'Cost center CC-9999 is not recognized. Contact the FinOps team.'}"), + ); + await openCreateModal(); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Test Team" } }); + fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); + await waitFor(() => { + expect((screen.getByPlaceholderText("Key") as HTMLInputElement).value).toBe("cost_center"); + }); + fireEvent.change(screen.getByPlaceholderText("Value"), { target: { value: "CC-9999" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith( + "Error creating the team: Cost center CC-9999 is not recognized. Contact the FinOps team.", + ); + }); + }); + + it("should show a skeleton in the metadata section while the schema is loading", async () => { + vi.mocked(useTeamMetadataSchema).mockReturnValue({ data: undefined, isLoading: true } as any); + await openCreateModal(); + + expect(screen.getByTestId("metadata-schema-skeleton")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /add key-value pair/i })).not.toBeInTheDocument(); + }); + + it("should re-seed declared keys when the create modal is closed and reopened", async () => { + await openCreateModal(); + + await waitFor(() => { + expect((screen.getByPlaceholderText("Key") as HTMLInputElement).value).toBe("cost_center"); + }); + fireEvent.click(screen.getByLabelText("Remove key-value pair")); + await waitFor(() => { + expect(screen.queryByPlaceholderText("Key")).not.toBeInTheDocument(); + }); + + fireEvent.click(screen.getByRole("button", { name: /^close$/i })); + await waitFor(() => { + expect(screen.queryByLabelText(/team name/i)).not.toBeInTheDocument(); + }); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect((screen.getByPlaceholderText("Key") as HTMLInputElement).value).toBe("cost_center"); + }); + }); +}); + describe("Teams - models dropdown options", () => { beforeEach(() => { vi.clearAllMocks(); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 582f75578ed..fc89751db6a 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -15,6 +15,8 @@ import { teamsTableKeys } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useTeamDetailRouting } from "@/app/(dashboard)/teams/detailNavigation"; import { TeamsTable } from "./TeamsPage/TeamsTable"; import AccessGroupSelector from "./common_components/AccessGroupSelector"; +import MetadataKeyValueFields, { metadataPairsToObject } from "./common_components/MetadataKeyValueFields"; +import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import PassThroughRoutesSelector from "./common_components/PassThroughRoutesSelector"; import AgentSelector from "./agent_management/AgentSelector"; import ModelAliasManager from "./common_components/ModelAliasManager"; @@ -28,6 +30,7 @@ import type { Team } from "./key_team_helpers/key_list"; import MCPServerSelector from "./mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "./mcp_server_management/MCPToolPermissions"; import NotificationsManager from "./molecules/notifications_manager"; +import { extractProxyErrorMessage } from "@/lib/http/client"; import { Organization, fetchMCPAccessGroups, getGuardrailsList, getPoliciesList, teamDeleteCall } from "./networking"; import NumericalInput from "./shared/numerical_input"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; @@ -125,6 +128,7 @@ const getOrganizationAlias = ( const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser = false }) => { const { data: organizationsData } = useOrganizations(); const organizations = organizationsData ?? null; + const { data: teamMetadataSchemaFields = [], isLoading: isTeamMetadataSchemaLoading } = useTeamMetadataSchema(); const queryClient = useQueryClient(); const refreshTeams = () => queryClient.invalidateQueries({ queryKey: teamsTableKeys.all }); const [currentOrg] = useState(null); @@ -322,25 +326,11 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser NotificationsManager.info("Creating Team"); - // Handle logging settings in metadata - if (loggingSettings.length > 0) { - let metadata = {}; - if (formValues.metadata) { - try { - metadata = JSON.parse(formValues.metadata); - } catch (e) { - console.warn("Invalid JSON in metadata field, starting with empty object"); - } - } - - // Add logging settings to metadata - metadata = { - ...metadata, - logging: loggingSettings.filter((config) => config.callback_name), // Only include configs with callback_name - }; - - formValues.metadata = JSON.stringify(metadata); - } + const metadataObject = { + ...metadataPairsToObject(formValues.metadata), + ...(loggingSettings.length > 0 ? { logging: loggingSettings.filter((config) => config.callback_name) } : {}), + }; + formValues.metadata = Object.keys(metadataObject).length > 0 ? JSON.stringify(metadataObject) : undefined; if (formValues.secret_manager_settings) { if (typeof formValues.secret_manager_settings === "string") { @@ -451,7 +441,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser } } catch (error) { console.error("Error creating the team:", error); - NotificationsManager.fromBackend("Error creating the team: " + error); + NotificationsManager.fromBackend("Error creating the team: " + extractProxyErrorMessage(error)); } }; @@ -593,6 +583,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser footer={null} onOk={handleOk} onCancel={handleCancel} + destroyOnHidden >
<> @@ -747,6 +738,16 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser + + + = ({ accessToken, userID, userRole, premiumUser > - - - { + it("returns an empty list for null or undefined metadata", () => { + expect(metadataObjectToPairs(null)).toEqual([]); + expect(metadataObjectToPairs(undefined)).toEqual([]); + }); + + it("keeps plain string values as-is", () => { + expect(metadataObjectToPairs({ department: "research" })).toEqual([{ key: "department", value: "research" }]); + }); + + it("serializes non-string values as JSON", () => { + expect( + metadataObjectToPairs({ + tier: 3, + beta: true, + config: { region: "us" }, + tags: ["a", "b"], + empty: null, + }), + ).toEqual([ + { key: "tier", value: "3" }, + { key: "beta", value: "true" }, + { key: "config", value: '{"region":"us"}' }, + { key: "tags", value: '["a","b"]' }, + { key: "empty", value: "null" }, + ]); + }); + + it("quotes string values that would otherwise parse as JSON, so types round-trip", () => { + expect(metadataObjectToPairs({ code: "42", flag: "true" })).toEqual([ + { key: "code", value: '"42"' }, + { key: "flag", value: '"true"' }, + ]); + }); + + it("filters out excluded keys", () => { + expect( + metadataObjectToPairs({ department: "research", logging: [{ callback_name: "langfuse" }] }, new Set(["logging"])), + ).toEqual([{ key: "department", value: "research" }]); + }); +}); + +describe("metadataPairsToObject", () => { + it("returns an empty object for undefined pairs", () => { + expect(metadataPairsToObject(undefined)).toEqual({}); + }); + + it("keeps plain text values as strings", () => { + expect(metadataPairsToObject([{ key: "department", value: "research" }])).toEqual({ department: "research" }); + }); + + it("parses JSON values into their typed form", () => { + expect( + metadataPairsToObject([ + { key: "tier", value: "3" }, + { key: "beta", value: "true" }, + { key: "config", value: '{"region":"us"}' }, + { key: "code", value: '"42"' }, + ]), + ).toEqual({ tier: 3, beta: true, config: { region: "us" }, code: "42" }); + }); + + it("skips rows without a key and defaults a missing value to an empty string", () => { + expect(metadataPairsToObject([{ key: "", value: "orphan" }, undefined, { key: "kept" }])).toEqual({ kept: "" }); + }); + + it("round-trips a mixed-type metadata object losslessly", () => { + const metadata = { + department: "research", + code: "42", + tier: 3, + beta: true, + config: { region: "us", replicas: 2 }, + }; + expect(metadataPairsToObject(metadataObjectToPairs(metadata))).toEqual(metadata); + }); +}); + +interface HarnessProps { + onFinish: (values: { metadata?: MetadataPair[] }) => void; + initialMetadata?: MetadataPair[]; + schemaFields?: TeamMetadataField[]; + schemaLoading?: boolean; +} + +const Harness: React.FC = ({ onFinish, initialMetadata, schemaFields, schemaLoading }) => { + const [form] = Form.useForm(); + return ( + + + + + ); +}; + +describe("MetadataKeyValueFields", () => { + it("renders one row per existing pair", () => { + render( + , + ); + + const keyInputs = screen.getAllByPlaceholderText("Key"); + const valueInputs = screen.getAllByPlaceholderText("Value"); + expect(keyInputs.map((input) => (input as HTMLInputElement).value)).toEqual(["department", "tier"]); + expect(valueInputs.map((input) => (input as HTMLInputElement).value)).toEqual(["research", "3"]); + }); + + it("adds a row and submits the entered pair", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + render(); + + await user.click(screen.getByRole("button", { name: /add key-value pair/i })); + await user.type(screen.getByPlaceholderText("Key"), "cost_center"); + await user.type(screen.getByPlaceholderText("Value"), "eng-1"); + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => { + expect(onFinish).toHaveBeenCalledWith({ metadata: [{ key: "cost_center", value: "eng-1" }] }); + }); + }); + + it("removes a row when its remove icon is clicked", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + render( + , + ); + + await user.click(screen.getAllByLabelText("Remove key-value pair")[0]); + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => { + expect(onFinish).toHaveBeenCalledWith({ metadata: [{ key: "tier", value: "3" }] }); + }); + }); + + it("blocks submission on duplicate keys", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + render( + , + ); + + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => { + expect(screen.getAllByText("Duplicate key").length).toBeGreaterThan(0); + }); + expect(onFinish).not.toHaveBeenCalled(); + }); + + it("blocks submission when a row is missing its key", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + render(); + + await user.click(screen.getByRole("button", { name: /add key-value pair/i })); + await user.type(screen.getByPlaceholderText("Value"), "orphan"); + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => { + expect(screen.getByText("Missing key")).toBeInTheDocument(); + }); + expect(onFinish).not.toHaveBeenCalled(); + }); +}); + +describe("MetadataKeyValueFields with a declared schema", () => { + const schema: TeamMetadataField[] = [ + { key: "cost_center", label: "Cost Center" }, + { key: "app_name", label: "Application Name" }, + ]; + + it("should prepopulate one ordinary editable pair row per declared key", async () => { + render(); + + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value)).toEqual([ + "cost_center", + "app_name", + ]); + }); + screen.getAllByPlaceholderText("Key").forEach((input) => expect(input).toBeEnabled()); + expect(screen.getAllByLabelText("Remove key-value pair")).toHaveLength(2); + }); + + it("should submit a prepopulated key with its typed value", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + render(); + + await user.type(await screen.findByPlaceholderText("Value"), "CC-1001"); + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => { + expect(onFinish).toHaveBeenCalledWith({ metadata: [{ key: "cost_center", value: "CC-1001" }] }); + }); + }); + + it("should not add a second row for keys already present in the form", async () => { + render( + , + ); + + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value)).toEqual([ + "cost_center", + "app_name", + ]); + }); + expect(screen.getAllByPlaceholderText("Value").map((input) => (input as HTMLInputElement).value)).toEqual([ + "CC-1001", + "", + ]); + }); + + it("should let the user remove a prepopulated row", async () => { + const user = userEvent.setup(); + render(); + + await screen.findAllByPlaceholderText("Key"); + await user.click(screen.getAllByLabelText("Remove key-value pair")[0]); + + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value)).toEqual([ + "app_name", + ]); + }); + }); + + it("should show a skeleton instead of the editor while the schema is loading", () => { + render(); + + expect(screen.getByTestId("metadata-schema-skeleton")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /add key-value pair/i })).not.toBeInTheDocument(); + }); + + it("should seed rows when the schema arrives after an initial loading state", async () => { + const onFinish = vi.fn(); + const { rerender } = render(); + + rerender(); + + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value)).toEqual([ + "cost_center", + "app_name", + ]); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.tsx b/ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.tsx new file mode 100644 index 00000000000..da085f95ad8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.tsx @@ -0,0 +1,135 @@ +import { MinusCircleOutlined, PlusOutlined } from "@ant-design/icons"; +import { Button, Form, FormInstance, Input, Skeleton, Space } from "antd"; +import React, { useEffect, useRef } from "react"; + +import { TeamMetadataField } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; + +export interface MetadataPair { + key: string; + value: string; +} + +function formatMetadataValue(value: unknown): string { + if (typeof value !== "string") { + return JSON.stringify(value) ?? ""; + } + try { + JSON.parse(value); + return JSON.stringify(value); + } catch { + return value; + } +} + +function parseMetadataValue(raw: string): unknown { + try { + return JSON.parse(raw); + } catch { + return raw; + } +} + +export function metadataObjectToPairs( + metadata: Record | null | undefined, + excludedKeys: ReadonlySet = new Set(), +): MetadataPair[] { + return Object.entries(metadata ?? {}) + .filter(([key]) => !excludedKeys.has(key)) + .map(([key, value]) => ({ key, value: formatMetadataValue(value) })); +} + +export function metadataPairsToObject( + pairs: readonly (Partial | undefined)[] | undefined, +): Record { + return Object.fromEntries( + (pairs ?? []) + .filter((pair): pair is Partial & { key: string } => Boolean(pair?.key)) + .map((pair) => [pair.key, parseMetadataValue(pair.value ?? "")]), + ); +} + +interface MetadataKeyValueFieldsProps { + form: FormInstance; + name?: string; + schemaFields?: readonly TeamMetadataField[]; + schemaLoading?: boolean; +} + +const MetadataKeyValueFields: React.FC = ({ + form, + name = "metadata", + schemaFields = [], + schemaLoading = false, +}) => { + const seededRef = useRef(false); + + useEffect(() => { + if (seededRef.current || schemaLoading || schemaFields.length === 0) return; + seededRef.current = true; + const pairs: (Partial | undefined)[] = form.getFieldValue(name) ?? []; + if (!Array.isArray(pairs)) return; + const existingKeys = new Set(pairs.map((pair) => pair?.key).filter(Boolean)); + const seeded = schemaFields + .filter((field) => !existingKeys.has(field.key)) + .map((field) => ({ key: field.key, value: "" })); + if (seeded.length > 0) { + form.setFieldValue(name, [...pairs, ...seeded]); + } + }, [form, name, schemaFields, schemaLoading]); + + if (schemaLoading) { + return ( +
+ +
+ ); + } + + return ( + + {(fields, { add, remove }) => ( + <> + {fields.map(({ key, name: fieldName, ...restField }) => ( + + { + if (!value) return Promise.resolve(); + const all: (Partial | undefined)[] = form.getFieldValue(name) ?? []; + const dupes = all.filter((entry) => entry?.key === value); + if (dupes.length > 1) { + return Promise.reject(new Error("Duplicate key")); + } + return Promise.resolve(); + }, + }, + ]} + > + + + + + + remove(fieldName)} + style={{ color: "#ef4444" }} + /> + + ))} + + + + + )} + + ); +}; + +export default MetadataKeyValueFields; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 03cf0e9583c..5a2d33ee4bd 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -38,7 +38,7 @@ import type { CoordinationRedisTestResponse, } from "@/app/(dashboard)/caching/_components/coordination_redis_settings/types"; import { MCP_TOOLS_PREVIEW_FORBIDDEN_MESSAGE } from "./mcp_tools/constants"; -import { createApiClient, deriveErrorMessage } from "@/lib/http/client"; +import { createApiClient, deriveErrorMessage, unwrapProxyErrorMessage } from "@/lib/http/client"; import { resolveApiBase } from "@/lib/http/resolveApiBase"; import { registerAuthHeaderNameGetter, @@ -2643,7 +2643,7 @@ export const teamUpdateCall = async ( const errorData = await response.text(); handleError(errorData); console.error("Error response from the server:", errorData); - NotificationsManager.fromBackend("Failed to update team settings: " + errorData); + NotificationsManager.fromBackend("Failed to update team settings: " + unwrapProxyErrorMessage(errorData)); throw new Error(errorData); } const data = (await response.json()) as { data: Team; team_id: string }; diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 712cff80649..513719a2ad9 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1,3 +1,4 @@ +import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import * as networking from "@/components/networking"; import { screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; @@ -26,6 +27,10 @@ vi.mock("@/components/utils/dataUtils", () => ({ formatNumberWithCommas: vi.fn((value: number) => value.toLocaleString()), })); +vi.mock("@/app/(dashboard)/hooks/teams/useTeamMetadataSchema", () => ({ + useTeamMetadataSchema: vi.fn(() => ({ data: [], isLoading: false })), +})); + vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAllProxyModels: vi.fn(), })); @@ -220,6 +225,7 @@ describe("TeamInfoView", () => { isFetching: false, refetch: vi.fn(), } as any); + vi.mocked(useTeamMetadataSchema).mockReturnValue({ data: [], isLoading: false } as any); vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); vi.mocked(networking.getPoliciesList).mockResolvedValue({ policies: [] }); @@ -893,6 +899,137 @@ describe("TeamInfoView", () => { }); }); + describe("metadata key-value editing", () => { + const openSettingsEditor = async (user: ReturnType) => { + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /edit settings/i })); + + await waitFor(() => { + expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); + }); + }; + + it("prefills pairs from team metadata, hides UI-managed keys, and round-trips typed values on save", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + metadata: { + department: "research", + tier: 3, + beta: true, + config: { region: "us" }, + logging: [{ callback_name: "langfuse", callback_type: "success", callback_vars: {} }], + guardrails: ["g1"], + disable_global_guardrails: false, + model_tpm_limit: { "gpt-4": 100 }, + }, + models: ["gpt-4"], + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await openSettingsEditor(user); + + const keyValues = screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value); + expect(keyValues).toEqual(["department", "tier", "beta", "config"]); + const valueValues = screen.getAllByPlaceholderText("Value").map((input) => (input as HTMLInputElement).value); + expect(valueValues).toEqual(["research", "3", "true", '{"region":"us"}']); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + const updateArg = vi.mocked(networking.teamUpdateCall).mock.calls[0][1]; + expect(updateArg.metadata).toMatchObject({ + department: "research", + tier: 3, + beta: true, + config: { region: "us" }, + logging: [{ callback_name: "langfuse", callback_type: "success", callback_vars: {} }], + }); + expect(updateArg.metadata).not.toHaveProperty("model_tpm_limit"); + expect(updateArg.model_tpm_limit).toEqual({ "gpt-4": 100 }); + }); + + it("includes a newly added pair in the team update", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] })); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await openSettingsEditor(user); + + await user.click(screen.getByRole("button", { name: /add key-value pair/i })); + await user.type(screen.getByPlaceholderText("Key"), "cost_center"); + await user.type(screen.getByPlaceholderText("Value"), "eng-1"); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].metadata).toMatchObject({ cost_center: "eng-1" }); + }); + + it("should keep declared keys as ordinary prefilled rows and submit the edited value", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(useTeamMetadataSchema).mockReturnValue({ + data: [ + { key: "cost_center", label: "Cost Center" }, + { key: "app_name", label: "Application Name" }, + ], + isLoading: false, + } as any); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + metadata: { cost_center: "CC-OLD", department: "research" }, + models: ["gpt-4"], + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await openSettingsEditor(user); + + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value)).toEqual([ + "cost_center", + "department", + "app_name", + ]); + }); + expect(screen.getAllByPlaceholderText("Value")[0]).toHaveValue("CC-OLD"); + + await user.clear(screen.getAllByPlaceholderText("Value")[0]); + await user.type(screen.getAllByPlaceholderText("Value")[0], "CC-NEW"); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].metadata).toMatchObject({ + cost_center: "CC-NEW", + department: "research", + app_name: "", + }); + }); + }); + describe("model aliases", () => { const openSettingsEditor = async (user: ReturnType) => { await waitFor(() => { diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 34570043f52..bbe5dc05a88 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -35,6 +35,11 @@ import { CheckIcon, CopyIcon } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; +import MetadataKeyValueFields, { + metadataObjectToPairs, + metadataPairsToObject, +} from "../common_components/MetadataKeyValueFields"; +import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import ModelAliasManager from "../common_components/ModelAliasManager"; import AgentSelector from "../agent_management/AgentSelector"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; @@ -66,6 +71,18 @@ import { import TeamMembersComponent from "./TeamMemberTab"; import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; +const UI_MANAGED_METADATA_KEYS: ReadonlySet = new Set([ + "logging", + "secret_manager_settings", + "soft_budget_alerting_emails", + "model_tpm_limit", + "model_rpm_limit", + "allowed_passthrough_routes", + "guardrails", + "opted_out_global_guardrails", + "disable_global_guardrails", +]); + export interface TeamMembership { user_id: string; team_id: string; @@ -203,6 +220,7 @@ const TeamInfoView: React.FC = ({ const [organization, setOrganization] = useState(null); const { userRole, userId } = useAuthorized(); const { data: userOrganizations = [] } = useOrganizations(); + const { data: teamMetadataSchemaFields = [], isLoading: isTeamMetadataSchemaLoading } = useTeamMetadataSchema(); const queryClient = useQueryClient(); // Check if user is org admin for this team's organization @@ -461,16 +479,7 @@ const TeamInfoView: React.FC = ({ if (!accessToken) return; setIsTeamSaving(true); - let parsedMetadata = {}; - try { - const rawMetadata = values.metadata ? JSON.parse(values.metadata) : {}; - // Exclude soft_budget_alerting_emails from parsed metadata since it's handled separately - const { soft_budget_alerting_emails, ...rest } = rawMetadata; - parsedMetadata = rest; - } catch (e) { - NotificationsManager.fromBackend("Invalid JSON in metadata field"); - return; - } + const parsedMetadata = metadataPairsToObject(values.metadata); let secretManagerSettings: Record | undefined; if (typeof values.secret_manager_settings === "string") { @@ -980,21 +989,7 @@ const TeamInfoView: React.FC = ({ soft_budget_alerting_emails: Array.isArray(info.metadata?.soft_budget_alerting_emails) ? info.metadata.soft_budget_alerting_emails.join(", ") : "", - metadata: info.metadata - ? JSON.stringify( - (({ - logging, - secret_manager_settings, - soft_budget_alerting_emails, - model_tpm_limit, - model_rpm_limit, - allowed_passthrough_routes, - ...rest - }) => rest)(info.metadata), - null, - 2, - ) - : "", + metadata: metadataObjectToPairs(info.metadata, UI_MANAGED_METADATA_KEYS), logging_settings: info.metadata?.logging || [], secret_manager_settings: info.metadata?.secret_manager_settings ? JSON.stringify(info.metadata.secret_manager_settings, null, 2) @@ -1170,6 +1165,17 @@ const TeamInfoView: React.FC = ({ + + + + = ({ /> - - - -