diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 6046f1bb581..abf010e0d65 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -255,26 +255,44 @@ class CustomGuardrail(CustomLogger): f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}" ) + @staticmethod + def _get_admin_metadata(data: dict) -> dict: + """Return merged admin-configured key and team metadata from the request data. + + The proxy may inject admin metadata (user_api_key_metadata, + user_api_key_team_metadata) into either ``metadata`` or + ``litellm_metadata`` depending on endpoint. Check both so a caller + cannot shadow admin config by pre-populating the other key. + Key-level settings override team-level. + """ + team_meta: dict = {} + key_meta: dict = {} + for key in ("metadata", "litellm_metadata"): + # Defensive: an unparsed JSON-string metadata could leak past the + # proxy's normal parse path; don't AttributeError on .get(). + meta = data.get(key) + if not isinstance(meta, dict): + continue + team_meta = meta.get("user_api_key_team_metadata") or team_meta + key_meta = meta.get("user_api_key_metadata") or key_meta + return {**team_meta, **key_meta} + def get_disable_global_guardrail(self, data: dict) -> Optional[bool]: """ - Returns True if the global guardrail should be disabled + Returns True if the global guardrail should be disabled. + + Reads from admin-configured key/team metadata only, not from + the request body, to prevent callers from disabling guardrails. """ - if "disable_global_guardrails" in data: - return data["disable_global_guardrails"] - metadata = data.get("litellm_metadata") or data.get("metadata", {}) - if "disable_global_guardrails" in metadata: - return metadata["disable_global_guardrails"] - return False + return self._get_admin_metadata(data).get("disable_global_guardrails", False) def get_opted_out_global_guardrails_from_metadata(self, data: dict) -> List[str]: """ Returns the list of global guardrail names the team/key has opted out of. + + Reads from admin-configured key/team metadata only. """ - if "opted_out_global_guardrails" in data: - value = data["opted_out_global_guardrails"] - return value if isinstance(value, list) else [] - metadata = data.get("litellm_metadata") or data.get("metadata", {}) - value = metadata.get("opted_out_global_guardrails") + value = self._get_admin_metadata(data).get("opted_out_global_guardrails") return value if isinstance(value, list) else [] def _is_valid_response_type(self, result: Any) -> bool: @@ -417,7 +435,9 @@ class CustomGuardrail(CustomLogger): """ requested_guardrails = self.get_guardrail_from_metadata(data) disable_global_guardrail = self.get_disable_global_guardrail(data) - opted_out_global_guardrails = self.get_opted_out_global_guardrails_from_metadata(data) + opted_out_global_guardrails = ( + self.get_opted_out_global_guardrails_from_metadata(data) + ) verbose_logger.debug( "inside should_run_guardrail for guardrail=%s event_type= %s guardrail_supported_event_hooks= %s requested_guardrails= %s self.default_on= %s", self.guardrail_name, @@ -426,7 +446,10 @@ class CustomGuardrail(CustomLogger): requested_guardrails, self.default_on, ) - if self.default_on is True and self.guardrail_name in opted_out_global_guardrails: + if ( + self.default_on is True + and self.guardrail_name in opted_out_global_guardrails + ): return False if self.default_on is True and disable_global_guardrail is not True: diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index 92c97a59924..563609af1e4 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -48,7 +48,6 @@ _supported_callback_params = [ "braintrust_host", "slack_webhook_url", "lunary_public_key", - "turn_off_message_logging", ] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 56958a88f6d..1621f04213f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -31,6 +31,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( RBAC_ROLES, CallInfo, @@ -327,15 +328,63 @@ def _global_proxy_budget_check( ) +_GUARDRAIL_MODIFICATION_KEYS: tuple = ( + "guardrails", + "disable_global_guardrails", + "disable_global_guardrail", + "opted_out_global_guardrails", +) + + def _guardrail_modification_check( request_body: dict, team_object: Optional[LiteLLM_TeamTable] ) -> None: - _request_metadata: dict = request_body.get("metadata", {}) or {} - if not _request_metadata.get("guardrails"): - return + """ + Reject user-supplied metadata flags that would modify guardrail behavior + unless the team has explicit permission. Checked keys include the plural + ``guardrails`` list plus the per-request toggles that influence whether + default-on guardrails run (``disable_global_guardrails``, + ``disable_global_guardrail`` singular, and ``opted_out_global_guardrails``). + User-supplied values for the bypass toggles are also silently ignored by + ``_get_admin_metadata`` at read time; this check adds defense in depth by + failing loudly at the auth layer so operators see an explicit 403 instead + of a confusing silent-ignore. + """ from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails + def _coerce_to_dict(container: Any) -> Optional[dict]: + """Accept dict or JSON-string (from multipart/form-data or extra_body). + + Without this, an attacker can smuggle guardrail keys past the check by + sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` — + ``isinstance(dict)`` on the string returns False, the check returns + no-modification, and ``add_litellm_data_to_request`` parses the string + to a dict downstream. + """ + if isinstance(container, dict): + return container + if isinstance(container, str): + parsed = safe_json_loads(container) + return parsed if isinstance(parsed, dict) else None + return None + + def _user_requested_modification(container: Any) -> bool: + coerced = _coerce_to_dict(container) + if coerced is None: + return False + return any(coerced.get(key) for key in _GUARDRAIL_MODIFICATION_KEYS) + + # Check both metadata keys — callers can populate either depending on the + # endpoint. Cover the top-level too so root-level injection is rejected. + modifies = ( + _user_requested_modification(request_body.get("metadata")) + or _user_requested_modification(request_body.get("litellm_metadata")) + or _user_requested_modification(request_body) + ) + if not modifies: + return + if not can_modify_guardrails(team_object): raise HTTPException( status_code=403, @@ -451,9 +500,9 @@ async def common_checks( # noqa: PLR0915 model=_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=valid_token.team_model_aliases - if valid_token - else None, + team_model_aliases=( + valid_token.team_model_aliases if valid_token else None + ), ): raise ProxyException( message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 64766bbaadd..12d9bff91e0 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -823,19 +823,30 @@ def get_end_user_id_from_request_body( user_from_body_user_field = request_body["user"] return str(user_from_body_user_field) + def _as_dict(value: Any) -> dict: + # metadata / litellm_metadata can arrive as JSON strings from + # multipart/form-data or extra_body; coerce so string-encoded + # payloads can't evade end-user attribution. + if isinstance(value, dict): + return value + if isinstance(value, str): + from litellm.litellm_core_utils.safe_json_loads import safe_json_loads + + parsed = safe_json_loads(value) + return parsed if isinstance(parsed, dict) else {} + return {} + # Check 4: 'litellm_metadata.user' in request_body (commonly Anthropic) - litellm_metadata = request_body.get("litellm_metadata") - if isinstance(litellm_metadata, dict): - user_from_litellm_metadata = litellm_metadata.get("user") - if user_from_litellm_metadata is not None: - return str(user_from_litellm_metadata) + litellm_metadata = _as_dict(request_body.get("litellm_metadata")) + user_from_litellm_metadata = litellm_metadata.get("user") + if user_from_litellm_metadata is not None: + return str(user_from_litellm_metadata) # Check 5: 'metadata.user_id' in request_body (another common pattern) - metadata_dict = request_body.get("metadata") - if isinstance(metadata_dict, dict): - user_id_from_metadata_field = metadata_dict.get("user_id") - if user_id_from_metadata_field is not None: - return str(user_id_from_metadata_field) + metadata_dict = _as_dict(request_body.get("metadata")) + user_id_from_metadata_field = metadata_dict.get("user_id") + if user_id_from_metadata_field is not None: + return str(user_id_from_metadata_field) # Check 6: 'safety_identifier' in request body (OpenAI Responses API parameter) # SECURITY NOTE: safety_identifier can be set by any caller in the request body. diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 1dd25262127..71abdfa5e9e 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -197,10 +197,10 @@ def check_file_size_under_limit( if llm_router is not None and request_data["model"] in router_model_names: try: - deployment: Optional[ - Deployment - ] = llm_router.get_deployment_by_model_group_name( - model_group_name=request_data["model"] + deployment: Optional[Deployment] = ( + llm_router.get_deployment_by_model_group_name( + model_group_name=request_data["model"] + ) ) if ( deployment @@ -426,7 +426,16 @@ def get_tags_from_request_body(request_body: dict) -> List[str]: List of tag names (strings), empty list if no valid tags found """ metadata_variable_name = get_metadata_variable_name_from_kwargs(request_body) - metadata = request_body.get(metadata_variable_name) or {} + metadata = request_body.get(metadata_variable_name) + # metadata can arrive as a JSON string from multipart/form-data or extra_body; + # coerce defensively so .get() below never raises AttributeError. + if isinstance(metadata, str): + from litellm.litellm_core_utils.safe_json_loads import safe_json_loads + + parsed = safe_json_loads(metadata) + metadata = parsed if isinstance(parsed, dict) else {} + elif not isinstance(metadata, dict): + metadata = {} tags_in_metadata: Any = metadata.get("tags", []) tags_in_request_body: Any = request_body.get("tags", []) combined_tags: List[str] = [] diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b0adf7aa6ee..7467bbae232 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -981,13 +981,16 @@ async def add_litellm_data_to_request( # noqa: PLR0915 # Init - Proxy Server Request # we do this as soon as entering so we track the original request ########################################################## - # Track arrival time for queue time metric + # Track arrival time for queue time metric. The body snapshot is filled + # in after the admin-injection strip below so the audit / spend-tracking + # consumers of proxy_server_request["body"] see the cleaned metadata + # rather than attacker-forged user_api_key_* fields. arrival_time = time.time() data["proxy_server_request"] = { "url": str(request.url), "method": request.method, "headers": _headers, - "body": copy.copy(data), # use copy instead of deepcopy + "body": None, # filled in post-strip; see below "arrival_time": arrival_time, # Track when request arrived at proxy } @@ -1069,9 +1072,10 @@ async def add_litellm_data_to_request( # noqa: PLR0915 verbose_proxy_logger.warning( f"Failed to parse 'metadata' as JSON dict. Received value: {data['metadata']}" ) - data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy( - data["metadata"] - ) + # requester_metadata is snapshotted AFTER the strip below so + # downstream consumers (e.g. PANW guardrail reading user_ip / + # profile_id) don't see attacker-injected admin slots preserved in + # the deepcopy. # Parse litellm_metadata if it's a string (e.g., from multipart/form-data or extra_body) if "litellm_metadata" in data and data["litellm_metadata"] is not None: @@ -1083,11 +1087,89 @@ async def add_litellm_data_to_request( # noqa: PLR0915 ) else: data["litellm_metadata"] = parsed_litellm_metadata - # Merge litellm_metadata into the metadata variable (preserving existing values) - if isinstance(data["litellm_metadata"], dict): - for key, value in data["litellm_metadata"].items(): - if key not in data[_metadata_variable_name]: - data[_metadata_variable_name][key] = value + + # Strip internal pipeline state and admin-injection slots from user input. + # Runs AFTER the string-to-dict parse above so JSON-string metadata (sent + # via multipart/form-data or extra_body) cannot smuggle admin fields past + # the isinstance(dict) guard. + # + # The proxy populates a family of ``user_api_key_*`` fields below + # (user_api_key_metadata, user_api_key_user_id, user_api_key_alias, + # user_api_key_spend, user_api_key_team_metadata, …) into + # data[_metadata_variable_name]. Because the proxy only writes to ONE of + # the two metadata dicts, a caller pre-populating any of these keys on + # the OTHER metadata dict would have their forged values surface in + # guardrails, spend tracking, audit logs, and identity resolution. Strip + # by prefix so new ``user_api_key_*`` fields added in the future are + # covered without per-key maintenance. + for _meta_key in ("metadata", "litellm_metadata"): + _user_meta = data.get(_meta_key) + if isinstance(_user_meta, dict): + _user_meta.pop("_pipeline_managed_guardrails", None) + for _k in [k for k in _user_meta if k.startswith("user_api_key_")]: + _user_meta.pop(_k, None) + + # Strip caller-supplied routing/budget tags unless the admin has opted + # this key or team in via metadata.allow_client_tags=True. Tags drive + # tag-based routing and tag budget attribution — accepting them from + # untrusted callers lets an attacker reach restricted deployments or + # misattribute spend to a victim team's tag. + _admin_allow_client_tags = False + for _admin_meta in ( + user_api_key_dict.metadata, + user_api_key_dict.team_metadata, + ): + if ( + isinstance(_admin_meta, dict) + and _admin_meta.get("allow_client_tags") is True + ): + _admin_allow_client_tags = True + break + if not _admin_allow_client_tags: + _stripped_from: List[str] = [] + for _meta_key in ("metadata", "litellm_metadata"): + _user_meta = data.get(_meta_key) + if isinstance(_user_meta, dict) and "tags" in _user_meta: + _user_meta.pop("tags", None) + _stripped_from.append(_meta_key) + # Also strip the root-level `tags` field. get_tags_from_request_body + # reads request_body["tags"] directly and feeds it to the policy + # engine, so leaving it in place here would let the strip-in-metadata + # above be trivially bypassed by moving the tags to the body root. + if "tags" in data: + data.pop("tags", None) + _stripped_from.append("tags (root)") + if _stripped_from: + verbose_proxy_logger.warning( + "Stripped caller-supplied tags from %s: this key/team does " + "not have `allow_client_tags: true` in its metadata. Set it " + "to opt into client-supplied routing/budget tags.", + ", ".join(_stripped_from), + ) + + # Fill in the proxy_server_request body snapshot now that metadata has + # been parsed and stripped. Consumers (standard_logging_payload, lago, + # spend_tracking_utils, streaming_iterator) read `body` to audit the + # request; taking the snapshot here ensures they see cleaned metadata. + data["proxy_server_request"]["body"] = copy.copy(data) + + # Snapshot the (now-cleaned) requester-supplied metadata for downstream + # consumers. Taking the deepcopy AFTER the strip prevents attacker- + # injected admin slots (user_api_key_*, tags without opt-in, + # _pipeline_managed_guardrails) from surviving in requester_metadata + # where guardrails and audit paths may read from it. + if "metadata" in data and isinstance(data["metadata"], dict): + data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy( + data["metadata"] + ) + + # Now merge litellm_metadata into the metadata variable (preserving existing + # values) — runs AFTER the strip so attacker injections in litellm_metadata + # cannot cross-contaminate the admin-authoritative metadata dict. + if "litellm_metadata" in data and isinstance(data["litellm_metadata"], dict): + for key, value in data["litellm_metadata"].items(): + if key not in data[_metadata_variable_name]: + data[_metadata_variable_name][key] = value data = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( data=data, @@ -1250,15 +1332,24 @@ async def add_litellm_data_to_request( # noqa: PLR0915 user_agent = request.headers["user-agent"] data[_metadata_variable_name]["user_agent"] = user_agent - # Check if using tag based routing + # Check if using tag based routing. The helper reads caller-controlled + # sources (x-litellm-tags header, data["tags"] root-level), so its result + # is still gated by the same allow_client_tags flag that gated the + # body-metadata tag strip above. Otherwise the strip is trivially + # bypassed by sending tags via header or at the root of the body. tags = LiteLLMProxyRequestSetup.add_request_tag_to_metadata( llm_router=llm_router, headers=_headers, data=data, ) - if tags is not None: + if tags is not None and _admin_allow_client_tags: data[_metadata_variable_name]["tags"] = tags + elif tags is not None: + verbose_proxy_logger.warning( + "Ignored caller-supplied tags from header/root body: this " + "key/team does not have `allow_client_tags: true` in its metadata." + ) # Team Callbacks controls callback_settings_obj = _get_dynamic_logging_metadata( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2d789b982da..1d61bebee44 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2201,9 +2201,11 @@ def run_ollama_serve(): with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - verbose_proxy_logger.debug(f""" + verbose_proxy_logger.debug( + f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` - """) + """ + ) def _get_process_rss_mb() -> Optional[float]: @@ -7157,7 +7159,10 @@ async def chat_completion( # noqa: PLR0915 global user_temperature, user_request_timeout, user_max_tokens, user_api_base data = await _read_request_body(request=request) if user_api_key_dict is not None: - if data.get("metadata") is None: + if not isinstance(data.get("metadata"), dict): + # Covers both missing and JSON-string metadata (multipart / + # extra_body); otherwise `data["metadata"][k] = v` below raises + # TypeError on a string value and 500s the request. data["metadata"] = {} if ( hasattr(user_api_key_dict, "user_id") @@ -11372,7 +11377,9 @@ async def async_queue_request( # if users are using user_api_key_auth, set `user` in `data` data["user"] = user_api_key_dict.user_id - if "metadata" not in data: + if not isinstance(data.get("metadata"), dict): + # Covers both missing and JSON-string metadata (multipart / + # extra_body); see above for the same guard upstream. data["metadata"] = {} data["metadata"]["user_api_key"] = user_api_key_dict.api_key data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 64dc5fe4741..be27b852478 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -29,6 +29,9 @@ from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.router_strategy.tag_based_routing import _get_tags_from_request_kwargs +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, +) from litellm.router_utils.cooldown_callbacks import ( _get_prometheus_logger_from_callbacks, ) @@ -100,9 +103,9 @@ class RouterBudgetLimiting(CustomLogger): self.dual_cache = dual_cache self.redis_increment_operation_queue: List[RedisPipelineIncrementOperation] = [] asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) - self.provider_budget_config: Optional[ - GenericBudgetConfigType - ] = provider_budget_config + self.provider_budget_config: Optional[GenericBudgetConfigType] = ( + provider_budget_config + ) self.deployment_budget_config: Optional[GenericBudgetConfigType] = None self.tag_budget_config: Optional[GenericBudgetConfigType] = None self._init_provider_budgets() @@ -175,7 +178,10 @@ class RouterBudgetLimiting(CustomLogger): spend_map=spend_map, potential_deployments=potential_deployments, request_tags=_get_tags_from_request_kwargs( - request_kwargs=request_kwargs + request_kwargs=request_kwargs, + metadata_variable_name=get_metadata_variable_name_from_kwargs( + request_kwargs or {} + ), ), ) @@ -304,6 +310,16 @@ class RouterBudgetLimiting(CustomLogger): deployment_configs: Dict[str, GenericBudgetInfo] = {} deployment_providers: List[Optional[str]] = [] + # Resolve tags once before the loop (loop-invariant) + _request_tags: List[str] = [] + if self.tag_budget_config: + _request_tags = _get_tags_from_request_kwargs( + request_kwargs=request_kwargs, + metadata_variable_name=get_metadata_variable_name_from_kwargs( + request_kwargs or {} + ), + ) + for deployment in healthy_deployments: # Check provider budgets if self.provider_budget_config: @@ -330,17 +346,14 @@ class RouterBudgetLimiting(CustomLogger): cache_keys.append( f"deployment_spend:{model_id}:{budget_config.budget_duration}" ) - # Check tag budgets - if self.tag_budget_config: - request_tags = _get_tags_from_request_kwargs( - request_kwargs=request_kwargs + + # Check tag budgets (outside loop — tags are per-request, not per-deployment) + for _tag in _request_tags: + _tag_budget_config = self._get_budget_config_for_tag(_tag) + if _tag_budget_config: + cache_keys.append( + f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}" ) - for _tag in request_tags: - _tag_budget_config = self._get_budget_config_for_tag(_tag) - if _tag_budget_config: - cache_keys.append( - f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}" - ) return ( cache_keys, provider_configs, @@ -459,7 +472,10 @@ class RouterBudgetLimiting(CustomLogger): response_cost=response_cost, ) - request_tags = _get_tags_from_request_kwargs(kwargs) + request_tags = _get_tags_from_request_kwargs( + kwargs, + metadata_variable_name=get_metadata_variable_name_from_kwargs(kwargs or {}), + ) if len(request_tags) > 0: for _tag in request_tags: _tag_budget_config = self._get_budget_config_for_tag(_tag) diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index 1188ce9d592..0163f3bbd4f 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -102,13 +102,11 @@ def _match_deployment( return {"matched_via": "tags", "matched_value": matched_value} # 2. Regex match against request headers. - # When match_any=False and the deployment has both plain tags and tag_regex, - # the strict tag check has already failed (step 1 returned None). Allow - # the regex to fire only when the deployment has NO plain tags, so we never - # use regex as a backdoor around the operator's strict-tag policy. - strict_tag_check_failed = ( - not match_any and bool(deployment_tags) and bool(request_tags) - ) + # When match_any=False and the deployment has plain tags, the strict tag + # check either didn't run (no request tags) or failed (step 1 returned + # None). Block the regex path so it cannot circumvent the operator's + # strict-tag policy. + strict_tag_check_failed = not match_any and bool(deployment_tags) if deployment_tag_regex and header_strings and not strict_tag_check_failed: regex_match = _is_valid_deployment_tag_regex( deployment_tag_regex, header_strings diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 9f5f14457e8..916e01b0f2c 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -162,9 +162,13 @@ async def test_add_key_or_team_level_spend_logs_metadata_to_request( print(f"team_sl_metadata: {team_sl_metadata}") mock_request.url.path = "/chat/completions" + # Opt the key into client-supplied tags so request_tags are preserved + # and merged with admin-configured key/team tags. Without this flag, + # request_tags would be stripped by add_litellm_data_to_request. key_metadata = { "tags": key_tags, "spend_logs_metadata": key_sl_metadata, + "allow_client_tags": True, } team_metadata = { "tags": team_tags, @@ -838,12 +842,13 @@ async def test_add_litellm_data_to_request_duplicate_tags( mock_request.headers = {} mock_request.state = State() - # Setup key with tags in metadata + # Setup key with tags in metadata. Opt into client-supplied tags so the + # request_tags are preserved for the merge under test. user_api_key_dict = UserAPIKeyAuth( api_key="test_api_key", user_id="test_user_id", org_id="test_org_id", - metadata={"tags": key_tags}, + metadata={"tags": key_tags, "allow_client_tags": True}, ) # Setup request data with tags diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 3a959d599b5..0c904e9df50 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -173,17 +173,16 @@ class TestCustomGuardrailShouldRunGuardrail: assert result is False def test_should_run_guardrail_with_disable_global_guardrail(self): - """Test that disable_global_guardrail disables a global guardrail when set to True""" + """Test that disable_global_guardrails only works from admin metadata""" from litellm.types.guardrails import GuardrailEventHooks - # Create a guardrail with default_on=True (global guardrail) custom_guardrail = CustomGuardrail( guardrail_name="global_guardrail", default_on=True, event_hook=GuardrailEventHooks.pre_call, ) - # Test 1: Global guardrail runs by default when default_on=True + # Test 1: Global guardrail runs by default data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}], @@ -193,7 +192,7 @@ class TestCustomGuardrailShouldRunGuardrail: ) assert result is True, "Global guardrail should run when default_on=True" - # Test 2: Global guardrail is disabled when disable_global_guardrail=True at root level + # Test 2: User-injected disable at root level is IGNORED data_with_disable_root = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}], @@ -203,23 +202,10 @@ class TestCustomGuardrailShouldRunGuardrail: data=data_with_disable_root, event_type=GuardrailEventHooks.pre_call ) assert ( - result is False - ), "Global guardrail should be disabled when disable_global_guardrail=True" + result is True + ), "User-injected disable_global_guardrails should be ignored" - # Test 3: Global guardrail is disabled when disable_global_guardrail=True in litellm_metadata - data_with_disable_litellm = { - "model": "gpt-3.5-turbo", - "messages": [{"role": "user", "content": "test"}], - "litellm_metadata": {"disable_global_guardrails": True}, - } - result = custom_guardrail.should_run_guardrail( - data=data_with_disable_litellm, event_type=GuardrailEventHooks.pre_call - ) - assert ( - result is False - ), "Global guardrail should be disabled when disable_global_guardrail=True in litellm_metadata" - - # Test 4: Global guardrail is disabled when disable_global_guardrail=True in metadata + # Test 3: User-injected disable in metadata is IGNORED data_with_disable_metadata = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}], @@ -228,25 +214,51 @@ class TestCustomGuardrailShouldRunGuardrail: result = custom_guardrail.should_run_guardrail( data=data_with_disable_metadata, event_type=GuardrailEventHooks.pre_call ) - assert ( - result is False - ), "Global guardrail should be disabled when disable_global_guardrail=True in metadata" + assert result is True, "User-injected metadata disable should be ignored" - # Test 5: Global guardrail runs when disable_global_guardrail=False - data_with_disable_false = { + # Test 4: Admin-configured disable via user_api_key_metadata IS respected + data_with_admin_disable = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}], - "disable_global_guardrails": False, + "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, } result = custom_guardrail.should_run_guardrail( - data=data_with_disable_false, event_type=GuardrailEventHooks.pre_call + data=data_with_admin_disable, event_type=GuardrailEventHooks.pre_call + ) + assert result is False, "Admin-configured disable should be respected" + + # Test 5: Admin config in metadata isn't shadowed by user-supplied litellm_metadata + data_cross_key = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "test"}], + "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, + "litellm_metadata": {"request_tags": ["user-supplied"]}, + } + result = custom_guardrail.should_run_guardrail( + data=data_cross_key, event_type=GuardrailEventHooks.pre_call ) assert ( - result is True - ), "Global guardrail should still run when disable_global_guardrail=False" + result is False + ), "Admin config in metadata must not be shadowed by user-supplied litellm_metadata" + + # Test 6: After the pre-call strip runs, user-injected + # user_api_key_metadata in the non-authoritative metadata key is gone. + # _get_admin_metadata must then surface admin config unchanged. + data_post_strip = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "test"}], + "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, + "litellm_metadata": {}, # post-strip: attacker payload removed + } + result = custom_guardrail.should_run_guardrail( + data=data_post_strip, event_type=GuardrailEventHooks.pre_call + ) + assert ( + result is False + ), "Admin config in metadata must be respected when other metadata key is empty" def test_should_run_guardrail_with_opted_out_global_guardrails(self): - """Test the per-guardrail opt-out list for global (default_on=True) guardrails""" + """Test that per-guardrail opt-out only works from admin metadata""" from litellm.types.guardrails import GuardrailEventHooks custom_guardrail = CustomGuardrail( @@ -255,7 +267,7 @@ class TestCustomGuardrailShouldRunGuardrail: event_hook=GuardrailEventHooks.pre_call, ) - # Test 1: guardrail in the opt-out list at root level → skipped + # Test 1: User-injected opt-out at root level is IGNORED data_root = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}], @@ -265,23 +277,10 @@ class TestCustomGuardrailShouldRunGuardrail: custom_guardrail.should_run_guardrail( data=data_root, event_type=GuardrailEventHooks.pre_call ) - is False + is True ) - # Test 2: guardrail in the opt-out list inside litellm_metadata → skipped - data_litellm = { - "model": "gpt-3.5-turbo", - "messages": [{"role": "user", "content": "test"}], - "litellm_metadata": {"opted_out_global_guardrails": ["global_guardrail"]}, - } - assert ( - custom_guardrail.should_run_guardrail( - data=data_litellm, event_type=GuardrailEventHooks.pre_call - ) - is False - ) - - # Test 3: guardrail in the opt-out list inside metadata → skipped + # Test 2: User-injected opt-out in metadata is IGNORED data_metadata = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}], @@ -291,7 +290,7 @@ class TestCustomGuardrailShouldRunGuardrail: custom_guardrail.should_run_guardrail( data=data_metadata, event_type=GuardrailEventHooks.pre_call ) - is False + is True ) # Test 4: a different guardrail in the opt-out list → still runs @@ -588,7 +587,9 @@ class TestGuardrailSensitiveFieldStripping: duration=1.0, ) - logged_response = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] + logged_response = request_data["metadata"][ + "standard_logging_guardrail_information" + ][0]["guardrail_response"] assert "secret_fields" not in logged_response assert "sk-live-SHOULD-NOT-APPEAR" not in json.dumps(logged_response) @@ -599,7 +600,12 @@ class TestGuardrailSensitiveFieldStripping: guardrail.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=[ - {"result": "ok", "secret_fields": {"raw_headers": {"authorization": "Bearer sk-secret"}}}, + { + "result": "ok", + "secret_fields": { + "raw_headers": {"authorization": "Bearer sk-secret"} + }, + }, {"result": "also_ok"}, ], request_data=request_data, @@ -608,6 +614,7 @@ class TestGuardrailSensitiveFieldStripping: ) import json + serialized = json.dumps(request_data) assert "secret_fields" not in serialized assert "sk-secret" not in serialized @@ -621,21 +628,21 @@ class TestCustomGuardrailPassthroughSupport: """ Test that async_post_call_success_deployment_hook handles raw httpx.Response objects from passthrough endpoints without crashing with TypeError. - + This tests Fix #3: TypeError: TypedDict does not support instance and class checks """ import httpx custom_guardrail = CustomGuardrail() - + # Mock the async_post_call_success_hook to return None (guardrail didn't modify response) custom_guardrail.async_post_call_success_hook = AsyncMock(return_value=None) - + # Create a mock httpx.Response object (typical passthrough response) mock_response = AsyncMock(spec=httpx.Response) mock_response.status_code = 200 mock_response.text = "Mock response" - + request_data = { "guardrails": ["test_guardrail"], "user_api_key_user_id": "test_user", @@ -644,14 +651,14 @@ class TestCustomGuardrailPassthroughSupport: "user_api_key_hash": "test_hash", "user_api_key_request_route": "passthrough_route", } - + # This should not raise TypeError: TypedDict does not support instance and class checks result = await custom_guardrail.async_post_call_success_deployment_hook( request_data=request_data, response=mock_response, call_type=CallTypes.allm_passthrough_route, ) - + # When result is None, should return the original response assert result == mock_response @@ -659,53 +666,53 @@ class TestCustomGuardrailPassthroughSupport: async def test_async_post_call_success_deployment_hook_with_none_call_type(self): """ Test that async_post_call_success_deployment_hook handles None call_type gracefully. - + This ensures that even if call_type is None (before fix #1), the guardrail doesn't crash. """ custom_guardrail = CustomGuardrail() - + # Mock the async_post_call_success_hook to return None custom_guardrail.async_post_call_success_hook = AsyncMock(return_value=None) - + mock_response = AsyncMock() - + request_data = { "guardrails": ["test_guardrail"], "user_api_key_user_id": "test_user", } - + # Call with None call_type - should not crash result = await custom_guardrail.async_post_call_success_deployment_hook( request_data=request_data, response=mock_response, call_type=None, ) - + # Should return the original response when result is None assert result == mock_response def test_is_valid_response_type_with_none(self): """ Test _is_valid_response_type helper method correctly identifies None as invalid. - + This is part of Fix #3: Safely handling TypedDict types that don't support isinstance checks. """ custom_guardrail = CustomGuardrail() - + # None should be invalid assert custom_guardrail._is_valid_response_type(None) is False def test_is_valid_response_type_with_typeddict_error(self): """ Test _is_valid_response_type gracefully handles TypeError from TypedDict. - + This tests Fix #3: When isinstance() is called with TypedDict types, it raises TypeError. The method should catch this and allow the response through. """ from litellm.types.utils import ModelResponse - + custom_guardrail = CustomGuardrail() - + # Create a valid LiteLLM response object response = ModelResponse( id="test-id", @@ -714,13 +721,12 @@ class TestCustomGuardrailPassthroughSupport: model="test-model", object="chat.completion", ) - + # This should return True (it's a valid response type or TypeError is caught) result = custom_guardrail._is_valid_response_type(response) assert result is True - class TestEventTypeLogging: """Tests for event_type logging in guardrail information.""" @@ -1014,7 +1020,9 @@ class TestTracingFieldsPopulation: guardrail_json_response="blocked", request_data=request_data, guardrail_status="guardrail_intervened", - tracing_detail=GuardrailTracingDetail(policy_template="EU AI Act Article 5"), + tracing_detail=GuardrailTracingDetail( + policy_template="EU AI Act Article 5" + ), ) slg_list = request_data["metadata"]["standard_logging_guardrail_information"] diff --git a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py index 91969a2b8e2..55f3c2ba3aa 100644 --- a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py +++ b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py @@ -82,13 +82,18 @@ def test_env_reference_in_litellm_params_metadata_raises(): def test_non_string_values_are_not_flagged(): kwargs = { "langsmith_sampling_rate": 0.5, - "turn_off_message_logging": True, } params = initialize_standard_callback_dynamic_params(kwargs) assert params.get("langsmith_sampling_rate") == 0.5 - assert params.get("turn_off_message_logging") is True + + +def test_turn_off_message_logging_not_extracted_from_request(): + """turn_off_message_logging is admin-only — must not be settable via request.""" + kwargs = {"turn_off_message_logging": True} + params = initialize_standard_callback_dynamic_params(kwargs) + assert params.get("turn_off_message_logging") is None def test_empty_kwargs_returns_empty_params(): diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index bd659ed518f..0391e224084 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -62,9 +62,9 @@ def reset_constants_module(): # Reload modules before test importlib.reload(constants) importlib.reload(auth_checks) - + yield - + # Reload modules after test to clean up importlib.reload(constants) importlib.reload(auth_checks) @@ -157,9 +157,9 @@ def test_experimental_ui_token_ignores_litellm_ui_session_duration( expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) now = get_utc_datetime() # Must be ~10 min, NOT 24h. If LITELLM_UI_SESSION_DURATION were incorrectly used, this would fail. - assert expires <= now + timedelta(minutes=11), ( - "Experimental UI must use 10-min expiry, not LITELLM_UI_SESSION_DURATION" - ) + assert expires <= now + timedelta( + minutes=11 + ), "Experimental UI must use 10-min expiry, not LITELLM_UI_SESSION_DURATION" def test_get_experimental_ui_login_jwt_auth_token_invalid( @@ -293,13 +293,15 @@ def test_get_cli_jwt_auth_token_custom_expiration( # Set custom expiration to 48 hours monkeypatch.setenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", "48") - + # Reload the constants module to pick up the new env var importlib.reload(constants) # Also reload auth_checks to pick up the new constant value importlib.reload(auth_checks) - - token = auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + + token = auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token( + valid_sso_user_defined_values + ) # Decrypt and verify token contents decrypted_token = decrypt_value_helper( @@ -315,7 +317,6 @@ def test_get_cli_jwt_auth_token_custom_expiration( assert expires <= get_utc_datetime() + timedelta(hours=48, minutes=1) - @pytest.mark.asyncio async def test_default_internal_user_params_with_get_user_object(monkeypatch): """Test that default_internal_user_params is used when creating a new user via get_user_object""" @@ -436,7 +437,9 @@ async def test_get_user_object_upsert_includes_user_email(): mock_prisma_client.db.litellm_usertable.create.assert_called_once() creation_args = mock_prisma_client.db.litellm_usertable.create.call_args[1]["data"] - assert "user_email" in creation_args, "user_email should be included when upserting a new user" + assert ( + "user_email" in creation_args + ), "user_email should be included when upserting a new user" assert creation_args["user_email"] == "test@example.com" assert creation_args["user_id"] == "new_test_user" @@ -463,7 +466,9 @@ def test_log_budget_lookup_failure_skips_user_not_found(): @pytest.mark.asyncio -@patch("litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock) +@patch( + "litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock +) async def test_get_team_db_check_calls_new_team_on_upsert(mock_new_team, monkeypatch): """ Test that _get_team_db_check correctly calls the `new_team` function @@ -497,8 +502,12 @@ async def test_get_team_db_check_calls_new_team_on_upsert(mock_new_team, monkeyp @pytest.mark.asyncio -@patch("litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock) -async def test_get_team_db_check_does_not_call_new_team_if_exists(mock_new_team, monkeypatch): +@patch( + "litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock +) +async def test_get_team_db_check_does_not_call_new_team_if_exists( + mock_new_team, monkeypatch +): """ Test that _get_team_db_check does NOT call the `new_team` function if the team already exists in the database. @@ -541,8 +550,9 @@ async def test_vector_store_access_check_early_returns( if vector_store_registry: vector_store_registry.get_vector_store_ids_to_run.return_value = None - with patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch( - "litellm.vector_store_registry", vector_store_registry + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.vector_store_registry", vector_store_registry), ): result = await vector_store_access_check( request_body=request_body, @@ -639,8 +649,9 @@ async def test_vector_store_access_check_with_permissions(): mock_vector_store_registry = MagicMock() mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-1"] - with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( - "litellm.vector_store_registry", mock_vector_store_registry + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.vector_store_registry", mock_vector_store_registry), ): result = await vector_store_access_check( request_body=request_body, @@ -653,8 +664,9 @@ async def test_vector_store_access_check_with_permissions(): # Test with denied access mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-3"] - with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( - "litellm.vector_store_registry", mock_vector_store_registry + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.vector_store_registry", mock_vector_store_registry), ): with pytest.raises(ProxyException) as exc_info: await vector_store_access_check( @@ -687,8 +699,9 @@ async def test_vector_store_access_check_with_team_permissions(): "team-store-allowed" ] - with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( - "litellm.vector_store_registry", mock_vector_store_registry + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.vector_store_registry", mock_vector_store_registry), ): result = await vector_store_access_check( request_body=request_body, @@ -702,8 +715,9 @@ async def test_vector_store_access_check_with_team_permissions(): "team-store-denied" ] - with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( - "litellm.vector_store_registry", mock_vector_store_registry + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.vector_store_registry", mock_vector_store_registry), ): with pytest.raises(ProxyException) as exc_info: await vector_store_access_check( @@ -1598,12 +1612,15 @@ async def test_custom_auth_common_checks_opt_in(): mock_request = MagicMock() # Default (no flag) — common_checks should NOT be called - with patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", - new_callable=AsyncMock, - ) as mock_common, patch( - "litellm.proxy.proxy_server.general_settings", - {}, + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, + patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ), ): mock_common.return_value = True result = await _run_post_custom_auth_checks( @@ -1616,12 +1633,15 @@ async def test_custom_auth_common_checks_opt_in(): mock_common.assert_not_called() # With flag=True — common_checks SHOULD be called - with patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", - new_callable=AsyncMock, - ) as mock_common, patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, + patch( + "litellm.proxy.proxy_server.general_settings", + {"custom_auth_run_common_checks": True}, + ), ): mock_common.return_value = True result = await _run_post_custom_auth_checks( @@ -1660,9 +1680,7 @@ async def test_virtual_key_budget_check_reads_from_spend_counter(): return 1.5 return fallback_spend - with patch( - "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend - ): + with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): with pytest.raises(litellm.BudgetExceededError) as exc_info: await _virtual_key_max_budget_check( valid_token=valid_token, @@ -1692,9 +1710,7 @@ async def test_virtual_key_budget_check_fallback_no_counter(): async def mock_get_current_spend(counter_key, fallback_spend): return fallback_spend - with patch( - "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend - ): + with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): with pytest.raises(litellm.BudgetExceededError) as exc_info: await _virtual_key_max_budget_check( valid_token=valid_token, @@ -1723,9 +1739,7 @@ async def test_team_budget_check_reads_from_spend_counter(): return 1.5 return fallback_spend - with patch( - "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend - ): + with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): with pytest.raises(litellm.BudgetExceededError) as exc_info: await _team_max_budget_check( team_object=team_object, @@ -1763,12 +1777,13 @@ async def test_team_member_budget_check_reads_from_spend_counter(): return 1.5 return fallback_spend - with patch( - "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend - ), patch( - "litellm.proxy.auth.auth_checks.get_team_membership", - new_callable=AsyncMock, - return_value=team_membership, + with ( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=team_membership, + ), ): with pytest.raises(litellm.BudgetExceededError) as exc_info: await _check_team_member_budget( @@ -1780,3 +1795,142 @@ async def test_team_member_budget_check_reads_from_spend_counter(): proxy_logging_obj=proxy_logging_obj, ) assert exc_info.value.current_cost == 1.5 + + +class TestGuardrailModificationCheck: + """Defense-in-depth: `_guardrail_modification_check` must 403 when the + caller's metadata attempts to modify any guardrail-related key and the + team lacks the `modify_guardrails` permission. Checks both the + historically-covered `guardrails` list and the bypass toggles that + `_get_admin_metadata` silently ignores at read time. + """ + + def _call(self, request_body): + from litellm.proxy.auth.auth_checks import _guardrail_modification_check + + team_object = MagicMock() + team_object.metadata = {} # no permission + return _guardrail_modification_check( + request_body=request_body, team_object=team_object + ) + + def test_noop_when_no_guardrail_keys_present(self): + # no-op — should return silently + self._call({"metadata": {"unrelated": "value"}}) + + def test_rejects_guardrails_list(self): + from fastapi import HTTPException + + with patch( + "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", + return_value=False, + ): + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"guardrails": ["custom"]}}) + assert exc.value.status_code == 403 + + def test_rejects_disable_global_guardrails_plural(self): + from fastapi import HTTPException + + with patch( + "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", + return_value=False, + ): + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"disable_global_guardrails": True}}) + assert exc.value.status_code == 403 + + def test_rejects_disable_global_guardrail_singular(self): + """VERIA-28's originally-reported singular-key typo variant.""" + from fastapi import HTTPException + + with patch( + "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", + return_value=False, + ): + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"disable_global_guardrail": True}}) + assert exc.value.status_code == 403 + + def test_rejects_opted_out_global_guardrails(self): + from fastapi import HTTPException + + with patch( + "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", + return_value=False, + ): + with pytest.raises(HTTPException) as exc: + self._call( + {"metadata": {"opted_out_global_guardrails": ["some_guardrail"]}} + ) + assert exc.value.status_code == 403 + + def test_rejects_injection_via_litellm_metadata_key(self): + """Caller can populate the OTHER metadata key; that must also 403.""" + from fastapi import HTTPException + + with patch( + "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", + return_value=False, + ): + with pytest.raises(HTTPException) as exc: + self._call({"litellm_metadata": {"disable_global_guardrails": True}}) + assert exc.value.status_code == 403 + + def test_rejects_root_level_injection(self): + """Top-level injection (`request_body["disable_global_guardrails"]`) + was VERIA-28's easiest variant to hit — keep it rejected.""" + from fastapi import HTTPException + + with patch( + "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", + return_value=False, + ): + with pytest.raises(HTTPException) as exc: + self._call({"disable_global_guardrails": True}) + assert exc.value.status_code == 403 + + def test_allows_when_team_has_permission(self): + with patch( + "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", + return_value=True, + ): + # no-op, should not raise + self._call({"metadata": {"disable_global_guardrails": True}}) + + def test_rejects_string_encoded_metadata_bypass(self): + """Regression: attacker sends metadata as JSON string to bypass the + isinstance(dict) guard. The check must coerce the string to dict + and evaluate guardrail modification keys inside it.""" + import json as _json + + from fastapi import HTTPException + + attacker_payload = {"disable_global_guardrails": True} + with patch( + "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", + return_value=False, + ): + with pytest.raises(HTTPException) as exc: + self._call({"metadata": _json.dumps(attacker_payload)}) + assert exc.value.status_code == 403 + + def test_rejects_string_encoded_litellm_metadata_bypass(self): + """Same bypass via the litellm_metadata key.""" + import json as _json + + from fastapi import HTTPException + + attacker_payload = {"guardrails": ["evaded"]} + with patch( + "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", + return_value=False, + ): + with pytest.raises(HTTPException) as exc: + self._call({"litellm_metadata": _json.dumps(attacker_payload)}) + assert exc.value.status_code == 403 + + def test_noop_when_string_is_not_json_object(self): + """Unparseable strings should not trigger a 403 — they have no keys.""" + self._call({"metadata": "not-json"}) + self._call({"metadata": '"just a string"'}) diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index a1484bc263b..c9f595626ed 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -835,3 +835,41 @@ def test_safe_get_request_headers_state_unavailable(): result = _safe_get_request_headers(mock_request) assert result == {"content-type": "application/json"} + + +class TestGetTagsFromRequestBodyStringCoerce: + """Regression: the auth-time tag helper used `metadata.get("tags", ...)` + directly, which raised AttributeError when metadata arrived as a JSON + string (multipart/form-data or extra_body). That turned into a DoS at + auth time and potentially bypassed tag-based RBAC if the caller caught + the exception and fell through with empty tags. + """ + + def test_json_string_metadata_is_coerced_to_dict(self): + from litellm.proxy.common_utils.http_parsing_utils import ( + get_tags_from_request_body, + ) + + metadata_json = json.dumps({"tags": ["a", "b"]}) + # Must not raise + tags = get_tags_from_request_body({"metadata": metadata_json}) + assert tags == ["a", "b"] + + def test_unparseable_string_metadata_is_ignored(self): + from litellm.proxy.common_utils.http_parsing_utils import ( + get_tags_from_request_body, + ) + + # Must not raise; must yield no metadata tags but keep root tags + tags = get_tags_from_request_body( + {"metadata": "not-json", "tags": ["root-only"]} + ) + assert tags == ["root-only"] + + def test_dict_metadata_still_works(self): + from litellm.proxy.common_utils.http_parsing_utils import ( + get_tags_from_request_body, + ) + + tags = get_tags_from_request_body({"metadata": {"tags": ["x"]}}) + assert tags == ["x"] diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 09d11388d84..49fff6de0ff 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1295,7 +1295,6 @@ def test_create_file_with_nested_litellm_metadata( "target_model_names": "gpt-3.5-turbo", "litellm_metadata[spend_logs_metadata][owner]": "john_doe", "litellm_metadata[spend_logs_metadata][team]": "engineering", - "litellm_metadata[tags]": "production", "litellm_metadata[environment]": "prod", }, headers={"Authorization": "Bearer test-key"}, @@ -1306,11 +1305,12 @@ def test_create_file_with_nested_litellm_metadata( result = response.json() assert result["id"] == "file-test-123" - # Verify nested metadata was correctly parsed + # Verify nested metadata was correctly parsed. + # Note: caller-supplied `tags` is stripped by default; test removed + # to keep the parsing test focused on parser correctness. assert "spend_logs_metadata" in captured_litellm_metadata assert captured_litellm_metadata["spend_logs_metadata"]["owner"] == "john_doe" assert captured_litellm_metadata["spend_logs_metadata"]["team"] == "engineering" - assert captured_litellm_metadata["tags"] == "production" assert captured_litellm_metadata["environment"] == "prod" diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index cf7e71b14d4..360ddd426a4 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -207,6 +207,580 @@ async def test_add_litellm_data_to_request_parses_string_metadata(): assert updated_data["metadata"]["generation_name"] == "gen123" +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_strips_admin_injection_slots(): + """User-supplied user_api_key_metadata / user_api_key_team_metadata / + _pipeline_managed_guardrails must be stripped from both metadata keys + before the proxy writes its own admin-populated values. Otherwise a + caller can shadow admin config via the non-`_metadata_variable_name` + metadata key (e.g. litellm_metadata while the proxy writes to metadata). + """ + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + # Caller tries to inject admin config into BOTH metadata keys + attacker_admin_payload = {"disable_global_guardrails": True} + data = { + "model": "gpt-3.5-turbo", + "metadata": { + "user_api_key_metadata": attacker_admin_payload, + "user_api_key_team_metadata": attacker_admin_payload, + "_pipeline_managed_guardrails": ["evaded"], + }, + "litellm_metadata": { + "user_api_key_metadata": attacker_admin_payload, + "user_api_key_team_metadata": attacker_admin_payload, + "_pipeline_managed_guardrails": ["evaded"], + }, + } + + real_admin_metadata = {"admin_flag": "from_proxy"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata=real_admin_metadata, + team_metadata=real_admin_metadata, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + # The key that matches `_metadata_variable_name` gets proxy-populated + # with the real admin payload; the OTHER key must not retain the + # attacker's injection. + populated = updated["metadata"] + assert populated["user_api_key_metadata"] == real_admin_metadata + assert populated["user_api_key_team_metadata"] == real_admin_metadata + assert "_pipeline_managed_guardrails" not in populated or populated[ + "_pipeline_managed_guardrails" + ] != ["evaded"] + + other = updated.get("litellm_metadata") or {} + assert other.get("user_api_key_metadata") in (None, {}, real_admin_metadata) + assert other.get("user_api_key_team_metadata") in (None, {}, real_admin_metadata) + assert "_pipeline_managed_guardrails" not in other + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_strips_all_user_api_key_prefix_keys(): + """Strip must cover the full user_api_key_* family, not a hand-maintained + list of 2-3 names. Proxy writes a dozen such fields (user_id, alias, + spend, team_id, request_route, …) and an attacker populating any of them + in the non-authoritative metadata key would otherwise forge identity / + spend in audit logs and guardrails.""" + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + attacker_injected = { + "user_api_key_user_id": "victim", + "user_api_key_alias": "admin-key", + "user_api_key_spend": 0.0, + "user_api_key_team_id": "victim-team", + "user_api_key_end_user_id": "victim-user", + "user_api_key_request_route": "/fake/route", + "user_api_key_hash": "fake-hash", + } + data = { + "model": "gpt-3.5-turbo", + "metadata": {**attacker_injected}, + "litellm_metadata": {**attacker_injected}, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + user_id="real-user", + metadata={}, + team_metadata={}, + spend=42.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + # The non-authoritative metadata dict must not retain ANY attacker-injected + # user_api_key_* key. + other = updated.get("litellm_metadata") or {} + attacker_leaks = [k for k in other if k.startswith("user_api_key_")] + assert attacker_leaks == [], f"Unexpected leaked keys: {attacker_leaks}" + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_string_metadata_does_not_crash(): + """Regression: pre-strip code that pre-populated data['metadata'][k]=v + before the string-to-dict parse would crash on JSON-string metadata. + The snapshot / strip / admin-population pipeline must survive metadata + arriving as a string.""" + import json as _json + + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "multipart/form-data"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = { + "model": "gpt-3.5-turbo", + "metadata": _json.dumps({"generation_name": "test"}), + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + # Must not raise TypeError / AttributeError. + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + # The parsed metadata should be a dict and the proxy snapshot body + # should have been taken AFTER the strip (so no leaked user_api_key_* + # from a raw string snapshot). + assert isinstance(updated["metadata"], dict) + assert updated["metadata"].get("generation_name") == "test" + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_proxy_server_request_body_is_post_strip(): + """Regression: proxy_server_request['body'] used to be snapshotted before + the admin-slot strip, so standard_logging_object and spend-tracking + readers saw attacker-injected payload. Snapshot must now be post-strip.""" + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = { + "model": "gpt-3.5-turbo", + "metadata": {"user_api_key_user_id": "victim"}, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + user_id="real-user", + metadata={}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + snapshot_body = updated["proxy_server_request"]["body"] + assert snapshot_body is not None + snapshot_metadata = snapshot_body.get("metadata") or {} + assert "user_api_key_user_id" not in snapshot_metadata or ( + snapshot_metadata["user_api_key_user_id"] != "victim" + ) + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_strips_string_encoded_admin_injection(): + """Regression: metadata arriving as a JSON string (multipart/form-data or + extra_body) must not bypass the admin-injection strip. The parse happens + AFTER receipt, so the strip has to run after the parse, not before. + """ + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "multipart/form-data"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + # Attacker encodes an admin-injection payload inside a JSON string. + attacker_payload = { + "user_api_key_metadata": {"disable_global_guardrails": True}, + "user_api_key_team_metadata": {"disable_global_guardrails": True}, + "_pipeline_managed_guardrails": ["evaded"], + } + data = { + "model": "gpt-3.5-turbo", + "metadata": json.dumps(attacker_payload), + "litellm_metadata": json.dumps(attacker_payload), + } + + real_admin_metadata = {"admin_flag": "from_proxy"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata=real_admin_metadata, + team_metadata=real_admin_metadata, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + populated = updated["metadata"] + # The real admin payload from user_api_key_dict wins. + assert populated["user_api_key_metadata"] == real_admin_metadata + assert populated["user_api_key_team_metadata"] == real_admin_metadata + assert populated.get("_pipeline_managed_guardrails") != ["evaded"] + + other = updated.get("litellm_metadata") or {} + # After the strip, litellm_metadata has no admin-injection slots. + assert "user_api_key_metadata" not in other + assert "user_api_key_team_metadata" not in other + assert "_pipeline_managed_guardrails" not in other + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_ignores_x_litellm_tags_header_without_permission(): + """Regression: the `x-litellm-tags` header bypassed the body-metadata + tag strip. Header tags must also be gated by `allow_client_tags`.""" + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = { + "Content-Type": "application/json", + "x-litellm-tags": "restricted-tier,victim-team", + } + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = {"model": "gpt-3.5-turbo"} + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert "tags" not in (updated.get("metadata") or {}) + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_ignores_root_level_tags_without_permission(): + """Regression: root-level `data["tags"]` bypassed the body-metadata + tag strip. Root-level tags must also be gated by `allow_client_tags`.""" + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = { + "model": "gpt-3.5-turbo", + "tags": ["restricted-tier", "victim-team"], + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert "tags" not in (updated.get("metadata") or {}) + # Also ensure the root-level tags are removed. get_tags_from_request_body + # reads request_body["tags"] directly, so leaving it in place would let + # the policy engine see caller-supplied tags even after the metadata + # strip. + assert "tags" not in updated + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_honors_header_tags_when_opted_in(): + """When allow_client_tags=True, header-supplied tags flow through.""" + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = { + "Content-Type": "application/json", + "x-litellm-tags": "production,ab-test", + } + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = {"model": "gpt-3.5-turbo"} + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"allow_client_tags": True}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["metadata"].get("tags") == ["production", "ab-test"] + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_strips_user_tags_without_permission(): + """Caller-supplied metadata.tags must be stripped when the key/team + metadata does not opt in via allow_client_tags=True. Otherwise an + attacker can reach restricted tag-routed deployments or attribute + spend to a victim team's tag.""" + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = { + "model": "gpt-3.5-turbo", + "metadata": {"tags": ["restricted-tier", "victim-team"]}, + "litellm_metadata": {"tags": ["also-stripped"]}, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert "tags" not in (updated.get("metadata") or {}) + assert "tags" not in (updated.get("litellm_metadata") or {}) + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_preserves_user_tags_when_key_opts_in(): + """When key.metadata.allow_client_tags=True, caller-supplied tags are + preserved and reach the router.""" + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = { + "model": "gpt-3.5-turbo", + "metadata": {"tags": ["opted-in-tag"]}, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"allow_client_tags": True}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["metadata"].get("tags") == ["opted-in-tag"] + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_preserves_user_tags_when_team_opts_in(): + """Team-level allow_client_tags is also honored (not just key-level).""" + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = { + "model": "gpt-3.5-turbo", + "metadata": {"tags": ["team-allowed"]}, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={"allow_client_tags": True}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["metadata"].get("tags") == ["team-allowed"] + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_user_spend_and_budget(): from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request @@ -221,7 +795,10 @@ async def test_add_litellm_data_to_request_user_spend_and_budget(): request_mock.client = MagicMock() request_mock.client.host = "127.0.0.1" - data = {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]} + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + } user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", @@ -275,9 +852,11 @@ async def test_add_litellm_data_to_request_audio_transcription_multipart(): "file": b"Fake audio bytes", } + # Opt the key in to client-supplied tags so the parsed tags from the + # JSON-string multipart body aren't stripped by the admin-injection strip. user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", - metadata={}, + metadata={"allow_client_tags": True}, team_metadata={}, spend=0.0, max_budget=100.0, @@ -1023,6 +1602,7 @@ def test_add_headers_to_llm_call_by_model_group_existing_headers_in_data(): # Restore original model_group_settings litellm.model_group_settings = original_model_group_settings + import json import time from typing import Optional @@ -1040,15 +1620,16 @@ class TestCustomLogger(CustomLogger): def __init__(self): self.standard_logging_object: Optional[StandardLoggingPayload] = None super().__init__() - + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): print(f"SUCCESS CALLBACK CALLED! kwargs keys: {list(kwargs.keys())}") self.standard_logging_object = kwargs.get("standard_logging_object") print(f"Captured standard_logging_object: {self.standard_logging_object}") - + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): print(f"FAILURE CALLBACK CALLED! kwargs keys: {list(kwargs.keys())}") + @pytest.mark.asyncio async def test_add_litellm_metadata_from_request_headers(): """ @@ -1065,8 +1646,16 @@ async def test_add_litellm_metadata_from_request_headers(): try: # Prepare test data (ensure no streaming, add mock_response and api_key to route to litellm.acompletion) - headers = {"x-litellm-spend-logs-metadata": '{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion", "timestamp": "2025-09-02T10:30:00Z"}'} - data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], "stream": False, "mock_response": "Hi", "api_key": "fake-key"} + headers = { + "x-litellm-spend-logs-metadata": '{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion", "timestamp": "2025-09-02T10:30:00Z"}' + } + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + "mock_response": "Hi", + "api_key": "fake-key", + } # Create mock request with headers mock_request = MagicMock(spec=Request) @@ -1078,9 +1667,7 @@ async def test_add_litellm_metadata_from_request_headers(): # Create mock user API key dict mock_user_api_key_dict = UserAPIKeyAuth( - api_key="test-key", - user_id="test-user", - org_id="test-org" + api_key="test-key", user_id="test-user", org_id="test-org" ) # Create mock proxy logging object @@ -1095,7 +1682,7 @@ async def test_add_litellm_metadata_from_request_headers(): async def mock_post_call_success_hook(*args, **kwargs): # Return the response unchanged - return kwargs.get('response', args[2] if len(args) > 2 else None) + return kwargs.get("response", args[2] if len(args) > 2 else None) mock_proxy_logging_obj.during_call_hook = mock_during_call_hook mock_proxy_logging_obj.pre_call_hook = mock_pre_call_hook @@ -1108,10 +1695,15 @@ async def test_add_litellm_metadata_from_request_headers(): general_settings = {} # Create mock select_data_generator with correct signature - def mock_select_data_generator(response=None, user_api_key_dict=None, request_data=None): + def mock_select_data_generator( + response=None, user_api_key_dict=None, request_data=None + ): async def mock_generator(): - yield "data: " + json.dumps({"choices": [{"delta": {"content": "Hello"}}]}) + "\n\n" + yield "data: " + json.dumps( + {"choices": [{"delta": {"content": "Hello"}}]} + ) + "\n\n" yield "data: [DONE]\n\n" + return mock_generator() # Create the processor @@ -1129,22 +1721,28 @@ async def test_add_litellm_metadata_from_request_headers(): select_data_generator=mock_select_data_generator, llm_router=None, model="gpt-4", - is_streaming_request=False + is_streaming_request=False, ) # Sleep for 3 seconds to allow logging to complete await asyncio.sleep(3) # Check if standard_logging_object was set - assert test_logger.standard_logging_object is not None, "standard_logging_object should be populated after LLM request" + assert ( + test_logger.standard_logging_object is not None + ), "standard_logging_object should be populated after LLM request" # Verify the logging object contains expected metadata standard_logging_obj = test_logger.standard_logging_object - print(f"Standard logging object captured: {json.dumps(standard_logging_obj, indent=4, default=str)}") + print( + f"Standard logging object captured: {json.dumps(standard_logging_obj, indent=4, default=str)}" + ) SPEND_LOGS_METADATA = standard_logging_obj["metadata"]["spend_logs_metadata"] - assert SPEND_LOGS_METADATA == dict(json.loads(headers["x-litellm-spend-logs-metadata"])), "spend_logs_metadata should be the same as the headers" + assert SPEND_LOGS_METADATA == dict( + json.loads(headers["x-litellm-spend-logs-metadata"]) + ), "spend_logs_metadata should be the same as the headers" finally: litellm.callbacks = original_callbacks @@ -1197,7 +1795,9 @@ def test_get_internal_user_header_from_mapping_returns_expected_header(): {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, ] - header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( + mappings + ) assert header_name == "X-OpenWebUI-User-Id" @@ -1205,7 +1805,9 @@ def test_get_internal_user_header_from_mapping_none_when_absent(): mappings = [ {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"} ] - header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( + mappings + ) assert header_name is None single = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"} @@ -1218,7 +1820,10 @@ def test_add_internal_user_from_user_mapping_sets_user_id_when_header_present(): headers = {"X-OpenWebUI-User-Id": "internal-user-123"} general_settings = { "user_header_mappings": [ - {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, + { + "header_name": "X-OpenWebUI-User-Id", + "litellm_user_role": "internal_user", + }, {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, ] } @@ -1312,7 +1917,7 @@ async def test_team_guardrails_append_to_key_guardrails(): metadata = updated_data.get("metadata", {}) guardrails = metadata.get("guardrails", []) - + assert "key-guardrail-1" in guardrails assert "key-guardrail-2" in guardrails assert "team-guardrail-1" in guardrails @@ -1341,7 +1946,7 @@ async def test_request_guardrails_do_not_override_key_guardrails(): metadata={"guardrails": ["key-guardrail-1"]}, team_metadata={}, ) - + # Test case: Request with empty guardrails should not result in empty guardrails data_with_empty = { "model": "gpt-3.5-turbo", @@ -1361,7 +1966,7 @@ async def test_request_guardrails_do_not_override_key_guardrails(): _metadata = updated_data_empty.get("metadata", {}) requested_guardrails = _metadata.get("guardrails", []) - + assert "guardrails" not in updated_data_empty assert "key-guardrail-1" in requested_guardrails assert len(requested_guardrails) == 1 @@ -1476,7 +2081,10 @@ def test_update_model_if_key_alias_exists(): assert data["model"] == "xai/grok-4-fast-non-reasoning" # Test case 2: Key alias doesn't exist - data = {"model": "unknown-model", "messages": [{"role": "user", "content": "Hello"}]} + data = { + "model": "unknown-model", + "messages": [{"role": "user", "content": "Hello"}], + } user_api_key_dict = UserAPIKeyAuth( api_key="test-key", aliases={"modelAlias": "xai/grok-4-fast-non-reasoning"}, @@ -1594,16 +2202,22 @@ async def test_embedding_header_forwarding_with_model_group(): # Verify that only x- prefixed headers (except x-stainless) were forwarded forwarded_headers = updated_data["headers"] - assert "X-Custom-Header" in forwarded_headers, "X-Custom-Header should be forwarded" + assert ( + "X-Custom-Header" in forwarded_headers + ), "X-Custom-Header should be forwarded" assert forwarded_headers["X-Custom-Header"] == "custom-value" assert "X-Request-ID" in forwarded_headers, "X-Request-ID should be forwarded" assert forwarded_headers["X-Request-ID"] == "test-request-123" # Verify that authorization header was NOT forwarded (sensitive header) - assert "Authorization" not in forwarded_headers, "Authorization header should not be forwarded" + assert ( + "Authorization" not in forwarded_headers + ), "Authorization header should not be forwarded" # Verify that Content-Type was NOT forwarded (doesn't start with x-) - assert "Content-Type" not in forwarded_headers, "Content-Type should not be forwarded" + assert ( + "Content-Type" not in forwarded_headers + ), "Content-Type should not be forwarded" # Verify original data fields are preserved assert updated_data["model"] == "local-openai/text-embedding-3-small" @@ -1659,8 +2273,9 @@ async def test_embedding_header_forwarding_without_model_group_config(): ) # Verify that headers were NOT added since model is not in forward list - assert "headers" not in updated_data or updated_data.get("headers") is None, \ - "Headers should not be forwarded for models not in forward_client_headers_to_llm_api list" + assert ( + "headers" not in updated_data or updated_data.get("headers") is None + ), "Headers should not be forwarded for models not in forward_client_headers_to_llm_api list" # Verify original data fields are preserved assert updated_data["model"] == "text-embedding-ada-002" @@ -1714,7 +2329,9 @@ async def test_add_guardrails_from_policy_engine(): attachment_registry = get_attachment_registry() attachment_registry._attachments = [ PolicyAttachment(policy="global-baseline", scope="*"), # applies to all - PolicyAttachment(policy="healthcare", teams=["healthcare-team"]), # applies to healthcare team + PolicyAttachment( + policy="healthcare", teams=["healthcare-team"] + ), # applies to healthcare team ] attachment_registry._initialized = True @@ -1757,7 +2374,10 @@ async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_po data = { "model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], - "policies": ["PII-POLICY-GLOBAL", "HIPAA-POLICY"], # Dynamic policies - should be accepted and removed + "policies": [ + "PII-POLICY-GLOBAL", + "HIPAA-POLICY", + ], # Dynamic policies - should be accepted and removed "metadata": {}, } @@ -1780,7 +2400,9 @@ async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_po ) # Verify that 'policies' was removed from the request body - assert "policies" not in data, "'policies' should be removed from request body to prevent forwarding to LLM provider" + assert ( + "policies" not in data + ), "'policies' should be removed from request body to prevent forwarding to LLM provider" # Verify that other fields are preserved assert "model" in data @@ -1869,7 +2491,9 @@ async def test_bearer_token_not_in_debug_logs(): from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ProxyConfig - secret_token = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.fakesignature" + secret_token = ( + "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.fakesignature" + ) mock_request = MagicMock(spec=Request) mock_request.headers = { @@ -1898,8 +2522,10 @@ async def test_bearer_token_not_in_debug_logs(): logger.setLevel(logging.DEBUG) try: - with patch("litellm.proxy.proxy_server.llm_router", None), \ - patch("litellm.proxy.proxy_server.premium_user", True): + with ( + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", True), + ): await add_litellm_data_to_request( data=data, request=mock_request, @@ -2020,9 +2646,7 @@ def test_resolve_project_model_specific_wins(): "gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}}, "defaultconfig": {"azure": {"litellm_credentials": "team-default"}}, } - result = _resolve_credential_from_model_config( - "gpt-4", project_config, team_config - ) + result = _resolve_credential_from_model_config("gpt-4", project_config, team_config) assert result == "proj-gpt4" @@ -2034,9 +2658,7 @@ def test_resolve_project_default_wins_over_team(): "gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}}, "defaultconfig": {"azure": {"litellm_credentials": "team-default"}}, } - result = _resolve_credential_from_model_config( - "gpt-4", project_config, team_config - ) + result = _resolve_credential_from_model_config("gpt-4", project_config, team_config) assert result == "proj-default" @@ -2091,12 +2713,8 @@ def test_apply_overrides_project_model_specific(setup_test_credentials): }, project_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-rec-azure"} - }, - "gpt-4-vision": { - "azure": {"litellm_credentials": "hotel-rec-vision"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-rec-azure"}}, + "gpt-4-vision": {"azure": {"litellm_credentials": "hotel-rec-vision"}}, } }, ) @@ -2123,12 +2741,8 @@ def test_apply_overrides_project_default(setup_test_credentials): }, project_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-rec-azure"} - }, - "gpt-4-vision": { - "azure": {"litellm_credentials": "hotel-rec-vision"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-rec-azure"}}, + "gpt-4-vision": {"azure": {"litellm_credentials": "hotel-rec-vision"}}, } }, ) @@ -2231,9 +2845,7 @@ def test_apply_overrides_missing_credential_name(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "gpt-4": { - "azure": {"litellm_credentials": "nonexistent-credential"} - } + "gpt-4": {"azure": {"litellm_credentials": "nonexistent-credential"}} } }, ) @@ -2272,9 +2884,7 @@ def test_apply_overrides_no_model_in_data(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "some-cred"} - } + "defaultconfig": {"azure": {"litellm_credentials": "some-cred"}} } }, ) @@ -2305,9 +2915,7 @@ def test_apply_overrides_clientside_api_version_preserved(setup_test_credentials api_key="test-key", team_metadata={ "model_config": { - "gpt-4-vision": { - "azure": {"litellm_credentials": "hotel-rec-vision"} - } + "gpt-4-vision": {"azure": {"litellm_credentials": "hotel-rec-vision"}} } }, )