diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 00000000000..1a29c0b6691 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,2 @@ +/ui/ @yuneng-jiang @ryan-crabbe-berri +/litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713230852_add_key_type_to_litellm_verification_token/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713230852_add_key_type_to_litellm_verification_token/migration.sql new file mode 100644 index 00000000000..708b7601346 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713230852_add_key_type_to_litellm_verification_token/migration.sql @@ -0,0 +1,6 @@ +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "key_type" TEXT; + +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "key_type" TEXT; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index fb4d8d0b5a3..a23cecc3911 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -422,6 +422,7 @@ model LiteLLM_VerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") @@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f575372fc3d..64d4dd578b2 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -239,6 +239,18 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_output_audio_tokens_metric"), ) + self.litellm_video_duration_seconds_metric = self._counter_factory( + "litellm_video_duration_seconds_metric", + "Seconds of video generated, from usage.duration_seconds on video generation calls", + labelnames=self.get_labels_for_metric("litellm_video_duration_seconds_metric"), + ) + + self.litellm_images_generated_metric = self._counter_factory( + "litellm_images_generated_metric", + "Number of images generated, from the image generation response", + labelnames=self.get_labels_for_metric("litellm_images_generated_metric"), + ) + # Remaining Budget for Team self.litellm_remaining_team_budget_metric = self._gauge_factory( "litellm_remaining_team_budget_metric", @@ -1336,6 +1348,12 @@ class PrometheusLogger(CustomLogger): label_context=label_context, ) + self._increment_media_generation_metrics( + standard_logging_payload=standard_logging_payload, + enum_values=enum_values, + label_context=label_context, + ) + # MCP tool call metrics self._increment_mcp_tool_call_metrics( standard_logging_payload=standard_logging_payload, @@ -1459,8 +1477,65 @@ class PrometheusLogger(CustomLogger): ), ] - for counter, metric_name, value in detail_metrics: - if not isinstance(value, (int, float)) or value <= 0: + PrometheusLogger._inc_sparse_usage_counters( + self, + detail_metrics, + enum_values=enum_values, + label_context=label_context, + ) + + def _increment_media_generation_metrics( + self, + standard_logging_payload: StandardLoggingPayload, + enum_values: UserAPIKeyLabelValues, + label_context: PrometheusLabelFactoryContext | None = None, + ) -> None: + """ + Increment video-seconds and images-generated counters from + ``standard_logging_payload["metadata"]["usage_object"]``. Video + providers report ``duration_seconds`` there; image generation calls + report ``output_image_count``. Both are sparse: only emitted when the + value is present and > 0, so token-only call types are unaffected. + """ + metadata = standard_logging_payload.get("metadata") or {} + usage_object = metadata.get("usage_object") if isinstance(metadata, dict) else None + if not isinstance(usage_object, dict): + return + + media_metrics: list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]] = [ + ( + self.litellm_video_duration_seconds_metric, + "litellm_video_duration_seconds_metric", + usage_object.get("duration_seconds"), + ), + ( + self.litellm_images_generated_metric, + "litellm_images_generated_metric", + usage_object.get("output_image_count"), + ), + ] + + PrometheusLogger._inc_sparse_usage_counters( + self, + media_metrics, + enum_values=enum_values, + label_context=label_context, + ) + + def _inc_sparse_usage_counters( + self, + counters_with_values: list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]], + enum_values: UserAPIKeyLabelValues, + label_context: PrometheusLabelFactoryContext | None = None, + ) -> None: + """ + Increment each ``(counter, metric_name, value)`` entry whose value is + a positive number. Non-numeric values (including booleans from + malformed provider usage dicts) and values <= 0 are skipped, keeping + scrape output sparse. + """ + for counter, metric_name, value in counters_with_values: + if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0: continue PrometheusLogger._inc_labeled_counter( self, @@ -1716,6 +1791,35 @@ class PrometheusLogger(CustomLogger): amount=float(response_cost), ) + @staticmethod + def _get_remaining_from_v3_rate_limit_headers( + standard_logging_payload: StandardLoggingPayload | None, + rate_limit_type: Literal["requests", "tokens"], + ) -> int | None: + """ + Read the per-(key, model) remaining value emitted by the v3 rate + limiter (``parallel_request_limiter_v3.py``), which writes + ``x-ratelimit-model_per_key-remaining-{requests,tokens}`` into + ``standard_logging_object.hidden_params.additional_headers`` instead + of the ``litellm-key-remaining-*`` metadata keys the legacy limiter + sets. The header carries no model group; it always refers to this + request's model group, which is what the gauges are labeled with. + Values are written in-process as plain ints (never HTTP-serialized + strings), so anything else is rejected rather than coerced. + """ + if standard_logging_payload is None: + return None + hidden_params = standard_logging_payload.get("hidden_params") + if hidden_params is None: + return None + additional_headers = hidden_params.get("additional_headers") + if additional_headers is None: + return None + value = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}") + if isinstance(value, bool) or not isinstance(value, int): + return None + return value + def _set_virtual_key_rate_limit_metrics( self, user_api_key: Optional[str], @@ -1733,11 +1837,20 @@ class PrometheusLogger(CustomLogger): model_group = get_model_group_from_litellm_kwargs(kwargs) remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}" remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}" + standard_logging_payload: StandardLoggingPayload | None = kwargs.get("standard_logging_object") remaining_requests = metadata.get(remaining_requests_variable_name) + if remaining_requests is None: + remaining_requests = self._get_remaining_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, rate_limit_type="requests" + ) if remaining_requests is None: remaining_requests = sys.maxsize remaining_tokens = metadata.get(remaining_tokens_variable_name) + if remaining_tokens is None: + remaining_tokens = self._get_remaining_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, rate_limit_type="tokens" + ) if remaining_tokens is None: remaining_tokens = sys.maxsize diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 352e55e9c23..b8ef9d8cca7 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -2,26 +2,8 @@ from typing import Optional from litellm.llms.openai.data_residency import infer_openai_data_residency -# Pre-define optional kwargs keys as frozenset for O(1) lookups -# These are extracted from kwargs only if present, avoiding unnecessary .get() calls -OPTIONAL_KWARGS_KEYS = frozenset( +AWS_CREDENTIAL_KWARGS_KEYS = frozenset( { - "azure_ad_token", - "tenant_id", - "client_id", - "client_secret", - "azure_username", - "azure_password", - "azure_scope", - "timeout", - "gcs_bucket_name", - "bucket_name", - "vertex_credentials", - "vertex_project", - "vertex_location", - "vertex_ai_project", - "vertex_ai_location", - "vertex_ai_credentials", "aws_region_name", "aws_access_key_id", "aws_secret_access_key", @@ -34,14 +16,40 @@ OPTIONAL_KWARGS_KEYS = frozenset( "aws_external_id", "aws_bedrock_runtime_endpoint", "aws_bedrock_project_id", - "tpm", - "rpm", - "itpm", - "otpm", - "use_xai_oauth", } ) +# Pre-define optional kwargs keys as frozenset for O(1) lookups +# These are extracted from kwargs only if present, avoiding unnecessary .get() calls +OPTIONAL_KWARGS_KEYS = ( + frozenset( + { + "azure_ad_token", + "tenant_id", + "client_id", + "client_secret", + "azure_username", + "azure_password", + "azure_scope", + "timeout", + "gcs_bucket_name", + "bucket_name", + "vertex_credentials", + "vertex_project", + "vertex_location", + "vertex_ai_project", + "vertex_ai_location", + "vertex_ai_credentials", + "tpm", + "rpm", + "itpm", + "otpm", + "use_xai_oauth", + } + ) + | AWS_CREDENTIAL_KWARGS_KEYS +) + # Backward-compatible alias for existing imports/tests. _OPTIONAL_KWARGS_KEYS = OPTIONAL_KWARGS_KEYS diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 936d79b22d6..461ab62b815 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -73,6 +73,7 @@ from litellm.litellm_core_utils.model_param_helper import ModelParamHelper from litellm.litellm_core_utils.redact_messages import ( redact_message_input_output_from_custom_logger, redact_message_input_output_from_logging, + redact_streaming_responses_for_custom_logger, ) from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.llms.base_llm.search.transformation import SearchResponse @@ -2576,6 +2577,9 @@ class Logging(LiteLLMLoggingBaseClass): model_call_details = callback.redact_standard_logging_payload_from_model_call_details( model_call_details=model_call_details ) + model_call_details = redact_streaming_responses_for_custom_logger( + model_call_details=model_call_details, custom_logger=callback + ) ################################## if self.stream is True: if "async_complete_streaming_response" in model_call_details: @@ -5208,10 +5212,15 @@ def get_standard_logging_object_payload( call_type = kwargs.get("call_type") cache_hit = kwargs.get("cache_hit", False) # Extract usage as a plain dict, avoiding Pydantic round-trip - usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict( + raw_usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict( response_obj=response_obj, combined_usage_object=cast(Optional[Usage], kwargs.get("combined_usage_object")), ) + usage_dict = ( + {**raw_usage_dict, "output_image_count": len(init_response_obj.data)} + if isinstance(init_response_obj, ImageResponse) and init_response_obj.data + else raw_usage_dict + ) id = response_obj.get("id", kwargs.get("litellm_call_id")) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index cc9264e93f8..6e8429839ad 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -38,10 +38,45 @@ def redact_message_input_output_from_custom_logger( litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger ): if hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True: - return perform_redaction(litellm_logging_obj.model_call_details, result) + return perform_redaction(litellm_logging_obj.model_call_details, result, redact_streaming_responses=False) return result +def redact_streaming_responses_for_custom_logger(model_call_details: dict, custom_logger: CustomLogger) -> dict: + """ + Returns a copy of model_call_details whose streaming response entries are redacted deepcopies + when the custom logger has opted out of message logging. The shared model_call_details is left + untouched so other callbacks still receive the unredacted response. + """ + if not (hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True): + return model_call_details + redacted_entries = { + streaming_key: _redacted_streaming_response_copy(model_call_details[streaming_key]) + for streaming_key in ("complete_streaming_response", "async_complete_streaming_response") + if model_call_details.get(streaming_key) is not None + } + if not redacted_entries: + return model_call_details + return {**model_call_details, **redacted_entries} + + +def _redacted_streaming_response_copy(streaming_response): + redacted_response = copy.deepcopy(streaming_response) + _redact_streaming_response(redacted_response) + return redacted_response + + +def _redact_streaming_response(streaming_response): + if hasattr(streaming_response, "choices"): + for choice in streaming_response.choices: + _redact_choice_content(choice) + redact_vertex_ai_metadata_from_logged_object(streaming_response) + elif hasattr(streaming_response, "output"): + _redact_responses_api_output(streaming_response.output) + if hasattr(streaming_response, "reasoning") and streaming_response.reasoning is not None: + streaming_response.reasoning = None + + def _redact_choice_content(choice): """Helper to redact content in a choice (message or delta).""" if isinstance(choice, litellm.Choices): @@ -150,9 +185,13 @@ def _redact_model_response_dict_choices(choices, redacted_str: str): _redact_choice_content(choice) -def perform_redaction(model_call_details: dict, result): +def perform_redaction(model_call_details: dict, result, redact_streaming_responses: bool = True): """ Performs the actual redaction on the logging object and result. + + redact_streaming_responses=False skips the in-place redaction of the shared streaming + response entries; per-callback redaction hands each opted-out callback its own redacted + copy via redact_streaming_responses_for_custom_logger instead. """ # Redact model_call_details model_call_details["messages"] = [{"role": "user", "content": "redacted-by-litellm"}] @@ -162,17 +201,9 @@ def perform_redaction(model_call_details: dict, result): redact_vertex_ai_metadata_from_litellm_params(model_call_details) # Redact streaming response - if model_call_details.get("stream", False) is True and "complete_streaming_response" in model_call_details: - _streaming_response = model_call_details["complete_streaming_response"] - if hasattr(_streaming_response, "choices"): - for choice in _streaming_response.choices: - _redact_choice_content(choice) - redact_vertex_ai_metadata_from_logged_object(_streaming_response) - elif hasattr(_streaming_response, "output"): - _redact_responses_api_output(_streaming_response.output) - # Redact reasoning field in ResponsesAPIResponse - if hasattr(_streaming_response, "reasoning") and _streaming_response.reasoning is not None: - _streaming_response.reasoning = None + if redact_streaming_responses and model_call_details.get("stream", False) is True: + for _streaming_key in ("complete_streaming_response", "async_complete_streaming_response"): + _redact_streaming_response(model_call_details.get(_streaming_key)) # Redact result if result is not None: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 1a052f457c5..172e54de98e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -198,14 +198,16 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def translate_tool_choice_to_responses_api( tool_choice: AnthropicMessagesToolChoice, - ) -> Dict[str, Any]: + ) -> Union[str, dict[str, Any]]: """Convert Anthropic tool_choice to Responses API tool_choice.""" tc_type = tool_choice.get("type") if tc_type == "any": - return {"type": "required"} + return "required" elif tc_type == "tool": return {"type": "function", "name": tool_choice.get("name", "")} - return {"type": "auto"} + elif tc_type == "none": + return "none" + return "auto" @staticmethod def translate_context_management_to_responses_api( diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index f449851b76f..df811f8d262 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -877,6 +877,15 @@ class BaseAWSLLM: "Resource": "*", "Condition": {"Bool": {"aws:SecureTransport": "true"}}, }, + { + "Sid": "BedrockMantleLiteLLM", + "Effect": "Allow", + "Action": [ + "bedrock-mantle:CreateInference", + ], + "Resource": "*", + "Condition": {"Bool": {"aws:SecureTransport": "true"}}, + }, ], } assume_role_params = { diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index d107ca7a0d7..3c2ae238a0b 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -20,6 +20,8 @@ from litellm.types.utils import LlmProviders from ..common_utils import OpenAIError +OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS = 16 + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -59,6 +61,19 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): key="supports_none_reasoning_effort", ) + @staticmethod + def _enforce_min_max_output_tokens(max_output_tokens: "int | None") -> "int | None": + """Raise sub-minimum max_output_tokens up to the OpenAI Responses API minimum. + + OpenAI's Responses API rejects max_output_tokens below 16 for every model + (not gpt-5 specific), so a client like Claude Code that sends a max_tokens=1 + warmup probe on model switch would otherwise 400. Values that are None or + already at/above the minimum are returned unchanged. + """ + if isinstance(max_output_tokens, int) and max_output_tokens < OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS: + return OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS + return max_output_tokens + def get_supported_openai_params(self, model: str) -> list: """ All OpenAI Responses API params are supported @@ -92,6 +107,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): """ params = dict(response_api_optional_params) + if "max_output_tokens" in params: + params["max_output_tokens"] = self._enforce_min_max_output_tokens(params.get("max_output_tokens")) + if self._is_gpt_5_model(model=model): temperature = params.get("temperature") if temperature is not None and temperature != 1: diff --git a/litellm/main.py b/litellm/main.py index 7d457d9cdd1..6fd68921fb0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -92,7 +92,10 @@ from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) -from litellm.litellm_core_utils.get_litellm_params import OPTIONAL_KWARGS_KEYS +from litellm.litellm_core_utils.get_litellm_params import ( + AWS_CREDENTIAL_KWARGS_KEYS, + OPTIONAL_KWARGS_KEYS, +) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, @@ -5322,7 +5325,7 @@ def completion( # type: ignore tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), - aws_bedrock_project_id=kwargs.get("aws_bedrock_project_id"), + **{key: kwargs[key] for key in AWS_CREDENTIAL_KWARGS_KEYS if key in kwargs}, ) cast(LiteLLMLoggingObj, logging).update_environment_variables( model=model, diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index d67726be584..519066b8266 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -36,6 +36,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): budget_reset_at: Optional[datetime] = None allowed_cache_controls: Optional[list] = [] allowed_routes: Optional[list] = [] + key_type: str | None = None permissions: Dict = {} model_spend: Dict = {} model_max_budget: Dict = {} diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index e300a22e5db..421f1dcfbea 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -18,6 +18,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti is_bridge_envelope_shaped, resolve_bridge_envelope, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, +) from litellm.proxy._types import ( UI_TEAM_ID, LiteLLM_TeamTable, @@ -543,7 +546,7 @@ class MCPRequestHandler: header_key = server.alias or server.server_name if header_key is None: raise HTTPException(status_code=500, detail="Server misconfigured: MCP server has no routable name") - admitted = await MCPRequestHandler._reload_admitted_key(result.identity.key_hash) + admitted = await MCPRequestHandler._reload_admitted_principal(result.identity) await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route) injected = {header_key: {"Authorization": result.upstream_authorization.get_secret_value()}} new_headers = {**(mcp_server_auth_headers or {}), **injected} @@ -572,6 +575,89 @@ class MCPRequestHandler: route=route, ) + @staticmethod + async def _reload_admitted_principal(identity: EnvelopeIdentity) -> UserAPIKeyAuth: + """Reload the live litellm record the envelope's subject references. + + Dispatches on the sealed subject type: a ``key_hash`` reloads the virtual key that + minted the envelope (the scripted two-header client that presents a litellm key at the + token endpoint), a ``user_id`` reloads the user that authenticated interactively (the + DCR client, whose SSO login at the bridged authorize yields a user, not a key). Both + return a ``UserAPIKeyAuth`` the caller runs through the centralized policy gate, so + team/project/org/budget/SCIM enforcement is identical to the principal presenting + itself directly.""" + match identity.subject_type: + case "key_hash": + return await MCPRequestHandler._reload_admitted_key(identity.subject) + case "user_id": + return await MCPRequestHandler._reload_admitted_user(identity.subject) + case _: + assert_never(identity.subject_type) + + @staticmethod + async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: + """Reload the live user an interactively-minted envelope references and admit them as + themselves. + + The DCR client authenticates via SSO at the bridged authorize, which yields a user + subject rather than a virtual key, so the envelope admits under the user's own + identity: the reloaded ``user_id`` and the user's own MCP object permission ride on the + returned ``UserAPIKeyAuth``, and the SAME ``get_allowed_mcp_servers`` the key path uses then + computes which servers the user may reach, so the user's litellm MCP grants and access groups + gate the request exactly as a key's do. Only the user's OWN object permission is bound: a + ``UserAPIKeyAuth`` carries a single ``team_id`` while a user may belong to many teams, so + team-inherited MCP grants for a user are a follow-up (they need a many-teams union + ``get_allowed_mcp_servers`` does not do off one auth object). The caller's centralized policy + gate enforces the user's live budget and org state, and a SCIM-deactivated owner fails closed. + + Error handling mirrors the key path's retryable-503 contract, but ``get_user_object`` defeats a + type-based check: where ``get_key_object`` raises a typed ``ProxyException`` for a missing key + and lets a DB outage propagate raw, ``get_user_object`` catches every DB failure and re-raises a + bare ``ValueError``, so a missing user and a real outage look identical and the original error + survives only as ``__context__``. ``_raise_503_if_db_unavailable`` therefore walks the cause + chain: a transient DB outage still surfaces as a retryable 503, while a missing user, or any + other non-outage resolution failure, fails closed as a 401 rather than an opaque 500. The + object-permission load shares this one boundary, so an outage there is classified the same + way (``get_object_permission`` itself swallows a failed load to ``None``, matching how + ``get_key_object`` best-effort-loads a key's object permission).""" + from litellm.proxy.auth.auth_checks import get_object_permission, get_user_object + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Server misconfigured: no database connection") + try: + user_object = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + # Resolve the user's own MCP object permission (get_user_object does not load it) so the shared + # get_allowed_mcp_servers can grant the user their litellm-granted servers. Reuses the same + # get_object_permission resolver the key and team paths use; no permission logic is duplicated. + object_permission = user_object.object_permission if user_object is not None else None + if user_object is not None and object_permission is None and user_object.object_permission_id: + object_permission = await get_object_permission( + object_permission_id=user_object.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except (ProxyException, HTTPException): + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + except Exception as e: # noqa: BLE001 # a DB outage anywhere in the resolution is a retryable 503, not an opaque 500; anything else fails closed as 401 + MCPRequestHandler._raise_503_if_db_unavailable(e) + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + if user_object is None: + raise HTTPException(status_code=401, detail="Invalid or expired credential") + if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: + raise HTTPException(status_code=401, detail="Invalid or expired credential") + return UserAPIKeyAuth( + user_id=user_object.user_id, + user_role=user_object.user_role, + object_permission=object_permission, + object_permission_id=user_object.object_permission_id, + ) + @staticmethod async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: """Reload the live key record an admitted envelope references and re-check live policy. @@ -615,10 +701,14 @@ class MCPRequestHandler: """Raise a retryable 503 when ``e`` means the auth database is unreachable, else return so the caller applies its own fail-closed mapping. A DB outage must not masquerade as an auth failure (401) or surface as an opaque 500; the caller retries. Mirrors ``UserAPIKeyAuthExceptionHandler``, - which renders a service-unavailable database error as 503 on the standard pipeline.""" + which renders a service-unavailable database error as 503 on the standard pipeline. + + Classifies across the ``__cause__``/``__context__`` chain, not just ``e`` itself: ``get_user_object`` + re-raises every DB failure as a bare ``ValueError``, so a type-based check on the top exception + would miss a real outage wrapped inside it.""" from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler - if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): raise HTTPException( status_code=503, detail="Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 8d1713a5911..b6e9a094cd5 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -12,7 +12,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response -from pydantic import BaseModel, SecretStr, ValidationError +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError from typing_extensions import assert_never from litellm._logging import verbose_logger @@ -24,6 +24,15 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( TokenEndpointAuthConfigError, build_token_endpoint_client_auth, ) +from litellm.proxy._experimental.mcp_server.faults import ( + CallerRejected, + CredentialSource, + UpstreamProtocolFault, + classify_upstream_dcr_rejection, + classify_upstream_token_rejection, + dcr_fault_detail, + render_token_fault, +) from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, @@ -41,7 +50,9 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer if TYPE_CHECKING: from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, EnvelopeKeys, + RefreshCredential, UpstreamTokenGrant, ) from litellm.proxy._types import LiteLLM_MCPServerTable, UserAPIKeyAuth @@ -98,6 +109,8 @@ def encode_state_with_base_url( code_challenge: Optional[str] = None, code_challenge_method: Optional[str] = None, client_redirect_uri: Optional[str] = None, + litellm_user_id: str | None = None, + mcp_server_id: str | None = None, ) -> str: """ Encode the base_url, original state, and PKCE parameters using encryption. @@ -108,6 +121,11 @@ def encode_state_with_base_url( code_challenge: PKCE code challenge from client code_challenge_method: PKCE code challenge method from client client_redirect_uri: Original redirect_uri from client + litellm_user_id: The SSO-authenticated litellm user captured at the bridge authorize + (interactive dcr_bridge oauth_delegate only); the callback seals it into the gateway + authorization code so the token mint can bind the envelope to this user + mcp_server_id: The bridge server the interactive flow targets, sealed alongside + litellm_user_id so the gateway code cannot be replayed against another server Returns: An encrypted string that encodes all values @@ -118,6 +136,8 @@ def encode_state_with_base_url( "code_challenge": code_challenge, "code_challenge_method": code_challenge_method, "client_redirect_uri": client_redirect_uri, + "litellm_user_id": litellm_user_id, + "mcp_server_id": mcp_server_id, } state_json = json.dumps(state_data, sort_keys=True) encrypted_state = encrypt_value_helper(state_json) @@ -145,6 +165,68 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data +_BRIDGE_AUTH_CODE_PREFIX = "llm_bcode_" + + +class _BridgeAuthorizationCode(BaseModel): + """The identity and upstream code the gateway seals into the authorization code it hands a DCR + client for an interactive dcr_bridge oauth_delegate sign-in, recovered at the token endpoint.""" + + model_config = ConfigDict(frozen=True) + upstream_code: str = Field(min_length=1) + litellm_user_id: str = Field(min_length=1) + mcp_server_id: str = Field(min_length=1) + + +def is_bridge_authorization_code(code: str) -> bool: + """Cheap prefix check that ``code`` is a gateway-sealed bridge authorization code rather than a + raw upstream code, so the token endpoint can route without decrypting.""" + return code.startswith(_BRIDGE_AUTH_CODE_PREFIX) + + +def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp_server_id: str) -> str: + """Seal the upstream authorization code and the SSO-captured litellm user into a gateway + authorization code. The DCR client only echoes this opaque value back at the token endpoint; the + gateway decrypts it there to recover the user (to bind the envelope) and the upstream code (to + exchange with the upstream), so a litellm identity captured in the browser at authorize survives + to the back-channel token call with nothing stored server-side. Encrypted with the repo's + authenticated symmetric helper (the same family the OAuth state uses), so the client can neither + read nor forge it.""" + payload = json.dumps( + {"upstream_code": upstream_code, "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id}, + sort_keys=True, + ) + return _BRIDGE_AUTH_CODE_PREFIX + encrypt_value_helper(payload) + + +def open_bridge_authorization_code(code: str) -> _BridgeAuthorizationCode | None: + """Recover the sealed identity and upstream code, or ``None`` when ``code`` is not a gateway + bridge code or does not decrypt / validate. Total over hostile input: a raw upstream code (the + scripted two-header path) returns ``None`` and the caller falls through to the existing + behavior.""" + if not is_bridge_authorization_code(code): + return None + decrypted = decrypt_value_helper( + code[len(_BRIDGE_AUTH_CODE_PREFIX) :], "bridge_authorization_code", return_original_value=False + ) + if not isinstance(decrypted, str): + return None + try: + return _BridgeAuthorizationCode.model_validate_json(decrypted) + except ValidationError: + return None + + +def _redirect_to_litellm_login(request: Request) -> RedirectResponse: + """Send an unauthenticated browser through litellm login before the interactive bridge authorize + can capture its identity. The bridge oauth_delegate flow seals the SSO user into the gateway code, + so a session is required; without one there is nothing to bind. After login the user re-initiates + the connection, which then finds the session cookie (the seamless return-to round-trip, which is + origin-validated against the control-plane URL, is a follow-up).""" + base_url = get_request_base_url(request) + return RedirectResponse(f"{base_url}/sso/key/generate") + + # LIT-4197: some upstream authorization servers reject an over-long ``state`` # (the encrypted OAuth session blob routinely exceeds their limit). The upstream # only needs an opaque value it echoes back on ``/callback``, so we forward a @@ -414,9 +496,23 @@ async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyR token = _litellm_key_from_request(request) if not token: return "no_active_key" - from litellm.proxy._types import ( # noqa: PLC0415 # inline import avoids a module-load circular import - ProxyException, - hash_token, + from litellm.proxy._types import hash_token # noqa: PLC0415 # inline import avoids a module-load circular import + + return await _reload_active_key_by_hash(hash_token(token)) + + +async def _reload_active_key_by_hash(key_hash: str) -> "_ResolvedKey | _KeyResolutionFailure": + """Reload the live key record for ``key_hash`` (cache first, then DB) and gate it on active state, + returning the resolved key or a precise failure. Shared by the token request's presented-key + resolution (:func:`_resolve_active_litellm_key`, which hashes the presented key) and the refresh + path (which already holds the hash sealed in the refresh envelope), so both re-validate identity + through one active-key gate and one failure classification. Classification mirrors admission's + ``_reload_admitted_key``: no DB connection is a gateway fault, a ``ProxyException`` / ``HTTPException`` + from ``get_key_object`` is an unknown or invalid key, a database-service-unavailable error is a + retryable outage, and anything else is an unexpected gateway fault. A blocked or expired key is + ``no_active_key``, so a revoked key can neither mint nor refresh a bridge envelope.""" + from litellm.proxy._types import ( + ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import ) from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import get_key_object, @@ -431,7 +527,6 @@ async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyR if prisma_client is None: return "unresolvable" - key_hash = hash_token(token) try: key_obj = await get_key_object( hashed_token=key_hash, @@ -444,7 +539,7 @@ async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyR if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc): return "unavailable" verbose_logger.debug( - "_resolve_active_litellm_key: unexpected key-resolution error (%s)", + "_reload_active_key_by_hash: unexpected key-resolution error (%s)", type(exc).__name__, ) return "unresolvable" @@ -453,6 +548,107 @@ async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyR return _ResolvedKey(key_hash=key_hash, key=key_obj) +async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | None": + """Re-validate a live litellm user by id, returning ``None`` when the user is active or a precise + failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a + user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a + deactivated user cannot keep refreshing, mirroring how admission re-validates the same user subject on + the egress side. No DB connection is a gateway fault (``unresolvable``) and a + database-service-unavailable error is a retryable outage (``unavailable``). Everything else fails + closed as ``no_active_key`` (the caller maps it to invalid_grant): a ``ProxyException`` / + ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` + catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look + identical, the original error surviving only as ``__context__``), so the outage check walks the cause + chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault.""" + from litellm.proxy._types import ( + ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import + ) + from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import + get_user_object, + ) + from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import + PrismaDBExceptionHandler, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + return "unresolvable" + try: + user_object = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + except (ProxyException, HTTPException): + return "no_active_key" + except Exception as exc: # noqa: BLE001 # a DB outage is retryable; a missing user (get_user_object's wrapped ValueError) or any other resolution failure fails closed as no_active_key, never a 500 + if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc): + return "unavailable" + verbose_logger.debug("_reload_active_user_by_id: user-resolution error (%s)", type(exc).__name__) + return "no_active_key" + if user_object is None: + return "no_active_key" + if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: + return "no_active_key" + return None + + +async def _key_owner_scim_deactivated(key: "UserAPIKeyAuth") -> bool: + """True only when the key's owning user was explicitly SCIM-deactivated, so a refresh revokes an + offboarded owner's key exactly as admission does via ``_reject_if_admitted_owner_scim_deactivated``. + A key with no owner, a missing owner record, or a failed lookup fails OPEN (returns ``False``), + matching admission and the standard builder: a key may outlive its owner record, and a transient DB + blip must not revoke a live key. Only an explicit ``scim_active`` of ``False`` gates renewal.""" + if key.user_id is None: + return False + from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import + get_user_object, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + return False + try: + owner = await get_user_object( + user_id=key.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + except Exception as exc: # noqa: BLE001 # fail open: a missing owner (get_user_object's wrapped ValueError) or a DB blip must not revoke a live key + verbose_logger.debug("refresh: key-owner SCIM lookup failed, not revoking (%s)", type(exc).__name__) + return False + return owner is not None and isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False + + +async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResolutionFailure | None": + """Re-validate that the subject sealed in a refresh envelope is still live, dispatching on its type: + a key_hash reloads the virtual key, a user_id reloads the user. Returns ``None`` when the subject is + active or a precise failure otherwise, so revocation gates renewal for either identity source the same + way admission gates the egress: a blocked or expired key, a SCIM-deactivated key owner (mirroring + admission's owner check, so an offboarded user cannot keep renewing a still-active key), and a + deactivated or deleted user all fail closed to ``no_active_key``.""" + match identity.subject_type: + case "key_hash": + reloaded = await _reload_active_key_by_hash(identity.subject) + if not isinstance(reloaded, _ResolvedKey): + return reloaded + if await _key_owner_scim_deactivated(reloaded.key): + return "no_active_key" + return None + case "user_id": + return await _reload_active_user_by_id(identity.subject) + case _: + assert_never(identity.subject_type) + + async def _extract_user_id_from_request(request: Request) -> str | None: """The litellm ``user_id`` for the token request, so a per-user token is stored under the same identity the egress later reads it by. Storage is best-effort, so every non-resolved outcome @@ -697,12 +893,31 @@ async def authorize_with_server( parsed = urlparse(redirect_uri) base_url = urlunparse(parsed._replace(query="")) request_base_url = get_request_base_url(request) + + # Interactive dcr_bridge oauth_delegate sign-in: this arm runs the gateway /callback and /token in + # the loop, so the gateway can capture the litellm user here (from the browser's UI session) and + # carry it to the back-channel token mint. Seal the SSO user and the target server into the state; + # the callback reads them back to mint the gateway authorization code. A DCR client cannot present a + # litellm key, so the browser session is the only identity source; without one there is nothing to + # bind, so send the user through login first. Every other oauth2 server keeps the identity-less state. + litellm_user_id: str | None = None + if mcp_server.is_dcr_bridge and mcp_server.is_oauth_delegate: + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import + _user_id_from_session_cookie, + ) + + litellm_user_id = _user_id_from_session_cookie(request) + if litellm_user_id is None: + return _redirect_to_litellm_login(request) + encoded_state = encode_state_with_base_url( base_url=base_url, original_state=state, code_challenge=code_challenge, code_challenge_method=code_challenge_method, client_redirect_uri=redirect_uri, + litellm_user_id=litellm_user_id, + mcp_server_id=mcp_server.server_id if litellm_user_id else None, ) relay_state = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES) @@ -812,7 +1027,7 @@ def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenG _BridgeMintError = Literal[ "no_identity", - "unsupported_grant", + "invalid_refresh", "identity_unavailable", "identity_unresolvable", "not_configured", @@ -824,11 +1039,14 @@ _BridgeMintError = Literal[ @dataclass(frozen=True, slots=True) class _BridgeMintReady: - """Everything the seal needs, resolved once before the exchange: the authorizing key hash and the - master-key-derived envelope keys. Passing this forward means identity resolution and key derivation - happen exactly once, and ``_finish_bridge_mint`` has no preconditions left that could fail.""" + """Everything the seal needs, resolved once before the exchange: the identity to bind the envelope + to and the master-key-derived envelope keys. The identity is a key_hash subject for the scripted + two-header client (resolved from the litellm key it presents) or a user_id subject for the + interactive SSO client (the user recovered from the gateway authorization code), so one phase-3 seal + serves both. Resolving identity here means ``_finish_bridge_mint`` has no preconditions left to + fail.""" - key_hash: str + identity: "EnvelopeIdentity" keys: "EnvelopeKeys" @@ -844,15 +1062,15 @@ def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse: status, code, desc = ( 400, "invalid_request", - "this server issues a gateway-bound credential; send a litellm credential " - "(x-litellm-api-key or Authorization) on the token request", + "this server issues a gateway-bound credential; complete the interactive sign-in, or " + "send a litellm credential (x-litellm-api-key or Authorization) on the token request", ) - case "unsupported_grant": + case "invalid_refresh": status, code, desc = ( 400, - "unsupported_grant_type", - "this server issues a gateway-bound credential and supports only the authorization_code " - "grant; re-run authorization_code to renew rather than refresh_token", + "invalid_grant", + "the refresh credential is not a valid, live refresh envelope for this server; " + "re-run authorization_code to obtain a new one", ) case "identity_unavailable": status, code, desc = ( @@ -923,45 +1141,130 @@ def _upstream_rejection_to_mint_error(rejection: _UpstreamGrantRejection) -> _Br assert_never(rejection) -async def _prepare_bridge_mint(request: Request, grant_type: str) -> "_BridgeMintReady | _BridgeMintError": - """Phase 1, BEFORE the upstream exchange: reject a grant this mint does not support, confirm the - gateway can mint (master_key set), resolve the litellm identity, and derive the envelope keys. - Returns a ready context or a precise failure value. Running before the exchange is what makes every - failure here fail closed without consuming the single-use code or rotating a refresh token. A bridge - server issues only envelopes and seals no upstream refresh_token, so the client holds none to - present: the refresh_token grant is rejected up front rather than exchanged (which could rotate the - upstream credential) and its result then discarded. Identity-resolution failures keep their origin - so the mapper statuses each truthfully.""" +async def _prepare_bridge_mint( + request: Request, + mcp_server: MCPServer, + bridge_identity: _BridgeAuthorizationCode | None = None, +) -> "_BridgeMintReady | _BridgeMintError": + """Phase 1 for the authorization_code grant, BEFORE the upstream exchange: confirm the gateway can + mint (master_key set), resolve the litellm identity, and derive the envelope keys. Returns a ready + context or a precise failure value. Running before the exchange is what makes every failure here fail + closed without consuming the single-use code. + + Two identity sources, one envelope. The interactive DCR client authenticates via SSO at the bridged + authorize, so its identity arrives as ``bridge_identity`` (the user recovered from the gateway + authorization code) and mints a user subject. The scripted two-header client presents a litellm key + on the token request instead, so its identity is the active key's hash and mints a key_hash subject. + A missing or invalid presented key keeps its resolution origin so the mapper statuses it truthfully; + neither source present is ``no_identity``. The refresh_token grant has its own phase-1 + (:func:`_prepare_bridge_refresh`), which recovers identity from the presented refresh envelope.""" from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import envelope_keys_from_master_key, ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + key_hash_identity, + user_identity, + ) from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import master_key, ) - if grant_type != "authorization_code": - return "unsupported_grant" if not master_key: return "not_configured" + keys = envelope_keys_from_master_key(master_key) + if bridge_identity is not None: + identity = user_identity(server_id=mcp_server.server_id, user_id=bridge_identity.litellm_user_id) + return _BridgeMintReady(identity=identity, keys=keys) resolved = await _resolve_active_litellm_key(request) if not isinstance(resolved, _ResolvedKey): return _key_resolution_failure_to_mint_error(resolved) - return _BridgeMintReady(key_hash=resolved.key_hash, keys=envelope_keys_from_master_key(master_key)) + identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=resolved.key_hash) + return _BridgeMintReady(identity=identity, keys=keys) + + +@dataclass(frozen=True, slots=True) +class _BridgeRefreshReady: + """A validated refresh request: the identity+keys to mint the renewed pair under, the upstream refresh + token (unwrapped from the client's refresh envelope) to exchange with the upstream IdP, and the scope + sealed alongside it at mint. The upstream refresh token is a ``SecretStr`` like every other credential + in this layer, so a repr or a traceback that captures this value never exposes the raw upstream refresh + token in plaintext. ``upstream_scope`` carries the originally-granted scope so the renewal re-requests + it when the client (a DCR/MCP client that typically omits scope on refresh) sends none, keeping the + renewed token's scope stable against an upstream that would otherwise narrow or drop it.""" + + ready: "_BridgeMintReady" + upstream_refresh_token: SecretStr + upstream_scope: str | None = None + + +def _refresh_key_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError: + """Lift an identity-resolution failure on the refresh path into the mint taxonomy. Unlike the mint + path, a resolved-but-inactive (or unknown) key is ``invalid_grant`` rather than ``invalid_request``: + the client did present an identity (sealed in the refresh envelope), but it is no longer live, so the + refresh is invalid and the client must re-authenticate. A transient outage is still 503 and a gateway + fault still 500, matching the mint path and admission.""" + match failure: + case "no_active_key": + return "invalid_refresh" + case "unavailable": + return "identity_unavailable" + case "unresolvable": + return "identity_unresolvable" + case _: + assert_never(failure) + + +async def _prepare_bridge_refresh( + mcp_server: MCPServer, refresh_value: str | None +) -> "_BridgeRefreshReady | _BridgeMintError": + """Phase 1 for the refresh_token grant, BEFORE the upstream exchange: open the client's refresh + envelope, re-validate the sealed litellm identity so a revoked key cannot keep refreshing, and + recover the upstream refresh token to exchange. Identity comes entirely from the sealed envelope, not + the HTTP request, so the request object is not needed here. The client presents a refresh envelope, + never a raw upstream refresh token, so a missing value, a non-envelope, an unopenable envelope, or one + minted for another server is ``invalid_grant``. Running before the exchange means a rejected refresh + never consumes or rotates the upstream refresh token.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import + BridgeRefreshOpened, + envelope_keys_from_master_key, + open_bridge_refresh_envelope, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + master_key, + ) + + if not master_key: + return "not_configured" + if not refresh_value: + return "invalid_refresh" + keys = envelope_keys_from_master_key(master_key) + opened = open_bridge_refresh_envelope(refresh_value, keys, datetime.now(timezone.utc), mcp_server.server_id) + if not isinstance(opened, BridgeRefreshOpened): + return "invalid_refresh" + failure = await _revalidate_active_subject(opened.identity) + if failure is not None: + return _refresh_key_failure_to_mint_error(failure) + return _BridgeRefreshReady( + ready=_BridgeMintReady(identity=opened.identity, keys=keys), + upstream_refresh_token=opened.refresh.refresh_token, + upstream_scope=opened.refresh.scope, + ) def _finish_bridge_mint( ready: "_BridgeMintReady", mcp_server: MCPServer, token_response: object, now: datetime ) -> "JSONResponse | _BridgeMintError": - """Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held envelope using - the pre-resolved identity and keys, so the client holds one bearer that admits it and forwards the - upstream token with nothing stored server-side. The only failures here are properties of the - upstream response (no usable token, an already-expired lifetime, or a token too large to seal), - returned as values.""" + """Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held access envelope + using the pre-resolved identity and keys, and, when the upstream returned a refresh token, seal a + long-lived refresh envelope alongside it so the client can renew without re-authenticating. Shared by + the authorization_code and refresh_token paths, so a renewal that the upstream rotates re-issues a + fresh refresh envelope. The only hard failures here are properties of the upstream access token (no + usable token, an already-expired lifetime, or a token too large to seal); a refresh token that cannot + be sealed degrades to an access-only response rather than failing the whole exchange.""" from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import build_bridge_token_response, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import - EnvelopeIdentity, SealedEnvelope, UpstreamTokenGrant, ) @@ -969,17 +1272,88 @@ def _finish_bridge_mint( grant = _bridge_grant_from_token_response(token_response) if not isinstance(grant, UpstreamTokenGrant): return _upstream_rejection_to_mint_error(grant) - identity = EnvelopeIdentity(server_id=mcp_server.server_id, key_hash=ready.key_hash) - sealed = build_bridge_token_response(identity, grant, ready.keys, now) + sealed = build_bridge_token_response(ready.identity, grant, ready.keys, now) if not isinstance(sealed, SealedEnvelope): return "too_large" # Report expires_in from the JWT's own second-truncated exp, rounding the elapsed portion up, so the # client is never told the bearer lives past the point admission (which uses that exp) rejects it. expires_in = max(0, int(sealed.expires_at.timestamp()) - math.ceil(now.timestamp())) - body = {"access_token": sealed.token.get_secret_value(), "token_type": "Bearer", "expires_in": expires_in} + refresh_envelope = _mint_refresh_envelope_value(ready.identity, token_response, ready.keys, now, mcp_server) + body = { + "access_token": sealed.token.get_secret_value(), + "token_type": "Bearer", + "expires_in": expires_in, + # A refresh envelope rides along only when the upstream returned a refresh token to seal; when it + # rotates on renewal, the client receives the new one and the old envelope's upstream token dies. + **({"refresh_token": refresh_envelope} if refresh_envelope is not None else {}), + } return JSONResponse(body, headers=TOKEN_NO_CACHE_HEADERS) +def _token_credential_source(mcp_server: MCPServer) -> CredentialSource: + """Mirrors the resolved-client rule in :func:`exchange_token_with_server`: when the server has a + stored client_id the gateway presents its own credentials upstream, so a credential rejection is + the operator's fault, not the caller's.""" + return "gateway_stored" if mcp_server.client_id else "caller_supplied" + + +def _upstream_refresh_credential(token_response: object) -> "RefreshCredential | None": + """Extract the upstream refresh grant from a token response, or ``None`` when there is none to seal. + Each field is isinstance-checked so nothing untyped reaches the refresh envelope; ``refresh_expires_in`` + (the refresh token's own lifetime, when the upstream reports it) is classified like ``expires_in`` and + bounds the refresh envelope's TTL. An upstream that reports the refresh token itself as already elapsed + (``refresh_expires_in`` non-positive) yields ``None`` rather than a refresh envelope: sealing a dead + token would hand the client a full-TTL-capped envelope the IdP will reject, so the exchange degrades to + an access-only response (the client re-authenticates at access expiry), mirroring how + :func:`_bridge_grant_from_token_response` refuses an already-elapsed access token instead of capping it.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + RefreshCredential, + ) + + if not isinstance(token_response, dict): + return None + refresh = token_response.get("refresh_token") + if not isinstance(refresh, str) or not refresh: + return None + lifetime = _classify_upstream_lifetime(token_response.get("refresh_expires_in")) + if lifetime == "expired": + return None + scope = token_response.get("scope") + return RefreshCredential( + refresh_token=SecretStr(refresh), + scope=scope if isinstance(scope, str) and scope else None, + expires_in=lifetime if isinstance(lifetime, int) else None, + ) + + +def _mint_refresh_envelope_value( + identity: "EnvelopeIdentity", token_response: object, keys: "EnvelopeKeys", now: datetime, mcp_server: MCPServer +) -> str | None: + """Seal the upstream refresh grant (if any) into a refresh envelope and return its bearer string, or + ``None`` when the upstream returned no refresh token or the refresh token is too large to seal. A + too-large refresh token degrades to an access-only response (logged) rather than failing an exchange + that already succeeded upstream: the client simply re-authenticates when the access envelope expires.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import + build_bridge_refresh_token_response, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + SealedEnvelope, + ) + + refresh_credential = _upstream_refresh_credential(token_response) + if refresh_credential is None: + return None + sealed = build_bridge_refresh_token_response(identity, refresh_credential, keys, now) + if isinstance(sealed, SealedEnvelope): + return sealed.token.get_secret_value() + verbose_logger.warning( + "bridge mint: the upstream refresh token is too large to seal into a refresh envelope for " + "server=%s; issuing an access-only response, so the client re-authenticates at access expiry", + mcp_server.server_id, + ) + return None + + async def exchange_token_with_server( request: Request, mcp_server: MCPServer, @@ -1014,25 +1388,61 @@ async def exchange_token_with_server( except TokenEndpointAuthConfigError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc + bridge_identity: _BridgeAuthorizationCode | None = None + bridge_mint_ready: _BridgeMintReady | None = None + bridge_upstream_refresh: SecretStr | None = None + bridge_upstream_scope: str | None = None + refresh_request_scope: str | None = None + is_bridge = mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge + if grant_type == "refresh_token": - if not refresh_token: + # Phase 1 for a bridge refresh: open the client's refresh envelope, re-validate the sealed + # identity, and unwrap the real upstream refresh token BEFORE building token_data, so the exchange + # sends the upstream token and never the envelope. A failure returns without touching the upstream. + if is_bridge: + prepared_refresh = await _prepare_bridge_refresh(mcp_server, refresh_token) + if not isinstance(prepared_refresh, _BridgeRefreshReady): + return _bridge_mint_error_response(prepared_refresh) + bridge_mint_ready = prepared_refresh.ready + bridge_upstream_refresh = prepared_refresh.upstream_refresh_token + bridge_upstream_scope = prepared_refresh.upstream_scope + # A bridge server sends the unwrapped upstream refresh token recovered from the client's refresh + # envelope above; every other server sends the client's own refresh token verbatim. + upstream_refresh_token = ( + bridge_upstream_refresh.get_secret_value() if bridge_upstream_refresh is not None else refresh_token + ) + if not upstream_refresh_token: raise HTTPException( status_code=400, detail="refresh_token is required for refresh_token grant", ) token_data: dict = { "grant_type": "refresh_token", - "refresh_token": refresh_token, + "refresh_token": upstream_refresh_token, **client_auth.body, } - if scope: - token_data["scope"] = scope + refresh_request_scope = scope or bridge_upstream_scope + if refresh_request_scope: + token_data["scope"] = refresh_request_scope else: if not code: raise HTTPException( status_code=400, detail="code is required for authorization_code grant", ) + # Interactive dcr_bridge oauth_delegate: the client presents the gateway authorization code the + # callback sealed. Recover the SSO user and the real upstream code from it; the upstream exchange + # below uses the upstream code, and the mint binds the envelope to the recovered user. Bind the + # sealed server to this request so a code minted for one bridge server cannot be spent at another. + # A raw upstream code (scripted path) opens to None and the code is used as-is. + bridge_identity = open_bridge_authorization_code(code) + if bridge_identity is not None: + if bridge_identity.mcp_server_id != mcp_server.server_id: + raise HTTPException( + status_code=400, + detail="Authorization code was issued for a different MCP server", + ) + code = bridge_identity.upstream_code bridge_token_relay = _dcr_bridge_relays_client_registration(mcp_server) if bridge_token_relay and not redirect_uri: raise HTTPException( @@ -1052,40 +1462,48 @@ async def exchange_token_with_server( } if code_verifier: token_data["code_verifier"] = code_verifier - - # Phase 1: for a bridge oauth_delegate mint, validate all preconditions and resolve identity+keys - # BEFORE the exchange below consumes the single-use upstream code, and carry the ready context to - # phase 3. A failure here returns without ever touching the upstream credential. - bridge_mint_ready: _BridgeMintReady | None = None - if mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge: - prepared = await _prepare_bridge_mint(request, grant_type) - if not isinstance(prepared, _BridgeMintReady): - return _bridge_mint_error_response(prepared) - bridge_mint_ready = prepared - + # Phase 1 for a bridge authorization_code mint: resolve identity (the SSO user recovered above, or + # the presented litellm key) and the envelope keys BEFORE the exchange consumes the single-use code. + if is_bridge: + prepared = await _prepare_bridge_mint(request, mcp_server, bridge_identity) + if not isinstance(prepared, _BridgeMintReady): + return _bridge_mint_error_response(prepared) + bridge_mint_ready = prepared async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) - response = await async_client.post( - mcp_server.token_url, - headers={"Accept": "application/json", **client_auth.headers}, - data=token_data, - ) + try: + response = await async_client.post( + mcp_server.token_url, + headers={"Accept": "application/json", **client_auth.headers}, + data=token_data, + ) + if response is not None: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + fault = classify_upstream_token_rejection( + exc.response, + credential_source=_token_credential_source(mcp_server), + log_context=mcp_server.server_id, + ) + upstream_rejected_bridge_refresh = ( + is_bridge + and grant_type == "refresh_token" + and isinstance(fault, CallerRejected) + and fault.code == "invalid_grant" + ) + if upstream_rejected_bridge_refresh: + verbose_logger.info( + "bridge refresh: the upstream rejected the sealed refresh token for server=%s with " + "invalid_grant (revoked or expired at the IdP); returning invalid_grant so the client " + "re-runs authorization_code rather than an opaque upstream error", + mcp_server.server_id, + ) + return _bridge_mint_error_response("invalid_refresh") + return render_token_fault(fault) if response is None: raise HTTPException( status_code=502, detail="MCP upstream token endpoint returned no response", ) - - try: - response.raise_for_status() - except httpx.HTTPStatusError as exc: - if "invalid_target" in exc.response.text: - verbose_logger.warning( - "MCP server %s: the upstream authorization server rejected the token request with " - "invalid_target; it may require RFC 8707 resource indicators, which the gateway " - "does not send yet (tracked as LIT-4339)", - mcp_server.server_id, - ) - raise token_response = response.json() # Validate token response against server-configured rules before any storage. @@ -1130,13 +1548,19 @@ async def exchange_token_with_server( # upstream token) instead of the raw upstream token, so the one bearer both admits the caller and # forwards the upstream credential. Only this mode mints; every other server returns the raw token. if bridge_mint_ready is not None: + if refresh_request_scope and isinstance(token_response, dict) and not token_response.get("scope"): + token_response = {**token_response, "scope": refresh_request_scope} # Phase 3: seal the upstream grant into the client-held envelope; failures map through the same # OAuth-shaped response as the phase-1 preconditions. minted = _finish_bridge_mint(bridge_mint_ready, mcp_server, token_response, datetime.now(timezone.utc)) return minted if isinstance(minted, JSONResponse) else _bridge_mint_error_response(minted) + raw_access_token = token_response.get("access_token") if isinstance(token_response, dict) else None + if not isinstance(raw_access_token, str) or not raw_access_token: + return render_token_fault(UpstreamProtocolFault(note="the upstream token response has no usable access_token")) + result = { - "access_token": token_response["access_token"], + "access_token": raw_access_token, "token_type": token_response.get("token_type", "Bearer"), } @@ -1392,21 +1816,6 @@ async def _persist_dcr_client_registration( return "failed" -_MAX_UPSTREAM_ERROR_CHARS = 500 - - -def _safe_upstream_error_detail(response: httpx.Response) -> str: - """Bounded plaintext summary of an upstream registration failure for the client. - - RFC 7591 error bodies are small JSON objects (``error`` / ``error_description``); relaying the - text lets the client read the real reason instead of a bare 500, and the length bound keeps a - hostile or oversized upstream body from bloating the gateway response.""" - body = response.text - if not body: - return response.reason_phrase or "upstream registration failed" - return body[:_MAX_UPSTREAM_ERROR_CHARS] - - async def register_client_with_server( request: Request, mcp_server: MCPServer, @@ -1466,19 +1875,24 @@ async def register_client_with_server( } async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Register) - response = await async_client.post( - mcp_server.registration_url, - headers=headers, - json=register_data, - ) + try: + response = await async_client.post( + mcp_server.registration_url, + headers=headers, + json=register_data, + ) + if response is not None: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + status_code, detail = dcr_fault_detail( + classify_upstream_dcr_rejection(exc.response, log_context=mcp_server.server_id) + ) + raise HTTPException(status_code=status_code, detail=detail) from exc if response is None: raise HTTPException( status_code=502, detail="MCP upstream registration endpoint returned no response", ) - if bridge_relay and response.status_code >= 400: - raise HTTPException(status_code=response.status_code, detail=_safe_upstream_error_detail(response)) - response.raise_for_status() token_response = response.json() @@ -1706,7 +2120,20 @@ async def callback( # states while permitting same-origin / allowlisted clients. redirect_uri = _get_validated_client_redirect_uri(request, state_data) - params = {"code": code, "state": original_state} + # Interactive dcr_bridge oauth_delegate: the state carries the litellm user the authorize step + # captured. Instead of forwarding the raw upstream code (which the client would present at the + # token endpoint with no way to prove who signed in), seal the user and the upstream code into a + # gateway authorization code and forward THAT. The token endpoint decrypts it to bind the + # envelope to this user. Every other flow forwards the raw code unchanged. + litellm_user_id = state_data.get("litellm_user_id") + mcp_server_id = state_data.get("mcp_server_id") + forwarded_code = code + if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id: + forwarded_code = seal_bridge_authorization_code( + upstream_code=code, litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id + ) + + params = {"code": forwarded_code, "state": original_state} complete_returned_url = _append_query_params(redirect_uri, params) response = RedirectResponse(url=complete_returned_url, status_code=302) _clear_oauth_state_cookie(response, request, state) diff --git a/litellm/proxy/_experimental/mcp_server/faults/__init__.py b/litellm/proxy/_experimental/mcp_server/faults/__init__.py new file mode 100644 index 00000000000..da078f0e242 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/faults/__init__.py @@ -0,0 +1,38 @@ +"""Typed fault values for upstream OAuth/DCR failures (phase 1 of the MCP error-handling framework). + +The invariant this package exists to enforce: an upstream failure is classified ONCE into a single +fault value, and the response status, wire error code, and prose are all derived from that value. +Deriving all three from one classification makes contradictory pairings (a caller-fault error code on +a server-fault status) unrepresentable, and gives the trust-boundary rule one enforcement point: +spec-defined machine fields may cross to callers, upstream prose and raw bodies go to server logs. +""" + +from litellm.proxy._experimental.mcp_server.faults.classify import ( + classify_upstream_dcr_rejection, + classify_upstream_token_rejection, +) +from litellm.proxy._experimental.mcp_server.faults.render_oauth import ( + dcr_fault_detail, + render_token_fault, +) +from litellm.proxy._experimental.mcp_server.faults.types import ( + CallerRejected, + CredentialSource, + GatewayRejected, + UpstreamOAuthFault, + UpstreamProtocolFault, + UpstreamReportedFault, +) + +__all__ = [ + "CallerRejected", + "CredentialSource", + "GatewayRejected", + "UpstreamOAuthFault", + "UpstreamProtocolFault", + "UpstreamReportedFault", + "classify_upstream_dcr_rejection", + "classify_upstream_token_rejection", + "dcr_fault_detail", + "render_token_fault", +] diff --git a/litellm/proxy/_experimental/mcp_server/faults/classify.py b/litellm/proxy/_experimental/mcp_server/faults/classify.py new file mode 100644 index 00000000000..8b3a09f8d8d --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/faults/classify.py @@ -0,0 +1,133 @@ +"""The single place that reads upstream OAuth/DCR failure responses. + +Every accessor here is total: an upstream that lies about its content encoding, sends an undecodable +body, or omits the spec fields yields a classified fault, never an exception. Nothing outside this +module should touch a failed upstream response's body. +""" + +from __future__ import annotations + +import httpx + +from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.faults.types import ( + GATEWAY_CAPABILITY_CODES, + GATEWAY_CREDENTIAL_CODES, + MAX_WIRE_FIELD_CHARS, + CallerRejected, + CredentialSource, + GatewayRejected, + UpstreamOAuthFault, + UpstreamProtocolFault, + UpstreamReportedFault, +) + + +def _safe_text(response: httpx.Response) -> str: + try: + return response.text + except Exception: + return "" + + +def _safe_json(response: httpx.Response) -> object: + try: + return response.json() + except Exception: + return None + + +def _bounded_field(value: object) -> str | None: + if not isinstance(value, str) or not value: + return None + return value[:MAX_WIRE_FIELD_CHARS] + + +def _log_out_of_contract(endpoint_kind: str, response: httpx.Response, log_context: str) -> None: + verbose_logger.warning( + "MCP upstream %s endpoint (%s) returned HTTP %s outside the OAuth error contract (first %s chars): %s", + endpoint_kind, + log_context, + response.status_code, + MAX_WIRE_FIELD_CHARS, + _safe_text(response)[:MAX_WIRE_FIELD_CHARS], + ) + + +def _classify_oauth_error_code( + code: str, + description: str | None, + error_uri: str | None, + credential_source: CredentialSource, + log_context: str, +) -> UpstreamOAuthFault: + """Blame assignment for a contract-conformant OAuth error code, shared by the token and DCR + classifiers. Codes by which the upstream blames itself keep that blame; ``invalid_target`` is a + gateway capability gap (RFC 8707 resource indicators, LIT-4339) no matter whose credentials were + presented; credential-indicting codes follow the credential source; everything else, including + codes we do not recognize, is the caller's to act on. The upstream's HTTP status is deliberately + never consulted: status derives from this classification at render time, which is what keeps + status and code from contradicting each other.""" + if code == "server_error" or code == "temporarily_unavailable": + return UpstreamReportedFault(code=code) + if code in GATEWAY_CAPABILITY_CODES: + verbose_logger.warning( + "MCP server %s: the upstream authorization server rejected the request with " + "invalid_target; it may require RFC 8707 resource indicators, which the gateway " + "does not send yet (tracked as LIT-4339)", + log_context, + ) + return GatewayRejected(code=code) + if credential_source == "gateway_stored" and code in GATEWAY_CREDENTIAL_CODES: + verbose_logger.warning( + "MCP server %s: upstream authorization server rejected the gateway's configured client " + "credentials (%s): %s", + log_context, + code, + description or "", + ) + return GatewayRejected(code=code) + return CallerRejected(code=code, description=description, error_uri=error_uri) + + +def classify_upstream_token_rejection( + response: httpx.Response, + credential_source: CredentialSource, + log_context: str, +) -> UpstreamOAuthFault: + """Classify a token-endpoint rejection into exactly one fault: a body with an RFC 6749 §5.2 + ``error`` field goes through blame assignment (:func:`_classify_oauth_error_code`); anything + without a usable ``error`` field is an upstream protocol fault.""" + parsed = _safe_json(response) + fields = parsed if isinstance(parsed, dict) else {} + code = _bounded_field(fields.get("error")) + if code is None: + _log_out_of_contract("token", response, log_context) + return UpstreamProtocolFault(note=f"upstream token endpoint returned HTTP {response.status_code}") + return _classify_oauth_error_code( + code, + description=_bounded_field(fields.get("error_description")), + error_uri=_bounded_field(fields.get("error_uri")), + credential_source=credential_source, + log_context=log_context, + ) + + +def classify_upstream_dcr_rejection(response: httpx.Response, log_context: str) -> UpstreamOAuthFault: + """Classify a dynamic-client-registration rejection. RFC 7591 §3.2.2 errors carry + ``error`` / ``error_description`` and go through the same blame assignment as token errors + (registration sends no client credentials, so credential codes stay caller-actionable); anything + without a usable ``error`` field is an upstream protocol fault.""" + parsed = _safe_json(response) + fields = parsed if isinstance(parsed, dict) else {} + code = _bounded_field(fields.get("error")) + if code is None: + _log_out_of_contract("registration", response, log_context) + return UpstreamProtocolFault(note=f"upstream registration failed with HTTP {response.status_code}") + return _classify_oauth_error_code( + code, + description=_bounded_field(fields.get("error_description")), + error_uri=None, + credential_source="caller_supplied", + log_context=log_context, + ) diff --git a/litellm/proxy/_experimental/mcp_server/faults/render_oauth.py b/litellm/proxy/_experimental/mcp_server/faults/render_oauth.py new file mode 100644 index 00000000000..89ce5011830 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/faults/render_oauth.py @@ -0,0 +1,89 @@ +"""Render upstream OAuth/DCR faults onto the wire. The only place that chooses statuses and bodies +for these faults, so every consumer emits the same contract: RFC 6749 §5.2-shaped JSON with the §5.1 +no-store headers on token endpoints, HTTPException details on registration. Status, code, and prose +all derive from the fault tag; exhaustive matches keep a new fault arm from shipping unrendered. +""" + +from __future__ import annotations + +from fastapi.responses import JSONResponse +from typing_extensions import assert_never + +from litellm.proxy._experimental.mcp_server.faults.types import UpstreamOAuthFault +from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS + + +def _gateway_rejected_description(code: str) -> str: + if code == "invalid_target": + return ( + "the upstream authorization server rejected the request (invalid_target); " + "it may require RFC 8707 resource indicators, which the gateway does not send yet" + ) + return ( + f"the upstream authorization server rejected the gateway's configured client credentials " + f"({code}); verify the MCP server's client_id and client_secret" + ) + + +def _upstream_reported_status_and_description(code: str) -> tuple[int, str]: + if code == "temporarily_unavailable": + return 503, "the upstream authorization server is temporarily unavailable; retry shortly" + return 502, "the upstream authorization server reported an internal error" + + +def render_token_fault(fault: UpstreamOAuthFault) -> JSONResponse: + """RFC 6749 §5.2 response for a token-endpoint fault. Caller-actionable rejections relay the + upstream's code on the status that code implies (401 for invalid_client per §5.2, else 400); + gateway-side faults are 502 ``server_error`` with gateway-authored prose so a caller is never + blamed for, or shown the internals of, a failure only the operator can fix.""" + match fault.tag: + case "caller_rejected": + content = { + "error": fault.code, + **({"error_description": fault.description} if fault.description else {}), + **({"error_uri": fault.error_uri} if fault.error_uri else {}), + } + status_code = 401 if fault.code == "invalid_client" else 400 + return JSONResponse(status_code=status_code, content=content, headers=TOKEN_NO_CACHE_HEADERS) + case "gateway_rejected": + return JSONResponse( + status_code=502, + content={ + "error": "server_error", + "error_description": _gateway_rejected_description(fault.code), + }, + headers=TOKEN_NO_CACHE_HEADERS, + ) + case "upstream_reported_fault": + status_code, description = _upstream_reported_status_and_description(fault.code) + return JSONResponse( + status_code=status_code, + content={"error": fault.code, "error_description": description}, + headers=TOKEN_NO_CACHE_HEADERS, + ) + case "upstream_protocol_fault": + return JSONResponse( + status_code=502, + content={"error": "server_error", "error_description": fault.note}, + headers=TOKEN_NO_CACHE_HEADERS, + ) + case _: + assert_never(fault.tag) + + +def dcr_fault_detail(fault: UpstreamOAuthFault) -> tuple[int, str]: + """Status and detail string for a registration fault, raised as HTTPException by the caller. + RFC 7591 §3.2.2 defines registration errors as 400, so a contract-conformant rejection is 400 + regardless of the status the upstream chose; everything else is a 502 upstream fault.""" + match fault.tag: + case "caller_rejected": + detail = f"{fault.code}: {fault.description}" if fault.description else fault.code + return 400, detail + case "gateway_rejected": + return 502, _gateway_rejected_description(fault.code) + case "upstream_reported_fault": + return _upstream_reported_status_and_description(fault.code) + case "upstream_protocol_fault": + return 502, fault.note + case _: + assert_never(fault.tag) diff --git a/litellm/proxy/_experimental/mcp_server/faults/types.py b/litellm/proxy/_experimental/mcp_server/faults/types.py new file mode 100644 index 00000000000..128b5e3e6cf --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/faults/types.py @@ -0,0 +1,79 @@ +"""Fault taxonomy for upstream OAuth token and DCR registration failures. + +Each fault is a frozen model on a ``tag`` literal. The tag alone decides the HTTP status, the wire +error code, and whose prose the caller sees, so those three facts can never disagree the way they can +when an upstream's status and error code are relayed independently. +""" + +from __future__ import annotations + +from typing import Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict + +MAX_WIRE_FIELD_CHARS = 500 +"""Bound on every upstream-derived string that crosses to a caller or into a log line.""" + +CredentialSource: TypeAlias = Literal["gateway_stored", "caller_supplied"] +"""Whose client credentials the gateway presented upstream: the MCP server's stored configuration or +credentials the caller supplied on the request. Decides whether a credential rejection is the +caller's problem to fix or the gateway operator's.""" + +GATEWAY_CREDENTIAL_CODES: frozenset[str] = frozenset({"invalid_client", "unauthorized_client"}) +"""RFC 6749 error codes that indict the OAuth client's credentials or grant authorization. When the +gateway presented its own stored credentials, these are gateway-side faults the caller cannot act on; +when the caller supplied the credentials, they are the caller's to fix.""" + +GATEWAY_CAPABILITY_CODES: frozenset[str] = frozenset({"invalid_target"}) +"""Codes that indict a gateway capability regardless of whose credentials were presented: +``invalid_target`` means the upstream wants RFC 8707 resource indicators, which the gateway does not +send yet (LIT-4339). Never the caller's fault.""" + +UPSTREAM_FAULT_CODES: frozenset[str] = frozenset({"server_error", "temporarily_unavailable"}) +"""Codes by which the upstream blames itself. Relaying them as caller faults would invert blame, so +they classify as upstream-reported faults and render on the 5xx their meaning implies.""" + + +class CallerRejected(BaseModel): + """The upstream spoke the OAuth error contract and the failure is actionable by our caller + (e.g. ``invalid_grant``: re-run authorization). The code and its bounded prose relay on the + 4xx status the code itself implies.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["caller_rejected"] = "caller_rejected" + code: str + description: str | None = None + error_uri: str | None = None + + +class GatewayRejected(BaseModel): + """The upstream rejected the request for a cause only the gateway operator can address: the + server's stored client credentials or a gateway capability gap. Not actionable by the caller: + rendered as 502 with gateway-authored prose naming the code; the upstream's prose goes to + server logs only.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["gateway_rejected"] = "gateway_rejected" + code: str + + +class UpstreamReportedFault(BaseModel): + """The upstream blamed itself in the OAuth vocabulary. Rendered on the 5xx the code implies + (``server_error`` 502, ``temporarily_unavailable`` 503) so blame and status agree.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["upstream_reported_fault"] = "upstream_reported_fault" + code: Literal["server_error", "temporarily_unavailable"] + + +class UpstreamProtocolFault(BaseModel): + """The upstream broke the error contract: no JSON ``error`` field, an undecodable body, or a + success response without a usable token. Rendered as 502 with a gateway-authored note; the + upstream body never crosses to the caller.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["upstream_protocol_fault"] = "upstream_protocol_fault" + note: str + + +UpstreamOAuthFault: TypeAlias = CallerRejected | GatewayRejected | UpstreamReportedFault | UpstreamProtocolFault diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py index 5530fbc46fd..c352f3a683e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py @@ -21,11 +21,16 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import EnvelopeKeys, EnvelopeMintError, OpenedEnvelope, + OpenedRefreshEnvelope, + RefreshCredential, SealedEnvelope, UpstreamTokenGrant, is_envelope, + is_refresh_envelope, mint_envelope, + mint_refresh_envelope, open_envelope, + open_refresh_envelope, ) _SIGNING_KEY_DOMAIN = b"litellm-mcp-bridge:envelope-signing:" @@ -92,6 +97,67 @@ def build_bridge_token_response( return mint_envelope(identity, grant, keys, now) +def build_bridge_refresh_token_response( + identity: EnvelopeIdentity, + refresh: RefreshCredential, + keys: EnvelopeKeys, + now: datetime, +) -> SealedEnvelope | EnvelopeMintError: + """Seal ``refresh`` for ``identity`` into the long-lived refresh envelope the token endpoint returns + alongside the access envelope, so the client can renew without re-authenticating. A thin, pure + wrapper over :func:`mint_refresh_envelope`; returns the mint error as a value for the caller to map. + """ + return mint_refresh_envelope(identity, refresh, keys, now) + + +class BridgeRefreshOpened(BaseModel): + """A valid refresh envelope presented to the token endpoint: the identity to re-validate and renew + under, and the upstream refresh grant to exchange.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["opened"] = "opened" + identity: EnvelopeIdentity + refresh: RefreshCredential + + +class BridgeRefreshInvalid(BaseModel): + """The presented refresh grant is not a valid refresh envelope for this server (not refresh-shaped, + will not open, or minted for a different server); the token endpoint fails the refresh closed.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["invalid"] = "invalid" + + +BridgeRefreshResult: TypeAlias = BridgeRefreshOpened | BridgeRefreshInvalid + + +def open_bridge_refresh_envelope( + refresh_value: str, + keys: EnvelopeKeys, + now: datetime, + expected_server_id: str, +) -> BridgeRefreshResult: + """Open a refresh envelope a bridge ``oauth_delegate`` client presented on a refresh_token grant. + + The token-endpoint mirror of :func:`resolve_bridge_envelope`: strips an optional ``Bearer`` scheme, + then returns ``BridgeRefreshOpened`` with the recovered identity and upstream refresh grant, or + ``BridgeRefreshInvalid`` for anything that is not a valid refresh envelope for this server. Never + raises; total over hostile input via :func:`open_refresh_envelope`. ``expected_server_id`` binds the + envelope to the server the request targets, so a refresh envelope minted for one server cannot renew + against another. A raw upstream refresh token (not envelope-shaped) is ``BridgeRefreshInvalid``: this + mode never hands the client a bare upstream refresh token, so it must never accept one. + """ + candidate = _strip_bearer(refresh_value) + if not is_refresh_envelope(candidate): + return BridgeRefreshInvalid() + opened = open_refresh_envelope(candidate, keys, now) + if not isinstance(opened, OpenedRefreshEnvelope): + return BridgeRefreshInvalid() + if opened.identity.server_id != expected_server_id: + return BridgeRefreshInvalid() + return BridgeRefreshOpened(identity=opened.identity, refresh=opened.refresh) + + class NotBridgeEnvelope(BaseModel): """The bearer is not an envelope; admission continues on its normal path.""" @@ -128,10 +194,12 @@ def _strip_bearer(value: str) -> str: def is_bridge_envelope_shaped(authorization_value: str) -> bool: - """Cheap, keyless test that an ``Authorization`` value carries an envelope (optional - ``Bearer`` scheme stripped). The admission edge engages the bridge arm only for an - envelope, so a plain upstream bearer falls through to normal oauth2 admission.""" - return is_envelope(_strip_bearer(authorization_value)) + """Cheap, keyless test that an ``Authorization`` value carries an envelope of either kind (optional + ``Bearer`` scheme stripped). The admission edge engages the bridge arm for an access envelope (to + admit) and for a refresh envelope (to reject it explicitly, since a refresh credential is never + usable at the tool-call edge); a plain upstream bearer falls through to normal oauth2 admission.""" + candidate = _strip_bearer(authorization_value) + return is_envelope(candidate) or is_refresh_envelope(candidate) def resolve_bridge_envelope( @@ -148,6 +216,10 @@ def resolve_bridge_envelope( envelope, and ``BridgeEnvelopeInvalid`` for an envelope-shaped bearer that will not open. Never raises: it is total over hostile input via :func:`open_envelope`. + A refresh envelope is ``BridgeEnvelopeInvalid`` here: it is a valid gateway credential but only ever + presented back to the token endpoint, never usable to authenticate a tool call, so admission must + fail it closed rather than let it fall through to another arm. + ``expected_server_id`` is the ``server_id`` of the MCP server the request targets; an opened envelope whose sealed ``server_id`` does not match is rejected as ``BridgeEnvelopeInvalid``. Binding here (rather than leaving it to the caller) prevents @@ -157,6 +229,8 @@ def resolve_bridge_envelope( unlike ``hmac.compare_digest`` on ``str``, does not raise on a non-ASCII server_id. """ candidate = _strip_bearer(authorization_value) + if is_refresh_envelope(candidate): + return BridgeEnvelopeInvalid() if not is_envelope(candidate): return NotBridgeEnvelope() opened = open_envelope(candidate, keys, now) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py index 517c2ef5c8f..9118a3e129d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -44,18 +44,33 @@ from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value ENVELOPE_PREFIX = "llm_env_" -"""Marker prefix on every serialized envelope so the edge can cheaply tell an envelope +"""Marker prefix on every serialized ACCESS envelope so the edge can cheaply tell an envelope from a raw upstream token before doing any cryptography.""" +REFRESH_ENVELOPE_PREFIX = "llm_refresh_" +"""Marker prefix on every serialized REFRESH envelope. A distinct prefix keeps the two credentials +routable without crypto and, together with the signed ``kind`` claim, stops one from being presented +where the other is expected: a refresh envelope carries a long-lived upstream refresh token and is only +ever presented back to the token endpoint, never forwarded upstream on a tool call.""" + ENVELOPE_ISSUER = "litellm-mcp-bridge" """``iss`` claim stamped into every envelope and required back on open.""" MAX_ENVELOPE_TTL_SECONDS = 3600 -"""Hard ceiling on envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)`` +"""Hard ceiling on ACCESS envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)`` (the cap alone when the upstream omits ``expires_in``), matching the 1h lifetime of the BYOK session bearer this module's signing approach is borrowed from: a client-held credential should never outlive a bounded window even when the upstream token does.""" +MAX_REFRESH_ENVELOPE_TTL_SECONDS = 1209600 +"""Hard ceiling on REFRESH envelope lifetime (14 days). A refresh envelope only renews the short-lived +access envelope, and each renewal re-validates the sealed litellm key (revocation gates it) and is +re-minted with a fresh window, so the practical bound is idle time, not a fixed session. ``exp`` is +``min(upstream refresh_expires_in, this cap)`` (the cap alone when the upstream omits it); if the +upstream refresh token dies first, the next renewal simply fails at the upstream and the client +re-authenticates. The value is deliberately far shorter than a typical upstream refresh-token lifetime +so a leaked refresh envelope is bounded even if the upstream would have honoured it for longer.""" + MAX_ENVELOPE_BYTES = 12288 """Size cap on the final serialized envelope (prefix + JWT, in bytes). Upstream JWTs commonly run 2-4KB; base64 plus encryption overhead roughly doubles that inside the @@ -66,21 +81,48 @@ typed error, never truncated.""" _ENVELOPE_JWT_ALGORITHM = "HS256" +EnvelopeKind = Literal["access", "refresh"] +"""Which credential an envelope is. Stamped into the signed claims and required to match on open, so a +signature-valid envelope of one kind cannot be replayed as the other even if its wire prefix is swapped +(the prefix is not part of the signed payload; this claim is).""" + + +EnvelopeSubjectType: TypeAlias = Literal["key_hash", "user_id"] +"""Discriminator for what litellm principal the envelope binds the grant to. + +``key_hash`` is a hashed virtual key (the scripted two-header client mints under the key it +presents at the token endpoint); ``user_id`` is a litellm user subject (the interactive DCR +client mints under the SSO-authenticated user, which is the only identity that browser login +yields). Admission reloads a key record for the first and a user record for the second, then +runs both through the same live-policy gate, so team/org/budget/revocation enforcement is +identical either way.""" + class EnvelopeIdentity(BaseModel): - """The litellm identity the envelope binds the inner grant to. + """The litellm principal the envelope binds the inner grant to. - ``key_hash`` is the hashed litellm key that authorized the mint, never a raw - credential (and the edge rejects a bare hash presented as a bearer). Admission - reloads the live key record by it, so the key's current team/org/object-permission - restrictions and its revocation state are enforced at use time rather than frozen at - mint time. ``server_id`` binds the envelope to one MCP server so it cannot be replayed - across a server boundary. + ``subject`` is the principal identifier and ``subject_type`` says how to resolve it: a + hashed litellm key (``key_hash``) or a litellm user id (``user_id``), never a raw + credential (and the edge rejects a bare hash or id presented as a bearer). Admission + reloads the live record by it, so the principal's current team/org restrictions and its + revocation state are enforced at use time rather than frozen at mint time. ``server_id`` + binds the envelope to one MCP server so it cannot be replayed across a server boundary. """ model_config = ConfigDict(frozen=True) server_id: str = Field(min_length=1) - key_hash: str = Field(min_length=1) + subject_type: EnvelopeSubjectType + subject: str = Field(min_length=1) + + +def key_hash_identity(server_id: str, key_hash: str) -> EnvelopeIdentity: + """The identity for the scripted client that mints under a presented virtual key.""" + return EnvelopeIdentity(server_id=server_id, subject_type="key_hash", subject=key_hash) + + +def user_identity(server_id: str, user_id: str) -> EnvelopeIdentity: + """The identity for the interactive DCR client that mints under its SSO user subject.""" + return EnvelopeIdentity(server_id=server_id, subject_type="user_id", subject=user_id) class UpstreamTokenGrant(BaseModel): @@ -99,6 +141,21 @@ class UpstreamTokenGrant(BaseModel): expires_in: int | None = Field(default=None, gt=0) +class RefreshCredential(BaseModel): + """The upstream refresh grant sealed inside a refresh envelope. + + Only the refresh token (plus the scope to re-request and the refresh token's own lifetime, when the + upstream reports it) is sealed; the access token is never in a refresh envelope. ``refresh_token`` is + a ``SecretStr`` so reprs never leak it, and ``expires_in`` (the refresh token's lifetime, not the + access token's) must be positive when present. + """ + + model_config = ConfigDict(frozen=True) + refresh_token: SecretStr = Field(min_length=1) + scope: str | None = None + expires_in: int | None = Field(default=None, gt=0) + + class EnvelopeKeys(BaseModel): """Injected key material: the HS256 signing key and the symmetric encryption key. @@ -121,13 +178,21 @@ class SealedEnvelope(BaseModel): class OpenedEnvelope(BaseModel): - """A validated envelope: the identity it was minted for and the recovered grant.""" + """A validated access envelope: the identity it was minted for and the recovered grant.""" model_config = ConfigDict(frozen=True) identity: EnvelopeIdentity grant: UpstreamTokenGrant +class OpenedRefreshEnvelope(BaseModel): + """A validated refresh envelope: the identity it was minted for and the recovered refresh grant.""" + + model_config = ConfigDict(frozen=True) + identity: EnvelopeIdentity + refresh: RefreshCredential + + class EnvelopeTooLarge(BaseModel): """The serialized envelope exceeded ``MAX_ENVELOPE_BYTES``; carries sizes only.""" @@ -199,8 +264,10 @@ class _EnvelopeClaims(BaseModel): iss: str iat: int exp: int + kind: EnvelopeKind server_id: str = Field(min_length=1) - key_hash: str = Field(min_length=1) + subject_type: EnvelopeSubjectType + subject: str = Field(min_length=1) grant: str = Field(min_length=1) @@ -213,11 +280,25 @@ class _GrantWire(BaseModel): expires_in: int | None = None +class _RefreshWire(BaseModel): + model_config = ConfigDict(frozen=True) + refresh_token: str + scope: str | None = None + expires_in: int | None = None + + def is_envelope(candidate: str) -> bool: - """Cheap prefix check so the edge can route envelopes vs raw tokens without crypto.""" + """Cheap prefix check for an ACCESS envelope so the edge can route envelopes vs raw tokens without + crypto. A refresh envelope has a different prefix and is not an access envelope.""" return candidate.startswith(ENVELOPE_PREFIX) +def is_refresh_envelope(candidate: str) -> bool: + """Cheap prefix check for a REFRESH envelope so the token endpoint can route a refresh grant that + carries an envelope vs a raw upstream refresh token without crypto.""" + return candidate.startswith(REFRESH_ENVELOPE_PREFIX) + + def mint_envelope( identity: EnvelopeIdentity, grant: UpstreamTokenGrant, @@ -231,23 +312,15 @@ def mint_envelope( serialized envelope exceeds ``MAX_ENVELOPE_BYTES``. """ expires_at = now + timedelta(seconds=_envelope_ttl_seconds(grant.expires_in)) - claims = _EnvelopeClaims( - iss=ENVELOPE_ISSUER, - iat=int(now.timestamp()), - exp=int(expires_at.timestamp()), - server_id=identity.server_id, - key_hash=identity.key_hash, - grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key), + return _seal( + kind="access", + prefix=ENVELOPE_PREFIX, + identity=identity, + grant_blob=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key), + expires_at=expires_at, + signing_key=keys.signing_key, + now=now, ) - token = ENVELOPE_PREFIX + jwt.encode( - claims.model_dump(), - keys.signing_key.get_secret_value(), - algorithm=_ENVELOPE_JWT_ALGORITHM, - ) - size_bytes = len(token.encode("utf-8")) - if size_bytes > MAX_ENVELOPE_BYTES: - return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES) - return SealedEnvelope(token=SecretStr(token), expires_at=expires_at) def open_envelope( @@ -263,35 +336,136 @@ def open_envelope( re-derived, so it is stale by up to the envelope's lifetime; callers that need a live remaining lifetime should use ``now`` against the upstream, not this field. """ - if not is_envelope(candidate): - return NotAnEnvelope() - # UTF-8 byte length is never below character length, so a character count already over the - # cap rejects an oversize candidate in O(1) without encoding it; the exact byte check then - # runs only on candidates already bounded to <= MAX_ENVELOPE_BYTES characters. - if len(candidate) > MAX_ENVELOPE_BYTES: - return MalformedPayload() - if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES: - return MalformedPayload() - claims = _decode_claims(candidate.removeprefix(ENVELOPE_PREFIX), keys.signing_key) + claims = _open_claims(candidate, prefix=ENVELOPE_PREFIX, expected_kind="access", keys=keys, now=now) if not isinstance(claims, _EnvelopeClaims): return claims - if now.timestamp() >= claims.exp: - return Expired() grant = _decrypt_grant(claims.grant, keys.encryption_key) if not isinstance(grant, UpstreamTokenGrant): return grant return OpenedEnvelope( - identity=EnvelopeIdentity(server_id=claims.server_id, key_hash=claims.key_hash), + identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject), grant=grant, ) +def mint_refresh_envelope( + identity: EnvelopeIdentity, + refresh: RefreshCredential, + keys: EnvelopeKeys, + now: datetime, +) -> SealedEnvelope | EnvelopeMintError: + """Seal ``refresh`` for ``identity`` into a long-lived, client-held refresh envelope. + + ``exp`` is ``min(refresh.expires_in, MAX_REFRESH_ENVELOPE_TTL_SECONDS)`` seconds from ``now`` (the + cap alone when the upstream omits the refresh lifetime). Sealing a distinct ``kind="refresh"`` claim + is what keeps a refresh envelope from ever opening as an access credential at the MCP edge. Returns + ``EnvelopeTooLarge`` when the serialized envelope exceeds ``MAX_ENVELOPE_BYTES``. + """ + expires_at = now + timedelta(seconds=_refresh_ttl_seconds(refresh.expires_in)) + return _seal( + kind="refresh", + prefix=REFRESH_ENVELOPE_PREFIX, + identity=identity, + grant_blob=_encrypt_grant_blob(_refresh_plaintext(refresh), keys.encryption_key), + expires_at=expires_at, + signing_key=keys.signing_key, + now=now, + ) + + +def open_refresh_envelope( + candidate: str, + keys: EnvelopeKeys, + now: datetime, +) -> OpenedRefreshEnvelope | EnvelopeOpenError: + """Validate a refresh ``candidate`` and recover the identity and inner refresh grant. + + Total over hostile input exactly like :func:`open_envelope`: every invalid, expired, tampered, + wrong-kind, or undecryptable candidate maps to a distinct ``EnvelopeOpenError`` variant, never a + raise. The ``kind="refresh"`` claim is required, so an access envelope re-prefixed as a refresh one + is rejected as ``MalformedPayload``. + """ + claims = _open_claims(candidate, prefix=REFRESH_ENVELOPE_PREFIX, expected_kind="refresh", keys=keys, now=now) + if not isinstance(claims, _EnvelopeClaims): + return claims + refresh = _decrypt_refresh(claims.grant, keys.encryption_key) + if not isinstance(refresh, RefreshCredential): + return refresh + return OpenedRefreshEnvelope( + identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject), + refresh=refresh, + ) + + +def _seal( + kind: EnvelopeKind, + prefix: str, + identity: EnvelopeIdentity, + grant_blob: str, + expires_at: datetime, + signing_key: SecretStr, + now: datetime, +) -> SealedEnvelope | EnvelopeTooLarge: + """Sign the claims for either envelope kind and enforce the size cap. Shared by both mints so the + JWT shape, issuer, and size guard cannot drift between access and refresh envelopes.""" + claims = _EnvelopeClaims( + iss=ENVELOPE_ISSUER, + iat=int(now.timestamp()), + exp=int(expires_at.timestamp()), + kind=kind, + server_id=identity.server_id, + subject_type=identity.subject_type, + subject=identity.subject, + grant=grant_blob, + ) + token = prefix + jwt.encode(claims.model_dump(), signing_key.get_secret_value(), algorithm=_ENVELOPE_JWT_ALGORITHM) + size_bytes = len(token.encode("utf-8")) + if size_bytes > MAX_ENVELOPE_BYTES: + return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES) + return SealedEnvelope(token=SecretStr(token), expires_at=expires_at) + + +def _open_claims( + candidate: str, + prefix: str, + expected_kind: EnvelopeKind, + keys: EnvelopeKeys, + now: datetime, +) -> _EnvelopeClaims | EnvelopeOpenError: + """Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an attacker-controlled + candidate, shared by both openers so the security gate is identical for access and refresh. Returns + the validated claims or a distinct ``EnvelopeOpenError``; never raises.""" + if not candidate.startswith(prefix): + return NotAnEnvelope() + # UTF-8 byte length is never below character length, so a character count already over the cap + # rejects an oversize candidate in O(1) without encoding it; the exact byte check then runs only on + # candidates already bounded to <= MAX_ENVELOPE_BYTES characters. + if len(candidate) > MAX_ENVELOPE_BYTES: + return MalformedPayload() + if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES: + return MalformedPayload() + claims = _decode_claims(candidate.removeprefix(prefix), keys.signing_key) + if not isinstance(claims, _EnvelopeClaims): + return claims + if claims.kind != expected_kind: + return MalformedPayload() + if now.timestamp() >= claims.exp: + return Expired() + return claims + + def _envelope_ttl_seconds(upstream_expires_in: int | None) -> int: if upstream_expires_in is None: return MAX_ENVELOPE_TTL_SECONDS return min(upstream_expires_in, MAX_ENVELOPE_TTL_SECONDS) +def _refresh_ttl_seconds(upstream_refresh_expires_in: int | None) -> int: + if upstream_refresh_expires_in is None: + return MAX_REFRESH_ENVELOPE_TTL_SECONDS + return min(upstream_refresh_expires_in, MAX_REFRESH_ENVELOPE_TTL_SECONDS) + + def _grant_plaintext(grant: UpstreamTokenGrant) -> str: wire = _GrantWire( access_token=grant.access_token.get_secret_value(), @@ -303,6 +477,15 @@ def _grant_plaintext(grant: UpstreamTokenGrant) -> str: return wire.model_dump_json(exclude_none=True) +def _refresh_plaintext(refresh: RefreshCredential) -> str: + wire = _RefreshWire( + refresh_token=refresh.refresh_token.get_secret_value(), + scope=refresh.scope, + expires_in=refresh.expires_in, + ) + return wire.model_dump_json(exclude_none=True) + + def _decode_claims( compact: str, signing_key: SecretStr, @@ -364,3 +547,22 @@ def _decrypt_grant( return UpstreamTokenGrant.model_validate_json(plaintext) except ValidationError: return MalformedPayload() + + +def _decrypt_refresh( + blob: str, + encryption_key: SecretStr, +) -> RefreshCredential | DecryptFailed | MalformedPayload: + from nacl.exceptions import CryptoError + + try: + plaintext = decrypt_value( + value=base64.urlsafe_b64decode(blob), + signing_key=encryption_key.get_secret_value(), + ) + except (CryptoError, ValueError): + return DecryptFailed() + try: + return RefreshCredential.model_validate_json(plaintext) + except ValidationError: + return MalformedPayload() diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 23fe7730c17..008fdd0d50b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1118,6 +1118,7 @@ class GenerateKeyRequest(KeyRequestBase): class GenerateKeyResponse(KeyRequestBase): key: str # type: ignore key_name: Optional[str] = None + key_type: str | None = None expires: Optional[datetime] = None user_id: Optional[str] = None token_id: Optional[str] = None diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 3a93896a206..e4c565e4464 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -8,6 +8,10 @@ from litellm.proxy._types import ( ) from litellm.secret_managers.main import str_to_bool +# Bounds the __cause__/__context__ walk in is_database_service_unavailable_error_in_chain. +# Real exception chains are a few links deep; the cap also makes the walk cycle-safe. +_MAX_EXCEPTION_CHAIN_DEPTH = 20 + class PrismaDBExceptionHandler: """ @@ -218,6 +222,32 @@ class PrismaDBExceptionHandler: ), ) + @staticmethod + def is_database_service_unavailable_error_in_chain(e: BaseException) -> bool: + """Like ``is_database_service_unavailable_error`` but also walks the + ``__cause__`` / ``__context__`` chain. + + ``is_database_service_unavailable_error`` classifies a single exception + by type, which a caller that catches a raw DB failure and re-raises a + domain exception of a different type defeats. ``get_user_object`` in + ``litellm/proxy/auth/auth_checks.py`` is the concrete case: it wraps + every DB error, a genuine outage included, in a bare ``ValueError`` + whose original error survives only as ``__context__``. A type check on + the ``ValueError`` misses the outage, so the caller would mistake an + infrastructure fault for an auth failure. Walking the chain recovers the + real signal, which is the PEP 3134 way to inspect a wrapped cause. + + The walk is depth-bounded, which also makes it cycle-safe. + """ + current: BaseException | None = e + for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH): + if not isinstance(current, Exception): + return False + if PrismaDBExceptionHandler.is_database_service_unavailable_error(current): + return True + current = current.__cause__ or current.__context__ + return False + @staticmethod def handle_db_exception(e: Exception): """ diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b128b0ea57e..a3679b84bd6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -468,7 +468,10 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: Handle the key type. """ key_type = data.key_type - data_json.pop("key_type", None) + if key_type is None: + data_json.pop("key_type", None) + return data_json + data_json["key_type"] = key_type.value if key_type == LiteLLMKeyType.LLM_API: data_json["allowed_routes"] = ["llm_api_routes"] elif key_type == LiteLLMKeyType.MANAGEMENT: @@ -3566,6 +3569,7 @@ async def generate_key_helper_fn( created_by: Optional[str] = None, updated_by: Optional[str] = None, allowed_routes: Optional[list] = None, + key_type: str | None = None, sso_user_id: Optional[str] = None, object_permission_id: Optional[str] = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable object_permission: Optional[LiteLLM_ObjectPermissionBase] = None, @@ -3706,6 +3710,7 @@ async def generate_key_helper_fn( "created_by": created_by, "updated_by": updated_by, "allowed_routes": allowed_routes or [], + "key_type": key_type, "object_permission_id": object_permission_id, "router_settings": router_settings_json, "access_group_ids": access_group_ids or [], diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index fb4d8d0b5a3..a23cecc3911 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -422,6 +422,7 @@ model LiteLLM_VerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") @@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") diff --git a/litellm/router.py b/litellm/router.py index 6539d3c0c43..6e8127110cc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7605,13 +7605,13 @@ class Router: model_name = entry.get("model_name") if isinstance(entry, dict) else entry.model_name if not model_name or not lp: continue - if model_name in self.adaptive_routers: - continue deployment = Deployment( model_name=model_name, litellm_params=(lp if not isinstance(lp, dict) else LiteLLM_Params(**lp)), model_info=(entry.get("model_info") if isinstance(entry, dict) else entry.model_info), ) + if model_name in self.adaptive_routers: + continue self.init_adaptive_router_deployment(deployment=deployment) for model_name, complexity_router in self.complexity_routers.items(): @@ -10707,56 +10707,39 @@ class Router: if self.routing_plugins: await self._run_routing_plugins(model=model, request_kwargs=request_kwargs, messages=messages) - ######################################################### - # Check if any auto-router should be used - ######################################################### - if model in self.auto_routers: - return await self.auto_routers[model].async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) + router_strategy = ( + self.auto_routers.get(model) + or self.complexity_routers.get(model) + or self.adaptive_routers.get(model) + or self.quality_routers.get(model) + ) + if router_strategy is None: + return None - ######################################################### - # Check if any complexity-router should be used - ######################################################### - if model in self.complexity_routers: - return await self.complexity_routers[model].async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) + pre_routing_hook_response = await router_strategy.async_pre_routing_hook( + model=model, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) - ######################################################### - # Check if an adaptive-router should be used - ######################################################### - adaptive_router = self.adaptive_routers.get(model) - if adaptive_router is not None: - return await adaptive_router.async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) + # `model` (the alias, e.g. "smart-router") is never the deployment actually + # called - apply the alias's own litellm_params (besides `model` itself, + # which is just the alias marker) to the request, since the tier/route + # deployment the hook selected won't have them. Router-only fields + # (tpm, rpm, weight, complexity_router_config, ...) are excluded from the + # actual outbound LLM call downstream by litellm.types.utils.all_litellm_params, + # not here. + if pre_routing_hook_response is not None: + alias_index = self.model_name_to_deployment_indices.get(model, []) + if alias_index: + alias_litellm_params = self.model_list[alias_index[0]].get("litellm_params", {}) + for key, value in alias_litellm_params.items(): + if key != "model" and value is not None: + request_kwargs.setdefault(key, value) - ######################################################### - # Check if any quality-router should be used - ######################################################### - if model in self.quality_routers: - return await self.quality_routers[model].async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) - - return None + return pre_routing_hook_response def get_available_deployment( self, diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index bebdbba90ef..bd3b300b558 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -18,7 +18,7 @@ from __future__ import annotations import asyncio import random import re -from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Literal, Union, cast from pydantic import BaseModel @@ -809,6 +809,44 @@ class ComplexityRouter(CustomLogger): return user_message, system_prompt + @staticmethod + def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: + """Metadata may land on `metadata` or `litellm_metadata` depending on the + endpoint, mirroring DeploymentAffinityCheck's precedence.""" + return [ + metadata + for metadata_key in ("litellm_metadata", "metadata") + if isinstance(metadata := request_kwargs.get(metadata_key), dict) + ] + + @staticmethod + def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: + """Resolve a client-supplied session_id.""" + for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): + session_id = metadata.get("session_id") + if session_id is not None: + return str(session_id) + return None + + @staticmethod + def _get_user_api_key_hash_from_request_kwargs(request_kwargs: dict) -> str | None: + """Resolve the proxy-derived API key hash, the same trust boundary + DeploymentAffinityCheck uses for its own key-based affinity (not the + client-supplied OpenAI `user` param, which isn't authenticated).""" + for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): + user_key = metadata.get("user_api_key_hash") + if user_key is not None: + return str(user_key) + return None + + def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) -> str: + # Namespace by the caller's API key hash so two different callers reusing the + # same client-supplied session_id can't poison each other's routing pin. Falls + # back to "unscoped" only when there's no authenticated caller to scope by + # (e.g. direct Router usage without the proxy layer). + caller_scope = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" + return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" + async def async_pre_routing_hook( self, model: str, @@ -816,10 +854,70 @@ class ComplexityRouter(CustomLogger): messages: list[dict[str, Any]] | None = None, input: Union[str, list] | None = None, specific_deployment: bool | None = False, - ) -> Optional[PreRoutingHookResponse]: + ) -> PreRoutingHookResponse | None: """ Pre-routing hook called before the routing decision. + When `session_affinity` is enabled and a session_id is resolvable on the request, + pins the model chosen on the session's first turn and reuses it for every later + turn, skipping classification entirely. Otherwise delegates to `_classify_and_route`. + """ + from litellm.types.router import PreRoutingHookResponse + + session_id = self._get_session_id_from_request_kwargs(request_kwargs) if self.config.session_affinity else None + cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None + + if cache_key is not None: + pinned_model = await self.litellm_router_instance.cache.async_get_cache(key=cache_key) + if isinstance(pinned_model, str): + # Refresh the TTL on every hit so an active session doesn't lose its + # pin mid-conversation just because it outlives the original write. + await self.litellm_router_instance.cache.async_set_cache( + key=cache_key, + value=pinned_model, + ttl=self.config.session_affinity_ttl_seconds, + ) + if self.config.adaptive: + from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + ) + + kwargs_metadata = request_kwargs.setdefault("metadata", {}) + if isinstance(kwargs_metadata, dict): + kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = pinned_model + verbose_router_logger.info( + f"ComplexityRouter: routing decision cause=session_affinity_pin, routed_model={pinned_model}" + ) + has_original_messages = messages is not None and len(messages) > 0 + return PreRoutingHookResponse( + model=pinned_model, + messages=messages if has_original_messages else None, + ) + + response = await self._classify_and_route( + model=model, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) + if cache_key is not None and response is not None: + await self.litellm_router_instance.cache.async_set_cache( + key=cache_key, + value=response.model, + ttl=self.config.session_affinity_ttl_seconds, + ) + return response + + async def _classify_and_route( + self, + model: str, + request_kwargs: dict, + messages: list[dict[str, Any]] | None = None, + input: Union[str, list] | None = None, + specific_deployment: bool | None = False, + ) -> PreRoutingHookResponse | None: + """ Classifies the request by complexity and returns the appropriate model. Supports chat completions (messages), Responses API (input), and other formats via the guardrail translation handler dispatch. diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index df699d1a059..e4bd36505e6 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -361,6 +361,20 @@ class ComplexityRouterConfig(BaseModel): description="Minimum cosine similarity for a semantic keyword match", ) + # Session affinity: pin the first turn's routed model for the rest of the session + session_affinity: bool = Field( + default=False, + description=( + "When True and a session_id is resolvable on the request, pin the model chosen on the " + "session's first turn and reuse it for every later turn, skipping re-classification." + ), + ) + session_affinity_ttl_seconds: int = Field( + default=3600, + gt=0, + description="TTL for the session affinity pin; refreshed on every cache hit", + ) + model_config = ConfigDict(extra="allow") # Allow additional fields @field_validator("tiers", mode="before") diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 69bccff701f..318ba1f5956 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -213,6 +213,8 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_input_audio_tokens_metric", "litellm_output_reasoning_tokens_metric", "litellm_output_audio_tokens_metric", + "litellm_video_duration_seconds_metric", + "litellm_images_generated_metric", "litellm_deployment_successful_fallbacks", "litellm_deployment_failed_fallbacks", "litellm_remaining_team_budget_metric", @@ -506,6 +508,9 @@ class PrometheusMetricLabels: litellm_output_reasoning_tokens_metric = litellm_output_tokens_metric litellm_output_audio_tokens_metric = litellm_output_tokens_metric + litellm_video_duration_seconds_metric = litellm_output_tokens_metric + litellm_images_generated_metric = litellm_output_tokens_metric + litellm_deployment_state = [ UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.MODEL_ID.value, @@ -717,6 +722,8 @@ class PrometheusMetricLabels: "litellm_input_tokens_metric", "litellm_total_tokens_metric", "litellm_output_tokens_metric", + "litellm_video_duration_seconds_metric", + "litellm_images_generated_metric", } ) # Managed batch metrics diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e33e2335525..90ea99ceb23 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3210,6 +3210,16 @@ all_litellm_params = ( "_litellm_tpm_reserved_model", "_litellm_tpm_reserved_scopes", "_litellm_tpm_reservation_released", + "auto_router_config_path", + "auto_router_config", + "auto_router_default_model", + "auto_router_embedding_model", + "complexity_router_config", + "complexity_router_default_model", + "adaptive_router_config", + "adaptive_router_default_model", + "quality_router_config", + "quality_router_default_model", ] + list(StandardCallbackDynamicParams.__annotations__.keys()) + list(CustomPricingLiteLLMParams.model_fields.keys()) diff --git a/schema.prisma b/schema.prisma index fb4d8d0b5a3..a23cecc3911 100644 --- a/schema.prisma +++ b/schema.prisma @@ -422,6 +422,7 @@ model LiteLLM_VerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") @@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index 0cfbb5b0b66..12f47331dbb 100644 --- a/tests/e2e/docker-compose.yml +++ b/tests/e2e/docker-compose.yml @@ -18,6 +18,12 @@ configs: type: redis host: redis port: 6379 + # OTEL v2 trace destination for the logging suite's trace-completeness + # tests: the arize_phoenix preset is OTLP with a configurable endpoint + # (PHOENIX_COLLECTOR_HTTP_ENDPOINT below points it at the jaeger service), + # so gen-AI spans export through a preset-owned provider - the code path + # where trace splits actually happen - with no cloud credentials needed. + callbacks: ["arize_phoenix"] router_settings: routing_strategy: simple-shuffle @@ -68,9 +74,14 @@ services: condition: service_healthy redis: condition: service_healthy + jaeger: + condition: service_healthy env_file: .env environment: LITELLM_MASTER_KEY: sk-1234 + LITELLM_OTEL_V2: "true" + PHOENIX_COLLECTOR_HTTP_ENDPOINT: http://jaeger:4318/v1/traces + PHOENIX_API_KEY: local-jaeger-noauth DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm UI_USERNAME: admin UI_PASSWORD: sk-1234 @@ -114,3 +125,15 @@ services: interval: 3s timeout: 3s retries: 20 + +# throwaway OTEL trace destination (OTLP ingest on 4318 inside the network, +# query API on host 16686 for test read-back; see E2E_OTEL_QUERY_URL) + jaeger: + image: jaegertracing/all-in-one:1.62.0 + ports: + - "16686:16686" + healthcheck: + test: ["CMD", "wget", "-qO-", "http://localhost:14269/"] + interval: 3s + timeout: 3s + retries: 20 diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 75bd715a23a..6bfec2b514e 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -24,6 +24,13 @@ CONTROL_PLANE_BASE_URL = os.environ.get( UI_USERNAME = os.environ.get("E2E_UI_USERNAME", "admin") UI_PASSWORD = os.environ.get("E2E_UI_PASSWORD", MASTER_KEY) +CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5") + +# Jaeger query API of the compose stack's OTEL trace destination (the `jaeger` +# service in docker-compose.yml maps it to host 16686). Trace-completeness tests +# read exported spans back through it. +OTEL_QUERY_URL = os.environ.get("E2E_OTEL_QUERY_URL", "http://localhost:16686").rstrip("/") + # Writes on the proxy are eventually consistent (e.g. spend rows flush on # proxy_batch_write_at, ~60s). Read-backs poll to this deadline, never sleep-once. POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) diff --git a/tests/e2e/logging/conftest.py b/tests/e2e/logging/conftest.py index 60492c9f0bd..22c70eea0dd 100644 --- a/tests/e2e/logging/conftest.py +++ b/tests/e2e/logging/conftest.py @@ -11,14 +11,8 @@ from collections.abc import Iterator import pytest -from logging_client import ( - LangfuseCreds, - LoggingClient, - PhoenixCreds, - build_logging_client, - load_langfuse_creds, - load_phoenix_creds, -) +from logging_client import LangfuseCreds, LoggingClient, build_logging_client, load_langfuse_creds +from otel_client import OtelReader, build_otel_reader def pytest_configure(config: pytest.Config) -> None: @@ -39,6 +33,12 @@ def client() -> Iterator[LoggingClient]: model_cleanup.teardown() +@pytest.fixture(scope="session") +def otel_reader() -> OtelReader: + """Read-back client for the compose stack's Jaeger trace destination.""" + return build_otel_reader() + + @pytest.fixture def datadog_creds() -> None: """Require Datadog shipping credentials. Hard-fail when absent; never skip.""" diff --git a/tests/e2e/logging/otel_client.py b/tests/e2e/logging/otel_client.py new file mode 100644 index 00000000000..f4a0e4fe102 --- /dev/null +++ b/tests/e2e/logging/otel_client.py @@ -0,0 +1,138 @@ +"""Jaeger read-back for the OTEL trace-completeness tests: typed models over the +Jaeger query API (the destination's own API - completeness is judged on what the +backend actually holds, never on "export succeeded" proxy-side). + +Traces are fetched server-side by the ``litellm.call_id`` tag the gen-AI span +carries (the request's x-litellm-call-id response header), so read-back is +immune to the query page filling up with unrelated traffic (background jobs, +other suites sharing the stack). Jaeger returns every span of a matching trace, +so the completeness assertions see the whole tree. A failed query is a hard +failure, never an empty result - an unreachable destination must not read as +"the trace never arrived". + +External reads go through ``e2e_http`` (the only module allowed to call +``requests.*``). +""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass + +import pytest +from pydantic import BaseModel, ConfigDict, Field + +from e2e_config import OTEL_QUERY_URL, POLL_INTERVAL, POLL_TIMEOUT +from e2e_http import URL, NoBody, Success, get + +#: OTEL resource service.name the proxy exports under (OTEL_SERVICE_NAME default). +JAEGER_SERVICE = "litellm" +#: Span tag carrying the request's x-litellm-call-id (stamped on the gen-AI span). +CALL_ID_TAG = "litellm.call_id" + + +class JaegerTag(BaseModel): + model_config = ConfigDict(extra="ignore") + + key: str + value: str | int | float | bool | None = None + + +class JaegerReference(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + ref_type: str = Field(alias="refType") + trace_id: str = Field(alias="traceID") + span_id: str = Field(alias="spanID") + + +class JaegerSpan(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + span_id: str = Field(alias="spanID") + operation_name: str = Field(alias="operationName") + start_time: int = Field(default=0, alias="startTime") + references: list[JaegerReference] = [] + tags: list[JaegerTag] = [] + + @property + def kind(self) -> str: + for tag in self.tags: + if tag.key == "span.kind": + return str(tag.value) + return "" + + +class JaegerTrace(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + trace_id: str = Field(alias="traceID") + spans: list[JaegerSpan] = [] + + def span_names(self) -> list[str]: + return sorted(span.operation_name for span in self.spans) + + +class JaegerTracesPage(BaseModel): + model_config = ConfigDict(extra="ignore") + + data: list[JaegerTrace] = [] + + +class _TracesQuery(BaseModel): + service: str + tags: str + limit: int = 20 + lookback: str = "1h" + + +def _settled(trace: JaegerTrace, names: set[str], prefixes: set[str]) -> bool: + present = set(trace.span_names()) + return names.issubset(present) and all( + any(name.startswith(prefix) for name in present) for prefix in prefixes + ) + + +@dataclass(frozen=True, slots=True) +class OtelReader: + query_url: str + + def traces_for_call(self, call_id: str) -> list[JaegerTrace]: + """Every trace holding a span tagged with this call id. Jaeger matches + spans server-side and returns their full traces; more than one hit for + one call IS the split-trace bug, so this never collapses to one.""" + result = get( + URL(f"{self.query_url}/api/traces"), + headers=NoBody(), + params=_TracesQuery(service=JAEGER_SERVICE, tags=json.dumps({CALL_ID_TAG: call_id})), + response_type=JaegerTracesPage, + timeout=30.0, + ) + match result: + case Success(data=page): + return page.data + case failure: + pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}") + + def poll_traces_for_call( + self, *, call_id: str, settled_names: set[str], settled_prefixes: set[str] + ) -> list[JaegerTrace]: + """Poll until exactly one trace holds the call and it carries every span + name in ``settled_names`` plus at least one name per prefix in + ``settled_prefixes`` (spans flush in batches, the cost write lands after + the response), then return the hits. At the deadline the last hits are + returned as-is so the caller's assertions report the real final state - + on a split trace this never settles and the orphan comes back.""" + deadline = time.monotonic() + POLL_TIMEOUT + hits: list[JaegerTrace] = [] + while time.monotonic() < deadline: + hits = self.traces_for_call(call_id) + if len(hits) == 1 and _settled(hits[0], settled_names, settled_prefixes): + return hits + time.sleep(POLL_INTERVAL) + return hits + + +def build_otel_reader() -> OtelReader: + return OtelReader(query_url=OTEL_QUERY_URL) diff --git a/tests/e2e/logging/test_otel_trace_e2e.py b/tests/e2e/logging/test_otel_trace_e2e.py new file mode 100644 index 00000000000..d445da3dff0 --- /dev/null +++ b/tests/e2e/logging/test_otel_trace_e2e.py @@ -0,0 +1,191 @@ +"""Live e2e: OTEL trace completeness on the admin-owned destination (LIT-3787). + +Covers logging.otel.success.exports_metric: a successful non-streaming call must +land at the OTEL destination as ONE connected trace - a single root SERVER span +with the auth phase, db lookups, and cost write under it, and the gen-AI CLIENT +span parented into the same tree. The regression this pins: the proxy publishing +the global TracerProvider before callbacks init made server spans export through +a different provider than the preset's gen-AI spans, so the destination received +the gen-AI span alone, dangling (fixed in #30590; verified failing at its parent +commit 1bd603d1ac). + +Both halves of the contract are asserted: the recorded state (the proxy reports +the OTEL v2 logger active via /health/readiness/details) and the enforced +behavior (the complete span tree at the destination, read back through the +destination's own query API - never proxy-side "export succeeded" logs). +""" + +from __future__ import annotations + +import time +from collections.abc import Callable + +import pytest +from pydantic import BaseModel, ConfigDict + +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from e2e_http import NoBody, StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from logging_client import LoggingClient +from otel_client import JaegerTrace, OtelReader + +pytestmark = pytest.mark.e2e + +MODEL = CHEAP_ANTHROPIC_MODEL +COST_SPAN = "batch_write_to_db _PROXY_track_cost_callback" +DB_SPAN_PREFIX = "postgres " +#: The active OTEL v2 logger's name in /health/readiness/details success_callbacks. +OTEL_V2_LOGGER_NAME = "OpenTelemetryV2" + + +class _ReadinessDetails(BaseModel): + model_config = ConfigDict(extra="ignore") + + success_callbacks: list[str] = [] + + +def _assert_otel_destination_configured(client: LoggingClient) -> None: + """Recorded state: the proxy reports the OTEL v2 logger among its active + callbacks, so a missing/failed destination config fails here, before any + traffic-based assertion can time out confusingly.""" + result = client.gateway.probe("/health/readiness/details", params=NoBody()) + assert result.status_code == 200, ( + f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}" + ) + details = _ReadinessDetails.model_validate_json(result.body) + assert OTEL_V2_LOGGER_NAME in details.success_callbacks, ( + f"the proxy must report the {OTEL_V2_LOGGER_NAME} callback active " + f"(LITELLM_OTEL_V2 + arize_phoenix preset in the compose config); got: {details.success_callbacks}" + ) + + +def _first_ok(client: LoggingClient, send: Callable[[], StreamingResponse]) -> StreamingResponse: + """First successful call on a fresh key. A fresh key may briefly 401 until + the data plane's auth cache picks it up, so retry on 401 to a deadline; a + 401 is rejected before the LLM call so it exports no gen-AI span and cannot + contaminate the trace assertions. Any other failure is behavior under test + and fails hard.""" + deadline = time.monotonic() + client.gateway.poll_timeout + while True: + outcome = send() + if outcome.ok: + return outcome + if outcome.status_code != 401 or time.monotonic() >= deadline: + require_successful_call(outcome) + time.sleep(client.gateway.poll_interval) + + +def _parent_ids(span_id: str, trace: JaegerTrace) -> list[str]: + span = next(s for s in trace.spans if s.span_id == span_id) + return [ref.span_id for ref in span.references if ref.ref_type == "CHILD_OF"] + + +def _chain_reaches(span_id: str, root_id: str, trace: JaegerTrace) -> bool: + """Walk parent references (within the trace) from span_id up to root_id.""" + seen: set[str] = set() + in_trace = {s.span_id for s in trace.spans} + current = span_id + while current not in seen: + if current == root_id: + return True + seen.add(current) + parents = [p for p in _parent_ids(current, trace) if p in in_trace] + if not parents: + return False + current = parents[0] + return False + + +def _assert_complete_trace(hits: list[JaegerTrace], *, route: str, genai_span: str) -> None: + """The enforced behavior: the destination holds exactly one trace for the + call, rooted at the SERVER span, with auth/db/cost children and the gen-AI + span all connected into that one tree - no dangling parent references.""" + assert hits, ( + "no trace for this call arrived at the destination within the deadline " + "(nothing tagged with its call id was found)" + ) + assert len(hits) == 1, ( + f"expected exactly ONE trace for the call, got {len(hits)}: " + f"{[(t.trace_id, t.span_names()) for t in hits]} - more than one trace for " + "one call is the split-trace bug (gen-AI span exported away from its root)" + ) + trace = hits[0] + names = trace.span_names() + in_trace = {span.span_id for span in trace.spans} + + dangling = [ + span.operation_name + for span in trace.spans + if span.references and not any(ref.span_id in in_trace for ref in span.references) + ] + assert not dangling, ( + f"span(s) {dangling} reference a parent that never reached the destination " + f"(orphaned trace); spans present: {names}" + ) + + roots = [span for span in trace.spans if not span.references] + assert len(roots) == 1, f"expected exactly one root span, got {[s.operation_name for s in roots]}; spans: {names}" + root = roots[0] + assert root.operation_name == f"POST {route}", ( + f"the root must be the SERVER span 'POST {route}', got {root.operation_name!r}" + ) + assert root.kind == "server", f"the root span must have kind=server, got {root.kind!r}" + + assert f"auth {route}" in names, f"auth phase span 'auth {route}' missing; spans: {names}" + assert any(name.startswith(DB_SPAN_PREFIX) for name in names), ( + f"no db ('{DB_SPAN_PREFIX}*') span in the trace; spans: {names}" + ) + assert COST_SPAN in names, f"cost write span {COST_SPAN!r} missing; spans: {names}" + + genai = next((span for span in trace.spans if span.operation_name == genai_span), None) + assert genai is not None, f"gen-AI span {genai_span!r} missing; spans: {names}" + assert genai.kind == "client", f"gen-AI span must have kind=client, got {genai.kind!r}" + assert _chain_reaches(genai.span_id, root.span_id, trace), ( + f"gen-AI span {genai_span!r} is in the trace but its parent chain does not " + f"reach the root SERVER span; spans: {names}" + ) + + +def _settled_names(*, route: str, genai_span: str) -> set[str]: + return {f"POST {route}", f"auth {route}", COST_SPAN, genai_span} + + +class TestOtelTraceCompleteness: + @pytest.mark.covers("logging.otel.success.exports_metric") + def test_chat_completions_exports_complete_trace( + self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager + ) -> None: + """This test verifies that a successful non-streaming + /chat/completions request produces one complete OTEL trace. + + The trace should have a single server root span for the incoming request, with + the authentication, database, and cost-recording work beneath it. The span for + the actual model call must also belong to that same trace, rather than being + exported separately with a missing parent. + + This matters because a split trace is easy to miss: all of the spans may still + arrive, but the model call appears without the surrounding request context. + That makes it difficult to understand where time was spent, connect the model + cost to the original request, or investigate a slow or failed call. + + /chat/completions is the main OpenAI-compatible route used by most customers, + so it is important that trace parenting works correctly on this path. + """ + route = "/chat/completions" + _assert_otel_destination_configured(client) + + key = client.key_with_alias(f"otel-trace-chat-{unique_marker()}", models=[MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = _first_ok( + client, lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16) + ) + assert outcome.call_id is not None, "success response must carry x-litellm-call-id" + + hits = otel_reader.poll_traces_for_call( + call_id=outcome.call_id, + settled_names=_settled_names(route=route, genai_span=f"chat {MODEL}"), + settled_prefixes={DB_SPAN_PREFIX}, + ) + _assert_complete_trace(hits, route=route, genai_span=f"chat {MODEL}") diff --git a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py index 891e5020f37..d56d5e51b04 100644 --- a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py +++ b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py @@ -2,7 +2,7 @@ import io import os import sys -from typing import Optional +from typing import Optional, Union sys.path.insert(0, os.path.abspath("../..")) @@ -12,6 +12,7 @@ import json import logging import time from unittest.mock import AsyncMock, patch +from datetime import datetime import httpx import pytest @@ -20,17 +21,24 @@ import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.responses.main import mock_responses_api_response -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import ( + ModelResponse, + ResponsesAPIResponse, + StandardLoggingPayload, + TextCompletionResponse, +) class TestCustomLogger(CustomLogger): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None + self.response_obj: Optional[Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse]] = None async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): standard_logging_payload = kwargs.get("standard_logging_object", None) self.logged_standard_logging_payload = standard_logging_payload + self.response_obj = response_obj @pytest.mark.asyncio @@ -108,6 +116,78 @@ async def test_dynamic_turn_off_message_logging_overrides_global_off(dynamic_tur assert standard_logging_payload["messages"][0]["content"] == expected_message_content +@pytest.mark.asyncio +async def test_redaction_with_custom_logger_streaming(): + """Test redaction of responses for custom logger callbacks""" + from litellm.litellm_core_utils.litellm_logging import Logging + + class LoggingWithoutSyncSuccessHandler(Logging): + def success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs): + pass + + litellm.turn_off_message_logging = True + test_custom_logger = TestCustomLogger() + + try: + litellm_logging_obj = LoggingWithoutSyncSuccessHandler( + model="gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + litellm_call_id="1234", + start_time=datetime.now(), + function_id="1234", + dynamic_async_success_callbacks=[test_custom_logger], + ) + + response = await litellm.acompletion( + model="gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello", + stream=True, + litellm_logging_obj=litellm_logging_obj, + ) + + # Consume the stream to trigger logging + chunks = [] + async for chunk in response: + chunks.append(chunk) + + await asyncio.sleep(1) + async_complete_streaming_response = test_custom_logger.response_obj + assert async_complete_streaming_response is not None + assert async_complete_streaming_response.choices[0].message.content == "redacted-by-litellm" + finally: + litellm.turn_off_message_logging = False + + +@pytest.mark.asyncio +async def test_streaming_redaction_scoped_to_opted_out_logger(): + """One logger opting out of message logging must not blank the response for other loggers""" + litellm.turn_off_message_logging = False + opted_out_logger = TestCustomLogger(message_logging=False) + compliant_logger = TestCustomLogger() + litellm.callbacks = [opted_out_logger, compliant_logger] + + try: + response = await litellm.acompletion( + model="gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello", + stream=True, + ) + async for _ in response: + pass + + await asyncio.sleep(1) + assert opted_out_logger.response_obj is not None + assert opted_out_logger.response_obj.choices[0].message.content == "redacted-by-litellm" + assert compliant_logger.response_obj is not None + assert compliant_logger.response_obj.choices[0].message.content == "hello" + finally: + litellm.callbacks = [] + + @pytest.mark.asyncio async def test_redaction_responses_api(): """Test redaction with ResponsesAPIResponse format""" diff --git a/tests/test_litellm/integrations/test_prometheus_media_generation_metrics.py b/tests/test_litellm/integrations/test_prometheus_media_generation_metrics.py new file mode 100644 index 00000000000..2ce92f08ef0 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_media_generation_metrics.py @@ -0,0 +1,180 @@ +""" +Unit tests for the video-seconds and images-generated Prometheus counters (LIT-4254). + +Video providers report ``duration_seconds`` inside the usage object that lands +on ``standard_logging_payload["metadata"]["usage_object"]``; image generation +calls report ``output_image_count`` there. Both counters are sparse: only +incremented when the value is present and > 0. +""" + +from typing import get_args +from unittest.mock import MagicMock + +import pytest + +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import ( + DEFINED_PROMETHEUS_METRICS, + PrometheusMetricLabels, + UserAPIKeyLabelValues, +) + +MEDIA_GENERATION_METRICS = [ + "litellm_video_duration_seconds_metric", + "litellm_images_generated_metric", +] + + +@pytest.fixture +def sample_enum_values(): + return UserAPIKeyLabelValues( + end_user="test-end-user", + hashed_api_key="test-key-hash", + api_key_alias="test-key-alias", + team="test-team", + team_alias="test-team-alias", + user="test-user", + model="sora-2", + ) + + +def _make_mock_logger(): + logger = MagicMock() + for name in MEDIA_GENERATION_METRICS: + setattr(logger, name, MagicMock()) + logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + return logger + + +class TestMediaGenerationMetricsRegistration: + def test_metrics_in_defined_prometheus_metrics(self): + defined = get_args(DEFINED_PROMETHEUS_METRICS) + for name in MEDIA_GENERATION_METRICS: + assert name in defined, f"{name} missing from DEFINED_PROMETHEUS_METRICS" + + def test_metric_labels_defined(self): + for name in MEDIA_GENERATION_METRICS: + assert hasattr(PrometheusMetricLabels, name), f"{name} missing from PrometheusMetricLabels" + + def test_metrics_share_output_token_label_set(self): + assert ( + PrometheusMetricLabels.litellm_video_duration_seconds_metric + == PrometheusMetricLabels.litellm_output_tokens_metric + ) + assert ( + PrometheusMetricLabels.litellm_images_generated_metric + == PrometheusMetricLabels.litellm_output_tokens_metric + ) + + def test_runtime_label_set_matches_output_tokens_metric(self): + """Full parity with litellm_output_tokens_metric, including the org labels + appended via _org_label_metrics, so existing token dashboards can be cloned.""" + expected = PrometheusMetricLabels.get_labels("litellm_output_tokens_metric") + for name in MEDIA_GENERATION_METRICS: + assert PrometheusMetricLabels.get_labels(name) == expected + + +class TestIncrementMediaGenerationMetrics: + def test_video_duration_incremented(self, sample_enum_values): + logger = _make_mock_logger() + payload = {"metadata": {"usage_object": {"duration_seconds": 8.0}}} + + PrometheusLogger._increment_media_generation_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + logger.litellm_video_duration_seconds_metric.labels().inc.assert_called_once_with(8.0) + logger.litellm_images_generated_metric.labels.assert_not_called() + + def test_image_count_incremented(self, sample_enum_values): + logger = _make_mock_logger() + payload = { + "metadata": { + "usage_object": { + "prompt_tokens": 18, + "completion_tokens": 391, + "total_tokens": 409, + "output_image_count": 2, + } + } + } + + PrometheusLogger._increment_media_generation_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + logger.litellm_images_generated_metric.labels().inc.assert_called_once_with(2.0) + logger.litellm_video_duration_seconds_metric.labels.assert_not_called() + + def test_token_only_usage_is_a_noop(self, sample_enum_values): + logger = _make_mock_logger() + payload = { + "metadata": { + "usage_object": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30, + } + } + } + + PrometheusLogger._increment_media_generation_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + for name in MEDIA_GENERATION_METRICS: + getattr(logger, name).labels.assert_not_called() + + @pytest.mark.parametrize("bad_value", [0, 0.0, None, -4.0, "4", True]) + def test_non_positive_or_non_numeric_values_are_ignored(self, sample_enum_values, bad_value): + logger = _make_mock_logger() + payload = { + "metadata": { + "usage_object": { + "duration_seconds": bad_value, + "output_image_count": bad_value, + } + } + } + + PrometheusLogger._increment_media_generation_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + for name in MEDIA_GENERATION_METRICS: + getattr(logger, name).labels.assert_not_called() + + def test_missing_usage_object_is_a_noop(self, sample_enum_values): + logger = _make_mock_logger() + + for payload in ({"metadata": {}}, {"metadata": None}, {"metadata": {"usage_object": "redacted"}}): + PrometheusLogger._increment_media_generation_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + for name in MEDIA_GENERATION_METRICS: + getattr(logger, name).labels.assert_not_called() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py index bb035c4c3ee..9c6d2e018ff 100644 --- a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py @@ -326,3 +326,148 @@ async def test_should_leave_rate_limit_labels_blank_for_non_rate_limit_failure() assert isinstance(enum_values, UserAPIKeyLabelValues) assert enum_values.rate_limit_category is None assert enum_values.rate_limit_type is None + + +def _logger_with_mock_virtual_key_gauges() -> PrometheusLogger: + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_remaining_api_key_requests_for_model = MagicMock() + logger.litellm_remaining_api_key_tokens_for_model = MagicMock() + logger.get_labels_for_metric = MagicMock(return_value=[]) + return logger + + +def _kwargs_with_v3_rate_limit_headers(additional_headers: dict) -> dict: + return { + "litellm_params": {"metadata": {"model_group": "gpt-4o-mini"}}, + "standard_logging_object": { + "metadata": {}, + "hidden_params": {"additional_headers": additional_headers}, + }, + } + + +def _set_virtual_key_metrics(logger: PrometheusLogger, kwargs: dict) -> None: + logger._set_virtual_key_rate_limit_metrics( + user_api_key="test-hash", + user_api_key_alias="test-alias", + kwargs=kwargs, + metadata=kwargs["litellm_params"]["metadata"], + model_id="model-123", + ) + + +def test_should_read_v3_remaining_headers_when_metadata_keys_absent(): + """ + Regression for LIT-2577: the default v3 rate limiter writes remaining + per-(key, model) values into + ``standard_logging_object.hidden_params.additional_headers`` as + ``x-ratelimit-model_per_key-remaining-{requests,tokens}`` and never sets + the legacy ``litellm-key-remaining-*`` metadata keys, so the gauges were + pinned to ``sys.maxsize``. + """ + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": 42, + "x-ratelimit-model_per_key-remaining-tokens": 900, + "x-ratelimit-model_per_key-limit-requests": 100, + "x-ratelimit-model_per_key-limit-tokens": 1000, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + 42 + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + 900 + ) + + +def test_should_prefer_legacy_metadata_keys_over_v3_headers(): + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": 42, + "x-ratelimit-model_per_key-remaining-tokens": 900, + } + ) + kwargs["litellm_params"]["metadata"].update( + { + "litellm-key-remaining-requests-gpt-4o-mini": 3, + "litellm-key-remaining-tokens-gpt-4o-mini": 200, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + 3 + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + 200 + ) + + +def test_should_treat_zero_v3_remaining_as_zero(): + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": 0, + "x-ratelimit-model_per_key-remaining-tokens": 0, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + 0 + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + 0 + ) + + +def test_should_keep_maxsize_sentinel_when_no_rate_limit_source_present(): + import sys + + logger = _logger_with_mock_virtual_key_gauges() + kwargs = { + "litellm_params": {"metadata": {"model_group": "gpt-4o-mini"}}, + "standard_logging_object": {"metadata": {}, "hidden_params": {}}, + } + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) + + +@pytest.mark.parametrize("bad_value", ["not-a-number", None, True]) +def test_should_ignore_non_int_v3_header_values(bad_value): + import sys + + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": bad_value, + "x-ratelimit-model_per_key-remaining-tokens": bad_value, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 0523ed7ecb1..ade2c677745 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3707,3 +3707,69 @@ def test_set_cost_breakdown_stores_reasoning_cost(): cost_for_built_in_tools_cost_usd_dollar=0.0, ) assert "reasoning_cost" not in no_reasoning.cost_breakdown + + +def _build_payload_for_media_response(logging_obj, init_response_obj, kwargs=None): + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.datetime.now() + return get_standard_logging_object_payload( + kwargs=kwargs or {"litellm_call_id": "media-call-id", "model": "test-model", "messages": []}, + init_response_obj=init_response_obj, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + +def test_image_response_sets_output_image_count_on_usage_object(logging_obj): + """Generated-image count must land on metadata.usage_object for callbacks (e.g. Prometheus).""" + from litellm.types.utils import ImageResponse + + response = ImageResponse(created=1, data=[{"url": "https://img/1"}, {"url": "https://img/2"}]) + + payload = _build_payload_for_media_response(logging_obj, response) + + assert payload is not None + assert payload["metadata"]["usage_object"]["output_image_count"] == 2 + + +def test_output_image_count_survives_message_redaction(logging_obj, monkeypatch): + """Redaction replaces the ImageResponse body, so the count must be captured pre-redaction.""" + import litellm + from litellm.types.utils import ImageResponse + + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + response = ImageResponse(created=1, data=[{"url": "https://img/1"}]) + + payload = _build_payload_for_media_response(logging_obj, response) + + assert payload is not None + assert payload["response"] == {"text": "redacted-by-litellm"} + assert payload["metadata"]["usage_object"]["output_image_count"] == 1 + + +def test_non_image_response_has_no_output_image_count(logging_obj): + payload = _build_payload_for_media_response( + logging_obj, {"id": "chatcmpl-1", "usage": {"prompt_tokens": 1, "completion_tokens": 2}} + ) + + assert payload is not None + assert "output_image_count" not in payload["metadata"]["usage_object"] + + +def test_zero_token_video_usage_preserves_duration_seconds(logging_obj): + """Video usage bills by duration; the payload must keep duration_seconds even with zero tokens.""" + payload = _build_payload_for_media_response( + logging_obj, {"id": "video-1", "usage": {"duration_seconds": 4.0}} + ) + + assert payload is not None + assert payload["metadata"]["usage_object"]["duration_seconds"] == 4.0 + assert payload["total_tokens"] == 0 + assert payload["completion_tokens"] == 0 diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 36f220f9a2c..0f7f492ddb6 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -10,9 +10,11 @@ from types import SimpleNamespace import pytest import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.redact_messages import ( _redact_responses_api_output, perform_redaction, + redact_streaming_responses_for_custom_logger, should_redact_message_logging, ) from litellm.responses.main import mock_responses_api_response @@ -442,3 +444,109 @@ class TestPerformRedaction: assert "vertex_ai_url_context_metadata" not in hidden_params assert "vertex_ai_safety_ratings" not in hidden_params assert "vertex_ai_citation_metadata" not in hidden_params + + def test_redact_async_complete_streaming_response(self): + """Test that async_complete_streaming_response is properly redacted.""" + response_obj = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret content", role="assistant") + ) + ] + ) + + model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "prompt": "hi", + "input": "hi", + "stream": True, + "async_complete_streaming_response": response_obj, + } + + perform_redaction(model_call_details, result=None) + + redacted_response = model_call_details["async_complete_streaming_response"] + assert redacted_response.choices[0].message.content == "redacted-by-litellm" + + def test_redact_complete_streaming_response(self): + """Test that complete_streaming_response is properly redacted.""" + response_obj = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret content", role="assistant") + ) + ] + ) + + model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "prompt": "hi", + "input": "hi", + "stream": True, + "complete_streaming_response": response_obj, + } + + perform_redaction(model_call_details, result=None) + + redacted_response = model_call_details["complete_streaming_response"] + assert redacted_response.choices[0].message.content == "redacted-by-litellm" + + def test_streaming_responses_untouched_when_disabled(self): + response_obj = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret content", role="assistant") + ) + ] + ) + + model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "prompt": "hi", + "input": "hi", + "stream": True, + "async_complete_streaming_response": response_obj, + } + + perform_redaction(model_call_details, result=None, redact_streaming_responses=False) + + assert response_obj.choices[0].message.content == "secret content" + + +class TestRedactStreamingResponsesForCustomLogger: + def _model_call_details(self): + response_obj = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret content", role="assistant") + ) + ] + ) + return { + "stream": True, + "async_complete_streaming_response": response_obj, + }, response_obj + + def test_opted_out_logger_gets_redacted_copy(self): + model_call_details, response_obj = self._model_call_details() + opted_out_logger = CustomLogger(message_logging=False) + + redacted_details = redact_streaming_responses_for_custom_logger( + model_call_details=model_call_details, custom_logger=opted_out_logger + ) + + redacted_response = redacted_details["async_complete_streaming_response"] + assert redacted_response.choices[0].message.content == "redacted-by-litellm" + assert response_obj.choices[0].message.content == "secret content" + assert model_call_details["async_complete_streaming_response"] is response_obj + + def test_compliant_logger_gets_shared_response(self): + model_call_details, response_obj = self._model_call_details() + compliant_logger = CustomLogger() + + result_details = redact_streaming_responses_for_custom_logger( + model_call_details=model_call_details, custom_logger=compliant_logger + ) + + assert result_details is model_call_details + assert response_obj.choices[0].message.content == "secret content" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 14e872485ec..606ff39b35e 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -99,17 +99,11 @@ class TestContextManagementConversion: } ) kwargs = _ADAPTER.translate_request(req) - assert kwargs["context_management"] == [ - {"type": "compaction", "compact_threshold": 100000} - ] + assert kwargs["context_management"] == [{"type": "compaction", "compact_threshold": 100000}] def test_translate_request_drops_anthropic_only_context_management(self): """context_management with only unknown edit types is omitted from kwargs.""" - req = _make_request( - context_management={ - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - } - ) + req = _make_request(context_management={"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}) kwargs = _ADAPTER.translate_request(req) assert "context_management" not in kwargs @@ -134,9 +128,7 @@ class TestOutputConfigStructuredOutput: def test_output_config_format_json_schema_converted(self): """output_config.format.json_schema is converted to OpenAI text.format.""" - req = _make_request( - output_config={"format": {"type": "json_schema", "schema": self._SCHEMA}} - ) + req = _make_request(output_config={"format": {"type": "json_schema", "schema": self._SCHEMA}}) kwargs = _ADAPTER.translate_request(req) assert "text" in kwargs fmt = kwargs["text"]["format"] @@ -153,9 +145,7 @@ class TestOutputConfigStructuredOutput: def test_output_format_still_works(self): """The original output_format field still takes precedence when present.""" - req = _make_request( - output_format={"type": "json_schema", "schema": self._SCHEMA} - ) + req = _make_request(output_format={"type": "json_schema", "schema": self._SCHEMA}) kwargs = _ADAPTER.translate_request(req) assert "text" in kwargs assert kwargs["text"]["format"]["type"] == "json_schema" @@ -250,9 +240,7 @@ class TestTranslateMessagesToResponsesInput: ] result = _translate_messages(messages) assert len(result) == 1 - assert result[0]["content"] == [ - {"type": "input_image", "image_url": "data:image/png;base64,abc123"} - ] + assert result[0]["content"] == [{"type": "input_image", "image_url": "data:image/png;base64,abc123"}] def test_user_url_image(self): """User message with URL image source becomes input_image with the URL.""" @@ -268,9 +256,7 @@ class TestTranslateMessagesToResponsesInput: } ] result = _translate_messages(messages) - assert result[0]["content"] == [ - {"type": "input_image", "image_url": "https://example.com/img.jpg"} - ] + assert result[0]["content"] == [{"type": "input_image", "image_url": "https://example.com/img.jpg"}] def test_user_base64_image_empty_data_skipped(self): """Base64 image with empty data is skipped (no URL can be formed).""" @@ -341,9 +327,7 @@ class TestTranslateMessagesToResponsesInput: messages = [ { "role": "user", - "content": [ - {"type": "tool_result", "tool_use_id": "call_null", "content": None} - ], + "content": [{"type": "tool_result", "tool_use_id": "call_null", "content": None}], } ] result = _translate_messages(messages) @@ -370,9 +354,7 @@ class TestTranslateMessagesToResponsesInput: } ] result = _translate_messages(messages) - assert result[0]["content"] == [ - {"type": "output_text", "text": "Here is the answer."} - ] + assert result[0]["content"] == [{"type": "output_text", "text": "Here is the answer."}] def test_assistant_tool_use_becomes_function_call(self): """Assistant tool_use block becomes a top-level function_call item.""" @@ -404,15 +386,11 @@ class TestTranslateMessagesToResponsesInput: messages = [ { "role": "assistant", - "content": [ - {"type": "thinking", "thinking": "Let me reason step by step."} - ], + "content": [{"type": "thinking", "thinking": "Let me reason step by step."}], } ] result = _translate_messages(messages) - assert result[0]["content"] == [ - {"type": "output_text", "text": "Let me reason step by step."} - ] + assert result[0]["content"] == [{"type": "output_text", "text": "Let me reason step by step."}] def test_assistant_empty_thinking_block_skipped(self): """Assistant thinking block with empty thinking text is skipped.""" @@ -584,28 +562,27 @@ class TestTranslateToolsToResponsesAPI: class TestTranslateToolChoiceToResponsesAPI: - """Anthropic tool_choice -> Responses API tool_choice.""" + """Anthropic tool_choice -> Responses API tool_choice. - def test_auto_maps_to_auto(self): - assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "auto"}) == { - "type": "auto" - } + The Responses API's tool_choice schema (openai.types.responses.tool_choice_options) + is a bare Literal["none", "auto", "required"] for these simple cases - not an + object like {"type": "auto"}. Sending the object shape to an OpenAI-compatible + server gets rejected with a pydantic validation error. + """ - def test_any_maps_to_required(self): - assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "any"}) == { - "type": "required" - } + def test_auto_maps_to_bare_string_auto(self): + assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "auto"}) == "auto" + + def test_any_maps_to_bare_string_required(self): + assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "any"}) == "required" + + def test_none_maps_to_bare_string_none(self): + assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "none"}) == "none" def test_specific_tool_maps_to_function(self): - result = _ADAPTER.translate_tool_choice_to_responses_api( - {"type": "tool", "name": "get_weather"} - ) + result = _ADAPTER.translate_tool_choice_to_responses_api({"type": "tool", "name": "get_weather"}) assert result == {"type": "function", "name": "get_weather"} - def test_unknown_type_defaults_to_auto(self): - result = _ADAPTER.translate_tool_choice_to_responses_api({"type": "none"}) - assert result == {"type": "auto"} - # --------------------------------------------------------------------------- # translate_thinking_to_reasoning @@ -616,17 +593,13 @@ class TestTranslateThinkingToReasoning: """Anthropic thinking param -> Responses API reasoning param.""" def test_budget_high_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning( - {"type": "enabled", "budget_tokens": 10000} - ) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}) # Default (reasoning_auto_summary=False): only effort, no summary assert result == {"effort": "high"} assert result is not None and "summary" not in result def test_budget_above_threshold_high_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning( - {"type": "enabled", "budget_tokens": 50000} - ) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 50000}) assert result is not None assert result["effort"] == "high" assert "summary" not in result @@ -652,9 +625,7 @@ class TestTranslateThinkingToReasoning: assert result is not None and "summary" not in result def test_budget_minimal_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning( - {"type": "enabled", "budget_tokens": 500} - ) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 500}) assert result == {"effort": "minimal"} assert result is not None and "summary" not in result @@ -707,9 +678,7 @@ class TestTranslateThinkingToReasoning: original = litellm.reasoning_auto_summary try: litellm.reasoning_auto_summary = True - result = _ADAPTER.translate_thinking_to_reasoning( - {"type": "enabled", "budget_tokens": 10000} - ) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}) assert result == {"effort": "high", "summary": "detailed"} finally: litellm.reasoning_auto_summary = original @@ -789,11 +758,7 @@ class TestTranslateRequestBroaderCoverage: assert kwargs["top_p"] == 0.9 def test_tools_translated(self): - req = _make_request( - tools=[ - {"name": "calculator", "description": "Does math.", "input_schema": {}} - ] - ) + req = _make_request(tools=[{"name": "calculator", "description": "Does math.", "input_schema": {}}]) kwargs = _ADAPTER.translate_request(req) assert len(kwargs["tools"]) == 1 assert kwargs["tools"][0]["name"] == "calculator" @@ -929,9 +894,7 @@ class TestTranslateResponse: def test_multiple_text_parts(self): """Multiple output_text parts become multiple text content blocks.""" - response = _make_mock_response( - output=[_make_output_message(["Part 1", "Part 2"])] - ) + response = _make_mock_response(output=[_make_output_message(["Part 1", "Part 2"])]) result: Any = _ADAPTER.translate_response(response) assert len(result["content"]) == 2 assert result["content"][0]["text"] == "Part 1" diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index 7e9c8a273ae..0cbdc518cc2 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -158,6 +158,54 @@ class TestClaudePlatformActionsCovered: ) +class TestBedrockMantleActionsCovered: + """LIT-3859: bedrock_mantle inference authorizes against the + ``bedrock-mantle`` action namespace, so the session-policy ceiling + must include it or every Mantle request via OIDC/WIF auth denies + with "no session policy allows the bedrock-mantle:CreateInference + action" even when the role's identity policy grants it.""" + + def test_bedrock_mantle_create_inference_present(self): + policy = _captured_policy() + all_actions: set = set() + for stmt in policy["Statement"]: + stmt_actions = stmt.get("Action") + if isinstance(stmt_actions, str): + all_actions.add(stmt_actions) + elif isinstance(stmt_actions, list): + all_actions.update(stmt_actions) + assert "bedrock-mantle:CreateInference" in all_actions, ( + "bedrock-mantle:CreateInference missing from session policy — " + "bedrock_mantle/* requests will 403 on OIDC/WIF auth" + ) + + def test_bedrock_mantle_statement_allows(self): + policy = _captured_policy() + stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM") + assert stmt["Effect"] == "Allow" + assert stmt["Resource"] == "*" + + def test_no_bedrock_mantle_wildcard(self): + policy = _captured_policy() + stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM") + actions = stmt["Action"] + if isinstance(actions, str): + actions = [actions] + assert "bedrock-mantle:*" not in actions, ( + "session policy must not grant bedrock-mantle:* — " + "the ceiling should match the documented action set" + ) + + def test_bedrock_mantle_statement_carries_secure_transport_condition(self): + policy = _captured_policy() + stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM") + cond = stmt.get("Condition") or {} + assert cond.get("Bool", {}).get("aws:SecureTransport") == "true", ( + "BedrockMantleLiteLLM must require aws:SecureTransport=true " + "to keep parity with the bedrock statement" + ) + + def _make_jwt(payload: dict) -> str: def _segment(data: dict) -> str: return base64.urlsafe_b64encode(json.dumps(data).encode()).rstrip(b"=").decode() diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index d389b54b3f1..151c51f1ca0 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -44,6 +44,60 @@ class TestOpenAIResponsesAPIConfig: # The function should return the params unchanged assert result == test_params + @pytest.mark.parametrize("max_output_tokens", [1, 15]) + def test_map_openai_params_clamps_max_output_tokens_below_minimum(self, max_output_tokens): + """OpenAI's Responses API rejects max_output_tokens < 16. + + Claude Code (via the Anthropic Messages -> Responses adapter) sends a + max_tokens=1 warmup probe when running `/model`, which produced: + "Invalid 'max_output_tokens': integer below minimum value. + Expected a value >= 16, but got 1 instead." + Clamp anything below the minimum up to 16 instead of erroring. + """ + result = self.config.map_openai_params( + response_api_optional_params={"max_output_tokens": max_output_tokens}, + model=self.model, + drop_params=False, + ) + + assert result["max_output_tokens"] == 16 + + def test_map_openai_params_preserves_max_output_tokens_at_or_above_minimum(self): + """Values already >= 16 must pass through untouched.""" + result = self.config.map_openai_params( + response_api_optional_params={"max_output_tokens": 256}, + model=self.model, + drop_params=False, + ) + + assert result["max_output_tokens"] == 256 + + def test_map_openai_params_leaves_max_output_tokens_absent(self): + """A request without max_output_tokens must not gain the key.""" + result = self.config.map_openai_params( + response_api_optional_params={"input": "hi"}, + model=self.model, + drop_params=False, + ) + + assert "max_output_tokens" not in result + + @pytest.mark.parametrize( + "value, expected", + [ + (1, 16), + (15, 16), + (16, 16), + (17, 17), + (256, 256), + (None, None), + ], + ) + def test_enforce_min_max_output_tokens(self, value, expected): + """Below the minimum clamps to 16; the boundary, larger values, and None + are returned unchanged so no previously-valid request regresses.""" + assert self.config._enforce_min_max_output_tokens(value) == expected + def validate_responses_api_request_params(self, params, expected_fields): """ Validate that the params dict has the expected structure of ResponsesAPIRequestParams diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index c785ac577f7..6f132aaae9c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -4910,6 +4910,7 @@ class TestMCPDcrBridgeDelegateAdmission: cls, *, key_hash=None, + user_id=None, server_id="bridge-server-id", access_token="inner-upstream-access-token", token_type="Bearer", @@ -4921,17 +4922,23 @@ class TestMCPDcrBridgeDelegateAdmission: envelope_keys_from_master_key, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( - EnvelopeIdentity, SealedEnvelope, UpstreamTokenGrant, + key_hash_identity, mint_envelope, + user_identity, ) from pydantic import SecretStr + identity = ( + user_identity(server_id=server_id, user_id=user_id) + if user_id is not None + else key_hash_identity(server_id=server_id, key_hash=key_hash or cls._KEY_HASH) + ) keys = envelope_keys_from_master_key(master_key or cls._MASTER_KEY) now = minted_at or datetime.now(timezone.utc) sealed = mint_envelope( - identity=EnvelopeIdentity(server_id=server_id, key_hash=key_hash or cls._KEY_HASH), + identity=identity, grant=UpstreamTokenGrant( access_token=SecretStr(access_token), token_type=token_type, @@ -4999,6 +5006,38 @@ class TestMCPDcrBridgeDelegateAdmission: stack.enter_context(patcher) yield get_key_object + @staticmethod + @contextlib.contextmanager + def _patch_user_reload(*, return_value=None, side_effect=None): + """Patch the user-subject reload path an interactively-minted envelope takes: the + ``get_user_object`` lookup ``_reload_admitted_user`` runs (which also drives the SCIM gate), + plus the ``prisma_client`` / ``user_api_key_cache`` globals. The centralized gate's own + fetches fail-safe to None under the MagicMock prisma, so an unblocked user admits. Yields the + ``get_user_object`` mock so a caller can assert the sealed user_id was the reload key.""" + get_user_object = AsyncMock(return_value=return_value, side_effect=side_effect) + with ( + patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + ): + yield get_user_object + + @staticmethod + def _wrapped_user_lookup_error(original: BaseException) -> ValueError: + """Reproduce get_user_object's real exception contract (litellm/proxy/auth/auth_checks.py): it + catches every DB failure in a broad ``except`` and re-raises a bare ``ValueError``, so the + original error (a missing-user Exception or a real outage) survives only as ``__context__``. + Injecting a raw ConnectionError/Exception instead would exercise a shape production never + produces and let a chain-blind outage classifier pass. That wrapping fidelity is itself pinned by + test_get_user_object_wraps_db_outage_as_valueerror_preserving_context in test_auth_checks.""" + try: + raise original + except BaseException: + try: + raise ValueError(f"User doesn't exist in db. Got error - {original}") + except ValueError as wrapped: + return wrapped + @staticmethod def _mcp_request(path="/mcp/bridge_delegate_server"): """A minimal ``Request`` for direct ``_admit_dcr_bridge_delegate`` calls, mirroring how @@ -5060,6 +5099,155 @@ class TestMCPDcrBridgeDelegateAdmission: "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} } + async def test_user_subject_envelope_admits_under_the_reloaded_user(self): + """An interactively-minted (user_id) envelope admits under the reloaded USER, not a key: the + reload is keyed by the sealed user_id, the admitted auth carries that user_id, the raw-key + pipeline is never invoked, and the inner upstream token is injected for egress. This is the + interactive-DCR admission the whole flow exists for.""" + envelope = self._mint_bridge_envelope(user_id="sso-user-7") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_user_reload( + return_value=MagicMock( + user_id="sso-user-7", + metadata={"scim_active": True}, + user_role=None, + object_permission=None, + object_permission_id=None, + ) + ) as get_user_object, + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + (auth_result, _h, _s, mcp_server_auth_headers, _o, _r) = await MCPRequestHandler.process_mcp_request(scope) + + assert get_user_object.await_args.kwargs["user_id"] == "sso-user-7" + assert auth_result.user_id == "sso-user-7" + mock_auth.assert_not_called() + assert mcp_server_auth_headers == { + "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} + } + + async def test_user_subject_envelope_carries_the_users_mcp_object_permission(self): + """The admitted user's own MCP object permission rides on the returned auth so the shared + get_allowed_mcp_servers grants the user their litellm-granted servers, rather than admitting a + bare user with no MCP access. Regression for the signed-in SSO client getting zero tools because + the reload dropped the user's object permission.""" + object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="op-user-7", mcp_servers=["bridge_delegate_server"] + ) + envelope = self._mint_bridge_envelope(user_id="sso-user-7") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ), + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_user_reload( + return_value=MagicMock( + user_id="sso-user-7", + metadata={"scim_active": True}, + user_role=None, + object_permission=object_permission, + object_permission_id="op-user-7", + ) + ), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + (auth_result, _h, _s, _headers, _o, _r) = await MCPRequestHandler.process_mcp_request(scope) + + assert auth_result.object_permission is not None + assert auth_result.object_permission.mcp_servers == ["bridge_delegate_server"] + + async def test_user_subject_envelope_missing_user_fails_closed_401(self): + """A user_id envelope whose user has since been deleted must fail closed with a 401, not a 500. + get_user_object catches the missing row and re-raises a bare ValueError (it does not return None + on the production path), so the reload must fail closed rather than let it propagate as an opaque + 500, and must not mistake the wrapped ValueError for a DB outage. Regression for the missing-user + path surfacing as a 500.""" + envelope = self._mint_bridge_envelope(user_id="ghost-user") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_user_reload(side_effect=self._wrapped_user_lookup_error(Exception())), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + + async def test_user_subject_envelope_db_outage_is_retryable_503(self): + """A transient database outage while reloading the envelope's user is a retryable 503, not an + opaque 500, matching the key path's contract so an interactive DCR client retries instead of + treating a live identity as invalid. get_user_object wraps the outage in a bare ValueError, so this + exercises the chain-aware classifier; a raw ConnectionError would falsely pass even the old + chain-blind check because it is an OSError. Regression for the user reload dropping the 503 arm.""" + envelope = self._mint_bridge_envelope(user_id="sso-user-7") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_user_reload( + side_effect=self._wrapped_user_lookup_error(ConnectionError("auth database unreachable")) + ), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 503 + + async def test_user_subject_envelope_scim_deactivated_user_fails_closed_401(self): + """SCIM-deactivating the envelope's user revokes it immediately: the reloaded user carries + scim_active False, so admission 401s rather than letting an offboarded user keep tool access + until the envelope expires.""" + envelope = self._mint_bridge_envelope(user_id="offboarded-user") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_user_reload(return_value=MagicMock(user_id="offboarded-user", metadata={"scim_active": False})), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + async def test_revoked_key_envelope_fails_closed_401(self): """An envelope whose key has since been deleted must fail closed: ``get_key_object`` raises for the missing row, so admission 401s instead of admitting the caller as an unrestricted diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_classify.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_classify.py new file mode 100644 index 00000000000..dc20d664a53 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_classify.py @@ -0,0 +1,147 @@ +"""Classification matrix for upstream OAuth/DCR rejections: who is blamed depends only on the §5.2 +code and whose credentials the gateway presented, never on the upstream's HTTP status.""" + +import httpx + +from litellm.proxy._experimental.mcp_server.faults.classify import ( + classify_upstream_dcr_rejection, + classify_upstream_token_rejection, +) +from litellm.proxy._experimental.mcp_server.faults.types import ( + CallerRejected, + GatewayRejected, + UpstreamProtocolFault, + UpstreamReportedFault, +) + + +def _response(status_code: int, *, json_body: object = None, text_body: str = "", headers: dict = None) -> httpx.Response: + request = httpx.Request("POST", "https://idp.example.com/token") + if json_body is not None: + return httpx.Response(status_code, json=json_body, request=request) + return httpx.Response(status_code, text=text_body, headers=headers or {}, request=request) + + +def test_caller_fault_code_classifies_as_caller_rejected_regardless_of_status(): + fault = classify_upstream_token_rejection( + _response(500, json_body={"error": "invalid_grant", "error_description": "Code expired."}), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, CallerRejected) + assert fault.code == "invalid_grant" + assert fault.description == "Code expired." + + +def test_credential_code_with_gateway_stored_credentials_indicts_gateway(): + fault = classify_upstream_token_rejection( + _response(401, json_body={"error": "invalid_client", "error_description": "not found"}), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, GatewayRejected) + assert fault.code == "invalid_client" + + +def test_credential_code_with_caller_supplied_credentials_stays_caller_fault(): + fault = classify_upstream_token_rejection( + _response(401, json_body={"error": "invalid_client"}), + credential_source="caller_supplied", + log_context="srv", + ) + assert isinstance(fault, CallerRejected) + assert fault.code == "invalid_client" + + +def test_unknown_code_relays_as_caller_rejected(): + fault = classify_upstream_token_rejection( + _response(400, json_body={"error": "slow_down", "error_description": "Polling too fast."}), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, CallerRejected) + assert fault.code == "slow_down" + + +def test_body_without_error_field_is_protocol_fault(): + fault = classify_upstream_token_rejection( + _response(404, text_body="not here"), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, UpstreamProtocolFault) + assert fault.note == "upstream token endpoint returned HTTP 404" + + +def test_unreadable_body_is_protocol_fault_not_exception(): + unreadable = httpx.Response( + 400, + stream=httpx.ByteStream(b"\x1f\x8bnot-gzip"), + headers={"content-encoding": "gzip"}, + request=httpx.Request("POST", "https://idp.example.com/token"), + ) + fault = classify_upstream_token_rejection(unreadable, credential_source="gateway_stored", log_context="srv") + assert isinstance(fault, UpstreamProtocolFault) + + +def test_wire_fields_are_bounded(): + fault = classify_upstream_token_rejection( + _response(400, json_body={"error": "invalid_request", "error_description": "x" * 5000}), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, CallerRejected) + assert len(fault.description) == 500 + + +def test_dcr_rejection_with_rfc7591_code_is_caller_rejected(): + fault = classify_upstream_dcr_rejection( + _response(400, json_body={"error": "invalid_redirect_uri", "error_description": "not allowed"}), + log_context="srv", + ) + assert isinstance(fault, CallerRejected) + assert fault.code == "invalid_redirect_uri" + + +def test_dcr_rejection_without_code_is_protocol_fault(): + fault = classify_upstream_dcr_rejection(_response(500, text_body="trace"), log_context="srv") + assert isinstance(fault, UpstreamProtocolFault) + assert fault.note == "upstream registration failed with HTTP 500" + + +def test_upstream_self_blame_codes_stay_upstream_faults(): + fault = classify_upstream_token_rejection( + _response(400, json_body={"error": "server_error", "error_description": "boom"}), + credential_source="caller_supplied", + log_context="srv", + ) + assert isinstance(fault, UpstreamReportedFault) + assert fault.code == "server_error" + + +def test_temporarily_unavailable_is_upstream_fault(): + fault = classify_upstream_token_rejection( + _response(503, json_body={"error": "temporarily_unavailable"}), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, UpstreamReportedFault) + assert fault.code == "temporarily_unavailable" + + +def test_invalid_target_is_gateway_fault_even_with_caller_credentials(): + fault = classify_upstream_token_rejection( + _response(400, json_body={"error": "invalid_target"}), + credential_source="caller_supplied", + log_context="srv", + ) + assert isinstance(fault, GatewayRejected) + assert fault.code == "invalid_target" + + +def test_dcr_server_error_code_is_not_blamed_on_caller(): + fault = classify_upstream_dcr_rejection( + _response(500, json_body={"error": "server_error"}), + log_context="srv", + ) + assert isinstance(fault, UpstreamReportedFault) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_render_oauth.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_render_oauth.py new file mode 100644 index 00000000000..78513e315a7 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_render_oauth.py @@ -0,0 +1,96 @@ +"""Rendering contract: status, wire code, and prose all derive from the fault tag, so a caller-fault +code can never ship on a server-fault status and gateway-side faults never carry provider prose.""" + +import json + +from litellm.proxy._experimental.mcp_server.faults.render_oauth import ( + dcr_fault_detail, + render_token_fault, +) +from litellm.proxy._experimental.mcp_server.faults.types import ( + CallerRejected, + GatewayRejected, + UpstreamProtocolFault, + UpstreamReportedFault, +) + + +def test_caller_rejected_renders_code_derived_status(): + response = render_token_fault(CallerRejected(code="invalid_grant", description="Code expired.")) + assert response.status_code == 400 + assert json.loads(response.body) == {"error": "invalid_grant", "error_description": "Code expired."} + assert response.headers["cache-control"] == "no-store" + + +def test_caller_rejected_invalid_client_renders_401(): + response = render_token_fault(CallerRejected(code="invalid_client")) + assert response.status_code == 401 + assert json.loads(response.body) == {"error": "invalid_client"} + + +def test_caller_rejected_includes_error_uri_only_when_present(): + response = render_token_fault( + CallerRejected(code="invalid_scope", description="bad scope", error_uri="https://idp.example.com/e") + ) + assert json.loads(response.body) == { + "error": "invalid_scope", + "error_description": "bad scope", + "error_uri": "https://idp.example.com/e", + } + + +def test_gateway_rejected_renders_502_with_gateway_prose(): + response = render_token_fault(GatewayRejected(code="invalid_client")) + assert response.status_code == 502 + body = json.loads(response.body) + assert body["error"] == "server_error" + assert "invalid_client" in body["error_description"] + assert "client_id and client_secret" in body["error_description"] + + +def test_gateway_invalid_target_prose_names_resource_indicators(): + response = render_token_fault(GatewayRejected(code="invalid_target")) + body = json.loads(response.body) + assert response.status_code == 502 + assert "RFC 8707" in body["error_description"] + + +def test_protocol_fault_renders_502_note(): + response = render_token_fault(UpstreamProtocolFault(note="upstream token endpoint returned HTTP 503")) + assert response.status_code == 502 + assert json.loads(response.body) == { + "error": "server_error", + "error_description": "upstream token endpoint returned HTTP 503", + } + + +def test_dcr_caller_rejection_is_400_per_rfc7591_regardless_of_upstream_status(): + status_code, detail = dcr_fault_detail(CallerRejected(code="invalid_client_metadata", description="bad grant types")) + assert status_code == 400 + assert detail == "invalid_client_metadata: bad grant types" + + +def test_dcr_protocol_fault_is_502(): + status_code, detail = dcr_fault_detail(UpstreamProtocolFault(note="upstream registration failed with HTTP 500")) + assert status_code == 502 + assert detail == "upstream registration failed with HTTP 500" + + +def test_upstream_reported_server_error_renders_502_with_matching_code(): + response = render_token_fault(UpstreamReportedFault(code="server_error")) + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +def test_upstream_reported_temporarily_unavailable_renders_503_with_matching_code(): + response = render_token_fault(UpstreamReportedFault(code="temporarily_unavailable")) + assert response.status_code == 503 + body = json.loads(response.body) + assert body["error"] == "temporarily_unavailable" + assert "retry" in body["error_description"] + + +def test_dcr_upstream_reported_fault_maps_to_5xx(): + status_code, detail = dcr_fault_detail(UpstreamReportedFault(code="server_error")) + assert status_code == 502 + assert "internal error" in detail diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index 82e8e2aae89..753a3d6a942 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -15,10 +15,14 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, BridgeEnvelopeInvalid, + BridgeRefreshInvalid, + BridgeRefreshOpened, NotBridgeEnvelope, + build_bridge_refresh_token_response, build_bridge_token_response, envelope_keys_from_master_key, is_bridge_envelope_shaped, + open_bridge_refresh_envelope, resolve_bridge_envelope, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( @@ -26,15 +30,17 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import EnvelopeIdentity, EnvelopeKeys, EnvelopeTooLarge, + RefreshCredential, SealedEnvelope, UpstreamTokenGrant, + key_hash_identity, mint_envelope, ) _NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) _MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789" _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" -_IDENTITY = EnvelopeIdentity(server_id="srv-456", key_hash="hashed-key-123") +_IDENTITY = key_hash_identity(server_id="srv-456", key_hash="hashed-key-123") _SERVER_ID = _IDENTITY.server_id @@ -48,6 +54,76 @@ def _sealed_token(keys: EnvelopeKeys, now: datetime = _NOW, identity: EnvelopeId return sealed.token.get_secret_value() +_UPSTREAM_REFRESH = "upstream-refresh-do-not-leak-9b2c" + + +def _sealed_refresh(keys: EnvelopeKeys, now: datetime = _NOW, identity: EnvelopeIdentity = _IDENTITY) -> str: + sealed = build_bridge_refresh_token_response( + identity, RefreshCredential(refresh_token=SecretStr(_UPSTREAM_REFRESH)), keys, now + ) + assert isinstance(sealed, SealedEnvelope) + return sealed.token.get_secret_value() + + +def test_open_bridge_refresh_envelope_round_trips_identity_and_refresh(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = open_bridge_refresh_envelope(_sealed_refresh(keys), keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeRefreshOpened) + assert result.identity == _IDENTITY + assert result.refresh.refresh_token.get_secret_value() == _UPSTREAM_REFRESH + + +def test_open_bridge_refresh_envelope_strips_bearer_scheme(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = open_bridge_refresh_envelope(f"Bearer {_sealed_refresh(keys)}", keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeRefreshOpened) + + +def test_open_bridge_refresh_envelope_rejects_wrong_server(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = open_bridge_refresh_envelope(_sealed_refresh(keys), keys, _NOW, "a-different-server") + assert isinstance(result, BridgeRefreshInvalid) + + +def test_open_bridge_refresh_envelope_rejects_non_refresh_bearers(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + # an access envelope is not a refresh envelope; a raw upstream refresh token is not one either + assert isinstance(open_bridge_refresh_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID), BridgeRefreshInvalid) + assert isinstance(open_bridge_refresh_envelope("raw-refresh-token", keys, _NOW, _SERVER_ID), BridgeRefreshInvalid) + + +def test_open_bridge_refresh_envelope_rejects_under_wrong_master_key(): + minted = envelope_keys_from_master_key(_MASTER_KEY) + other = envelope_keys_from_master_key(_MASTER_KEY + "-rotated") + result = open_bridge_refresh_envelope(_sealed_refresh(minted), other, _NOW, _SERVER_ID) + assert isinstance(result, BridgeRefreshInvalid) + + +def test_refresh_envelope_is_never_admitted_at_the_tool_call_edge(): + """A refresh envelope must never authenticate a tool call. The admission edge engages the bridge arm + for it (is_bridge_envelope_shaped is true for either envelope kind), and the consumer rejects it as + BridgeEnvelopeInvalid, which admission fails closed (401): a refresh credential is only ever + presented back to the token endpoint.""" + keys = envelope_keys_from_master_key(_MASTER_KEY) + refresh = _sealed_refresh(keys) + assert is_bridge_envelope_shaped(refresh) is True + assert is_bridge_envelope_shaped(f"Bearer {refresh}") is True + result = resolve_bridge_envelope(refresh, keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeInvalid) + + +def test_refresh_jwt_wearing_the_access_prefix_is_rejected_at_the_edge(): + """Belt-and-suspenders against a swapped wire prefix: a refresh JWT re-prefixed as an access envelope + opens far enough to hit the signed kind claim, which rejects it, so admission fails closed rather + than forwarding a refresh credential's contents upstream.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import REFRESH_ENVELOPE_PREFIX + + keys = envelope_keys_from_master_key(_MASTER_KEY) + swapped = ENVELOPE_PREFIX + _sealed_refresh(keys).removeprefix(REFRESH_ENVELOPE_PREFIX) + result = resolve_bridge_envelope(swapped, keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeInvalid) + + def test_key_derivation_is_deterministic(): assert envelope_keys_from_master_key(_MASTER_KEY) == envelope_keys_from_master_key(_MASTER_KEY) @@ -138,7 +214,7 @@ def test_resolve_envelope_minted_for_another_server_is_invalid(): captured or misrouted envelope cannot forward one server's upstream credential to another. The valid access token stays sealed; the mismatch alone fails the resolve.""" keys = envelope_keys_from_master_key(_MASTER_KEY) - other_server_identity = EnvelopeIdentity(server_id="srv-OTHER", key_hash=_IDENTITY.key_hash) + other_server_identity = key_hash_identity(server_id="srv-OTHER", key_hash=_IDENTITY.subject) token = _sealed_token(keys, identity=other_server_identity) result = resolve_bridge_envelope(token, keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeInvalid) @@ -155,7 +231,7 @@ def test_resolve_non_ascii_server_id_stays_total_and_does_not_raise(): unicode server_id); it stays total and returns a typed result. A matching non-ASCII id admits, a mismatching one is BridgeEnvelopeInvalid, and neither raises.""" keys = envelope_keys_from_master_key(_MASTER_KEY) - unicode_identity = EnvelopeIdentity(server_id="srv-café", key_hash=_IDENTITY.key_hash) + unicode_identity = key_hash_identity(server_id="srv-café", key_hash=_IDENTITY.subject) token = _sealed_token(keys, identity=unicode_identity) assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-café"), BridgeEnvelopeAdmitted) assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-cafe"), BridgeEnvelopeInvalid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py index b44f3f84cc9..ae196c9080b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py @@ -24,6 +24,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ENVELOPE_PREFIX, MAX_ENVELOPE_BYTES, MAX_ENVELOPE_TTL_SECONDS, + MAX_REFRESH_ENVELOPE_TTL_SECONDS, + REFRESH_ENVELOPE_PREFIX, BadSignature, DecryptFailed, EnvelopeIdentity, @@ -33,11 +35,18 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import MalformedPayload, NotAnEnvelope, OpenedEnvelope, + OpenedRefreshEnvelope, + RefreshCredential, SealedEnvelope, UpstreamTokenGrant, is_envelope, + is_refresh_envelope, + key_hash_identity, mint_envelope, + mint_refresh_envelope, open_envelope, + open_refresh_envelope, + user_identity, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value @@ -51,7 +60,7 @@ _WRONG_SIGNING = EnvelopeKeys(signing_key=SecretStr(_OTHER_SIGNING_KEY), encrypt _WRONG_ENCRYPTION = EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr(_OTHER_ENCRYPTION_KEY)) _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" _REFRESH_TOKEN = "upstream-refresh-token-do-not-leak-1d0aa4b7" -_IDENTITY = EnvelopeIdentity(server_id="srv-456", key_hash="hashed-key-123") +_IDENTITY = key_hash_identity(server_id="srv-456", key_hash="hashed-key-123") def _full_grant() -> UpstreamTokenGrant: @@ -137,17 +146,99 @@ def test_minimal_grant_round_trips_without_none_leakage_into_claims(): def test_claim_layout_and_no_plaintext_token_in_envelope(): token = _sealed_token(_full_grant()) claims = _unverified_claims(token) - assert set(claims) == {"iss", "iat", "exp", "server_id", "key_hash", "grant"} + assert set(claims) == {"iss", "iat", "exp", "kind", "server_id", "subject_type", "subject", "grant"} assert claims["iss"] == ENVELOPE_ISSUER assert claims["iat"] == int(_NOW.timestamp()) assert claims["exp"] == int(_NOW.timestamp()) + 600 + assert claims["kind"] == "access" assert claims["server_id"] == "srv-456" - assert claims["key_hash"] == "hashed-key-123" + assert claims["subject_type"] == "key_hash" + assert claims["subject"] == "hashed-key-123" assert _ACCESS_TOKEN not in token assert _ACCESS_TOKEN not in json.dumps(claims) assert _REFRESH_TOKEN not in json.dumps(claims) +def _refresh_credential() -> RefreshCredential: + return RefreshCredential(refresh_token=SecretStr(_REFRESH_TOKEN), scope="read:tools", expires_in=None) + + +def _sealed_refresh_token(refresh: RefreshCredential | None = None, keys: EnvelopeKeys = _KEYS) -> str: + sealed = mint_refresh_envelope(_IDENTITY, refresh or _refresh_credential(), keys, _NOW) + assert isinstance(sealed, SealedEnvelope) + return sealed.token.get_secret_value() + + +def test_refresh_envelope_round_trips_identity_and_refresh_token(): + token = _sealed_refresh_token() + assert is_refresh_envelope(token) + assert not is_envelope(token) + opened = open_refresh_envelope(token, _KEYS, _NOW) + assert isinstance(opened, OpenedRefreshEnvelope) + assert opened.identity == _IDENTITY + assert opened.refresh.refresh_token.get_secret_value() == _REFRESH_TOKEN + assert opened.refresh.scope == "read:tools" + + +def test_refresh_envelope_ttl_is_min_of_upstream_refresh_lifetime_and_cap(): + short = mint_refresh_envelope( + _IDENTITY, RefreshCredential(refresh_token=SecretStr("r"), expires_in=120), _KEYS, _NOW + ) + assert isinstance(short, SealedEnvelope) + assert short.expires_at == _NOW + timedelta(seconds=120) + capped = mint_refresh_envelope( + _IDENTITY, + RefreshCredential(refresh_token=SecretStr("r"), expires_in=MAX_REFRESH_ENVELOPE_TTL_SECONDS + 86400), + _KEYS, + _NOW, + ) + assert isinstance(capped, SealedEnvelope) + assert capped.expires_at == _NOW + timedelta(seconds=MAX_REFRESH_ENVELOPE_TTL_SECONDS) + default = mint_refresh_envelope(_IDENTITY, RefreshCredential(refresh_token=SecretStr("r")), _KEYS, _NOW) + assert isinstance(default, SealedEnvelope) + assert default.expires_at == _NOW + timedelta(seconds=MAX_REFRESH_ENVELOPE_TTL_SECONDS) + + +def test_access_and_refresh_envelopes_do_not_cross_open(): + access = _sealed_token(_full_grant()) + refresh = _sealed_refresh_token() + # each opener rejects the other kind's prefix outright + assert isinstance(open_refresh_envelope(access, _KEYS, _NOW), NotAnEnvelope) + assert isinstance(open_envelope(refresh, _KEYS, _NOW), NotAnEnvelope) + + +def test_prefix_swap_is_rejected_by_the_signed_kind_claim(): + # the wire prefix is not signed, so swap it; the signed kind claim must still reject the cross-use + refresh = _sealed_refresh_token() + swapped_to_access = ENVELOPE_PREFIX + refresh.removeprefix(REFRESH_ENVELOPE_PREFIX) + assert isinstance(open_envelope(swapped_to_access, _KEYS, _NOW), MalformedPayload) + access = _sealed_token(_full_grant()) + swapped_to_refresh = REFRESH_ENVELOPE_PREFIX + access.removeprefix(ENVELOPE_PREFIX) + assert isinstance(open_refresh_envelope(swapped_to_refresh, _KEYS, _NOW), MalformedPayload) + + +def test_refresh_envelope_total_over_hostile_input(): + token = _sealed_refresh_token() + # expired against the injected clock + assert isinstance( + open_refresh_envelope(token, _KEYS, _NOW + timedelta(seconds=MAX_REFRESH_ENVELOPE_TTL_SECONDS)), Expired + ) + # wrong signing key + assert isinstance(open_refresh_envelope(token, _WRONG_SIGNING, _NOW), BadSignature) + # right signature, wrong encryption key + assert isinstance(open_refresh_envelope(token, _WRONG_ENCRYPTION, _NOW), DecryptFailed) + # not an envelope at all + assert isinstance(open_refresh_envelope("raw-upstream-refresh-token", _KEYS, _NOW), NotAnEnvelope) + + +def test_refresh_envelope_never_leaks_the_refresh_token_in_plaintext(): + token = _sealed_refresh_token() + assert _REFRESH_TOKEN not in token + claims = jwt.decode(token.removeprefix(REFRESH_ENVELOPE_PREFIX), options={"verify_signature": False}) + assert claims["kind"] == "refresh" + assert _REFRESH_TOKEN not in json.dumps(claims) + + @pytest.mark.parametrize( "expires_in, expected_ttl", [ @@ -226,11 +317,11 @@ def test_wrong_issuer_is_malformed_payload(): def test_missing_identity_claim_is_malformed_payload(): claims = _unverified_claims(_sealed_token(_full_grant())) - forged = _forge({key: value for key, value in claims.items() if key != "key_hash"}) + forged = _forge({key: value for key, value in claims.items() if key != "subject"}) assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) -@pytest.mark.parametrize("identity_claim", ["server_id", "key_hash"]) +@pytest.mark.parametrize("identity_claim", ["server_id", "subject"]) def test_signed_empty_identity_claim_is_malformed_payload_not_a_raise(identity_claim): claims = _unverified_claims(_sealed_token(_full_grant())) forged = _forge({**claims, identity_claim: ""}) @@ -463,9 +554,11 @@ def test_non_positive_expires_in_is_rejected_at_construction_without_leaking(): def test_empty_identity_and_key_fields_are_rejected_at_construction(): with pytest.raises(ValidationError): - EnvelopeIdentity(server_id="", key_hash="hashed-key-123") + EnvelopeIdentity(server_id="", subject_type="key_hash", subject="hashed-key-123") with pytest.raises(ValidationError): - EnvelopeIdentity(server_id="srv-456", key_hash="") + EnvelopeIdentity(server_id="srv-456", subject_type="key_hash", subject="") + with pytest.raises(ValidationError): + EnvelopeIdentity(server_id="srv-456", subject_type="not-a-subject-type", subject="x") with pytest.raises(ValidationError): EnvelopeKeys(signing_key=SecretStr(""), encryption_key=SecretStr(_ENCRYPTION_KEY)) with pytest.raises(ValidationError): @@ -474,6 +567,20 @@ def test_empty_identity_and_key_fields_are_rejected_at_construction(): UpstreamTokenGrant(access_token=SecretStr(""), token_type="Bearer") +def test_user_subject_identity_round_trips(): + """The user_id subject variant seals and opens with its discriminator intact, so the edge can + tell an interactively-minted (user) envelope from a scripted (key_hash) one and reload the right + kind of record.""" + identity = user_identity(server_id="srv-456", user_id="user-42") + sealed = mint_envelope(identity, _full_grant(), _KEYS, _NOW) + assert isinstance(sealed, SealedEnvelope) + opened = open_envelope(sealed.token.get_secret_value(), _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + assert opened.identity.server_id == "srv-456" + assert opened.identity.subject_type == "user_id" + assert opened.identity.subject == "user-42" + + def test_public_models_are_frozen(): sealed = mint_envelope(_IDENTITY, _full_grant(), _KEYS, _NOW) assert isinstance(sealed, SealedEnvelope) @@ -484,4 +591,4 @@ def test_public_models_are_frozen(): with pytest.raises(ValidationError): opened.grant = _minimal_grant() with pytest.raises(ValidationError): - _IDENTITY.key_hash = "someone-elses-hash" + _IDENTITY.subject = "someone-elses-hash" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 68466e624ec..47a49b96327 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4274,6 +4274,9 @@ async def test_register_bridge_relay_surfaces_upstream_error_not_500(): error_response = MagicMock() error_response.status_code = 400 error_response.text = '{"error":"invalid_redirect_uri","error_description":"redirect_uri not allowed"}' + error_response.json = MagicMock( + return_value={"error": "invalid_redirect_uri", "error_description": "redirect_uri not allowed"} + ) error_response.raise_for_status = MagicMock( side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response) ) @@ -4307,9 +4310,10 @@ async def test_register_bridge_relay_surfaces_upstream_error_not_500(): @pytest.mark.asyncio -async def test_register_non_bridge_upstream_error_still_raises_500(): - """Non-bridge DCR keeps its pre-change behavior: raise_for_status propagates so the flag-off - contract is byte-identical; only the bridge relay arm relays the upstream status.""" +async def test_register_non_bridge_upstream_error_relays_status_not_500(): + """A non-bridge DCR rejection must relay the upstream status and RFC 7591 error body just like + the bridge relay arm; a raw HTTPStatusError would escape to the global handler and surface as an + opaque 500 that hides the real reason from the create-flow UI.""" import httpx from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -4319,6 +4323,7 @@ async def test_register_non_bridge_upstream_error_still_raises_500(): error_response = MagicMock() error_response.status_code = 400 error_response.text = '{"error":"invalid_client_metadata"}' + error_response.json = MagicMock(return_value={"error": "invalid_client_metadata"}) error_response.raise_for_status = MagicMock( side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response) ) @@ -4338,7 +4343,7 @@ async def test_register_non_bridge_upstream_error_still_raises_500(): return_value=False, ), ): - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(HTTPException) as exc: await register_client_with_server( request=_bridge_mock_request(), mcp_server=oauth2_server, @@ -4348,6 +4353,9 @@ async def test_register_non_bridge_upstream_error_still_raises_500(): token_endpoint_auth_method=None, ) + assert exc.value.status_code == 400 + assert "invalid_client_metadata" in str(exc.value.detail) + @pytest.mark.asyncio async def test_register_bridge_relay_never_persists(): @@ -4365,7 +4373,7 @@ async def test_register_bridge_relay_never_persists(): _BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" -async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_client_out=None): +async def _exchange_for_bridge_server(server, upstream_body, key_hash, code="auth-code", fake_client_out=None): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _ResolvedKey, exchange_token_with_server, @@ -4398,13 +4406,17 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie request=_bridge_mock_request(), mcp_server=server, grant_type="authorization_code", - code="auth-code", + code=code, redirect_uri="https://claude.ai/api/mcp/auth_callback", client_id="dcr-client-123", client_secret=None, code_verifier="verifier", ) - if server.is_oauth_delegate and server.is_dcr_bridge: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import is_bridge_authorization_code + + # The key_hash path resolves the presented litellm key; the interactive SSO path recovers identity + # from the gateway authorization code instead, so it never awaits the resolver. + if server.is_oauth_delegate and server.is_dcr_bridge and not is_bridge_authorization_code(code): key_resolver.assert_awaited_once() else: key_resolver.assert_not_awaited() @@ -4441,10 +4453,198 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) opened = resolve_bridge_envelope(token, keys, datetime.now(timezone.utc), server.server_id) assert isinstance(opened, BridgeEnvelopeAdmitted) - assert opened.identity.key_hash == "hashed-litellm-key-77" + assert opened.identity.subject_type == "key_hash" + assert opened.identity.subject == "hashed-litellm-key-77" assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN" +def test_bridge_authorization_code_round_trips_and_rejects_hostile_input(): + """The gateway authorization code seals and recovers the upstream code and the SSO user, and is + total over hostile input: a raw upstream code (scripted path) opens to None, and a tampered or + non-gateway value opens to None rather than raising.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + is_bridge_authorization_code, + open_bridge_authorization_code, + seal_bridge_authorization_code, + ) + + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + sealed = seal_bridge_authorization_code( + upstream_code="up-code", litellm_user_id="sso-user-9", mcp_server_id="srv-1" + ) + assert is_bridge_authorization_code(sealed) + opened = open_bridge_authorization_code(sealed) + assert opened is not None + assert opened.upstream_code == "up-code" + assert opened.litellm_user_id == "sso-user-9" + assert opened.mcp_server_id == "srv-1" + assert open_bridge_authorization_code("raw-upstream-code") is None + assert open_bridge_authorization_code(sealed[:-4] + "aaaa") is None + + +@pytest.mark.asyncio +async def test_interactive_bridge_token_exchange_mints_user_subject_envelope(): + """An interactive dcr_bridge oauth_delegate exchange (the client presents the gateway code the + callback sealed, and NO litellm key) mints an envelope bound to the SSO-captured user: it opens + to a user_id subject, and the upstream exchange used the real upstream code recovered from the + gateway code, not the sealed wrapper.""" + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_bridge_authorization_code, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + BridgeEnvelopeAdmitted, + envelope_keys_from_master_key, + resolve_bridge_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + gateway_code = seal_bridge_authorization_code( + upstream_code="REAL-UPSTREAM-CODE", litellm_user_id="sso-user-42", mcp_server_id=server.server_id + ) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + captured: dict = {} + response = await _exchange_for_bridge_server( + server, upstream, key_hash=None, code=gateway_code, fake_client_out=captured + ) + + token = json.loads(response.body)["access_token"] + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = resolve_bridge_envelope(token, keys, datetime.now(timezone.utc), server.server_id) + assert isinstance(opened, BridgeEnvelopeAdmitted) + assert opened.identity.subject_type == "user_id" + assert opened.identity.subject == "sso-user-42" + assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN" + assert captured["client"].post.call_args.kwargs["data"]["code"] == "REAL-UPSTREAM-CODE" + + +@pytest.mark.asyncio +async def test_interactive_bridge_gateway_code_for_another_server_is_rejected_400(): + """A gateway authorization code is bound to the server it was minted for: presenting it at another + server's token endpoint is a 400, so a code cannot be replayed across a server boundary.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_bridge_authorization_code, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + gateway_code = seal_bridge_authorization_code( + upstream_code="up-code", litellm_user_id="sso-user-42", mcp_server_id="a-different-server-id" + ) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + with pytest.raises(HTTPException) as exc: + await _exchange_for_bridge_server(server, upstream, key_hash=None, code=gateway_code) + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_interactive_bridge_authorize_seals_sso_user_into_state(): + """On the short-circuit bridge oauth_delegate arm, authorize captures the SSO user from the UI + session cookie and seals it (and the target server) into the encrypted OAuth state, so the + callback can later mint a user-bound gateway code; it still proceeds to the upstream redirect.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate, client_id="admin-client", registration_url=None) + captured: dict = {} + + def _capture(**kwargs): + captured.update(kwargs) + return "mocked_encrypted_state" + + with ( + patch( + "litellm.proxy._experimental.mcp_server.byok_oauth_endpoints._user_id_from_session_cookie", + return_value="sso-user-42", + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encode_state_with_base_url", + side_effect=_capture, + ), + ): + response = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="ignored", + redirect_uri="http://127.0.0.1:60108/callback", + state="s", + code_challenge="chal", + code_challenge_method="S256", + ) + + assert captured["litellm_user_id"] == "sso-user-42" + assert captured["mcp_server_id"] == server.server_id + assert "/sso/key/generate" not in response.headers["location"] + + +@pytest.mark.asyncio +async def test_interactive_bridge_authorize_without_session_redirects_to_login(): + """Without a UI session there is no identity to bind, so the short-circuit bridge oauth_delegate + authorize sends the browser through litellm login instead of proceeding to the upstream.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate, client_id="admin-client", registration_url=None) + with patch( + "litellm.proxy._experimental.mcp_server.byok_oauth_endpoints._user_id_from_session_cookie", + return_value=None, + ): + response = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="ignored", + redirect_uri="http://127.0.0.1:60108/callback", + state="s", + code_challenge="chal", + code_challenge_method="S256", + ) + assert "/sso/key/generate" in response.headers["location"] + + +@pytest.mark.asyncio +async def test_interactive_bridge_callback_seals_user_into_gateway_code(): + """When the OAuth state carries the captured SSO user, the callback forwards a gateway + authorization code (sealing the user and upstream code) to the client instead of the raw upstream + code, so the client's later token call can prove who signed in.""" + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + is_bridge_authorization_code, + ) + + state_data = { + "original_state": "client-state", + "client_redirect_uri": "http://127.0.0.1:60108/cb", + "base_url": "http://127.0.0.1:60108/cb", + "litellm_user_id": "sso-user-42", + "mcp_server_id": "bridge_srv", + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_encoded_oauth_state", + return_value="enc", + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash", + return_value=state_data, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._get_validated_client_redirect_uri", + return_value="http://127.0.0.1:60108/cb", + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await callback(request=_bridge_mock_request(), code="REAL-UPSTREAM-CODE", state="relay") + + forwarded_code = parse_qs(urlparse(response.headers["location"]).query)["code"][0] + assert is_bridge_authorization_code(forwarded_code) + + @pytest.mark.asyncio async def test_oauth_delegate_bridge_token_exchange_fails_closed_without_litellm_identity(): """Without a resolvable litellm identity on the token request, the exchange must not mint an @@ -4480,11 +4680,11 @@ async def test_bridge_envelope_too_large_upstream_token_is_502(): @pytest.mark.asyncio -async def test_bridge_envelope_does_not_seal_upstream_refresh_token(): - """The upstream refresh_token is never sealed into the client-held envelope: the edge never - consumes it and a long-lived upstream credential should not live in the client bearer. The opened - envelope's grant carries no refresh token even when the upstream returned one, and neither does - the response body.""" +async def test_bridge_access_envelope_never_carries_upstream_refresh_token(): + """The upstream refresh token is never sealed into the ACCESS envelope, the bearer forwarded upstream + on every tool call: the opened access grant carries no refresh token even when the upstream returned + one, and the raw refresh token never appears in the access envelope. It rides only in the separate + refresh envelope returned as the response's refresh_token, encrypted, never in plaintext.""" from datetime import datetime, timezone from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( @@ -4506,21 +4706,22 @@ async def test_bridge_envelope_does_not_seal_upstream_refresh_token(): response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") body = json.loads(response.body) - assert "refresh_token" not in body assert "UPSTREAM-REFRESH" not in body["access_token"] keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) opened = open_envelope(body["access_token"], keys, datetime.now(timezone.utc)) assert isinstance(opened, OpenedEnvelope) assert opened.grant.refresh_token is None + # the refresh token rides only in the separate, encrypted refresh envelope, never in plaintext + assert body["refresh_token"].startswith("llm_refresh_") + assert "UPSTREAM-REFRESH" not in body["refresh_token"] @pytest.mark.asyncio -async def test_bridge_refresh_grant_is_rejected_before_upstream(): - """A bridge oauth_delegate server issues only envelopes and seals no upstream refresh_token, so the - client never holds one to present. _prepare_bridge_mint rejects the refresh_token grant up front - with unsupported_grant_type, BEFORE any upstream exchange, so a stray refresh request can never - rotate or consume the client's upstream refresh credential; renewal is re-running - authorization_code. This is checked before identity resolution, so it holds even with a valid key.""" +async def test_bridge_refresh_grant_with_non_envelope_is_invalid_grant_before_upstream(): + """A bridge oauth_delegate client only ever holds a refresh envelope, never a raw upstream refresh + token, so a refresh_token grant carrying a bare (non-envelope) value is invalid_grant, rejected in + _prepare_bridge_refresh BEFORE any upstream exchange. Rejecting before the exchange means a bad + refresh request can never consume or rotate an upstream refresh token.""" from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server from litellm.types.mcp import MCPAuth @@ -4547,10 +4748,564 @@ async def test_bridge_refresh_grant_is_rejected_before_upstream(): ) assert response.status_code == 400 - assert json.loads(response.body)["error"] == "unsupported_grant_type" + assert json.loads(response.body)["error"] == "invalid_grant" fake_http_client.post.assert_not_called() +def _mint_test_refresh_envelope( + server_id="bridge_srv", key_hash="hashed-litellm-key-77", upstream_refresh="UPSTREAM-REFRESH", identity=None, + scope=None, +): + """Mint a refresh envelope the way the producer does, for driving the refresh_token grant in tests. + Defaults to a key_hash subject; pass ``identity`` to seal a specific subject (e.g. a user_id), and + ``scope`` to seal the scope to re-request on refresh.""" + from datetime import datetime, timezone + + from pydantic import SecretStr + + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + build_bridge_refresh_token_response, + envelope_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + RefreshCredential, + SealedEnvelope, + key_hash_identity, + ) + + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + identity = identity if identity is not None else key_hash_identity(server_id=server_id, key_hash=key_hash) + sealed = build_bridge_refresh_token_response( + identity, RefreshCredential(refresh_token=SecretStr(upstream_refresh), scope=scope), keys, + datetime.now(timezone.utc), + ) + assert isinstance(sealed, SealedEnvelope) + return sealed.token.get_secret_value() + + +async def _refresh_for_bridge_server( + server, refresh_envelope_value, upstream_body, revalidate_result=None, fake_client_out=None +): + """Drive a refresh_token grant for a bridge server: the client presents ``refresh_envelope_value``, + the sealed subject re-validates to ``revalidate_result`` (``None`` when the key or user is still + active, or a failure literal like "no_active_key" when revoked/deactivated), and the upstream returns + ``upstream_body``. Patching the single subject-revalidation dispatch covers both a key_hash and a + user_id refresh envelope. Returns the response; the captured client exposes the POST call so a test + can assert what refresh token was actually sent upstream.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + + fake_http_response = MagicMock() + fake_http_response.json.return_value = upstream_body + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + if fake_client_out is not None: + fake_client_out["client"] = fake_http_client + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._revalidate_active_subject", + new=AsyncMock(return_value=revalidate_result), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + return await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="dcr-client-123", + client_secret=None, + code_verifier=None, + refresh_token=refresh_envelope_value, + ) + + +@pytest.mark.asyncio +async def test_bridge_mint_returns_refresh_envelope_that_opens_to_upstream_refresh(): + """When the upstream returns a refresh token, the authorization_code mint returns a refresh envelope + alongside the access envelope. The refresh envelope is a distinct llm_refresh_ credential that opens + (under the same keys and server_id) to the upstream refresh token, so the client can renew later.""" + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + envelope_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + OpenedRefreshEnvelope, + open_refresh_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 3600, "refresh_token": "R-UP"} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + refresh_env = body["refresh_token"] + assert refresh_env.startswith("llm_refresh_") + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = open_refresh_envelope(refresh_env, keys, datetime.now(timezone.utc)) + assert isinstance(opened, OpenedRefreshEnvelope) + assert opened.identity.server_id == server.server_id + assert opened.refresh.refresh_token.get_secret_value() == "R-UP" + + +@pytest.mark.asyncio +async def test_bridge_mint_omits_refresh_envelope_when_upstream_has_no_refresh(): + """No refresh envelope is issued when the upstream returns no refresh token, so the response carries + only the access envelope; the client re-authenticates at access expiry (nothing to renew with).""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + assert "refresh_token" not in body + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_sends_unwrapped_upstream_token_and_renews(): + """A refresh_token grant carrying a valid refresh envelope renews: the exchange unwraps the envelope + and sends the REAL upstream refresh token upstream (never the envelope), then returns a fresh access + envelope. This is the flow that lets the client renew without re-authenticating.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="UPSTREAM-REFRESH") + upstream = {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600} + captured: dict = {} + response = await _refresh_for_bridge_server(server, refresh_env, upstream, None, fake_client_out=captured) + + assert response.status_code == 200 + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + # the upstream exchange received the unwrapped upstream refresh token, never the client's envelope + sent = captured["client"].post.call_args.kwargs["data"] + assert sent["grant_type"] == "refresh_token" + assert sent["refresh_token"] == "UPSTREAM-REFRESH" + assert not sent["refresh_token"].startswith("llm_refresh_") + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_rotates_refresh_envelope_wrapping_new_upstream_token(): + """When the upstream rotates the refresh token on renewal, the client receives a new refresh envelope + that wraps the NEW upstream refresh token, so the rotation is carried through faithfully.""" + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + envelope_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + OpenedRefreshEnvelope, + open_refresh_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="OLD-UP-REFRESH") + upstream = { + "access_token": "NEW-ACCESS", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "NEW-UP-REFRESH", + } + response = await _refresh_for_bridge_server(server, refresh_env, upstream, None) + + body = json.loads(response.body) + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = open_refresh_envelope(body["refresh_token"], keys, datetime.now(timezone.utc)) + assert isinstance(opened, OpenedRefreshEnvelope) + assert opened.refresh.refresh_token.get_secret_value() == "NEW-UP-REFRESH" + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_with_revoked_key_is_invalid_grant_before_upstream(): + """A valid refresh envelope whose sealed litellm key has since been revoked cannot keep refreshing: + the reload gate reports no_active_key and the refresh is invalid_grant, returned BEFORE the upstream + exchange so the upstream refresh token is never consumed. Revocation kills renewal.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id) + captured: dict = {} + response = await _refresh_for_bridge_server( + server, refresh_env, {"access_token": "NEW"}, "no_active_key", fake_client_out=captured + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + captured["client"].post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_refresh_envelope_for_another_server_is_invalid_grant(): + """A refresh envelope minted for one server cannot renew against another: the sealed server_id must + match the server the refresh targets, so a cross-server refresh envelope is invalid_grant and never + reaches the upstream exchange.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + foreign_env = _mint_test_refresh_envelope(server_id="some-other-server") + captured: dict = {} + response = await _refresh_for_bridge_server( + server, foreign_env, {"access_token": "NEW"}, None, fake_client_out=captured + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + captured["client"].post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_renews_a_user_subject_envelope(): + """The interactive SSO client mints a user_id-subject envelope, so its refresh envelope carries a + user subject too. Renewing it re-validates the user (still active here), unwraps the upstream refresh + token, and returns a fresh access envelope that opens back to the same user_id subject; the upstream + exchange received the real upstream refresh token, not the client's envelope.""" + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + BridgeEnvelopeAdmitted, + envelope_keys_from_master_key, + resolve_bridge_envelope, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import user_identity + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + user_env = _mint_test_refresh_envelope( + identity=user_identity(server_id=server.server_id, user_id="sso-user-42"), upstream_refresh="UP-REFRESH-USER" + ) + upstream = {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600} + captured: dict = {} + response = await _refresh_for_bridge_server(server, user_env, upstream, None, fake_client_out=captured) + + assert response.status_code == 200 + body = json.loads(response.body) + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = resolve_bridge_envelope(body["access_token"], keys, datetime.now(timezone.utc), server.server_id) + assert isinstance(opened, BridgeEnvelopeAdmitted) + assert opened.identity.subject_type == "user_id" + assert opened.identity.subject == "sso-user-42" + assert captured["client"].post.call_args.kwargs["data"]["refresh_token"] == "UP-REFRESH-USER" + + +@pytest.mark.asyncio +async def test_bridge_refresh_re_requests_the_sealed_scope_when_client_omits_it(): + """A DCR/MCP client omits scope on the refresh request, so the gateway must re-request the scope sealed + at mint; dropping it lets a stricter upstream narrow the renewed token. The upstream POST must carry + the sealed scope even though the client sent none. Regression for the dropped sealed refresh scope.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope( + server_id=server.server_id, upstream_refresh="UP-REFRESH", scope="read:tools write:tools" + ) + captured: dict = {} + response = await _refresh_for_bridge_server( + server, refresh_env, {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600}, None, + fake_client_out=captured, + ) + + assert response.status_code == 200 + assert captured["client"].post.call_args.kwargs["data"]["scope"] == "read:tools write:tools" + + +@pytest.mark.asyncio +async def test_bridge_refresh_re_seals_scope_when_upstream_omits_it_so_the_chain_keeps_it(): + """RFC 6749 5.1 lets an upstream omit scope in a refresh response when it is unchanged. The re-minted + refresh envelope must still seal the scope that was requested, otherwise the NEXT refresh loses it and + a stricter upstream could narrow the token. The returned refresh envelope carries the scope even though + the upstream response had none, and a second refresh off it still re-requests the scope.""" + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + envelope_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + OpenedRefreshEnvelope, + open_refresh_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope( + server_id=server.server_id, upstream_refresh="UP-1", scope="mcp:read mcp:write" + ) + # the upstream rotates the refresh token but OMITS scope (valid when unchanged) + upstream_no_scope = {"access_token": "NEW", "token_type": "Bearer", "expires_in": 3600, "refresh_token": "UP-2"} + + captured: dict = {} + r1 = await _refresh_for_bridge_server(server, refresh_env, upstream_no_scope, None, fake_client_out=captured) + assert r1.status_code == 200 + assert captured["client"].post.call_args.kwargs["data"]["scope"] == "mcp:read mcp:write" + + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + new_env = json.loads(r1.body)["refresh_token"] + opened = open_refresh_envelope(new_env, keys, datetime.now(timezone.utc)) + assert isinstance(opened, OpenedRefreshEnvelope) + assert opened.refresh.scope == "mcp:read mcp:write" + + captured2: dict = {} + r2 = await _refresh_for_bridge_server(server, new_env, upstream_no_scope, None, fake_client_out=captured2) + assert r2.status_code == 200 + assert captured2["client"].post.call_args.kwargs["data"]["scope"] == "mcp:read mcp:write" + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_with_deactivated_user_is_invalid_grant_before_upstream(): + """A user_id-subject refresh envelope whose user has since been deactivated (SCIM offboarding, or + the user no longer exists) cannot keep refreshing: subject re-validation reports no_active_key and + the refresh is invalid_grant, returned BEFORE the upstream exchange. Revocation kills renewal for the + user subject exactly as it does for the key subject.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import user_identity + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + user_env = _mint_test_refresh_envelope(identity=user_identity(server_id=server.server_id, user_id="gone-user")) + captured: dict = {} + response = await _refresh_for_bridge_server( + server, user_env, {"access_token": "NEW"}, "no_active_key", fake_client_out=captured + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + captured["client"].post.assert_not_called() + + +@pytest.mark.asyncio +async def test_revalidate_active_subject_dispatches_on_subject_type(): + """Subject re-validation routes a key_hash envelope to the key reload and a user_id envelope to the + user reload, so revocation gates renewal for either identity source through one dispatch point.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _ResolvedKey, + _revalidate_active_subject, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity, user_identity + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + new=AsyncMock(return_value=_ResolvedKey(key_hash="kh", key=MagicMock())), + ) as key_reload, + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_user_by_id", + new=AsyncMock(return_value=None), + ) as user_reload, + ): + assert await _revalidate_active_subject(key_hash_identity(server_id="s", key_hash="kh")) is None + key_reload.assert_awaited_once_with("kh") + user_reload.assert_not_awaited() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + new=AsyncMock(), + ) as key_reload2, + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_user_by_id", + new=AsyncMock(return_value="no_active_key"), + ) as user_reload2, + ): + assert await _revalidate_active_subject(user_identity(server_id="s", user_id="u42")) == "no_active_key" + user_reload2.assert_awaited_once_with("u42") + key_reload2.assert_not_awaited() + + +def test_upstream_refresh_credential_expired_refresh_token_is_not_sealed(): + """An upstream that reports its refresh token already elapsed (refresh_expires_in non-positive) must + not be sealed: _upstream_refresh_credential returns None so the exchange degrades to an access-only + response, mirroring how the access grant refuses an already-elapsed access token rather than capping a + dead token to the full refresh TTL. A live or unspecified lifetime still yields a credential.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _upstream_refresh_credential + + assert _upstream_refresh_credential({"access_token": "A", "refresh_token": "R", "refresh_expires_in": 0}) is None + assert _upstream_refresh_credential({"refresh_token": "R", "refresh_expires_in": -5}) is None + live = _upstream_refresh_credential({"refresh_token": "R", "refresh_expires_in": 1800}) + assert live is not None and live.expires_in == 1800 + unspecified = _upstream_refresh_credential({"refresh_token": "R"}) + assert unspecified is not None and unspecified.expires_in is None + + +@pytest.mark.asyncio +async def test_bridge_refresh_upstream_invalid_grant_maps_to_invalid_grant(): + """When the sealed upstream refresh token has been revoked or expired at the IdP, the upstream returns + 400 invalid_grant. The bridge refresh path maps that to an RFC 6749 invalid_grant response so the OAuth + client re-runs authorization_code, rather than surfacing the opaque upstream error it cannot act on.""" + import httpx + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="LIVE-ENVELOPE-REFRESH") + + error_response = MagicMock() + error_response.status_code = 400 + error_response.text = '{"error": "invalid_grant", "error_description": "refresh token expired"}' + error_response.json = MagicMock(return_value={"error": "invalid_grant", "error_description": "refresh token expired"}) + error_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response) + ) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=error_response) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._revalidate_active_subject", + new=AsyncMock(return_value=None), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="dcr-client-123", + client_secret=None, + code_verifier=None, + refresh_token=refresh_env, + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + + +@pytest.mark.asyncio +async def test_bridge_refresh_upstream_error_detection_parses_json_not_substring(): + """The upstream invalid_grant detection reads the classified RFC 6749 5.2 error code, not a substring + of the body. An upstream error whose code is not invalid_grant (here invalid_client, with the string + invalid_grant only inside error_description) must NOT be mistaken for a dead refresh token: it renders + as the classified upstream rejection rather than triggering a spurious authorization_code re-run.""" + import httpx + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="UP") + + error_response = MagicMock() + error_response.status_code = 400 + error_response.text = '{"error": "invalid_client", "error_description": "this is not an invalid_grant problem"}' + error_response.json = MagicMock( + return_value={"error": "invalid_client", "error_description": "this is not an invalid_grant problem"} + ) + error_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response) + ) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=error_response) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._revalidate_active_subject", + new=AsyncMock(return_value=None), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="dcr-client-123", + client_secret=None, + code_verifier=None, + refresh_token=refresh_env, + ) + + assert response.status_code == 401 + body = json.loads(response.body) + assert body["error"] == "invalid_client" + + +@pytest.mark.asyncio +async def test_revalidate_key_subject_revoked_when_owner_scim_deactivated(proxy_globals): + """A key_hash refresh envelope whose key is still active but whose OWNING user was SCIM-deactivated must + fail closed to no_active_key, mirroring how admission's _reject_if_admitted_owner_scim_deactivated + revokes an offboarded owner's key. Without this, an offboarded user keeps renewing a live key.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _ResolvedKey, _revalidate_active_subject + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + resolved = _ResolvedKey(key_hash="kh", key=MagicMock(user_id="offboarded-owner")) + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + new=AsyncMock(return_value=resolved), + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(return_value=MagicMock(metadata={"scim_active": False})), + ), + ): + result = await _revalidate_active_subject(key_hash_identity(server_id="s", key_hash="kh")) + + assert result == "no_active_key" + + +@pytest.mark.asyncio +async def test_revalidate_key_subject_active_owner_renews_and_missing_owner_fails_open(proxy_globals): + """The key-owner SCIM gate blocks only an explicit scim_active False: an active owner renews (None), and + a missing owner (get_user_object's wrapped ValueError) fails OPEN, since a key may outlive its owner + record and a transient blip must not revoke a live key.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _ResolvedKey, _revalidate_active_subject + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + resolved = _ResolvedKey(key_hash="kh", key=MagicMock(user_id="live-owner")) + identity = key_hash_identity(server_id="s", key_hash="kh") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + new=AsyncMock(return_value=resolved), + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(return_value=MagicMock(metadata={"scim_active": True})), + ), + ): + assert await _revalidate_active_subject(identity) is None + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + new=AsyncMock(return_value=resolved), + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=_wrapped_user_lookup_error(Exception())), + ), + ): + assert await _revalidate_active_subject(identity) is None + + @pytest.mark.asyncio async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset(): """master_key is validated BEFORE the upstream exchange (in _prepare_bridge_mint), so a @@ -4726,10 +5481,11 @@ def test_bridge_reported_expires_in_can_be_zero_at_jwt_exp_boundary(): from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( envelope_keys_from_master_key, ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity from litellm.types.mcp import MCPAuth ready = _BridgeMintReady( - key_hash="hashed-litellm-key-77", + identity=key_hash_identity(server_id="bridge_srv", key_hash="hashed-litellm-key-77"), keys=envelope_keys_from_master_key(_BRIDGE_MASTER_KEY), ) response = _finish_bridge_mint( @@ -5360,6 +6116,60 @@ async def test_resolve_active_litellm_key_no_database_is_unresolvable(proxy_glob assert await _resolve_active_litellm_key(request) == "unresolvable" +def _wrapped_user_lookup_error(original: BaseException) -> ValueError: + """Reproduce get_user_object's real exception contract (litellm/proxy/auth/auth_checks.py): it + catches every DB failure in a broad ``except`` and re-raises a bare ``ValueError``, so the original + error (a missing-user Exception or a real outage) survives only as ``__context__``. Injecting a raw + ConnectionError/Exception instead would exercise a shape production never produces and let a + chain-blind outage classifier pass. The wrapping fidelity is pinned by + test_get_user_object_wraps_db_outage_as_valueerror_preserving_context in test_auth_checks.""" + try: + raise original + except BaseException: + try: + raise ValueError(f"User doesn't exist in db. Got error - {original}") + except ValueError as wrapped: + return wrapped + + +@pytest.mark.asyncio +async def test_reload_active_user_by_id_missing_user_is_no_active_key(proxy_globals): + """A user_id refresh envelope whose user has been deleted must fail closed to no_active_key (the + refresh path maps it to invalid_grant), not unresolvable/500. get_user_object catches the missing row + and re-raises a bare ValueError, so a missing user must not be misclassified as a DB outage or an + opaque gateway fault.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _reload_active_user_by_id + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + with patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=_wrapped_user_lookup_error(Exception())), + ): + assert await _reload_active_user_by_id("gone-user") == "no_active_key" + + +@pytest.mark.asyncio +async def test_reload_active_user_by_id_db_outage_is_unavailable(proxy_globals): + """A transient DB outage while re-validating the user on refresh is a retryable outage, distinct from + a missing user, so the refresh path surfaces "unavailable" (a 503) rather than blaming the caller. + get_user_object wraps the outage in a bare ValueError, so this exercises the chain-aware classifier; a + raw ConnectionError would falsely pass even a chain-blind check because it is an OSError.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _reload_active_user_by_id + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + with patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=_wrapped_user_lookup_error(ConnectionError("user database unreachable"))), + ): + assert await _reload_active_user_by_id("sso-user-7") == "unavailable" + + @pytest.mark.asyncio async def test_token_endpoint_uses_client_secret_basic_when_configured(): """LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the @@ -5960,3 +6770,364 @@ async def test_token_exchange_pairs_client_secret_with_server_client_id(): sent = mock_async_client.post.call_args.kwargs["data"] assert sent["client_id"] == "persisted-client" assert "client_secret" not in sent + + +def _upstream_token_response(status_code: int, *, json_body: object = None, text_body: str = "") -> "httpx.Response": + import httpx + + request = httpx.Request("POST", "https://oauth2.googleapis.com/token") + if json_body is not None: + return httpx.Response(status_code, json=json_body, request=request) + return httpx.Response(status_code, text=text_body, request=request) + + +async def _exchange_with_upstream_response(upstream_response, *, server_client_id="web-client.apps.googleusercontent.com"): + """Run the raw (non-bridge) authorization_code exchange against a canned upstream token-endpoint + response and return what the gateway would hand the client. ``server_client_id=None`` models the + caller-supplied-credentials flow (no stored client on the server).""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="gcal", + name="gcal", + server_name="gcal", + alias="gcal", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=server_client_id, + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + ) + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=upstream_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ): + return await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://litellm.example.com/ui/mcp/oauth/callback", + client_id="web-client.apps.googleusercontent.com", + client_secret=None, + code_verifier="verifier", + ) + + +@pytest.mark.asyncio +async def test_token_exchange_gateway_credential_rejection_is_502_with_gateway_prose(): + """When the gateway presented the server's stored client credentials and the IdP rejected them + (Google refusing a secret-less or unknown client), the fault is the operator's, not the caller's: + 502 server_error with gateway-authored prose naming the code, and the IdP's own prose stays in + server logs. Before the framework this either 500ed raw or relayed provider prose verbatim.""" + response = await _exchange_with_upstream_response( + _upstream_token_response( + 401, json_body={"error": "invalid_client", "error_description": "The OAuth client was not found."} + ) + ) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body["error"] == "server_error" + assert "invalid_client" in body["error_description"] + assert "client_id and client_secret" in body["error_description"] + assert "The OAuth client was not found." not in body["error_description"] + assert response.headers["cache-control"] == "no-store" + + +@pytest.mark.asyncio +async def test_token_exchange_caller_supplied_credential_rejection_relays_code(): + """When the caller supplied the client credentials themselves (no stored client on the server), + an invalid_client rejection is theirs to act on: the §5.2 code relays on the 401 that code + implies.""" + response = await _exchange_with_upstream_response( + _upstream_token_response( + 401, json_body={"error": "invalid_client", "error_description": "The OAuth client was not found."} + ), + server_client_id=None, + ) + + assert response.status_code == 401 + body = json.loads(response.body) + assert body == {"error": "invalid_client", "error_description": "The OAuth client was not found."} + + +@pytest.mark.asyncio +async def test_token_exchange_status_derives_from_error_code_not_upstream_status(): + """An upstream that pairs a caller-fault code with a server-fault status (invalid_grant on a 500) + must not produce a contradictory response: status derives from the classified fault, so the + caller sees 400 invalid_grant and knows to re-authorize rather than blaming the gateway.""" + response = await _exchange_with_upstream_response( + _upstream_token_response(500, json_body={"error": "invalid_grant", "error_description": "Code expired."}) + ) + + assert response.status_code == 400 + body = json.loads(response.body) + assert body == {"error": "invalid_grant", "error_description": "Code expired."} + + +@pytest.mark.asyncio +async def test_token_exchange_relays_only_rfc6749_error_fields(): + """Only error / error_description / error_uri cross the gateway; any other upstream body field is + dropped so an arbitrary rejection payload cannot ride the relay to the client.""" + response = await _exchange_with_upstream_response( + _upstream_token_response( + 400, + json_body={ + "error": "invalid_grant", + "error_description": "Code was already redeemed.", + "error_uri": "https://idp.example.com/errors/invalid_grant", + "internal_trace": "should never reach the client", + }, + ) + ) + + assert response.status_code == 400 + body = json.loads(response.body) + assert set(body.keys()) == {"error", "error_description", "error_uri"} + + +@pytest.mark.asyncio +async def test_token_exchange_maps_out_of_contract_rejection_to_502(): + """A rejection outside the §5.2 contract (no JSON error field, or a status the token-endpoint + contract does not define) is an upstream fault; 502 keeps it from being misread as a caller + mistake while the description still names the upstream status.""" + response = await _exchange_with_upstream_response( + _upstream_token_response(503, text_body="upstream maintenance") + ) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body["error"] == "server_error" + assert "HTTP 503" in body["error_description"] + + +@pytest.mark.asyncio +async def test_token_exchange_bounds_relayed_error_fields(): + """Relayed §5.2 fields are length-bounded so a hostile or broken upstream cannot bloat the + gateway response.""" + response = await _exchange_with_upstream_response( + _upstream_token_response(400, json_body={"error": "invalid_request", "error_description": "x" * 5000}) + ) + + assert response.status_code == 400 + body = json.loads(response.body) + assert len(body["error_description"]) == 500 + + +@pytest.mark.asyncio +async def test_token_exchange_200_without_access_token_is_502_not_keyerror(): + """A 200 whose body has no usable access_token used to KeyError into a 500; the raw arm now + answers 502 with the same wording as the bridge arm's no_upstream_token rejection.""" + response = await _exchange_with_upstream_response( + _upstream_token_response(200, json_body={"token_type": "Bearer"}) + ) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body["error"] == "server_error" + assert "access_token" in body["error_description"] + + +@pytest.mark.asyncio +async def test_token_exchange_relays_rejection_when_http_client_raises(): + """litellm's AsyncHTTPHandler.post raise_for_status()es internally and raises MaskedHTTPStatusError + at call time, so in production the rejection escapes from the post call itself rather than from the + explicit raise_for_status; the relay must catch it there too (proven live: a mock returning the + error response passed while the real proxy still 500ed).""" + import httpx + + rejection = _upstream_token_response( + 401, json_body={"error": "invalid_client", "error_description": "The OAuth client was not found."} + ) + raising_client = MagicMock() + raising_client.post = AsyncMock( + side_effect=httpx.HTTPStatusError("Client error '401 Unauthorized'", request=rejection.request, response=rejection) + ) + + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="gcal", + name="gcal", + server_name="gcal", + alias="gcal", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="web-client.apps.googleusercontent.com", + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + ) + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=raising_client, + ): + response = await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://litellm.example.com/ui/mcp/oauth/callback", + client_id="web-client.apps.googleusercontent.com", + client_secret=None, + code_verifier="verifier", + ) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body["error"] == "server_error" + assert "invalid_client" in body["error_description"] + + +@pytest.mark.asyncio +async def test_register_relays_rejection_when_http_client_raises(): + """Same live mechanism as the token exchange: the DCR rejection escapes from the post call itself, + so the register relay must catch it there, not only from the explicit raise_for_status.""" + import httpx + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client_with_server, + ) + + rejection = httpx.Response( + 400, + json={"error": "invalid_client_metadata"}, + request=httpx.Request("POST", "https://idp.example.com/register"), + ) + raising_client = MagicMock() + raising_client.post = AsyncMock( + side_effect=httpx.HTTPStatusError("Client error '400 Bad Request'", request=rejection.request, response=rejection) + ) + + oauth2_server = _bridge_server(auth_type=MCPAuth.oauth2, dcr_bridge=None) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=raising_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reuse_persisted_dcr_client_if_available", + new_callable=AsyncMock, + return_value=False, + ), + ): + with pytest.raises(HTTPException) as exc: + await register_client_with_server( + request=_bridge_mock_request(), + mcp_server=oauth2_server, + client_name="Claude", + grant_types=None, + response_types=None, + token_endpoint_auth_method=None, + ) + + assert exc.value.status_code == 400 + assert "invalid_client_metadata" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_token_exchange_never_relays_out_of_contract_body_to_client(): + """These endpoints serve unauthenticated OAuth clients, so a non-RFC6749 upstream body (HTML + error page, proxy banner, stack trace) must stay in server logs; the client sees only the + upstream status.""" + response = await _exchange_with_upstream_response( + _upstream_token_response(404, text_body="Error 404 stack trace: secret internals") + ) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 404"} + + +@pytest.mark.asyncio +async def test_register_never_relays_out_of_contract_body_to_client(): + """Same trust boundary for DCR: a non-RFC7591 rejection body is logged server-side and the + client detail names only the status.""" + import httpx + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client_with_server, + ) + + rejection = httpx.Response( + 500, + text="Tomcat stack trace with internals", + request=httpx.Request("POST", "https://idp.example.com/register"), + ) + raising_client = MagicMock() + raising_client.post = AsyncMock( + side_effect=httpx.HTTPStatusError("Server error '500'", request=rejection.request, response=rejection) + ) + + oauth2_server = _bridge_server(auth_type=MCPAuth.oauth2, dcr_bridge=None) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=raising_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reuse_persisted_dcr_client_if_available", + new_callable=AsyncMock, + return_value=False, + ), + ): + with pytest.raises(HTTPException) as exc: + await register_client_with_server( + request=_bridge_mock_request(), + mcp_server=oauth2_server, + client_name="Claude", + grant_types=None, + response_types=None, + token_endpoint_auth_method=None, + ) + + assert exc.value.status_code == 502 + assert str(exc.value.detail) == "upstream registration failed with HTTP 500" + assert "Tomcat" not in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): + """An upstream whose failure body cannot be read (unconsumed stream, lying content-encoding) + makes response.text/.json raise; the classifier must stay total so the caller still gets the + §5.2-shaped 502 instead of the opaque 500 this change set out to remove.""" + import httpx + + unreadable = httpx.Response( + 400, + stream=httpx.ByteStream(b"\x1f\x8bnot-actually-gzip"), + headers={"content-encoding": "gzip"}, + request=httpx.Request("POST", "https://oauth2.googleapis.com/token"), + ) + + response = await _exchange_with_upstream_response(unreadable) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d12ff20ee5b..3433d7dc2d3 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -701,6 +701,38 @@ async def test_default_internal_user_params_with_get_user_object(monkeypatch): assert creation_args["user_role"] == "internal_user" +@pytest.mark.asyncio +async def test_get_user_object_wraps_db_outage_as_valueerror_preserving_context(): + """Pin get_user_object's exception contract: it catches every DB failure in a broad except and + re-raises a bare ValueError, so a real outage survives only as __context__ rather than as the + exception type. The MCP dcr_bridge admission and refresh paths depend on this to tell a transient + outage (retry, 503) from a missing user (fail closed), which is why they classify across the cause + chain instead of the top exception's type. If this wrapping ever changes, that classification must + change with it, so this test guards the contract the callers rely on.""" + from unittest.mock import AsyncMock, MagicMock, patch + + mock_prisma_client = MagicMock() + mock_prisma_client.db = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + side_effect=ConnectionError("can't reach database server") + ) + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): + with pytest.raises(ValueError) as exc_info: + await get_user_object( + user_id="outage-contract-probe-user", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + user_id_upsert=False, + proxy_logging_obj=None, + ) + + assert isinstance(exc_info.value.__context__, ConnectionError) + + @pytest.mark.asyncio async def test_get_user_object_upsert_includes_user_email(): """Test that user_email is included when creating a new user via get_user_object upsert""" diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 0634a01326c..23099177812 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -286,6 +286,46 @@ def test_is_database_service_unavailable_error_excludes_non_infra(error): ) +def _wrapped_like_get_user_object(original): + """Reproduce get_user_object's exception contract (litellm/proxy/auth/auth_checks.py): it catches + every DB failure in a broad ``except`` and re-raises a bare ``ValueError``, so the original error + survives only as ``__context__``. Building it by raising inside an ``except`` sets ``__context__`` + exactly as production does.""" + try: + raise original + except BaseException: + try: + raise ValueError("User doesn't exist in db. Got error - x") + except ValueError as wrapped: + return wrapped + + +def test_is_database_service_unavailable_error_in_chain_sees_through_wrapping(): + """The chain-aware classifier must see a real outage that a caller wrapped in a different type. + get_user_object turns a connection error into a bare ValueError whose type check reads as non-infra, + so the single-exception check returns False and only the chain walk recovers the outage. A missing + user (whose wrapped cause is a plain Exception) must stay non-infra on both.""" + outage = _wrapped_like_get_user_object(ConnectionError("can't reach database server")) + missing_user = _wrapped_like_get_user_object(Exception()) + + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(outage) is False + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(outage) is True + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(missing_user) is False + # parity: a raw outage with no wrapper is still an outage, and a plain ValueError is not + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(ConnectionError("boom")) is True + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(ValueError("nope")) is False + + +def test_is_database_service_unavailable_error_in_chain_terminates_on_a_cause_cycle(): + """The walk must terminate on a pathological __cause__ cycle rather than hang. Neither link is an + outage, so the bounded walk returns False instead of looping forever.""" + first = ValueError("first") + second = ValueError("second") + first.__cause__ = second + second.__cause__ = first + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(first) is False + + def test_is_database_service_unavailable_error_asyncpg(monkeypatch): """asyncpg connection/interface errors map to service-unavailable. asyncpg is not a hard dependency, so inject a stand-in module to exercise the diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index dd9cbcf5232..81840745d0e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -608,3 +608,320 @@ class TestValidateFiniteSpend: with pytest.raises(HTTPException) as exc_info: validate_finite_spend(bad) assert exc_info.value.status_code == 400 + + +class TestValidateFiniteSpendErrorDetail: + """The 400 for non-finite spend must carry the exact {"error": } body.""" + + def test_rejection_detail_is_exact(self): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + with pytest.raises(HTTPException) as exc_info: + validate_finite_spend(float("nan")) + + assert exc_info.value.detail == { + "error": "spend must be a finite number. Received: nan" + } + + +class TestRequireCallerUserIdErrorDetail: + """The 403 for a service-account key must carry the exact error body.""" + + def test_rejection_detail_is_exact(self): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.common_utils import ( + require_caller_user_id_for_non_admin, + ) + + service_account_key = UserAPIKeyAuth( + user_id=None, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + with pytest.raises(HTTPException) as exc_info: + require_caller_user_id_for_non_admin(service_account_key) + + assert exc_info.value.detail == { + "error": "Service-account keys cannot query user analytics. Use a user-bound key, or call as a proxy admin." + } + + +class TestCheckPassthroughRoutesCallerPermission: + """Only proxy admins may set allowed_passthrough_routes (top-level or under + metadata); non-admins get a 403 naming the entity.""" + + def _non_admin(self): + return UserAPIKeyAuth( + user_id="u1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + + def test_top_level_routes_rejected_with_default_entity(self): + from fastapi import HTTPException + from pydantic import BaseModel + + from litellm.proxy.management_endpoints.common_utils import ( + _check_passthrough_routes_caller_permission, + ) + + class _RouteData(BaseModel): + allowed_passthrough_routes: list | None = None + metadata: dict | None = None + + data = _RouteData(allowed_passthrough_routes=["/v1/foo"]) + with pytest.raises(HTTPException) as exc_info: + _check_passthrough_routes_caller_permission(data, self._non_admin()) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == { + "error": "Only proxy admins can set `allowed_passthrough_routes` on a key." + } + + def test_metadata_routes_rejected_with_default_entity(self): + from fastapi import HTTPException + from pydantic import BaseModel + + from litellm.proxy.management_endpoints.common_utils import ( + _check_passthrough_routes_caller_permission, + ) + + class _RouteData(BaseModel): + allowed_passthrough_routes: list | None = None + metadata: dict | None = None + + data = _RouteData(metadata={"allowed_passthrough_routes": ["/v1/foo"]}) + with pytest.raises(HTTPException) as exc_info: + _check_passthrough_routes_caller_permission(data, self._non_admin()) + + assert exc_info.value.detail == { + "error": "Only proxy admins can set `metadata.allowed_passthrough_routes` on a key." + } + + def test_tolerates_data_missing_passthrough_and_metadata_fields(self): + from pydantic import BaseModel + + from litellm.proxy.management_endpoints.common_utils import ( + _check_passthrough_routes_caller_permission, + ) + + class _Bare(BaseModel): + unrelated: str = "x" + + assert ( + _check_passthrough_routes_caller_permission(_Bare(), self._non_admin()) + is None + ) + + +class TestIsUserOrgAdminForTeam: + """The caller must be looked up with its exact identity; a nulled or omitted + lookup argument would silently mis-resolve org-admin status.""" + + @pytest.mark.asyncio + async def test_get_user_object_called_with_caller_identity(self): + from litellm.proxy.management_endpoints.common_utils import ( + _is_user_org_admin_for_team, + ) + + team = LiteLLM_TeamTable( + team_id="t1", organization_id="org1", members_with_roles=[] + ) + key = UserAPIKeyAuth( + user_id="u1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + fake_prisma, fake_cache, fake_logging = MagicMock(), MagicMock(), MagicMock() + mock_get_user = AsyncMock(return_value=None) + + with patch( + "litellm.proxy.proxy_server.prisma_client", fake_prisma + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", fake_cache + ), patch( + "litellm.proxy.proxy_server.proxy_logging_obj", fake_logging + ), patch( + "litellm.proxy.auth.auth_checks.get_user_object", mock_get_user + ): + result = await _is_user_org_admin_for_team(key, team) + + assert result is False + mock_get_user.assert_awaited_once_with( + user_id="u1", + prisma_client=fake_prisma, + user_api_key_cache=fake_cache, + user_id_upsert=False, + proxy_logging_obj=fake_logging, + ) + + +class TestTeamMemberHasPermission: + def test_requires_caller_to_be_a_team_member(self): + from litellm.proxy.management_endpoints.common_utils import ( + _team_member_has_permission, + ) + + team = LiteLLM_TeamTable( + team_id="t1", + team_member_permissions=["/key/generate"], + members_with_roles=[Member(user_id="someone-else", role="user")], + ) + key = UserAPIKeyAuth( + user_id="u1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + assert _team_member_has_permission(key, team, "/key/generate") is False + + +class TestUserHasAdminPrivilegesGuard: + @pytest.mark.asyncio + async def test_no_user_lookup_when_prisma_is_none(self): + """With no DB the guard short-circuits before any user lookup.""" + auth = UserAPIKeyAuth( + user_id="user1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + mock_get_user = AsyncMock(return_value=None) + with patch("litellm.proxy.auth.auth_checks.get_user_object", mock_get_user): + result = await _user_has_admin_privileges( + user_api_key_dict=auth, prisma_client=None + ) + assert result is False + mock_get_user.assert_not_called() + + @pytest.mark.asyncio + async def test_org_admin_membership_grants_privileges(self): + """With DB + user_id present, an ORG_ADMIN membership yields True.""" + auth = UserAPIKeyAuth( + user_id="user1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + now = datetime.now(timezone.utc) + user_obj = LiteLLM_UserTable( + user_id="user1", + organization_memberships=[ + LiteLLM_OrganizationMembershipTable( + user_id="user1", + organization_id="org1", + user_role=LitellmUserRoles.ORG_ADMIN.value, + created_at=now, + updated_at=now, + ) + ], + ) + mock_get_user = AsyncMock(return_value=user_obj) + with patch("litellm.proxy.auth.auth_checks.get_user_object", mock_get_user): + result = await _user_has_admin_privileges( + user_api_key_dict=auth, prisma_client=MagicMock() + ) + assert result is True + + +class TestAdminCanInviteUserGuard: + @pytest.mark.asyncio + async def test_no_user_lookup_when_prisma_is_none(self): + auth = UserAPIKeyAuth( + user_id="admin1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + mock_get_user = AsyncMock(return_value=None) + with patch("litellm.proxy.auth.auth_checks.get_user_object", mock_get_user): + result = await admin_can_invite_user( + target_user_id="target1", + user_api_key_dict=auth, + prisma_client=None, + ) + assert result is False + mock_get_user.assert_not_called() + + @pytest.mark.asyncio + async def test_org_admin_can_invite_user_in_shared_org(self): + now = datetime.now(timezone.utc) + auth = UserAPIKeyAuth( + user_id="admin1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + + def membership(role): + return LiteLLM_OrganizationMembershipTable( + user_id="x", + organization_id="org1", + user_role=role, + created_at=now, + updated_at=now, + ) + + admin_obj = LiteLLM_UserTable( + user_id="admin1", + organization_memberships=[membership(LitellmUserRoles.ORG_ADMIN.value)], + ) + target_obj = LiteLLM_UserTable( + user_id="target1", + organization_memberships=[membership(LitellmUserRoles.INTERNAL_USER.value)], + ) + mock_get_user = AsyncMock(side_effect=[admin_obj, target_obj]) + with patch("litellm.proxy.auth.auth_checks.get_user_object", mock_get_user): + result = await admin_can_invite_user( + target_user_id="target1", + user_api_key_dict=auth, + prisma_client=MagicMock(), + ) + assert result is True + + +class TestTeamAdminCanInviteUserQuery: + @pytest.mark.asyncio + async def test_find_many_queries_admin_teams_with_exact_where(self): + mock_prisma = MagicMock() + mock_auth = MagicMock() + mock_auth.user_id = "admin" + admin_user = LiteLLM_UserTable(user_id="admin", teams=["t1", "t2"]) + target_user = LiteLLM_UserTable(user_id="target", teams=["t2"]) + + def make_team(tid): + obj = MagicMock() + obj.team_id = tid + obj.model_dump = lambda: { + "team_id": tid, + "members_with_roles": [{"user_id": "admin", "role": "admin"}], + } + return obj + + find_many = AsyncMock(return_value=[make_team("t1"), make_team("t2")]) + mock_prisma.db.litellm_teamtable.find_many = find_many + + await _team_admin_can_invite_user( + user_api_key_dict=mock_auth, + admin_user_obj=admin_user, + target_user_obj=target_user, + prisma_client=mock_prisma, + ) + + find_many.assert_awaited_once_with(where={"team_id": {"in": ["t1", "t2"]}}) + + +class TestSetObjectMetadataFieldPremiumArg: + def test_premium_check_receives_the_field_name(self): + team = LiteLLM_TeamTable(team_id="t1", metadata={}) + with patch( + "litellm.proxy.management_endpoints.common_utils._premium_user_check" + ) as mock_premium: + _set_object_metadata_field(team, "guardrails", ["g1"]) + mock_premium.assert_called_once_with("guardrails") + + +class TestUpdateMetadataFieldMove: + def test_none_valued_field_is_not_moved_into_metadata(self): + """A None value must leave the field untouched (guard requires non-None).""" + from litellm.proxy.management_endpoints.common_utils import ( + _update_metadata_field, + ) + + updated_kv = {"guardrails": None} + _update_metadata_field(updated_kv=updated_kv, field_name="guardrails") + assert updated_kv == {"guardrails": None} + + def test_set_premium_field_is_moved_into_metadata(self): + updated_kv = {"guardrails": ["g1"]} + with patch( + "litellm.proxy.management_endpoints.common_utils._premium_user_check" + ): + _update_metadata_fields(updated_kv) + assert "guardrails" not in updated_kv + assert updated_kv["metadata"]["guardrails"] == ["g1"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index a57e18df6ef..bc463d5e75d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -405,3 +405,98 @@ class TestResolveModelForCostLookup: assert resolved_model == "azure/openai/gpt-5.3-codex" assert provider is None + + def test_returns_custom_llm_provider_on_base_model_path(self): + """base_model path: the custom_llm_provider from litellm_params is + returned as the second tuple element, unchanged.""" + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "my-azure-model", + "litellm_params": { + "model": "azure/my-deployment", + "base_model": "azure/gpt-4o", + "custom_llm_provider": "azure", + }, + "model_info": {"id": "test-id"}, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + resolved_model, provider = _resolve_model_for_cost_lookup("my-azure-model") + + assert resolved_model == "azure/gpt-4o" + assert provider == "azure" + + def test_returns_custom_llm_provider_on_resolved_model_path(self): + """resolved-model path (no base_model): the custom_llm_provider from + litellm_params is returned alongside litellm_params.model.""" + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "openai/gpt-4", + "custom_llm_provider": "openai", + }, + "model_info": {"id": "test-id"}, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4") + + assert resolved_model == "openai/gpt-4" + assert provider == "openai" + + def test_resolves_base_model_when_deployment_has_no_litellm_params(self): + """A deployment can omit litellm_params entirely; base_model from + model_info must still resolve (the .get default must be {} not None, + else the later litellm_params.get(...) raises and resolution is lost).""" + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "my-azure-model", + "model_info": {"base_model": "azure/gpt-4o"}, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + resolved_model, provider = _resolve_model_for_cost_lookup("my-azure-model") + + assert resolved_model == "azure/gpt-4o" + assert provider is None + + def test_resolves_model_when_deployment_has_no_model_info(self): + """A deployment can omit model_info entirely; litellm_params.model must + still resolve (the .get default must be {} not None, else the earlier + model_info.get(...) raises and resolution is lost).""" + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4") + + assert resolved_model == "openai/gpt-4" + assert provider is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index db6d3489830..c9abdf09a5d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -12272,6 +12272,55 @@ async def test_bulk_update_team_keys_team_member_no_permission(monkeypatch): mock.update_data.assert_not_called() +def test_handle_key_type_persists_key_type_and_derives_routes(): + """`handle_key_type` keeps `key_type` in the payload (so it is persisted on + the token) while still deriving the `allowed_routes` preset. Regression for + the UI showing scoped keys as "All Proxy Models": the frontend now reads the + persisted `key_type` instead of reverse-mapping the preset string.""" + from litellm.proxy._types import GenerateKeyRequest, LiteLLMKeyType + from litellm.proxy.management_endpoints.key_management_endpoints import ( + handle_key_type, + ) + + cases = { + LiteLLMKeyType.MANAGEMENT: ("management", ["management_routes"]), + LiteLLMKeyType.READ_ONLY: ("read_only", ["info_routes"]), + LiteLLMKeyType.LLM_API: ("llm_api", ["llm_api_routes"]), + } + for key_type, (expected_type, expected_routes) in cases.items(): + data = GenerateKeyRequest(key_type=key_type) + out = handle_key_type(data, {"key_type": key_type}) + assert out["key_type"] == expected_type + assert out["allowed_routes"] == expected_routes + + +def test_handle_key_type_default_persists_type_without_forcing_routes(): + """`default` is persisted but must not overwrite an explicit `allowed_routes` + (e.g. a SCIM key created with `["/scim/*"]` and no explicit key_type).""" + from litellm.proxy._types import GenerateKeyRequest, LiteLLMKeyType + from litellm.proxy.management_endpoints.key_management_endpoints import ( + handle_key_type, + ) + + data = GenerateKeyRequest(key_type=LiteLLMKeyType.DEFAULT) + out = handle_key_type(data, {"allowed_routes": ["/scim/*"], "key_type": LiteLLMKeyType.DEFAULT}) + assert out["key_type"] == "default" + assert out["allowed_routes"] == ["/scim/*"] + + +def test_handle_key_type_none_drops_key_type(): + """When no `key_type` is supplied the payload must not carry a `key_type` + entry, so old keys stay `null` and the frontend keeps its route fallback.""" + from litellm.proxy._types import GenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + handle_key_type, + ) + + data = GenerateKeyRequest(key_type=None) + out = handle_key_type(data, {"key_type": None}) + assert "key_type" not in out + + # ---- pydantic-layer validation ------------------------------------------- diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index da02b774e41..f7c9f343f80 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -14,13 +14,12 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from pydantic import ValidationError -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path import litellm from litellm import Router from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache from litellm.router_strategy.complexity_router.complexity_router import ( ComplexityRouter, DimensionScore, @@ -130,9 +129,7 @@ class TestTokenScoring: tier, score, signals = complexity_router.classify("What is Python?") # Should be classified as SIMPLE due to short length and simple indicator assert tier == ComplexityTier.SIMPLE - assert any("short" in s.lower() for s in signals) or any( - "simple" in s.lower() for s in signals - ) + assert any("short" in s.lower() for s in signals) or any("simple" in s.lower() for s in signals) def test_long_prompt_positive_score(self, complexity_router): """Long prompts should get positive scores (complex indicator).""" @@ -143,9 +140,7 @@ class TestTokenScoring: tier, score, signals = complexity_router.classify(long_prompt) # Should have positive score and detect long token count or technical terms assert score > 0, f"Expected positive score for long prompt, got {score}" - assert any("long" in s.lower() for s in signals) or any( - "technical" in s.lower() for s in signals - ) + assert any("long" in s.lower() for s in signals) or any("technical" in s.lower() for s in signals) class TestCodePresenceScoring: @@ -220,9 +215,7 @@ class TestMultiStepPatterns: def test_first_then_pattern(self, complexity_router): """'First...then' patterns should increase complexity.""" - prompt = ( - "First analyze the data, then create a visualization, then write a report" - ) + prompt = "First analyze the data, then create a visualization, then write a report" tier, score, signals = complexity_router.classify(prompt) assert any("multi-step" in s.lower() for s in signals) @@ -266,9 +259,7 @@ class TestTierAssignment: ) tier, score, signals = complexity_router.classify(prompt) # Should detect technical terms - assert any( - "technical" in s.lower() for s in signals - ), f"Expected technical signals, got {signals}" + assert any("technical" in s.lower() for s in signals), f"Expected technical signals, got {signals}" # Score should be positive due to technical content assert score > 0, f"Expected positive score, got {score}" @@ -468,13 +459,9 @@ class TestConfigOverrides: complexity_router_config=config, ) # With very low thresholds, even neutral prompts should be COMPLEX or higher - tier, score, signals = router.classify( - "Explain how HTTP works with REST APIs and distributed systems" - ) + tier, score, signals = router.classify("Explain how HTTP works with REST APIs and distributed systems") # With boundaries this low, should be at least MEDIUM (anything above -0.5) - assert ( - tier != ComplexityTier.SIMPLE - ), f"Expected non-SIMPLE tier, got {tier} with score {score}" + assert tier != ComplexityTier.SIMPLE, f"Expected non-SIMPLE tier, got {tier} with score {score}" def test_custom_token_thresholds(self, mock_router_instance): """Test custom token thresholds work correctly.""" @@ -499,9 +486,7 @@ class TestConfigOverrides: long_prompt = "This is a test prompt " * 30 # ~120 tokens tier, score, signals = router.classify(long_prompt) # Should get token length signal indicating "long" - assert any( - "long" in s.lower() if s else False for s in signals - ), f"Expected 'long' signal, got {signals}" + assert any("long" in s.lower() if s else False for s in signals), f"Expected 'long' signal, got {signals}" class TestCustomTechnicalKeywords: @@ -516,9 +501,7 @@ class TestCustomTechnicalKeywords: ) assert router.technical_keywords == DEFAULT_TECHNICAL_KEYWORDS + ["udp", "kafka"] - def test_custom_keywords_appended_to_technical_keywords_override( - self, mock_router_instance - ): + def test_custom_keywords_appended_to_technical_keywords_override(self, mock_router_instance): """Custom keywords should be appended to a technical_keywords override.""" router = ComplexityRouter( model_name="test-router", @@ -535,9 +518,7 @@ class TestCustomTechnicalKeywords: router = ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, - complexity_router_config={ - "custom_technical_keywords": ["TCP", "udp", "UDP", "kafka"] - }, + complexity_router_config={"custom_technical_keywords": ["TCP", "udp", "UDP", "kafka"]}, ) lowered = [kw.lower() for kw in router.technical_keywords] assert lowered == [kw.lower() for kw in DEFAULT_TECHNICAL_KEYWORDS] + [ @@ -560,9 +541,7 @@ class TestCustomTechnicalKeywords: assert router_absent.technical_keywords == DEFAULT_TECHNICAL_KEYWORDS assert router_none.technical_keywords == DEFAULT_TECHNICAL_KEYWORDS - def test_prompt_with_only_custom_keywords_scores_technical( - self, mock_router_instance, basic_config - ): + def test_prompt_with_only_custom_keywords_scores_technical(self, mock_router_instance, basic_config): """A prompt matching only custom keywords should score higher on technicalTerms.""" prompt = "Configure udp multicast between kafka brokers" baseline_router = ComplexityRouter( @@ -581,9 +560,7 @@ class TestCustomTechnicalKeywords: _, baseline_score, baseline_signals = baseline_router.classify(prompt) _, custom_score, custom_signals = custom_router.classify(prompt) assert not any("technical" in s.lower() for s in baseline_signals) - assert any( - "technical" in s.lower() for s in custom_signals - ), f"Expected technical signal, got {custom_signals}" + assert any("technical" in s.lower() for s in custom_signals), f"Expected technical signal, got {custom_signals}" assert custom_score > baseline_score @@ -776,9 +753,7 @@ class TestKeywordFalsePositives: prompt = "What is the capital of France?" tier, score, signals = complexity_router.classify(prompt) # Should NOT detect code presence from 'api' in 'capital' - assert not any( - "code" in s.lower() for s in signals - ), "False positive: got code signal from 'capital'" + assert not any("code" in s.lower() for s in signals), "False positive: got code signal from 'capital'" # Should be SIMPLE (definition question) assert tier == ComplexityTier.SIMPLE @@ -787,9 +762,7 @@ class TestKeywordFalsePositives: prompt = "Explain digital marketing strategies" tier, score, signals = complexity_router.classify(prompt) # Should NOT detect code presence from 'git' in 'digital' - assert not any( - "code" in s.lower() for s in signals - ), "False positive: got code signal from 'digital'" + assert not any("code" in s.lower() for s in signals), "False positive: got code signal from 'digital'" def test_try_not_in_entry(self, complexity_router): """'try' should not match in 'entry'.""" @@ -803,43 +776,33 @@ class TestKeywordFalsePositives: """'error' should not match in 'terrorism'.""" prompt = "The country is dealing with terrorism" tier, score, signals = complexity_router.classify(prompt) - assert not any( - "code" in s.lower() for s in signals - ), "False positive: got code signal from 'terrorism'" + assert not any("code" in s.lower() for s in signals), "False positive: got code signal from 'terrorism'" def test_class_not_in_classical(self, complexity_router): """'class' should not match in 'classical'.""" prompt = "I enjoy listening to classical music" tier, score, signals = complexity_router.classify(prompt) - assert not any( - "code" in s.lower() for s in signals - ), "False positive: got code signal from 'classical'" + assert not any("code" in s.lower() for s in signals), "False positive: got code signal from 'classical'" def test_merge_not_in_emerged(self, complexity_router): """'merge' should not match in 'emerged'.""" prompt = "A new leader emerged from the crowd" tier, score, signals = complexity_router.classify(prompt) - assert not any( - "code" in s.lower() for s in signals - ), "False positive: got code signal from 'emerged'" + assert not any("code" in s.lower() for s in signals), "False positive: got code signal from 'emerged'" def test_actual_api_keyword_detected(self, complexity_router): """Actual 'api' usage should be detected.""" prompt = "How do I call the REST api endpoint?" tier, score, signals = complexity_router.classify(prompt) # Should detect code presence from actual 'api' usage - assert any( - "code" in s.lower() for s in signals - ), f"Expected code signal for 'api', got {signals}" + assert any("code" in s.lower() for s in signals), f"Expected code signal for 'api', got {signals}" def test_actual_git_keyword_detected(self, complexity_router): """Actual 'git' usage should be detected.""" prompt = "How do I use git to commit changes?" tier, score, signals = complexity_router.classify(prompt) # Should detect code presence from actual 'git' usage - assert any( - "code" in s.lower() for s in signals - ), f"Expected code signal for 'git', got {signals}" + assert any("code" in s.lower() for s in signals), f"Expected code signal for 'git', got {signals}" class TestEdgeCases: @@ -859,9 +822,7 @@ class TestEdgeCases: # Should have positive score due to length assert score > 0, f"Expected positive score for very long prompt, got {score}" # Should detect long token count - assert any( - "long" in s.lower() for s in signals - ), f"Expected 'long' signal, got {signals}" + assert any("long" in s.lower() for s in signals), f"Expected 'long' signal, got {signals}" def test_unicode_prompt(self, complexity_router): """Test handling of unicode characters.""" @@ -879,9 +840,7 @@ class TestEdgeCases: """ tier, score, signals = complexity_router.classify(prompt) # The "step N" pattern should be detected - assert any( - "multi-step" in s.lower() for s in signals - ), f"Expected multi-step signal, got {signals}" + assert any("multi-step" in s.lower() for s in signals), f"Expected multi-step signal, got {signals}" class TestRouterComplexityDeploymentMethods: @@ -1019,9 +978,7 @@ class TestAsyncPreRoutingHookMultiFormat: assert result.messages is not None @pytest.mark.asyncio - async def test_should_route_with_responses_api_string_input( - self, complexity_router - ): + async def test_should_route_with_responses_api_string_input(self, complexity_router): """Test routing with Responses API string input via handler dispatch.""" from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, @@ -1109,9 +1066,7 @@ class TestAsyncPreRoutingHookMultiFormat: assert result.model is not None @pytest.mark.asyncio - async def test_should_return_none_when_no_messages_or_input( - self, complexity_router - ): + async def test_should_return_none_when_no_messages_or_input(self, complexity_router): """Test that None is returned when neither messages nor input is available.""" result = await complexity_router.async_pre_routing_hook( model="test-model", @@ -1122,9 +1077,7 @@ class TestAsyncPreRoutingHookMultiFormat: assert result is None @pytest.mark.asyncio - async def test_should_prefer_original_messages_over_conversion( - self, complexity_router - ): + async def test_should_prefer_original_messages_over_conversion(self, complexity_router): """Test that original messages are used when both messages and input are available.""" messages = [{"role": "user", "content": "What is 2+2?"}] result = await complexity_router.async_pre_routing_hook( @@ -1136,9 +1089,7 @@ class TestAsyncPreRoutingHookMultiFormat: assert result.messages == messages @pytest.mark.asyncio - async def test_should_include_instructions_in_classification( - self, complexity_router - ): + async def test_should_include_instructions_in_classification(self, complexity_router): """Test that Responses API instructions influence classification via system message.""" from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, @@ -1175,9 +1126,7 @@ class TestExtractUserMessageAndSystemPrompt: {"role": "assistant", "content": "Hi!"}, {"role": "user", "content": "How are you?"}, ] - user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( - messages - ) + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt(messages) assert user_msg == "How are you?" assert sys_prompt == "You are helpful." @@ -1187,9 +1136,7 @@ class TestExtractUserMessageAndSystemPrompt: {"role": "system", "content": "You are helpful."}, {"role": "assistant", "content": "Hi!"}, ] - user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( - messages - ) + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt(messages) assert user_msg is None assert sys_prompt == "You are helpful." @@ -1207,17 +1154,13 @@ class TestExtractUserMessageAndSystemPrompt: ], } ] - user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( - messages - ) + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt(messages) assert user_msg == "Describe this image" assert sys_prompt is None def test_should_handle_empty_messages(self): """Test with empty messages list.""" - user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( - [] - ) + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt([]) assert user_msg is None assert sys_prompt is None @@ -1282,17 +1225,13 @@ class TestLLMClassifier: assert tier == ComplexityTier.SIMPLE @pytest.mark.asyncio - async def test_aclassify_llm_success_routes_by_llm_verdict( - self, llm_complexity_router, mock_router_instance - ): + async def test_aclassify_llm_success_routes_by_llm_verdict(self, llm_complexity_router, mock_router_instance): """A well-formed structured LLM response should decide the tier directly. Uses a prompt that heuristic scoring alone would classify as SIMPLE, to prove the LLM verdict -- not the heuristic scorer -- is what decided the tier. """ - mock_router_instance.acompletion = AsyncMock( - return_value=_llm_response('{"tier": "COMPLEX"}') - ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) tier, score, signals = await llm_complexity_router.aclassify("hi") assert tier == ComplexityTier.COMPLEX assert "llm-classifier:COMPLEX" in signals @@ -1311,13 +1250,9 @@ class TestLLMClassifier: sees no user_api_key/team_id/user_id and silently drops all spend logging and budget accounting for the classifier call. """ - mock_router_instance.acompletion = AsyncMock( - return_value=_llm_response('{"tier": "SIMPLE"}') - ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) request_metadata = {"user_api_key": "sk-abc", "user_api_key_team_id": "team-1"} - await llm_complexity_router.aclassify( - "hi", request_kwargs={"litellm_metadata": request_metadata} - ) + await llm_complexity_router.aclassify("hi", request_kwargs={"litellm_metadata": request_metadata}) call_kwargs = mock_router_instance.acompletion.call_args.kwargs assert call_kwargs["metadata"] == request_metadata @@ -1333,18 +1268,14 @@ class TestLLMClassifier: business touching, so it must be stripped while the rest of the attribution metadata (key/team) is preserved. """ - mock_router_instance.acompletion = AsyncMock( - return_value=_llm_response('{"tier": "SIMPLE"}') - ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) request_metadata = { "user_api_key": "sk-abc", "user_api_key_team_id": "team-1", "user_api_key_budget_reservation": {"reserved_cost": 1.0}, "user_api_key_auth": {"models": ["gpt-4o"], "budget_reservation": {"reserved_cost": 1.0}}, } - await llm_complexity_router.aclassify( - "hi", request_kwargs={"litellm_metadata": request_metadata} - ) + await llm_complexity_router.aclassify("hi", request_kwargs={"litellm_metadata": request_metadata}) call_kwargs = mock_router_instance.acompletion.call_args.kwargs # user_api_key_budget_reservation is stripped (budget enforcement) while # user_api_key_auth is kept so _filter_deployments_by_model_access_groups @@ -1391,13 +1322,9 @@ class TestLLMClassifier: assert tier == ComplexityTier.SIMPLE @pytest.mark.asyncio - async def test_pre_routing_hook_uses_llm_classifier_end_to_end( - self, llm_complexity_router, mock_router_instance - ): + async def test_pre_routing_hook_uses_llm_classifier_end_to_end(self, llm_complexity_router, mock_router_instance): """The full pre-routing hook should route using the LLM classifier's verdict.""" - mock_router_instance.acompletion = AsyncMock( - return_value=_llm_response('{"tier": "REASONING"}') - ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}')) request_metadata = {"user_api_key": "sk-abc", "user_api_key_team_id": "team-1"} result = await llm_complexity_router.async_pre_routing_hook( model="test-model", @@ -1410,6 +1337,186 @@ class TestLLMClassifier: assert call_kwargs["metadata"] == request_metadata +class TestRouterPreRoutingAliasOverrides: + """ + Regression tests for: litellm_params configured on a complexity-router alias + entry (e.g. `cache_control_injection_points`, `drop_params`) were silently + dropped, because `async_pre_routing_hook` swaps `model` from the alias name + to the selected tier's model *before* the deployment lookup - so the actual + outbound call only ever merges in the tier deployment's own litellm_params, + never the alias's. + """ + + def _make_router(self) -> Router: + return Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "drop_params": True, + "cache_control_injection_points": [{"location": "message", "role": "system"}], + "complexity_router_config": { + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + } + }, + "complexity_router_default_model": "gpt-4o", + }, + }, + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + }, + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o"}, + }, + ] + ) + + @pytest.mark.asyncio + async def test_alias_litellm_params_applied_to_request_kwargs(self): + """cache_control_injection_points/drop_params set on the alias entry + reach the outbound request even though the tier deployment is what + actually gets called.""" + router = self._make_router() + request_kwargs: Dict = {} + + result = await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert result is not None + assert request_kwargs["drop_params"] is True + assert request_kwargs["cache_control_injection_points"] == [{"location": "message", "role": "system"}] + + @pytest.mark.asyncio + async def test_alias_overrides_exclude_only_model(self): + """`model` (the alias marker, e.g. auto_router/complexity_router) is + excluded since it's never a real provider model. Router-only fields + like complexity_router_config DO flow through into request_kwargs at + this layer - they're filtered from the actual outbound LLM call + downstream by litellm.types.utils.all_litellm_params instead, not by + the router's pre-routing hook. See test_router_init_only_params_are_ + never_sent_to_a_provider for the guard on that downstream filter.""" + router = self._make_router() + request_kwargs: Dict = {} + + await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert "model" not in request_kwargs + assert request_kwargs["complexity_router_config"] == { + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + } + } + assert request_kwargs["complexity_router_default_model"] == "gpt-4o" + + def test_router_init_only_params_are_never_sent_to_a_provider(self): + """The router's pre-routing hook only excludes `model` (see + test_alias_overrides_exclude_only_model above) - every other alias + litellm_param, including router-init-only fields like + complexity_router_config, flows into request_kwargs unfiltered. That's + only safe because litellm.completion()/acompletion() itself strips + anything listed in all_litellm_params before building the provider + request. If one of these keys is ever removed from that list, it + ships raw to the real provider as extra_body - verified live via + litellm.completion(..., complexity_router_config={...}) landing in + extra_body before this list included it.""" + from litellm.types.utils import all_litellm_params + + router_init_only_params = ( + "auto_router_config_path", + "auto_router_config", + "auto_router_default_model", + "auto_router_embedding_model", + "complexity_router_config", + "complexity_router_default_model", + "adaptive_router_config", + "adaptive_router_default_model", + "quality_router_config", + "quality_router_default_model", + ) + for param in router_init_only_params: + assert param in all_litellm_params, ( + f"{param} must stay in litellm.types.utils.all_litellm_params - " + "removing it means it ships raw to the real provider as extra_body" + ) + + @pytest.mark.asyncio + async def test_caller_supplied_kwargs_are_not_overwritten(self): + """A value the caller already passed for this request takes + precedence over the alias's configured default.""" + router = self._make_router() + request_kwargs: Dict = {"drop_params": False} + + await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert request_kwargs["drop_params"] is False + + @pytest.mark.asyncio + async def test_non_alias_model_is_untouched(self): + """A plain (non-router-alias) model name is not affected by the + alias-override merge at all.""" + router = self._make_router() + request_kwargs: Dict = {} + + result = await router.async_pre_routing_hook( + model="gpt-4o-mini", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert result is None + assert request_kwargs == {} + + @pytest.mark.asyncio + async def test_adaptive_router_alias_overrides_survive_reload(self): + """Alias litellm_params are read fresh from self.model_list at request + time (not cached at init), so a set_model_list() reload (e.g. + /config/reload) - which rebuilds self.model_list but leaves an + already-built AdaptiveRouter alone - can't leave them stale.""" + model_list = [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "drop_params": True, + "adaptive_router_config": {"available_models": ["gpt-4o-mini"]}, + }, + }, + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + }, + ] + router = Router(model_list=model_list) + router.set_model_list(model_list) + assert "smart-router" in router.adaptive_routers + + request_kwargs: Dict = {} + await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert request_kwargs["drop_params"] is True + + class TestAdaptiveSoftFloors: def test_adaptive_defaults_use_cost_weighted_cold_policy(self): config = ComplexityRouterConfig( @@ -1430,9 +1537,7 @@ class TestAdaptiveSoftFloors: "model": "openai/gpt-4o-mini", "input_cost_per_token": 0.00000015, }, - "model_info": { - "adaptive_router_preferences": {"quality_tier": 1, "strengths": []} - }, + "model_info": {"adaptive_router_preferences": {"quality_tier": 1, "strengths": []}}, }, { "model_name": "premium", @@ -1440,9 +1545,7 @@ class TestAdaptiveSoftFloors: "model": "openai/gpt-4o", "input_cost_per_token": 0.000005, }, - "model_info": { - "adaptive_router_preferences": {"quality_tier": 3, "strengths": []} - }, + "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, }, ] router.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]} @@ -1467,9 +1570,7 @@ class TestAdaptiveSoftFloors: with pytest.raises(ValidationError): ComplexityRouterConfig(adaptive=True, tiers={"SIMPLE": []}) - def test_cold_start_randomly_samples_unobserved_classified_tier_models( - self, adaptive_router_instance - ): + def test_cold_start_randomly_samples_unobserved_classified_tier_models(self, adaptive_router_instance): cr = ComplexityRouter( model_name="hybrid", litellm_router_instance=adaptive_router_instance, @@ -1498,9 +1599,7 @@ class TestAdaptiveSoftFloors: "premium", } - def test_get_model_for_tier_list_without_adaptive_random_choice( - self, mock_router_instance - ): + def test_get_model_for_tier_list_without_adaptive_random_choice(self, mock_router_instance): router = ComplexityRouter( model_name="test", litellm_router_instance=mock_router_instance, @@ -1519,9 +1618,7 @@ class TestAdaptiveSoftFloors: choice.assert_called_once_with(pool) assert router.get_model_for_tier(ComplexityTier.MEDIUM) == "mid" - def test_soft_floor_prefers_home_tier_when_posteriors_equal( - self, adaptive_router_instance, hybrid_config - ): + def test_soft_floor_prefers_home_tier_when_posteriors_equal(self, adaptive_router_instance, hybrid_config): from litellm.router_strategy.adaptive_router.bandit import BanditCell from litellm.types.router import RequestType @@ -1533,9 +1630,7 @@ class TestAdaptiveSoftFloors: adaptive = cr._ensure_adaptive_router() assert adaptive is not None for model in ("cheap", "premium"): - adaptive._cells[(RequestType.GENERAL, model)] = BanditCell( - alpha=5.0, beta=5.0 - ) + adaptive._cells[(RequestType.GENERAL, model)] = BanditCell(alpha=5.0, beta=5.0) # Equal quality samples; home-tier penalty should favor cheap for SIMPLE. with patch( @@ -1545,9 +1640,7 @@ class TestAdaptiveSoftFloors: picked = cr._soft_floor_pick(ComplexityTier.SIMPLE, "hi") assert picked == "cheap" - def test_soft_floor_allows_cross_tier_when_posterior_dominates( - self, adaptive_router_instance, hybrid_config - ): + def test_soft_floor_allows_cross_tier_when_posterior_dominates(self, adaptive_router_instance, hybrid_config): from litellm.router_strategy.adaptive_router.bandit import BanditCell from litellm.types.router import RequestType @@ -1558,12 +1651,8 @@ class TestAdaptiveSoftFloors: ) adaptive = cr._ensure_adaptive_router() assert adaptive is not None - adaptive._cells[(RequestType.GENERAL, "cheap")] = BanditCell( - alpha=1.0, beta=20.0 - ) - adaptive._cells[(RequestType.GENERAL, "premium")] = BanditCell( - alpha=20.0, beta=1.0 - ) + adaptive._cells[(RequestType.GENERAL, "cheap")] = BanditCell(alpha=1.0, beta=20.0) + adaptive._cells[(RequestType.GENERAL, "premium")] = BanditCell(alpha=20.0, beta=1.0) with patch( "litellm.router_strategy.adaptive_router.bandit.thompson_sample", @@ -1572,9 +1661,7 @@ class TestAdaptiveSoftFloors: picked = cr._soft_floor_pick(ComplexityTier.SIMPLE, "hi") assert picked == "premium" - def test_reused_model_has_zero_distance_in_each_configured_tier( - self, adaptive_router_instance - ): + def test_reused_model_has_zero_distance_in_each_configured_tier(self, adaptive_router_instance): from litellm.router_strategy.adaptive_router.bandit import BanditCell from litellm.types.router import RequestType @@ -1593,9 +1680,7 @@ class TestAdaptiveSoftFloors: adaptive = cr._ensure_adaptive_router() assert adaptive is not None for model in ("cheap", "premium"): - adaptive._cells[(RequestType.GENERAL, model)] = BanditCell( - alpha=6.0, beta=5.0 - ) + adaptive._cells[(RequestType.GENERAL, model)] = BanditCell(alpha=6.0, beta=5.0) request_kwargs: Dict = {"metadata": {}} with patch( @@ -1604,20 +1689,14 @@ class TestAdaptiveSoftFloors: ): cr._soft_floor_pick(ComplexityTier.MEDIUM, "hi", request_kwargs) - candidates = request_kwargs["metadata"]["adaptive_router_decision"][ - "candidates" - ] - assert { - candidate["model"]: candidate["tier_distance"] for candidate in candidates - } == { + candidates = request_kwargs["metadata"]["adaptive_router_decision"]["candidates"] + assert {candidate["model"]: candidate["tier_distance"] for candidate in candidates} == { "cheap": 0, "premium": 0, } @pytest.mark.asyncio - async def test_pre_routing_hook_adaptive_stashes_chosen_model( - self, adaptive_router_instance, hybrid_config - ): + async def test_pre_routing_hook_adaptive_stashes_chosen_model(self, adaptive_router_instance, hybrid_config): cr = ComplexityRouter( model_name="hybrid", litellm_router_instance=adaptive_router_instance, @@ -1631,10 +1710,7 @@ class TestAdaptiveSoftFloors: ) assert result is not None assert result.model in {"cheap", "premium"} - assert ( - request_kwargs["metadata"].get("adaptive_router_chosen_model") - == result.model - ) + assert request_kwargs["metadata"].get("adaptive_router_chosen_model") == result.model decision = request_kwargs["metadata"]["adaptive_router_decision"] assert decision["phase"] == "cold_start" assert decision["classified_tier"] == "SIMPLE" @@ -1657,9 +1733,7 @@ class TestLexicalKeywordTierRules: } @pytest.mark.asyncio - async def test_matching_rule_overrides_scoring( - self, mock_router_instance, rule_config - ): + async def test_matching_rule_overrides_scoring(self, mock_router_instance, rule_config): """A prompt hitting a rule keyword routes to that tier, not the scored tier.""" router = ComplexityRouter( model_name="test-router", @@ -1745,9 +1819,7 @@ class TestLexicalKeywordTierRules: assert router._lexical_tier_override("nothing relevant here") is None @pytest.mark.asyncio - async def test_no_rule_match_falls_back_to_scoring( - self, mock_router_instance, basic_config - ): + async def test_no_rule_match_falls_back_to_scoring(self, mock_router_instance, basic_config): """A prompt that matches no rule is classified by the scorer as usual.""" config = { **basic_config, @@ -1768,9 +1840,7 @@ class TestLexicalKeywordTierRules: assert result is not None assert result.model == "gpt-4o-mini" # SIMPLE via scoring, rule did not fire - def test_word_boundary_avoids_substring_false_positive( - self, mock_router_instance, basic_config - ): + def test_word_boundary_avoids_substring_false_positive(self, mock_router_instance, basic_config): """A single-word rule keyword must not match inside a larger word.""" config = { **basic_config, @@ -1788,10 +1858,7 @@ class TestLexicalKeywordTierRules: def _make_embedding_response(vectors: List[List[float]]) -> "litellm.EmbeddingResponse": return litellm.EmbeddingResponse( model="fake-embed", - data=[ - {"embedding": vec, "index": idx, "object": "embedding"} - for idx, vec in enumerate(vectors) - ], + data=[{"embedding": vec, "index": idx, "object": "embedding"} for idx, vec in enumerate(vectors)], object="list", ) @@ -1818,8 +1885,7 @@ class FakeEmbeddingRouter: def _vectors(self, docs: List[str]) -> List[List[float]]: return [ - [1.0, 0.0] if any(marker in doc.lower() for marker in self._CLUSTER_MARKERS) else [0.0, 1.0] - for doc in docs + [1.0, 0.0] if any(marker in doc.lower() for marker in self._CLUSTER_MARKERS) else [0.0, 1.0] for doc in docs ] @staticmethod @@ -2358,9 +2424,7 @@ class TestRoutingDecisionCauseLogging: verbose_router_logger.removeHandler(caplog.handler) @pytest.mark.asyncio - async def test_literal_keyword_match_logs_its_cause( - self, mock_router_instance, basic_config, router_log_capture - ): + async def test_literal_keyword_match_logs_its_cause(self, mock_router_instance, basic_config, router_log_capture): config = { **basic_config, "keyword_tier_rules": [{"keywords": ["deploy to k8s"], "tier": "REASONING"}], @@ -2406,9 +2470,7 @@ class TestRoutingDecisionCauseLogging: assert "cause=literal_keyword_match" not in router_log_capture.text @pytest.mark.asyncio - async def test_complexity_scorer_logs_its_cause( - self, mock_router_instance, basic_config, router_log_capture - ): + async def test_complexity_scorer_logs_its_cause(self, mock_router_instance, basic_config, router_log_capture): # No keyword rules -> the scorer decides, and its line must be tagged as such. router = ComplexityRouter( model_name="test-router", @@ -2424,3 +2486,217 @@ class TestRoutingDecisionCauseLogging: assert "score=" in router_log_capture.text assert "cause=literal_keyword_match" not in router_log_capture.text assert "cause=semantic_keyword_match" not in router_log_capture.text + + +class TestSessionAffinity: + """Test the opt-in session_affinity sticky-routing behavior.""" + + REASONING_MESSAGE = [ + { + "role": "user", + "content": "Let's think step by step and reason through this problem carefully.", + } + ] + SIMPLE_MESSAGE = [{"role": "user", "content": "Hello!"}] + + @pytest.fixture + def session_affinity_config(self, basic_config) -> Dict: + return {**basic_config, "session_affinity": True} + + @staticmethod + def _request_kwargs(session_id: str) -> Dict: + return {"metadata": {"session_id": session_id}} + + @pytest.mark.asyncio + async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): + """Regression: session_affinity defaults to False, so a shared session_id must + not pin the model -- each turn is still classified independently.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=basic_config, + ) + request_kwargs = self._request_kwargs("session-1") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" + assert second.model == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_pins_model_after_first_turn(self, mock_router_instance, session_affinity_config): + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + request_kwargs = self._request_kwargs("session-1") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + assert first.model == "o1-preview" + + with patch.object(router, "aclassify", wraps=router.aclassify) as spy_aclassify: + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + spy_aclassify.assert_not_called() + # Pinned to the first turn's model, not re-classified down to SIMPLE. + assert second.model == "o1-preview" + + @pytest.mark.asyncio + async def test_different_sessions_classify_independently(self, mock_router_instance, session_affinity_config): + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + reasoning = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("session-a"), messages=self.REASONING_MESSAGE + ) + simple = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("session-b"), messages=self.SIMPLE_MESSAGE + ) + assert reasoning.model == "o1-preview" + assert simple.model == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_respects_ttl_seconds(self, mock_router_instance, basic_config): + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value=None) + mock_router_instance.cache = cache + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "session_affinity": True, + "session_affinity_ttl_seconds": 120, + }, + ) + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("session-1"), messages=self.SIMPLE_MESSAGE + ) + cache.async_set_cache.assert_called_once() + call_kwargs = cache.async_set_cache.call_args.kwargs + assert call_kwargs["ttl"] == 120 + assert call_kwargs["value"] == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_ttl_refreshed_on_cache_hit(self, mock_router_instance, basic_config): + """Regression: a pinned turn must refresh the TTL, not just the first write -- + otherwise a session outliving session_affinity_ttl_seconds silently loses its pin.""" + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value="o1-preview") + mock_router_instance.cache = cache + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "session_affinity": True, + "session_affinity_ttl_seconds": 90, + }, + ) + result = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("session-1"), messages=self.SIMPLE_MESSAGE + ) + assert result.model == "o1-preview" + cache.async_set_cache.assert_called_once() + call_kwargs = cache.async_set_cache.call_args.kwargs + assert call_kwargs["value"] == "o1-preview" + assert call_kwargs["ttl"] == 90 + + @pytest.mark.asyncio + async def test_different_api_keys_do_not_share_pin(self, mock_router_instance, session_affinity_config): + """A session_id is client-supplied and unauthenticated; two different callers + (API keys) reusing the same session_id must not poison each other's pin.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + caller_a_kwargs = {"metadata": {"session_id": "shared-session", "user_api_key_hash": "key-a"}} + caller_b_kwargs = {"metadata": {"session_id": "shared-session", "user_api_key_hash": "key-b"}} + + pinned_for_a = await router.async_pre_routing_hook( + model="test-model", request_kwargs=caller_a_kwargs, messages=self.REASONING_MESSAGE + ) + assert pinned_for_a.model == "o1-preview" + + # Caller B reuses the same session_id but has a different API key; its trivial + # message must classify fresh, not inherit caller A's REASONING-tier pin. + result_for_b = await router.async_pre_routing_hook( + model="test-model", request_kwargs=caller_b_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert result_for_b.model == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_no_session_id_falls_back_to_reclassify(self, mock_router_instance, session_affinity_config): + cache = AsyncMock() + mock_router_instance.cache = cache + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + result = await router.async_pre_routing_hook( + model="test-model", request_kwargs={}, messages=self.SIMPLE_MESSAGE + ) + assert result.model == "gpt-4o-mini" + cache.async_get_cache.assert_not_called() + cache.async_set_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_adaptive_pinned_turn_still_stamps_chosen_model_metadata(self, mock_router_instance): + """Regression: skipping classification on a pinned turn must not break the + adaptive bandit's reward-feedback loop, which only records a turn's outcome + when ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY is present in the request metadata.""" + mock_router_instance.cache = DualCache() + mock_router_instance.model_list = [ + { + "model_name": "cheap", + "litellm_params": {"model": "openai/gpt-4o-mini", "input_cost_per_token": 0.0}, + "model_info": {}, + }, + ] + mock_router_instance.model_name_to_deployment_indices = {"cheap": [0]} + router = ComplexityRouter( + model_name="hybrid", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "adaptive": True, + "session_affinity": True, + "tiers": { + "SIMPLE": ["cheap"], + "MEDIUM": ["cheap"], + "COMPLEX": ["cheap"], + "REASONING": ["cheap"], + }, + "default_model": "cheap", + }, + ) + first = await router.async_pre_routing_hook( + model="hybrid", + request_kwargs=self._request_kwargs("session-1"), + messages=[{"role": "user", "content": "hi"}], + ) + assert first.model == "cheap" + + request_kwargs_2 = self._request_kwargs("session-1") + with patch.object(router, "aclassify", wraps=router.aclassify) as spy_aclassify: + second = await router.async_pre_routing_hook( + model="hybrid", + request_kwargs=request_kwargs_2, + messages=[{"role": "user", "content": "hi again"}], + ) + spy_aclassify.assert_not_called() + assert second.model == "cheap" + assert request_kwargs_2["metadata"]["adaptive_router_chosen_model"] == "cheap" diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 28cf4fa0744..4611aafa3c1 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2081,3 +2081,79 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): assert response.usage.prompt_tokens > 0 assert response.usage.completion_tokens > 0 assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "aws_credential_kwargs", + [ + { + "aws_session_name": "litellm-gcp", + "aws_role_name": "arn:aws:iam::123456789012:role/litellm-bedrock-role", + "aws_web_identity_token": "oidc/google/108963886734710037768", + }, + { + "aws_access_key_id": "AKIASTATICKEYFORTEST", + "aws_secret_access_key": "static-secret-key", + "aws_session_token": "static-session-token", + }, + ], + ids=["web_identity", "static_keys"], +) +async def test_acompletion_forwards_aws_credentials_through_responses_bridge( + respx_mock: respx.MockRouter, monkeypatch, aws_credential_kwargs: dict +): + from botocore.credentials import Credentials + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + original_disable_aiohttp = litellm.disable_aiohttp_transport + try: + litellm.disable_aiohttp_transport = True + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + get_credentials_mock = MagicMock(return_value=Credentials("fake-key", "fake-secret")) + monkeypatch.setattr(BaseAWSLLM, "get_credentials", get_credentials_mock) + + respx_mock.post("https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses").respond( + json={ + "id": "resp_123", + "object": "response", + "created_at": 1760144904, + "status": "completed", + "model": "openai.gpt-5.4", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + } + ) + + response = await litellm.acompletion( + model="bedrock_mantle/openai.gpt-5.4", + messages=[{"role": "user", "content": "hi"}], + api_base="https://bedrock-mantle.us-east-2.api.aws/v1", + aws_region_name="us-east-2", + num_retries=0, + **aws_credential_kwargs, + ) + + assert response.choices[0].message.content == "ok" + credential_kwargs = get_credentials_mock.call_args.kwargs + assert credential_kwargs["aws_region_name"] == "us-east-2" + for key, value in aws_credential_kwargs.items(): + assert credential_kwargs[key] == value + authorization = respx_mock.calls.last.request.headers["Authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256") + assert "fake-key" in authorization + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts index b30bb8aca7b..e7f67d7367f 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts @@ -96,7 +96,9 @@ test.describe("Proxy Admin - Teams", () => { const teamRow = page.locator("tr", { hasText: E2E_TEAM_DELETE_ALIAS }).first(); await expect(teamRow).toBeVisible({ timeout: 10_000 }); - await teamRow.locator("svg, img").last().click(); + // Actions live in a kebab menu: open it, then click "Delete team". + await teamRow.locator('[data-testid^="team-actions-"]').click(); + await page.getByTestId("team-action-delete").click(); const modal = page.locator(".ant-modal:visible"); await expect(modal).toBeVisible({ timeout: 5_000 }); diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index ba2e2a375ca..fa2287bdaef 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -515,6 +515,152 @@ "count": 2 } }, + "src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.test.tsx": { + "unused-imports/no-unused-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx": { + "react-hooks/immutability": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/MCPSubmissionsTab.tsx": { + "react-hooks/set-state-in-effect": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 1 + }, + "unused-imports/no-unused-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/OpenAPIQuickPicker.tsx": { + "react-hooks/set-state-in-effect": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.tsx": { + "no-nested-ternary": { + "count": 3 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx": { + "no-nested-ternary": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/create_mcp_server.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 4 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx": { + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/static-components": { + "count": 4 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_connection_status.tsx": { + "no-nested-ternary": { + "count": 3 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_discovery.tsx": { + "react-hooks/set-state-in-effect": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_config.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_display.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/immutability": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 5 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_tools.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 2 + } + }, "src/app/(dashboard)/memory/_components/MemoryView.tsx": { "react-hooks/set-state-in-effect": { "count": 1 @@ -1497,13 +1643,13 @@ }, "src/components/Teams.tsx": { "no-nested-ternary": { - "count": 4 + "count": 2 }, "no-restricted-imports": { "count": 1 }, "react-hooks/set-state-in-effect": { - "count": 4 + "count": 3 } }, "src/components/ToolDetail.tsx": { @@ -1544,14 +1690,6 @@ "count": 1 } }, - "src/components/VirtualKeysPage/VirtualKeysTable.tsx": { - "no-nested-ternary": { - "count": 2 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/components/activity_metrics.tsx": { "no-nested-ternary": { "count": 1 @@ -1796,11 +1934,6 @@ "count": 1 } }, - "src/components/common_components/user_search_modal.tsx": { - "react-hooks/use-memo": { - "count": 1 - } - }, "src/components/constants.tsx": { "react-hooks/set-state-in-effect": { "count": 1 @@ -1860,45 +1993,11 @@ "count": 1 } }, - "src/components/mcp_tools/ByokCredentialModal.tsx": { - "no-restricted-syntax": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.test.tsx": { - "unused-imports/no-unused-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx": { - "react-hooks/immutability": { - "count": 2 - } - }, - "src/app/(dashboard)/mcp-servers/_components/MCPSubmissionsTab.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/mcp_tools/MCPToolArgumentsForm.tsx": { "no-nested-ternary": { "count": 5 } }, - "src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - }, - "unused-imports/no-unused-imports": { - "count": 2 - } - }, "src/components/mcp_tools/McpCrudPermissionPanel.tsx": { "no-nested-ternary": { "count": 3 @@ -1907,123 +2006,6 @@ "count": 1 } }, - "src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/OpenAPIQuickPicker.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.tsx": { - "no-nested-ternary": { - "count": 3 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx": { - "no-nested-ternary": { - "count": 2 - } - }, - "src/app/(dashboard)/mcp-servers/_components/create_mcp_server.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 4 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/static-components": { - "count": 4 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_connection_status.tsx": { - "no-nested-ternary": { - "count": 3 - }, - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_discovery.tsx": { - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_config.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_display.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/immutability": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 5 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_tools.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, "src/components/model_add/AddCredentialModal.tsx": { "no-restricted-imports": { "count": 1 @@ -2117,9 +2099,6 @@ "src/components/molecules/filter.tsx": { "no-nested-ternary": { "count": 2 - }, - "react-hooks/use-memo": { - "count": 1 } }, "src/components/molecules/models/columns.test.tsx": { @@ -2181,9 +2160,6 @@ }, "react-hooks/set-state-in-effect": { "count": 4 - }, - "react-hooks/use-memo": { - "count": 1 } }, "src/components/organization/organization_view.tsx": { @@ -2545,4 +2521,4 @@ "count": 1 } } -} +} \ No newline at end of file diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index b8f6265441b..91cf705060e 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -14,7 +14,7 @@ "@base-ui/react": "^1.6.0", "@headlessui/tailwindcss": "0.2.2", "@heroicons/react": "1.0.6", - "@tanstack/react-pacer": "0.2.0", + "@tanstack/react-pacer": "0.22.1", "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", "@tremor/react": "3.18.7", @@ -51,7 +51,6 @@ "@testing-library/jest-dom": "6.9.1", "@testing-library/react": "16.3.2", "@testing-library/user-event": "14.6.1", - "@types/lodash": "4.17.23", "@types/node": "20.19.37", "@types/react": "18.2.48", "@types/react-copy-to-clipboard": "5.0.7", @@ -3619,11 +3618,31 @@ "tailwindcss": "4.3.2" } }, - "node_modules/@tanstack/pacer": { - "version": "0.2.0", - "resolved": "https://registry.npmjs.org/@tanstack/pacer/-/pacer-0.2.0.tgz", - "integrity": "sha512-fUJs3NpSwtAL/tfq8kuYdgvm9HbbJvHsOG6aHY2dFDfff0NBFNwjvyGreWZZRPs2zgoIbr4nOk+rRV7aQgmf+A==", + "node_modules/@tanstack/devtools-event-client": { + "version": "0.4.4", + "resolved": "https://registry.npmjs.org/@tanstack/devtools-event-client/-/devtools-event-client-0.4.4.tgz", + "integrity": "sha512-6T5Yop/793YI+H+5J8Hsyj4kCih9sl4t3ElLgKioW5hk3ocn+ZdSJ94tT7vL7uabxSugWYBZlOTMPzEw2puvQw==", "license": "MIT", + "bin": { + "intent": "bin/intent.js" + }, + "engines": { + "node": ">=18" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + } + }, + "node_modules/@tanstack/pacer": { + "version": "0.21.1", + "resolved": "https://registry.npmjs.org/@tanstack/pacer/-/pacer-0.21.1.tgz", + "integrity": "sha512-hB01dd4rlsYcTCNP7wK186jgAe6K5qimgM1Y5Jtvz+9PUaILvpmeLLjmQNUNSO1l23lIt+CeQR6mO1mjlPvRtQ==", + "license": "MIT", + "dependencies": { + "@tanstack/devtools-event-client": "^0.4.3", + "@tanstack/store": "^0.11.0" + }, "engines": { "node": ">=18" }, @@ -3643,12 +3662,13 @@ } }, "node_modules/@tanstack/react-pacer": { - "version": "0.2.0", - "resolved": "https://registry.npmjs.org/@tanstack/react-pacer/-/react-pacer-0.2.0.tgz", - "integrity": "sha512-KU5GtjkKSeNdYCilen5Dc+Pu/6BPQbsQshKrUUjrg7URyJIiGBCz6ZZFre1QjDz/aeUeqUJWMWSm+2Dsh64v+w==", + "version": "0.22.1", + "resolved": "https://registry.npmjs.org/@tanstack/react-pacer/-/react-pacer-0.22.1.tgz", + "integrity": "sha512-CenQqK0GluSPIrnsG1yuD7w5uMSQ/4lI9AcGEFxBrRd66r260boWcYRIsS5+eHtXb238FoZYhKmJPGlhRzmHRw==", "license": "MIT", "dependencies": { - "@tanstack/pacer": "0.2.0" + "@tanstack/pacer": "0.21.1", + "@tanstack/react-store": "^0.11.0" }, "engines": { "node": ">=18" @@ -3678,6 +3698,24 @@ "react": "^18 || ^19" } }, + "node_modules/@tanstack/react-store": { + "version": "0.11.0", + "resolved": "https://registry.npmjs.org/@tanstack/react-store/-/react-store-0.11.0.tgz", + "integrity": "sha512-tX4YXh3PDkmpvGQWkWqKpzs/MSqbtuwY9dWdWhtV9Q50PmO+jOkUKIWIX4G85dwt7lxdHLXsiaEKPdKmC8F41w==", + "license": "MIT", + "dependencies": { + "@tanstack/store": "0.11.0", + "use-sync-external-store": "^1.6.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + }, + "peerDependencies": { + "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0", + "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" + } + }, "node_modules/@tanstack/react-table": { "version": "8.21.3", "resolved": "https://registry.npmjs.org/@tanstack/react-table/-/react-table-8.21.3.tgz", @@ -3715,6 +3753,16 @@ "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, + "node_modules/@tanstack/store": { + "version": "0.11.0", + "resolved": "https://registry.npmjs.org/@tanstack/store/-/store-0.11.0.tgz", + "integrity": "sha512-WlzzCt3xi0G6pCAJu1U+2jiECwabETDpQDi3hfkFZvJii9AuZqEKbOiVarX1/bWhTNjU486yQtJCCasi/0q+Cw==", + "license": "MIT", + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + } + }, "node_modules/@tanstack/table-core": { "version": "8.21.3", "resolved": "https://registry.npmjs.org/@tanstack/table-core/-/table-core-8.21.3.tgz", @@ -4060,13 +4108,6 @@ "dev": true, "license": "MIT" }, - "node_modules/@types/lodash": { - "version": "4.17.23", - "resolved": "https://registry.npmjs.org/@types/lodash/-/lodash-4.17.23.tgz", - "integrity": "sha512-RDvF6wTulMPjrNdCoYRC8gNR880JNGT8uB+REUpC2Ns4pRqQJhGz90wh7rgdXDPpCczF3VGktDuFGVnz8zP7HA==", - "dev": true, - "license": "MIT" - }, "node_modules/@types/mdast": { "version": "4.0.4", "resolved": "https://registry.npmjs.org/@types/mdast/-/mdast-4.0.4.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 759fc8fb2bc..7aea571ea8b 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -30,7 +30,7 @@ "@base-ui/react": "^1.6.0", "@headlessui/tailwindcss": "0.2.2", "@heroicons/react": "1.0.6", - "@tanstack/react-pacer": "0.2.0", + "@tanstack/react-pacer": "0.22.1", "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", "@tremor/react": "3.18.7", @@ -67,7 +67,6 @@ "@testing-library/jest-dom": "6.9.1", "@testing-library/react": "16.3.2", "@testing-library/user-event": "14.6.1", - "@types/lodash": "4.17.23", "@types/node": "20.19.37", "@types/react": "18.2.48", "@types/react-copy-to-clipboard": "5.0.7", @@ -95,7 +94,6 @@ "js-yaml": "4.2.0", "glob": "13.0.0", "minimatch": "10.2.4", - "lodash": "4.18.1", "ws": "8.21.0", "braces": "3.0.3", "axios": "1.13.6", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx index dc70aaa409b..4ee6332c54e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx @@ -1,6 +1,6 @@ import React from "react"; import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { screen, waitFor } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "@/../tests/test-utils"; import AgentCardDiscovery from "./agent_card_discovery"; @@ -243,6 +243,54 @@ describe("AgentCardDiscovery", () => { expect(selection.selected_card.capabilities.streaming).toBe(false); }); + it("does not fire discovery before the debounce wait and fires once with the last URL", async () => { + mockDiscover.mockResolvedValue({ + url: "https://last.example.com", + agent_card: sampleCard, + }); + renderWithProviders(); + const input = screen.getByPlaceholderText("https://upstream-agent.example.com"); + + act(() => { + fireEvent.change(input, { target: { value: "https://first.example.com" } }); + }); + act(() => { + vi.advanceTimersByTime(399); + }); + expect(mockDiscover).not.toHaveBeenCalled(); + + act(() => { + fireEvent.change(input, { target: { value: "https://last.example.com" } }); + }); + act(() => { + vi.advanceTimersByTime(399); + }); + expect(mockDiscover).not.toHaveBeenCalled(); + + await act(async () => { + await vi.advanceTimersByTimeAsync(1); + }); + expect(mockDiscover).toHaveBeenCalledTimes(1); + expect(mockDiscover).toHaveBeenCalledWith("tok", "https://last.example.com", undefined); + }); + + it("fires no discovery when unmounted mid-wait", () => { + const { unmount } = renderWithProviders(); + const input = screen.getByPlaceholderText("https://upstream-agent.example.com"); + + act(() => { + fireEvent.change(input, { target: { value: "https://first.example.com" } }); + }); + act(() => { + vi.advanceTimersByTime(200); + }); + unmount(); + act(() => { + vi.advanceTimersByTime(2000); + }); + expect(mockDiscover).not.toHaveBeenCalled(); + }); + it("blocks discover when no access token is provided", async () => { const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime }); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx index 5ea7458f643..e979b2dbe3f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx @@ -11,6 +11,7 @@ import { SearchOutlined, } from "@ant-design/icons"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { DiscoveredAgentCard, discoverAgentCardCall } from "@/components/networking"; import { ALLOWED_CAPABILITY_KEYS, @@ -22,6 +23,8 @@ import { const { Text, Paragraph } = Typography; const { Panel } = Collapse; +const DISCOVERY_DEBOUNCE_WAIT_MS = 400; + export interface DiscoveredAgentCardSelection { /** Full upstream card the proxy fetched, unmodified. */ raw_card: DiscoveredAgentCard; @@ -171,6 +174,14 @@ const AgentCardDiscovery: React.FC = ({ // eslint-disable-next-line react-hooks/exhaustive-deps }, [accessToken, effectiveUrl, isParentDriven, discoveryMode, discoveryParamsKey]); + const debouncedDiscover = useDebouncedCallback( + () => { + if (!accessToken || !effectiveUrl.trim()) return; + void handleDiscover(); + }, + { wait: DISCOVERY_DEBOUNCE_WAIT_MS }, + ); + // Auto-discover when the URL (or parent plan) becomes available. Debounce // is applied uniformly so rapid changes from a watched parent form (e.g. // typing into a LangGraph api_base / assistant_id field) don't fire one @@ -186,11 +197,8 @@ const AgentCardDiscovery: React.FC = ({ return; } - const timer = window.setTimeout(() => { - void handleDiscover(); - }, 400); - return () => window.clearTimeout(timer); - }, [accessToken, effectiveUrl, handleDiscover]); + debouncedDiscover(); + }, [accessToken, effectiveUrl, handleDiscover, debouncedDiscover]); const toggleSkill = (id: string, checked: boolean) => { setSelectedSkillIds((prev) => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx index 07cdcbc888c..4217a765732 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -1,6 +1,8 @@ "use client"; import React, { useState, useEffect, useCallback } from "react"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { SearchIcon, PlusIcon, @@ -728,6 +730,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { rejected: 0, }); const [search, setSearch] = useState(""); + const [searchDebounced] = useDebouncedValue(search, { wait: DEBOUNCE_WAIT_MS }); const [statusFilter, setStatusFilter] = useState<"all" | GuardrailStatus>("all"); const [selectedId, setSelectedId] = useState(null); const [expandedHeaders, setExpandedHeaders] = useState>(new Set()); @@ -737,16 +740,10 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { } | null>(null); const [isLoading, setIsLoading] = useState(true); const [error, setError] = useState(null); - const [searchDebounced, setSearchDebounced] = useState(""); const [isSubmitModalOpen, setIsSubmitModalOpen] = useState(false); const [submitForm] = Form.useForm(); const registerGuardrail = useRegisterGuardrail(); - useEffect(() => { - const t = setTimeout(() => setSearchDebounced(search), 300); - return () => clearTimeout(t); - }, [search]); - const fetchSubmissions = useCallback(async () => { if (!accessToken) { setIsLoading(false); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index 490e5e3dad1..f532c44ffd7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -84,6 +84,24 @@ export const teamListCall = async ( } }; +export const teamsTableKeys = createQueryKeys("teamsTable"); + +export const useTeamsTable = ( + page: number, + pageSize: number, + options: TeamListCallOptions = {}, +): UseQueryResult => { + const { accessToken } = useAuthorized(); + + return useQuery({ + queryKey: teamsTableKeys.list({ page, limit: pageSize, ...options }), + queryFn: async () => await teamListCall(accessToken!, page, pageSize, options), + enabled: Boolean(accessToken), + staleTime: 30000, + placeholderData: keepPreviousData, + }); +}; + const teamKeys = createQueryKeys("teams"); export const useTeams = (): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index 0afb4bd9314..f186fef22da 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -684,7 +684,6 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) refetch(); setByokModalServer(null); }} - accessToken={accessToken || ""} /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index 560a4ec15d0..4bf4af0c7f6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -296,7 +296,7 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te } return ( -
+
{/* Model Management Header */} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index 576729d18d3..20c9e805a0d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -14,11 +14,13 @@ import { useQueryClient } from "@tanstack/react-query"; import { Grid, TabPanel } from "@tremor/react"; import { Badge, Button, Select, Skeleton, Space, Typography } from "antd"; import ModelSettingsModal from "@/components/model_dashboard/ModelSettingsModal/ModelSettingsModal"; -import debounce from "lodash/debounce"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { useEffect, useMemo, useState } from "react"; import { useModelsInfo } from "../../hooks/models/useModels"; import { transformModelData } from "../utils/modelDataTransformer"; type ModelViewMode = "all" | "current_team"; + +const SEARCH_DEBOUNCE_WAIT_MS = 200; const { Text } = Typography; interface AllModelsTabProps { @@ -59,23 +61,17 @@ const AllModelsTab = ({ const [sorting, setSorting] = useState([]); const [isModelSettingsModalVisible, setIsModelSettingsModalVisible] = useState(false); - // Debounce search input - const debouncedUpdateSearch = useMemo( - () => - debounce((value: string) => { - setDebouncedSearch(value); - // Reset to page 1 when search changes - setCurrentPage(1); - setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 })); - }, 200), - [], + const debouncedUpdateSearch = useDebouncedCallback( + (value: string) => { + setDebouncedSearch(value); + setCurrentPage(1); + setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 })); + }, + { wait: SEARCH_DEBOUNCE_WAIT_MS }, ); useEffect(() => { debouncedUpdateSearch(modelNameSearch); - return () => { - debouncedUpdateSearch.cancel(); - }; }, [modelNameSearch, debouncedUpdateSearch]); // Determine teamId to pass to the query - only pass if not "personal" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx index d3af5b62668..87d8010759d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx @@ -191,7 +191,7 @@ const OrganizationsTable: React.FC = ({ } return ( -
+
{(userRole === "Admin" || userRole === "Org Admin") && ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 689abf66b41..d8f927b9a63 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -77,6 +77,7 @@ import { A2ATaskMetadata, MessageType } from "@/components/chat_ui/types"; import { useCodeInterpreter } from "../../hooks/useCodeInterpreter"; import { useChatHistory } from "../../hooks/useChatHistory"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; const { TextArea } = Input; const { Dragger } = Upload; @@ -99,6 +100,8 @@ interface ChatUIProps { const MCP_SUPPORTED_ENDPOINTS = new Set([EndpointType.CHAT, EndpointType.RESPONSES, EndpointType.MCP]); +const CUSTOM_MODEL_DEBOUNCE_WAIT_MS = 500; + const ChatUI: React.FC = ({ accessToken, token, @@ -185,7 +188,9 @@ const ChatUI: React.FC = ({ const [modelInfo, setModelInfo] = useState([]); const [agentInfo, setAgentInfo] = useState([]); const [selectedAgent, setSelectedAgent] = useState(undefined); - const customModelTimeout = useRef(null); + const debouncedSetSelectedModel = useDebouncedCallback((value: string) => setSelectedModel(value), { + wait: CUSTOM_MODEL_DEBOUNCE_WAIT_MS, + }); const [endpointType, setEndpointType] = useState( () => sessionStorage.getItem("endpointType") || EndpointType.CHAT, ); @@ -1255,16 +1260,7 @@ const ChatUI: React.FC = ({ { - // Using setTimeout to create a simple debounce effect - if (customModelTimeout.current) { - clearTimeout(customModelTimeout.current); - } - - customModelTimeout.current = setTimeout(() => { - setSelectedModel(value); - }, 500); // 500ms delay after typing stops - }} + onValueChange={debouncedSetSelectedModel} /> )}
@@ -2186,7 +2182,6 @@ const ChatUI: React.FC = ({ loadMCPServers(); setByokModalServer(null); }} - accessToken={accessToken || ""} /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx index 65157206f21..3b396924aba 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx @@ -1,7 +1,9 @@ "use client"; import NotificationsManager from "@/components/molecules/notifications_manager"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { ClearOutlined, DeleteOutlined, FilePdfOutlined, PlusOutlined } from "@ant-design/icons"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { Button, Input, Select, Tooltip } from "antd"; import { useEffect, useMemo, useState } from "react"; import { v4 as uuidv4 } from "uuid"; @@ -105,14 +107,8 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }: disabledPersonalKeyCreation ? "custom" : "session", ); const [customApiKey, setCustomApiKey] = useState(""); - const [debouncedCustomApiKey, setDebouncedCustomApiKey] = useState(""); + const [debouncedCustomApiKey] = useDebouncedValue(customApiKey, { wait: DEBOUNCE_WAIT_MS }); const [customProxyBaseUrl] = useState(() => sessionStorage.getItem("customProxyBaseUrl") || ""); - useEffect(() => { - const timer = setTimeout(() => { - setDebouncedCustomApiKey(customApiKey); - }, 300); - return () => clearTimeout(timer); - }, [customApiKey]); useEffect(() => { return () => { if (uploadedFilePreviewUrl) { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.test.ts index b61f27cf717..762e17a6f91 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.test.ts @@ -1,5 +1,5 @@ import { renderHook, act } from "@testing-library/react"; -import { describe, it, expect, beforeEach, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { useChatHistory } from "./useChatHistory"; describe("useChatHistory", () => { @@ -499,6 +499,80 @@ describe("useChatHistory", () => { }); }); + describe("debounced chatHistory persistence", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + }); + + it("should not write chatHistory to sessionStorage before the debounce wait elapses", () => { + const setItemSpy = vi.spyOn(Storage.prototype, "setItem"); + const { result } = renderHook(() => useChatHistory({ simplified: false })); + + act(() => { + result.current.updateTextUI("user", "hello"); + }); + act(() => { + vi.advanceTimersByTime(499); + }); + + expect(setItemSpy.mock.calls.filter(([key]) => key === "chatHistory")).toHaveLength(0); + + setItemSpy.mockRestore(); + }); + + it("should write chatHistory exactly once with the last value after the wait", () => { + const setItemSpy = vi.spyOn(Storage.prototype, "setItem"); + const { result } = renderHook(() => useChatHistory({ simplified: false })); + + act(() => { + result.current.updateTextUI("user", "h"); + }); + act(() => { + vi.advanceTimersByTime(300); + }); + act(() => { + result.current.updateTextUI("user", "i"); + }); + act(() => { + vi.advanceTimersByTime(499); + }); + + expect(setItemSpy.mock.calls.filter(([key]) => key === "chatHistory")).toHaveLength(0); + + act(() => { + vi.advanceTimersByTime(1); + }); + + const writes = setItemSpy.mock.calls.filter(([key]) => key === "chatHistory"); + expect(writes).toHaveLength(1); + expect(JSON.parse(writes[0][1])).toEqual([{ role: "user", content: "hi" }]); + + setItemSpy.mockRestore(); + }); + + it("should not write chatHistory when unmounted mid-wait", () => { + const setItemSpy = vi.spyOn(Storage.prototype, "setItem"); + const { result, unmount } = renderHook(() => useChatHistory({ simplified: false })); + + act(() => { + result.current.updateTextUI("user", "hello"); + }); + unmount(); + act(() => { + vi.advanceTimersByTime(1000); + }); + + expect(setItemSpy.mock.calls.filter(([key]) => key === "chatHistory")).toHaveLength(0); + + setItemSpy.mockRestore(); + }); + }); + describe("simplified mode session isolation", () => { it("should not hydrate messageTraceId from sessionStorage in simplified mode", () => { sessionStorage.setItem("messageTraceId", "trace-from-playground"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts index 6e2a263f599..7ee30757524 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts @@ -1,9 +1,12 @@ import React, { useState, useEffect } from "react"; +import { useDebouncer } from "@tanstack/react-pacer/debouncer"; import { MessageType, A2ATaskMetadata } from "@/components/chat_ui/types"; import { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; import { MCPEvent } from "@/components/mcp_tools/types"; import { truncateString } from "@/utils/textUtils"; +const CHAT_HISTORY_PERSIST_WAIT_MS = 500; + export interface UseChatHistoryReturn { // State chatHistory: MessageType[]; @@ -64,20 +67,20 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat return saved ? JSON.parse(saved) : true; // Default to API session management }); - // Debounced chatHistory persistence - useEffect(() => { - if (simplified) return; // Do not persist chat history in simplified (embedded) mode - // When chatHistory is empty (e.g. after clearChatHistory removed the key), - // don't re-write an empty array back into sessionStorage. - if (chatHistory.length === 0) return; - const handler = setTimeout(() => { - sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory)); - }, 500); // Debounce by 500ms + const persistDebouncer = useDebouncer( + (history: MessageType[]) => { + sessionStorage.setItem("chatHistory", JSON.stringify(history)); + }, + { wait: CHAT_HISTORY_PERSIST_WAIT_MS }, + ); - return () => { - clearTimeout(handler); - }; - }, [chatHistory, simplified]); + useEffect(() => { + if (simplified || chatHistory.length === 0) { + persistDebouncer.cancel(); + return; + } + persistDebouncer.maybeExecute(chatHistory); + }, [chatHistory, simplified, persistDebouncer]); // messageTraceId/responsesSessionId/useApiSessionManagement persistence useEffect(() => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx index 043d229e3c1..9c5a5592d5e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx @@ -115,7 +115,7 @@ const TagManagement: React.FC = ({ accessToken, userID, userRole }) => }, [accessToken]); return ( -
+
{selectedTagId ? ( = ({ teams, organizations }) => { // Debounced search for user selector const [userSearchInput, setUserSearchInput] = useState(""); const [debouncedUserSearch, setDebouncedUserSearch] = useDebouncedState("", { - wait: 300, + wait: DEBOUNCE_WAIT_MS, }); const { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index 18709e05df8..db3b17d6af3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -16,6 +16,7 @@ import { import OnboardingModal, { InvitationLink } from "@/components/onboarding_link"; import { updateExistingKeys } from "@/utils/dataUtils"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { isAdminRole, isProxyAdminRole } from "@/utils/roles"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { useQuery, useQueryClient } from "@tanstack/react-query"; @@ -86,7 +87,7 @@ const ViewUserDashboard: React.FC = ({ const [userToDelete, setUserToDelete] = useState(null); const [activeTab, setActiveTab] = useState("users"); const [filters, setFilters] = useState(initialFilters); - const [debouncedFilters, setDebouncedFilters, debouncer] = useDebouncedState(filters, { wait: 300 }); + const [debouncedFilters, setDebouncedFilters, debouncer] = useDebouncedState(filters, { wait: DEBOUNCE_WAIT_MS }); const [isInvitationLinkModalVisible, setIsInvitationLinkModalVisible] = useState(false); const [invitationLinkData, setInvitationLinkData] = useState(null); const [baseUrl, setBaseUrl] = useState(null); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx index fd02effc972..565e645f7b5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx @@ -138,7 +138,7 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID />
) : ( -
+

Vector Store Management

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.test.tsx index e73abfe7cd6..1aee3fcc8ab 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.test.tsx @@ -4,7 +4,10 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import WorkflowRuns from "./WorkflowRuns"; -vi.mock("@/components/networking", () => ({ proxyBaseUrl: "" })); +vi.mock("@/components/networking", () => ({ + proxyBaseUrl: "", + getGlobalLitellmHeaderName: () => "x-litellm-api-key", +})); interface FakeRun { run_id: string; @@ -78,4 +81,18 @@ describe("WorkflowRuns (migrated onto shared DataTable)", () => { expect(await screen.findByText("No workflow runs yet")).toBeInTheDocument(); }); + + it("sends the configured litellm key header on every fetch instead of hardcoding Authorization", async () => { + const user = userEvent.setup(); + const fetchSpy = mockFetch(RUNS); + vi.stubGlobal("fetch", fetchSpy); + render(); + + await user.click(await screen.findByText("First run")); + + await waitFor(() => expect(fetchSpy).toHaveBeenCalledTimes(3)); + for (const [url, init] of fetchSpy.mock.calls as [string, RequestInit][]) { + expect(init.headers, url).toEqual({ "x-litellm-api-key": "Bearer tok" }); + } + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.tsx index 9afa07251c2..7354b7479f4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.tsx @@ -2,7 +2,7 @@ import React, { useState, useEffect, useCallback, useMemo } from "react"; import { Button, Collapse, Drawer, Empty, Spin, Tooltip, Typography } from "antd"; import { ReloadOutlined } from "@ant-design/icons"; import type { ColumnDef, ColumnFiltersState } from "@tanstack/react-table"; -import { proxyBaseUrl } from "@/components/networking"; +import { getGlobalLitellmHeaderName, proxyBaseUrl } from "@/components/networking"; import { DataTable, DataTableFilterDrawer, @@ -507,7 +507,7 @@ const WorkflowRuns: React.FC = ({ accessToken }) => { setLoadingRuns(true); try { const res = await fetch(`${proxyBaseUrl ?? ""}/v1/workflows/runs?limit=100`, { - headers: { Authorization: `Bearer ${accessToken}` }, + headers: { [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}` }, }); if (!res.ok) throw new Error(`HTTP ${res.status}`); const data = await res.json(); @@ -531,10 +531,10 @@ const WorkflowRuns: React.FC = ({ accessToken }) => { const base = proxyBaseUrl ?? ""; const [evRes, msgRes] = await Promise.all([ fetch(`${base}/v1/workflows/runs/${run.run_id}/events`, { - headers: { Authorization: `Bearer ${accessToken}` }, + headers: { [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}` }, }), fetch(`${base}/v1/workflows/runs/${run.run_id}/messages`, { - headers: { Authorization: `Bearer ${accessToken}` }, + headers: { [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}` }, }), ]); const evData = evRes.ok ? await evRes.json() : { events: [] }; diff --git a/ui/litellm-dashboard/src/app/globals.css b/ui/litellm-dashboard/src/app/globals.css index 0b555fcdd52..4589d0f528a 100644 --- a/ui/litellm-dashboard/src/app/globals.css +++ b/ui/litellm-dashboard/src/app/globals.css @@ -11,6 +11,25 @@ @custom-variant dark (&:where(.dark, .dark *)); +/* shadcn Base UI primitives reference these variants; upstream omits them (shadcn-ui/ui#9196) */ +@custom-variant data-open (&:where([data-state="open"], [data-open]:not([data-open="false"]))); +@custom-variant data-closed (&:where([data-state="closed"], [data-closed]:not([data-closed="false"]))); +@custom-variant data-checked (&:where([data-state="checked"], [data-checked]:not([data-checked="false"]))); +@custom-variant data-unchecked (&:where([data-state="unchecked"], [data-unchecked]:not([data-unchecked="false"]))); +@custom-variant data-selected (&:where([data-selected="true"])); +@custom-variant data-disabled (&:where([data-disabled="true"], [data-disabled]:not([data-disabled="false"]))); +@custom-variant data-active (&:where([data-state="active"], [data-active]:not([data-active="false"]))); +@custom-variant data-horizontal (&:where([data-orientation="horizontal"])); +@custom-variant data-vertical (&:where([data-orientation="vertical"])); + +@utility no-scrollbar { + -ms-overflow-style: none; + scrollbar-width: none; + &::-webkit-scrollbar { + display: none; + } +} + :root { --radius: 0.5rem; --background: oklch(1 0 0); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index d3dd914f5c5..eb47f775d90 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -382,7 +382,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, } return ( -
+
{publicPage == false ? (
{/* Header with Title, Description and URL */} diff --git a/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx b/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx index d42d5ab324a..1d19ba3255d 100644 --- a/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx +++ b/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx @@ -1,4 +1,5 @@ import { useInfiniteKeyAliases } from "@/app/(dashboard)/hooks/keys/useKeyAliases"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { LoadingOutlined } from "@ant-design/icons"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { Select } from "antd"; @@ -16,7 +17,6 @@ export interface PaginatedKeyAliasSelectProps { } const SCROLL_THRESHOLD = 0.8; -const DEBOUNCE_MS = 300; export const PaginatedKeyAliasSelect = ({ value, @@ -30,7 +30,7 @@ export const PaginatedKeyAliasSelect = ({ }: PaginatedKeyAliasSelectProps) => { const [searchInput, setSearchInput] = useState(""); const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { - wait: DEBOUNCE_MS, + wait: DEBOUNCE_WAIT_MS, }); const teamId = allFilters?.["Team ID"] || undefined; diff --git a/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx b/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx index da3ecf77ded..a77b2cf561e 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx @@ -1,4 +1,5 @@ import { useInfiniteModelInfo } from "@/app/(dashboard)/hooks/models/useModels"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { LoadingOutlined } from "@ant-design/icons"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { Select, Space, Typography } from "antd"; @@ -17,7 +18,6 @@ export interface PaginatedModelSelectProps { } const SCROLL_THRESHOLD = 0.8; -const DEBOUNCE_MS = 300; export const PaginatedModelSelect = ({ value, @@ -30,7 +30,7 @@ export const PaginatedModelSelect = ({ }: PaginatedModelSelectProps) => { const [searchInput, setSearchInput] = useState(""); const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { - wait: DEBOUNCE_MS, + wait: DEBOUNCE_WAIT_MS, }); const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteModelInfo( diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index ee84e17d0c9..f885b582dd2 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -5,11 +5,20 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; import { fetchMCPAccessGroups, getGuardrailsList, teamCreateCall } from "./networking"; import Teams from "./Teams"; -import { teamListCall } from "@/app/(dashboard)/hooks/teams/useTeams"; const mockTeamInfoView = vi.fn(); const mockUseOrganizations = vi.fn(); +// The teams grid is unit-tested in TeamsPage/TeamsTable.test.tsx. Here we stub it and drive its callbacks +// directly so we can test the Teams shell wiring (delete modal, detail view) without the real DataTable. +let mockTeamsTableProps: any = null; +vi.mock("./TeamsPage/TeamsTable", () => ({ + TeamsTable: (props: any) => { + mockTeamsTableProps = props; + return
; + }, +})); + vi.mock("./networking", () => ({ teamCreateCall: vi.fn(), teamDeleteCall: vi.fn(), @@ -19,8 +28,9 @@ vi.mock("./networking", () => ({ getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), })); +// Teams invalidates teamsTableKeys on mutations; the selected team is passed up from the table. vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ - teamListCall: vi.fn().mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 0 }), + teamsTableKeys: { all: ["teamsTable"] }, })); vi.mock("./molecules/notifications_manager", () => ({ @@ -116,6 +126,21 @@ vi.mock("./common_components/AccessGroupSelector", () => ({ ), })); +const baseTableTeam = { + team_id: "1", + team_alias: "Test Team", + organization_id: "org-123", + models: ["gpt-4"], + max_budget: 100, + budget_duration: "1d", + tpm_limit: 1000, + rpm_limit: 1000, + created_at: new Date().toISOString(), + keys: [], + members_with_roles: [], + spend: 0, +}; + const createQueryClient = () => { return new QueryClient({ defaultOptions: { @@ -131,10 +156,16 @@ const renderWithQueryClient = (component: React.ReactElement) => { return render({component}); }; +// Re-establish safe defaults before every test (clearAllMocks keeps return values, so restore them here). +beforeEach(() => { + mockTeamsTableProps = null; +}); + describe("Teams - handleCreate organization handling", () => { beforeEach(() => { vi.clearAllMocks(); mockTeamInfoView.mockClear(); + mockTeamsTableProps = null; vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue([]); vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); @@ -142,7 +173,6 @@ describe("Teams - handleCreate organization handling", () => { }); it("should not include organization_id when it's an empty string", async () => { - const mockAccessToken = "test-token"; const formValues: Record = { team_alias: "Test Team", organization_id: "", // Empty string @@ -168,7 +198,6 @@ describe("Teams - handleCreate organization handling", () => { models: [], }; - // Simulate the handleCreate logic let organizationId = formValues?.organization_id || null; if (organizationId === "" || typeof organizationId !== "string") { formValues.organization_id = null; @@ -182,11 +211,10 @@ describe("Teams - handleCreate organization handling", () => { it("should trim and keep valid organization_id string", async () => { const formValues: Record = { team_alias: "Test Team", - organization_id: " org-123 ", // String with whitespace + organization_id: " org-123 ", models: [], }; - // Simulate the handleCreate logic let organizationId = formValues?.organization_id || null; if (organizationId === "" || typeof organizationId !== "string") { formValues.organization_id = null; @@ -204,7 +232,6 @@ describe("Teams - handleCreate organization handling", () => { models: [], }; - // Simulate the handleCreate logic let organizationId = formValues?.organization_id || null; if (organizationId === "" || typeof organizationId !== "string") { formValues.organization_id = null; @@ -223,7 +250,6 @@ describe("Teams - handleCreate organization handling", () => { max_budget: 100, }; - // Simulate the handleCreate logic let organizationId = formValues?.organization_id || null; if (organizationId === "" || typeof organizationId !== "string") { formValues.organization_id = null; @@ -231,18 +257,13 @@ describe("Teams - handleCreate organization handling", () => { formValues.organization_id = organizationId.trim(); } - // Verify the structure expect(formValues).toEqual({ team_alias: "Test Team", organization_id: null, models: ["gpt-4"], max_budget: 100, }); - - // Verify we're not sending an empty string expect(formValues.organization_id).not.toBe(""); - - // Verify it's explicitly null, not undefined expect(formValues.organization_id).toBeNull(); }); @@ -259,7 +280,6 @@ describe("Teams - handleCreate organization handling", () => { models: [], }; - // Simulate the handleCreate logic with currentOrg fallback let organizationId = formValues?.organization_id || currentOrg?.organization_id; if (organizationId === "" || typeof organizationId !== "string") { formValues.organization_id = null; @@ -270,204 +290,27 @@ describe("Teams - handleCreate organization handling", () => { expect(formValues.organization_id).toBe("fallback-org-id"); }); - it("should not include organizations as an empty array in the request payload", async () => { - const mockTeamCreateCall = vi.mocked(teamCreateCall); - const mockAccessToken = "test-token"; - - const formValues = { - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - organizations: [], // This should never be sent - }; - - // Remove organizations key if it's empty - if (Array.isArray(formValues.organizations) && formValues.organizations.length === 0) { - delete (formValues as any).organizations; - } - - // Verify organizations key is removed - expect(formValues).not.toHaveProperty("organizations"); - expect(formValues).toEqual({ - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - }); - }); - - it("should handle organization_id validation for org admins", () => { - // This test simulates the validation that should happen for org admins - const isOrgAdmin = true; - const formValues: Record = { - team_alias: "Test Team", - // organization_id is missing/undefined - }; - - // For org admins, organization_id should be required - const hasOrganization = - formValues.organization_id !== undefined && - formValues.organization_id !== null && - formValues.organization_id !== ""; - - if (isOrgAdmin && !hasOrganization) { - // This should trigger validation error - expect(hasOrganization).toBe(false); - } - }); - - it("should allow null organization_id for global admins", () => { - const isAdmin = true; - const formValues: Record = { - team_alias: "Test Team", - organization_id: null, - models: [], - }; - - // Global admins can create teams without an organization - if (isAdmin) { - expect(formValues.organization_id).toBeNull(); - // This is valid for admins - } - }); - - it("should ensure organization_id is never an empty list", () => { - const invalidFormValues: Record = { - team_alias: "Test Team", - organization_id: [], // Wrong type - should be string or null - }; - - // Type check: organization_id should never be an array - expect(Array.isArray(invalidFormValues.organization_id)).toBe(true); - - // Correct it to null - if (Array.isArray(invalidFormValues.organization_id)) { - invalidFormValues.organization_id = null; - } - - expect(invalidFormValues.organization_id).toBeNull(); - expect(Array.isArray(invalidFormValues.organization_id)).toBe(false); - }); - - it("should clear the delete modal when the cancel button is clicked", async () => { + it("opens the delete modal when the table's delete action fires", async () => { mockUseOrganizations.mockReturnValue({ data: [] }); - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); renderWithQueryClient(); - await waitFor(() => { - expect(screen.getByTestId("delete-team-button")).toBeInTheDocument(); - }); - const deleteTeamButton = screen.getByTestId("delete-team-button"); - act(() => { - fireEvent.click(deleteTeamButton); + + await waitFor(() => expect(mockTeamsTableProps).not.toBeNull()); + await act(async () => { + mockTeamsTableProps.onDeleteTeam(baseTableTeam); }); + expect(screen.getByText("Delete Team?")).toBeInTheDocument(); }); }); -describe("Teams - empty state", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - it("should display empty state message when teams array is empty", async () => { - vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); - renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("No teams yet")).toBeInTheDocument(); - }); - expect( - screen.getByText("Create your first team to organize members and manage access to models."), - ).toBeInTheDocument(); - }); - - it("should display empty state message when teams is null", async () => { - vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); - renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("No teams yet")).toBeInTheDocument(); - }); - expect( - screen.getByText("Create your first team to organize members and manage access to models."), - ).toBeInTheDocument(); - }); - - it("should not display empty state when teams array has items", async () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("Test Team")).toBeInTheDocument(); - }); - expect(screen.queryByText("No teams yet")).not.toBeInTheDocument(); - expect( - screen.queryByText("Create your first team to organize members and manage access to models."), - ).not.toBeInTheDocument(); - }); -}); - describe("Teams - helper functions", () => { describe("getAdminOrganizations", () => { it("should return all organizations for Admin role", () => { const organizations = [ - { - organization_id: "org-1", - organization_alias: "Org 1", - models: [], - members: [], - }, - { - organization_id: "org-2", - organization_alias: "Org 2", - models: [], - members: [], - }, + { organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }, + { organization_id: "org-2", organization_alias: "Org 2", models: [], members: [] }, ]; - // Simulate getAdminOrganizations logic for Admin const userRole = "Admin"; const result = userRole === "Admin" ? organizations : []; @@ -477,7 +320,6 @@ describe("Teams - helper functions", () => { it("should return only org_admin organizations for Org Admin role", () => { const userID = "user-123"; - const userRole = "Org Admin"; const organizations = [ { organization_id: "org-1", @@ -499,7 +341,6 @@ describe("Teams - helper functions", () => { }, ]; - // Simulate getAdminOrganizations logic const result = organizations.filter((org) => org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), ); @@ -519,7 +360,6 @@ describe("Teams - helper functions", () => { }, ]; - // Simulate getAdminOrganizations logic const result = organizations.filter((org) => org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), ); @@ -531,8 +371,7 @@ describe("Teams - helper functions", () => { describe("canCreateOrManageTeams", () => { it("should return true for Admin role", () => { const userRole = "Admin"; - const result = userRole === "Admin"; - expect(result).toBe(true); + expect(userRole === "Admin").toBe(true); }); it("should return true for org_admin in any organization", () => { @@ -577,6 +416,7 @@ describe("Teams - helper functions", () => { describe("Teams - premium props", () => { beforeEach(() => { + vi.clearAllMocks(); mockTeamInfoView.mockClear(); vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue([]); vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); @@ -584,38 +424,14 @@ describe("Teams - premium props", () => { mockUseOrganizations.mockReturnValue({ data: [] }); }); - it("passes premiumUser flag to TeamInfoView", async () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "team-123456789", - team_alias: "Premium Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); + it("passes premiumUser flag to TeamInfoView when a team is opened", async () => { + const premiumTeam = { ...baseTableTeam, team_id: "team-123456789", team_alias: "Premium Team" }; renderWithQueryClient(); - const teamIdElement = await screen.findByText("team-123456789"); - act(() => { - fireEvent.click(teamIdElement); - }); + await waitFor(() => expect(mockTeamsTableProps).not.toBeNull()); + act(() => mockTeamsTableProps.onSelectTeam(premiumTeam)); await waitFor(() => expect(mockTeamInfoView).toHaveBeenCalled()); - expect(mockTeamInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ premiumUser: true })); }); }); @@ -627,114 +443,22 @@ describe("Teams - Default Team Settings tab visibility", () => { }); it("should show Default Team Settings tab for Admin role", () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); renderWithQueryClient(); - expect(screen.getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument(); }); it("should show Default Team Settings tab for proxy_admin role", () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); renderWithQueryClient(); - expect(screen.getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument(); }); it("should not show Default Team Settings tab for proxy_admin_viewer role", () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); renderWithQueryClient(); - expect(screen.queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument(); }); it("should not show Default Team Settings tab for Admin Viewer role", () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); renderWithQueryClient(); - expect(screen.queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument(); }); }); @@ -761,7 +485,6 @@ describe("Teams - access_group_ids in team create", () => { }); it("should pass access_group_ids to teamCreateCall when creating team", async () => { - vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); renderWithQueryClient(); const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; @@ -773,25 +496,19 @@ describe("Teams - access_group_ids in team create", () => { expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); }); - const teamNameInput = screen.getByLabelText(/team name/i); - fireEvent.change(teamNameInput, { target: { value: "Test Team" } }); + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Test Team" } }); + fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); - const modelsInput = screen.getByTestId("create-team-models-select"); - fireEvent.change(modelsInput, { target: { value: "gpt-4" } }); - - const additionalSettingsAccordion = screen.getByText("Additional Settings"); - fireEvent.click(additionalSettingsAccordion); + fireEvent.click(screen.getByText("Additional Settings")); await waitFor(() => { expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); }); - const accessGroupInput = screen.getByTestId("access-group-selector"); - fireEvent.change(accessGroupInput, { target: { value: "ag-1,ag-2" } }); + fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } }); const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); - const createTeamSubmitButton = createTeamSubmitButtons[createTeamSubmitButtons.length - 1]; - fireEvent.click(createTeamSubmitButton); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); await waitFor(() => { expect(teamCreateCall).toHaveBeenCalledWith( @@ -814,9 +531,6 @@ describe("Teams - models dropdown options", () => { }); it("should not render all-proxy-models option in models select", async () => { - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); - - vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); renderWithQueryClient(); await waitFor(() => { @@ -831,207 +545,7 @@ describe("Teams - models dropdown options", () => { await waitFor(() => { expect(screen.getByLabelText(/models/i)).toBeInTheDocument(); }); - const allProxyModelsOption = screen.queryByText("All Proxy Models"); - expect(allProxyModelsOption).not.toBeInTheDocument(); - }); -}); - -describe("Teams - organization alias display", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - it("should display organization alias instead of organization id", async () => { - const mockOrganizations = [ - { - organization_id: "org-123", - organization_alias: "Test Organization", - budget_id: "budget-1", - metadata: {}, - models: [], - spend: 0, - model_spend: {}, - created_at: new Date().toISOString(), - created_by: "user-1", - updated_at: new Date().toISOString(), - updated_by: "user-1", - litellm_budget_table: null, - teams: null, - users: null, - members: null, - }, - ]; - - mockUseOrganizations.mockReturnValue({ data: mockOrganizations }); - - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("Test Organization")).toBeInTheDocument(); - }); - expect(screen.queryByText("org-123")).not.toBeInTheDocument(); - }); - - it("should display organization id when alias is not found", async () => { - mockUseOrganizations.mockReturnValue({ data: [] }); - - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-unknown", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("org-unknown")).toBeInTheDocument(); - }); - }); - - it("should display N/A when organization_id is null", async () => { - mockUseOrganizations.mockReturnValue({ data: [] }); - - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: null, - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - await waitFor(() => { - // When organization_id is null, the table shows "—" in the Organization column - expect(screen.getAllByText("—").length).toBeGreaterThan(0); - }); - }); -}); - -describe("Teams - Resources column keys badge", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - it("renders keys_count from the v2 payload in the Resources badge", async () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Team With Keys", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - keys_count: 3, - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - const { container } = renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("Team With Keys")).toBeInTheDocument(); - }); - const cyanTag = container.querySelector(".ant-tag-cyan"); - expect(cyanTag).not.toBeNull(); - expect(cyanTag?.textContent).toContain("3"); - }); - - it("falls back to keys.length when keys_count is absent", async () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "2", - team_alias: "Legacy Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [{ token: "t1" }, { token: "t2" }], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - const { container } = renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("Legacy Team")).toBeInTheDocument(); - }); - const cyanTag = container.querySelector(".ant-tag-cyan"); - expect(cyanTag).not.toBeNull(); - expect(cyanTag?.textContent).toContain("2"); + expect(screen.queryByText("All Proxy Models")).not.toBeInTheDocument(); }); }); @@ -1042,39 +556,16 @@ describe("Teams - delete team warning copy", () => { }); const openDeleteModal = async (team: any) => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [team], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); renderWithQueryClient(); - await waitFor(() => { - expect(screen.getByTestId("delete-team-button")).toBeInTheDocument(); - }); - act(() => { - fireEvent.click(screen.getByTestId("delete-team-button")); + await waitFor(() => expect(mockTeamsTableProps).not.toBeNull()); + await act(async () => { + mockTeamsTableProps.onDeleteTeam(team); }); expect(screen.getByText("Delete Team?")).toBeInTheDocument(); }; - const baseTeam = { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - members_with_roles: [], - spend: 0, - }; - it("warns that the team's models are deleted when the team has keys", async () => { - await openDeleteModal({ ...baseTeam, keys: [], keys_count: 5 }); + await openDeleteModal({ ...baseTableTeam, keys: [], keys_count: 5 }); expect(screen.getByText(/Warning: This team has 5 keys associated with it/i)).toHaveTextContent( /along with any models created for this team/i, @@ -1085,7 +576,7 @@ describe("Teams - delete team warning copy", () => { }); it("still warns about model deletion in the confirmation message when the team has no keys", async () => { - await openDeleteModal({ ...baseTeam, keys: [], keys_count: 0 }); + await openDeleteModal({ ...baseTableTeam, keys: [], keys_count: 0 }); expect(screen.queryByText(/Warning: This team has/i)).not.toBeInTheDocument(); expect(screen.getByText(/Are you sure you want to delete this team/i)).toHaveTextContent( @@ -1101,7 +592,6 @@ describe("Teams - LIT-2530 organization stays optional for proxy admin with a si vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4"]); vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); - vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); vi.mocked(teamCreateCall).mockResolvedValue({ team_id: "new-team-1", team_alias: "No Org Team", diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index e83d1acf4e5..57a5e1cbff0 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -3,38 +3,16 @@ import AvailableTeamsPanel from "@/components/team/available_teams"; import TeamInfoView from "@/components/team/TeamInfo"; import TeamSSOSettings from "@/components/TeamSSOSettings"; import { isProxyAdminRole } from "@/utils/roles"; -import { InfoCircleOutlined, PlusOutlined, TeamOutlined, ReloadOutlined } from "@ant-design/icons"; +import { InfoCircleOutlined } from "@ant-design/icons"; import { Accordion, AccordionBody, AccordionHeader, TextInput } from "@tremor/react"; -import { - Button, - Card, - Flex, - Form, - Input, - Layout, - Modal, - Pagination, - Progress, - Select, - Space, - Switch, - Table, - Tabs, - Tag, - theme, - Tooltip, - Typography, - message, -} from "antd"; -import type { ColumnsType } from "antd/es/table"; -import type { SorterResult } from "antd/es/table/interface"; -import { KeyIcon, LayersIcon, SearchIcon, UsersIcon } from "lucide-react"; -import React, { useEffect, useMemo, useRef, useState } from "react"; -import { AntDLoadingSpinner } from "@/components/ui/AntDLoadingSpinner"; -import { DateCell, IdCell } from "@/components/shared/table_cells"; -import OrganizationDropdown from "./common_components/OrganizationDropdown"; -import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; -import { teamListCall as v2TeamListCall, type TeamsResponse } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { Button, Form, Input, Layout, Modal, Select, Switch, Tabs, theme, Tooltip, Typography } from "antd"; +import { Plus, Users } from "lucide-react"; +import React, { useEffect, useState } from "react"; +import { useQueryClient } from "@tanstack/react-query"; +import { PageHeader } from "@/components/shared/PageHeader"; +import { Button as UIButton } from "@/components/ui/button"; +import { teamsTableKeys } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { TeamsTable } from "./TeamsPage/TeamsTable"; import AccessGroupSelector from "./common_components/AccessGroupSelector"; import PassThroughRoutesSelector from "./common_components/PassThroughRoutesSelector"; import AgentSelector from "./agent_management/AgentSelector"; @@ -45,7 +23,7 @@ import { fetchAvailableModelsForTeamOrKey, unfurlWildcardModelsInList, } from "./key_team_helpers/fetch_available_models_team_key"; -import type { KeyResponse, Team } from "./key_team_helpers/key_list"; +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"; @@ -61,13 +39,6 @@ interface TeamProps { premiumUser?: boolean; } -interface FilterState { - search: string; - organization_id: string; - sort_by: string; - sort_order: "asc" | "desc"; -} - interface EditTeamModalProps { visible: boolean; onCancel: () => void; @@ -75,21 +46,10 @@ interface EditTeamModalProps { onSubmit: (data: FormData) => void; // Assuming FormData is the type of data to be submitted } -import { updateExistingKeys } from "@/utils/dataUtils"; import DeleteResourceModal from "./common_components/DeleteResourceModal"; -import { Member, teamCreateCall } from "./networking"; +import { teamCreateCall } from "./networking"; import { ModelSelect } from "./ModelSelect/ModelSelect"; -interface TeamInfo { - members_with_roles: Member[]; -} - -interface PerTeamInfo { - keys: KeyResponse[]; - keys_count: number; - team_info: TeamInfo; -} - const getOrganizationModels = (organization: Organization | null, userModels: string[]) => { let tempModelsToPick = []; @@ -164,70 +124,17 @@ const getOrganizationAlias = ( const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser = false }) => { const { data: organizationsData } = useOrganizations(); const organizations = organizationsData ?? null; - const [teams, setTeams] = useState(null); - const [isLoading, setIsLoading] = useState(true); - const [fetchError, setFetchError] = useState(null); - const [currentPage, setCurrentPage] = useState(1); - const [pageSize, setPageSize] = useState(10); - const [totalTeams, setTotalTeams] = useState(0); - const [currentOrg, setCurrentOrg] = useState(null); + const queryClient = useQueryClient(); + const refreshTeams = () => queryClient.invalidateQueries({ queryKey: teamsTableKeys.all }); + const [currentOrg] = useState(null); const [currentOrgForCreateTeam, setCurrentOrgForCreateTeam] = useState(null); - const [filters, setFilters] = useState({ - search: "", - organization_id: "", - sort_by: "created_at", - sort_order: "desc", - }); - const searchDebounceRef = useRef | null>(null); - const [isSearching, setIsSearching] = useState(false); - - const fetchTeamsV2 = async ( - opts: { - page?: number; - size?: number; - sortBy?: string; - sortOrder?: string; - organizationID?: string; - search?: string; - } = {}, - ) => { - if (!accessToken) return; - const page = opts.page ?? currentPage; - const size = opts.size ?? pageSize; - const sortBy = opts.sortBy ?? filters.sort_by; - const sortOrder = opts.sortOrder ?? filters.sort_order; - const organizationID = opts.organizationID ?? filters.organization_id; - const search = opts.search ?? filters.search; - - setIsLoading(true); - setFetchError(null); - try { - const response: TeamsResponse = await v2TeamListCall(accessToken, page, size, { - organizationID: organizationID || null, - search: search || null, - userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null, - sortBy: sortBy || null, - sortOrder: sortOrder || null, - }); - setTeams(response.teams ?? []); - setTotalTeams(response.total ?? 0); - } catch (err: any) { - setFetchError(err?.message || "Failed to fetch teams"); - } finally { - setIsLoading(false); - } - }; - - useEffect(() => { - fetchTeamsV2(); - }, [accessToken]); const [form] = Form.useForm(); const [memberForm] = Form.useForm(); const [value, setValue] = useState(""); const [editModalVisible, setEditModalVisible] = useState(false); - const [selectedTeam, setSelectedTeam] = useState(null); + const [selectedTeam, setSelectedTeam] = useState(null); const [selectedTeamId, setSelectedTeamId] = useState(null); const [editTeam, setEditTeam] = useState(false); @@ -238,7 +145,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [teamToDelete, setTeamToDelete] = useState(null); const [modelsToPick, setModelsToPick] = useState([]); - const [perTeamInfo, setPerTeamInfo] = useState>({}); const [isTeamDeleting, setIsTeamDeleting] = useState(false); // Add this state near the other useState declarations const [guardrailsList, setGuardrailsList] = useState([]); @@ -325,30 +231,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser fetchMcpAccessGroups(); }, [accessToken]); - useEffect(() => { - const fetchTeamInfo = () => { - if (!teams) return; - - const newPerTeamInfo = teams.reduce( - (acc, team) => { - acc[team.team_id] = { - keys: team.keys || [], - keys_count: team.keys_count ?? team.keys?.length ?? 0, - team_info: { - members_with_roles: team.members_with_roles || [], - }, - }; - return acc; - }, - {} as Record, - ); - - setPerTeamInfo(newPerTeamInfo); - }; - - fetchTeamInfo(); - }, [teams]); - const handleOk = () => { setIsTeamModalVisible(false); form.resetFields(); @@ -386,14 +268,14 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser }; const confirmDelete = async () => { - if (teamToDelete == null || teams == null || accessToken == null) { + if (teamToDelete == null || accessToken == null) { return; } try { setIsTeamDeleting(true); await teamDeleteCall(accessToken, teamToDelete.team_id); - await fetchTeamsV2(); + await refreshTeams(); NotificationsManager.success("Team deleted successfully"); } catch (error) { NotificationsManager.fromBackend("Error deleting the team: " + error); @@ -425,13 +307,11 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser }; fetchUserModels(); - }, [accessToken, userID, userRole, teams]); + }, [accessToken, userID, userRole]); const handleCreate = async (formValues: Record) => { try { if (accessToken != null) { - const newTeamAlias = formValues?.team_alias; - const existingTeamAliases = teams?.map((t) => t.team_alias) ?? []; let organizationId = formValues?.organization_id || currentOrg?.organization_id; if (organizationId === "" || typeof organizationId !== "string") { formValues.organization_id = null; @@ -439,11 +319,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser formValues.organization_id = organizationId.trim(); } - // Remove guardrails from top level since it's now in metadata - if (existingTeamAliases.includes(newTeamAlias)) { - throw new Error(`Team alias ${newTeamAlias} already exists, please pick another alias`); - } - NotificationsManager.info("Creating Team"); // Handle logging settings in metadata @@ -565,10 +440,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser await teamCreateCall(accessToken, formValues); NotificationsManager.success("Team created"); - await fetchTeamsV2({ - page: currentPage, - size: pageSize, - }); + await refreshTeams(); form.resetFields(); setLoggingSettings([]); setModelAliases({}); @@ -595,352 +467,31 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser return false; }; - const handleSearchChange = (value: string) => { - if (searchDebounceRef.current) clearTimeout(searchDebounceRef.current); - setIsSearching(true); - searchDebounceRef.current = setTimeout(async () => { - try { - setFilters((prev) => ({ ...prev, search: value })); - setCurrentPage(1); - await fetchTeamsV2({ page: 1, search: value }); - } finally { - setIsSearching(false); - } - }, 300); - }; - - const handleFilterChange = async (key: keyof FilterState, value: string) => { - const newFilters = { ...filters, [key]: value }; - setFilters(newFilters); - setCurrentPage(1); - if (!accessToken) return; - try { - const response: TeamsResponse = await v2TeamListCall(accessToken, 1, pageSize, { - organizationID: newFilters.organization_id || null, - search: newFilters.search || null, - userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null, - sortBy: newFilters.sort_by || null, - sortOrder: newFilters.sort_order || null, - }); - setTeams(response.teams ?? []); - setTotalTeams(response.total ?? 0); - } catch (error) { - console.error("Error fetching teams:", error); - } - }; - - const handleFilterReset = () => { - if (searchDebounceRef.current) clearTimeout(searchDebounceRef.current); - setIsSearching(false); - const resetFilters: FilterState = { - search: "", - organization_id: "", - sort_by: "created_at", - sort_order: "desc", - }; - setFilters(resetFilters); - setCurrentPage(1); - fetchTeamsV2({ page: 1, organizationID: "", search: "", sortBy: "created_at", sortOrder: "desc" }); - }; - const { token } = theme.useToken(); - const { Title, Text } = Typography; + const { Text } = Typography; const { Content } = Layout; - const handleRetry = () => { - fetchTeamsV2(); - }; - - const handleTableSort = ( - _pagination: unknown, - _filters: unknown, - sorter: SorterResult | SorterResult[], - ) => { - const s = Array.isArray(sorter) ? sorter[0] : sorter; - const sortBy = s.order ? (s.columnKey as string) : "created_at"; - const sortOrder = s.order === "ascend" ? "asc" : s.order === "descend" ? "desc" : "desc"; - setFilters((prev) => ({ ...prev, sort_by: sortBy, sort_order: sortOrder })); - fetchTeamsV2({ sortBy, sortOrder }); - }; - - const teamColumns: ColumnsType = useMemo( - () => [ - { - title: "Team ID", - dataIndex: "team_id", - key: "team_id", - width: 170, - ellipsis: true, - render: (id: string) => ( - setSelectedTeamId(teamId)} dataTestId="team-id-cell" /> - ), - }, - { - title: "Team Alias", - dataIndex: "team_alias", - key: "team_alias", - ellipsis: true, - sorter: true, - render: (alias: string | undefined) => ( - - {alias || ( - - — - - )} - - ), - }, - { - title: "Organization", - key: "organization", - width: 160, - ellipsis: true, - render: (_: unknown, record: Team) => { - const orgAlias = getOrganizationAlias(record.organization_id, organizations); - return record.organization_id ? ( - - {orgAlias} - - ) : ( - — - ); - }, - }, - { - title: "Resources", - key: "resources", - width: 240, - render: (_: unknown, record: Team) => { - const memberCount = perTeamInfo?.[record.team_id]?.team_info?.members_with_roles?.length ?? 0; - const modelCount = record.models?.length ?? 0; - const keyCount = perTeamInfo?.[record.team_id]?.keys_count ?? 0; - return ( - - - - - - {memberCount} - - - - - - - - {modelCount} - - - - - - - - {keyCount} - - - - - ); - }, - }, - { - title: "Spend / Budget", - key: "spend", - width: 200, - sorter: true, - render: (_: unknown, record: Team) => { - const spendVal = record.spend ?? 0; - const budgetVal = record.max_budget; - const spendStr = `$${spendVal.toLocaleString(undefined, { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`; - const budgetStr = - budgetVal != null - ? `$${budgetVal.toLocaleString(undefined, { minimumFractionDigits: 2, maximumFractionDigits: 2 })}` - : "Unlimited"; - const percent = budgetVal != null && budgetVal > 0 ? Math.min((spendVal / budgetVal) * 100, 100) : null; - return ( - - - {spendStr} - - {" / "} - {budgetStr} - - - {percent != null && ( - = 90 ? "#ff4d4f" : percent >= 70 ? "#faad14" : "#1677ff"} - style={{ marginBottom: 0 }} - /> - )} - - ); - }, - }, - { - title: "Created", - dataIndex: "created_at", - key: "created_at", - width: 130, - ellipsis: true, - sorter: true, - render: (date: string | undefined) => , - }, - { - title: "Actions", - key: "actions", - width: 120, - align: "right" as const, - render: (_: unknown, record: Team) => ( - - { - navigator.clipboard - .writeText(record.team_id) - .then(() => message.success("Team ID copied")) - .catch(() => message.error("Failed to copy")); - }} - /> - {userRole === "Admin" && ( - <> - { - setSelectedTeamId(record.team_id); - setEditTeam(true); - }} - /> - handleDelete(record)} - /> - - )} - - ), - }, - ], - [userRole, perTeamInfo, organizations], - ); - - const displayTeams = useMemo(() => teams ?? [], [teams]); - - const renderTeamsContent = () => { - if (isLoading) { - return ( - - - - ); - } - - if (fetchError) { - return ( - - - Failed to load teams - - - {fetchError} - - - - ); - } - - return ( - - columns={teamColumns} - dataSource={displayTeams} - rowKey="team_id" - pagination={false} - onChange={handleTableSort} - locale={{ - emptyText: ( -
- -
- No teams yet -
-
- - Create your first team to organize members and manage access to models. - -
- {canCreateOrManageTeams(userRole, userID, organizations) && ( - - )} -
- ), - }} - scroll={{ x: 1000 }} - size="middle" - /> - ); - }; - const tabItems = [ { key: "your-teams", label: "Your Teams", children: ( <> - - - - } - suffix={isSearching ? : null} - placeholder="Search teams by name or ID..." - onChange={(e) => handleSearchChange(e.target.value)} - allowClear - style={{ maxWidth: 400 }} - /> - handleFilterChange("organization_id", value || "")} - loading={isLoading} - /> - - { - setCurrentPage(page); - setPageSize(size); - fetchTeamsV2({ page, size }); - }} - size="small" - showTotal={(total) => `${total} teams`} - showSizeChanger - pageSizeOptions={["10", "20", "50"]} - /> - - - {renderTeamsContent()} - + { + setSelectedTeam(team); + setSelectedTeamId(team.team_id); + setEditTeam(false); + }} + onEditTeam={(team) => { + setSelectedTeam(team); + setSelectedTeamId(team.team_id); + setEditTeam(true); + }} + onDeleteTeam={handleDelete} + /> = ({ accessToken, userID, userRole, premiumUser {selectedTeamId ? ( { - setTeams((teams) => { - if (teams == null) { - return teams; - } - return teams.map((team) => { - if (data.team_id === team.team_id) { - return updateExistingKeys(team, data); - } - return team; - }); - }); - fetchTeamsV2(); + onUpdate={() => { + refreshTeams(); }} onClose={() => { + setSelectedTeam(null); setSelectedTeamId(null); setEditTeam(false); }} accessToken={accessToken} - is_team_admin={is_team_admin(teams?.find((team) => team.team_id === selectedTeamId))} + is_team_admin={is_team_admin(selectedTeam)} is_proxy_admin={userRole == "Admin"} userModels={userModels} editTeam={editTeam} @@ -1018,25 +559,21 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser /> ) : ( <> - - - - <TeamOutlined style={{ marginRight: 8 }} /> - Teams - - Manage teams, members, and their access to models and budgets - - {canCreateOrManageTeams(userRole, userID, organizations) && ( - - )} - +
+ } + title="Teams" + subtitle="Manage teams, members, and their access to models and budgets" + actions={ + canCreateOrManageTeams(userRole, userID, organizations) ? ( + setIsTeamModalVisible(true)} data-testid="create-team-button"> + + Create Team + + ) : undefined + } + /> +
diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx new file mode 100644 index 00000000000..9469be12128 --- /dev/null +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx @@ -0,0 +1,330 @@ +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, MockedFunction, vi } from "vitest"; + +import { renderWithProviders } from "../../../tests/test-utils"; +import { Team } from "../key_team_helpers/key_list"; +import { TeamsResponse, useTeamsTable } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { TeamsTable } from "./TeamsTable"; + +// Resolve debounced values synchronously so an applied filter lands in the useTeamsTable query within the test tick. +vi.mock("@tanstack/react-pacer/debouncer", async () => { + const React = await vi.importActual("react"); + return { + useDebouncedValue: (value: unknown) => [value, { cancel: vi.fn(), flush: vi.fn() }], + useDebouncedState: (initial: unknown) => { + const [value, setValue] = React.useState(initial); + return [value, setValue, { cancel: vi.fn(), flush: vi.fn() }]; + }, + }; +}); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(() => ({ + accessToken: "test-token", + userId: "test-user", + userRole: "Admin", + premiumUser: true, + token: "test-token", + })), +})); + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useTeamsTable: vi.fn(), + teamsTableKeys: { all: ["teamsTable"] }, +})); + +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganizations: vi.fn().mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "Test Organization" }], + }), +})); + +const mockTeam: Team = { + team_id: "team-1", + team_alias: "Acme Team", + models: ["gpt-4", "gpt-3.5-turbo", "claude-3", "claude-3-5-sonnet"], + max_budget: 100, + budget_duration: "1mo", + tpm_limit: 5000, + rpm_limit: 500, + organization_id: "org-1", + created_at: "2024-10-01T10:00:00Z", + updated_at: "2024-11-01T10:00:00Z", + keys: [], + keys_count: 3, + members_with_roles: [ + { user_id: "u1", user_email: "a@x.com", role: "admin" }, + { user_id: "u2", user_email: "b@x.com", role: "user" }, + ] as unknown as Team["members_with_roles"], + spend: 42.5, +}; + +const mockUseTeamsTable = useTeamsTable as MockedFunction; + +const teamsResult = (teams: Team[], data: Partial = {}, extra: Record = {}) => + ({ + data: { + teams, + total: teams.length, + page: 1, + page_size: 50, + total_pages: 1, + ...data, + } as TeamsResponse, + isPending: false, + isFetching: false, + isError: false, + refetch: vi.fn(), + ...extra, + }) as any; + +const noop = () => {}; + +const renderTable = (props: Partial> = {}) => + renderWithProviders( + , + ); + +const openFilters = () => fireEvent.click(screen.getByRole("button", { name: "Filters" })); +const lastOptions = () => mockUseTeamsTable.mock.calls[mockUseTeamsTable.mock.calls.length - 1][2] ?? {}; + +beforeEach(() => { + vi.clearAllMocks(); + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam])); +}); + +it("renders a team row with alias, organization, and spend/budget", async () => { + renderTable(); + + await waitFor(() => { + expect(screen.getByText("Acme Team")).toBeInTheDocument(); + expect(screen.getByText("Test Organization")).toBeInTheDocument(); + expect(screen.getByText("$42.5000")).toBeInTheDocument(); + expect(screen.getByText("of $100")).toBeInTheDocument(); + }); +}); + +it("renders the Resources cell with member, model, and key counts", () => { + renderTable(); + + expect(screen.getByTitle("2 members")).toBeInTheDocument(); + expect(screen.getByTitle("4 models")).toBeInTheDocument(); + expect(screen.getByTitle("3 keys")).toBeInTheDocument(); +}); + +it("shows 'No teams found' when the list is empty", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([])); + renderTable(); + expect(screen.getByText("No teams found")).toBeInTheDocument(); +}); + +it("shows a loading state on initial load and hides the data", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([], {}, { data: null, isPending: true, isFetching: true })); + renderTable(); + + expect(screen.getByText("Loading teams...")).toBeInTheDocument(); + expect(screen.queryByText("Acme Team")).not.toBeInTheDocument(); +}); + +describe("sort contract – only backend-sortable columns are sortable", () => { + it("requests the default created_at descending sort on first render", () => { + renderTable(); + expect(lastOptions()).toMatchObject({ sortBy: "created_at", sortOrder: "desc" }); + }); + + it("sorts by the backend team_alias field (not the label) when the Team header is clicked", async () => { + renderTable(); + fireEvent.click(screen.getByText("Team").closest("button") as HTMLElement); + + await waitFor(() => { + expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "team_alias" })); + }); + }); + + it("does not make Spend / Budget sortable (the backend rejects sort_by=spend)", () => { + renderTable(); + expect(screen.getByText("Spend / Budget").closest("button")).toBeNull(); + // Team and Created are the only sortable headers. + expect(screen.getByText("Team").closest("button")).not.toBeNull(); + expect(screen.getByText("Created").closest("button")).not.toBeNull(); + }); +}); + +describe("server-side filtering maps controls to the right query params", () => { + it("sends no filter params when nothing is applied", () => { + renderTable(); + expect(lastOptions()).toMatchObject({ organizationID: undefined, team_alias: undefined, teamID: undefined }); + }); + + it("threads an applied Team alias filter into the query", async () => { + renderTable(); + openFilters(); + + fireEvent.change(await screen.findByPlaceholderText(/Enter team alias/), { target: { value: "acme" } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => { + expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ team_alias: "acme" })); + }); + }); + + it("threads an applied Team ID filter into the query", async () => { + renderTable(); + openFilters(); + + fireEvent.change(await screen.findByPlaceholderText(/Enter team ID/), { target: { value: "team-xyz" } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => { + expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ teamID: "team-xyz" })); + }); + }); + + it("threads the toolbar search into the search param", async () => { + renderTable(); + fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "platform" } }); + + await waitFor(() => { + expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: "platform" })); + }); + }); +}); + +describe("non-admin scoping", () => { + it("scopes the list to the current user when the role is not an admin role", () => { + renderTable({ userRole: "Internal User", userID: "user-42" }); + expect(lastOptions()).toMatchObject({ userID: "user-42" }); + }); + + it("does not scope by user for the Admin role", () => { + renderTable({ userRole: "Admin", userID: "admin-1" }); + expect(lastOptions()).toMatchObject({ userID: undefined }); + }); +}); + +describe("row actions", () => { + it("opens the team detail when the team cell is clicked", () => { + const onSelectTeam = vi.fn(); + renderTable({ onSelectTeam }); + + fireEvent.click(screen.getByText("Acme Team")); + expect(onSelectTeam).toHaveBeenCalledWith(expect.objectContaining({ team_id: "team-1" })); + }); + + it("offers Edit and Delete to an Admin and wires them to the callbacks", async () => { + const onEditTeam = vi.fn(); + const onDeleteTeam = vi.fn(); + const user = userEvent.setup(); + renderTable({ userRole: "Admin", onEditTeam, onDeleteTeam }); + + await user.click(screen.getByTestId("team-actions-team-1")); + + await user.click(await screen.findByText("Edit team")); + expect(onEditTeam).toHaveBeenCalledWith(expect.objectContaining({ team_id: "team-1" })); + + await user.click(screen.getByTestId("team-actions-team-1")); + await user.click(await screen.findByText("Delete team")); + expect(onDeleteTeam).toHaveBeenCalledWith(expect.objectContaining({ team_id: "team-1" })); + }); + + it("hides Edit and Delete from a non-admin, leaving only Copy team ID", async () => { + const user = userEvent.setup(); + renderTable({ userRole: "Internal User" }); + + await user.click(screen.getByTestId("team-actions-team-1")); + + expect(await screen.findByText("Copy team ID")).toBeInTheDocument(); + expect(screen.queryByText("Edit team")).not.toBeInTheDocument(); + expect(screen.queryByText("Delete team")).not.toBeInTheDocument(); + }); +}); + +describe("pagination total comes from the query response", () => { + it("shows the total count and page count from the response", async () => { + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam], { total: 137, total_pages: 3 })); + renderTable(); + + await waitFor(() => { + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-50 of 137"); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 3"); + }); + }); +}); + +describe("refresh control", () => { + it("calls refetch when clicked", () => { + const refetch = vi.fn(); + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam], {}, { refetch })); + renderTable(); + + fireEvent.click(screen.getByTestId("datatable-refresh")); + expect(refetch).toHaveBeenCalledTimes(1); + }); + + it("keeps rows visible but disables refresh while a background fetch is in flight", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam], {}, { isFetching: true })); + renderTable(); + + expect(screen.getByTestId("datatable-refresh")).toBeDisabled(); + expect(screen.getByText("Acme Team")).toBeInTheDocument(); + }); +}); + +describe("column rendering details", () => { + it("shows the organization alias when the id resolves, and the raw id when it does not", async () => { + mockUseTeamsTable.mockReturnValue( + teamsResult([ + { ...mockTeam, team_id: "a", organization_id: "org-1" }, + { ...mockTeam, team_id: "b", team_alias: "Orphan Team", organization_id: "org-unknown" }, + ]), + ); + renderTable(); + + await waitFor(() => { + expect(screen.getByText("Test Organization")).toBeInTheDocument(); + expect(screen.getByText("org-unknown")).toBeInTheDocument(); + }); + }); + + it("renders an em dash for a team with no organization", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([{ ...mockTeam, organization_id: null as unknown as string }])); + renderTable(); + expect(screen.getByText("—")).toBeInTheDocument(); + }); + + it("falls back to keys.length when keys_count is absent", () => { + mockUseTeamsTable.mockReturnValue( + teamsResult([ + { + ...mockTeam, + keys_count: undefined, + keys: [{ token: "t1" }, { token: "t2" }] as unknown as Team["keys"], + }, + ]), + ); + renderTable(); + expect(screen.getByTitle("2 keys")).toBeInTheDocument(); + }); +}); + +describe("hidden-by-default columns", () => { + it("hides Members, Models, Rate Limits, and Updated until toggled on", async () => { + const user = userEvent.setup(); + renderTable(); + + expect(screen.queryByText("Rate Limits")).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Columns" })); + const menu = await screen.findByRole("menu"); + expect(within(menu).getByText("Rate Limits")).toBeInTheDocument(); + expect(within(menu).getByText("Updated")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx new file mode 100644 index 00000000000..3b75db52b16 --- /dev/null +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx @@ -0,0 +1,201 @@ +"use client"; + +import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { useTeamsTable } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { + DataTable, + DataTableFilterDrawer, + DataTableFilterField, + DataTableToolbar, +} from "@/components/shared/DataTable"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { Input } from "@/components/ui/input"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import React, { useCallback, useMemo, useState } from "react"; + +import { Team } from "../key_team_helpers/key_list"; +import { getTeamTableColumns, TEAM_TABLE_HIDDEN_COLUMNS } from "./teamTableColumns"; + +interface TeamsTableProps { + userRole: string | null; + userID: string | null; + onSelectTeam: (team: Team) => void; + onEditTeam: (team: Team) => void; + onDeleteTeam: (team: Team) => void; +} + +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +const toSortOrder = (sorting: SortingState): "asc" | "desc" | undefined => { + const active = sorting[0]; + if (!active) return undefined; + return active.desc ? "desc" : "asc"; +}; + +const FILTER_LABELS: Record = { + org_id: "Organization", + alias: "Team alias", + team_id: "Team ID", +}; + +export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDeleteTeam }: TeamsTableProps) { + const { data: fetchedOrganizations } = useOrganizations(); + const organizations = useMemo(() => fetchedOrganizations ?? [], [fetchedOrganizations]); + + const [sorting, setSorting] = useState(DEFAULT_SORTING); + const [tablePagination, setTablePagination] = useState({ pageIndex: 0, pageSize: 50 }); + const [columnFilters, setColumnFilters] = useState([]); + const [filtersOpen, setFiltersOpen] = useState(false); + const [searchInput, setSearchInput] = useState(""); + const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); + + const getFilterValue = useCallback( + (columnId: string): string | undefined => { + const entry = columnFilters.find((filter) => filter.id === columnId); + return typeof entry?.value === "string" && entry.value.trim() ? entry.value.trim() : undefined; + }, + [columnFilters], + ); + + const isAdminView = userRole === "Admin" || userRole === "Admin Viewer"; + + const teamListOptions = { + organizationID: getFilterValue("org_id"), + team_alias: getFilterValue("alias"), + teamID: getFilterValue("team_id"), + search: searchQuery.trim() || undefined, + userID: isAdminView ? undefined : userID ?? undefined, + sortBy: sorting[0]?.id, + sortOrder: toSortOrder(sorting), + }; + + const { + data: teamsResponse, + isPending: isLoading, + isFetching, + refetch, + } = useTeamsTable(tablePagination.pageIndex + 1, tablePagination.pageSize, teamListOptions); + + const teamList = useMemo(() => teamsResponse?.teams ?? [], [teamsResponse]); + const rowCount = teamsResponse?.total ?? 0; + + const handleSearchChange = useCallback((value: string) => { + setSearchInput(value); + setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); + + const handleSortingChange = useCallback>((updaterOrValue) => { + setSorting(updaterOrValue); + setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); + + const handleColumnFiltersChange = useCallback>((updaterOrValue) => { + setColumnFilters(updaterOrValue); + setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); + + const columns = useMemo(() => { + const columnDeps = { organizations, userRole, onSelectTeam, onEditTeam, onDeleteTeam }; + return getTeamTableColumns(columnDeps); + }, [organizations, userRole, onSelectTeam, onEditTeam, onDeleteTeam]); + + const orgOptions = useMemo( + () => + organizations + .filter((org) => org.organization_id) + .map((org) => { + const id = org.organization_id as string; + return { label: org.organization_alias || id, value: id, sublabel: org.organization_alias ? id : undefined }; + }), + [organizations], + ); + + const formatFilterValue = useCallback( + (columnId: string, value: unknown): string => { + const raw = String(value); + if (columnId === "org_id") { + return organizations.find((org) => org.organization_id === raw)?.organization_alias || raw; + } + return raw; + }, + [organizations], + ); + + return ( + row.team_id} + defaultColumnVisibility={TEAM_TABLE_HIDDEN_COLUMNS} + sortingMode="server" + sorting={sorting} + onSortingChange={handleSortingChange} + paginationMode="server" + pagination={tablePagination} + onPaginationChange={setTablePagination} + rowCount={rowCount} + filterMode="server" + columnFilters={columnFilters} + onColumnFiltersChange={handleColumnFiltersChange} + enableColumnResizing + columnResizeMode="onChange" + isLoading={isLoading} + loadingMessage="Loading teams..." + noDataMessage="No teams found" + maxBodyHeight="calc(75vh - 210px)" + size="compact" + toolbar={(table) => ( + <> + refetch?.()} + isRefreshing={isFetching} + onOpenFilters={() => setFiltersOpen(true)} + filterLabels={FILTER_LABELS} + formatFilterValue={formatFilterValue} + /> + + {({ get, set }) => ( + <> + + set("org_id", value)} + placeholder="Select an organization…" + emptyText="No organizations found" + /> + + + set("alias", event.target.value)} + placeholder="Enter team alias…" + /> + + + set("team_id", event.target.value)} + placeholder="Enter team ID…" + /> + + + )} + + + )} + /> + ); +} diff --git a/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx new file mode 100644 index 00000000000..ecf7387ee83 --- /dev/null +++ b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx @@ -0,0 +1,287 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { Copy, KeyRound, Layers, MoreHorizontal, Pencil, Trash2, Users } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell, SpendBudgetCell } from "@/components/shared/table_cells"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { Skeleton } from "@/components/ui/skeleton"; +import { cn } from "@/lib/cva.config"; +import { copyToClipboard, formatNumberWithCommas } from "@/utils/dataUtils"; + +import { Team } from "../key_team_helpers/key_list"; +import { Organization } from "../networking"; + +interface ResourceTone { + icon: typeof Users; + className: string; +} + +const RESOURCE_TONES: Record<"members" | "models" | "keys", ResourceTone> = { + members: { icon: Users, className: "bg-violet-50 text-violet-700 ring-violet-600/20" }, + models: { icon: Layers, className: "bg-sky-50 text-sky-700 ring-sky-600/20" }, + keys: { icon: KeyRound, className: "bg-emerald-50 text-emerald-700 ring-emerald-600/20" }, +}; + +const teamMemberCount = (team: Team): number => team.members_count ?? team.members_with_roles?.length ?? 0; +const teamModelCount = (team: Team): number => team.models?.length ?? 0; +const teamKeyCount = (team: Team): number => team.keys_count ?? team.keys?.length ?? 0; + +function ResourcesCell({ team }: { team: Team }) { + const items = [ + { key: "members" as const, label: "members", count: teamMemberCount(team) }, + { key: "models" as const, label: "models", count: teamModelCount(team) }, + { key: "keys" as const, label: "keys", count: teamKeyCount(team) }, + ]; + + return ( +
+ {items.map((item) => { + const tone = RESOURCE_TONES[item.key]; + const Icon = tone.icon; + return ( + + + {item.count} + + ); + })} +
+ ); +} + +function RateLimitLine({ label, value }: { label: string; value: number | null }) { + return ( +
+ {label} + {value != null ? formatNumberWithCommas(value) : "Unlimited"} +
+ ); +} + +interface TeamRowActionsProps { + team: Team; + canManage: boolean; + onEditTeam: (team: Team) => void; + onDeleteTeam: (team: Team) => void; +} + +function TeamRowActions({ team, canManage, onEditTeam, onDeleteTeam }: TeamRowActionsProps) { + const handleCopy = () => { + void copyToClipboard(team.team_id, "Team ID copied"); + }; + + return ( + + + + + + {canManage && ( + onEditTeam(team)} data-testid="team-action-edit"> + + Edit team + + )} + + + Copy team ID + + {canManage && ( + <> + + onDeleteTeam(team)} data-testid="team-action-delete"> + + Delete team + + + )} + + + ); +} + +interface TeamTableColumnsDeps { + organizations: Organization[]; + userRole: string | null; + onSelectTeam: (team: Team) => void; + onEditTeam: (team: Team) => void; + onDeleteTeam: (team: Team) => void; +} + +export const getTeamTableColumns = ({ + organizations, + userRole, + onSelectTeam, + onEditTeam, + onDeleteTeam, +}: TeamTableColumnsDeps): ColumnDef[] => { + const canManage = userRole === "Admin"; + + return [ + { + id: "team_alias", + accessorKey: "team_alias", + meta: { + title: "Team", + renderSkeleton: () => ( +
+ + +
+ ), + }, + header: ({ column }) => , + size: 260, + enableSorting: true, + cell: ({ row }) => { + const team = row.original; + const hasAlias = Boolean(team.team_alias); + return ( + onSelectTeam(team)} + /> + ); + }, + }, + { + id: "organization_alias", + accessorKey: "organization_id", + meta: { title: "Organization" }, + header: "Organization", + size: 160, + enableSorting: false, + cell: (info) => { + const orgId = info.getValue() as string | null; + if (!orgId) return —; + const org = organizations.find((o) => o.organization_id === orgId); + const displayValue = org?.organization_alias || orgId; + const width = info.cell.column.getSize(); + return ( + + {displayValue} + + ); + }, + }, + { + id: "resources", + meta: { + title: "Resources", + renderSkeleton: () => ( +
+ + + +
+ ), + }, + header: "Resources", + size: 210, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "spend", + accessorKey: "spend", + meta: { title: "Spend / Budget", skeleton: "meter" }, + header: "Spend / Budget", + size: 200, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "created_at", + accessorKey: "created_at", + meta: { title: "Created" }, + header: ({ column }) => , + size: 130, + enableSorting: true, + cell: (info) => , + }, + { + id: "members", + meta: { title: "Members" }, + header: "Members", + size: 110, + enableSorting: false, + cell: ({ row }) => {teamMemberCount(row.original)}, + }, + { + id: "models", + meta: { title: "Models" }, + header: "Models", + size: 100, + enableSorting: false, + cell: ({ row }) => {teamModelCount(row.original)}, + }, + { + id: "rate_limits", + meta: { title: "Rate Limits", skeleton: "twoLine" }, + header: "Rate Limits", + size: 140, + enableSorting: false, + cell: ({ row }) => ( +
+ + +
+ ), + }, + { + id: "updated_at", + accessorKey: "updated_at", + meta: { title: "Updated" }, + header: "Updated", + size: 130, + enableSorting: false, + cell: (info) => , + }, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 60, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, + ]; +}; + +export const TEAM_TABLE_HIDDEN_COLUMNS: Record = { + members: false, + models: false, + rate_limits: false, + updated_at: false, +}; diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 02f5d588149..dd08fe00656 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -16,6 +16,8 @@ vi.mock("@tanstack/react-pacer/debouncer", async () => { const [value, setValue] = React.useState(initial); return [value, setValue, { cancel: vi.fn(), flush: vi.fn() }]; }, + useDebouncedCallback: (fn: (...args: unknown[]) => void) => fn, + useDebouncer: (fn: (...args: unknown[]) => void) => ({ maybeExecute: fn, cancel: vi.fn(), flush: vi.fn() }), }; }); @@ -63,7 +65,7 @@ const mockKey: KeyResponse = { key_alias: "Test Key Alias", spend: 5.5, max_budget: 100, - expires: "2024-12-31T23:59:59Z", + expires: "2999-12-31T23:59:59Z", models: ["gpt-3.5-turbo", "gpt-4"], aliases: {}, config: {}, @@ -154,6 +156,8 @@ const keysResult = (keys: KeyResponse[], data: Partial = {}, extra ...extra, }) as any; +const openFilters = () => fireEvent.click(screen.getByRole("button", { name: "Filters" })); + beforeEach(() => { vi.clearAllMocks(); @@ -170,6 +174,12 @@ it("should render VirtualKeysTable component", () => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); +it("renders the page header with the create-key action slot", () => { + renderWithProviders(Create New Key} />); + expect(screen.getByRole("heading", { name: "Virtual Keys" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Create New Key" })).toBeInTheDocument(); +}); + it("should display key information correctly", async () => { renderWithProviders(); @@ -177,6 +187,7 @@ it("should display key information correctly", async () => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); expect(screen.getByText("Test Team")).toBeInTheDocument(); expect(screen.getByText("$5.5000")).toBeInTheDocument(); + expect(screen.getByText("of $100")).toBeInTheDocument(); }); }); @@ -188,14 +199,49 @@ it("should display user email correctly", async () => { }); }); -it("should show loading message only on initial load (isPending)", () => { +it("shows the user alias over the email in the visible cell when both exist", async () => { + mockUseKeys.mockReturnValue( + keysResult([{ ...mockKey, user: { user_id: "user-1", user_email: "user@example.com", user_alias: "The User" } }]), + ); + + renderWithProviders(); + + const row = (await screen.findByText("Test Key Alias")).closest("tr") as HTMLElement; + expect(within(row).getByText("The User")).toBeInTheDocument(); + expect(within(row).queryByText("user@example.com")).not.toBeInTheDocument(); +}); + +it("shows created_by_user alias over email in the Created By column when it is enabled", async () => { + mockUseKeys.mockReturnValue( + keysResult([ + { + ...mockKey, + created_by: "some-uuid", + created_by_user: { user_id: "some-uuid", user_email: "creator@example.com", user_alias: "The Creator" }, + }, + ]), + ); + const user = userEvent.setup(); + renderWithProviders(); + + // Created By is hidden by default; turn it on via the Columns menu. + await user.click(screen.getByRole("button", { name: "Columns" })); + await user.click(await screen.findByText("Created By")); + await user.keyboard("{Escape}"); + + const row = (await screen.findByText("Test Key Alias")).closest("tr") as HTMLElement; + expect(within(row).getByText("The Creator")).toBeInTheDocument(); + expect(within(row).queryByText("creator@example.com")).not.toBeInTheDocument(); +}); + +it("should show a loading state on the initial load and hide the data", () => { mockUseKeys.mockReturnValue(keysResult([], {}, { data: null, isPending: true, isFetching: true })); renderWithProviders(); - expect(screen.getByText("🚅 Loading keys...")).toBeInTheDocument(); + expect(screen.getByText("Loading keys...")).toBeInTheDocument(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); - expect(screen.queryByText("Test Team")).not.toBeInTheDocument(); }); it("should show 'No keys found' message when the key list is empty", () => { @@ -206,61 +252,98 @@ it("should show 'No keys found' message when the key list is empty", () => { expect(screen.getByText("No keys found")).toBeInTheDocument(); }); -it("should handle models with more than 3 entries to trigger expansion UI", () => { +it("collapses models beyond the visible limit into a '+N more' badge", () => { mockUseKeys.mockReturnValue( keysResult([{ ...mockKey, models: ["gpt-3.5-turbo", "gpt-4", "gpt-4-turbo", "claude-3", "claude-3-5-sonnet"] }]), ); renderWithProviders(); - expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); + expect(screen.getByText("+2 more")).toBeInTheDocument(); }); -it("should render table headers correctly", () => { +it("should render the redesigned table headers", () => { renderWithProviders(); - expect(screen.getByText("Key ID")).toBeInTheDocument(); - expect(screen.getByText("Key Alias")).toBeInTheDocument(); + expect(screen.getByText("Key")).toBeInTheDocument(); expect(screen.getByText("Team")).toBeInTheDocument(); expect(screen.getByText("Models")).toBeInTheDocument(); - expect(screen.getByText("Spend (USD)")).toBeInTheDocument(); + expect(screen.getByText("Spend", { selector: "[data-sort-field='spend']" })).toBeInTheDocument(); + expect(screen.getByText("Budget", { selector: "[data-sort-field='max_budget']" })).toBeInTheDocument(); }); -it("should handle column resizing hover events", () => { +it("sorts by the backend key_alias field (not the column label) when the Key header is clicked", async () => { renderWithProviders(); - const headerCell = document.querySelector("[data-header-id]") as HTMLElement; - expect(headerCell).toBeInTheDocument(); + const keyHeader = screen.getByText("Key").closest("button") as HTMLElement; + fireEvent.click(keyHeader); - const resizer = headerCell?.querySelector(".resizer") as HTMLElement; - expect(resizer).toBeInTheDocument(); - expect(resizer.style.opacity).toBe("0"); - - fireEvent.mouseEnter(headerCell); - expect(resizer.style.opacity).toBe("0.5"); - - fireEvent.mouseLeave(headerCell); - expect(resizer.style.opacity).toBe("0"); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "key_alias" })); + }); }); -it("should open KeyInfoView when clicking on a key ID button", async () => { +it("sorts by the backend max_budget field when 'Budget descending' is chosen from the Spend / Budget menu", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByTestId("sort-trigger-spend")); + await user.click(await screen.findByText("Budget descending")); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ sortBy: "max_budget", sortOrder: "desc" }), + ); + }); +}); + +it("emphasizes the active field in the Spend / Budget header so the sorted column reads without opening the menu", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByTestId("sort-trigger-spend")); + await user.click(await screen.findByText("Budget descending")); + + await waitFor(() => { + expect(screen.getByText("Budget", { selector: "[data-sort-field='max_budget']" }).className).toContain( + "font-semibold", + ); + }); + expect(screen.getByText("Spend", { selector: "[data-sort-field='spend']" }).className).toContain( + "text-muted-foreground", + ); +}); + +it("sorts by spend ascending when 'Spend ascending' is chosen from the Spend / Budget menu", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByTestId("sort-trigger-spend")); + await user.click(await screen.findByText("Spend ascending")); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "spend", sortOrder: "asc" })); + }); +}); + +it("should open KeyInfoView when clicking the key cell", async () => { renderWithProviders(); await waitFor(() => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); - expect(screen.getByText(/Showing.*results/)).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toBeInTheDocument(); - const keyIdButton = screen.getByText("sk-1234567890abcdef"); - fireEvent.click(keyIdButton); + fireEvent.click(screen.getByText("Test Key Alias")); await waitFor(() => { expect(screen.getByText("Back to Keys")).toBeInTheDocument(); - expect(screen.getByText("Created At")).toBeInTheDocument(); }); - expect(screen.queryByText(/Showing.*results/)).not.toBeInTheDocument(); + expect(screen.queryByTestId("pagination-range")).not.toBeInTheDocument(); }); it("should display 'Default Proxy Admin' for user_id when value is 'default_user_id'", async () => { @@ -282,44 +365,6 @@ it("should display 'Default Proxy Admin' for user_id when value is 'default_user }); }); -it("should display created_by_user email in 'Created By' column when available", async () => { - mockUseKeys.mockReturnValue( - keysResult([ - { - ...mockKey, - created_by: "some-uuid-1234", - created_by_user: { user_id: "some-uuid-1234", user_email: "creator@example.com", user_alias: null }, - }, - ]), - ); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("creator@example.com")).toBeInTheDocument(); - }); -}); - -it("should display created_by_user alias over email when both are available", async () => { - mockUseKeys.mockReturnValue( - keysResult([ - { - ...mockKey, - created_by: "some-uuid-1234", - created_by_user: { user_id: "some-uuid-1234", user_email: "creator@example.com", user_alias: "The Creator" }, - }, - ]), - ); - - renderWithProviders(); - - // Scope to the key's row so we assert the visible cell value: the hover popover that - // also holds the email is portaled out of the row, not the displayed "Created By" text. - const row = (await screen.findByText("Test Key Alias")).closest("tr") as HTMLElement; - expect(within(row).getByText("The Creator")).toBeInTheDocument(); - expect(within(row).queryByText("creator@example.com")).not.toBeInTheDocument(); -}); - it("should render table without crashing when models is null", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, models: null as unknown as string[] }])); @@ -327,6 +372,7 @@ it("should render table without crashing when models is null", async () => { await waitFor(() => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); }); }); @@ -341,13 +387,14 @@ it("should display 'Unknown' for last_active when value is null", async () => { }); describe("server-side filtering – the LIT-4080 regression guard", () => { - it("threads an active User ID filter into the useKeys query so any refetch keeps it", async () => { + it("threads an applied User ID filter into the useKeys query so any refetch keeps it", async () => { renderWithProviders(); - fireEvent.click(screen.getByRole("button", { name: "Filters" })); + openFilters(); - const userIdInput = await screen.findByPlaceholderText("Enter User ID..."); + const userIdInput = await screen.findByPlaceholderText(/Enter User ID/); fireEvent.change(userIdInput, { target: { value: "user-42" } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ userID: "user-42" })); @@ -361,18 +408,19 @@ describe("server-side filtering – the LIT-4080 regression guard", () => { expect(lastCall[2] ?? {}).toMatchObject({ userID: undefined, teamID: undefined, keyHash: undefined }); }); - it("drops the filter from the useKeys query when Reset Filters is clicked", async () => { + it("drops the filter from the useKeys query when it is cleared", async () => { renderWithProviders(); - fireEvent.click(screen.getByRole("button", { name: "Filters" })); - const userIdInput = await screen.findByPlaceholderText("Enter User ID..."); + openFilters(); + const userIdInput = await screen.findByPlaceholderText(/Enter User ID/); fireEvent.change(userIdInput, { target: { value: "user-42" } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ userID: "user-42" })); }); - fireEvent.click(screen.getByRole("button", { name: "Reset Filters" })); + fireEvent.click(screen.getByTestId("datatable-clear-filters")); await waitFor(() => { const lastCall = mockUseKeys.mock.calls[mockUseKeys.mock.calls.length - 1]; @@ -388,8 +436,8 @@ describe("pagination display – total count comes from useKeys", () => { renderWithProviders(); await waitFor(() => { - expect(screen.getByText("Showing 1 - 50 of 509 results")).toBeInTheDocument(); - expect(screen.getByText("Page 1 of 11")).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-50 of 509"); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 11"); }); }); @@ -399,57 +447,44 @@ describe("pagination display – total count comes from useKeys", () => { renderWithProviders(); await waitFor(() => { - expect(screen.getByText("Showing 1 - 1 of 1 results")).toBeInTheDocument(); - expect(screen.getByText("Page 1 of 1")).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-1 of 1"); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 1"); }); }); }); -describe("refetch button", () => { - it("should show Fetch button in normal state", () => { +describe("refresh button", () => { + it("renders an enabled refresh control in the normal state", () => { renderWithProviders(); - const fetchButton = screen.getByTitle("Fetch data"); - expect(fetchButton).toBeInTheDocument(); - expect(fetchButton).not.toBeDisabled(); - expect(screen.getByText("Fetch")).toBeInTheDocument(); + const refresh = screen.getByTestId("datatable-refresh"); + expect(refresh).toBeInTheDocument(); + expect(refresh).not.toBeDisabled(); }); - it("should show Fetching state and keep table data visible during refetch", () => { + it("disables the refresh control while a fetch is in flight but keeps data visible", () => { mockUseKeys.mockReturnValue(keysResult([mockKey], {}, { isFetching: true })); renderWithProviders(); - expect(screen.getByText("Fetching")).toBeInTheDocument(); - expect(screen.getByTitle("Fetch data")).toBeDisabled(); + expect(screen.getByTestId("datatable-refresh")).toBeDisabled(); expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); - expect(screen.queryByText("🚅 Loading keys...")).not.toBeInTheDocument(); }); - it("should call refetch when Fetch button is clicked", () => { + it("calls refetch when the refresh control is clicked", () => { const mockRefetch = vi.fn(); mockUseKeys.mockReturnValue(keysResult([mockKey], {}, { refetch: mockRefetch })); renderWithProviders(); - fireEvent.click(screen.getByTitle("Fetch data")); + fireEvent.click(screen.getByTestId("datatable-refresh")); expect(mockRefetch).toHaveBeenCalledTimes(1); }); - - it("should show Fetch button enabled on error so user can retry", () => { - mockUseKeys.mockReturnValue(keysResult([], {}, { data: null, isError: true })); - - renderWithProviders(); - - const fetchButton = screen.getByTitle("Fetch data"); - expect(fetchButton).not.toBeDisabled(); - expect(screen.getByText("Fetch")).toBeInTheDocument(); - }); }); -describe("Status column reflects key.blocked / scim_blocked metadata", () => { - it("should render Active for a non-blocked key", async () => { +describe("Status column reflects blocked / expiry / scim metadata", () => { + it("renders Active for a non-blocked, unexpired key", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, blocked: false, metadata: {} }])); renderWithProviders(); @@ -459,7 +494,19 @@ describe("Status column reflects key.blocked / scim_blocked metadata", () => { }); }); - it("should render Blocked when key.blocked is true", async () => { + it("renders Expired when the expiry date has passed", async () => { + mockUseKeys.mockReturnValue( + keysResult([{ ...mockKey, blocked: false, metadata: {}, expires: "2020-01-01T00:00:00Z" }]), + ); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByTestId(`key-status-${mockKey.token_id}`)).toHaveTextContent("Expired"); + }); + }); + + it("renders Blocked when key.blocked is true", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, blocked: true, metadata: {} }])); renderWithProviders(); @@ -470,7 +517,7 @@ describe("Status column reflects key.blocked / scim_blocked metadata", () => { expect(screen.queryByText(/Blocked by SCIM/i)).not.toBeInTheDocument(); }); - it("should mark a SCIM-blocked key with the SCIM tooltip reason", async () => { + it("marks a SCIM-blocked key with the SCIM tooltip reason", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, blocked: true, metadata: { scim_blocked: true } }])); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index cae6dc54df5..112b6cbcba4 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -1,810 +1,251 @@ "use client"; -import { useKeys, KeyListCallOptions } from "@/app/(dashboard)/hooks/keys/useKeys"; + +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useAllTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; -import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; -import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { - ColumnDef, - flexRender, - getCoreRowModel, - PaginationState, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { Badge, Icon, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Text } from "@tremor/react"; -import { InfoCircleOutlined, SyncOutlined } from "@ant-design/icons"; -import { Button as AntButton, Popover, Skeleton, Typography } from "antd"; -import { DateCell, IdCell, MoneyCell, StatusBadge } from "@/components/shared/table_cells"; -import React, { useDeferredValue, useMemo, useState } from "react"; -import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; -import { PaginatedKeyAliasSelect } from "../KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect"; + DataTable, + DataTableFilterDrawer, + DataTableFilterField, + DataTableToolbar, +} from "@/components/shared/DataTable"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { PageHeader } from "@/components/shared/PageHeader"; +import { Input } from "@/components/ui/input"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import { KeyRound } from "lucide-react"; +import React, { useCallback, useMemo, useState } from "react"; + import { KeyResponse, Team } from "../key_team_helpers/key_list"; -import FilterComponent, { FilterOption } from "../molecules/filter"; -import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; import KeyInfoView from "../templates/key_info_view"; +import { getKeyTableColumns, KEY_TABLE_HIDDEN_COLUMNS } from "./keyTableColumns"; -type KeyFilterState = { - "Team ID": string; - "Organization ID": string; - "Key Alias": string; - "User ID": string; - "Key Hash": string; +interface VirtualKeysTableProps { + headerActions?: React.ReactNode; +} + +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +const toSortOrder = (sorting: SortingState): "asc" | "desc" | undefined => { + const active = sorting[0]; + if (!active) return undefined; + return active.desc ? "desc" : "asc"; }; -const DEFAULT_KEY_FILTERS: KeyFilterState = { - "Team ID": "", - "Organization ID": "", - "Key Alias": "", - "User ID": "", - "Key Hash": "", +const FILTER_LABELS: Record = { + team_id: "Team", + org_id: "Organization", + user_id: "User ID", + key_hash: "Key ID", }; -type KeyListFilterOptions = Pick< - KeyListCallOptions, - "teamID" | "organizationID" | "selectedKeyAlias" | "userID" | "keyHash" ->; +export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { + const { data: fetchedOrganizations } = useOrganizations(); + const organizations = useMemo(() => fetchedOrganizations ?? [], [fetchedOrganizations]); + const { data: fetchedTeams } = useAllTeams(); + const allTeams = useMemo(() => fetchedTeams ?? [], [fetchedTeams]); -const toKeyListFilters = (filters: KeyFilterState): KeyListFilterOptions => ({ - teamID: filters["Team ID"].trim() || undefined, - organizationID: filters["Organization ID"].trim() || undefined, - selectedKeyAlias: filters["Key Alias"].trim() || undefined, - userID: filters["User ID"].trim() || undefined, - keyHash: filters["Key Hash"].trim() || undefined, -}); - -export function VirtualKeysTable() { - const { data: fetchedOrganizations, isLoading: isOrgsLoading } = useOrganizations(); - const resolvedOrganizations = useMemo(() => fetchedOrganizations ?? [], [fetchedOrganizations]); const [selectedKey, setSelectedKey] = useState(null); - const [sorting, setSorting] = React.useState([{ id: "created_at", desc: true }]); - const [tablePagination, setTablePagination] = React.useState({ - pageIndex: 0, - pageSize: 50, - }); - const [filters, setFilters] = useState(DEFAULT_KEY_FILTERS); - const [debouncedFilters] = useDebouncedValue(filters, { wait: 300 }); + const [sorting, setSorting] = useState(DEFAULT_SORTING); + const [tablePagination, setTablePagination] = useState({ pageIndex: 0, pageSize: 50 }); + const [columnFilters, setColumnFilters] = useState([]); + const [filtersOpen, setFiltersOpen] = useState(false); + const [searchInput, setSearchInput] = useState(""); + const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); - const sortBy = sorting.length > 0 ? sorting[0].id : null; - const sortOrder = sorting.length > 0 ? (sorting[0].desc ? "desc" : "asc") : null; + const getFilterValue = useCallback( + (columnId: string): string | undefined => { + const entry = columnFilters.find((filter) => filter.id === columnId); + return typeof entry?.value === "string" && entry.value.trim() ? entry.value.trim() : undefined; + }, + [columnFilters], + ); + + const sortBy = sorting[0]?.id; + const sortOrder = toSortOrder(sorting); + + const keyListOptions = { + teamID: getFilterValue("team_id"), + organizationID: getFilterValue("org_id"), + selectedKeyAlias: searchQuery.trim() || undefined, + userID: getFilterValue("user_id"), + keyHash: getFilterValue("key_hash"), + sortBy, + sortOrder, + expand: "user", + }; const { data: keys, isPending: isLoading, isFetching, - isError, refetch, - } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, { - ...toKeyListFilters(debouncedFilters), - sortBy: sortBy || undefined, - sortOrder: sortOrder || undefined, - expand: "user", - }); - const [expandedAccordions, setExpandedAccordions] = useState>({}); + } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, keyListOptions); const keyList = useMemo(() => keys?.keys ?? [], [keys]); + const rowCount = keys?.total_count ?? 0; - const { data: fetchedTeams, isLoading: isTeamsLoading } = useAllTeams(); - const allTeams = useMemo(() => fetchedTeams ?? [], [fetchedTeams]); - - // Defer the transition so the button stays in loading state until the table - // has rendered with the new data (mirrors the spend-logs pattern) - const isFetchingDeferred = useDeferredValue(isFetching); - const isButtonLoading = (isFetching || isFetchingDeferred) && !isError; - - const handleRefresh = () => { - refetch(); - }; - - const handleFilterChange = (newFilters: Record) => { - setFilters({ - "Team ID": newFilters["Team ID"] || "", - "Organization ID": newFilters["Organization ID"] || "", - "Key Alias": newFilters["Key Alias"] || "", - "User ID": newFilters["User ID"] || "", - "Key Hash": newFilters["Key Hash"] || "", - }); + const handleSearchChange = useCallback((value: string) => { + setSearchInput(value); setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); - }; + }, []); - const handleFilterReset = () => { - setFilters(DEFAULT_KEY_FILTERS); + const handleSortingChange = useCallback>((updaterOrValue) => { + setSorting(updaterOrValue); setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); - }; + }, []); - const totalCount = keys?.total_count ?? 0; + const handleColumnFiltersChange = useCallback>((updaterOrValue) => { + setColumnFilters(updaterOrValue); + setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); - const columns: ColumnDef[] = useMemo( - () => [ - { - id: "expander", - header: () => null, - size: 40, - enableSorting: false, - cell: ({ row }) => - row.getCanExpand() ? ( - - ) : null, - }, - { - id: "token", - accessorKey: "token", - header: "Key ID", - size: 100, - enableSorting: true, - cell: (info) => setSelectedKey(info.row.original)} />, - }, - { - id: "key_alias", - accessorKey: "key_alias", - header: "Key Alias", - size: 150, - enableSorting: true, - cell: (info) => { - const value = info.getValue() as string; - const width = info.cell.column.getSize(); - return ( - - {value ?? "-"} - - ); - }, - }, - { - id: "status", - header: "Status", - size: 100, - enableSorting: false, - cell: ({ row }) => { - const key = row.original; - if (key.blocked !== true) { - return ; - } - const isScimBlocked = (key.metadata as Record | null | undefined)?.scim_blocked === true; - const reason = isScimBlocked - ? "Blocked by SCIM (external identity provider deactivated or deleted the owning user)." - : "Blocked. Requests using this key will be rejected with 401."; - return ( - - ); - }, - }, - { - id: "key_name", - accessorKey: "key_name", - header: "Secret Key", - size: 120, - enableSorting: false, - cell: (info) => {info.getValue() as string}, - }, - { - id: "team_alias", - accessorKey: "team_id", - header: "Team", - size: 120, - enableSorting: false, - cell: (info) => { - const teamId = info.getValue() as string | null; - if (!teamId) return "-"; - const team = allTeams.find((t) => t.team_id === teamId); - const displayValue = team?.team_alias || teamId; - const width = info.cell.column.getSize(); - return ( - - {displayValue} - - ); - }, - }, - { - id: "organization_alias", - accessorKey: "org_id", - header: "Organization", - size: 140, - enableSorting: false, - cell: (info) => { - const orgId = info.getValue() as string | null; - if (!orgId) return "-"; - const org = resolvedOrganizations.find((o) => o.organization_id === orgId); - const displayValue = org?.organization_alias || orgId; - const width = info.cell.column.getSize(); - return ( - - {displayValue} - - ); - }, - }, - { - id: "user", - accessorKey: "user", - header: () => ( - - User - - - - - ), - size: 160, - enableSorting: false, - cell: ({ row }) => { - const key = row.original; - const userAlias = key.user?.user_alias ?? null; - const userEmail = key.user?.user_email ?? key.user_email ?? null; - const userId = key.user_id ?? null; - const isDefaultAdmin = userId === "default_user_id"; - const displayValue = userAlias || userEmail || userId; - const width = 160; - - const popoverContent = ( -
- {[ - { label: "User Alias", value: userAlias }, - { label: "User Email", value: userEmail }, - { label: "User ID", value: userId }, - ].map(({ label, value }) => ( -
- {label} - {value ? ( - - {value} - - ) : ( - - - )} -
- ))} -
- ); - - if (isDefaultAdmin && !userAlias && !userEmail) { - return ( - - - - - - ); - } - - return ( - - - {displayValue || "-"} - - - ); - }, - }, - { - id: "created_at", - accessorKey: "created_at", - header: "Created At", - size: 120, - enableSorting: true, - cell: (info) => , - }, - { - id: "created_by", - accessorKey: "created_by", - header: "Created By", - size: 160, - enableSorting: false, - cell: (info) => { - const userId = info.getValue() as string | null; - if (!userId) return "-"; - const key = info.row.original; - const createdByUser = key.created_by_user; - const userAlias = createdByUser?.user_alias ?? null; - const userEmail = createdByUser?.user_email ?? null; - const isDefaultAdmin = userId === "default_user_id"; - const displayValue = userAlias || userEmail || userId; - const width = 160; - - const popoverContent = ( -
- {[ - { label: "User Alias", value: userAlias }, - { label: "User Email", value: userEmail }, - { label: "User ID", value: userId }, - ].map(({ label, value }) => ( -
- {label} - {value ? ( - - {value} - - ) : ( - - - )} -
- ))} -
- ); - - if (isDefaultAdmin && !userAlias && !userEmail) { - return ( - - - - - - ); - } - - return ( - - - {displayValue} - - - ); - }, - }, - { - id: "updated_at", - accessorKey: "updated_at", - header: "Updated At", - size: 120, - enableSorting: true, - cell: (info) => , - }, - { - id: "last_active", - accessorKey: "last_active", - header: () => ( - - Last Active - - - - - ), - size: 130, - enableSorting: false, - cell: (info) => , - }, - { - id: "expires", - accessorKey: "expires", - header: "Expires", - size: 120, - enableSorting: false, - cell: (info) => , - }, - { - id: "spend", - accessorKey: "spend", - header: "Spend (USD)", - size: 100, - enableSorting: true, - cell: (info) => , - }, - { - id: "max_budget", - accessorKey: "max_budget", - header: "Budget (USD)", - size: 110, - enableSorting: true, - cell: (info) => { - const maxBudget = info.getValue() as number | null; - if (maxBudget !== null) { - return `$${formatNumberWithCommas(maxBudget)}`; - } - const teamId = info.row.original.team_id; - const team = allTeams.find((t) => t.team_id === teamId); - if (team?.max_budget != null) { - return `$${formatNumberWithCommas(team.max_budget)} (Team)`; - } - return "Unlimited"; - }, - }, - { - id: "budget_reset_at", - accessorKey: "budget_reset_at", - header: "Budget Reset", - size: 130, - enableSorting: false, - cell: (info) => , - }, - { - id: "models", - accessorKey: "models", - header: "Models", - size: 200, - enableSorting: false, - cell: (info) => { - const models = info.getValue() as string[]; - return ( -
- {Array.isArray(models) ? ( -
- {models.length === 0 ? ( - - All Proxy Models - - ) : ( - <> -
- {models.length > 3 && ( -
- { - setExpandedAccordions((prev) => ({ - ...prev, - [info.row.id]: !prev[info.row.id], - })); - }} - /> -
- )} -
- {models.slice(0, 3).map((model, index) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} - {models.length > 3 && !expandedAccordions[info.row.id] && ( - - - +{models.length - 3} {models.length - 3 === 1 ? "more model" : "more models"} - - - )} - {expandedAccordions[info.row.id] && ( -
- {models.slice(3).map((model, index) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} -
- )} -
-
- - )} -
- ) : null} -
- ); - }, - }, - { - id: "rate_limits", - header: "Rate Limits", - size: 140, - enableSorting: false, - cell: ({ row }) => { - const key = row.original; - return ( -
-
TPM: {key.tpm_limit !== null ? key.tpm_limit : "Unlimited"}
-
RPM: {key.rpm_limit !== null ? key.rpm_limit : "Unlimited"}
-
- ); - }, - }, - ], - [allTeams, resolvedOrganizations], + const columns = useMemo( + () => getKeyTableColumns({ allTeams, organizations, onSelectKey: setSelectedKey }), + [allTeams, organizations], ); - const filterOptions: FilterOption[] = [ - { - name: "Team ID", - label: "Team ID", - isSearchable: true, - loading: isTeamsLoading, - searchFn: async (searchText: string) => { - if (!allTeams || allTeams.length === 0) return []; + const teamOptions = useMemo( + () => + allTeams.map((team) => ({ + label: team.team_alias || team.team_id, + value: team.team_id, + sublabel: team.team_alias ? team.team_id : undefined, + })), + [allTeams], + ); - const filteredTeams = allTeams.filter( - (team) => - team.team_id.toLowerCase().includes(searchText.toLowerCase()) || - (team.team_alias && team.team_alias.toLowerCase().includes(searchText.toLowerCase())), - ); + const orgOptions = useMemo( + () => + organizations + .filter((org) => org.organization_id) + .map((org) => { + const id = org.organization_id as string; + return { label: org.organization_alias || id, value: id, sublabel: org.organization_alias ? id : undefined }; + }), + [organizations], + ); - return filteredTeams.map((team) => ({ - label: `${team.team_alias || team.team_id} (${team.team_id})`, - value: team.team_id, - })); - }, + const formatFilterValue = useCallback( + (columnId: string, value: unknown): string => { + const raw = String(value); + if (columnId === "team_id") { + return allTeams.find((team) => team.team_id === raw)?.team_alias || raw; + } + if (columnId === "org_id") { + return organizations.find((org) => org.organization_id === raw)?.organization_alias || raw; + } + return raw; }, - { - name: "Organization ID", - label: "Organization ID", - isSearchable: true, - loading: isOrgsLoading, - searchFn: async (searchText: string) => { - if (!resolvedOrganizations || resolvedOrganizations.length === 0) return []; + [allTeams, organizations], + ); - const filteredOrgs = resolvedOrganizations.filter( - (org) => org.organization_id?.toLowerCase().includes(searchText.toLowerCase()) ?? false, - ); - - return filteredOrgs - .filter((org) => org.organization_id !== null && org.organization_id !== undefined) - .map((org) => ({ - label: `${org.organization_id || "Unknown"} (${org.organization_id})`, - value: org.organization_id as string, - })); - }, - }, - { - name: "Key Alias", - label: "Key Alias", - customComponent: PaginatedKeyAliasSelect, - }, - { - name: "User ID", - label: "User ID", - isSearchable: false, - }, - { - name: "Key Hash", - label: "Key ID", - isSearchable: false, - }, - ]; - - const table = useReactTable({ - data: keyList, - columns: columns.filter((col) => col.id !== "expander"), - columnResizeMode: "onChange", - columnResizeDirection: "ltr", - state: { - sorting, - pagination: tablePagination, - }, - onSortingChange: (updaterOrValue) => { - const newSorting = typeof updaterOrValue === "function" ? updaterOrValue(sorting) : updaterOrValue; - setSorting(newSorting); - setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); - }, - onPaginationChange: setTablePagination, - getCoreRowModel: getCoreRowModel(), - enableSorting: true, - manualSorting: true, - manualPagination: true, - pageCount: Math.ceil(totalCount / tablePagination.pageSize), - }); - - const { pageIndex, pageSize } = table.getState().pagination; - const start = pageIndex * pageSize + 1; - const end = Math.min((pageIndex + 1) * pageSize, totalCount); - const rangeLabel = `${start} - ${end}`; - return ( -
- {selectedKey ? ( + if (selectedKey) { + return ( +
setSelectedKey(null)} keyData={selectedKey} teams={allTeams} + onDelete={refetch} /> - ) : ( -
-
- + ); + } + + return ( +
+ } + title="Virtual Keys" + subtitle="Every key that authenticates requests to the gateway." + actions={headerActions} + /> + row.token} + defaultColumnVisibility={KEY_TABLE_HIDDEN_COLUMNS} + sortingMode="server" + sorting={sorting} + onSortingChange={handleSortingChange} + paginationMode="server" + pagination={tablePagination} + onPaginationChange={setTablePagination} + rowCount={rowCount} + filterMode="server" + columnFilters={columnFilters} + onColumnFiltersChange={handleColumnFiltersChange} + enableColumnResizing + columnResizeMode="onChange" + isLoading={isLoading} + loadingMessage="Loading keys..." + noDataMessage="No keys found" + maxBodyHeight="calc(75vh - 210px)" + size="compact" + toolbar={(table) => ( + <> + refetch?.()} + isRefreshing={isFetching} + onOpenFilters={() => setFiltersOpen(true)} + filterLabels={FILTER_LABELS} + formatFilterValue={formatFilterValue} /> -
- -
-
- {isLoading ? ( - - ) : ( - - Showing {rangeLabel} of {totalCount} results - + + {({ get, set }) => ( + <> + + set("team_id", value)} + placeholder="Select a team…" + emptyText="No teams found" + /> + + + set("org_id", value)} + placeholder="Select an organization…" + emptyText="No organizations found" + /> + + + set("user_id", event.target.value)} + placeholder="Enter User ID…" + /> + + + set("key_hash", event.target.value)} + placeholder="Enter Key ID…" + /> + + )} - - } - onClick={handleRefresh} - disabled={isButtonLoading} - title="Fetch data" - > - {isButtonLoading ? "Fetching" : "Fetch"} - -
- -
- {isLoading ? ( - - ) : ( - - Page {pageIndex + 1} of {table.getPageCount()} - - )} - - {isLoading ? ( - - ) : ( - - )} - - {isLoading ? ( - - ) : ( - - )} -
-
-
-
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - { - const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`); - if (resizer) { - (resizer as HTMLElement).style.opacity = "0.5"; - } - }} - onMouseLeave={() => { - const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`); - if (resizer && !header.column.getIsResizing()) { - (resizer as HTMLElement).style.opacity = "0"; - } - }} - onClick={header.column.getCanSort() ? header.column.getToggleSortingHandler() : undefined} - > -
-
- {header.isPlaceholder - ? null - : flexRender(header.column.columnDef.header, header.getContext())} -
- {header.id !== "actions" && header.column.getCanSort() && ( -
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
- )} -
header.column.resetSize()} - onMouseDown={header.getResizeHandler()} - onTouchStart={header.getResizeHandler()} - className={`resizer ${table.options.columnResizeDirection} ${header.column.getIsResizing() ? "isResizing" : ""}`} - style={{ - position: "absolute", - right: 0, - top: 0, - height: "100%", - width: "5px", - background: header.column.getIsResizing() ? "#3b82f6" : "transparent", - cursor: "col-resize", - userSelect: "none", - touchAction: "none", - opacity: header.column.getIsResizing() ? 1 : 0, - }} - /> -
- - ))} - - ))} - - - {isLoading ? ( - - -
-

🚅 Loading keys...

-
- - - ) : keyList.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - 3 ? "px-0" : ""}`} - > - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No keys found

-
- - - )} - -
-
-
-
-
- )} + + + )} + />
); } diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx new file mode 100644 index 00000000000..133ff89a898 --- /dev/null +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -0,0 +1,364 @@ +"use client"; + +import { InfoCircleOutlined } from "@ant-design/icons"; +import { ColumnDef } from "@tanstack/react-table"; +import { Popover, Typography } from "antd"; + +import { DataTableMultiSortHeader, DataTableSortHeader, type DataTableSortField } from "@/components/shared/DataTable"; +import { Skeleton } from "@/components/ui/skeleton"; +import { + DateCell, + IdCell, + IdentityCell, + ModelsCell, + SpendBudgetCell, + StatusBadge, + type StatusTone, +} from "@/components/shared/table_cells"; + +import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; +import { KeyResponse, Team } from "../key_team_helpers/key_list"; +import { Organization } from "../networking"; + +interface KeyStatus { + tone: StatusTone; + label: string; + tooltip?: string; +} + +const SPEND_BUDGET_SORT_FIELDS: DataTableSortField[] = [ + { id: "spend", label: "Spend" }, + { id: "max_budget", label: "Budget" }, +]; + +const getKeyStatus = (key: KeyResponse): KeyStatus => { + if (key.blocked === true) { + const isScimBlocked = (key.metadata as Record | null | undefined)?.scim_blocked === true; + return { + tone: "error", + label: "Blocked", + tooltip: isScimBlocked + ? "Blocked by SCIM (external identity provider deactivated or deleted the owning user)." + : "Blocked. Requests using this key will be rejected with 401.", + }; + } + const expiresAt = key.expires ? Date.parse(key.expires) : Number.NaN; + if (!Number.isNaN(expiresAt) && expiresAt < Date.now()) { + return { tone: "warning", label: "Expired", tooltip: "This key has passed its expiry date." }; + } + return { tone: "success", label: "Active" }; +}; + +const UserPopoverCell = ({ + userAlias, + userEmail, + userId, + width, +}: { + userAlias: string | null; + userEmail: string | null; + userId: string | null; + width: number; +}) => { + const displayValue = userAlias || userEmail || userId; + const isDefaultAdmin = userId === "default_user_id"; + + const popoverContent = ( +
+ {[ + { label: "User Alias", value: userAlias }, + { label: "User Email", value: userEmail }, + { label: "User ID", value: userId }, + ].map(({ label, value }) => ( +
+ {label} + {value ? ( + + {value} + + ) : ( + - + )} +
+ ))} +
+ ); + + if (isDefaultAdmin && !userAlias && !userEmail) { + return ( + + + + + + ); + } + + return ( + + + {displayValue || "-"} + + + ); +}; + +const InfoHeader = ({ label, tooltip }: { label: string; tooltip: string }) => ( + + {label} + + + + +); + +interface KeyTableColumnsDeps { + allTeams: Team[]; + organizations: Organization[]; + onSelectKey: (key: KeyResponse) => void; +} + +export const getKeyTableColumns = ({ + allTeams, + organizations, + onSelectKey, +}: KeyTableColumnsDeps): ColumnDef[] => [ + { + id: "key_alias", + accessorKey: "key_alias", + meta: { + title: "Key", + renderSkeleton: () => ( +
+ +
+ + +
+
+ ), + }, + header: ({ column }) => , + size: 260, + enableSorting: true, + cell: ({ row }) => { + const status = getKeyStatus(row.original); + return ( + + } + onClick={() => onSelectKey(row.original)} + /> + ); + }, + }, + { + id: "token", + accessorKey: "token", + meta: { title: "Key ID" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: (info) => onSelectKey(info.row.original)} />, + }, + { + id: "team_alias", + accessorKey: "team_id", + meta: { title: "Team" }, + header: "Team", + size: 120, + enableSorting: false, + cell: (info) => { + const teamId = info.getValue() as string | null; + if (!teamId) return "-"; + const team = allTeams.find((t) => t.team_id === teamId); + const displayValue = team?.team_alias || teamId; + const width = info.cell.column.getSize(); + return ( + + {displayValue} + + ); + }, + }, + { + id: "organization_alias", + accessorKey: "org_id", + meta: { title: "Organization" }, + header: "Organization", + size: 140, + enableSorting: false, + cell: (info) => { + const orgId = info.getValue() as string | null; + if (!orgId) return "-"; + const org = organizations.find((o) => o.organization_id === orgId); + const displayValue = org?.organization_alias || orgId; + const width = info.cell.column.getSize(); + return ( + + {displayValue} + + ); + }, + }, + { + id: "user", + accessorKey: "user", + meta: { title: "User" }, + header: () => ( + + ), + size: 160, + enableSorting: false, + cell: ({ row }) => { + const key = row.original; + return ( + + ); + }, + }, + { + id: "created_at", + accessorKey: "created_at", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: (info) => , + }, + { + id: "created_by", + accessorKey: "created_by", + meta: { title: "Created By" }, + header: "Created By", + size: 160, + enableSorting: false, + cell: (info) => { + const userId = info.getValue() as string | null; + if (!userId) return "-"; + const createdByUser = info.row.original.created_by_user; + return ( + + ); + }, + }, + { + id: "updated_at", + accessorKey: "updated_at", + meta: { title: "Updated At" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: (info) => , + }, + { + id: "last_active", + accessorKey: "last_active", + meta: { title: "Last Active" }, + header: () => ( + + ), + size: 130, + enableSorting: false, + cell: (info) => , + }, + { + id: "expires", + accessorKey: "expires", + meta: { title: "Expires" }, + header: "Expires", + size: 120, + enableSorting: false, + cell: (info) => , + }, + { + id: "spend", + accessorKey: "spend", + meta: { title: "Spend / Budget", skeleton: "meter" }, + header: ({ table }) => , + size: 180, + enableSorting: true, + cell: ({ row }) => { + const teamId = row.original.team_id; + const team = allTeams.find((t) => t.team_id === teamId); + return ( + + ); + }, + }, + { + id: "budget_reset_at", + accessorKey: "budget_reset_at", + meta: { title: "Budget Reset" }, + header: "Budget Reset", + size: 130, + enableSorting: false, + cell: (info) => , + }, + { + id: "models", + accessorKey: "models", + meta: { title: "Models", skeleton: "chips" }, + header: "Models", + size: 220, + enableSorting: false, + cell: (info) => ( + + ), + }, + { + id: "rate_limits", + meta: { title: "Rate Limits" }, + header: "Rate Limits", + size: 140, + enableSorting: false, + cell: ({ row }) => { + const key = row.original; + return ( +
+
TPM: {key.tpm_limit !== null ? key.tpm_limit : "Unlimited"}
+
RPM: {key.rpm_limit !== null ? key.rpm_limit : "Unlimited"}
+
+ ); + }, + }, +]; + +export const KEY_TABLE_HIDDEN_COLUMNS: Record = { + token: false, + organization_alias: false, + created_by: false, + updated_at: false, + expires: false, + budget_reset_at: false, + rate_limits: false, +}; diff --git a/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.test.tsx b/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.test.tsx index e96936e179a..0fc9d9fd6fb 100644 --- a/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.test.tsx @@ -56,4 +56,23 @@ describe("FilterInput", () => { expect(input.value).toBe("a"); }); + + it("should not call onChange when unmounted mid-debounce", () => { + const onChange = vi.fn(); + const { unmount } = render(); + + const input = screen.getByPlaceholderText("Search..."); + + act(() => { + fireEvent.change(input, { target: { value: "test" } }); + }); + + unmount(); + + act(() => { + vi.advanceTimersByTime(300); + }); + + expect(onChange).not.toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.tsx b/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.tsx index 3590619ba25..abb6d591386 100644 --- a/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.tsx +++ b/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.tsx @@ -1,8 +1,9 @@ import { cx } from "@/lib/cva.config"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { Input } from "antd"; -import debounce from "lodash/debounce"; import { LucideIcon } from "lucide-react"; -import React, { useCallback, useEffect, useMemo, useState } from "react"; +import React, { useEffect, useState } from "react"; interface FilterInputProps { placeholder?: string; @@ -13,8 +14,6 @@ interface FilterInputProps { style?: React.CSSProperties; } -const DEBOUNCE_DELAY = 300; - export const FilterInput: React.FC = ({ placeholder, value, onChange, icon: Icon, className }) => { const [localValue, setLocalValue] = useState(value); @@ -22,22 +21,13 @@ export const FilterInput: React.FC = ({ placeholder, value, on setLocalValue(value); }, [value]); - const debouncedOnChange = useMemo(() => debounce((val: string) => onChange(val), DEBOUNCE_DELAY), [onChange]); + const debouncedOnChange = useDebouncedCallback((val: string) => onChange(val), { wait: DEBOUNCE_WAIT_MS }); - useEffect(() => { - return () => { - debouncedOnChange.cancel(); - }; - }, [debouncedOnChange]); - - const handleChange = useCallback( - (e: React.ChangeEvent) => { - const newValue = e.target.value; - setLocalValue(newValue); - debouncedOnChange(newValue); - }, - [debouncedOnChange], - ); + const handleChange = (e: React.ChangeEvent) => { + const newValue = e.target.value; + setLocalValue(newValue); + debouncedOnChange(newValue); + }; return ( ({ + fetchAvailableModels: vi.fn().mockResolvedValue([]), +})); + +const openCustomModelInput = () => { + const selector = document.querySelector(".ant-select-selector"); + expect(selector).toBeTruthy(); + act(() => { + fireEvent.mouseDown(selector!); + }); + act(() => { + fireEvent.click(screen.getByText("Enter custom model")); + }); + return screen.getByPlaceholderText("Enter custom model name"); +}; + +describe("ModelSelector custom model debounce", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + act(() => { + vi.runOnlyPendingTimers(); + }); + vi.useRealTimers(); + }); + + it("does not call onChange before the debounce wait elapses", () => { + const onChange = vi.fn(); + render(); + + const input = openCustomModelInput(); + + act(() => { + fireEvent.change(input, { target: { value: "gpt-4o" } }); + }); + + expect(onChange).not.toHaveBeenCalled(); + + act(() => { + vi.advanceTimersByTime(499); + }); + + expect(onChange).not.toHaveBeenCalled(); + }); + + it("calls onChange exactly once with the last typed value after the wait", () => { + const onChange = vi.fn(); + render(); + + const input = openCustomModelInput(); + + act(() => { + fireEvent.change(input, { target: { value: "g" } }); + fireEvent.change(input, { target: { value: "gp" } }); + fireEvent.change(input, { target: { value: "gpt-5.2" } }); + }); + + expect(onChange).not.toHaveBeenCalled(); + + act(() => { + vi.advanceTimersByTime(500); + }); + + expect(onChange).toHaveBeenCalledTimes(1); + expect(onChange).toHaveBeenCalledWith("gpt-5.2"); + }); + + it("does not call onChange when unmounted mid-wait", () => { + const onChange = vi.fn(); + const { unmount } = render(); + + const input = openCustomModelInput(); + + act(() => { + fireEvent.change(input, { target: { value: "gpt-4o" } }); + }); + + unmount(); + + act(() => { + vi.advanceTimersByTime(500); + }); + + expect(onChange).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx index 1121647c041..f2621cd1acb 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx @@ -1,9 +1,12 @@ -import React, { useState, useEffect, useRef } from "react"; +import React, { useState, useEffect } from "react"; import { TextInput, Text } from "@tremor/react"; import { Select } from "antd"; import { RobotOutlined } from "@ant-design/icons"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; +const MODEL_SELECT_DEBOUNCE_MS = 500; + interface ModelSelectorProps { accessToken: string; value?: string; @@ -30,7 +33,6 @@ const ModelSelector: React.FC = ({ const [selectedModel, setSelectedModel] = useState(value); const [showCustomModelInput, setShowCustomModelInput] = useState(false); const [modelInfo, setModelInfo] = useState([]); - const customModelTimeout = useRef(null); useEffect(() => { setSelectedModel(value); @@ -67,19 +69,13 @@ const ModelSelector: React.FC = ({ } }; - const handleCustomModelChange = (value: string) => { - // Using setTimeout to create a simple debounce effect - if (customModelTimeout.current) { - clearTimeout(customModelTimeout.current); - } - - customModelTimeout.current = setTimeout(() => { + const debouncedSelect = useDebouncedCallback( + (value: string) => { setSelectedModel(value); - if (onChange) { - onChange(value); - } - }, 500); // 500ms delay after typing stops - }; + onChange?.(value); + }, + { wait: MODEL_SELECT_DEBOUNCE_MS }, + ); return (
@@ -109,7 +105,7 @@ const ModelSelector: React.FC = ({ )} diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx new file mode 100644 index 00000000000..a70b7602e5b --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx @@ -0,0 +1,97 @@ +import { act, fireEvent, render, screen } from "@testing-library/react"; +import type { ReactNode } from "react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import type { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; +import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "./RouterSettingsAccordion"; + +vi.mock("../networking", () => ({ + getRouterSettingsCall: vi.fn().mockResolvedValue({}), +})); + +vi.mock("@/components/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn().mockResolvedValue([]), +})); + +vi.mock("../Settings/RouterSettings/Fallbacks/FallbackSelectionForm", () => ({ + FallbackSelectionForm: () => null, +})); + +vi.mock("@tremor/react", () => ({ + TabGroup: ({ children }: { children: ReactNode }) =>
{children}
, + TabList: ({ children }: { children: ReactNode }) =>
{children}
, + Tab: ({ children }: { children: ReactNode }) =>
{children}
, + TabPanels: ({ children }: { children: ReactNode }) =>
{children}
, + TabPanel: ({ children }: { children: ReactNode }) =>
{children}
, +})); + +vi.mock("../router_settings/RouterSettingsForm", () => ({ + default: ({ + value, + onChange, + }: { + value: RouterSettingsFormValue; + onChange: (value: RouterSettingsFormValue) => void; + }) => ( +
+ + +
+ ), +})); + +describe("RouterSettingsAccordion", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + }); + + const flushInitialPropagation = async (onChange: ReturnType) => { + await act(async () => { + vi.advanceTimersByTime(100); + }); + onChange.mockClear(); + }; + + it("debounces propagation and calls onChange once with the last value", async () => { + const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>(); + render(); + await flushInitialPropagation(onChange); + + fireEvent.click(screen.getByText("set-least-busy")); + act(() => { + vi.advanceTimersByTime(50); + }); + fireEvent.click(screen.getByText("set-usage-based")); + + expect(onChange).not.toHaveBeenCalled(); + + act(() => { + vi.advanceTimersByTime(99); + }); + expect(onChange).not.toHaveBeenCalled(); + + act(() => { + vi.advanceTimersByTime(1); + }); + expect(onChange).toHaveBeenCalledTimes(1); + expect(onChange.mock.calls[0][0].router_settings.routing_strategy).toBe("usage-based-routing"); + }); + + it("does not call onChange when unmounted mid-wait", async () => { + const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>(); + const { unmount } = render(); + await flushInitialPropagation(onChange); + + fireEvent.click(screen.getByText("set-least-busy")); + unmount(); + + act(() => { + vi.advanceTimersByTime(500); + }); + expect(onChange).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx index 0aa274b5749..08b917e302f 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx @@ -1,5 +1,6 @@ import React, { useEffect, useState, useImperativeHandle, forwardRef, useRef } from "react"; import { TabPanel, TabPanels, TabGroup, TabList, Tab } from "@tremor/react"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { getRouterSettingsCall } from "../networking"; import RouterSettingsForm, { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; import { Fallbacks } from "../Settings/RouterSettings/Fallbacks/AddFallbacks"; @@ -35,6 +36,8 @@ export interface RouterSettingsAccordionRef { getValue: () => RouterSettingsAccordionValue; } +const PROPAGATE_WAIT_MS = 100; + const RouterSettingsAccordion = forwardRef( ({ accessToken, value, onChange, modelData }, ref) => { const [formValue, setFormValue] = useState({ @@ -304,21 +307,26 @@ const RouterSettingsAccordion = forwardRef { - if (!onChange) { - return; - } - - const timeoutId = setTimeout(() => { + const debouncedPropagate = useDebouncedCallback( + () => { + if (!onChange) { + return; + } isInternalUpdateRef.current = true; const finalRouterSettings = buildRouterSettings(); onChange({ router_settings: finalRouterSettings, }); - }, 100); + }, + { wait: PROPAGATE_WAIT_MS }, + ); - return () => clearTimeout(timeoutId); + // Update parent when form values change (with debounce to avoid infinite loops) + useEffect(() => { + if (!onChange) { + return; + } + debouncedPropagate(); // eslint-disable-next-line react-hooks/exhaustive-deps }, [formValue, fallbacks]); diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx index 84140135db7..7d27886c7f5 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx @@ -3,6 +3,7 @@ import { Select, Typography } from "antd"; import { LoadingOutlined } from "@ant-design/icons"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { useInfiniteTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { Team } from "../key_team_helpers/key_list"; const { Text } = Typography; @@ -19,7 +20,6 @@ interface TeamDropdownProps { } const SCROLL_THRESHOLD = 0.8; -const DEBOUNCE_MS = 300; const TeamDropdown: React.FC = ({ value, @@ -31,7 +31,7 @@ const TeamDropdown: React.FC = ({ }) => { const [searchInput, setSearchInput] = useState(""); const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { - wait: DEBOUNCE_MS, + wait: DEBOUNCE_WAIT_MS, }); const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteTeams( diff --git a/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx b/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx index d91f83c589b..3a48b5f7b50 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx @@ -3,6 +3,7 @@ import { Select, Typography } from "antd"; import { LoadingOutlined } from "@ant-design/icons"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { useInfiniteTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { Team } from "../key_team_helpers/key_list"; const { Text } = Typography; @@ -17,7 +18,6 @@ interface TeamMultiSelectProps { } const SCROLL_THRESHOLD = 0.8; -const DEBOUNCE_MS = 300; const TeamMultiSelect: React.FC = ({ value = [], @@ -29,7 +29,7 @@ const TeamMultiSelect: React.FC = ({ }) => { const [searchInput, setSearchInput] = useState(""); const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { - wait: DEBOUNCE_MS, + wait: DEBOUNCE_WAIT_MS, }); const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteTeams( diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx new file mode 100644 index 00000000000..72b0e10e5d6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx @@ -0,0 +1,68 @@ +import { act, fireEvent, render, screen, within } from "@testing-library/react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import UserSearchModal from "./user_search_modal"; +import { userFilterUICall } from "@/components/networking"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; + +vi.mock("@/components/networking", () => ({ + userFilterUICall: vi.fn().mockResolvedValue([]), +})); + +const renderModal = () => + render(); + +const getEmailSearchInput = () => within(screen.getByTestId("member-email-search")).getByRole("combobox"); + +describe("UserSearchModal", () => { + beforeEach(() => { + vi.useFakeTimers(); + vi.mocked(userFilterUICall).mockClear(); + }); + + afterEach(() => { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + }); + + it("debounces the user search and fires exactly once with the last typed value", async () => { + renderModal(); + const input = getEmailSearchInput(); + + act(() => { + fireEvent.change(input, { target: { value: "a" } }); + fireEvent.change(input, { target: { value: "ab" } }); + fireEvent.change(input, { target: { value: "abc" } }); + }); + + act(() => { + vi.advanceTimersByTime(DEBOUNCE_WAIT_MS - 1); + }); + expect(userFilterUICall).not.toHaveBeenCalled(); + + await act(async () => { + vi.advanceTimersByTime(1); + await Promise.resolve(); + }); + + expect(userFilterUICall).toHaveBeenCalledTimes(1); + const params = vi.mocked(userFilterUICall).mock.calls[0][1]; + expect(params.get("user_email")).toBe("abc"); + }); + + it("does not fire the search when unmounted mid-wait", () => { + const { unmount } = renderModal(); + const input = getEmailSearchInput(); + + act(() => { + fireEvent.change(input, { target: { value: "abc" } }); + }); + + unmount(); + + act(() => { + vi.advanceTimersByTime(DEBOUNCE_WAIT_MS * 2); + }); + + expect(userFilterUICall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx index 866d7cbec7f..fafafd8e5d5 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx @@ -1,8 +1,9 @@ -import { useState, useCallback } from "react"; +import { useState } from "react"; import { Modal, Form, Button, Select, Tooltip } from "antd"; import { UserAddOutlined } from "@ant-design/icons"; -import debounce from "lodash/debounce"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { userFilterUICall } from "@/components/networking"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; interface User { user_id: string; user_email: string; @@ -93,9 +94,9 @@ const UserSearchModal: React.FC = ({ } }; - const debouncedSearch = useCallback( - debounce((text: string, fieldName: "user_email" | "user_id") => fetchUsers(text, fieldName), 300), - [], + const debouncedSearch = useDebouncedCallback( + (text: string, fieldName: "user_email" | "user_id") => fetchUsers(text, fieldName), + { wait: DEBOUNCE_WAIT_MS }, ); const handleSearch = (value: string, fieldName: "user_email" | "user_id"): void => { diff --git a/ui/litellm-dashboard/src/components/key_scope.test.ts b/ui/litellm-dashboard/src/components/key_scope.test.ts new file mode 100644 index 00000000000..7c20a1baece --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_scope.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it } from "vitest"; + +import { deriveKeyModelScope } from "./key_scope"; + +describe("deriveKeyModelScope", () => { + it("treats unrestricted keys (null/empty allowed_routes) as full model access", () => { + expect(deriveKeyModelScope(null)).toEqual({ hasModelAccess: true, label: null }); + expect(deriveKeyModelScope(undefined)).toEqual({ hasModelAccess: true, label: null }); + expect(deriveKeyModelScope([])).toEqual({ hasModelAccess: true, label: null }); + }); + + it("classifies SCIM keys as no model access", () => { + expect(deriveKeyModelScope(["/scim/*"])).toEqual({ hasModelAccess: false, label: "SCIM" }); + expect(deriveKeyModelScope(["/scim/v2/Users", "/scim/v2/Groups"])).toEqual({ + hasModelAccess: false, + label: "SCIM", + }); + }); + + it("classifies management-only keys as no model access", () => { + expect(deriveKeyModelScope(["management_routes"])).toEqual({ hasModelAccess: false, label: "Management" }); + }); + + it("classifies read-only keys as no model access", () => { + expect(deriveKeyModelScope(["info_routes"])).toEqual({ hasModelAccess: false, label: "Read-only" }); + }); + + it("leaves LLM-API and custom scopes with model access (default rendering)", () => { + expect(deriveKeyModelScope(["llm_api_routes"])).toEqual({ hasModelAccess: true, label: null }); + expect(deriveKeyModelScope(["/chat/completions"])).toEqual({ hasModelAccess: true, label: null }); + expect(deriveKeyModelScope(["management_routes", "llm_api_routes"])).toEqual({ + hasModelAccess: true, + label: null, + }); + }); + + it("prefers a persisted key_type over allowed_routes for the no-inference buckets", () => { + expect(deriveKeyModelScope([], "management")).toEqual({ hasModelAccess: false, label: "Management" }); + expect(deriveKeyModelScope([], "read_only")).toEqual({ hasModelAccess: false, label: "Read-only" }); + expect(deriveKeyModelScope(["some_future_mgmt_preset"], "management")).toEqual({ + hasModelAccess: false, + label: "Management", + }); + }); + + it("falls back to allowed_routes for null/default/llm_api key_type", () => { + expect(deriveKeyModelScope(["/scim/*"], null)).toEqual({ hasModelAccess: false, label: "SCIM" }); + expect(deriveKeyModelScope(["/scim/*"], "default")).toEqual({ hasModelAccess: false, label: "SCIM" }); + expect(deriveKeyModelScope([], "default")).toEqual({ hasModelAccess: true, label: null }); + expect(deriveKeyModelScope([], "llm_api")).toEqual({ hasModelAccess: true, label: null }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/key_scope.ts b/ui/litellm-dashboard/src/components/key_scope.ts new file mode 100644 index 00000000000..01dc595b4f1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_scope.ts @@ -0,0 +1,49 @@ +export interface KeyModelScope { + hasModelAccess: boolean; + label: string | null; +} + +const MANAGEMENT_ROUTES_PRESET = "management_routes"; +const INFO_ROUTES_PRESET = "info_routes"; +const SCIM_ROUTE_PREFIX = "/scim"; + +const MANAGEMENT_SCOPE: KeyModelScope = { hasModelAccess: false, label: "Management" }; +const READ_ONLY_SCOPE: KeyModelScope = { hasModelAccess: false, label: "Read-only" }; +const SCIM_SCOPE: KeyModelScope = { hasModelAccess: false, label: "SCIM" }; +const FULL_MODEL_ACCESS: KeyModelScope = { hasModelAccess: true, label: null }; + +const isScimRoute = (route: string): boolean => route.startsWith(SCIM_ROUTE_PREFIX); + +const isOnlyPreset = (allowedRoutes: string[], preset: string): boolean => + allowedRoutes.length === 1 && allowedRoutes[0] === preset; + +export const deriveKeyModelScope = ( + allowedRoutes: string[] | null | undefined, + keyType?: string | null, +): KeyModelScope => { + if (keyType === "management") { + return MANAGEMENT_SCOPE; + } + + if (keyType === "read_only") { + return READ_ONLY_SCOPE; + } + + if (!Array.isArray(allowedRoutes) || allowedRoutes.length === 0) { + return FULL_MODEL_ACCESS; + } + + if (allowedRoutes.every(isScimRoute)) { + return SCIM_SCOPE; + } + + if (isOnlyPreset(allowedRoutes, MANAGEMENT_ROUTES_PRESET)) { + return MANAGEMENT_SCOPE; + } + + if (isOnlyPreset(allowedRoutes, INFO_ROUTES_PRESET)) { + return READ_ONLY_SCOPE; + } + + return FULL_MODEL_ACCESS; +}; diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 285daa8f156..e1c1fcb232c 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -12,8 +12,10 @@ export interface Team { rpm_limit: number | null; organization_id: string; created_at: string; + updated_at?: string | null; keys: KeyResponse[]; keys_count?: number; + members_count?: number; members_with_roles: Member[]; spend: number; access_group_ids?: string[]; @@ -45,6 +47,7 @@ export interface KeyResponse { budget_reset_at: string; allowed_cache_controls: string[]; allowed_routes: string[]; + key_type: string | null; permissions: Record; model_spend: Record; model_max_budget: Record; diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index c24cebe7e7d..92dd3c849ef 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -7,9 +7,9 @@ import { getProxyBaseUrl } from "@/components/networking"; import { useTheme } from "@/contexts/ThemeContext"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; +import { ScrollArea } from "@/components/ui/scroll-area"; import { Sidebar, - SidebarContent, SidebarFooter, SidebarGroup, SidebarGroupLabel, @@ -608,15 +608,17 @@ const Sidebar_: React.FC = ({
- - {visibleGroups.map((group, gi) => ( - - {gi > 0 && } - {group.groupLabel} - {group.items.map((item) => renderItem(item))} - - ))} - + + + {isAdminRole(userRole) && ( diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.test.tsx new file mode 100644 index 00000000000..021aec5f85f --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.test.tsx @@ -0,0 +1,72 @@ +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { registerAuthHeaderNameGetter, registerAuthTokenGetter, registerBaseUrlGetter } from "@/lib/http/runtime"; +import { ByokCredentialModal } from "./ByokCredentialModal"; +import type { MCPServer } from "./types"; + +const fetchSpy = vi.hoisted(() => { + const spy = vi.fn<(request: Request) => Promise>(); + vi.stubGlobal("fetch", spy); + return spy; +}); + +vi.mock("@/components/molecules/message_manager", () => ({ + default: { success: vi.fn(), error: vi.fn() }, +})); + +const SERVER = { server_id: "srv-1", alias: "Linear", server_name: "Linear" } as MCPServer; + +const jsonResponse = (body: unknown, status = 200) => + new Response(JSON.stringify(body), { status, headers: { "Content-Type": "application/json" } }); + +async function fillAndSubmit(user: ReturnType) { + await user.click(screen.getByText("Continue to Authentication")); + await user.type(screen.getByPlaceholderText("Enter your API key"), "linear-key"); + await user.click(screen.getByRole("button", { name: /Connect & Authorize/ })); +} + +beforeEach(() => { + fetchSpy.mockReset(); + registerBaseUrlGetter(() => ""); + registerAuthTokenGetter(() => "sk-session"); +}); + +describe("ByokCredentialModal", () => { + it("saves the credential with the session's configured litellm key header, not a hardcoded Authorization", async () => { + registerAuthHeaderNameGetter(() => "x-litellm-api-key"); + fetchSpy.mockResolvedValue(jsonResponse({ server_id: "srv-1", has_credential: true })); + const onSuccess = vi.fn(); + const user = userEvent.setup(); + render( {}} onSuccess={onSuccess} />); + + await fillAndSubmit(user); + + await waitFor(() => expect(onSuccess).toHaveBeenCalledWith("srv-1")); + const request = fetchSpy.mock.calls[0][0]; + expect(request.method).toBe("POST"); + expect(new URL(request.url).pathname).toBe("/v1/mcp/server/srv-1/user-credential"); + expect(request.headers.get("x-litellm-api-key")).toBe("Bearer sk-session"); + expect(request.headers.get("Authorization")).toBeNull(); + expect(await request.json()).toEqual({ credential: "linear-key", save: true }); + }); + + it("surfaces the backend's detail.error message when the save fails", async () => { + registerAuthHeaderNameGetter(() => "Authorization"); + fetchSpy.mockResolvedValue( + jsonResponse({ detail: { error: "This MCP server does not support BYOK credentials" } }, 400), + ); + const MessageManager = (await import("@/components/molecules/message_manager")).default; + const onSuccess = vi.fn(); + const user = userEvent.setup(); + render( {}} onSuccess={onSuccess} />); + + await fillAndSubmit(user); + + await waitFor(() => + expect(MessageManager.error).toHaveBeenCalledWith("This MCP server does not support BYOK credentials"), + ); + expect(onSuccess).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx index cb07db871fd..f36de019aa5 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx @@ -3,6 +3,8 @@ import React, { useState } from "react"; import { Modal, Input, Switch } from "antd"; import MessageManager from "@/components/molecules/message_manager"; +import { fetchClient } from "@/lib/http/api"; +import { ApiError } from "@/lib/http/client"; import { KeyOutlined, LockOutlined, @@ -14,21 +16,22 @@ import { } from "@ant-design/icons"; import { MCPServer } from "./types"; +const byokSaveErrorMessage = (e: unknown): string => { + if (e instanceof ApiError) { + const detail = (e.body as { detail?: { error?: string } } | null)?.detail?.error; + if (detail) return detail; + } + return e instanceof Error && e.message ? e.message : "Failed to connect"; +}; + interface ByokCredentialModalProps { server: MCPServer; open: boolean; onClose: () => void; onSuccess: (serverId: string) => void; - accessToken: string; } -export const ByokCredentialModal: React.FC = ({ - server, - open, - onClose, - onSuccess, - accessToken, -}) => { +export const ByokCredentialModal: React.FC = ({ server, open, onClose, onSuccess }) => { const [step, setStep] = useState<1 | 2>(1); const [apiKey, setApiKey] = useState(""); const [saveKey, setSaveKey] = useState(true); @@ -52,23 +55,15 @@ export const ByokCredentialModal: React.FC = ({ } setLoading(true); try { - const response = await fetch(`/v1/mcp/server/${server.server_id}/user-credential`, { - method: "POST", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, - }, - body: JSON.stringify({ credential: apiKey.trim(), save: saveKey }), + await fetchClient.POST("/v1/mcp/server/{server_id}/user-credential", { + params: { path: { server_id: server.server_id } }, + body: { credential: apiKey.trim(), save: saveKey }, }); - if (!response.ok) { - const err = await response.json(); - throw new Error(err?.detail?.error || "Failed to save credential"); - } MessageManager.success(`Connected to ${serverDisplayName}`); onSuccess(server.server_id); handleClose(); - } catch (e: any) { - MessageManager.error(e.message || "Failed to connect"); + } catch (e) { + MessageManager.error(byokSaveErrorMessage(e)); } finally { setLoading(false); } diff --git a/ui/litellm-dashboard/src/components/molecules/filter.test.tsx b/ui/litellm-dashboard/src/components/molecules/filter.test.tsx index d956cd93168..66c671b12d4 100644 --- a/ui/litellm-dashboard/src/components/molecules/filter.test.tsx +++ b/ui/litellm-dashboard/src/components/molecules/filter.test.tsx @@ -547,6 +547,44 @@ describe("FilterComponent", () => { }); }); + it("cancels a pending debounced search when the component unmounts mid-type", async () => { + const user = userEvent.setup({ delay: null }); + const mockSearchFn = vi.fn().mockResolvedValue([{ label: "Result", value: "result" }]); + + const options: FilterOption[] = [ + { + name: "model", + label: "Model", + isSearchable: true, + searchFn: mockSearchFn, + }, + ]; + + const { unmount } = renderWithProviders( + , + ); + + await user.click(screen.getByRole("button", { name: "Filters" })); + + await waitFor(() => { + expect(mockSearchFn).toHaveBeenCalledWith(""); + }); + + vi.clearAllMocks(); + + const modelLabel = screen.getByText("Model"); + const modelSelect = within(modelLabel.closest("div")!).getByRole("combobox"); + await user.click(modelSelect); + await user.type(modelSelect, "test"); + + expect(mockSearchFn).not.toHaveBeenCalled(); + + unmount(); + + await new Promise((resolve) => setTimeout(resolve, 400)); + expect(mockSearchFn).not.toHaveBeenCalled(); + }); + it("should reset all filter values when reset button is clicked", async () => { const user = userEvent.setup({ delay: null }); renderWithProviders( diff --git a/ui/litellm-dashboard/src/components/molecules/filter.tsx b/ui/litellm-dashboard/src/components/molecules/filter.tsx index 45ad3a6b9ca..8218de41a19 100644 --- a/ui/litellm-dashboard/src/components/molecules/filter.tsx +++ b/ui/litellm-dashboard/src/components/molecules/filter.tsx @@ -1,6 +1,7 @@ +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { FilterIcon } from "@heroicons/react/outline"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { Button, Input, Select } from "antd"; -import debounce from "lodash/debounce"; import React, { useCallback, useEffect, useState } from "react"; export interface FilterOptionCustomComponentProps { @@ -54,8 +55,8 @@ const FilterComponent: React.FC = ({ [key: string]: boolean; }>({}); - const debouncedSearch = useCallback( - debounce(async (value: string, option: FilterOption) => { + const debouncedSearch = useDebouncedCallback( + async (value: string, option: FilterOption) => { if (!option.isSearchable || !option.searchFn) return; setSearchLoadingMap((prev) => ({ ...prev, [option.name]: true })); @@ -68,8 +69,8 @@ const FilterComponent: React.FC = ({ } finally { setSearchLoadingMap((prev) => ({ ...prev, [option.name]: false })); } - }, 300), - [], + }, + { wait: DEBOUNCE_WAIT_MS }, ); // Load initial options for searchable filters diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index b9d04a00f61..e6ee4de2735 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -530,3 +530,67 @@ describe("buildModelGroupTestRequest", () => { expect(body).toEqual({ model: "text-embedding-3-small", input: "test from litellm" }); }); }); + +describe("testMCPToolsListRequest auth headers", () => { + const originalFetch = global.fetch; + + const captureFetch = () => { + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + status: 200, + headers: { get: () => "application/json" }, + json: vi.fn().mockResolvedValue({ tools: [] }), + } as any); + global.fetch = mockFetch as any; + return mockFetch; + }; + + const sentHeaders = (mockFetch: ReturnType): Record => + (mockFetch.mock.calls[0][1] as RequestInit).headers as Record; + + afterEach(() => { + Networking.setGlobalLitellmHeaderName("Authorization"); + global.fetch = originalFetch; + }); + + it("sends the litellm key under a custom litellm_key_header_name even when an upstream OAuth token uses Authorization", async () => { + Networking.setGlobalLitellmHeaderName("x-litellm-key"); + const mockFetch = captureFetch(); + + await Networking.testMCPToolsListRequest("sk-key", {}, "upstream-oauth-token"); + + const headers = sentHeaders(mockFetch); + expect(headers["x-litellm-key"]).toBe("Bearer sk-key"); + expect(headers["Authorization"]).toBe("Bearer upstream-oauth-token"); + }); + + it("Bearer-prefixes x-litellm-api-key when it is the configured key header (raw values fail _get_bearer_token)", async () => { + Networking.setGlobalLitellmHeaderName("x-litellm-api-key"); + const mockFetch = captureFetch(); + + await Networking.testMCPToolsListRequest("sk-key", {}, "upstream-oauth-token"); + + const headers = sentHeaders(mockFetch); + expect(headers["x-litellm-api-key"]).toBe("Bearer sk-key"); + expect(headers["Authorization"]).toBe("Bearer upstream-oauth-token"); + }); + + it("never clobbers the upstream OAuth token on default deployments", async () => { + const mockFetch = captureFetch(); + + await Networking.testMCPToolsListRequest("sk-key", {}, "upstream-oauth-token"); + + const headers = sentHeaders(mockFetch); + expect(headers["Authorization"]).toBe("Bearer upstream-oauth-token"); + expect(headers["x-litellm-api-key"]).toBe("sk-key"); + }); + + it("sends the litellm key as the bearer on default deployments without an OAuth token", async () => { + const mockFetch = captureFetch(); + + await Networking.testMCPToolsListRequest("sk-key", {}); + + const headers = sentHeaders(mockFetch); + expect(headers["Authorization"]).toBe("Bearer sk-key"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 10d7f12604d..7e3af46b931 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6650,6 +6650,9 @@ export const testMCPToolsListRequest = async ( }; if (accessToken) { headers["x-litellm-api-key"] = accessToken; + if (globalLitellmHeaderName.toLowerCase() !== "authorization") { + headers[globalLitellmHeaderName] = `Bearer ${accessToken}`; + } } if (oauthAccessToken) { headers["Authorization"] = `Bearer ${oauthAccessToken}`; @@ -6844,7 +6847,11 @@ export const exchangeMcpOAuthToken = async ({ const data = await response.json(); if (!response.ok) { - const errorMessage = deriveErrorMessage(data) || data?.detail || "OAuth token exchange failed"; + const oauthErrorMessage = + typeof data?.error === "string" && typeof data?.error_description === "string" + ? `${data.error}: ${data.error_description}` + : undefined; + const errorMessage = oauthErrorMessage || deriveErrorMessage(data) || data?.detail || "OAuth token exchange failed"; throw new Error(errorMessage); } return data; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx index 0fe8adb70e1..f84b95b8d4d 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx @@ -1,7 +1,8 @@ import { act, fireEvent, within } from "@testing-library/react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; import { Team } from "../key_team_helpers/key_list"; +import { userFilterUICall } from "../networking"; import CreateKey from "./create_key_button"; const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall, teamDropdownTeamsRef } = @@ -134,14 +135,16 @@ vi.mock("antd", () => { const Select = ({ children, onChange, + onSearch, options, ...props }: { children?: any; onChange?: (value: string) => void; + onSearch?: (value: string) => void; options?: Array<{ value: string; label: string }>; - }) => - React.createElement( + }) => { + const select = React.createElement( "select", { ...props, @@ -151,6 +154,21 @@ vi.mock("antd", () => { options?.map((opt: any) => React.createElement("option", { key: opt.value, value: opt.value }, opt.label)), ); + if (!onSearch) { + return select; + } + + return React.createElement( + React.Fragment, + null, + React.createElement("input", { + "data-testid": "select-search-input", + onChange: (event: React.ChangeEvent) => onSearch(event.target.value), + }), + select, + ); + }; + Select.Option = ({ children, ...props }: { children?: any }) => React.createElement("option", props, children); const Input = (props: any) => React.createElement("input", props); @@ -641,6 +659,80 @@ describe("CreateKey", () => { }); }); + describe("user search debounce", () => { + const mockUserFilterUICall = vi.mocked(userFilterUICall); + + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + }); + + const renderUserSearch = () => { + const view = renderWithProviders( + , + ); + return { input: screen.getByTestId("select-search-input"), unmount: view.unmount }; + }; + + it("should not fire the search before the wait elapses", () => { + const { input } = renderUserSearch(); + + act(() => { + fireEvent.change(input, { target: { value: "alice" } }); + }); + + expect(mockUserFilterUICall).not.toHaveBeenCalled(); + + act(() => { + vi.advanceTimersByTime(299); + }); + + expect(mockUserFilterUICall).not.toHaveBeenCalled(); + }); + + it("should fire exactly one search carrying the last value after the wait", async () => { + const { input } = renderUserSearch(); + + act(() => { + fireEvent.change(input, { target: { value: "a" } }); + vi.advanceTimersByTime(100); + fireEvent.change(input, { target: { value: "al" } }); + vi.advanceTimersByTime(100); + fireEvent.change(input, { target: { value: "alice" } }); + }); + + expect(mockUserFilterUICall).not.toHaveBeenCalled(); + + await act(async () => { + vi.advanceTimersByTime(300); + }); + + expect(mockUserFilterUICall).toHaveBeenCalledTimes(1); + const params = mockUserFilterUICall.mock.calls[0][1] as URLSearchParams; + expect(params.get("user_email")).toBe("alice"); + }); + + it("should fire nothing when unmounted mid-wait", () => { + const { input, unmount } = renderUserSearch(); + + act(() => { + fireEvent.change(input, { target: { value: "alice" } }); + }); + + unmount(); + + act(() => { + vi.advanceTimersByTime(1000); + }); + + expect(mockUserFilterUICall).not.toHaveBeenCalled(); + }); + }); + describe("tags dropdown", () => { it("should populate tags dropdown with options from useTags hook", async () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index ef2ddab70ed..4bfd869f17a 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -10,8 +10,9 @@ import { InfoCircleOutlined } from "@ant-design/icons"; import { useQueryClient } from "@tanstack/react-query"; import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react"; import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip, Typography } from "antd"; -import debounce from "lodash/debounce"; -import React, { useCallback, useEffect, useState } from "react"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; +import React, { useEffect, useState } from "react"; import { rolesWithWriteAccess } from "../../utils/roles"; import AgentSelector from "../agent_management/AgentSelector"; import { mapDisplayToInternalNames } from "../callback_info_helpers"; @@ -688,14 +689,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp } }; - const debouncedSearch = useCallback( - debounce((text: string) => fetchUsers(text), 300), - [accessToken], - ); - - const handleUserSearch = (value: string): void => { - debouncedSearch(value); - }; + const handleUserSearch = useDebouncedCallback((text: string) => fetchUsers(text), { wait: DEBOUNCE_WAIT_MS }); const handleUserSelect = (_value: string, option: UserOption): void => { const selectedUser = option.user; diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 6f4ad17313b..46811e0106b 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -566,7 +566,7 @@ const Settings: React.FC = ({ accessToken, userRole, userID, } return ( -
+
diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index ef0d842ad3e..0d3468d4ada 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -5,7 +5,7 @@ import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; import { DataTable } from "./DataTable"; -import { DataTableSortHeader } from "./DataTableSortHeader"; +import { DataTableMultiSortHeader, DataTableSortHeader } from "./DataTableSortHeader"; import { DataTableViewOptions } from "./DataTableViewOptions"; interface Person { @@ -55,6 +55,23 @@ const dropdownSortColumns: ColumnDef[] = [ }, ]; +const multiSortColumns: ColumnDef[] = [ + { + id: "spend", + accessorKey: "name", + header: ({ table }) => ( + + ), + cell: ({ row }) => {row.original.name}, + }, +]; + const nameEmailColumns: ColumnDef[] = [ { accessorKey: "name", @@ -165,6 +182,60 @@ describe("DataTable sorting", () => { await user.click(await screen.findByText("Reset")); expect(names()).toEqual(["Charlie", "Alice", "Bob"]); }); + + it("multi-sort header emits the chosen field id (not the column id) as the sort key", async () => { + const user = userEvent.setup(); + const onSortingChange = vi.fn(); + render( + , + ); + + await user.click(screen.getByTestId("sort-trigger-spend")); + await user.click(await screen.findByText("Budget descending")); + expect(onSortingChange).toHaveBeenLastCalledWith([{ id: "max_budget", desc: true }]); + + await user.click(screen.getByTestId("sort-trigger-spend")); + await user.click(await screen.findByText("Spend ascending")); + expect(onSortingChange).toHaveBeenLastCalledWith([{ id: "spend", desc: false }]); + }); + + it("multi-sort header reflects the active field and direction, and Reset clears it", async () => { + const user = userEvent.setup(); + const onSortingChange = vi.fn(); + render( + , + ); + + await user.click(screen.getByTestId("sort-trigger-spend")); + // The header trigger shows the active (descending) indicator while sorted by a field it owns. + expect(screen.getByTestId("sort-trigger-spend").querySelector("[data-sort-indicator='desc']")).not.toBeNull(); + + await user.click(await screen.findByText("Reset")); + expect(onSortingChange).toHaveBeenLastCalledWith([]); + }); +}); + +describe("DataTable layout", () => { + it("stretches the table to fill the container when resizing is on, so hidden columns leave no right-side gap", () => { + const { container } = render(); + + const table = container.querySelector("table"); + expect(table).not.toBeNull(); + // width pins the natural column total (horizontal scroll on overflow); minWidth:100% fills the gap on underflow. + expect(table?.style.minWidth).toBe("100%"); + }); }); describe("DataTable pagination", () => { @@ -303,6 +374,38 @@ describe("DataTable loading", () => { // per-column widths differ instead of every cell sharing one fixed width expect(new Set(bars.map((bar) => bar.className)).size).toBeGreaterThan(1); }); + + it("renders shape-specific skeletons for badge, chips, and meter columns", () => { + const columns: ColumnDef[] = [ + { id: "badge", header: "Badge", meta: { skeleton: "badge" }, cell: () => null }, + { id: "chips", header: "Chips", meta: { skeleton: "chips" }, cell: () => null }, + { id: "meter", header: "Meter", meta: { skeleton: "meter" }, cell: () => null }, + ]; + render(); + + const firstRow = screen.getAllByTestId("skeleton-row").at(0); + const cells = Array.from(firstRow?.querySelectorAll("td") ?? []); + const barsIn = (cell: Element | undefined) => cell?.querySelectorAll('[data-slot="skeleton"]').length ?? 0; + + // badge = a single pill, chips = three pills, meter = value bar + track bar + expect(barsIn(cells[0])).toBe(1); + expect(cells[0]?.querySelector('[data-slot="skeleton"]')?.className).toContain("rounded-full"); + expect(barsIn(cells[1])).toBe(3); + expect(barsIn(cells[2])).toBe(2); + }); + + it("uses a column's renderSkeleton override when provided", () => { + const columns: ColumnDef[] = [ + { + id: "custom", + header: "Custom", + meta: { renderSkeleton: () =>
loading
}, + cell: () => null, + }, + ]; + render(); + expect(screen.getAllByTestId("custom-skeleton").length).toBeGreaterThan(0); + }); }); describe("DataTable column visibility", () => { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index 758ca5a597b..d55831b87ca 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -338,7 +338,11 @@ const SKELETON_WIDTHS = ["w-[58%]", "w-[44%]", "w-[70%]", "w-[50%]", "w-[64%]", function SkeletonCell({ column, index }: { column: Column | undefined; index: number }) { const meta = column?.columnDef.meta; const width = SKELETON_WIDTHS[index % SKELETON_WIDTHS.length]; - if (meta?.skeleton === "twoLine") { + const shape = meta?.skeleton; + if (meta?.renderSkeleton !== undefined) { + return <>{meta.renderSkeleton()}; + } + if (shape === "twoLine") { return (
@@ -346,6 +350,26 @@ function SkeletonCell({ column, index }: { column: Column
); } + if (shape === "badge") { + return ; + } + if (shape === "chips") { + return ( +
+ + + +
+ ); + } + if (shape === "meter") { + return ( +
+ + +
+ ); + } return ; } @@ -508,7 +532,7 @@ export function DataTable(props: DataTableProps { if (paginationSlot !== undefined) { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.tsx index 1cf09ce4f47..8988bd6c2d0 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.tsx @@ -1,8 +1,8 @@ "use client"; import { Menu } from "@base-ui/react/menu"; -import type { Column, SortDirection } from "@tanstack/react-table"; -import { ChevronDown, ChevronsUpDown, ChevronUp, X } from "lucide-react"; +import type { Column, SortDirection, Table } from "@tanstack/react-table"; +import { Check, ChevronDown, ChevronsUpDown, ChevronUp, X } from "lucide-react"; import type * as React from "react"; import { cn } from "@/lib/cva.config"; @@ -94,3 +94,101 @@ export function DataTableSortHeader({ ); } + +export interface DataTableSortField { + /** Backend sort column, sent verbatim as the sorting state id (e.g. "spend", "max_budget"). */ + id: string; + label: string; +} + +interface DataTableMultiSortHeaderProps { + table: Table; + fields: DataTableSortField[]; + className?: string; +} + +/** + * Sort header for a column that merges several backend-sortable fields into one cell + * (e.g. a combined Spend / Budget cell). The header label is the field labels joined by " / ", + * with the field currently driving the sort emphasized so the active column reads at a glance + * without opening the menu. The chevron opens a menu offering each field in both directions. + */ +export function DataTableMultiSortHeader({ table, fields, className }: DataTableMultiSortHeaderProps) { + const active = table.getState().sorting[0]; + const activeField = active !== undefined && fields.some((field) => field.id === active.id) ? active : undefined; + const activeDirection: SortDirection = activeField?.desc === true ? "desc" : "asc"; + const sorted: false | SortDirection = activeField === undefined ? false : activeDirection; + + const options = fields.flatMap((field) => [ + { key: `${field.id}-asc`, id: field.id, desc: false, label: `${field.label} ascending`, Icon: ChevronUp }, + { key: `${field.id}-desc`, id: field.id, desc: true, label: `${field.label} descending`, Icon: ChevronDown }, + ]); + + const segmentClass = (isActive: boolean): string => { + if (isActive) return "font-semibold text-foreground"; + if (activeField) return "text-muted-foreground"; + return ""; + }; + + const labelSegments = fields.flatMap((field, index) => { + const isActive = activeField?.id === field.id; + const segment = ( + + {field.label} + + ); + if (index === 0) return [segment]; + return [ + + {" / "} + , + segment, + ]; + }); + + return ( +
+ {labelSegments} + + field.label).join(" or ")}`} + onClick={(event) => event.stopPropagation()} + className={cn( + "inline-flex size-6 items-center justify-center rounded-md hover:bg-muted", + sorted ? "text-primary" : "text-muted-foreground", + )} + > + + + } + /> + + + + {options.map((option) => { + const isActive = activeField?.id === option.id && activeField.desc === option.desc; + return ( + table.setSorting([{ id: option.id, desc: option.desc }])} + > + {option.label} + {isActive && } + + ); + })} + table.setSorting([])}> + Reset + + + + + +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/columnMeta.ts b/ui/litellm-dashboard/src/components/shared/DataTable/columnMeta.ts index 0f14c277c6f..eff4e0cb7db 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/columnMeta.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/columnMeta.ts @@ -1,4 +1,5 @@ import type { RowData } from "@tanstack/react-table"; +import type * as React from "react"; import type { ColumnPinnedSide, DataTableSkeletonShape } from "./types"; @@ -10,5 +11,7 @@ declare module "@tanstack/react-table" { title?: string; pinned?: ColumnPinnedSide; skeleton?: DataTableSkeletonShape; + /** Full control over this column's loading skeleton, for cells the built-in shapes can't mirror. */ + renderSkeleton?: () => React.ReactNode; } } diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts index c4218f6051a..1ee1eed1258 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts @@ -5,7 +5,12 @@ export { DataTableFilterDrawer, DataTableFilterField, type FilterDraft } from ". export { DataTablePagination, DEFAULT_PAGE_SIZE_OPTIONS } from "./DataTablePagination"; export { DataTableToolbar } from "./DataTableToolbar"; export { DataTableViewOptions } from "./DataTableViewOptions"; -export { DataTableSortHeader, type DataTableSortVariant } from "./DataTableSortHeader"; +export { + DataTableSortHeader, + DataTableMultiSortHeader, + type DataTableSortVariant, + type DataTableSortField, +} from "./DataTableSortHeader"; export type { DataTablePaginationProps } from "./DataTablePagination"; export type { ColumnPinnedSide, diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts index f5130b4c823..672ab512ef4 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts @@ -18,7 +18,7 @@ export type FilterMode = "none" | "client" | "server"; export type ColumnResizeMode = "onEnd" | "onChange"; export type DataTableSize = "compact" | "default"; export type ColumnPinnedSide = "left" | "right"; -export type DataTableSkeletonShape = "text" | "twoLine"; +export type DataTableSkeletonShape = "text" | "twoLine" | "badge" | "chips" | "meter"; export interface DataTableProps { data: TData[]; diff --git a/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx b/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx new file mode 100644 index 00000000000..f7a313271da --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx @@ -0,0 +1,31 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; + +import { PageHeader } from "./PageHeader"; + +describe("PageHeader", () => { + it("renders the title as a heading", () => { + render(); + expect(screen.getByRole("heading", { name: "Virtual Keys" })).toBeInTheDocument(); + }); + + it("renders the subtitle, icon, and actions when provided", () => { + render( + } + actions={} + />, + ); + expect(screen.getByText("Every key that authenticates requests")).toBeInTheDocument(); + expect(screen.getByTestId("icon")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Create New Key" })).toBeInTheDocument(); + }); + + it("omits the optional slots when not provided", () => { + render(); + expect(screen.queryByRole("button")).not.toBeInTheDocument(); + expect(document.querySelector("p")).toBeNull(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/PageHeader.tsx b/ui/litellm-dashboard/src/components/shared/PageHeader.tsx new file mode 100644 index 00000000000..e314e8e8bc2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/PageHeader.tsx @@ -0,0 +1,25 @@ +"use client"; + +import * as React from "react"; + +interface PageHeaderProps { + title: React.ReactNode; + subtitle?: React.ReactNode; + icon?: React.ReactNode; + actions?: React.ReactNode; +} + +export function PageHeader({ title, subtitle, icon, actions }: PageHeaderProps) { + return ( +
+
+ {icon != null && {icon}} +
+

{title}

+ {subtitle != null &&

{subtitle}

} +
+
+ {actions != null &&
{actions}
} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx new file mode 100644 index 00000000000..acf50d282b4 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx @@ -0,0 +1,64 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { SearchSelect } from "./SearchSelect"; + +const OPTIONS = [ + { label: "Acme Prod", value: "team-1" }, + { label: "Growth", value: "team-2" }, + { label: "Data Team", value: "team-3" }, +]; + +describe("SearchSelect", () => { + it("renders the placeholder when nothing is selected", () => { + render(); + expect(screen.getByPlaceholderText("Select Team…")).toBeInTheDocument(); + }); + + it("shows the selected option's label in the field", () => { + render(); + expect(screen.getByRole("combobox")).toHaveValue("Growth"); + }); + + it("shows a clear control only when a value is selected", () => { + const { rerender } = render(); + expect(document.querySelector('[data-slot="combobox-clear"]')).toBeNull(); + rerender(); + expect(document.querySelector('[data-slot="combobox-clear"]')).not.toBeNull(); + }); + + it("filters the options client-side as you type", async () => { + const user = userEvent.setup(); + render(); + const input = screen.getByRole("combobox"); + await user.click(input); + await user.type(input, "grow"); + expect(await screen.findByText("Growth")).toBeInTheDocument(); + expect(screen.queryByText("Acme Prod")).not.toBeInTheDocument(); + }); + + it("renders a muted sublabel and matches it when searching", async () => { + const user = userEvent.setup(); + render( + , + ); + const input = screen.getByRole("combobox"); + await user.click(input); + expect(await screen.findByText("team-abc-123")).toBeInTheDocument(); + await user.type(input, "abc-123"); + expect(await screen.findByText("Acme Prod")).toBeInTheDocument(); + }); + + it("selects an option and reports its value", async () => { + const onValueChange = vi.fn(); + const user = userEvent.setup(); + render(); + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("Growth")); + expect(onValueChange).toHaveBeenCalledWith("team-2"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx new file mode 100644 index 00000000000..eff5cfc804d --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx @@ -0,0 +1,78 @@ +"use client"; + +import { + Combobox, + ComboboxContent, + ComboboxEmpty, + ComboboxInput, + ComboboxItem, + ComboboxList, +} from "@/components/ui/combobox"; + +export interface SearchSelectOption { + label: string; + value: string; + /** Optional muted second line (e.g. an id); also matched when searching. */ + sublabel?: string; +} + +interface SearchSelectProps { + options: SearchSelectOption[]; + value?: string; + onValueChange: (value: string) => void; + placeholder?: string; + emptyText?: string; + disabled?: boolean; + className?: string; +} + +const matchesQuery = (option: SearchSelectOption, query: string): boolean => { + const q = query.trim().toLowerCase(); + if (!q) return true; + return option.label.toLowerCase().includes(q) || (option.sublabel?.toLowerCase().includes(q) ?? false); +}; + +export function SearchSelect({ + options, + value, + onValueChange, + placeholder = "Select…", + emptyText = "No results", + disabled = false, + className, +}: SearchSelectProps) { + const selected = options.find((option) => option.value === value) ?? null; + + return ( + onValueChange(item?.value ?? "")} + isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value} + itemToStringLabel={(item: SearchSelectOption) => item.label} + filter={matchesQuery} + disabled={disabled} + > + + + {emptyText} + + {(item: SearchSelectOption) => ( + + + {item.label} + {item.sublabel != null && item.sublabel !== "" && ( + {item.sublabel} + )} + + + )} + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.test.tsx new file mode 100644 index 00000000000..f3ec203adbb --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.test.tsx @@ -0,0 +1,41 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { IdentityCell } from "./identity_cell"; + +describe("IdentityCell", () => { + it("renders the title", () => { + render(); + expect(screen.getByText("prod-gateway")).toBeInTheDocument(); + }); + + it("renders the subtitle and an inline badge together", () => { + render(Active} />); + expect(screen.getByText("sk-...v0Pw")).toBeInTheDocument(); + expect(screen.getByText("Active")).toBeInTheDocument(); + }); + + it("omits the subtitle row when there is no subtitle or badge", () => { + render(); + expect(document.querySelector("span.font-mono")).toBeNull(); + }); + + it("renders a static div (no button) when not clickable", () => { + render(); + expect(screen.queryByRole("button")).not.toBeInTheDocument(); + }); + + it("renders a clickable button that signals interactivity and fires onClick", async () => { + const onClick = vi.fn(); + const user = userEvent.setup(); + render(); + const button = screen.getByRole("button"); + expect(button.querySelector(".lucide-chevron-right")).not.toBeNull(); + // The clickable area must read as clickable: a hover background and a pointer cursor. + expect(button.className).toContain("hover:bg-muted"); + expect(button.className).toContain("cursor-pointer"); + await user.click(button); + expect(onClick).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.tsx new file mode 100644 index 00000000000..97263436454 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.tsx @@ -0,0 +1,51 @@ +"use client"; + +import { ChevronRight } from "lucide-react"; +import * as React from "react"; + +import { cn } from "@/lib/cva.config"; + +interface IdentityCellProps { + title: React.ReactNode; + subtitle?: React.ReactNode; + badge?: React.ReactNode; + onClick?: () => void; + className?: string; + titleClassName?: string; +} + +export function IdentityCell({ title, subtitle, badge, onClick, className, titleClassName }: IdentityCellProps) { + const hasSubtitleRow = (subtitle != null && subtitle !== "") || badge != null; + + const body = ( +
+ {title} + {hasSubtitleRow && ( + + {subtitle != null && subtitle !== "" && ( + {subtitle} + )} + {badge} + + )} +
+ ); + + if (onClick != null) { + return ( + + ); + } + + return
{body}
; +} diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/index.ts b/ui/litellm-dashboard/src/components/shared/table_cells/index.ts index e189413d43d..9fdd04d169c 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/index.ts +++ b/ui/litellm-dashboard/src/components/shared/table_cells/index.ts @@ -1,5 +1,8 @@ export { CellTooltip } from "./cell_tooltip"; export { DateCell, formatCellDate, formatFullTimestamp, type DatePrecision } from "./date_cell"; export { IdCell, type IdCellVariant } from "./id_cell"; +export { IdentityCell } from "./identity_cell"; +export { ModelsCell } from "./models_cell"; export { MoneyCell } from "./money_cell"; +export { SpendBudgetCell } from "./spend_budget_cell"; export { StatusBadge, type StatusTone } from "./status_badge"; diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx new file mode 100644 index 00000000000..3a618114613 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx @@ -0,0 +1,69 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it } from "vitest"; + +import { ModelsCell } from "./models_cell"; + +describe("ModelsCell", () => { + it("shows 'All Proxy Models' when the list is empty, null, or undefined", () => { + const { rerender } = render(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + rerender(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + rerender(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + }); + + it("shows 'No model access' for scope-restricted keys with an empty model list", () => { + const { rerender } = render(); + expect(screen.getByText("No model access")).toBeInTheDocument(); + expect(screen.queryByText("All Proxy Models")).not.toBeInTheDocument(); + rerender(); + expect(screen.getByText("No model access")).toBeInTheDocument(); + rerender(); + expect(screen.getByText("No model access")).toBeInTheDocument(); + }); + + it("still shows 'All Proxy Models' for empty models when the key is not scope-restricted", () => { + render(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + expect(screen.queryByText("No model access")).not.toBeInTheDocument(); + }); + + it("uses a persisted key_type to render 'No model access' regardless of allowed_routes", () => { + const { rerender } = render(); + expect(screen.getByText("No model access")).toBeInTheDocument(); + expect(screen.queryByText("All Proxy Models")).not.toBeInTheDocument(); + rerender(); + expect(screen.getByText("No model access")).toBeInTheDocument(); + }); + + it("renders every model with no overflow badge when at or below the limit", () => { + render(); + expect(screen.getByText("gpt-4o")).toBeInTheDocument(); + expect(screen.getByText("claude-sonnet-4-5")).toBeInTheDocument(); + expect(screen.getByText("o3-mini")).toBeInTheDocument(); + expect(screen.queryByText(/more$/)).not.toBeInTheDocument(); + }); + + it("collapses models beyond the limit into a '+N more' badge", () => { + render(); + expect(screen.getByText("a")).toBeInTheDocument(); + expect(screen.getByText("b")).toBeInTheDocument(); + expect(screen.queryByText("c")).not.toBeInTheDocument(); + expect(screen.getByText("+3 more")).toBeInTheDocument(); + }); + + it("reveals the hidden models in a tooltip on hover", async () => { + const user = userEvent.setup(); + render(); + await user.hover(screen.getByText("+2 more")); + expect(await screen.findByText("c")).toBeInTheDocument(); + expect(await screen.findByText("d")).toBeInTheDocument(); + }); + + it("labels the all-proxy-models wildcard", () => { + render(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx new file mode 100644 index 00000000000..81928d6fe87 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx @@ -0,0 +1,72 @@ +"use client"; + +import { deriveKeyModelScope } from "@/components/key_scope"; +import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; +import { Badge } from "@/components/ui/badge"; + +import { CellTooltip } from "./cell_tooltip"; + +interface ModelsCellProps { + models: string[] | null | undefined; + maxVisible?: number; + allowedRoutes?: string[] | null; + keyType?: string | null; +} + +const WILDCARD_MODEL = "all-proxy-models"; + +const formatModel = (model: string): string => { + if (model === WILDCARD_MODEL) { + return "All Proxy Models"; + } + const name = getModelDisplayName(model); + return name.length > 30 ? `${name.slice(0, 30)}...` : name; +}; + +export function ModelsCell({ models, maxVisible = 3, allowedRoutes, keyType }: ModelsCellProps) { + if (!Array.isArray(models) || models.length === 0) { + const scope = deriveKeyModelScope(allowedRoutes, keyType); + if (!scope.hasModelAccess) { + return ( + + No model access + + } + /> + ); + } + return All Proxy Models; + } + + const visible = models.slice(0, maxVisible); + const overflow = models.slice(maxVisible); + + return ( +
+ {visible.map((model, index) => ( + + {formatModel(model)} + + ))} + {overflow.length > 0 && ( + + {overflow.map((model, index) => ( + {formatModel(model)} + ))} +
+ } + trigger={ + + +{overflow.length} more + + } + /> + )} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx new file mode 100644 index 00000000000..707441aef1d --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx @@ -0,0 +1,53 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; + +import { SpendBudgetCell } from "./spend_budget_cell"; + +const indicator = (container: HTMLElement) => container.querySelector('[data-slot="meter-indicator"]'); + +describe("SpendBudgetCell", () => { + it("shows Unlimited and renders no meter when there is no budget", () => { + const { container } = render(); + expect(screen.getByText("· Unlimited")).toBeInTheDocument(); + expect(screen.queryByRole("meter")).not.toBeInTheDocument(); + expect(indicator(container)).toBeNull(); + }); + + it("shows $0.00 for zero or undefined spend, never a hyphen", () => { + const { rerender } = render(); + expect(screen.getByText("$0.00")).toBeInTheDocument(); + expect(screen.queryByText("-")).not.toBeInTheDocument(); + rerender(); + expect(screen.getByText("$0.00")).toBeInTheDocument(); + expect(screen.queryByText("-")).not.toBeInTheDocument(); + }); + + it("renders a meter carrying the spend and budget when a budget exists", () => { + render(); + const meter = screen.getByRole("meter"); + expect(meter).toHaveAttribute("aria-valuenow", "25"); + expect(meter).toHaveAttribute("aria-valuemax", "100"); + expect(screen.getByText("of $100")).toBeInTheDocument(); + }); + + it("keeps the default tone below 80% usage", () => { + const { container } = render(); + expect(indicator(container)?.className).toContain("bg-primary"); + }); + + it("switches to the warning tone at 80% usage", () => { + const { container } = render(); + expect(indicator(container)?.className).toContain("bg-amber-500"); + }); + + it("switches to the over tone above 100% usage", () => { + const { container } = render(); + expect(indicator(container)?.className).toContain("bg-destructive"); + }); + + it("falls back to the team budget and labels it", () => { + render(); + expect(screen.getByText("of $200 (Team)")).toBeInTheDocument(); + expect(screen.getByRole("meter")).toHaveAttribute("aria-valuemax", "200"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx new file mode 100644 index 00000000000..10956f23b1c --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx @@ -0,0 +1,44 @@ +"use client"; + +import { Meter, MeterIndicator, MeterTrack } from "@/components/ui/meter"; +import { formatNumberWithCommas, getSpendString } from "@/utils/dataUtils"; + +interface SpendBudgetCellProps { + spend: number | null | undefined; + maxBudget: number | null | undefined; + teamMaxBudget?: number | null; +} + +const meterTone = (pct: number): "default" | "warning" | "over" => { + if (pct > 100) return "over"; + if (pct >= 80) return "warning"; + return "default"; +}; + +export function SpendBudgetCell({ spend, maxBudget, teamMaxBudget }: SpendBudgetCellProps) { + const spendValue = typeof spend === "number" && !Number.isNaN(spend) ? spend : 0; + const budget = maxBudget ?? teamMaxBudget ?? null; + const isTeamBudget = maxBudget == null && teamMaxBudget != null; + const hasBudget = typeof budget === "number" && budget > 0; + const pct = hasBudget ? (spendValue / budget) * 100 : 0; + + const spendText = spendValue > 0 ? getSpendString(spendValue, 4) : "$0.00"; + const budgetLabel = + budget === null ? "· Unlimited" : `of $${formatNumberWithCommas(budget)}${isTeamBudget ? " (Team)" : ""}`; + + return ( +
+
+ {spendText}{" "} + {budgetLabel} +
+ {hasBudget && ( + + + + + + )} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index 207e6f2ccfe..66690e5478f 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -9,6 +9,7 @@ import { DataTableToolbar, } from "@/components/shared/DataTable"; import { Input } from "@/components/ui/input"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { ColumnDef, ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; @@ -17,6 +18,7 @@ import { Popover, Tooltip, Typography } from "antd"; import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; import React, { useCallback, useEffect, useMemo, useState } from "react"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; +import { deriveKeyModelScope } from "../key_scope"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; import KeyInfoView from "../templates/key_info_view"; @@ -43,7 +45,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi const [columnFilters, setColumnFilters] = useState([]); const [filtersOpen, setFiltersOpen] = useState(false); const [searchInput, setSearchInput] = useState(""); - const [searchQuery] = useDebouncedValue(searchInput, { wait: 300 }); + const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); const handleSearchChange = useCallback((value: string) => { setSearchInput(value); @@ -338,14 +340,24 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi enableSorting: false, cell: (info) => { const models = info.getValue() as string[]; + const scope = deriveKeyModelScope(info.row.original.allowed_routes, info.row.original.key_type); + const emptyModelsBadge = !scope.hasModelAccess ? ( + + + No model access + + + ) : ( + + All Proxy Models + + ); return (
{Array.isArray(models) ? (
{models.length === 0 ? ( - - All Proxy Models - + emptyModelsBadge ) : ( <>
diff --git a/ui/litellm-dashboard/src/components/ui/combobox.tsx b/ui/litellm-dashboard/src/components/ui/combobox.tsx new file mode 100644 index 00000000000..2854928140e --- /dev/null +++ b/ui/litellm-dashboard/src/components/ui/combobox.tsx @@ -0,0 +1,266 @@ +"use client"; + +import * as React from "react"; +import { Combobox as ComboboxPrimitive } from "@base-ui/react"; + +import { cn } from "@/lib/cva.config"; +import { Button } from "@/components/ui/button"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; +import { ChevronDownIcon, XIcon, CheckIcon } from "lucide-react"; + +const Combobox = ComboboxPrimitive.Root; + +function ComboboxValue({ ...props }: ComboboxPrimitive.Value.Props) { + return ; +} + +const ComboboxTrigger = React.forwardRef< + React.ComponentRef, + ComboboxPrimitive.Trigger.Props +>(({ className, children, ...props }, ref) => { + return ( + + {children} + + + ); +}); +ComboboxTrigger.displayName = "ComboboxTrigger"; + +function ComboboxClear({ className, ...props }: ComboboxPrimitive.Clear.Props) { + return ( + } + className={cn(className)} + {...props} + > + + + ); +} + +function ComboboxInput({ + className, + children, + disabled = false, + showTrigger = true, + showClear = false, + ...props +}: ComboboxPrimitive.Input.Props & { + showTrigger?: boolean; + showClear?: boolean; +}) { + return ( + + } {...props} /> + + {showTrigger && ( + } + data-slot="input-group-button" + className="group-has-data-[slot=combobox-clear]/input-group:hidden data-pressed:bg-transparent" + disabled={disabled} + /> + )} + {showClear && } + + {children} + + ); +} + +function ComboboxContent({ + className, + side = "bottom", + sideOffset = 6, + align = "start", + alignOffset = 0, + anchor, + ...props +}: ComboboxPrimitive.Popup.Props & + Pick) { + return ( + + + + + + ); +} + +function ComboboxList({ className, ...props }: ComboboxPrimitive.List.Props) { + return ( + + ); +} + +function ComboboxItem({ className, children, ...props }: ComboboxPrimitive.Item.Props) { + return ( + + {children} + } + > + + + + ); +} + +function ComboboxGroup({ className, ...props }: ComboboxPrimitive.Group.Props) { + return ; +} + +function ComboboxLabel({ className, ...props }: ComboboxPrimitive.GroupLabel.Props) { + return ( + + ); +} + +function ComboboxCollection({ ...props }: ComboboxPrimitive.Collection.Props) { + return ; +} + +function ComboboxEmpty({ className, ...props }: ComboboxPrimitive.Empty.Props) { + return ( + + ); +} + +function ComboboxSeparator({ className, ...props }: ComboboxPrimitive.Separator.Props) { + return ( + + ); +} + +function ComboboxChips({ + className, + ...props +}: React.ComponentPropsWithRef & ComboboxPrimitive.Chips.Props) { + return ( + + ); +} + +function ComboboxChip({ + className, + children, + showRemove = true, + ...props +}: ComboboxPrimitive.Chip.Props & { + showRemove?: boolean; +}) { + return ( + + {children} + {showRemove && ( + } + className="-ml-1 opacity-50 hover:opacity-100" + data-slot="combobox-chip-remove" + > + + + )} + + ); +} + +function ComboboxChipsInput({ className, ...props }: ComboboxPrimitive.Input.Props) { + return ( + + ); +} + +function useComboboxAnchor() { + return React.useRef(null); +} + +export { + Combobox, + ComboboxInput, + ComboboxContent, + ComboboxList, + ComboboxItem, + ComboboxGroup, + ComboboxLabel, + ComboboxCollection, + ComboboxEmpty, + ComboboxSeparator, + ComboboxChips, + ComboboxChip, + ComboboxChipsInput, + ComboboxTrigger, + ComboboxValue, + useComboboxAnchor, +}; diff --git a/ui/litellm-dashboard/src/components/ui/dropdown-menu.tsx b/ui/litellm-dashboard/src/components/ui/dropdown-menu.tsx new file mode 100644 index 00000000000..03fefba7d19 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ui/dropdown-menu.tsx @@ -0,0 +1,254 @@ +"use client"; + +import * as React from "react"; +import { Menu as MenuPrimitive } from "@base-ui/react/menu"; + +import { cn } from "@/lib/cva.config"; +import { ChevronRightIcon, CheckIcon } from "lucide-react"; + +function DropdownMenu({ ...props }: MenuPrimitive.Root.Props) { + return ; +} + +function DropdownMenuPortal({ ...props }: MenuPrimitive.Portal.Props) { + return ; +} + +function DropdownMenuTrigger({ ...props }: MenuPrimitive.Trigger.Props) { + return ; +} + +function DropdownMenuContent({ + align = "start", + alignOffset = 0, + side = "bottom", + sideOffset = 4, + className, + ...props +}: MenuPrimitive.Popup.Props & Pick) { + return ( + + + + + + ); +} + +function DropdownMenuGroup({ ...props }: MenuPrimitive.Group.Props) { + return ; +} + +function DropdownMenuLabel({ + className, + inset, + ...props +}: MenuPrimitive.GroupLabel.Props & { + inset?: boolean; +}) { + return ( + + ); +} + +function DropdownMenuItem({ + className, + inset, + variant = "default", + ...props +}: MenuPrimitive.Item.Props & { + inset?: boolean; + variant?: "default" | "destructive"; +}) { + return ( + + ); +} + +function DropdownMenuSub({ ...props }: MenuPrimitive.SubmenuRoot.Props) { + return ; +} + +function DropdownMenuSubTrigger({ + className, + inset, + children, + ...props +}: MenuPrimitive.SubmenuTrigger.Props & { + inset?: boolean; +}) { + return ( + + {children} + + + ); +} + +function DropdownMenuSubContent({ + align = "start", + alignOffset = -3, + side = "right", + sideOffset = 0, + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function DropdownMenuCheckboxItem({ + className, + children, + checked, + inset, + ...props +}: MenuPrimitive.CheckboxItem.Props & { + inset?: boolean; +}) { + return ( + + + + + + + {children} + + ); +} + +function DropdownMenuRadioGroup({ ...props }: MenuPrimitive.RadioGroup.Props) { + return ; +} + +function DropdownMenuRadioItem({ + className, + children, + inset, + ...props +}: MenuPrimitive.RadioItem.Props & { + inset?: boolean; +}) { + return ( + + + + + + + {children} + + ); +} + +function DropdownMenuSeparator({ className, ...props }: MenuPrimitive.Separator.Props) { + return ( + + ); +} + +function DropdownMenuShortcut({ className, ...props }: React.ComponentProps<"span">) { + return ( + + ); +} + +export { + DropdownMenu, + DropdownMenuPortal, + DropdownMenuTrigger, + DropdownMenuContent, + DropdownMenuGroup, + DropdownMenuLabel, + DropdownMenuItem, + DropdownMenuCheckboxItem, + DropdownMenuRadioGroup, + DropdownMenuRadioItem, + DropdownMenuSeparator, + DropdownMenuShortcut, + DropdownMenuSub, + DropdownMenuSubTrigger, + DropdownMenuSubContent, +}; diff --git a/ui/litellm-dashboard/src/components/ui/input-group.tsx b/ui/litellm-dashboard/src/components/ui/input-group.tsx new file mode 100644 index 00000000000..8ee9b7f17bd --- /dev/null +++ b/ui/litellm-dashboard/src/components/ui/input-group.tsx @@ -0,0 +1,140 @@ +"use client"; + +import * as React from "react"; +import { type VariantProps } from "cva"; + +import { cn, cva } from "@/lib/cva.config"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Textarea } from "@/components/ui/textarea"; + +function InputGroup({ className, ...props }: React.ComponentProps<"div">) { + return ( +
[data-align=block-end]]:h-auto has-[>[data-align=block-end]]:flex-col has-[>[data-align=block-start]]:h-auto has-[>[data-align=block-start]]:flex-col has-[>textarea]:h-auto dark:bg-input/30 dark:has-[[data-slot][aria-invalid=true]]:ring-destructive/40 has-[>[data-align=block-end]]:[&>input]:pt-3 has-[>[data-align=block-start]]:[&>input]:pb-3 has-[>[data-align=inline-end]]:[&>input]:pr-1.5 has-[>[data-align=inline-start]]:[&>input]:pl-1.5", + className, + )} + {...props} + /> + ); +} + +const inputGroupAddonVariants = cva({ + base: "flex h-auto cursor-text items-center justify-center gap-2 py-1.5 text-sm font-medium text-muted-foreground select-none group-data-[disabled=true]/input-group:opacity-50 [&>kbd]:rounded-[calc(var(--radius)-5px)] [&>svg:not([class*='size-'])]:size-4", + variants: { + align: { + "inline-start": "order-first pl-2 has-[>button]:-ml-1 has-[>kbd]:ml-[-0.15rem]", + "inline-end": "order-last pr-2 has-[>button]:-mr-1 has-[>kbd]:mr-[-0.15rem]", + "block-start": + "order-first w-full justify-start px-2.5 pt-2 group-has-[>input]/input-group:pt-2 [.border-b]:pb-2", + "block-end": "order-last w-full justify-start px-2.5 pb-2 group-has-[>input]/input-group:pb-2 [.border-t]:pt-2", + }, + }, + defaultVariants: { + align: "inline-start", + }, +}); + +function InputGroupAddon({ + className, + align = "inline-start", + ...props +}: React.ComponentProps<"div"> & VariantProps) { + return ( +
{ + if ((e.target as HTMLElement).closest("button")) { + return; + } + e.currentTarget.parentElement?.querySelector("input")?.focus(); + }} + {...props} + /> + ); +} + +const inputGroupButtonVariants = cva({ + base: "flex items-center gap-2 text-sm shadow-none", + variants: { + size: { + xs: "h-6 gap-1 rounded-[calc(var(--radius)-5px)] px-1.5 [&>svg:not([class*='size-'])]:size-3.5", + sm: "", + "icon-xs": "size-6 rounded-[calc(var(--radius)-5px)] p-0 has-[>svg]:p-0", + "icon-sm": "size-8 p-0 has-[>svg]:p-0", + }, + }, + defaultVariants: { + size: "xs", + }, +}); + +const InputGroupButton = React.forwardRef< + React.ComponentRef, + Omit, "size" | "type"> & + VariantProps & { + type?: "button" | "submit" | "reset"; + } +>(({ className, type = "button", variant = "ghost", size = "xs", ...props }, ref) => { + return ( +