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/.github/actions/setup-uv-with-retries/action.yml b/.github/actions/setup-uv-with-retries/action.yml new file mode 100644 index 00000000000..1627038dc3d --- /dev/null +++ b/.github/actions/setup-uv-with-retries/action.yml @@ -0,0 +1,47 @@ +name: "Set up uv with retries" +description: >- + Install uv via astral-sh/setup-uv, retrying on transient failures. Even with + an exact pinned version, the action resolves the artifact URL by fetching + https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a + single request with no retry, timeout, or fallback, so one connection-level + network error ("fetch failed") fails the whole job before any test runs. + Retrying the full step covers the manifest fetch and the binary download. + +inputs: + version: + description: "uv version to install" + required: true + +runs: + using: composite + steps: + - name: Set up uv (attempt 1) + id: attempt-1 + continue-on-error: true + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + version: ${{ inputs.version }} + + - name: Wait before attempt 2 + if: steps.attempt-1.outcome == 'failure' + shell: bash + run: sleep 15 + + - name: Set up uv (attempt 2) + id: attempt-2 + if: steps.attempt-1.outcome == 'failure' + continue-on-error: true + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + version: ${{ inputs.version }} + + - name: Wait before attempt 3 + if: steps.attempt-2.outcome == 'failure' + shell: bash + run: sleep 30 + + - name: Set up uv (attempt 3) + if: steps.attempt-2.outcome == 'failure' + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + version: ${{ inputs.version }} diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 9fd81b27f3b..92230fc8892 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -63,7 +63,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/auto_update_price_and_context_window.yml b/.github/workflows/auto_update_price_and_context_window.yml index da96ccffc93..5e2820356a0 100644 --- a/.github/workflows/auto_update_price_and_context_window.yml +++ b/.github/workflows/auto_update_price_and_context_window.yml @@ -18,7 +18,7 @@ jobs: with: persist-credentials: false - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" - name: Create Pull Request diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index 439126aa1ee..9c24bad00f1 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -31,7 +31,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index a1772102b89..54a8e53d7a3 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -37,7 +37,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/mutation-test.yml b/.github/workflows/mutation-test.yml index 183f12f969c..6684952b998 100644 --- a/.github/workflows/mutation-test.yml +++ b/.github/workflows/mutation-test.yml @@ -39,7 +39,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/oss_daily_guardrails.yml b/.github/workflows/oss_daily_guardrails.yml index 950c51c9b60..f9dc746ee05 100644 --- a/.github/workflows/oss_daily_guardrails.yml +++ b/.github/workflows/oss_daily_guardrails.yml @@ -35,7 +35,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 872a1799d98..9d28ca211cf 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -38,7 +38,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 0b5b0e9b976..09406d77634 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -33,7 +33,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" @@ -172,7 +172,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index 5b5290880c1..a5a4e722133 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -32,7 +32,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-semgrep.yml b/.github/workflows/test-semgrep.yml index f0dcb9887be..6e9f5e42fa2 100644 --- a/.github/workflows/test-semgrep.yml +++ b/.github/workflows/test-semgrep.yml @@ -31,7 +31,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-terraform-provider.yml b/.github/workflows/test-terraform-provider.yml index 03d8ff3461c..058a2538c15 100644 --- a/.github/workflows/test-terraform-provider.yml +++ b/.github/workflows/test-terraform-provider.yml @@ -74,7 +74,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-unit-documentation.yml b/.github/workflows/test-unit-documentation.yml index 03f9f0a510b..c12a289ce9f 100644 --- a/.github/workflows/test-unit-documentation.yml +++ b/.github/workflows/test-unit-documentation.yml @@ -42,7 +42,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-unit-proxy-legacy.yml b/.github/workflows/test-unit-proxy-legacy.yml index 0068e80e584..bcbf365babf 100644 --- a/.github/workflows/test-unit-proxy-legacy.yml +++ b/.github/workflows/test-unit-proxy-legacy.yml @@ -59,7 +59,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 85ccbef752f..97571a4576d 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.49" +version = "0.1.50" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.49" +version = "0.1.50" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", 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-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index a54db2ace65..b67d9d8570a 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.76" +version = "0.4.77" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.76" +version = "0.4.77" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index c8bfcabc64e..856556f7c56 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -515,6 +515,22 @@ class CustomGuardrail(CustomLogger): return True return False + def uses_apply_guardrail_interface(self) -> bool: + return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail + + def _deployment_pre_call_target(self) -> "CustomLogger": + if not self.uses_apply_guardrail_interface(): + return self + try: + from litellm.proxy.utils import unified_guardrail + except ImportError as e: + raise ImportError( + f"Guardrail {self.guardrail_name or type(self).__name__} implements apply_guardrail, which needs " + "the litellm proxy dependencies to run at the deployment level. " + "Install them with: pip install 'litellm[proxy]'" + ) from e + return unified_guardrail + async def async_pre_call_deployment_hook( self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] ) -> Optional[dict]: @@ -533,7 +549,10 @@ class CustomGuardrail(CustomLogger): # CHECK IF GUARDRAIL REJECTS THE REQUEST if call_type == CallTypes.completion or call_type == CallTypes.acompletion: - result = await self.async_pre_call_hook( + target = self._deployment_pre_call_target() + if target is not self: + kwargs["guardrail_to_apply"] = self + result = await target.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( user_id=kwargs.get("user_api_key_user_id"), team_id=kwargs.get("user_api_key_team_id"), @@ -543,7 +562,7 @@ class CustomGuardrail(CustomLogger): ), cache=dc, data=kwargs, - call_type=call_type.value or "acompletion", # type: ignore + call_type="completion" if call_type == CallTypes.completion else "acompletion", ) if result is not None and isinstance(result, dict): 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/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 1f5b76f3d0a..0ec1f3eae13 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1441,24 +1441,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): output_key=param, ) elif param == "response_format" and isinstance(value, dict): - if any( - substring in model - for substring in { - "sonnet-4.5", - "sonnet-4-5", - "opus-4.1", - "opus-4-1", - "opus-4.5", - "opus-4-5", - "opus-4.6", - "opus-4-6", - "opus-4.7", - "opus-4-7", - "sonnet-4.6", - "sonnet-4-6", - "sonnet_4.6", - "sonnet_4_6", - } + if AnthropicConfig._supports_model_capability( + model, + "supports_native_structured_output", + self._resolved_provider, ): _output_format = self.map_response_format_to_anthropic_output_format(value) if _output_format is not None: diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 0bcf34a45d6..e006662ec4d 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -340,11 +340,15 @@ class AnthropicModelInfo(BaseLLMModelInfo): def _get_model_capability(model: str, key: str) -> Optional[bool]: """Read boolean capability ``key`` from the model map, or None when no entry declares it.""" + from litellm.utils import _get_bundled_model_cost_map + try: - for cand in AnthropicModelInfo._model_map_lookup_candidates(model): - value = litellm.model_cost.get(cand, {}).get(key) - if isinstance(value, bool): - return value + candidates = AnthropicModelInfo._model_map_lookup_candidates(model) + for model_cost in (litellm.model_cost, _get_bundled_model_cost_map()): + for cand in candidates: + value = model_cost.get(cand, {}).get(key) + if isinstance(value, bool): + return value except Exception: pass return None diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 00941587753..05679bf39ab 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -375,6 +375,36 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): else: optional_params.pop("output_config", None) + @staticmethod + def _drop_incompatible_temperature_for_thinking( + model: str, optional_params: dict, custom_llm_provider: str + ) -> None: + """Anthropic rejects any ``temperature`` other than 1 while extended thinking + is enabled ("temperature may only be set to 1 when thinking is enabled"). + + Clients like Claude Code send ``thinking``/``output_config.effort`` together + with a pinned ``temperature`` (e.g. the safety classifier uses ``temperature=0`` + for determinism). When the request lands on a non-adaptive model, the effort + interface is reshaped above into legacy ``thinking={type: enabled}`` (or kept + as ``output_config.effort`` on Opus 4.5), and the leftover ``temperature`` would + 400. Preserving the thinking the caller asked for wins over an unhonorable + sampling value (Anthropic forces ``temperature=1`` under thinking regardless), + so drop it and let the API default apply. + + Adaptive models (4.6+) own this natively and are left untouched. + """ + if AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider): + return + temperature = optional_params.get("temperature") + if temperature is None or temperature == 1: + return + thinking = optional_params.get("thinking") + output_config = optional_params.get("output_config") + thinking_enabled = isinstance(thinking, dict) and thinking.get("type") == "enabled" + effort_enabled = isinstance(output_config, dict) and output_config.get("effort") is not None + if thinking_enabled or effort_enabled: + optional_params.pop("temperature", None) + def transform_anthropic_messages_request( self, model: str, @@ -415,6 +445,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): custom_llm_provider=self._resolved_provider, ) + self._drop_incompatible_temperature_for_thinking( + model=model, + optional_params=anthropic_messages_optional_request_params, + custom_llm_provider=self._resolved_provider, + ) + system_param = anthropic_messages_optional_request_params.get("system") if self.should_strip_billing_metadata() and system_param is not None: filtered_system = self._filter_billing_headers_from_system(system_param) 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/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a6f59abc140..c99347e189c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11331,6 +11331,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11362,6 +11363,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true @@ -11424,6 +11426,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, @@ -11479,6 +11482,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11506,6 +11510,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11559,6 +11564,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true @@ -11586,6 +11592,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true @@ -11614,6 +11621,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "provider_specific_entry": { @@ -11648,6 +11656,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "provider_specific_entry": { @@ -11682,6 +11691,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, @@ -11718,6 +11728,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, @@ -11788,6 +11799,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, 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/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py new file mode 100644 index 00000000000..19048e2eb7c --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -0,0 +1,694 @@ +"""Bridge token flow: litellm identity resolution and the DCR-bridge oauth_delegate mint/refresh pipeline.""" + +import math +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Literal, Optional + +from fastapi import HTTPException, Request +from fastapi.responses import JSONResponse +from pydantic import SecretStr +from typing_extensions import assert_never + +from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +if TYPE_CHECKING: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _BridgeAuthorizationCode + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, + EnvelopeKeys, + RefreshCredential, + UpstreamTokenGrant, + ) + from litellm.proxy._types import UserAPIKeyAuth + + +def _litellm_key_from_request(request: Request) -> Optional[str]: + """Return the LiteLLM API key presented on the request, or ``None``. + + Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code + send) as well as ``Authorization``; either may carry a bare token or ``Bearer ``. + ``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry + an OAuth/upstream bearer. + """ + for header_value in ( + request.headers.get("x-litellm-api-key"), + request.headers.get("Authorization") or request.headers.get("authorization"), + ): + if not header_value: + continue + value = header_value.strip() + if value.lower().startswith("bearer "): + value = value[7:].strip() + if value: + return value + return None + + +def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool: + """``True`` when the presented key is neither blocked nor past its expiry. + + The OAuth token endpoint is unauthenticated, so the presented key is validated here before it is + trusted; a revoked or expired key must not mint a bridge envelope or write a stored credential. + ``get_key_object`` resolves a row without these checks (the main ``user_api_key_auth`` pipeline + enforces them downstream, which this endpoint bypasses), so they are applied here. Deleted keys + are already rejected upstream, where ``get_key_object`` raises on a row that no longer exists. + + This is an active-state gate only; it deliberately does not require a ``user_id``. A valid + team-scoped or service-account key has no ``user_id`` yet is a legitimate credential, so gating + on ``user_id`` presence would wrongly reject it. Callers that need the user (the per-user token + store) derive it separately via :func:`_active_key_user_id`. + + Total by design: ``expires`` is typed ``str | datetime``, and an unparseable string would make + ``datetime.fromisoformat`` raise. Since the callers run this outside their key-resolution + ``try``, an uncaught parse error would surface as a 500 instead of the endpoint's fail-closed + behavior, so a malformed expiry is treated as inactive (return ``False``) rather than raising. + """ + if key_obj.blocked is True: + return False + expires = key_obj.expires + if expires is not None: + if isinstance(expires, datetime): + expiry = expires + else: + try: + expiry = datetime.fromisoformat(expires) + except (ValueError, TypeError): + return False + if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None: + expiry = expiry.replace(tzinfo=timezone.utc) + if expiry < datetime.now(timezone.utc): + return False + return True + + +def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> str | None: + """The active key's ``user_id``, or ``None`` when the key is blocked/expired or simply has no + ``user_id`` (a team-scoped or service-account key). Used only by the per-user token store, which + needs a user to key the stored credential; the bridge mint uses the key hash and does not.""" + return key_obj.user_id if _key_is_active(key_obj) else None + + +@dataclass(frozen=True, slots=True) +class _ResolvedKey: + """An active litellm key resolved from the token request: its hash (the value ``get_key_object`` + and the cache/DB layer key the record by) and the live record.""" + + key_hash: str + key: "UserAPIKeyAuth" + + +_KeyResolutionFailure = Literal["no_active_key", "unavailable", "unresolvable"] +"""Why a token request yielded no active litellm key, kept distinct so a caller statuses each truthfully +instead of blaming the client for a gateway problem: +- ``no_active_key``: none was presented, or the presented key is unknown / blocked / expired (the + caller's request is at fault) +- ``unavailable``: the auth database was transiently unreachable while resolving (retryable) +- ``unresolvable``: the gateway cannot resolve identity right now (no DB connection, or an unexpected + error) -- a gateway fault, not the caller's +The classification mirrors admission's ``_reload_admitted_key`` so the mint (ingress) and admission +(egress) never disagree on the status of the same outage.""" + + +async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyResolutionFailure": + """Resolve the presented litellm key to an active key record, or say precisely why not. + + Single resolution path the OAuth token endpoint reuses, resolving authoritatively via + ``get_key_object`` (cache first, then DB). The failure is a value, not a bare ``None``, so a caller + can tell "the client sent no usable credential" (a request error) apart from "the gateway could not + check" (an infrastructure error) and status each truthfully; collapsing both to ``None`` is what let + a DB outage read as a 400. A resolved key is still gated by ``_key_is_active``, so a blocked or + expired key is ``no_active_key`` while a valid team-scoped or service-account key (no ``user_id``) + resolves. 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.""" + token = _litellm_key_from_request(request) + if not token: + return "no_active_key" + 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, + ) + 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: + key_obj = await get_key_object( + hashed_token=key_hash, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except (ProxyException, HTTPException): + return "no_active_key" + except Exception as exc: # noqa: BLE001 # classify: a DB outage is retryable, anything else is an opaque gateway fault + if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc): + return "unavailable" + verbose_logger.debug( + "_reload_active_key_by_hash: unexpected key-resolution error (%s)", + type(exc).__name__, + ) + return "unresolvable" + if not _key_is_active(key_obj): + return "no_active_key" + 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 + (including a transient DB outage) collapses to ``None`` here and the caller simply skips the store; + the bridge mint, which must status those outcomes differently, consumes + :func:`_resolve_active_litellm_key` directly.""" + resolved = await _resolve_active_litellm_key(request) + if not isinstance(resolved, _ResolvedKey): + return None + return _active_key_user_id(resolved.key) + + +_UpstreamGrantRejection = Literal["no_access_token", "expired_lifetime"] +"""Why an upstream token response cannot back a bridge envelope: +- ``no_access_token``: the response carries no usable ``access_token`` +- ``expired_lifetime``: the response reports a parseable, non-positive ``expires_in``, i.e. an upstream + token that is already dead, so sealing it would forward a bearer the edge cannot use +An absent or unparseable ``expires_in`` is NOT a rejection; the lifetime is merely unknown and the +envelope caps it, the by-design behaviour for an upstream that omits the field.""" + + +def _classify_upstream_lifetime(raw_expires_in: object) -> "int | Literal['unspecified', 'expired']": + """Classify an upstream ``expires_in`` into a positive number of seconds, ``"unspecified"`` (absent + or unparseable, so the envelope caps it), or ``"expired"`` (a non-positive value the upstream reports + as already elapsed). Telling "we do not know the lifetime" apart from "the upstream says it is + already dead" is what stops an explicitly-expired token from silently receiving the envelope's 1h + cap. The expired decision is made on the parsed numeric value, not on ``int(...)`` of it, so a + positive sub-second lifetime in ``(0, 1)`` is not truncated to ``0`` and misread as elapsed; the + envelope works in whole seconds, so such a lifetime clamps up to its 1s floor. ``bool`` is excluded + (an ``int`` subclass but never a real lifetime), and the conversions can raise on ``NaN`` / + ``Infinity`` / oversized input, which reads as unparseable rather than surfacing as a 500.""" + if raw_expires_in is None or isinstance(raw_expires_in, bool) or not isinstance(raw_expires_in, (int, float, str)): + return "unspecified" + try: + numeric = float(raw_expires_in) + seconds = int(numeric) + except (ValueError, TypeError, OverflowError): + return "unspecified" + if numeric <= 0: + return "expired" + return max(1, seconds) + + +def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenGrant | _UpstreamGrantRejection": + """Validate an upstream OAuth token response into a typed grant, or say why it cannot back an + envelope. Each field is isinstance-checked so nothing untyped from ``response.json()`` reaches the + grant. ``expires_in`` is read three ways (see :func:`_classify_upstream_lifetime`): an unknown + lifetime leaves the grant ``expires_in`` ``None`` for the envelope to cap, a positive value is + honoured, and an explicit already-elapsed value is a rejection rather than a silent fall-through to + the cap.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + UpstreamTokenGrant, + ) + + if not isinstance(token_response, dict): + return "no_access_token" + access = token_response.get("access_token") + if not isinstance(access, str) or not access: + return "no_access_token" + lifetime = _classify_upstream_lifetime(token_response.get("expires_in")) + if lifetime == "expired": + return "expired_lifetime" + token_type = token_response.get("token_type") + scope = token_response.get("scope") + return UpstreamTokenGrant( + access_token=SecretStr(access), + token_type=token_type if isinstance(token_type, str) and token_type else "Bearer", + # The upstream refresh_token is deliberately NOT sealed: the edge never consumes it (it forwards + # only token_type + access_token), so it would be dead weight embedding a long-lived upstream + # credential in the client-held bearer, and it enlarges the envelope. Refresh support is a + # follow-up (a dedicated refresh-envelope); the client re-runs authorization_code at the cap. + refresh_token=None, + scope=scope if isinstance(scope, str) and scope else None, + expires_in=lifetime if isinstance(lifetime, int) else None, + ) + + +# --------------------------------------------------------------------------- +# DCR-bridge oauth_delegate mint: a three-phase pipeline whose failures are values. +# +# prepare (before the upstream exchange) -> validate every precondition and resolve identity+keys +# exchange (the single-use upstream code is consumed here, in exchange_token_with_server) +# finish (after the exchange) -> seal the upstream grant into the client-held envelope +# +# Every precondition lives in ``prepare``, which runs BEFORE the exchange, so no failure can burn the +# single-use code or rotate a refresh token, for either grant type -- that whole class of bug is gone +# by construction rather than guarded case by case. Failures are values mapped to an OAuth-shaped +# response in one place (``_bridge_mint_error_response``), so status codes and the RFC 6749 §5.2 body +# shape are uniform. Adding a failure mode is a new literal plus a match arm the type checker forces. +# --------------------------------------------------------------------------- + +_BridgeMintError = Literal[ + "no_identity", + "invalid_refresh", + "identity_unavailable", + "identity_unresolvable", + "not_configured", + "no_upstream_token", + "upstream_token_expired", + "too_large", +] + + +@dataclass(frozen=True, slots=True) +class _BridgeMintReady: + """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.""" + + identity: "EnvelopeIdentity" + keys: "EnvelopeKeys" + + +def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse: + """Map a bridge-mint failure value to its token-endpoint response: one place, RFC 6749 §5.2 shape + (top-level ``error``, no-store headers) for every case, with a status truthful about where the + failure is. The caller's request is 400, a transient gateway outage is 503, a gateway + misconfiguration is 500, and an upstream problem is 502. The identity-resolution statuses match how + admission statuses the same conditions on the egress side, so mint and admit never disagree under + one outage.""" + match error: + case "no_identity": + status, code, desc = ( + 400, + "invalid_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 "invalid_refresh": + status, code, desc = ( + 400, + "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 = ( + 503, + "temporarily_unavailable", + "the authentication database is temporarily unreachable; retry shortly", + ) + case "identity_unresolvable": + status, code, desc = ( + 500, + "server_error", + "the gateway could not resolve the litellm identity for this request", + ) + case "not_configured": + status, code, desc = ( + 500, + "server_error", + "the gateway is not configured to mint a gateway-bound credential (master_key is not set)", + ) + case "no_upstream_token": + status, code, desc = ( + 502, + "server_error", + "the upstream token response has no usable access_token", + ) + case "upstream_token_expired": + status, code, desc = ( + 502, + "server_error", + "the upstream token response reports an already-expired lifetime", + ) + case "too_large": + status, code, desc = ( + 502, + "server_error", + "the upstream token is too large to seal into a gateway-bound credential", + ) + case _: + assert_never(error) + return JSONResponse( + status_code=status, content={"error": code, "error_description": desc}, headers=TOKEN_NO_CACHE_HEADERS + ) + + +def _key_resolution_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError: + """Lift an identity-resolution failure into the mint taxonomy, preserving origin so the status stays + truthful: the caller's missing credential is 400, a transient DB outage is 503, and a gateway that + cannot resolve identity is 500.""" + match failure: + case "no_active_key": + return "no_identity" + case "unavailable": + return "identity_unavailable" + case "unresolvable": + return "identity_unresolvable" + case _: + assert_never(failure) + + +def _upstream_rejection_to_mint_error(rejection: _UpstreamGrantRejection) -> _BridgeMintError: + """Lift an upstream-response rejection into the mint taxonomy; both are upstream faults (502).""" + match rejection: + case "no_access_token": + return "no_upstream_token" + case "expired_lifetime": + return "upstream_token_expired" + case _: + assert_never(rejection) + + +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 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) + 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 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 + SealedEnvelope, + UpstreamTokenGrant, + ) + + grant = _bridge_grant_from_token_response(token_response) + if not isinstance(grant, UpstreamTokenGrant): + return _upstream_rejection_to_mint_error(grant) + 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())) + 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 _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 diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ebceb320906..54aff86aab2 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -10,7 +10,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, ValidationError +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -21,6 +21,24 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( TokenEndpointAuthConfigError, build_token_endpoint_client_auth, ) +from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _bridge_mint_error_response, + _BridgeMintReady, + _BridgeRefreshReady, + _extract_user_id_from_request, + _finish_bridge_mint, + _prepare_bridge_mint, + _prepare_bridge_refresh, +) +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, @@ -37,7 +55,7 @@ from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer if TYPE_CHECKING: - from litellm.proxy._types import LiteLLM_MCPServerTable, UserAPIKeyAuth + from litellm.proxy._types import LiteLLM_MCPServerTable # TTL cache for upstream OAuth metadata fetched from pass-through MCP servers. # Keeps us from hammering the upstream IdP on each discovery request. @@ -91,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. @@ -101,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 @@ -111,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) @@ -138,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 @@ -304,90 +393,6 @@ def _validate_token_response( ) -def _litellm_key_from_request(request: Request) -> Optional[str]: - """Return the LiteLLM API key presented on the request, or ``None``. - - Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code - send) as well as ``Authorization``; either may carry a bare token or ``Bearer ``. - ``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry - an OAuth/upstream bearer. - """ - for header_value in ( - request.headers.get("x-litellm-api-key"), - request.headers.get("Authorization") or request.headers.get("authorization"), - ): - if not header_value: - continue - value = header_value.strip() - if value.lower().startswith("bearer "): - value = value[7:].strip() - if value: - return value - return None - - -def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> Optional[str]: - """The key's ``user_id``, or ``None`` if the key is blocked or expired. - - The OAuth token endpoint is unauthenticated, so the presented key is validated here before its - identity is trusted to key a stored credential; a revoked or expired key must not be able to - write or overwrite the per-user OAuth token. ``get_key_object`` resolves a row without these - checks (the main ``user_api_key_auth`` pipeline enforces them downstream, which this endpoint - bypasses), so they are applied here. Deleted keys are already rejected upstream, where - ``get_key_object`` raises on a row that no longer exists. - """ - if key_obj.blocked is True: - return None - expires = key_obj.expires - if expires is not None: - expiry = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires) - if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None: - expiry = expiry.replace(tzinfo=timezone.utc) - if expiry < datetime.now(timezone.utc): - return None - return key_obj.user_id - - -async def _extract_user_id_from_request(request: Request) -> Optional[str]: - """Resolve the LiteLLM ``user_id`` at the OAuth token endpoint so a per-user token is stored - under the same identity the egress later reads it by (``user_api_key_auth.user_id``). - - Resolves authoritatively via ``get_key_object`` (cache first, then DB) instead of a raw cache - peek. On a multi-replica gateway the token-exchange request can land on a worker whose in-memory - cache never saw the key, and a cross-replica Redis hit deserializes to a plain ``dict`` rather - than a ``UserAPIKeyAuth``; the previous code read only ``Authorization`` and did - ``getattr(cached, "user_id")`` with no ``model_type`` rehydration and no DB fallback, so it - silently returned ``None`` and the token was never persisted, which makes the egress 401 on every - reconnect. The resolved key is validated (``_active_key_user_id``) before its identity is trusted, - so a blocked or expired key cannot write. Returns ``None`` when no key is present, the key cannot - be resolved, or it is blocked/expired. - """ - token = _litellm_key_from_request(request) - if not token: - return None - try: - from litellm.proxy._types import hash_token # noqa: PLC0415 - from litellm.proxy.auth.auth_checks import get_key_object # noqa: PLC0415 - from litellm.proxy.proxy_server import ( # noqa: PLC0415 - prisma_client, - user_api_key_cache, - ) - - key_obj = await get_key_object( - hashed_token=hash_token(token), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - return _active_key_user_id(key_obj) - except Exception as exc: - verbose_logger.debug( - "_extract_user_id_from_request: could not resolve a LiteLLM user_id for the presented " - "key (%s); per-user token will not be stored server-side.", - type(exc).__name__, - ) - return None - - async def _store_per_user_token_server_side( server: MCPServer, user_id: str, @@ -620,12 +625,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) @@ -654,6 +678,13 @@ async def authorize_with_server( return response +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" + + async def exchange_token_with_server( request: Request, mcp_server: MCPServer, @@ -688,25 +719,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( @@ -726,32 +793,49 @@ async def exchange_token_with_server( } if code_verifier: token_data["code_verifier"] = code_verifier - + # 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() - access_token = token_response["access_token"] # Validate token response against server-configured rules before any storage. # This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc. @@ -791,8 +875,23 @@ async def exchange_token_with_server( mcp_server.server_id, ) + # A DCR-bridge oauth_delegate server hands the client a gateway-bound envelope (identity plus the + # 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": access_token, + "access_token": raw_access_token, "token_type": token_response.get("token_type", "Bearer"), } @@ -1048,21 +1147,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, @@ -1122,19 +1206,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() @@ -1362,7 +1451,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/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5090aa7d7d5..68a61b85175 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3719,8 +3719,15 @@ if MCP_AVAILABLE: headers={"www-authenticate": upstream_www_authenticate}, ) + def _get_authorization_header_from_scope(scope: Scope) -> Optional[str]: + """First ``Authorization`` header value in the ASGI scope, or None.""" + for key, value in scope.get("headers", []): + if key.lower() == b"authorization": + return value.decode("latin-1") + return None + def _scope_has_authorization_header(scope: Scope) -> bool: - return any(key.lower() == b"authorization" for key, _ in scope.get("headers", [])) + return _get_authorization_header_from_scope(scope) is not None def _get_forwarded_auth_from_scope(scope: Scope) -> Optional[str]: """Return the upstream-bound ``Authorization`` header value, or None. @@ -3733,17 +3740,24 @@ if MCP_AVAILABLE: ``MCPRequestHandler.process_mcp_request``), and forwarding it upstream would leak the proxy key to a third-party MCP server. """ - authorization = None - has_litellm_key_header = False - for key, value in scope.get("headers", []): - key_lower = key.lower() - if key_lower == b"authorization": - authorization = value.decode("latin-1") - elif key_lower == b"x-litellm-api-key": - has_litellm_key_header = True + has_litellm_key_header = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", [])) if not has_litellm_key_header: return None - return authorization + return _get_authorization_header_from_scope(scope) + + def _is_delegate_upstream_probe_target(server: MCPServer) -> bool: + """Whether ``server`` is an interactive delegate-auth server whose client-supplied + token should be preflighted upstream. + + Mirrors the anonymous-delegate gate in ``get_allowed_mcp_servers``: the flow is + resolved via ``effective_oauth2_flow`` so an unstamped M2M-shape row fails closed + (its stored client credentials drive egress; the caller's bearer is irrelevant). + """ + return ( + server.auth_type == MCPAuth.oauth2 + and server.delegate_auth_to_upstream is True + and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" + ) async def _probe_upstream_auth( url: str, @@ -3805,7 +3819,7 @@ if MCP_AVAILABLE: mcp_servers: Optional[List[str]], client_ip: Optional[str], ) -> None: - """Probe pass-through upstream servers in parallel before the MCP session starts. + """Probe pass-through and delegate-auth upstream servers in parallel before the MCP session starts. Only servers the caller's key is already authorized to reach are probed — the list is derived from _get_allowed_mcp_servers so that a user cannot @@ -3813,11 +3827,42 @@ if MCP_AVAILABLE: The MCP SDK commits HTTP 200 headers before invoking handlers, so a 401 can only be returned before that point. This function raises HTTPException(401) - with a WWW-Authenticate header if any upstream rejects the client token. + with a WWW-Authenticate header if any upstream rejects the client token, or 403 + if the upstream accepts it but forbids the caller. Fails-open: network errors are logged and the request is allowed through. + + Delegate-auth servers (``auth_type=oauth2`` + ``delegate_auth_to_upstream``) + are probed with the caller's bare ``Authorization`` bearer. That bearer is only + an upstream token (never a LiteLLM key) when admission took the delegate bypass, + so the delegate target is resolved through ``get_mcp_server_by_name`` -- the same + resolver admission used -- rather than the wider allowed-server prefix/access-group + matching. A name that only reaches a delegate server via server_id or an access + group would have been admitted as a real LiteLLM key, so probing it would leak that + key upstream; requiring the admission-resolver match closes that gap. Without the + probe a rejected token is absorbed by the tools/list handler and masked as an empty + tool list. Gated to single-server routes so one rejected token cannot 401 a + multi-server aggregate connect, matching the OBO preflight gating; the challenge + echoes the requested name so aliased routes get the same resource_metadata URL as + the tokenless preemptive challenge. """ forwarded_auth = _get_forwarded_auth_from_scope(scope) - if not forwarded_auth: + requested_single_target = mcp_servers[0] if mcp_servers is not None and len(mcp_servers) == 1 else None + # The bare Authorization header (no x-litellm-api-key) is a valid upstream token + # only when admission classified it as one, i.e. the single requested name resolves + # to a delegate server under admission's own resolver. Resolve it the same way here + # so a server_id- or access-group-named delegate (which admission would have treated + # as a LiteLLM key) is never probed with that key. + delegate_server = ( + global_mcp_server_manager.get_mcp_server_by_name(requested_single_target, client_ip=client_ip) + if requested_single_target + else None + ) + delegate_auth = ( + _get_authorization_header_from_scope(scope) + if delegate_server is not None and _is_delegate_upstream_probe_target(delegate_server) + else None + ) + if not forwarded_auth and not delegate_auth: return # Use the authorized server set, not the raw user-supplied names, so that @@ -3827,33 +3872,49 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, client_ip=client_ip, ) - passthrough_servers = [ - srv - for srv in allowed_servers - # Restrict to genuine OAuth pass-through servers (auth_type none + - # Authorization in extra_headers). Gateway-managed OAuth2 servers - # must not receive the ``resource_metadata=`` challenge emitted - # below — they require ``authorization_uri=`` pointing at the - # gateway AS metadata. ``is_oauth_passthrough`` already requires - # ``auth_type in (None, MCPAuth.none)``, which is mutually - # exclusive with ``has_client_credentials`` (oauth2 + M2M flow), - # so M2M servers are implicitly excluded here. - if srv.is_oauth_passthrough - ] - if not passthrough_servers: + passthrough_targets: Tuple[Tuple[MCPServer, str, str], ...] = ( + tuple( + (srv, forwarded_auth, srv.name) + for srv in allowed_servers + # Restrict to genuine OAuth pass-through servers (auth_type none + + # Authorization in extra_headers). Gateway-managed OAuth2 servers + # must not receive the ``resource_metadata=`` challenge emitted + # below — they require ``authorization_uri=`` pointing at the + # gateway AS metadata. ``is_oauth_passthrough`` already requires + # ``auth_type in (None, MCPAuth.none)``, which is mutually + # exclusive with ``has_client_credentials`` (oauth2 + M2M flow), + # so M2M servers are implicitly excluded here. + if srv.is_oauth_passthrough + ) + if forwarded_auth + else () + ) + # Probe the admission-resolved delegate server only when the caller is actually + # authorized for it (present in the IP-filtered allowed set), keyed by server_id. + delegate_targets: Tuple[Tuple[MCPServer, str, str], ...] = ( + tuple( + (srv, delegate_auth, requested_single_target) + for srv in allowed_servers + if delegate_server is not None and srv.server_id == delegate_server.server_id + ) + if delegate_auth and requested_single_target + else () + ) + probe_targets = passthrough_targets + delegate_targets + if not probe_targets: return probe_results = await asyncio.gather( - *[_probe_upstream_auth(srv.url or "", forwarded_auth) for srv in passthrough_servers] + *[_probe_upstream_auth(srv.url or "", auth_header) for srv, auth_header, _ in probe_targets] ) - for srv, (probe_status, _) in zip(passthrough_servers, probe_results): + for (srv, _, challenge_server_name), (probe_status, _) in zip(probe_targets, probe_results): if probe_status == 401: # Token is missing or expired: keep pass-through clients on the # protected-resource discovery flow so they re-authorize against # the upstream IdP metadata proxied by LiteLLM. www_authenticate = _get_passthrough_www_authenticate( scope=scope, - server_name=srv.name, + server_name=challenge_server_name, invalid_token=True, ) raise HTTPException( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 23fe7730c17..5e3ea4b7dcb 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 @@ -2421,6 +2422,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "is active as a reminder that hard enforcement is relaxed." ), ) + skip_user_budget_on_team_key: bool | None = Field( + None, + description=( + "If True, restores the legacy behavior where a user's personal " + "max_budget is NOT enforced when their key belongs to a team; only " + "the team (and team-member) budgets apply. Defaults to False, meaning " + "the user's personal max_budget is always enforced regardless of " + "whether the key belongs to a team (see GitHub issue #12905)." + ), + ) user_url_validation: Optional[bool] = Field( None, description=( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 93811812901..00f6d44e25a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -626,26 +626,29 @@ async def common_checks( ) async def _user_max_budget_check() -> None: - # 4.1 personal budget, if personal key - if ( - (team_object is None or team_object.team_id is None) - and user_object is not None - and user_object.max_budget is not None - ): - from litellm.proxy.proxy_server import get_current_spend + if user_object is None or user_object.max_budget is None: + return + skip_for_team = ( + general_settings.get("skip_user_budget_on_team_key") is True + and team_object is not None + and team_object.team_id is not None + ) + if skip_for_team: + return + from litellm.proxy.proxy_server import get_current_spend - user_budget = user_object.max_budget - user_spend = await get_current_spend( - counter_key=f"spend:user:{user_object.user_id}", - fallback_spend=user_object.spend or 0.0, + user_budget = user_object.max_budget + user_spend = await get_current_spend( + counter_key=f"spend:user:{user_object.user_id}", + fallback_spend=user_object.spend or 0.0, + max_budget=user_budget, + ) + if math.isfinite(user_budget) and user_spend >= user_budget: + raise litellm.BudgetExceededError( + current_cost=user_spend, max_budget=user_budget, + message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", ) - if math.isfinite(user_budget) and user_spend >= user_budget: - raise litellm.BudgetExceededError( - current_cost=user_spend, - max_budget=user_budget, - message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", - ) # Each scope reads a distinct counter key with no cross-scope ordering # dependency, so the per-scope Redis-first reads run concurrently instead diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2613510bd0c..744d8182715 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2442,6 +2442,7 @@ async def _reserve_budget_after_common_checks( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, + skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True, ) 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/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 6cdf49818e6..8861a9e3bbc 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -34,6 +34,12 @@ def is_text_content_call_type(call_type: str) -> bool: TEXT_PART_TYPES: FrozenSet[str] = frozenset({"text", "input_text", "output_text"}) +# Responses-API item types whose ``output`` field carries user/tool text +# that guardrails should inspect. ``function_call_output`` is the +# built-in shape; ``custom_tool_call_output`` is the custom-tool +# counterpart (see ``ChatCompletionCustomToolCallOutput``). +_OUTPUT_ITEM_TYPES: frozenset[str] = frozenset({"function_call_output", "custom_tool_call_output"}) + def _iter_text_parts_in_content(content: Any) -> Iterator[str]: """Yield text fragments from a ``message.content`` value (string or @@ -72,7 +78,7 @@ def _coerce_input_to_messages(input_value: Any) -> List[Dict[str, Any]]: messages.append({"role": item.get("role") or "user", "content": [item]}) elif "content" in item: messages.append({"role": item.get("role") or "user", "content": item["content"]}) - elif item.get("type") == "function_call_output" and "output" in item: + elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: messages.append({"role": item.get("role") or "tool", "content": item["output"]}) return messages @@ -157,7 +163,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: input_value[idx] = {**item, "text": visit(item["text"])} elif "content" in item: item["content"] = _rewrite_content(item["content"]) - elif item.get("type") == "function_call_output" and "output" in item: + elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: item["output"] = _rewrite_content(item["output"]) return visited diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 9c4cef2f06b..31ce0cc2214 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -772,7 +772,13 @@ class LassoGuardrail(CustomGuardrail): data: Request data (used for conversation_id generation and tools extraction) cache: Cache instance for storing conversation_id (optional for post-call) """ - payload: Dict[str, Any] = {"messages": messages, "messageType": message_type} + payload: Dict[str, Any] = { + "messages": messages, + "messageType": message_type, + # Drives the "Used By" badge on Lasso Application API Keys: every call from this + # integration is attributed as "litellm" on the keys list. + "source": {"type": "litellm"}, + } # Add optional parameters if available if self.user_id: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b6342f4fa1a..b839426fcda 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -28,11 +28,20 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( + CallTypes, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) from litellm.utils import get_end_user_id_for_cost_tracking +_PASS_THROUGH_CALL_TYPES: frozenset[str] = frozenset( + { + CallTypes.pass_through.value, + CallTypes.llm_passthrough_route.value, + CallTypes.allm_passthrough_route.value, + } +) + class _ProxyDBLogger(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -219,11 +228,13 @@ class _ProxyDBLogger(CustomLogger): verbose_proxy_logger.debug( f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}" ) + call_type: Optional[str] = kwargs.get("call_type") if _should_track_cost_callback( user_api_key=user_api_key, user_id=user_id, team_id=team_id, end_user_id=end_user_id, + call_type=call_type, ): ## UPDATE DATABASE await _update_database_and_spend_counters( @@ -412,9 +423,15 @@ def _should_track_cost_callback( user_id: Optional[str], team_id: Optional[str], end_user_id: Optional[str], + call_type: Optional[str] = None, ) -> bool: """ Determine if the cost callback should be tracked based on the kwargs + + Pass-through endpoints can be configured with ``auth=false``, which leaves + the request with no key/user/team/end-user to attribute spend to. Those + requests still forward real provider traffic that operators expect to see + in request/usage logs, so they are tracked even when unauthenticated. """ # don't run track cost callback if user opted into disabling spend @@ -423,7 +440,7 @@ def _should_track_cost_callback( if user_api_key is not None or user_id is not None or team_id is not None or end_user_id is not None: return True - return False + return call_type in _PASS_THROUGH_CALL_TYPES def _get_budget_reservation_from_metadata(metadata: dict) -> Optional[dict]: 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/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f468f5ec30b..797f4600857 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3550,7 +3550,7 @@ async def team_info( try: team_info: Optional[BaseModel] = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id}, - include={"object_permission": True}, + include={"litellm_model_table": True, "object_permission": True}, ) if team_info is None: raise Exception diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 065464aa565..d8015bb8031 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -4034,30 +4034,41 @@ class MicrosoftSSOHandler: base_url = MicrosoftSSOHandler.get_graph_api_base_url() # Endpoint to get app role assignments for the given service principal endpoint = f"/servicePrincipals/{service_principal_id}/appRoleAssignedTo" - url = base_url + endpoint + next_link: str | None = base_url + endpoint headers = { "Authorization": f"Bearer {access_token}", "Content-Type": "application/json", } - response = await async_client.get(url, headers=headers) - response_json = response.json() - verbose_proxy_logger.debug(f"Response from service principal app role assigned to: {response_json}") group_ids: List[str] = [] service_principal_teams: List[MicrosoftServicePrincipalTeam] = [] + page_count = 0 - for _object in response_json.get("value", []): - if _object.get("principalType") == "Group": - # Append the group ID to the list - group_ids.append(_object.get("principalId")) - # Append the service principal team to the list - service_principal_teams.append( - MicrosoftServicePrincipalTeam( - principalDisplayName=_object.get("principalDisplayName"), - principalId=_object.get("principalId"), + while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: + response = await async_client.get(next_link, headers=headers) + response_json = response.json() + verbose_proxy_logger.debug(f"Response from service principal app role assigned to: {response_json}") + + for _object in response_json.get("value", []): + if _object.get("principalType") == "Group": + # Append the group ID to the list + group_ids.append(_object.get("principalId")) + # Append the service principal team to the list + service_principal_teams.append( + MicrosoftServicePrincipalTeam( + principalDisplayName=_object.get("principalDisplayName"), + principalId=_object.get("principalId"), + ) ) - ) + + next_link = response_json.get("@odata.nextLink") + page_count += 1 + + if next_link is not None and page_count >= MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: + verbose_proxy_logger.warning( + f"Reached maximum page limit of {MicrosoftSSOHandler.MAX_GRAPH_API_PAGES}. Some service principal group assignments may not be included." + ) return group_ids, service_principal_teams diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 37f3d6e49e0..a0b07ece4dd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14805,6 +14805,7 @@ async def get_config_list( "forward_client_headers_to_llm_api": {"type": "Boolean"}, "mcp_required_fields": {"type": "List"}, "cancel_on_disconnect": {"type": "Boolean"}, + "skip_user_budget_on_team_key": {"type": "Boolean"}, } return_val = [] 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/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 38fd0d3f343..e1a093a4a48 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -123,6 +123,7 @@ async def reserve_budget_for_request( proxy_logging_obj: ProxyLogging, end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, + skip_user_budget_on_team_key: bool = False, ) -> Optional[dict]: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None @@ -141,6 +142,7 @@ async def reserve_budget_for_request( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, + skip_user_budget_on_team_key=skip_user_budget_on_team_key, ) if not counters: return None @@ -296,6 +298,7 @@ async def _get_budget_counters( proxy_logging_obj: ProxyLogging, end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, + skip_user_budget_on_team_key: bool = False, ) -> List[_BudgetCounter]: counters: List[_BudgetCounter] = [] @@ -344,8 +347,9 @@ async def _get_budget_counters( ) ) + is_team_key = team_object is not None and team_object.team_id is not None if ( - (team_object is None or team_object.team_id is None) + not (is_team_key and skip_user_budget_on_team_key) and user_object is not None and user_object.user_id is not None and user_object.max_budget is not None diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 92dd5a513b9..12f9be970c7 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -261,7 +261,7 @@ async def aresponses_api_with_mcp( pre_processed_mcp_tools=original_mcp_tools, ) - return LiteLLM_Proxy_MCP_Handler._create_mcp_streaming_response( + mcp_streaming_response = LiteLLM_Proxy_MCP_Handler._create_mcp_streaming_response( input=input, model=model, all_tools=all_tools, @@ -272,6 +272,10 @@ async def aresponses_api_with_mcp( tool_server_map=tool_server_map, **kwargs, ) + await mcp_streaming_response._create_initial_response_iterator() + if mcp_streaming_response._initial_creation_error is not None: + raise mcp_streaming_response._initial_creation_error + return mcp_streaming_response # Determine if we should auto-execute tools should_auto_execute = bool(mcp_tools_with_litellm_proxy) and LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index c705a04963c..d6b45855be5 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -5,6 +5,8 @@ from litellm._uuid import uuid from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, + ErrorEvent, + ErrorEventError, MCPCallArgumentsDeltaEvent, MCPCallArgumentsDoneEvent, MCPCallCompletedEvent, @@ -316,6 +318,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Cache the response ID to ensure consistency across all events self._cached_response_id: Optional[str] = None + self._initial_creation_error: Exception | None = None + self._stream_error: Exception | None = None + self._error_event_emitted = False + self._last_sequence_number = 0 + def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" from typing import Dict, Optional @@ -380,10 +387,31 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): return LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(self.mcp_tools_with_litellm_proxy) + def _make_stream_error_event(self) -> ResponsesAPIStreamingResponse: + err = self._stream_error + status_code = getattr(err, "status_code", None) + return ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=self._last_sequence_number + 1, + error=ErrorEventError( + type="mcp_gateway_error", + code=str(status_code) if status_code is not None else "internal_error", + message=str(err) if err is not None else "MCP gateway stream failed", + param=None, + ), + ) + def __aiter__(self): return self async def __anext__(self) -> ResponsesAPIStreamingResponse: + chunk = await self._anext_impl() + sequence_number = getattr(chunk, "sequence_number", None) + if isinstance(sequence_number, int) and sequence_number > self._last_sequence_number: + self._last_sequence_number = sequence_number + return chunk + + async def _anext_impl(self) -> ResponsesAPIStreamingResponse: """ Phase-based streaming: 1. initial_response - Stream the first LLM response (includes response.created, response.in_progress, response.output_item.added) @@ -438,10 +466,16 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.phase = "continue_initial_response" return await self.__anext__() self.phase = "finished" + if self._stream_error is not None and not self._error_event_emitted: + self._error_event_emitted = True + return self._make_stream_error_event() raise StopAsyncIteration # Phase 6: Finished if self.phase == "finished": + if self._stream_error is not None and not self._error_event_emitted: + self._error_event_emitted = True + return self._make_stream_error_event() raise StopAsyncIteration # Should not reach here @@ -530,6 +564,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined] + if self._cached_response_id is None and hasattr(chunk, "response"): + new_response = getattr(chunk, "response", None) + new_response_id = getattr(new_response, "id", None) if new_response is not None else None + if new_response_id: + self._cached_response_id = new_response_id + # Ensure response ID consistency - update chunk if needed if self._cached_response_id and hasattr(chunk, "response"): response_obj = getattr(chunk, "response", None) @@ -589,6 +629,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): traceback.print_exc() self.base_iterator = None + self._initial_creation_error = e + self._stream_error = e # Don't set phase to "finished" here — let __anext__ emit any # pre-generated MCP discovery events before ending the iteration. @@ -761,6 +803,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if hasattr(follow_up_response, "__aiter__"): self.base_iterator = follow_up_response self.collected_response = None + self._cached_response_id = None except Exception as e: verbose_logger.error(f"Error creating follow-up iterator: {e}") @@ -768,6 +811,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): traceback.print_exc() self.base_iterator = None + self._stream_error = e def __iter__(self): return self 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/model_prices_and_context_window.json b/model_prices_and_context_window.json index 08a452be844..8f2de0d0398 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11331,6 +11331,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11362,6 +11363,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true @@ -11424,6 +11426,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, @@ -11479,6 +11482,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11506,6 +11510,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11559,6 +11564,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true @@ -11586,6 +11592,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true @@ -11614,6 +11621,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "provider_specific_entry": { @@ -11648,6 +11656,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "provider_specific_entry": { @@ -11682,6 +11691,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, @@ -11718,6 +11728,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, @@ -11788,6 +11799,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, diff --git a/pyproject.toml b/pyproject.toml index c2a0e85d13a..f890cc976f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.93.0" +version = "1.94.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -62,8 +62,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.26.0,<2.0", - "litellm-proxy-extras==0.4.76", - "litellm-enterprise==0.1.49", + "litellm-proxy-extras==0.4.77", + "litellm-enterprise==0.1.50", "RestrictedPython>=8.1,<9.0", "rich>=13.9.4,<14.0", "polars>=1.38.1,<2.0", @@ -205,7 +205,7 @@ ci = [ # protobuf, Pillow is a compiled C extension). "tenacity==8.5.0", "google-generativeai==0.8.6", - "Pillow==12.2.0", + "Pillow==12.3.0", # Azure batch E2E tests still import psycopg2 directly. "psycopg2-binary==2.9.11", "pytest-codspeed==4.3.0", @@ -264,6 +264,8 @@ constraint-dependencies = [ "aiohttp>=3.14.1,<4.0", "packaging>=24.0", "soupsieve>=2.8.4", + "httplib2>=0.32.0", + "setuptools>=83.0.0", ] override-dependencies = [ # a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0. @@ -284,7 +286,7 @@ members = ["enterprise", "litellm-proxy-extras"] profile = "black" [tool.commitizen] -version = "1.93.0" +version = "1.94.0" version_files = [ "pyproject.toml:^version", ] 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/CLAUDE.md b/tests/e2e/CLAUDE.md index 88992038cb7..6b1b9950656 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -11,7 +11,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `embeddings/` - the `/embeddings` endpoint across providers - `batches/` - the `/batches` endpoint (placeholder until the first test lands) - `realtime/` - realtime websocket sessions, including the pipecat audio path -- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window) and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) +- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials; also the dashboard UI behavior on top of them, driven through the proxy-served UI at /ui with playwright (optional dep behind importorskip) - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 7d54f05656e..85d9315b8c6 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -17,6 +17,7 @@ from __future__ import annotations import json import time +from datetime import datetime, timedelta, timezone from typing import Callable import pytest @@ -49,7 +50,7 @@ from e2e_http import ( unwrap, ) from lifecycle import ResourceManager -from models import KeyGenerateBody, SpendLogRow, SpendLogsParams +from models import KeyGenerateBody, SpendLogRow pytestmark = pytest.mark.e2e @@ -349,6 +350,10 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( the file-read path fires while the batch itself is not blocked. ``resources.key()`` cannot set limits, so the key is minted on the gateway directly and its delete deferred. + + Snapshots read /spend/logs/v2 over a bounded window around the test instead + of the unpaginated /spend/logs whole-table read, which grows with the + environment and OOMed the e2e runner on stage. """ user_id = f"e2e-batch-rl-{unique_marker()}" key = client.gateway.generate_key( @@ -356,8 +361,13 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( ) resources.defer(lambda: client.gateway.delete_key(key)) + window_start = datetime.now(timezone.utc) - timedelta(hours=1) + window_end = window_start + timedelta(hours=2) before = frozenset( - row.request_id for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams())) + row.request_id + for row in unattributed_rows( + client.gateway.spend_logs_window(start=window_start, end=window_end) + ) ) file = unwrap( @@ -379,7 +389,9 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( new_orphans = [ row - for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams())) + for row in unattributed_rows( + client.gateway.spend_logs_window(start=window_start, end=window_end) + ) if row.request_id not in before ] assert not new_orphans, ( diff --git a/tests/e2e/coverage_registry/logging.yaml b/tests/e2e/coverage_registry/logging.yaml index 65ab8f0096f..c4b26387c13 100644 --- a/tests/e2e/coverage_registry/logging.yaml +++ b/tests/e2e/coverage_registry/logging.yaml @@ -8,7 +8,7 @@ - {id: logging.datadog.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/datadog/datadog.py", rationale: "Powers dashboards/alerts; cardinality regressions common"} - {id: logging.datadog.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions], source: "integrations/datadog/datadog.py", rationale: "Failure metrics for alerting/SLO"} - {id: logging.prometheus.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/prometheus.py", rationale: "Standard OSS metrics; per-key cardinality (existing e2e)"} -- {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"} +- {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"} - {id: logging.otel.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions, messages], source: "integrations/otel/logger.py", rationale: "Error spans for observability continuity"} - {id: logging.braintrust.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/braintrust_logging.py", rationale: "Evals platform spend"} - {id: logging.langsmith.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/langsmith.py", rationale: "LangChain ecosystem"} 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..6e6c30709de 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -24,6 +24,14 @@ 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") +CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.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/e2e_gateway.py b/tests/e2e/e2e_gateway.py index 05f83ecc085..d40b96d60fa 100644 --- a/tests/e2e/e2e_gateway.py +++ b/tests/e2e/e2e_gateway.py @@ -12,6 +12,7 @@ import time import warnings from collections.abc import Callable from dataclasses import dataclass +from datetime import datetime from e2e_http import ( NoBody, @@ -50,6 +51,8 @@ from models import ( OcrResponse, SpendLogRow, SpendLogs, + SpendLogsPage, + SpendLogsPageParams, SpendLogsParams, ) from e2e_config import ( @@ -255,6 +258,28 @@ class Gateway: case _: return [] + def spend_logs_window(self, *, start: datetime, end: datetime) -> list[SpendLogRow]: + def fetch(page: int) -> SpendLogsPage: + return unwrap( + self.transport.get( + "/spend/logs/v2", + headers=self.transport.master, + params=SpendLogsPageParams( + start_date=start.strftime("%Y-%m-%d %H:%M:%S"), + end_date=end.strftime("%Y-%m-%d %H:%M:%S"), + page=page, + page_size=100, + ), + response_type=SpendLogsPage, + ) + ) + + first = fetch(1) + return [ + *first.data, + *(row for page in range(2, first.total_pages + 1) for row in fetch(page).data), + ] + def poll_logs_for_key( self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None ) -> list[SpendLogRow]: diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 32005faed4b..ad53b2b4aa2 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -109,14 +109,16 @@ class StreamingResponse(BaseModel): """Raw outcome for calls whose body is provider-native or streamed: status, the x-litellm-call-id header, the x-litellm-response-cost header (StandardLogging response_cost), the content-type (which tells streaming `text/event-stream` from - non-streaming `application/json`), and the body. SpendLogs.request_id is the - completion body id, not call_id. Used by passthrough and streaming, where one - validated JSON model does not fit.""" + non-streaming `application/json`), the response headers (lowercased names, e.g. + the x-ratelimit-* pacing headers and retry-after on a 429), and the body. + SpendLogs.request_id is the completion body id, not call_id. Used by passthrough + and streaming, where one validated JSON model does not fit.""" status_code: int call_id: str | None = None # x-litellm-call-id header response_cost: float | None = None # x-litellm-response-cost header content_type: str | None = None + headers: dict[str, str] = {} body: str chunks: int = 0 # streamed events (0 for non-streaming) @@ -276,12 +278,14 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon call_id = _hdr(resp, "x-litellm-call-id") response_cost = _parse_response_cost(resp) content_type = _hdr(resp, "content-type") + headers = {name.lower(): value for name, value in resp.headers.items()} if not stream or not (200 <= resp.status_code < 300): return StreamingResponse( status_code=resp.status_code, call_id=call_id, response_cost=response_cost, content_type=content_type, + headers=headers, body=resp.text, ) lines = cast("Iterator[bytes]", resp.iter_lines()) @@ -291,6 +295,7 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon call_id=call_id, response_cost=response_cost, content_type=content_type, + headers=headers, body="", chunks=chunks, ) diff --git a/tests/e2e/logging/conftest.py b/tests/e2e/logging/conftest.py index 567355f23ac..43d279602ef 100644 --- a/tests/e2e/logging/conftest.py +++ b/tests/e2e/logging/conftest.py @@ -11,6 +11,7 @@ import os import pytest 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: @@ -28,6 +29,12 @@ def client() -> LoggingClient: return build_logging_client() +@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/logging_client.py b/tests/e2e/logging/logging_client.py index 6071e657bd3..bffdf71ed80 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -35,6 +35,7 @@ from e2e_http import ( unwrap, ) from models import ( + AnthropicMessagesBody, ChatBody, ChatMessage, ChatResponse, @@ -75,6 +76,14 @@ WEATHER_TOOL = ChatTool( ) +class ResponsesRequestBody(BaseModel): + """OpenAI Responses API /v1/responses request (non-streaming).""" + + model: str + input: str + max_output_tokens: int + + class TeamCallbackBody(BaseModel): callback_name: Literal["langfuse_otel", "langfuse", "langsmith", "gcs"] callback_type: Literal["success", "failure", "success_and_failure"] @@ -455,6 +464,32 @@ class LoggingClient: json=body, ) + def messages_raw(self, key: str, model: str, text: str, *, max_tokens: int = 16) -> StreamingResponse: + """Non-streaming POST /v1/messages (Anthropic-native body): raw outcome + judged by status/body/headers, for tests that need x-litellm-call-id.""" + return self.gateway.transport.send( + "/v1/messages", + headers=self.gateway.transport.bearer(key), + json=AnthropicMessagesBody( + model=model, + max_tokens=max_tokens, + messages=[ChatMessage(role="user", content=text)], + ), + ) + + def responses_raw( + self, key: str, model: str, text: str, *, max_output_tokens: int = 64 + ) -> StreamingResponse: + """Non-streaming POST /v1/responses (OpenAI Responses API): raw outcome + judged by status/body/headers, for tests that need x-litellm-call-id. + max_output_tokens caps reasoning-model output cost; a capped response is + still a 200 and still exports the trace.""" + return self.gateway.transport.send( + "/v1/responses", + headers=self.gateway.transport.bearer(key), + json=ResponsesRequestBody(model=model, input=text, max_output_tokens=max_output_tokens), + ) + def scrape_metrics(self) -> str: return self.gateway.probe("/metrics", params=NoBody()).body 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..90887ea5510 --- /dev/null +++ b/tests/e2e/logging/test_otel_trace_e2e.py @@ -0,0 +1,259 @@ +"""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, CHEAP_OPENAI_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", exercised_on=["chat_completions"]) + 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}") + + @pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["messages"]) + def test_messages_exports_complete_trace( + self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager + ) -> None: + """This test verifies that one successful non-streaming /v1/messages request + produces exactly one complete OTEL trace. + + The trace must have a single root span named "POST /v1/messages". The + authentication, database, cost-writing, and model-call spans must all belong to + the same trace and have valid parent relationships leading back to that root. + + The model-call span is expected to be named "chat ". The test fails if + the request is split across multiple traces, if any span references a missing + parent, or if the model-call span cannot be connected back to the root.""" + route = "/v1/messages" + _assert_otel_destination_configured(client) + + key = client.key_with_alias(f"otel-trace-messages-{unique_marker()}", models=[MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = _first_ok( + client, lambda: client.messages_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}") + + @pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["responses"]) + def test_responses_exports_complete_trace( + self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager + ) -> None: + """This test verifies that one successful non-streaming /v1/responses request + produces exactly one complete OTEL trace. + + The trace must have a single root span named "POST /v1/responses". The + authentication, database, cost-writing, and model-call spans must all belong to + the same trace and have valid parent relationships leading back to that root. + + The model-call span is expected to be named "chat ". The test fails if + the request is split across multiple traces, if any span references a missing + parent, or if the model-call span cannot be connected back to the root.""" + route = "/v1/responses" + _assert_otel_destination_configured(client) + + key = client.key_with_alias(f"otel-trace-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = _first_ok( + client, + lambda: client.responses_raw(key, CHEAP_OPENAI_MODEL, f"reply with one word {marker}"), + ) + assert outcome.call_id is not None, "success response must carry x-litellm-call-id" + + genai_span = f"chat {CHEAP_OPENAI_MODEL}" + hits = otel_reader.poll_traces_for_call( + call_id=outcome.call_id, + settled_names=_settled_names(route=route, genai_span=genai_span), + settled_prefixes={DB_SPAN_PREFIX}, + ) + _assert_complete_trace(hits, route=route, genai_span=genai_span) diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index cbd5db0d59f..7a05a5c520f 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -111,7 +111,7 @@ class TestKeyRoutes: key = _generate_key( client, resources, - KeyGenerateBody(models=["gemini-2.5-flash"], key_alias=alias, tpm_limit=424242), + KeyGenerateBody(models=["gemini-2.5-flash"], key_alias=alias, tpm_limit=424242, rpm_limit=424243), ) info = client.gateway.key_info(key) @@ -122,6 +122,9 @@ class TestKeyRoutes: assert info.tpm_limit == 424242, ( f"/key/info reports tpm_limit {info.tpm_limit}, configured 424242" ) + assert info.rpm_limit == 424243, ( + f"/key/info reports rpm_limit {info.rpm_limit}, configured 424243" + ) _poll_chat_ok(client, key, "gemini-2.5-flash") _assert_model_denied( diff --git a/tests/e2e/models.py b/tests/e2e/models.py index e32f2709181..bf90426188f 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -8,7 +8,7 @@ from __future__ import annotations from typing import Literal -from pydantic import BaseModel, ConfigDict, RootModel +from pydantic import BaseModel, ConfigDict, RootModel, model_validator # ---------- keys ---------- @@ -82,6 +82,7 @@ class KeyInfo(BaseModel): key_alias: str | None = None models: list[str] = [] tpm_limit: int | None = None + rpm_limit: int | None = None team_id: str | None = None spend: float | None = None max_budget: float | None = None @@ -255,6 +256,16 @@ class SpendLogsParams(BaseModel): request_id: str | None = None api_key: str | None = None + @model_validator(mode="after") + def require_filter(self) -> SpendLogsParams: + if self.request_id is None and self.api_key is None: + raise ValueError( + "unfiltered /spend/logs returns the entire spend table and OOMs the " + "runner on long-lived environments; filter by request_id or api_key, " + "or use Gateway.spend_logs_window for a bounded /spend/logs/v2 read" + ) + return self + class SpendLogsPageParams(BaseModel): """Query for /spend/logs/v2, which requires an explicit date window and diff --git a/tests/e2e/quota_management/ratelimit/conftest.py b/tests/e2e/quota_management/ratelimit/conftest.py new file mode 100644 index 00000000000..4a5a73bb5e4 --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/conftest.py @@ -0,0 +1,15 @@ +"""Quota-management suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +live in the parent tests/e2e/conftest.py. QuotaClient holds the shared Gateway, +so the `resources` fixture cleans up keys through it. +""" + +import pytest + +from quota_client import QuotaClient, build_client + + +@pytest.fixture(scope="session") +def client() -> QuotaClient: + return build_client() diff --git a/tests/e2e/quota_management/ratelimit/quota_client.py b/tests/e2e/quota_management/ratelimit/quota_client.py new file mode 100644 index 00000000000..806ab1d1557 --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/quota_client.py @@ -0,0 +1,32 @@ +"""Client for the quota-management suite: the shared Gateway plus raw chat +calls judged by HTTP status, body, and headers (a rate-limit block is a 429 +whose body and retry-after header carry the contract, not a typed success +model).""" + +from __future__ import annotations + +from dataclasses import dataclass + +from e2e_gateway import Gateway, build_gateway +from e2e_http import StreamingResponse +from models import ChatBody, ChatMessage + + +@dataclass(frozen=True, slots=True) +class QuotaClient: + gateway: Gateway + + def chat(self, key: str, model: str, content: str, *, max_tokens: int = 16) -> StreamingResponse: + return self.gateway.transport.send( + "/chat/completions", + headers=self.gateway.transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + ), + ) + + +def build_client() -> QuotaClient: + return QuotaClient(gateway=build_gateway()) diff --git a/tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py b/tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py new file mode 100644 index 00000000000..a6c15b79bb1 --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py @@ -0,0 +1,254 @@ +"""Live e2e: key-level rpm/tpm rate limits on the gateway. + +Covers quota_management.ratelimit.*: a key generated with rpm_limit/tpm_limit gets a +429 once the limit is crossed inside one window (blocks_over_limit), serves +again once the window rolls and no sooner (resets_after_window), and successful responses +report x-ratelimit-* limit/remaining headers so clients can pace +(headers_report_remaining). Each test asserts both halves of the contract: the +recorded state (/key/info echoes the configured limit) and the enforced +behavior (the 429, the recovery, or the headers on live traffic). + +The v3 limiter counts a request against the rpm budget at the pre-call hook, +before model routing, so every call that clears auth consumes budget whether or +not it ultimately succeeds. The tpm budget is reserved pre-call from an estimate +(message chars // 4 + max_tokens) and reconciled to the body's actual +usage.total_tokens after the call, so a block may legitimately fire before the +actual spend crosses the limit; the tpm test asserts the exact contract on both +sides (a 429 only once the blocked call's reservation exceeds the remaining +budget, and no later than the first call after actual spend reaches the limit). +All calls of one test must land inside a single window +(LITELLM_RATE_LIMIT_WINDOW_SIZE, 60s default), which real chat latency +comfortably allows. + +The window opens at the pre-call hook of the first counted call, which happens +after that call is sent, so the send timestamp of the winning first call is a +lower bound on the window start. The reset test uses it to reject an early +reset: recovery must not arrive before the full window has elapsed from that +send, less a small tolerance for the limiter's integer-second window +arithmetic. +""" + +from __future__ import annotations + +import re +import time +from dataclasses import dataclass + +import pytest +from pydantic import BaseModel, ConfigDict, ValidationError + +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from models import KeyGenerateBody +from quota_client import QuotaClient + +pytestmark = pytest.mark.e2e + +MODEL = CHEAP_ANTHROPIC_MODEL +TPM_LIMIT = 60 +CHAT_MAX_TOKENS = 16 +RESERVATION_CHARS_PER_TOKEN = 4 +WINDOW_SECONDS = 60 +RESET_TOLERANCE_SECONDS = 5 +LAST_CALL_LATENCY_MARGIN_SECONDS = 10 + + +@dataclass(frozen=True, slots=True) +class _FirstOk: + sent_at: float + response: StreamingResponse + + +class _ChatUsage(BaseModel): + model_config = ConfigDict(extra="ignore") + + total_tokens: int + + +class _ChatBodyWithUsage(BaseModel): + model_config = ConfigDict(extra="ignore") + + usage: _ChatUsage + + +def _total_tokens(outcome: StreamingResponse) -> int: + try: + return _ChatBodyWithUsage.model_validate_json(outcome.body).usage.total_tokens + except ValidationError: + pytest.fail(f"successful chat body must report usage.total_tokens, got: {outcome.body[:300]}") + + +def _reserved_tokens(content: str) -> int: + return max(1, len(content) // RESERVATION_CHARS_PER_TOKEN) + CHAT_MAX_TOKENS + + +def _remaining_from_429(body: str) -> int: + found = re.search(r"Remaining: (\d+)", body) + if found is None: + pytest.fail(f"429 body must report the remaining budget, got: {body[:300]}") + return int(found.group(1)) + + +@dataclass(frozen=True, slots=True) +class _BlockedByReservation: + outcome: StreamingResponse + content: str + spent: int + + +@dataclass(frozen=True, slots=True) +class _CrossedLimit: + spent: int + + +def _spend_until_blocked_or_crossed( + client: QuotaClient, key: str, first: _FirstOk +) -> _BlockedByReservation | _CrossedLimit: + """Drive chat traffic, summing each body's actual usage.total_tokens, until + the limiter blocks (which the reservation may do before the actual spend + crosses the limit) or the actual spend reaches the limit.""" + window_deadline = first.sent_at + WINDOW_SECONDS - LAST_CALL_LATENCY_MARGIN_SECONDS + spent = _total_tokens(first.response) + while spent < TPM_LIMIT: + assert time.monotonic() < window_deadline, ( + f"spent only {spent} of {TPM_LIMIT} tokens before the {WINDOW_SECONDS}s window could roll; " + "the exact-crossing assertion needs every call inside one window" + ) + content = f"reply with one word {unique_marker()}" + outcome = client.chat(key, MODEL, content, max_tokens=CHAT_MAX_TOKENS) + if outcome.status_code == 429: + return _BlockedByReservation(outcome=outcome, content=content, spent=spent) + require_successful_call(outcome) + spent += _total_tokens(outcome) + return _CrossedLimit(spent=spent) + + +def _limited_key( + client: QuotaClient, + resources: ResourceManager, + *, + rpm_limit: int | None = None, + tpm_limit: int | None = None, +) -> str: + key = client.gateway.generate_key(KeyGenerateBody(models=[MODEL], rpm_limit=rpm_limit, tpm_limit=tpm_limit)) + resources.defer(lambda: client.gateway.delete_key(key)) + return key + + +def _chat(client: QuotaClient, key: str) -> StreamingResponse: + return client.chat(key, MODEL, f"reply with one word {unique_marker()}") + + +def _first_ok(client: QuotaClient, key: str) -> _FirstOk: + """First successful call on a fresh key, which opens the rate-limit window; + `sent_at` is captured just before the winning send, so the window opened no + earlier than it. 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 never reaches the + rate limiter, so only the successful call consumes budget. Any other failure + is behavior under test and fails hard.""" + deadline = time.monotonic() + client.gateway.poll_timeout + while True: + sent_at = time.monotonic() + outcome = _chat(client, key) + if outcome.ok: + return _FirstOk(sent_at=sent_at, response=outcome) + if outcome.status_code != 401 or time.monotonic() >= deadline: + require_successful_call(outcome) + time.sleep(client.gateway.poll_interval) + + +def _assert_rate_limited(outcome: StreamingResponse, limit_type: str) -> None: + assert outcome.status_code == 429, ( + f"expected a 429 {limit_type} rate-limit block, got {outcome.status_code}: {outcome.body[:300]}" + ) + assert "Rate limit exceeded for api_key" in outcome.body, ( + f"429 body must name the api_key scope, got: {outcome.body[:300]}" + ) + assert f"Limit type: {limit_type}" in outcome.body, ( + f"429 body must carry 'Limit type: {limit_type}', got: {outcome.body[:300]}" + ) + retry_after = outcome.headers.get("retry-after") + assert retry_after is not None and retry_after.isdigit() and int(retry_after) > 0, ( + f"429 must carry a positive integer retry-after header, got {retry_after!r}" + ) + + +class TestKeyRateLimits: + @pytest.mark.covers("quota_management.ratelimit.rpm.blocks_over_limit") + def test_rpm_limit_blocks_over_limit(self, client: QuotaClient, resources: ResourceManager) -> None: + key = _limited_key(client, resources, rpm_limit=3) + info = client.gateway.key_info(key) + assert info.rpm_limit == 3, f"/key/info reports rpm_limit {info.rpm_limit}, configured 3" + + _ = _first_ok(client, key) + for _ in range(2): + require_successful_call(_chat(client, key)) + + _assert_rate_limited(_chat(client, key), "requests") + + @pytest.mark.covers("quota_management.ratelimit.tpm.blocks_over_limit") + def test_tpm_limit_blocks_over_limit(self, client: QuotaClient, resources: ResourceManager) -> None: + key = _limited_key(client, resources, tpm_limit=TPM_LIMIT) + info = client.gateway.key_info(key) + assert info.tpm_limit == TPM_LIMIT, f"/key/info reports tpm_limit {info.tpm_limit}, configured {TPM_LIMIT}" + + first = _first_ok(client, key) + match _spend_until_blocked_or_crossed(client, key, first): + case _BlockedByReservation(outcome=outcome, content=content, spent=spent): + _assert_rate_limited(outcome, "tokens") + remaining = _remaining_from_429(outcome.body) + reserved = _reserved_tokens(content) + assert reserved > remaining, ( + f"blocked while the call still fit: {remaining} of {TPM_LIMIT} tokens remained but the call " + f"reserved only {reserved} ({spent} actual tokens spent so far)" + ) + case _CrossedLimit(): + _assert_rate_limited(_chat(client, key), "tokens") + + @pytest.mark.covers("quota_management.ratelimit.rpm.resets_after_window") + def test_rpm_limit_resets_after_window(self, client: QuotaClient, resources: ResourceManager) -> None: + key = _limited_key(client, resources, rpm_limit=1) + + first = _first_ok(client, key) + _assert_rate_limited(_chat(client, key), "requests") + + deadline = time.monotonic() + client.gateway.poll_timeout + while time.monotonic() < deadline: + attempt_sent_at = time.monotonic() + outcome = _chat(client, key) + if outcome.ok: + window_age = attempt_sent_at - first.sent_at + assert window_age >= WINDOW_SECONDS - RESET_TOLERANCE_SECONDS, ( + f"the key recovered {window_age:.1f}s after the window opened, before the " + f"{WINDOW_SECONDS}s window (less {RESET_TOLERANCE_SECONDS}s tolerance) elapsed; " + "the limiter reset early instead of after the window" + ) + return + assert outcome.status_code == 429, ( + f"while the window drains only 429s are acceptable, got {outcome.status_code}: {outcome.body[:300]}" + ) + time.sleep(client.gateway.poll_interval) + pytest.fail("a blocked key never recovered after the rate-limit window elapsed") + + @pytest.mark.covers("quota_management.ratelimit.rpm.headers_report_remaining") + def test_headers_report_limit_and_remaining(self, client: QuotaClient, resources: ResourceManager) -> None: + key = _limited_key(client, resources, rpm_limit=5, tpm_limit=100000) + + first = _first_ok(client, key).response + assert first.headers.get("x-ratelimit-api_key-limit-requests") == "5", ( + f"success response must report the key's request limit, headers: " + f"{ {k: v for k, v in first.headers.items() if 'ratelimit' in k} }" + ) + assert first.headers.get("x-ratelimit-api_key-remaining-requests") == str(5 - 1), ( + f"first call against rpm_limit=5 must leave {5 - 1} remaining, got " + f"{first.headers.get('x-ratelimit-api_key-remaining-requests')!r}" + ) + assert first.headers.get("x-ratelimit-api_key-limit-tokens") == "100000", ( + f"success response must report the key's token limit, got " + f"{first.headers.get('x-ratelimit-api_key-limit-tokens')!r}" + ) + remaining_tokens = first.headers.get("x-ratelimit-api_key-remaining-tokens") + assert remaining_tokens is not None and remaining_tokens.isdigit() and int(remaining_tokens) < 100000, ( + f"one call must leave remaining tokens reported and below the limit, got {remaining_tokens!r}" + ) diff --git a/tests/e2e/test_e2e_gateway.py b/tests/e2e/test_e2e_gateway.py index 9a9aa2fd2cc..70e846ad8d5 100644 --- a/tests/e2e/test_e2e_gateway.py +++ b/tests/e2e/test_e2e_gateway.py @@ -1,17 +1,23 @@ """Unit coverage for the Gateway model-management surface (create_model / -delete_model). +delete_model) and the bounded spend read-back (spend_logs_window). The batches conftest and several llm_translation tests register deployments at runtime through gateway.create_model; when that method went missing, every batch test errored at fixture setup (AttributeError) before a single request reached the proxy. This pins the surface with a typed fake Transport so a rename or signature drift fails here instead of in a live stage run. + +spend_logs_window exists because the unpaginated /spend/logs whole-table read +grew past the e2e runner's memory limit on stage and OOMKilled every run; these +tests pin its /spend/logs/v2 pagination and that SpendLogsParams can no longer +express the unfiltered read. """ from dataclasses import dataclass, field +from datetime import datetime, timezone import pytest -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from batches.batch_client import BatchClient from e2e_gateway import Gateway @@ -30,6 +36,9 @@ from models import ( ModelNewBody, ModelNewResponse, ModelsListResponse, + SpendLogsPage, + SpendLogsPageParams, + SpendLogsParams, ) @@ -46,6 +55,8 @@ class _RecordingTransport: servable_after_gets: int = 0 models_error: UnknownApiError | None = None model_gets: int = 0 + spend_total: int = 0 + spend_gets: list[SpendLogsPageParams] = field(default_factory=list) _created: list[str] = field(default_factory=list) def post[R: BaseModel]( @@ -91,6 +102,22 @@ class _RecordingTransport: return Success( data=response_type.model_validate({"data": [{"id": name} for name in visible]}) ) + if path == "/spend/logs/v2" and response_type is SpendLogsPage: + assert isinstance(params, SpendLogsPageParams) + self.spend_gets.append(params) + offset = (params.page - 1) * params.page_size + count = min(params.page_size, max(self.spend_total - offset, 0)) + return Success( + data=response_type.model_validate( + { + "data": [{"request_id": f"req-{offset + i}"} for i in range(count)], + "total": self.spend_total, + "page": params.page, + "page_size": params.page_size, + "total_pages": (self.spend_total + params.page_size - 1) // params.page_size, + } + ) + ) raise AssertionError(f"unexpected get: {path}") def delete[R: BaseModel]( @@ -202,3 +229,45 @@ def test_gateway_delete_model_posts_the_model_id() -> None: assert path == "/model/delete" assert isinstance(body, ModelDeleteBody) assert body.id == "registered-id" + + +WINDOW_START = datetime(2026, 7, 14, 12, 0, 0, tzinfo=timezone.utc) +WINDOW_END = datetime(2026, 7, 14, 14, 0, 0, tzinfo=timezone.utc) + + +def test_gateway_spend_logs_window_pages_through_every_row_in_the_window() -> None: + transport = _RecordingTransport(spend_total=250) + gateway = Gateway(transport=transport) + + rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END) + + assert len(rows) == 250 + assert len({row.request_id for row in rows}) == 250 + assert [params.page for params in transport.spend_gets] == [1, 2, 3] + assert all(params.start_date == "2026-07-14 12:00:00" for params in transport.spend_gets) + assert all(params.end_date == "2026-07-14 14:00:00" for params in transport.spend_gets) + + +def test_gateway_spend_logs_window_stops_at_an_exact_page_boundary() -> None: + transport = _RecordingTransport(spend_total=200) + gateway = Gateway(transport=transport) + + rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END) + + assert len(rows) == 200 + assert [params.page for params in transport.spend_gets] == [1, 2] + + +def test_gateway_spend_logs_window_returns_empty_for_an_empty_window() -> None: + transport = _RecordingTransport(spend_total=0) + gateway = Gateway(transport=transport) + + rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END) + + assert rows == [] + assert [params.page for params in transport.spend_gets] == [1] + + +def test_spend_logs_params_rejects_the_unfiltered_whole_table_read() -> None: + with pytest.raises(ValidationError, match="spend_logs_window"): + SpendLogsParams() 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/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 958b028c542..5471d2668e4 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -219,8 +219,8 @@ async def test_aaauser_personal_budgets(key_ownership): """ Set a personal budget on a user - - have it only apply when key belongs to user -> raises BudgetExceededError - - if key belongs to team, have key respect team budget -> allows call to go through + User budget is enforced regardless of key ownership (personal or team). + Both cases should raise BudgetExceededError when the user is over budget. """ import asyncio import time @@ -229,7 +229,12 @@ async def test_aaauser_personal_budgets(key_ownership): from starlette.datastructures import URL import litellm - from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth + from litellm.proxy._types import ( + LiteLLM_UserTable, + ProxyErrorTypes, + ProxyException, + UserAPIKeyAuth, + ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import hash_token, user_api_key_cache @@ -273,14 +278,9 @@ async def test_aaauser_personal_budgets(key_ownership): == valid_token ) - try: + with pytest.raises(ProxyException) as exc_info: await user_api_key_auth(request=request, api_key="Bearer " + user_key) - - if key_ownership == "user_key": - pytest.fail("Expected this call to fail. User is over limit.") - except Exception: - if key_ownership == "team_key": - pytest.fail("Expected this call to work. Key is below team budget.") + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index d300f326b9e..9289dece83f 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1614,3 +1614,105 @@ class TestGuardrailInterventionClassification: slg = request_data["metadata"]["standard_logging_guardrail_information"][0] assert slg["guardrail_status"] == "guardrail_intervened" + + +class _ApplyStyleGuardrail(CustomGuardrail): + """Overrides only apply_guardrail, like openai_moderation; async_pre_call_hook stays the CustomLogger no-op.""" + + def __init__(self, block: bool): + from litellm.types.guardrails import GuardrailEventHooks + + super().__init__( + guardrail_name="apply-style-guardrail", + event_hook=GuardrailEventHooks.pre_call, + default_on=False, + ) + self.block = block + self.apply_called = False + self.seen_texts = None + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + from fastapi import HTTPException + + self.apply_called = True + self.seen_texts = inputs.get("texts") + if self.block: + raise HTTPException(status_code=400, detail={"error": "Violated moderation policy"}) + return inputs + + +class TestApplyGuardrailStyleDeploymentDispatch: + """LIT-4217 regression: model-level guardrails that implement only the + unified apply_guardrail interface must execute in + async_pre_call_deployment_hook instead of silently hitting the + async_pre_call_hook no-op.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", [CallTypes.completion, CallTypes.acompletion]) + async def test_blocks_when_requested_via_model_level_guardrails(self, call_type): + from fastapi import HTTPException + + guardrail = _ApplyStyleGuardrail(block=True) + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "flagged content"}], + "guardrails": ["apply-style-guardrail"], + "metadata": {}, + } + + with pytest.raises(HTTPException): + await guardrail.async_pre_call_deployment_hook(kwargs, call_type) + + assert guardrail.apply_called is True + assert guardrail.seen_texts == ["flagged content"] + + @pytest.mark.asyncio + async def test_pass_path_runs_guardrail_and_strips_dispatch_key(self): + guardrail = _ApplyStyleGuardrail(block=False) + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "guardrails": ["apply-style-guardrail"], + "metadata": {}, + } + + result = await guardrail.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert guardrail.apply_called is True + assert result is not None + assert "guardrail_to_apply" not in result + assert result["messages"] == [{"role": "user", "content": "hello"}] + + @pytest.mark.asyncio + async def test_skips_when_not_requested(self): + guardrail = _ApplyStyleGuardrail(block=True) + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "guardrails": ["some-other-guardrail"], + "metadata": {}, + } + + result = await guardrail.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert guardrail.apply_called is False + assert result is not None + + @pytest.mark.asyncio + async def test_fails_closed_when_proxy_extras_missing(self): + import sys + from unittest.mock import patch + + guardrail = _ApplyStyleGuardrail(block=True) + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "flagged content"}], + "guardrails": ["apply-style-guardrail"], + "metadata": {}, + } + + with patch.dict(sys.modules, {"litellm.proxy.utils": None}): + with pytest.raises(ImportError, match="litellm\\[proxy\\]"): + await guardrail.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert guardrail.apply_called is False 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/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 43ef7fcd971..94a4a3fc945 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -957,15 +957,15 @@ def test_anthropic_structured_output_beta_header(): @pytest.mark.parametrize( "model_name", [ - "claude-opus-4-6-20250918", - "claude-opus-4.6-20250918", + "claude-opus-4-8", + "claude-opus-4-6-20260205", "claude-opus-4-5-20251101", "claude-opus-4.5-20251101", ], ) def test_opus_uses_native_structured_output(model_name): """ - Test that Opus 4.5 and 4.6 models use native Anthropic structured outputs + Test that supported Opus models use native Anthropic structured outputs (output_format) rather than the tool-based workaround. """ config = AnthropicConfig() @@ -1005,6 +1005,43 @@ def test_opus_uses_native_structured_output(model_name): assert optional_params.get("json_mode") is True +def test_native_structured_output_uses_bundled_capability_when_remote_map_lags( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = "claude-opus-4-8" + monkeypatch.setattr( + litellm, + "model_cost", + {model: {"supports_response_schema": True}}, + ) + litellm.get_model_info.cache_clear() + + try: + optional_params = AnthropicConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "answer", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + }, + }, + } + }, + optional_params={}, + model=model, + drop_params=False, + ) + finally: + litellm.get_model_info.cache_clear() + + assert "output_format" in optional_params + assert "tools" not in optional_params + + def test_non_structured_output_model_uses_tool_workaround(): """ Test that models NOT in the native structured output list still use the diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py index 06d3effcfbb..5254808e315 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py @@ -2,6 +2,7 @@ import pytest from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, ) @@ -174,6 +175,81 @@ def test_unrecognized_effort_raises_clean_400(): assert exc_info.value.status_code == 400 +def test_pinned_temperature_dropped_when_adaptive_downgraded_to_enabled(): + """Regression (#33203): Claude Code's safety classifier sends adaptive thinking + + temperature=0 to Haiku 4.5. The adaptive interface is downgraded to legacy enabled + thinking, but Anthropic rejects "temperature may only be set to 1 when thinking is + enabled". The pinned temperature must be dropped so the request succeeds while the + downgraded thinking is preserved.""" + params = _claude_code_payload(effort="medium") + params["temperature"] = 0 + result = _transform("claude-haiku-4-5", params) + + assert result["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + } + assert "temperature" not in result + + +def test_temperature_one_preserved_with_enabled_thinking(): + """temperature=1 is compatible with extended thinking, so it must be kept.""" + params = _claude_code_payload(effort="medium") + params["temperature"] = 1 + result = _transform("claude-haiku-4-5", params) + + assert result["thinking"]["type"] == "enabled" + assert result["temperature"] == 1 + + +def test_pinned_temperature_preserved_when_thinking_dropped(): + """When thinking is dropped entirely (non-reasoning model), there is no thinking + conflict, so a pinned temperature must survive untouched.""" + params = _claude_code_payload(effort="medium") + params["temperature"] = 0 + result = _transform("claude-3-5-haiku-latest", params) + + assert "thinking" not in result + assert result["temperature"] == 0 + + +def test_pinned_temperature_preserved_for_adaptive_model(): + """Adaptive models (4.6+) own the thinking/temperature relationship natively, so + the passthrough must not strip a pinned temperature for them.""" + params = _claude_code_payload(effort="high") + params["temperature"] = 0 + result = _transform("claude-sonnet-4-6", params) + + assert result["thinking"] == {"type": "adaptive"} + assert result["temperature"] == 0 + + +def test_pinned_temperature_dropped_for_opus_4_5_effort(): + """Opus 4.5 keeps native output_config.effort (extended thinking), which is equally + incompatible with a pinned non-1 temperature, so the temperature must be dropped.""" + params = _claude_code_payload(effort="medium") + params["temperature"] = 0 + result = _transform("claude-opus-4-5", params) + + assert result["output_config"] == {"effort": "medium"} + assert "temperature" not in result + + +def test_reasoning_effort_with_pinned_temperature_drops_temperature(): + """The reasoning_effort alias synthesizes legacy enabled thinking on a non-adaptive + model; a co-pinned non-1 temperature must be dropped to avoid the Anthropic 400.""" + result = _transform( + "claude-haiku-4-5", + {"max_tokens": 8192, "reasoning_effort": "low", "temperature": 0}, + ) + + assert result["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, + } + assert "temperature" not in result + + def test_non_adaptive_request_without_effort_is_untouched(): """A non-adaptive model receiving a request with no adaptive interface (no effort, no adaptive thinking) must pass through untouched.""" 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 c55a631c7b3..d1aceffa968 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(): @@ -4362,6 +4370,1247 @@ async def test_register_bridge_relay_never_persists(): mock_persist.assert_not_called() +_BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" + + +async def _exchange_for_bridge_server(server, upstream_body, key_hash, code="auth-code", fake_client_out=None): + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _ResolvedKey + 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) + # The mint consumes _resolve_active_litellm_key's tagged result: an active key resolves to a + # _ResolvedKey carrying its hash; a request with no usable credential resolves to "no_active_key". + resolution = _ResolvedKey(key_hash=key_hash, key=MagicMock()) if key_hash is not None else "no_active_key" + key_resolver = AsyncMock(return_value=resolution) + 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.bridge_token_flow._resolve_active_litellm_key", + new=key_resolver, + ), + 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="authorization_code", + code=code, + redirect_uri="https://claude.ai/api/mcp/auth_callback", + client_id="dcr-client-123", + client_secret=None, + code_verifier="verifier", + ) + 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() + return response + + +@pytest.mark.asyncio +async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token(): + """A dcr_bridge oauth_delegate token exchange returns a gateway-bound envelope, not the raw + upstream token: the response access_token opens (under the same master-key-derived keys and the + server_id) to the caller's identity and the upstream Authorization, and the raw upstream token + never appears in the bearer the client receives.""" + 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.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "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) + token = body["access_token"] + assert body["token_type"] == "Bearer" + assert body["expires_in"] > 0 + assert token.startswith("llm_env_") + assert "UPSTREAM-SECRET-TOKEN" not in token + assert "refresh_token" not in body + + 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 == "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 + identity-less envelope. It returns an RFC 6749 §5.2-shaped invalid_request (error at the top + level, not wrapped in detail) BEFORE exchanging the upstream code, so the single-use code is not + burned and the client can retry.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + 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, fake_client_out=captured) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + # identity resolution failed first, so the upstream single-use code was never exchanged (not burned) + captured["client"].post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_envelope_too_large_upstream_token_is_502(): + """An upstream token too large to seal into the envelope is an upstream-payload condition, so the + mint surfaces a 502 (as an RFC 6749 §5.2 error body, not a raised HTTPException) rather than a 500: + build_bridge_token_response returns EnvelopeTooLarge as a value, _finish_bridge_mint returns the + "too_large" failure, and _bridge_mint_error_response maps it to a truthful status.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "x" * 40000, "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +@pytest.mark.asyncio +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 ( + envelope_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + OpenedEnvelope, + open_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": "UPSTREAM-REFRESH", + } + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.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_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 + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock() + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + 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="client-refresh-token", + ) + + assert response.status_code == 400 + 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.bridge_token_flow._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.bridge_token_flow 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.bridge_token_flow._reload_active_key_by_hash", + new=AsyncMock(return_value=_ResolvedKey(key_hash="kh", key=MagicMock())), + ) as key_reload, + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._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.bridge_token_flow._reload_active_key_by_hash", + new=AsyncMock(), + ) as key_reload2, + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._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.bridge_token_flow 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.bridge_token_flow._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.bridge_token_flow._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.bridge_token_flow 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.bridge_token_flow._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.bridge_token_flow 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.bridge_token_flow._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.bridge_token_flow._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 + misconfigured gateway returns a 500 server_error without consuming the single-use code, avoiding + the burn-then-fail the pre-exchange phase exists to prevent. The failure is returned as an RFC 6749 + error body, not raised.""" + 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) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock() + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._resolve_active_litellm_key", + new=AsyncMock(return_value="no_active_key"), + ), + patch("litellm.proxy.proxy_server.master_key", None), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://claude.ai/api/mcp/auth_callback", + client_id="dcr-client-123", + client_secret=None, + code_verifier="verifier", + ) + + assert response.status_code == 500 + assert json.loads(response.body)["error"] == "server_error" + fake_http_client.post.assert_not_called() + + +async def _prepare_only_bridge_exchange(resolver_result): + """Drive exchange_token_with_server for a bridge oauth_delegate authorization_code request with the + identity resolver stubbed to a given tagged result, returning (response, post_mock) so a test can + assert the mapped status and that the single-use code was never exchanged.""" + 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) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock() + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._resolve_active_litellm_key", + new=AsyncMock(return_value=resolver_result), + ), + 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="authorization_code", + code="auth-code", + redirect_uri="https://claude.ai/api/mcp/auth_callback", + client_id="dcr-client-123", + client_secret=None, + code_verifier="verifier", + ) + return response, fake_http_client.post + + +@pytest.mark.asyncio +async def test_bridge_mint_db_outage_is_503_before_upstream(): + """A DB outage while resolving identity is a retryable gateway failure, so the mint returns 503 + temporarily_unavailable WITHOUT consuming the single-use code, matching how admission statuses the + same outage on the egress side. Collapsing every resolution failure to None used to blame the + client with 400 invalid_request for an infrastructure problem.""" + response, post = await _prepare_only_bridge_exchange("unavailable") + assert response.status_code == 503 + assert json.loads(response.body)["error"] == "temporarily_unavailable" + post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_mint_unresolvable_identity_is_500_before_upstream(): + """An unresolvable identity (no DB connection, or an unexpected resolution error) is a gateway + fault, so the mint returns 500 server_error before the exchange, a status distinct from both the + caller's 400 and the transient 503, matching admission's 500-vs-503 split for the same conditions.""" + response, post = await _prepare_only_bridge_exchange("unresolvable") + assert response.status_code == 500 + assert json.loads(response.body)["error"] == "server_error" + post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_mint_upstream_expired_lifetime_is_502(): + """An upstream token response reporting an already-elapsed lifetime (a parseable non-positive + expires_in) is rejected with 502 rather than sealed into an hour-long envelope around a dead + bearer. Regression for expires_in<=0 silently falling through to the 1h cap.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 0} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +@pytest.mark.asyncio +async def test_bridge_mint_positive_sub_second_lifetime_mints_not_502(): + """A positive fractional expires_in in (0, 1) is a live token, not an elapsed one, so it mints a + (1s-floored) envelope rather than being truncated to 0 and rejected with 502 after the single-use + code was already consumed. Regression for classifying a sub-second remaining lifetime as expired.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 0.5} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 200 + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + assert body["expires_in"] >= 0 + + +@pytest.mark.asyncio +async def test_bridge_mint_unknown_lifetime_is_capped_not_rejected(): + """An absent or unparseable expires_in leaves the lifetime unknown, which the envelope caps (never + inventing a longer life than the upstream stated); it is NOT rejected. Only an explicitly-dead + lifetime fails, so a metadata glitch on an otherwise-valid token still mints a bounded envelope.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": "not-a-number"} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 200 + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + assert 0 < body["expires_in"] <= 3600 + + +@pytest.mark.asyncio +async def test_bridge_reported_expires_in_does_not_overstate_jwt_exp(): + """The reported expires_in is derived from the envelope JWT's second-truncated exp (rounding the + elapsed portion up), so the client is never told the bearer lives past the point admission expires + it. Regression for the sub-second overstatement of the raw (expires_at - now) delta.""" + import time + + import jwt as _jwt + + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 300} + before = int(time.time()) + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + body = json.loads(response.body) + claims = _jwt.decode(body["access_token"].removeprefix("llm_env_"), options={"verify_signature": False}) + # projecting the reported lifetime from a time no later than the mint must not exceed the JWT exp + assert before + body["expires_in"] <= claims["exp"] + + +def test_bridge_reported_expires_in_can_be_zero_at_jwt_exp_boundary(): + from datetime import datetime, timezone + + from fastapi.responses import JSONResponse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _BridgeMintReady, + _finish_bridge_mint, + ) + 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( + 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( + ready=ready, + mcp_server=_bridge_server(auth_type=MCPAuth.oauth_delegate), + token_response={"access_token": "UP", "expires_in": 1}, + now=datetime.fromtimestamp(100.25, tz=timezone.utc), + ) + + assert isinstance(response, JSONResponse) + assert json.loads(response.body)["expires_in"] == 0 + + +def test_classify_upstream_lifetime(): + """expires_in from an IdP may be an int, a float (3600.0), or a numeric string ("3600"); each + coerces to a positive number of seconds. Absent or unparseable input (bool, non-numeric, NaN/inf, + oversized) is "unspecified" so the envelope caps it, while a parseable non-positive value is + "expired": the upstream reporting an already-dead token, which the mint must reject rather than + silently give the 1h cap.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _classify_upstream_lifetime + + assert _classify_upstream_lifetime(300) == 300 + assert _classify_upstream_lifetime(300.0) == 300 + assert _classify_upstream_lifetime("300") == 300 + assert _classify_upstream_lifetime(" 300 ") == 300 + # explicit, parseable, non-positive -> the upstream says the token is already dead + assert _classify_upstream_lifetime(0) == "expired" + assert _classify_upstream_lifetime(-5) == "expired" + assert _classify_upstream_lifetime(-0.5) == "expired" + # a positive sub-second lifetime is alive, not elapsed; it clamps up to the envelope's 1s floor + # rather than truncating to 0 and being misread as expired + assert _classify_upstream_lifetime(0.5) == 1 + assert _classify_upstream_lifetime(0.001) == 1 + # a positive value >= 1 truncates toward zero (never overstating the stated lifetime) + assert _classify_upstream_lifetime(1.9) == 1 + # unknown lifetime -> cap (never invent a longer life than the upstream stated) + assert _classify_upstream_lifetime(None) == "unspecified" + assert _classify_upstream_lifetime(True) == "unspecified" + assert _classify_upstream_lifetime("nope") == "unspecified" + # hostile numerics must not raise (int(float(...)) can OverflowError) -> unspecified + assert _classify_upstream_lifetime("inf") == "unspecified" + assert _classify_upstream_lifetime("1e999") == "unspecified" + assert _classify_upstream_lifetime("-inf") == "unspecified" + assert _classify_upstream_lifetime("nan") == "unspecified" + assert _classify_upstream_lifetime(float("inf")) == "unspecified" + assert _classify_upstream_lifetime(10**400) == "unspecified" + + +def test_bridge_grant_honors_and_rejects_upstream_lifetime(): + """The grant validator honors a positive lifetime, leaves an unknown one None for the envelope to + cap, and rejects an explicitly-expired one with "expired_lifetime" so a dead upstream token is + never sealed into an hour-long envelope.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _bridge_grant_from_token_response + + def grant(v): + return _bridge_grant_from_token_response({"access_token": "x", "expires_in": v}) + + assert grant(300).expires_in == 300 + assert grant(120.0).expires_in == 120 + # unknown lifetime backs a grant whose expires_in the envelope caps; it is not a rejection + assert grant("nope").expires_in is None + assert _bridge_grant_from_token_response({"access_token": "x"}).expires_in is None + # an explicitly already-dead lifetime is rejected, not silently capped at 1h + assert grant(0) == "expired_lifetime" + assert grant(-5) == "expired_lifetime" + + +@pytest.mark.asyncio +async def test_bridge_token_exchange_honors_short_float_expires_in_ttl(): + """A short float expires_in from the upstream caps the envelope TTL, so the client-held envelope + does not outlive the upstream token. Before coercion a float was dropped and the envelope + defaulted to the 1h cap (3600), which would forward a stale bearer after the upstream token + expired.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 120.0} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert json.loads(response.body)["expires_in"] <= 120 + + +@pytest.mark.asyncio +async def test_oauth_delegate_bridge_token_exchange_missing_access_token_is_502_not_keyerror(): + """When the upstream token response has no access_token, a dcr_bridge oauth_delegate exchange + returns a clean 502 error body rather than raising a KeyError. _finish_bridge_mint asks + _bridge_grant_from_token_response for a typed grant, gets None, and returns the "no_upstream_token" + failure, which maps to 502; nothing indexes token_response["access_token"] on the bridge path.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"token_type": "Bearer", "expires_in": 3600} + + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +@pytest.mark.asyncio +async def test_true_passthrough_bridge_token_exchange_returns_raw_upstream_token(): + """Only oauth_delegate mints. A true_passthrough dcr_bridge server relays the raw upstream token + to the client, since that mode has no litellm identity to bind and the caller owns the token.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.true_passthrough) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "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"] == "UPSTREAM-SECRET-TOKEN" + assert not body["access_token"].startswith("llm_env_") + + +@pytest.mark.asyncio +async def test_non_bridge_oauth_delegate_token_exchange_returns_raw_upstream_token(): + """An oauth_delegate server without dcr_bridge keeps the pre-change contract: the raw upstream + token is returned, so flag-off behavior is byte-identical.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate, dcr_bridge=None) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "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"] == "UPSTREAM-SECRET-TOKEN" + + async def _exchange_persistence_attempted_for_auth_type(auth_type) -> bool: """Run exchange_token_with_server for a server of ``auth_type`` and report whether it attempted to persist the exchanged token server-side. The client-forwarded token modes must not persist: @@ -4580,7 +5829,7 @@ async def test_extract_user_id_reads_x_litellm_api_key_header(proxy_globals): """The LiteLLM key arrives on x-litellm-api-key (what Claude Desktop/Code send), not Authorization. Reading only Authorization dropped the identity, so the per-user token was never stored and the egress 401'd forever. Resolution must honor x-litellm-api-key.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth, hash_token @@ -4605,7 +5854,7 @@ async def test_extract_user_id_rehydrates_cross_replica_dict_cache(proxy_globals """Cross-replica, async_get_cache hands back a serialized dict, not a UserAPIKeyAuth. Resolution must rehydrate it; the old getattr(cached, "user_id") returned None on a dict, which is exactly why a multi-replica gateway never found the stored token.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import hash_token @@ -4626,7 +5875,7 @@ async def test_extract_user_id_falls_back_to_db_on_cache_miss(proxy_globals): """A cache miss must read the key from the DB rather than returning None; the old code did a cache-only peek and skipped the DB, so any replica that hadn't just authenticated the key failed to store the token.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth @@ -4648,7 +5897,7 @@ async def test_extract_user_id_falls_back_to_db_on_cache_miss(proxy_globals): @pytest.mark.asyncio async def test_extract_user_id_none_without_litellm_key(proxy_globals): """No LiteLLM key on the request resolves to None without consulting the resolver.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -4665,7 +5914,7 @@ async def test_extract_user_id_rejects_blocked_key(proxy_globals): """A blocked LiteLLM key must not resolve an identity. get_key_object returns the DB row without checking blocked/expiry (the main auth pipeline does, and the public token endpoint bypasses it), so a revoked key could otherwise overwrite the stored per-user OAuth token for its user.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth @@ -4687,7 +5936,7 @@ async def test_extract_user_id_rejects_expired_key(proxy_globals): """An expired LiteLLM key must not resolve an identity, for the same reason as a blocked key.""" from datetime import datetime, timedelta, timezone - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth @@ -4706,6 +5955,219 @@ async def test_extract_user_id_rejects_expired_key(proxy_globals): assert await _extract_user_id_from_request(request) is None +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_returns_resolved_key_for_active_key(proxy_globals): + """The dcr_bridge mint seals the hash of the authorizing key so admission can reload the live + record. For an active key the resolver returns exactly hash_token(key), the same value + get_key_object and the whole cache/DB layer key the record by, so the sealed reference resolves + back to this key at admission.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + _ResolvedKey, + ) + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key = "sk-alice-key" + cache = UserApiKeyCache() + await cache.async_set_cache( + hash_token(key), + UserAPIKeyAuth(token=hash_token(key), user_id="alice"), + model_type=UserAPIKeyAuth, + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = object() + + request = _token_request({"x-litellm-api-key": f"Bearer {key}"}) + resolved = await _resolve_active_litellm_key(request) + assert isinstance(resolved, _ResolvedKey) + assert resolved.key_hash == hash_token(key) + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_resolves_key_without_user_id(proxy_globals): + """A valid team-scoped or service-account key has no user_id but is a legitimate credential, so it + must still resolve to a hash and be able to mint a bridge envelope. Gating the resolver on user_id + presence wrongly rejected these keys with invalid_request; the active-state gate now checks only + blocked and expiry, and the key hash (not the user) is what the mint seals. The per-user token + store still gets no user for such a key, since there is none to key a stored credential by.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _extract_user_id_from_request, + _ResolvedKey, + _resolve_active_litellm_key, + ) + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key = "sk-team-scoped-key" + cache = UserApiKeyCache() + await cache.async_set_cache( + hash_token(key), + UserAPIKeyAuth(token=hash_token(key), user_id=None, team_id="team-x"), + model_type=UserAPIKeyAuth, + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = object() + + request = _token_request({"x-litellm-api-key": f"Bearer {key}"}) + resolved = await _resolve_active_litellm_key(request) + assert isinstance(resolved, _ResolvedKey) + assert resolved.key_hash == hash_token(key) + assert await _extract_user_id_from_request(request) is None + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_rejects_blocked_key(proxy_globals): + """A blocked key must not yield a hash, so no gateway-bound envelope is minted for a revoked key; + the mint fails closed with invalid_request instead.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _FakePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + return UserAPIKeyAuth(token=token, user_id="blocked-user", blocked=True) + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _FakePrisma() + + request = _token_request({"x-litellm-api-key": "sk-blocked-key"}) + assert await _resolve_active_litellm_key(request) == "no_active_key" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_fails_closed_on_malformed_expiry(proxy_globals): + """A key whose stored expires string does not parse must fail closed to no-hash (the mint then + returns invalid_request), not surface an unhandled 500. The active-state check runs outside the + resolver's try, so it must be total over a bad expires rather than letting datetime.fromisoformat + raise. Before the fix this raised a ValueError instead of returning None.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _FakePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + return UserAPIKeyAuth(token=token, user_id="u", expires="not-a-parseable-date") + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _FakePrisma() + + request = _token_request({"x-litellm-api-key": "sk-bad-expiry-key"}) + assert await _resolve_active_litellm_key(request) == "no_active_key" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_no_active_key_without_litellm_key(proxy_globals): + """No LiteLLM key on the request yields no hash without consulting the resolver.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + request = _token_request({"content-type": "application/json"}) + assert await _resolve_active_litellm_key(request) == "no_active_key" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_db_outage_is_unavailable(proxy_globals): + """A database outage while resolving the presented key is a retryable infrastructure failure, not + the caller's fault, so the resolver reports "unavailable" (the mint statuses it 503) rather than + collapsing it to the same value as a missing credential. is_database_service_unavailable_error + classifies a connection error (an OSError) as an outage, matching admission's egress-side handling.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _OutagePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + raise ConnectionError("connection refused") + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _OutagePrisma() + + request = _token_request({"x-litellm-api-key": "sk-during-outage"}) + assert await _resolve_active_litellm_key(request) == "unavailable" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_no_database_is_unresolvable(proxy_globals): + """With no database connection configured the gateway cannot verify the presented key at all, so + the resolver reports "unresolvable" (the mint statuses it 500) instead of blaming the caller. + Mirrors admission, which 500s a missing prisma_client on the egress side.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = None + + request = _token_request({"x-litellm-api-key": "sk-no-db"}) + 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.bridge_token_flow 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.bridge_token_flow 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 @@ -5306,3 +6768,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/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 32996905166..a983ac3ff48 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -5428,6 +5428,457 @@ def test_get_forwarded_auth_from_scope_skips_when_no_litellm_key_header(): assert _get_forwarded_auth_from_scope(scope) is None +def _delegate_auth_mcp_server(server_id: str = "delegate-1") -> MCPServer: + return MCPServer( + server_id=server_id, + name="delegate_test", + url="http://upstream:9401/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + oauth2_flow="authorization_code", + ) + + +def _delegate_scope(headers: list) -> dict: + return { + "type": "http", + "method": "POST", + "path": "/mcp/delegate_test", + "scheme": "http", + "server": ("localhost", 4000), + "headers": headers, + } + + +def _patch_delegate_resolver(server: MCPServer, *resolvable_names: str): + """Patch the admission-parity resolver the delegate probe gates on. Returns + ``server`` only for names admission's ``get_mcp_server_by_name`` would match + (alias / server_name / name); every other name (server_id, access group) yields + None, exactly as the real resolver does.""" + + def _resolve(name, client_ip=None): + return server if name in resolvable_names else None + + return patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + side_effect=_resolve, + ) + + +@pytest.mark.asyncio +async def test_delegate_bad_token_gets_connect_time_401(): + """Regression (LIT-4194): a rejected upstream token on a delegate-auth server + must fail the connect with 401 + ``error="invalid_token"``, not be absorbed + into HTTP 200 + an empty tool list by the tools/list handler. + + Delegate-mode clients send only ``Authorization`` (no ``x-litellm-api-key``), + so ``_get_forwarded_auth_from_scope`` returns None and, before the fix, the + preflight returned early without probing. + """ + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server() + scope = _delegate_scope([(b"authorization", b"Bearer bogus-token")]) + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, 'Bearer realm="upstream", error="invalid_token"')), + ) as probe: + with pytest.raises(HTTPException) as exc_info: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + assert exc_info.value.status_code == 401 + challenge = exc_info.value.headers["www-authenticate"] + assert 'error="invalid_token"' in challenge + assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge + probe.assert_awaited_once() + probe_url, probe_auth = probe.call_args.args + assert probe_url == "http://upstream:9401/mcp" + assert probe_auth == "Bearer bogus-token" + + +@pytest.mark.asyncio +async def test_delegate_valid_token_passes_preflight(): + """An upstream-accepted token must not be blocked by the delegate preflight.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server() + scope = _delegate_scope([(b"authorization", b"Bearer good-token")]) + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(200, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + probe.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_delegate_valid_token_forbidden_returns_403(): + """An upstream that accepts the token but forbids the caller (403) must surface + as a bare 403 with no ``WWW-Authenticate`` re-auth hint (a fresh token with the + same scopes would loop), not as an invalid_token challenge.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server() + scope = _delegate_scope([(b"authorization", b"Bearer scoped-out-token")]) + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(403, None)), + ): + with pytest.raises(HTTPException) as exc_info: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + assert exc_info.value.status_code == 403 + assert not (exc_info.value.headers or {}) + + +@pytest.mark.asyncio +async def test_delegate_tokenless_request_not_probed(): + """Tokenless delegate requests are the preemptive challenge's job; the + preflight must not probe upstream with an empty credential.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server() + scope = _delegate_scope([(b"content-type", b"application/json")]) + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delegate_preflight_skipped_on_multi_server_routes(): + """The delegate probe is gated to single-server routes so one rejected token + cannot 401 a multi-server aggregate connect (matching the OBO preflight).""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + servers = [_delegate_auth_mcp_server("delegate-1"), _delegate_auth_mcp_server("delegate-2")] + scope = _delegate_scope([(b"authorization", b"Bearer bogus-token")]) + + with _patch_delegate_resolver(servers[0], "delegate_test", "other_server"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=servers), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test", "other_server"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_bare_authorization_never_probes_passthrough_servers(): + """A bare ``Authorization`` header may be a LiteLLM key (backward-compat), so + only delegate servers (where admission classified it as an upstream token) + may be probed with it; ``is_oauth_passthrough`` servers still require the + unambiguous ``x-litellm-api-key`` + ``Authorization`` pair.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + passthrough_server = MCPServer( + server_id="pt-1", + name="pt_server", + url="http://upstream:9402/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + oauth_passthrough=True, + extra_headers=["Authorization"], + ) + scope = _delegate_scope([(b"authorization", b"Bearer ambiguous-token")]) + + with _patch_delegate_resolver(passthrough_server, "pt_server"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[passthrough_server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["pt_server"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delegate_not_probed_when_named_only_via_server_id(): + """Security regression (LIT-4194): a delegate server reachable by the requested + name only through its server_id (or an access group) is admitted as a real + LiteLLM key by ``process_mcp_request`` (its ``get_mcp_server_by_name`` misses), + so the bare ``Authorization`` header is that LiteLLM key. The probe must resolve + the target through the SAME resolver and therefore skip it, never forwarding the + key upstream, even though the widened allowed-server set still contains it.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server(server_id="delegate-secret-id") + # Admission's resolver matches alias/server_name/name only, never server_id: the + # requested server_id resolves to None here, mirroring the real divergence. + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/delegate-secret-id", + "scheme": "http", + "server": ("localhost", 4000), + "headers": [(b"authorization", b"Bearer sk-litellm-proxy-key")], + } + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(user_id="u1", api_key="hashed-sk"), + mcp_servers=["delegate-secret-id"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delegate_preflight_with_unpatched_probe(): + """Integration across the preflight and the unpatched ``_probe_upstream_auth``, + mocked only at the httpx-client boundary (tests/test_litellm is mocked-only; the + real-network proof lives in the PR's live-proxy evidence). The mock honors the + ``AsyncHTTPHandler.post`` contract by raising ``httpx.HTTPStatusError`` on the + upstream 401, so the production ``except httpx.HTTPStatusError`` branch is the one + exercised. A rejected token surfaces as the connect-time 401 challenge; an + accepted token passes untouched, and the caller's bearer reaches the delegate URL.""" + import httpx + + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + accepted = MagicMock() + accepted.status_code = 200 + accepted.headers = {} + rejected = MagicMock() + rejected.status_code = 401 + rejected.headers = {"www-authenticate": 'Bearer realm="stub-upstream", error="invalid_token"'} + + async def respond_by_token(url=None, headers=None, json=None, timeout=None, **kwargs): + if headers.get("Authorization") == "Bearer good-token": + return accepted + raise httpx.HTTPStatusError( + "401 Unauthorized", + request=httpx.Request("POST", url), + response=rejected, + ) + + mock_client = MagicMock() + mock_client.post = AsyncMock(side_effect=respond_by_token) + + server = _delegate_auth_mcp_server() + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server.get_async_httpx_client", + return_value=mock_client, + ): + with pytest.raises(HTTPException) as exc_info: + await _check_passthrough_upstream_auth( + scope=_delegate_scope([(b"authorization", b"Bearer bogus-token")]), + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + await _check_passthrough_upstream_auth( + scope=_delegate_scope([(b"authorization", b"Bearer good-token")]), + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + assert exc_info.value.status_code == 401 + challenge = exc_info.value.headers["www-authenticate"] + assert 'error="invalid_token"' in challenge + assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge + probed_urls = [call.kwargs["url"] for call in mock_client.post.await_args_list] + assert probed_urls == ["http://upstream:9401/mcp", "http://upstream:9401/mcp"] + + +@pytest.mark.asyncio +async def test_delegate_challenge_echoes_requested_alias(): + """An alias-routed delegate request must be probed, and the challenge must echo + the requested alias (not the canonical server name) so the resource_metadata + URL matches what the tokenless preemptive challenge emits for the same route.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server().model_copy(update={"alias": "dt-alias"}) + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/dt-alias", + "scheme": "http", + "server": ("localhost", 4000), + "headers": [(b"authorization", b"Bearer bogus-token")], + } + + with _patch_delegate_resolver(server, "dt-alias"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')), + ): + with pytest.raises(HTTPException) as exc_info: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["dt-alias"], + client_ip=None, + ) + + challenge = exc_info.value.headers["www-authenticate"] + assert 'error="invalid_token"' in challenge + assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/dt-alias"' in challenge + + +@pytest.mark.asyncio +async def test_delegate_probe_not_fanned_out_to_access_group_members(): + """A single access-group name passes the one-target route gate but must not fan + the delegate probe out to group-expanded member servers; the group name resolves + to no server under admission's resolver, so no probe fires.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + group_member = _delegate_auth_mcp_server() + + with _patch_delegate_resolver(group_member, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[group_member]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=_delegate_scope([(b"authorization", b"Bearer bogus-token")]), + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["prod_tools_group"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +def test_is_delegate_upstream_probe_target_fails_closed_on_m2m_shape(): + """An unstamped M2M-shape row (null ``oauth2_flow`` + client credentials) + resolves to ``client_credentials`` and must not be probed with the caller's + bearer; its stored client credentials drive egress instead.""" + from litellm.proxy._experimental.mcp_server.server import ( + _is_delegate_upstream_probe_target, + ) + + assert _is_delegate_upstream_probe_target(_delegate_auth_mcp_server()) is True + + m2m_shape = MCPServer( + server_id="delegate-m2m", + name="delegate_m2m", + url="http://upstream:9401/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + oauth2_flow=None, + token_url="http://idp:9000/token", + client_id="client", + client_secret="secret", + ) + assert _is_delegate_upstream_probe_target(m2m_shape) is False + + non_delegate = MCPServer( + server_id="oauth2-plain", + name="oauth2_plain", + url="http://upstream:9401/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + ) + assert _is_delegate_upstream_probe_target(non_delegate) is False + + @pytest.mark.asyncio async def test_create_mcp_client_sampling_disabled_by_default(): """Sampling callback must be None when allow_sampling is not set (default False).""" diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d12ff20ee5b..27f43c4948f 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""" @@ -4413,3 +4445,85 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): request=MagicMock(spec=Request), ) assert "User=u1" in str(over.value) + + +@pytest.mark.asyncio +async def test_user_budget_enforced_on_team_key(): + """User budget must be enforced even when the key belongs to a team. + + Previously _user_max_budget_check skipped enforcement for team keys, + letting a user with a $100 personal budget spend unlimited through a + team key. This regression test ensures that is no longer the case. + """ + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0) + team = LiteLLM_TeamTable(team_id="t1", max_budget=2100.0) + token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1") + + async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): + return 999.0 if counter_key == "spend:user:u1" else 0.0 + + async def _no_membership(*a, **kw): + return None + + with patch("litellm.proxy.proxy_server.prisma_client", None), patch( + "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): + with pytest.raises(litellm.BudgetExceededError) as over: + await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + assert "User=u1" in str(over.value) + + +@pytest.mark.asyncio +async def test_skip_user_budget_on_team_key_flag_restores_old_behavior(): + """Setting skip_user_budget_on_team_key=True skips user budget for team keys. + + This is the opt-in escape hatch that restores the legacy behavior where + user budgets were not enforced when the key belonged to a team. + """ + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0) + team = LiteLLM_TeamTable(team_id="t1", max_budget=2100.0) + token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1") + + async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): + return 999.0 if counter_key == "spend:user:u1" else 0.0 + + async def _no_membership(*a, **kw): + return None + + with patch("litellm.proxy.proxy_server.prisma_client", None), patch( + "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): + result = await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={"skip_user_budget_on_team_key": True}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + assert result is True 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/guardrails/guardrail_hooks/test_lasso.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py index 5a84b6ebecd..16185cadbdf 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -693,6 +693,8 @@ class TestLassoGuardrail: assert prompt_payload["messages"] == messages assert prompt_payload["userId"] == "test-user" assert prompt_payload["sessionId"] == "test-conversation" + # Every call is attributed to the "litellm" integration for the "Used By" badge. + assert prompt_payload["source"] == {"type": "litellm"} # Test COMPLETION payload completion_messages = [{"role": "assistant", "content": "Test response"}] @@ -703,6 +705,7 @@ class TestLassoGuardrail: assert completion_payload["messages"] == completion_messages assert completion_payload["userId"] == "test-user" assert completion_payload["sessionId"] == "test-conversation" + assert completion_payload["source"] == {"type": "litellm"} def test_header_preparation(self): """Test header preparation.""" diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 34d92505359..3dfb98c12ea 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -524,3 +524,42 @@ def test_apply_redacted_messages_back_skips_input_when_not_string(): data = {"input": [{"type": "text", "text": "leak"}]} apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) assert data["input"] == [{"type": "text", "text": "leak"}] + + +# ------------------------------------------------------------------- +# LIT-4302: custom_tool_call_output walking +# ------------------------------------------------------------------- + +def test_iter_message_text_walks_custom_tool_call_output(): + """custom_tool_call_output items should yield their output text.""" + data = { + "input": [ + {"type": "custom_tool_call_output", "output": "tool-secret"}, + ] + } + from litellm.proxy.guardrails._content_utils import iter_message_text + texts = list(iter_message_text(data)) + assert "tool-secret" in texts + + +def test_walk_user_text_redacts_custom_tool_call_output(): + """walk_user_text should rewrite text inside custom_tool_call_output.""" + data = { + "input": [ + {"type": "custom_tool_call_output", "output": "PII-data"}, + ] + } + count = walk_user_text(data, lambda t: t.replace("PII-data", "[MASKED]")) + assert count >= 1 + assert data["input"][0]["output"] == "[MASKED]" + + +def test_build_inspection_messages_custom_tool_call_output(): + """build_inspection_messages should include custom_tool_call_output text.""" + data = { + "input": [ + {"type": "custom_tool_call_output", "output": "custom-tool-leak"}, + ] + } + msgs = build_inspection_messages(data) + assert any("custom-tool-leak" in m["content"] for m in msgs) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 813a0c5e38f..f289148101a 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -14,6 +14,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( _ProxyDBLogger, _get_budget_reservation_from_metadata, + _should_track_cost_callback, _update_database_and_spend_counters, ) @@ -1177,3 +1178,88 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com" ) + + +@pytest.mark.parametrize( + "call_type, expected", + [ + ("pass_through_endpoint", True), + ("llm_passthrough_route", True), + ("allm_passthrough_route", True), + ("acompletion", False), + ("call_mcp_tool", False), + (None, False), + ], +) +def test_should_track_cost_callback_pass_through_without_owner(call_type, expected): + """Regression for LIT-3782: unauthenticated pass-through requests (auth=false) + carry no key/user/team/end-user, yet must still be tracked so they land in + LiteLLM_SpendLogs. Other call types with no owner stay untracked.""" + assert ( + _should_track_cost_callback( + user_api_key=None, + user_id=None, + team_id=None, + end_user_id=None, + call_type=call_type, + ) + is expected + ) + + +@pytest.mark.parametrize( + "call_type, expect_spend_log", + [ + ("pass_through_endpoint", True), + ("acompletion", False), + (None, False), + ], +) +@pytest.mark.asyncio +async def test_track_cost_callback_logs_unauthenticated_pass_through_request( + call_type, expect_spend_log +): + """Regression for LIT-3782: a pass-through request with auth=false reaches the + cost callback with no key/user/team/end-user. Before the fix the spend-log + write was skipped and the request never appeared in request/usage logs. It + must now be written for pass-through call types while other unauthenticated + calls remain skipped.""" + logger = _ProxyDBLogger() + + kwargs = { + "call_type": call_type, + "model": "unknown", + "litellm_params": {"metadata": {}}, + "standard_logging_object": { + "response_cost": 0.0, + "request_tags": None, + }, + "stream": False, + } + + with ( + patch( + "litellm.proxy.proxy_server.increment_spend_counters", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.proxy_server.update_cache", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging, + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( + 1 if expect_spend_log else 0 + ) 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/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 59b08a7ec4e..f7b2df45a85 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -18,6 +18,7 @@ sys.path.insert( from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth from litellm.proxy._types import ( LiteLLM_BudgetTableFull, + LiteLLM_ModelTable, LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, LiteLLM_OrganizationTableWithMembers, @@ -30,6 +31,7 @@ from litellm.proxy._types import ( ProxyErrorTypes, ProxyException, TeamMemberAddRequest, + UpdateTeamRequest, ) from litellm.proxy.management_endpoints.team_endpoints import ( user_api_key_auth, # Assuming this dependency is needed @@ -40,6 +42,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, + _update_model_table, _validate_and_populate_member_user_info, _verify_team_access, delete_team, @@ -9436,6 +9439,96 @@ async def test_team_info_forwards_key_limit_to_get_data(): assert mock_prisma.get_data.await_args.kwargs["limit"] == 7 +@pytest.mark.asyncio +async def test_team_info_returns_model_aliases(): + """/team/info must join LiteLLM_ModelTable so the response exposes the team's + current model aliases; without the ``litellm_model_table`` include the field + comes back null and the Admin UI can never display them. + """ + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + team_row = LiteLLM_TeamTable( + team_id="team-1", + litellm_model_table=LiteLLM_ModelTable( + id=1, + model_aliases={"gpt-4o": "gpt-4o-team-1"}, + created_by="admin", + updated_by="admin", + ), + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object( + team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[]) + ), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + include = mock_prisma.db.litellm_teamtable.find_unique.await_args.kwargs["include"] + assert include["litellm_model_table"] is True + + litellm_model_table = response["team_info"].litellm_model_table + assert litellm_model_table is not None + assert litellm_model_table.model_aliases == {"gpt-4o": "gpt-4o-team-1"} + + +@pytest.mark.asyncio +async def test_update_model_table_clears_aliases_with_empty_map(): + """``model_aliases={}`` on /team/update must persist an empty map (json.dumps({})) + so existing aliases are cleared, while ``model_aliases=None`` must be a no-op that + leaves the model table untouched. + """ + mock_prisma = MagicMock() + mock_prisma.db.litellm_modeltable.create = AsyncMock() + mock_prisma.db.litellm_modeltable.upsert = AsyncMock( + return_value=MagicMock(id="model-123") + ) + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" + ) + + returned_model_id = await _update_model_table( + data=UpdateTeamRequest(team_id="team-1", model_aliases={}), + model_id="model-123", + prisma_client=mock_prisma, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="default_user_id", + ) + + mock_prisma.db.litellm_modeltable.upsert.assert_awaited_once() + upsert_kwargs = mock_prisma.db.litellm_modeltable.upsert.await_args.kwargs + assert upsert_kwargs["where"] == {"id": "model-123"} + assert upsert_kwargs["data"]["update"]["model_aliases"] == json.dumps({}) + assert upsert_kwargs["data"]["create"]["model_aliases"] == json.dumps({}) + assert returned_model_id == "model-123" + + mock_prisma.db.litellm_modeltable.create.reset_mock() + mock_prisma.db.litellm_modeltable.upsert.reset_mock() + + noop_model_id = await _update_model_table( + data=UpdateTeamRequest(team_id="team-1", model_aliases=None), + model_id="model-123", + prisma_client=mock_prisma, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="default_user_id", + ) + + mock_prisma.db.litellm_modeltable.create.assert_not_called() + mock_prisma.db.litellm_modeltable.upsert.assert_not_called() + assert noop_model_id == "model-123" + + class TestEmitTeamMembersMetric: """The _emit_team_members_metric seam between the team handlers and Prometheus.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 045e15f8b8b..2a4e2ed6b25 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -461,6 +461,53 @@ async def test_get_group_ids_from_service_principal_uses_configured_graph_endpoi ] +@pytest.mark.asyncio +async def test_get_group_ids_from_service_principal_paginates_through_all_pages(): + # Arrange + page_one = { + "@odata.nextLink": "https://graph.microsoft.com/v1.0/servicePrincipals/sp-123/appRoleAssignedTo?$skiptoken=page2", + "value": [ + { + "principalType": "Group", + "principalId": "group-on-page-1", + "principalDisplayName": "Group On Page 1", + } + ], + } + page_two = { + "value": [ + { + "principalType": "Group", + "principalId": "group-on-page-2", + "principalDisplayName": "Group On Page 2", + } + ], + } + responses = [page_one, page_two] + + async def mock_get(url, *args, **kwargs): + mock = MagicMock() + mock.json.return_value = responses.pop(0) + return mock + + async_client = MagicMock() + async_client.get = mock_get + + # Act + group_ids, teams = await MicrosoftSSOHandler.get_group_ids_from_service_principal( + service_principal_id="sp-123", + async_client=async_client, + access_token="mock_token", + ) + + # Assert + assert group_ids == ["group-on-page-1", "group-on-page-2"] + assert [team["principalId"] for team in teams] == [ + "group-on-page-1", + "group-on-page-2", + ] + + def test_get_group_ids_from_graph_api_response(): # Arrange mock_response = MicrosoftGraphAPIUserGroupResponse( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 3da103683ba..540f017ee88 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -576,6 +576,83 @@ async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_ await release_budget_reservation(reservation) +@pytest.mark.asyncio +async def test_should_reserve_user_budget_counter_for_team_key(spend_counter_state): + """A user's personal budget must be reserved even when the key belongs to a team. + + Regression for GitHub issue #12905: previously the reservation path skipped the + user spend counter whenever the key had a team, so a team key could overshoot the + user's personal max_budget under concurrency. + """ + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-user-on-team", + spend=0.0, + user_id="user-on-team", + team_id="team-no-budget", + ) + team_object = LiteLLM_TeamTable(team_id="team-no-budget", spend=0.0, max_budget=None) + user_object = LiteLLM_UserTable(user_id="user-on-team", spend=0.0, max_budget=5.0) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.3, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=user_object, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team") == pytest.approx(0.3) + + await release_budget_reservation(reservation) + + +@pytest.mark.asyncio +async def test_should_skip_user_budget_counter_for_team_key_when_flag_set(spend_counter_state): + """skip_user_budget_on_team_key=True restores the legacy behavior where a user's + personal budget is not reserved for a team key.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-user-on-team-skip", + spend=0.0, + user_id="user-on-team-skip", + team_id="team-no-budget-skip", + ) + team_object = LiteLLM_TeamTable(team_id="team-no-budget-skip", spend=0.0, max_budget=None) + user_object = LiteLLM_UserTable(user_id="user-on-team-skip", spend=0.0, max_budget=5.0) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.3, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=user_object, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + skip_user_budget_on_team_key=True, + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team-skip") is None + + await release_budget_reservation(reservation) + + @pytest.mark.asyncio async def test_should_seed_org_counter_from_with_budget_cache(spend_counter_state): counter_cache, key_cache = spend_counter_state diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 603d5cc15b7..2f0924e9192 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8649,6 +8649,38 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): app.dependency_overrides.clear() +def test_get_config_list_includes_skip_user_budget_on_team_key(monkeypatch): + """Related to #12905: the opt-out flag must be discoverable via /config/list so + it renders as a Boolean toggle on the Admin UI General Settings table. This + requires both the ConfigGeneralSettings field and the allowed_args entry.""" + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + assert resp.status_code == 200, resp.text + fields = {item["field_name"]: item for item in resp.json()} + assert "skip_user_budget_on_team_key" in fields + assert fields["skip_user_budget_on_team_key"]["field_type"] == "Boolean" + finally: + app.dependency_overrides.clear() + + def test_get_config_list_includes_budget_exceeded_throttle_percentage(monkeypatch): """The throttle fraction is a litellm_settings scalar surfaced on the General Settings table as a Float field so it sits with the other global limits; it diff --git a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py index cdace5f6327..24edf12fffe 100644 --- a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py @@ -39,8 +39,13 @@ def _output_item_added_chunk(): return SimpleNamespace(type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED) -def _completed_chunk(output): - response = ResponsesAPIResponse(id="resp-1", created_at=0, output=output) +def _created_chunk(response_id: str): + response = ResponsesAPIResponse(id=response_id, created_at=0, output=[]) + return SimpleNamespace(type=ResponsesAPIStreamEvents.RESPONSE_CREATED, response=response) + + +def _completed_chunk(output, response_id: str = "resp-1"): + response = ResponsesAPIResponse(id=response_id, created_at=0, output=output) return SimpleNamespace(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=response) @@ -52,12 +57,12 @@ def _text_message(text: str): return {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]} -def _tool_call_stream(call_id: str, tool_name: str) -> _FakeAsyncStream: - return _FakeAsyncStream([_completed_chunk([_function_call(call_id, tool_name)])]) +def _tool_call_stream(call_id: str, tool_name: str, response_id: str = "resp-1") -> _FakeAsyncStream: + return _FakeAsyncStream([_completed_chunk([_function_call(call_id, tool_name)], response_id=response_id)]) -def _text_only_stream(text: str) -> _FakeAsyncStream: - return _FakeAsyncStream([_completed_chunk([_text_message(text)])]) +def _text_only_stream(text: str, response_id: str = "resp-1") -> _FakeAsyncStream: + return _FakeAsyncStream([_completed_chunk([_text_message(text)], response_id=response_id)]) def _mock_mcp_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: @@ -170,3 +175,85 @@ async def test_tool_call_rounds_are_capped(monkeypatch): for call in aresponses_mock.call_args_list[:-1]: assert "tools" in call.kwargs assert "tools" not in aresponses_mock.call_args_list[-1].kwargs + + +@pytest.mark.asyncio +async def test_continuation_id_is_final_round_not_interim_tool_call(monkeypatch): + """ + Regression test for the broken `previous_response_id` continuation after a + gateway-executed MCP tool call. Each auto-execute round is a distinct + upstream response: the interim round holds only the model's function_call + (no tool output), the final round holds the answer. The client must be + handed the FINAL round's response id, because that is the one whose stored + chain includes the function_call_output. Pinning every event to the interim + round's id made the next turn continue from a response with a dangling + function_call, which the provider rejects with + "No tool output found for function call ...". + """ + _mock_mcp_environment(monkeypatch) + + aresponses_mock = AsyncMock(side_effect=[_text_only_stream("The first item is Alpha.", response_id="resp-final")]) + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + + iterator = _make_iterator( + [ + _created_chunk("resp-interim"), + _output_item_added_chunk(), + _completed_chunk([_function_call("call_1", "read_wiki_contents")], response_id="resp-interim"), + ] + ) + + chunks = [chunk async for chunk in iterator] + completed = [c for c in chunks if getattr(c, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED] + + assert completed[-1].response.output[0]["content"][0]["text"] == "The first item is Alpha." + assert completed[-1].response.id == "resp-final" + assert completed[-1].response.id != "resp-interim" + + +@pytest.mark.asyncio +async def test_follow_up_call_failure_emits_terminal_error_event(monkeypatch): + """ + Regression test for the silent-swallow path: when the follow-up LLM call + raises, the stream must emit a terminal `error` event instead of ending + silently after the tool events (which surfaced to clients as a successful + but empty completion). + """ + _mock_mcp_environment(monkeypatch) + + aresponses_mock = AsyncMock(side_effect=RuntimeError("boom from provider")) + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + + iterator = _make_iterator( + [ + _output_item_added_chunk(), + _completed_chunk([_function_call("call_1", "read_wiki_contents")]), + ] + ) + + chunks = [chunk async for chunk in iterator] + error_events = [c for c in chunks if getattr(c, "type", None) == ResponsesAPIStreamEvents.ERROR] + + assert len(error_events) == 1 + assert error_events[0].error.type == "mcp_gateway_error" + assert "boom from provider" in error_events[0].error.message + + +@pytest.mark.asyncio +async def test_initial_call_failure_is_stashed_for_eager_reraise(monkeypatch): + """ + Regression test: a failing initial LLM call must be stashed as + `_initial_creation_error` so `aresponses_api_with_mcp` can re-raise it as a + real 4xx before any SSE bytes are written, instead of returning HTTP 200 + with an empty stream. + """ + _mock_mcp_environment(monkeypatch) + + aresponses_mock = AsyncMock(side_effect=RuntimeError("initial boom")) + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + + iterator = _make_iterator([_output_item_added_chunk()]) + await iterator._create_initial_response_iterator() + + assert iterator._initial_creation_error is not None + assert "initial boom" in str(iterator._initial_creation_error) 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_claude_opus_4_8_config.py b/tests/test_litellm/test_claude_opus_4_8_config.py index 32f7d249e05..8eead8a9c84 100644 --- a/tests/test_litellm/test_claude_opus_4_8_config.py +++ b/tests/test_litellm/test_claude_opus_4_8_config.py @@ -95,6 +95,8 @@ def test_opus_4_8_model_pricing_and_capabilities(): assert info["supports_tool_choice"] is True assert info["supports_vision"] is True + assert model_data["claude-opus-4-8"]["supports_native_structured_output"] is True + def test_opus_4_8_bedrock_regional_model_pricing(): model_data = _load_root_cost_map() @@ -165,6 +167,7 @@ def test_opus_4_8_present_in_bundled_backup(): "azure_ai/claude-opus-4-8", ): assert model_name in backup, f"Missing from backup cost map: {model_name}" + assert backup["claude-opus-4-8"]["supports_native_structured_output"] is True def test_opus_4_8_registered_for_bedrock_converse(): 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/helpers/navigation.ts b/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts index 6b5e7f7bb8d..a58ece16f9c 100644 --- a/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts +++ b/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts @@ -41,7 +41,7 @@ export async function dismissFeedbackPopup(page: PlaywrightPage): Promise /** * Click on a team ID in the table. Team IDs are rendered differently depending * on the component version — try button first (Tremor Button), fall back to - * clickable span (OldTeams Typography.Text). + * clickable span (Teams Typography.Text). */ export async function clickTeamId(page: PlaywrightPage, teamId: string): Promise { const cell = page.locator("td").filter({ hasText: teamId }).first(); 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 3f6f4da25d0..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 @@ -1400,22 +1546,6 @@ "count": 1 } }, - "src/components/OldTeams.test.tsx": { - "max-nested-callbacks": { - "count": 4 - } - }, - "src/components/OldTeams.tsx": { - "no-nested-ternary": { - "count": 4 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 4 - } - }, "src/components/SCIM.tsx": { "no-restricted-imports": { "count": 1 @@ -1506,6 +1636,22 @@ "count": 1 } }, + "src/components/Teams.test.tsx": { + "max-nested-callbacks": { + "count": 4 + } + }, + "src/components/Teams.tsx": { + "no-nested-ternary": { + "count": 2 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 3 + } + }, "src/components/ToolDetail.tsx": { "unused-imports/no-unused-imports": { "count": 2 @@ -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 ? ( ; + return ; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index a95d48aa75b..d2f75609d18 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -8,6 +8,7 @@ import { DownOutlined, ExportOutlined, InfoCircleOutlined, LoadingOutlined, RightOutlined } from "@ant-design/icons"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { Card, Col, @@ -94,7 +95,7 @@ const UsagePage: React.FC = ({ 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/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx deleted file mode 100644 index 0b6e5786aaf..00000000000 --- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx +++ /dev/null @@ -1,1148 +0,0 @@ -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; -import React from "react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; -import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; -import { fetchMCPAccessGroups, getGuardrailsList, teamCreateCall } from "./networking"; -import OldTeams from "./OldTeams"; -import { teamListCall } from "@/app/(dashboard)/hooks/teams/useTeams"; - -const mockTeamInfoView = vi.fn(); -const mockUseOrganizations = vi.fn(); - -vi.mock("./networking", () => ({ - teamCreateCall: vi.fn(), - teamDeleteCall: vi.fn(), - fetchMCPAccessGroups: vi.fn(), - v2TeamListCall: vi.fn(), - getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), - getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), -})); - -vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ - teamListCall: vi.fn().mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 0 }), -})); - -vi.mock("./molecules/notifications_manager", () => ({ - default: { - info: vi.fn(), - success: vi.fn(), - error: vi.fn(), - fromBackend: vi.fn(), - }, -})); - -vi.mock("./key_team_helpers/fetch_available_models_team_key", () => ({ - fetchAvailableModelsForTeamOrKey: vi.fn(), - getModelDisplayName: vi.fn((model: string) => model), - unfurlWildcardModelsInList: vi.fn((teamModels: string[], allModels: string[]) => { - const wildcardDisplayNames: string[] = []; - const expandedModels: string[] = []; - - teamModels.forEach((teamModel) => { - if (teamModel.endsWith("/*")) { - const provider = teamModel.replace("/*", ""); - const matchingModels = allModels.filter((model) => model.startsWith(provider + "/")); - expandedModels.push(...matchingModels); - wildcardDisplayNames.push(teamModel); - } else { - expandedModels.push(teamModel); - } - }); - - return [...wildcardDisplayNames, ...expandedModels].filter((item, index, array) => array.indexOf(item) === index); - }), -})); - -vi.mock("@/components/team/TeamInfo", () => ({ - __esModule: true, - default: (props: any) => { - mockTeamInfoView(props); - return
; - }, -})); - -vi.mock("./ModelSelect/ModelSelect", () => { - const ModelSelect = React.forwardRef(({ value, onChange, dataTestId, id }: any, ref: any) => { - return ( - { - if (onChange) { - const newVal = e.target.value - ? e.target.value - .split(",") - .map((s: string) => s.trim()) - .filter(Boolean) - : []; - onChange(newVal); - } - }} - /> - ); - }); - ModelSelect.displayName = "ModelSelect"; - return { - ModelSelect, - }; -}); - -vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ - useOrganizations: () => mockUseOrganizations(), -})); - -vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ - useAccessGroups: vi.fn().mockReturnValue({ - data: [ - { access_group_id: "ag-1", access_group_name: "Group 1" }, - { access_group_id: "ag-2", access_group_name: "Group 2" }, - ], - isLoading: false, - isError: false, - }), -})); - -vi.mock("./common_components/AccessGroupSelector", () => ({ - default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( - onChange?.(e.target.value ? e.target.value.split(",").map((s) => s.trim()) : [])} - /> - ), -})); - -const createQueryClient = () => { - return new QueryClient({ - defaultOptions: { - queries: { - retry: false, - }, - }, - }); -}; - -const renderWithQueryClient = (component: React.ReactElement) => { - const queryClient = createQueryClient(); - return render({component}); -}; - -describe("OldTeams - handleCreate organization handling", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockTeamInfoView.mockClear(); - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue([]); - vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); - vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); - mockUseOrganizations.mockReturnValue({ data: null }); - }); - - 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 - models: [], - }; - - // Simulate the handleCreate logic - let organizationId = formValues?.organization_id || null; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - expect(formValues.organization_id).toBeNull(); - expect(formValues.organization_id).not.toBe(""); - }); - - it("should set organization_id to null when it's not a string type", async () => { - const formValues: Record = { - team_alias: "Test Team", - organization_id: undefined, - models: [], - }; - - // Simulate the handleCreate logic - let organizationId = formValues?.organization_id || null; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - expect(formValues.organization_id).toBeNull(); - }); - - it("should trim and keep valid organization_id string", async () => { - const formValues: Record = { - team_alias: "Test Team", - organization_id: " org-123 ", // String with whitespace - models: [], - }; - - // Simulate the handleCreate logic - let organizationId = formValues?.organization_id || null; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - expect(formValues.organization_id).toBe("org-123"); - }); - - it("should keep valid organization_id without modification", async () => { - const formValues: Record = { - team_alias: "Test Team", - organization_id: "f874bb43-b898-4813-beca-4054d224eafc", - models: [], - }; - - // Simulate the handleCreate logic - let organizationId = formValues?.organization_id || null; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - expect(formValues.organization_id).toBe("f874bb43-b898-4813-beca-4054d224eafc"); - }); - - it("should not send organization_id field when converting empty string to null", async () => { - const formValues: Record = { - team_alias: "Test Team", - organization_id: "", - models: ["gpt-4"], - max_budget: 100, - }; - - // Simulate the handleCreate logic - let organizationId = formValues?.organization_id || null; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - 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(); - }); - - it("should handle when currentOrg is used as fallback", async () => { - const currentOrg = { - organization_id: "fallback-org-id", - organization_alias: "Fallback Org", - models: [], - members: [], - }; - - const formValues: Record = { - team_alias: "Test Team", - 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; - } else { - formValues.organization_id = organizationId.trim(); - } - - 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 () => { - 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); - }); - expect(screen.getByText("Delete Team?")).toBeInTheDocument(); - }); -}); - -describe("OldTeams - 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("OldTeams - 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: [], - }, - ]; - - // Simulate getAdminOrganizations logic for Admin - const userRole = "Admin"; - const result = userRole === "Admin" ? organizations : []; - - expect(result).toEqual(organizations); - expect(result.length).toBe(2); - }); - - 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", - organization_alias: "Org 1", - models: [], - members: [{ user_id: "user-123", user_role: "org_admin" }], - }, - { - organization_id: "org-2", - organization_alias: "Org 2", - models: [], - members: [{ user_id: "user-456", user_role: "org_admin" }], - }, - { - organization_id: "org-3", - organization_alias: "Org 3", - models: [], - members: [{ user_id: "user-123", user_role: "member" }], - }, - ]; - - // Simulate getAdminOrganizations logic - const result = organizations.filter((org) => - org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), - ); - - expect(result.length).toBe(1); - expect(result[0].organization_id).toBe("org-1"); - }); - - it("should return empty array when user is not admin of any organization", () => { - const userID = "user-999"; - const organizations = [ - { - organization_id: "org-1", - organization_alias: "Org 1", - models: [], - members: [{ user_id: "user-123", user_role: "org_admin" }], - }, - ]; - - // Simulate getAdminOrganizations logic - const result = organizations.filter((org) => - org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), - ); - - expect(result.length).toBe(0); - }); - }); - - describe("canCreateOrManageTeams", () => { - it("should return true for Admin role", () => { - const userRole = "Admin"; - const result = userRole === "Admin"; - expect(result).toBe(true); - }); - - it("should return true for org_admin in any organization", () => { - const userID = "user-123"; - const organizations = [ - { - organization_id: "org-1", - organization_alias: "Org 1", - models: [], - members: [{ user_id: "user-123", user_role: "org_admin" }], - }, - ]; - - const result = organizations.some((org) => - org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), - ); - - expect(result).toBe(true); - }); - - it("should return false when user has no admin permissions", () => { - const userID = "user-123"; - const userRole: string = "User"; - const organizations = [ - { - organization_id: "org-1", - organization_alias: "Org 1", - models: [], - members: [{ user_id: "user-123", user_role: "member" }], - }, - ]; - - const isAdmin = userRole === "Admin"; - const isOrgAdmin = organizations.some((org) => - org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), - ); - - expect(isAdmin || isOrgAdmin).toBe(false); - }); - }); -}); - -describe("OldTeams - premium props", () => { - beforeEach(() => { - mockTeamInfoView.mockClear(); - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue([]); - vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); - vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); - 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, - }); - renderWithQueryClient(); - - const teamIdElement = await screen.findByText("team-123456789"); - act(() => { - fireEvent.click(teamIdElement); - }); - - await waitFor(() => expect(mockTeamInfoView).toHaveBeenCalled()); - - expect(mockTeamInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ premiumUser: true })); - }); -}); - -describe("OldTeams - Default Team Settings tab visibility", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - 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(); - }); -}); - -describe("OldTeams - access_group_ids in team create", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockTeamInfoView.mockClear(); - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); - vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); - vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); - vi.mocked(teamCreateCall).mockResolvedValue({ - team_id: "new-team-1", - team_alias: "Test Team", - models: ["gpt-4"], - organization_id: null, - keys: [], - members_with_roles: [], - spend: 0, - }); - mockUseOrganizations.mockReturnValue({ - data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }], - }); - }); - - 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]; - act(() => { - fireEvent.click(createButton); - }); - - await waitFor(() => { - expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); - }); - - const teamNameInput = screen.getByLabelText(/team name/i); - fireEvent.change(teamNameInput, { target: { value: "Test Team" } }); - - const modelsInput = screen.getByTestId("create-team-models-select"); - fireEvent.change(modelsInput, { target: { value: "gpt-4" } }); - - const additionalSettingsAccordion = screen.getByText("Additional Settings"); - fireEvent.click(additionalSettingsAccordion); - - 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" } }); - - const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); - const createTeamSubmitButton = createTeamSubmitButtons[createTeamSubmitButtons.length - 1]; - fireEvent.click(createTeamSubmitButton); - - await waitFor(() => { - expect(teamCreateCall).toHaveBeenCalledWith( - "test-token", - expect.objectContaining({ - team_alias: "Test Team", - models: ["gpt-4"], - access_group_ids: ["ag-1", "ag-2"], - }), - ); - }); - }); -}); - -describe("OldTeams - models dropdown options", () => { - beforeEach(() => { - vi.clearAllMocks(); - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - 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(() => { - expect(fetchAvailableModelsForTeamOrKey).toHaveBeenCalled(); - }); - - const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; - act(() => { - fireEvent.click(createButton); - }); - - await waitFor(() => { - expect(screen.getByLabelText(/models/i)).toBeInTheDocument(); - }); - const allProxyModelsOption = screen.queryByText("All Proxy Models"); - expect(allProxyModelsOption).not.toBeInTheDocument(); - }); -}); - -describe("OldTeams - 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("OldTeams - 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"); - }); -}); - -describe("OldTeams - delete team warning copy", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - 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")); - }); - 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 }); - - expect(screen.getByText(/Warning: This team has 5 keys associated with it/i)).toHaveTextContent( - /along with any models created for this team/i, - ); - expect(screen.getByText(/Are you sure you want to delete this team/i)).toHaveTextContent( - /any models created for it/i, - ); - }); - - it("still warns about model deletion in the confirmation message when the team has no keys", async () => { - await openDeleteModal({ ...baseTeam, 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( - /any models created for it/i, - ); - }); -}); - -describe("OldTeams - LIT-2530 organization stays optional for proxy admin with a single org", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockTeamInfoView.mockClear(); - 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", - models: ["gpt-4"], - organization_id: null, - keys: [], - members_with_roles: [], - spend: 0, - }); - mockUseOrganizations.mockReturnValue({ - data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }], - }); - }); - - it("creates a team with no organization when exactly one organization exists", async () => { - renderWithQueryClient(); - - const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; - act(() => { - fireEvent.click(createButton); - }); - - await waitFor(() => { - expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); - }); - - fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "No Org Team" } }); - fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); - - const submitButtons = screen.getAllByRole("button", { name: /create team/i }); - fireEvent.click(submitButtons[submitButtons.length - 1]); - - await waitFor(() => { - expect(teamCreateCall).toHaveBeenCalledWith( - "test-token", - expect.objectContaining({ team_alias: "No Org Team", organization_id: null }), - ); - }); - }); -}); diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx new file mode 100644 index 00000000000..7065b1a5fb6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -0,0 +1,661 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; +import { fetchMCPAccessGroups, getGuardrailsList, teamCreateCall } from "./networking"; +import Teams from "./Teams"; + +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(), + fetchMCPAccessGroups: vi.fn(), + v2TeamListCall: vi.fn(), + getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), + 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", () => ({ + teamsTableKeys: { all: ["teamsTable"] }, +})); + +vi.mock("./molecules/notifications_manager", () => ({ + default: { + info: vi.fn(), + success: vi.fn(), + error: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +vi.mock("./key_team_helpers/fetch_available_models_team_key", () => ({ + fetchAvailableModelsForTeamOrKey: vi.fn(), + getModelDisplayName: vi.fn((model: string) => model), + unfurlWildcardModelsInList: vi.fn((teamModels: string[], allModels: string[]) => { + const wildcardDisplayNames: string[] = []; + const expandedModels: string[] = []; + + teamModels.forEach((teamModel) => { + if (teamModel.endsWith("/*")) { + const provider = teamModel.replace("/*", ""); + const matchingModels = allModels.filter((model) => model.startsWith(provider + "/")); + expandedModels.push(...matchingModels); + wildcardDisplayNames.push(teamModel); + } else { + expandedModels.push(teamModel); + } + }); + + return [...wildcardDisplayNames, ...expandedModels].filter((item, index, array) => array.indexOf(item) === index); + }), +})); + +vi.mock("@/components/team/TeamInfo", () => ({ + __esModule: true, + default: (props: any) => { + mockTeamInfoView(props); + return
; + }, +})); + +vi.mock("./ModelSelect/ModelSelect", () => { + const ModelSelect = React.forwardRef(({ value, onChange, dataTestId, id }: any, ref: any) => { + return ( + { + if (onChange) { + const newVal = e.target.value + ? e.target.value + .split(",") + .map((s: string) => s.trim()) + .filter(Boolean) + : []; + onChange(newVal); + } + }} + /> + ); + }); + ModelSelect.displayName = "ModelSelect"; + return { + ModelSelect, + }; +}); + +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganizations: () => mockUseOrganizations(), +})); + +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: vi.fn().mockReturnValue({ + data: [ + { access_group_id: "ag-1", access_group_name: "Group 1" }, + { access_group_id: "ag-2", access_group_name: "Group 2" }, + ], + isLoading: false, + isError: false, + }), +})); + +vi.mock("./common_components/AccessGroupSelector", () => ({ + default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( + onChange?.(e.target.value ? e.target.value.split(",").map((s) => s.trim()) : [])} + /> + ), +})); + +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: { + queries: { + retry: false, + }, + }, + }); +}; + +const renderWithQueryClient = (component: React.ReactElement) => { + const queryClient = createQueryClient(); + 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: [] }); + mockUseOrganizations.mockReturnValue({ data: null }); + }); + + it("should not include organization_id when it's an empty string", async () => { + const formValues: Record = { + team_alias: "Test Team", + organization_id: "", // Empty string + models: [], + }; + + // Simulate the handleCreate logic + let organizationId = formValues?.organization_id || null; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues.organization_id).toBeNull(); + expect(formValues.organization_id).not.toBe(""); + }); + + it("should set organization_id to null when it's not a string type", async () => { + const formValues: Record = { + team_alias: "Test Team", + organization_id: undefined, + models: [], + }; + + let organizationId = formValues?.organization_id || null; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues.organization_id).toBeNull(); + }); + + it("should trim and keep valid organization_id string", async () => { + const formValues: Record = { + team_alias: "Test Team", + organization_id: " org-123 ", + models: [], + }; + + let organizationId = formValues?.organization_id || null; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues.organization_id).toBe("org-123"); + }); + + it("should keep valid organization_id without modification", async () => { + const formValues: Record = { + team_alias: "Test Team", + organization_id: "f874bb43-b898-4813-beca-4054d224eafc", + models: [], + }; + + let organizationId = formValues?.organization_id || null; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues.organization_id).toBe("f874bb43-b898-4813-beca-4054d224eafc"); + }); + + it("should not send organization_id field when converting empty string to null", async () => { + const formValues: Record = { + team_alias: "Test Team", + organization_id: "", + models: ["gpt-4"], + max_budget: 100, + }; + + let organizationId = formValues?.organization_id || null; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues).toEqual({ + team_alias: "Test Team", + organization_id: null, + models: ["gpt-4"], + max_budget: 100, + }); + expect(formValues.organization_id).not.toBe(""); + expect(formValues.organization_id).toBeNull(); + }); + + it("should handle when currentOrg is used as fallback", async () => { + const currentOrg = { + organization_id: "fallback-org-id", + organization_alias: "Fallback Org", + models: [], + members: [], + }; + + const formValues: Record = { + team_alias: "Test Team", + models: [], + }; + + let organizationId = formValues?.organization_id || currentOrg?.organization_id; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues.organization_id).toBe("fallback-org-id"); + }); + + it("opens the delete modal when the table's delete action fires", async () => { + mockUseOrganizations.mockReturnValue({ data: [] }); + renderWithQueryClient(); + + await waitFor(() => expect(mockTeamsTableProps).not.toBeNull()); + await act(async () => { + mockTeamsTableProps.onDeleteTeam(baseTableTeam); + }); + + expect(screen.getByText("Delete Team?")).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: [] }, + ]; + + const userRole = "Admin"; + const result = userRole === "Admin" ? organizations : []; + + expect(result).toEqual(organizations); + expect(result.length).toBe(2); + }); + + it("should return only org_admin organizations for Org Admin role", () => { + const userID = "user-123"; + const organizations = [ + { + organization_id: "org-1", + organization_alias: "Org 1", + models: [], + members: [{ user_id: "user-123", user_role: "org_admin" }], + }, + { + organization_id: "org-2", + organization_alias: "Org 2", + models: [], + members: [{ user_id: "user-456", user_role: "org_admin" }], + }, + { + organization_id: "org-3", + organization_alias: "Org 3", + models: [], + members: [{ user_id: "user-123", user_role: "member" }], + }, + ]; + + const result = organizations.filter((org) => + org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), + ); + + expect(result.length).toBe(1); + expect(result[0].organization_id).toBe("org-1"); + }); + + it("should return empty array when user is not admin of any organization", () => { + const userID = "user-999"; + const organizations = [ + { + organization_id: "org-1", + organization_alias: "Org 1", + models: [], + members: [{ user_id: "user-123", user_role: "org_admin" }], + }, + ]; + + const result = organizations.filter((org) => + org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), + ); + + expect(result.length).toBe(0); + }); + }); + + describe("canCreateOrManageTeams", () => { + it("should return true for Admin role", () => { + const userRole = "Admin"; + expect(userRole === "Admin").toBe(true); + }); + + it("should return true for org_admin in any organization", () => { + const userID = "user-123"; + const organizations = [ + { + organization_id: "org-1", + organization_alias: "Org 1", + models: [], + members: [{ user_id: "user-123", user_role: "org_admin" }], + }, + ]; + + const result = organizations.some((org) => + org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), + ); + + expect(result).toBe(true); + }); + + it("should return false when user has no admin permissions", () => { + const userID = "user-123"; + const userRole: string = "User"; + const organizations = [ + { + organization_id: "org-1", + organization_alias: "Org 1", + models: [], + members: [{ user_id: "user-123", user_role: "member" }], + }, + ]; + + const isAdmin = userRole === "Admin"; + const isOrgAdmin = organizations.some((org) => + org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), + ); + + expect(isAdmin || isOrgAdmin).toBe(false); + }); + }); +}); + +describe("Teams - premium props", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue([]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + it("passes premiumUser flag to TeamInfoView when a team is opened", async () => { + const premiumTeam = { ...baseTableTeam, team_id: "team-123456789", team_alias: "Premium Team" }; + renderWithQueryClient(); + + await waitFor(() => expect(mockTeamsTableProps).not.toBeNull()); + act(() => mockTeamsTableProps.onSelectTeam(premiumTeam)); + + await waitFor(() => expect(mockTeamInfoView).toHaveBeenCalled()); + expect(mockTeamInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ premiumUser: true })); + }); +}); + +describe("Teams - Create Team CTA is grouped with the tabs on the left", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + it("renders the Create Team button inside the tab bar, ahead of the tabs", () => { + const { container } = renderWithQueryClient(); + + const createButton = screen.getByTestId("create-team-button"); + const tabNav = container.querySelector(".ant-tabs-nav"); + + // The CTA lives in the tab bar's left slot, not the standalone page header. + expect(tabNav).not.toBeNull(); + expect(tabNav!.contains(createButton)).toBe(true); + + // It reads as the left end of the cluster: it precedes the first tab in DOM order. + const firstTab = screen.getByRole("tab", { name: "Your Teams" }); + expect(createButton.compareDocumentPosition(firstTab) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + }); + + it("omits the Create Team CTA for a role that cannot manage teams", () => { + renderWithQueryClient(); + expect(screen.queryByTestId("create-team-button")).not.toBeInTheDocument(); + }); +}); + +describe("Teams - Default Team Settings tab visibility", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + it("should show Default Team Settings tab for Admin role", () => { + renderWithQueryClient(); + expect(screen.getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument(); + }); + + it("should show Default Team Settings tab for proxy_admin role", () => { + renderWithQueryClient(); + expect(screen.getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument(); + }); + + it("should not show Default Team Settings tab for proxy_admin_viewer role", () => { + renderWithQueryClient(); + expect(screen.queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument(); + }); + + it("should not show Default Team Settings tab for Admin Viewer role", () => { + renderWithQueryClient(); + expect(screen.queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument(); + }); +}); + +describe("Teams - access_group_ids in team create", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(teamCreateCall).mockResolvedValue({ + team_id: "new-team-1", + team_alias: "Test Team", + models: ["gpt-4"], + organization_id: null, + keys: [], + members_with_roles: [], + spend: 0, + }); + mockUseOrganizations.mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }], + }); + }); + + it("should pass access_group_ids to teamCreateCall when creating team", async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Test Team" } }); + fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); + + fireEvent.click(screen.getByText("Additional Settings")); + + await waitFor(() => { + expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + team_alias: "Test Team", + models: ["gpt-4"], + access_group_ids: ["ag-1", "ag-2"], + }), + ); + }); + }); +}); + +describe("Teams - models dropdown options", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + it("should not render all-proxy-models option in models select", async () => { + renderWithQueryClient(); + + await waitFor(() => { + expect(fetchAvailableModelsForTeamOrKey).toHaveBeenCalled(); + }); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/models/i)).toBeInTheDocument(); + }); + expect(screen.queryByText("All Proxy Models")).not.toBeInTheDocument(); + }); +}); + +describe("Teams - delete team warning copy", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + const openDeleteModal = async (team: any) => { + renderWithQueryClient(); + await waitFor(() => expect(mockTeamsTableProps).not.toBeNull()); + await act(async () => { + mockTeamsTableProps.onDeleteTeam(team); + }); + expect(screen.getByText("Delete Team?")).toBeInTheDocument(); + }; + + it("warns that the team's models are deleted when the team has keys", async () => { + 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, + ); + expect(screen.getByText(/Are you sure you want to delete this team/i)).toHaveTextContent( + /any models created for it/i, + ); + }); + + it("still warns about model deletion in the confirmation message when the team has no keys", async () => { + 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( + /any models created for it/i, + ); + }); +}); + +describe("Teams - LIT-2530 organization stays optional for proxy admin with a single org", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(teamCreateCall).mockResolvedValue({ + team_id: "new-team-1", + team_alias: "No Org Team", + models: ["gpt-4"], + organization_id: null, + keys: [], + members_with_roles: [], + spend: 0, + }); + mockUseOrganizations.mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }], + }); + }); + + it("creates a team with no organization when exactly one organization exists", async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "No Org Team" } }); + fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); + + const submitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(submitButtons[submitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ team_alias: "No Org Team", organization_id: null }), + ); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx similarity index 71% rename from ui/litellm-dashboard/src/components/OldTeams.tsx rename to ui/litellm-dashboard/src/components/Teams.tsx index e83d1acf4e5..2627f2c1d7f 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.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,27 +559,28 @@ 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" + /> +
- + + 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..513054aae7a 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,21 @@ it("should render VirtualKeysTable component", () => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); +it("left-anchors the create-key CTA below the title, between the header and the table toolbar", () => { + renderWithProviders(Create New Key} />); + + const heading = screen.getByRole("heading", { name: "Virtual Keys" }); + const ctas = screen.getAllByRole("button", { name: "Create New Key" }); + expect(ctas).toHaveLength(1); + const cta = ctas[0]; + const search = screen.getByPlaceholderText(/Search by key alias/); + + // The CTA follows the title row... + expect(heading.compareDocumentPosition(cta) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + // ...and precedes the table's search toolbar, so it sits in its own row above the table. + expect(cta.compareDocumentPosition(search) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); +}); + it("should display key information correctly", async () => { renderWithProviders(); @@ -177,6 +196,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 +208,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 +261,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 +374,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 +381,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 +396,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 +417,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 +445,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 +456,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 +503,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 +526,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..fe929dc0179 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." + /> + {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/add_model/AdaptiveRoutingConfig.tsx b/ui/litellm-dashboard/src/components/add_model/AdaptiveRoutingConfig.tsx new file mode 100644 index 00000000000..720b6f88db3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AdaptiveRoutingConfig.tsx @@ -0,0 +1,134 @@ +import { Card, InputNumber, Radio, Slider, Space, Switch, Typography } from "antd"; +import React from "react"; +import { + AdaptiveEligible, + ComplexityRouterConfigValue, + DEFAULT_ADAPTIVE_WEIGHTS, + DEFAULT_TIER_DISTANCE_PENALTY, +} from "./ComplexityRouterConfig"; + +const { Text } = Typography; + +interface AdaptiveRoutingConfigProps { + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; +} + +const AdaptiveRoutingConfig: React.FC = ({ value, onChange }) => { + const adaptiveWeights = value.adaptive_weights ?? DEFAULT_ADAPTIVE_WEIGHTS; + const adaptiveEligible = value.adaptive_eligible ?? "all"; + const tierDistancePenalty = value.tier_distance_penalty ?? DEFAULT_TIER_DISTANCE_PENALTY; + + const handleAdaptiveToggle = (adaptive: boolean) => { + const nextValue: ComplexityRouterConfigValue = { + ...value, + adaptive, + adaptive_weights: adaptiveWeights, + adaptive_eligible: adaptiveEligible, + tier_distance_penalty: tierDistancePenalty, + }; + onChange(nextValue); + }; + + const handleQualityWeightChange = (qualityPercent: number) => { + const quality = qualityPercent / 100; + onChange({ ...value, adaptive_weights: { quality, cost: Math.round((1 - quality) * 100) / 100 } }); + }; + + const handleAdaptiveEligibleChange = (eligible: AdaptiveEligible) => { + onChange({ ...value, adaptive_eligible: eligible }); + }; + + const handleTierDistancePenaltyChange = (penalty: number | null) => { + onChange({ ...value, tier_distance_penalty: penalty ?? DEFAULT_TIER_DISTANCE_PENALTY }); + }; + + return ( + <> +
+ + Enable adaptive bandit selection +
+ + When disabled, each request always uses the model assigned to its classified tier. + + + + + How Adaptive Routing Works + + + It learns from how each conversation actually goes: does the user have to rephrase or correct the model, does + it get stuck repeating itself, does it run out of tool calls, does the user seem satisfied. Combined with + cost, this live feedback shifts future routing toward the models that are actually working well, and improves + as more conversations come in. Until there's enough feedback, it defaults to the classified tier's + model. + + + + {value.adaptive && ( +
+
+ + Quality vs. Cost ({Math.round(adaptiveWeights.quality * 100)}% quality /{" "} + {Math.round(adaptiveWeights.cost * 100)}% cost) + + `${v}% quality / ${100 - (v ?? 0)}% cost` }} + /> + + Higher quality weight favors more capable (pricier) models; higher cost weight favors cheaper models when + the bandit has feedback to act on. Recommended: 30% quality / 70% cost split. + +
+ +
+ + Eligible Model Pool + + handleAdaptiveEligibleChange(e.target.value)} + className="w-full" + > + + + All tiers (soft floor){" "} + — router can pick across tiers, depending on the best fit for the prompt + + + Classified tier only{" "} + — router can only pick models within tier + + + +
+ + {adaptiveEligible === "all" && ( +
+ + Tier Distance Penalty + + + + Score penalty applied per tier-step away from the classified tier. + +
+ )} +
+ )} + + ); +}; + +export default AdaptiveRoutingConfig; diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx new file mode 100644 index 00000000000..92df8029edc --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -0,0 +1,171 @@ +import { InfoCircleOutlined } from "@ant-design/icons"; +import { Select as AntdSelect, Card, InputNumber, Radio, Space, Tooltip, Typography } from "antd"; +import React from "react"; +import { ClassifierType, ComplexityRouterConfigValue, DEFAULT_CLASSIFIER_TIMEOUT_MS } from "./ComplexityRouterConfig"; + +const { Text } = Typography; + +interface ClassificationMethodConfigProps { + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; + modelOptions: { value: string; label: string }[]; + customTechnicalKeywords?: string[]; + onCustomTechnicalKeywordsChange?: (keywords: string[]) => void; + showValidationErrors?: boolean; +} + +const ClassificationMethodConfig: React.FC = ({ + value, + onChange, + modelOptions, + customTechnicalKeywords, + onCustomTechnicalKeywordsChange, + showValidationErrors = false, +}) => { + const classifierModelMissing = + showValidationErrors && value.classifier_type === "llm" && !value.classifier_llm_config?.model; + + const handleClassifierTypeChange = (classifierType: ClassifierType) => { + onChange({ + ...value, + classifier_type: classifierType, + classifier_llm_config: + classifierType === "llm" + ? value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS } + : undefined, + }); + }; + + const handleClassifierModelChange = (model: string) => { + onChange({ + ...value, + classifier_llm_config: { + model, + timeout_ms: value.classifier_llm_config?.timeout_ms ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, + }, + }); + }; + + const handleClassifierTimeoutChange = (timeoutMs: number | null) => { + onChange({ + ...value, + classifier_llm_config: { + model: value.classifier_llm_config?.model ?? "", + timeout_ms: timeoutMs ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, + }, + }); + }; + + return ( + <> + handleClassifierTypeChange(e.target.value)} + className="w-full" + > + + + Heuristic{" "} + (default) — rule-based scoring, no API calls, <1ms latency + + + LLM Classifier{" "} + — use a model to decide the tier (e.g. a small/fast model) + + + + + {value.classifier_type === "llm" && ( +
+
+ + Classifier Model + + + {classifierModelMissing && ( + + A classifier model is required + + )} +
+
+ + Timeout (ms) + + + + Falls back to the heuristic scorer if the classifier call errors, times out, or returns an unparseable + response. + +
+
+ )} + + {value.classifier_type === "heuristic" && ( +
+
+ Custom Technical Keywords + + + +
+ + Optional: Add terms to the built-in list to improve classification accuracy on the technical dimension. + (e.g., udp, kafka, terraform). + + onCustomTechnicalKeywordsChange?.(keywords)} + placeholder="Type a keyword and press Enter, or paste a comma-separated list" + tokenSeparators={[","]} + open={false} + suffixIcon={null} + style={{ width: "100%" }} + allowClear + /> +
+ )} + + + + How Classification Works + + + The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical + terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the + tier: + +
    +
  • + SIMPLE: Score < 0.15 +
    • +
  • + MEDIUM: Score 0.15 - 0.35 +
    • +
  • + COMPLEX: Score 0.35 - 0.60 +
    • +
  • + REASONING: Score > 0.60 (or 2+ reasoning markers) +
    • +
+ + + ); +}; + +export default ClassificationMethodConfig; diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index a34a8709918..e1f90296770 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -12,10 +12,10 @@ const mockModelInfo = [ const defaultValue: ComplexityRouterConfigValue = { tiers: { - SIMPLE: "gpt-3.5-turbo", - MEDIUM: "gpt-3.5-turbo", - COMPLEX: "gpt-4", - REASONING: "claude-3-opus", + SIMPLE: ["gpt-3.5-turbo"], + MEDIUM: ["gpt-3.5-turbo"], + COMPLEX: ["gpt-4"], + REASONING: ["claude-3-opus"], }, classifier_type: "heuristic", }; @@ -58,11 +58,13 @@ describe("ComplexityRouterConfig", () => { it("should display the how classification works section", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText("How Classification Works")).toBeInTheDocument(); }); it("should show score thresholds in the classification section", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText(/Score < 0.15/)).toBeInTheDocument(); expect(screen.getByText(/Score 0.15 - 0.35/)).toBeInTheDocument(); expect(screen.getByText(/Score 0.35 - 0.60/)).toBeInTheDocument(); @@ -107,6 +109,7 @@ describe("ComplexityRouterConfig", () => { it("should render the custom technical keywords field", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText("Custom Technical Keywords")).toBeInTheDocument(); }); @@ -118,6 +121,7 @@ describe("ComplexityRouterConfig", () => { onCustomTechnicalKeywordsChange={vi.fn()} />, ); + fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText("udp")).toBeInTheDocument(); expect(screen.getByText("kafka")).toBeInTheDocument(); }); @@ -132,14 +136,16 @@ describe("ComplexityRouterConfig", () => { onCustomTechnicalKeywordsChange={onCustomTechnicalKeywordsChange} />, ); - const keywordsCard = screen.getByText("Custom Technical Keywords").closest(".ant-card") as HTMLElement; - const input = within(keywordsCard).getByRole("combobox"); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + const keywordsSection = screen.getByText("Custom Technical Keywords").closest("div")?.parentElement as HTMLElement; + const input = within(keywordsSection).getByRole("combobox"); await user.type(input, "udp,"); expect(onCustomTechnicalKeywordsChange).toHaveBeenCalledWith(["udp"]); }); it("should render an empty state when no keyword tier rules exist", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); expect(screen.getByText("Keyword Tier Overrides")).toBeInTheDocument(); expect(screen.getByText("No keyword tier overrides configured")).toBeInTheDocument(); }); @@ -158,6 +164,7 @@ describe("ComplexityRouterConfig", () => { const user = userEvent.setup(); const onKeywordTierRulesChange = vi.fn(); renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); await user.click(screen.getByRole("button", { name: /add keyword rule/i })); expect(onKeywordTierRulesChange).toHaveBeenCalledTimes(1); const newRules = onKeywordTierRulesChange.mock.calls[0][0]; @@ -175,6 +182,7 @@ describe("ComplexityRouterConfig", () => { onKeywordTierRulesChange={onKeywordTierRulesChange} />, ); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); expect(screen.getByText("invoice")).toBeInTheDocument(); expect(screen.getByText("refund")).toBeInTheDocument(); @@ -184,6 +192,7 @@ describe("ComplexityRouterConfig", () => { it("should not show embedding model or match score fields when semantic matching is disabled", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); expect(screen.getByText("Semantic keyword matching")).toBeInTheDocument(); expect(screen.queryByText("Embedding model")).not.toBeInTheDocument(); expect(screen.queryByText("Minimum match score")).not.toBeInTheDocument(); @@ -191,6 +200,7 @@ describe("ComplexityRouterConfig", () => { it("should show embedding model and match score fields when semantic matching is enabled", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); expect(screen.getByText("Embedding model")).toBeInTheDocument(); expect(screen.getByText("Minimum match score")).toBeInTheDocument(); }); @@ -205,6 +215,7 @@ describe("ComplexityRouterConfig", () => { onSemanticMatchingEnabledChange={onSemanticMatchingEnabledChange} />, ); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); await user.click(screen.getByRole("switch")); expect(onSemanticMatchingEnabledChange).toHaveBeenCalledWith(true, expect.anything()); }); @@ -252,7 +263,7 @@ describe("ComplexityRouterConfig", () => { renderWithProviders( , ); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index c31ee41a6ec..855a1b27df9 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -1,19 +1,22 @@ import { InfoCircleOutlined } from "@ant-design/icons"; -import { Select as AntdSelect, Card, Collapse, Divider, InputNumber, Radio, Space, Tooltip, Typography } from "antd"; +import { Select as AntdSelect, Card, Collapse, Divider, Space, Tooltip, Typography } from "antd"; import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; +import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; +import ClassificationMethodConfig from "./ClassificationMethodConfig"; import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules"; import SemanticKeywordMatching from "./SemanticKeywordMatching"; const { Text } = Typography; export const DEFAULT_CLASSIFIER_TIMEOUT_MS = 3000; +export const DEFAULT_TIER_DISTANCE_PENALTY = 0.5; export interface ComplexityTiers { - SIMPLE: string; - MEDIUM: string; - COMPLEX: string; - REASONING: string; + SIMPLE: string[]; + MEDIUM: string[]; + COMPLEX: string[]; + REASONING: string[]; } export interface ClassifierLLMConfig { @@ -23,10 +26,23 @@ export interface ClassifierLLMConfig { export type ClassifierType = "heuristic" | "llm"; +export interface AdaptiveRouterWeights { + quality: number; + cost: number; +} + +export const DEFAULT_ADAPTIVE_WEIGHTS: AdaptiveRouterWeights = { quality: 0.3, cost: 0.7 }; + +export type AdaptiveEligible = "all" | "classified_tier"; + export interface ComplexityRouterConfigValue { tiers: ComplexityTiers; classifier_type: ClassifierType; classifier_llm_config?: ClassifierLLMConfig; + adaptive?: boolean; + adaptive_weights?: AdaptiveRouterWeights; + tier_distance_penalty?: number; + adaptive_eligible?: AdaptiveEligible; } interface ComplexityRouterConfigProps { @@ -95,44 +111,10 @@ const ComplexityRouterConfig: React.FC = ({ label: model.model_group, })); - const classifierModelMissing = - showValidationErrors && value.classifier_type === "llm" && !value.classifier_llm_config?.model; - - const handleTierChange = (tier: keyof ComplexityTiers, model: string) => { + const handleTierChange = (tier: keyof ComplexityTiers, models: string[]) => { onChange({ ...value, - tiers: { ...value.tiers, [tier]: model }, - }); - }; - - const handleClassifierTypeChange = (classifierType: ClassifierType) => { - onChange({ - ...value, - classifier_type: classifierType, - classifier_llm_config: - classifierType === "llm" - ? (value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS }) - : undefined, - }); - }; - - const handleClassifierModelChange = (model: string) => { - onChange({ - ...value, - classifier_llm_config: { - model, - timeout_ms: value.classifier_llm_config?.timeout_ms ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, - }, - }); - }; - - const handleClassifierTimeoutChange = (timeoutMs: number | null) => { - onChange({ - ...value, - classifier_llm_config: { - model: value.classifier_llm_config?.model ?? "", - timeout_ms: timeoutMs ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, - }, + tiers: { ...value.tiers, [tier]: models }, }); }; @@ -142,20 +124,20 @@ const ComplexityRouterConfig: React.FC = ({ Complexity Tier Configuration - + The complexity router automatically classifies requests by complexity using rule-based scoring (no API calls, - <1ms latency). Configure which model handles each tier. + <1ms latency). Configure which model(s) handle each tier. {(Object.keys(TIER_DESCRIPTIONS) as Array).map((tier, index) => { const tierInfo = TIER_DESCRIPTIONS[tier]; - const tierMissing = showValidationErrors && !value.tiers[tier]; + const tierMissing = showValidationErrors && value.tiers[tier].length === 0; return (
{index > 0 && } @@ -172,14 +154,21 @@ const ComplexityRouterConfig: React.FC = ({ Examples: {tierInfo.examples} handleTierChange(tier, model)} - placeholder={`Select model for ${tierInfo.label.toLowerCase()} queries`} + onChange={(models) => handleTierChange(tier, models)} + placeholder={`Select model(s) for ${tierInfo.label.toLowerCase()} queries`} showSearch style={{ width: "100%" }} options={modelOptions} status={tierMissing ? "error" : undefined} /> + {value.tiers[tier].length > 1 && ( + + Multiple models selected — the router randomly picks among them per request (or Thompson-samples + within the pool when adaptive routing is on). + + )} {tierMissing && ( This tier is required @@ -205,148 +194,61 @@ const ComplexityRouterConfig: React.FC = ({ ), children: ( - <> - handleClassifierTypeChange(e.target.value)} - className="w-full" - > - - - Heuristic{" "} - (default) — rule-based scoring, no API calls, <1ms latency - - - LLM Classifier{" "} - — use a model to decide the tier (e.g. a small/fast model) - - - - - {value.classifier_type === "llm" && ( -
-
- - Classifier Model - - - {classifierModelMissing && ( - - A classifier model is required - - )} -
-
- - Timeout (ms) - - - - Falls back to the heuristic scorer if the classifier call errors, times out, or returns an - unparseable response. - -
-
- )} - + ), }, + { + key: "adaptive", + label: ( + + Advanced: Adaptive Routing + + ), + children: , + }, + ...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange + ? [ + { + key: "keyword-semantic", + label: ( + + Advanced: Keyword/Semantic Matching + + ), + children: ( + <> + {onKeywordTierRulesChange && ( + + )} + {onKeywordTierRulesChange && onSemanticMatchingEnabledChange && ( + + )} + {onSemanticMatchingEnabledChange && ( + + )} + + ), + }, + ] + : []), ]} /> - - - - -
- - Custom Technical Keywords - - - - -
- - Optional: Add terms to the built-in list to improve classification accuracy on the technical dimension. (e.g., - udp, kafka, terraform). - - onCustomTechnicalKeywordsChange?.(keywords)} - placeholder="Type a keyword and press Enter, or paste a comma-separated list" - tokenSeparators={[","]} - open={false} - suffixIcon={null} - style={{ width: "100%" }} - allowClear - /> - - - - - - - How Classification Works - - - The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical - terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the - tier: - -
    -
  • - SIMPLE: Score < 0.15 -
    • -
  • - MEDIUM: Score 0.15 - 0.35 -
    • -
  • - COMPLEX: Score 0.35 - 0.60 -
    • -
  • - REASONING: Score > 0.60 (or 2+ reasoning markers) -
    • -
- - - {/* Keyword-tier and semantic sections only render when their change handlers are - wired (the add-router flow). The edit-auto-router modal doesn't pass them yet, so - they stay hidden there rather than rendering interactive-but-dead controls. */} - {onKeywordTierRulesChange && ( - <> - - - - )} - - {onSemanticMatchingEnabledChange && ( - <> - - - - )}
); }; diff --git a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx index c7583427af6..c2252843b29 100644 --- a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx +++ b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx @@ -1,5 +1,5 @@ import { InfoCircleOutlined } from "@ant-design/icons"; -import { Card, InputNumber, Select as AntdSelect, Switch, Tooltip, Typography } from "antd"; +import { InputNumber, Select as AntdSelect, Switch, Tooltip, Typography } from "antd"; import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; @@ -36,7 +36,7 @@ const SemanticKeywordMatching: React.FC = ({ const embeddingModelMissing = showValidationErrors && !embeddingModel; return ( - +
@@ -86,7 +86,7 @@ const SemanticKeywordMatching: React.FC = ({
)} - +
); }; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index a74eab0abdd..8724c27b41a 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -8,7 +8,11 @@ import { all_admin_roles } from "@/utils/roles"; import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder from "./RouterConfigBuilder"; -import ComplexityRouterConfig, { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import ComplexityRouterConfig, { + ComplexityRouterConfigValue, + DEFAULT_ADAPTIVE_WEIGHTS, + DEFAULT_TIER_DISTANCE_PENALTY, +} from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; import { DEFAULT_MATCH_THRESHOLD } from "./SemanticKeywordMatching"; import { @@ -38,7 +42,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc const [routerType, setRouterType] = useState("recommended"); const [complexityRouterConfig, setComplexityRouterConfig] = useState({ - tiers: { SIMPLE: "", MEDIUM: "", COMPLEX: "", REASONING: "" }, + tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", }); @@ -89,6 +93,10 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc tiers, classifier_type: classifierType, classifier_llm_config: classifierLlmConfig, + adaptive = false, + adaptive_weights: adaptiveWeights = DEFAULT_ADAPTIVE_WEIGHTS, + tier_distance_penalty: tierDistancePenalty = DEFAULT_TIER_DISTANCE_PENALTY, + adaptive_eligible: adaptiveEligible = "all", } = complexityRouterConfig; const missingTiersError = getMissingTiersError(tiers); @@ -111,7 +119,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc return; } - const defaultModel = tiers.MEDIUM || tiers.SIMPLE || tiers.COMPLEX || tiers.REASONING; + const defaultModel = tiers.MEDIUM[0] || tiers.SIMPLE[0] || tiers.COMPLEX[0] || tiers.REASONING[0]; form.setFieldsValue({ custom_llm_provider: "auto_router", @@ -132,6 +140,10 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc semanticMatchingEnabled, embeddingModel, matchThreshold, + adaptive, + adaptiveWeights, + tierDistancePenalty, + adaptiveEligible, }; const submitValues = { diff --git a/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.test.ts b/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.test.ts index 01b6470f17f..f8ab4bab903 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.test.ts @@ -1,10 +1,10 @@ import { buildAutoRouterTestTargets } from "./build_auto_router_test_targets"; const tiers = { - SIMPLE: "gpt-4o-mini", - MEDIUM: "claude-sonnet-4", - COMPLEX: "claude-sonnet-4", - REASONING: "o3", + SIMPLE: ["gpt-4o-mini"], + MEDIUM: ["claude-sonnet-4"], + COMPLEX: ["claude-sonnet-4"], + REASONING: ["o3"], }; describe("buildAutoRouterTestTargets", () => { @@ -17,9 +17,21 @@ describe("buildAutoRouterTestTargets", () => { ]); }); + it("emits a target per model when a tier has more than one, and dedups across tiers", () => { + const targets = buildAutoRouterTestTargets({ + tiers: { SIMPLE: ["gpt-4o-mini", "claude-sonnet-4"], MEDIUM: ["claude-sonnet-4"], COMPLEX: [], REASONING: [] }, + semanticMatchingEnabled: false, + embeddingModel: undefined, + }); + expect(targets).toEqual([ + { labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }, + { labels: ["SIMPLE", "MEDIUM"], modelGroup: "claude-sonnet-4", mode: "chat" }, + ]); + }); + it("drops empty/whitespace tiers", () => { const targets = buildAutoRouterTestTargets({ - tiers: { SIMPLE: "gpt-4o-mini", MEDIUM: "", COMPLEX: " ", REASONING: "" }, + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [" "], REASONING: [] }, semanticMatchingEnabled: false, embeddingModel: undefined, }); @@ -29,7 +41,7 @@ describe("buildAutoRouterTestTargets", () => { it("returns [] when no tier is configured", () => { expect( buildAutoRouterTestTargets({ - tiers: { SIMPLE: "", MEDIUM: "", COMPLEX: "", REASONING: "" }, + tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, semanticMatchingEnabled: false, embeddingModel: undefined, }), @@ -38,7 +50,7 @@ describe("buildAutoRouterTestTargets", () => { it("appends an embedding target only when semantic matching is on and a model is set", () => { const targets = buildAutoRouterTestTargets({ - tiers: { SIMPLE: "gpt-4o-mini", MEDIUM: "", COMPLEX: "", REASONING: "" }, + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] }, semanticMatchingEnabled: true, embeddingModel: "voyage-3-5", }); @@ -50,7 +62,7 @@ describe("buildAutoRouterTestTargets", () => { it("omits the embedding target when semantic matching is on but no model is chosen", () => { const targets = buildAutoRouterTestTargets({ - tiers: { SIMPLE: "gpt-4o-mini", MEDIUM: "", COMPLEX: "", REASONING: "" }, + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] }, semanticMatchingEnabled: true, embeddingModel: undefined, }); @@ -59,7 +71,7 @@ describe("buildAutoRouterTestTargets", () => { it("omits the embedding target when a model is set but semantic matching is off", () => { const targets = buildAutoRouterTestTargets({ - tiers: { SIMPLE: "gpt-4o-mini", MEDIUM: "", COMPLEX: "", REASONING: "" }, + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] }, semanticMatchingEnabled: false, embeddingModel: "voyage-3-5", }); diff --git a/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.ts b/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.ts index b2a3cc10012..708a25c16f0 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.ts @@ -29,9 +29,11 @@ export const buildAutoRouterTestTargets = ({ embeddingModel, }: BuildAutoRouterTestTargetsParams): AutoRouterTestTarget[] => { const groupedByModel = TIER_ORDER.reduce>((acc, tier) => { - const modelGroup = tiers[tier]?.trim(); - if (!modelGroup) return acc; - return { ...acc, [modelGroup]: [...(acc[modelGroup] ?? []), tier] }; + return (tiers[tier] ?? []).reduce((tierAcc, rawModel) => { + const modelGroup = rawModel?.trim(); + if (!modelGroup) return tierAcc; + return { ...tierAcc, [modelGroup]: [...(tierAcc[modelGroup] ?? []), tier] }; + }, acc); }, {}); const tierTargets: AutoRouterTestTarget[] = Object.entries(groupedByModel).map(([modelGroup, labels]) => ({ diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index e5b547d8240..85a15ffad45 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -6,10 +6,10 @@ import { } from "./build_complexity_router_config"; const tiers = { - SIMPLE: "gpt-4o-mini", - MEDIUM: "gpt-4o", - COMPLEX: "claude-sonnet-4", - REASONING: "o1-preview", + SIMPLE: ["gpt-4o-mini"], + MEDIUM: ["gpt-4o"], + COMPLEX: ["claude-sonnet-4"], + REASONING: ["o1-preview"], }; const baseParams: BuildComplexityRouterConfigParams = { @@ -21,6 +21,10 @@ const baseParams: BuildComplexityRouterConfigParams = { semanticMatchingEnabled: false, embeddingModel: undefined, matchThreshold: 0.5, + adaptive: false, + adaptiveWeights: { quality: 0.3, cost: 0.7 }, + tierDistancePenalty: 0.5, + adaptiveEligible: "all", }; describe("buildComplexityRouterConfig", () => { @@ -29,6 +33,14 @@ describe("buildComplexityRouterConfig", () => { expect(config).toEqual({ tiers, classifier_type: "heuristic" }); }); + it("passes through a tier configured with more than one model as a pool", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + tiers: { ...tiers, SIMPLE: ["gpt-4o-mini", "gpt-4o", "claude-haiku-4-5"] }, + }); + expect(config.tiers.SIMPLE).toEqual(["gpt-4o-mini", "gpt-4o", "claude-haiku-4-5"]); + }); + it("includes classifier_llm_config only when classifier_type is llm", () => { const config = buildComplexityRouterConfig({ ...baseParams, @@ -123,6 +135,47 @@ describe("buildComplexityRouterConfig", () => { const config = buildComplexityRouterConfig(params); expect(config.keyword_tier_rules).toBeUndefined(); }); + + it("omits adaptive fields when adaptive is disabled even if weights linger in state", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + adaptive: false, + adaptiveWeights: { quality: 0.9, cost: 0.1 }, + tierDistancePenalty: 2, + adaptiveEligible: "classified_tier", + }); + expect(config.adaptive).toBeUndefined(); + expect(config.adaptive_weights).toBeUndefined(); + expect(config.tier_distance_penalty).toBeUndefined(); + expect(config.adaptive_eligible).toBeUndefined(); + }); + + it("includes tier_distance_penalty when adaptive is enabled with eligible='all'", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + adaptive: true, + adaptiveWeights: { quality: 0.6, cost: 0.4 }, + tierDistancePenalty: 0.75, + adaptiveEligible: "all", + }); + expect(config.adaptive).toBe(true); + expect(config.adaptive_weights).toEqual({ quality: 0.6, cost: 0.4 }); + expect(config.tier_distance_penalty).toBe(0.75); + expect(config.adaptive_eligible).toBe("all"); + }); + + it("omits tier_distance_penalty when eligible='classified_tier', since the penalty doesn't apply there", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + adaptive: true, + adaptiveWeights: { quality: 0.6, cost: 0.4 }, + tierDistancePenalty: 0.75, + adaptiveEligible: "classified_tier", + }); + expect(config.adaptive).toBe(true); + expect(config.adaptive_eligible).toBe("classified_tier"); + expect(config.tier_distance_penalty).toBeUndefined(); + }); }); describe("getMissingTiersError", () => { @@ -131,23 +184,27 @@ describe("getMissingTiersError", () => { }); it("names the specific missing tier when only one is blank", () => { - expect(getMissingTiersError({ ...tiers, REASONING: "" })).toBe( + expect(getMissingTiersError({ ...tiers, REASONING: [] })).toBe( "Select a model for the following tier(s): REASONING", ); }); it("names multiple missing tiers in SIMPLE/MEDIUM/COMPLEX/REASONING order", () => { - expect(getMissingTiersError({ ...tiers, SIMPLE: "", REASONING: "" })).toBe( + expect(getMissingTiersError({ ...tiers, SIMPLE: [], REASONING: [] })).toBe( "Select a model for the following tier(s): SIMPLE, REASONING", ); }); it("names all four tiers when none are filled", () => { - const noTiers = { SIMPLE: "", MEDIUM: "", COMPLEX: "", REASONING: "" }; + const noTiers = { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }; expect(getMissingTiersError(noTiers)).toBe( "Select a model for the following tier(s): SIMPLE, MEDIUM, COMPLEX, REASONING", ); }); + + it("treats a tier with more than one model as filled", () => { + expect(getMissingTiersError({ ...tiers, SIMPLE: ["gpt-4o-mini", "gpt-4o"] })).toBeNull(); + }); }); describe("getSemanticConfigError", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index a4a8ee6b074..3c3f21163b3 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -1,12 +1,11 @@ import { KeywordTierRule } from "./KeywordTierRules"; -import { ClassifierLLMConfig, ClassifierType } from "./ComplexityRouterConfig"; - -export interface ComplexityTiers { - SIMPLE: string; - MEDIUM: string; - COMPLEX: string; - REASONING: string; -} +import { + AdaptiveEligible, + AdaptiveRouterWeights, + ClassifierLLMConfig, + ClassifierType, + ComplexityTiers, +} from "./ComplexityRouterConfig"; export interface BuildComplexityRouterConfigParams { tiers: ComplexityTiers; @@ -17,6 +16,10 @@ export interface BuildComplexityRouterConfigParams { semanticMatchingEnabled: boolean; embeddingModel: string | undefined; matchThreshold: number; + adaptive: boolean; + adaptiveWeights: AdaptiveRouterWeights; + tierDistancePenalty: number; + adaptiveEligible: AdaptiveEligible; } export interface ComplexityRouterConfigPayload { @@ -28,12 +31,16 @@ export interface ComplexityRouterConfigPayload { semantic_keyword_matching?: boolean; embedding_model?: string; match_threshold?: number; + adaptive?: boolean; + adaptive_weights?: AdaptiveRouterWeights; + tier_distance_penalty?: number; + adaptive_eligible?: AdaptiveEligible; } const TIER_KEYS: Array = ["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"]; export const getMissingTiersError = (tiers: ComplexityTiers): string | null => { - const missing = TIER_KEYS.filter((tier) => !tiers[tier]); + const missing = TIER_KEYS.filter((tier) => tiers[tier].length === 0); if (missing.length === 0) return null; return `Select a model for the following tier(s): ${missing.join(", ")}`; }; @@ -43,7 +50,8 @@ export const getSemanticConfigError = ({ embeddingModel, keywordTierRules, }: Pick): - string | null => { + | string + | null => { if (!semanticMatchingEnabled) return null; if (!embeddingModel) return "Select an embedding model to use semantic keyword matching"; if (keywordTierRules.length === 0) return "Add at least one keyword tier rule to use semantic keyword matching"; @@ -61,6 +69,10 @@ export const buildComplexityRouterConfig = ({ semanticMatchingEnabled, embeddingModel, matchThreshold, + adaptive, + adaptiveWeights, + tierDistancePenalty, + adaptiveEligible, }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => { // Trim keywords and drop empty ones; drop any rule left with no keywords. Clicking // "Add keyword rule" seeds a rule with an empty keywords list, so without this an @@ -81,5 +93,11 @@ export const buildComplexityRouterConfig = ({ embedding_model: embeddingModel, match_threshold: matchThreshold, }), + ...(adaptive && { + adaptive: true, + adaptive_weights: adaptiveWeights, + ...(adaptiveEligible === "all" && { tier_distance_penalty: tierDistancePenalty }), + adaptive_eligible: adaptiveEligible, + }), }; }; 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/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts new file mode 100644 index 00000000000..cd8093928d5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -0,0 +1,104 @@ +import { buildUpdatedComplexityRouterConfig } from "./edit_auto_router_modal"; + +const storedConfigValue = { + tiers: { + SIMPLE: "old-simple", + MEDIUM: "old-medium", + COMPLEX: "old-complex", + REASONING: "old-reasoning", + }, + classifier_type: "llm", + classifier_llm_config: { model: "old-classifier", timeout_ms: 1200 }, + custom_technical_keywords: ["kafka", "terraform"], + keyword_tier_rules: [{ keywords: ["invoice", "refund"], tier: "MEDIUM" }], + semantic_keyword_matching: true, + embedding_model: "voyage-4-large", + match_threshold: 0.65, + adaptive: true, + adaptive_weights: { quality: 0.3, cost: 0.7 }, + tier_distance_penalty: 0.8, + adaptive_eligible: "all", +}; + +const storedConfig = JSON.stringify(storedConfigValue); + +const tiers = { + SIMPLE: ["gpt-4o-mini"], + MEDIUM: ["gpt-4o-mini"], + COMPLEX: ["anthropic-sonnet-4-5"], + REASONING: ["anthropic-sonnet-4-5"], +}; + +const classifiedTierValue = { + tiers, + classifier_type: "heuristic" as const, + adaptive: true, + adaptive_weights: { quality: 0.4, cost: 0.6 }, + tier_distance_penalty: 0.8, + adaptive_eligible: "classified_tier" as const, +}; + +const expectedClassifiedTierConfig = { + tiers, + classifier_type: "heuristic", + custom_technical_keywords: ["kafka", "terraform"], + keyword_tier_rules: [{ keywords: ["invoice", "refund"], tier: "MEDIUM" }], + semantic_keyword_matching: true, + embedding_model: "voyage-4-large", + match_threshold: 0.65, + adaptive: true, + adaptive_weights: { quality: 0.4, cost: 0.6 }, + adaptive_eligible: "classified_tier", +}; + +const adaptiveDisabledValue = { + tiers, + classifier_type: "heuristic" as const, + adaptive: false, +}; + +const expectedAdaptiveDisabledConfig = { + tiers, + classifier_type: "heuristic", + custom_technical_keywords: ["kafka", "terraform"], + keyword_tier_rules: [{ keywords: ["invoice", "refund"], tier: "MEDIUM" }], + semantic_keyword_matching: true, + embedding_model: "voyage-4-large", + match_threshold: 0.65, +}; + +describe("buildUpdatedComplexityRouterConfig", () => { + it("preserves unrelated options and omits the penalty for classified-tier routing", () => { + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, classifiedTierValue); + + expect(updatedConfig).toEqual(expectedClassifiedTierConfig); + }); + + it("removes managed adaptive and classifier fields when they are disabled", () => { + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, adaptiveDisabledValue); + + expect(updatedConfig).toEqual(expectedAdaptiveDisabledConfig); + }); + + it("updates custom technical keywords when they are edited", () => { + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, classifiedTierValue, ["postgres"]); + + expect(updatedConfig.custom_technical_keywords).toEqual(["postgres"]); + }); + + it("removes custom technical keywords when they are cleared", () => { + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, classifiedTierValue, []); + + expect(updatedConfig.custom_technical_keywords).toBeUndefined(); + }); + + it("preserves a tier configured with more than one model as a pool", () => { + const multiModelValue = { + ...classifiedTierValue, + tiers: { ...tiers, SIMPLE: ["gpt-4o-mini", "claude-haiku-4-5"] }, + }; + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, multiModelValue); + + expect(updatedConfig.tiers).toMatchObject({ SIMPLE: ["gpt-4o-mini", "claude-haiku-4-5"] }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 9c97809bfde..ec54c9b7bad 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -4,13 +4,23 @@ import { Text, TextInput } from "@tremor/react"; import { modelAvailableCall, modelPatchUpdateCall } from "../networking"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder from "../add_model/RouterConfigBuilder"; -import ComplexityRouterConfig, { ComplexityRouterConfigValue } from "../add_model/ComplexityRouterConfig"; +import ComplexityRouterConfig, { + ComplexityRouterConfigValue, + DEFAULT_ADAPTIVE_WEIGHTS, + DEFAULT_TIER_DISTANCE_PENALTY, +} from "../add_model/ComplexityRouterConfig"; import NotificationsManager from "../molecules/notifications_manager"; const isComplexityRouterModel = (modelData: any): boolean => modelData?.litellm_params?.model?.startsWith("auto_router/complexity_router") || modelData?.litellm_params?.complexity_router_config != null; +const normalizeTierModels = (value: unknown): string[] => { + if (Array.isArray(value)) return value; + if (typeof value === "string" && value) return [value]; + return []; +}; + interface EditAutoRouterModalProps { isVisible: boolean; onCancel: () => void; @@ -20,6 +30,57 @@ interface EditAutoRouterModalProps { userRole: string; } +const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ + "tiers", + "classifier_type", + "classifier_llm_config", + "adaptive", + "adaptive_weights", + "tier_distance_penalty", + "adaptive_eligible", +]); + +const toRecord = (value: unknown): Record => { + const parsed: unknown = typeof value === "string" ? JSON.parse(value) : value; + return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed) + ? (parsed as Record) + : {}; +}; + +export const buildUpdatedComplexityRouterConfig = ( + storedConfig: unknown, + value: ComplexityRouterConfigValue, + customTechnicalKeywords?: string[], +): Record => { + const preservedConfig = Object.fromEntries( + Object.entries(toRecord(storedConfig)).filter( + ([key]) => + !MANAGED_COMPLEXITY_ROUTER_KEYS.has(key) && + (customTechnicalKeywords === undefined || key !== "custom_technical_keywords"), + ), + ); + const adaptiveEligible = value.adaptive_eligible ?? "all"; + + return { + ...preservedConfig, + tiers: value.tiers, + classifier_type: value.classifier_type, + ...(value.classifier_type === "llm" ? { classifier_llm_config: value.classifier_llm_config } : {}), + ...(customTechnicalKeywords && + customTechnicalKeywords.length > 0 && { + custom_technical_keywords: customTechnicalKeywords, + }), + ...(value.adaptive && { + adaptive: true, + adaptive_weights: value.adaptive_weights ?? DEFAULT_ADAPTIVE_WEIGHTS, + ...(adaptiveEligible === "all" && { + tier_distance_penalty: value.tier_distance_penalty ?? DEFAULT_TIER_DISTANCE_PENALTY, + }), + adaptive_eligible: adaptiveEligible, + }), + }; +}; + const EditAutoRouterModal: React.FC = ({ isVisible, onCancel, @@ -35,8 +96,9 @@ const EditAutoRouterModal: React.FC = ({ const [showCustomDefaultModel, setShowCustomDefaultModel] = useState(false); const [showCustomEmbeddingModel, setShowCustomEmbeddingModel] = useState(false); const [routerConfig, setRouterConfig] = useState(null); + const [customTechnicalKeywords, setCustomTechnicalKeywords] = useState([]); const [complexityRouterConfig, setComplexityRouterConfig] = useState({ - tiers: { SIMPLE: "", MEDIUM: "", COMPLEX: "", REASONING: "" }, + tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", }); const isComplexityRouter = isComplexityRouterModel(modelData); @@ -85,14 +147,21 @@ const EditAutoRouterModal: React.FC = ({ setComplexityRouterConfig({ tiers: { - SIMPLE: parsedConfig.tiers?.SIMPLE || "", - MEDIUM: parsedConfig.tiers?.MEDIUM || "", - COMPLEX: parsedConfig.tiers?.COMPLEX || "", - REASONING: parsedConfig.tiers?.REASONING || "", + SIMPLE: normalizeTierModels(parsedConfig.tiers?.SIMPLE), + MEDIUM: normalizeTierModels(parsedConfig.tiers?.MEDIUM), + COMPLEX: normalizeTierModels(parsedConfig.tiers?.COMPLEX), + REASONING: normalizeTierModels(parsedConfig.tiers?.REASONING), }, classifier_type: parsedConfig.classifier_type || "heuristic", classifier_llm_config: parsedConfig.classifier_llm_config, + adaptive: parsedConfig.adaptive || false, + adaptive_weights: parsedConfig.adaptive_weights, + tier_distance_penalty: parsedConfig.tier_distance_penalty, + adaptive_eligible: parsedConfig.adaptive_eligible || "all", }); + setCustomTechnicalKeywords( + Array.isArray(parsedConfig.custom_technical_keywords) ? parsedConfig.custom_technical_keywords : [], + ); form.setFieldsValue({ auto_router_name: modelData.model_name, @@ -138,7 +207,7 @@ const EditAutoRouterModal: React.FC = ({ if (isComplexityRouter) { const { tiers, classifier_type, classifier_llm_config } = complexityRouterConfig; - if (Object.values(tiers).filter(Boolean).length === 0) { + if (Object.values(tiers).every((models) => models.length === 0)) { NotificationsManager.fromBackend("Please select at least one model for a complexity tier"); return; } @@ -147,14 +216,14 @@ const EditAutoRouterModal: React.FC = ({ return; } - const defaultModel = tiers.MEDIUM || tiers.SIMPLE || tiers.COMPLEX || tiers.REASONING; + const defaultModel = tiers.MEDIUM[0] || tiers.SIMPLE[0] || tiers.COMPLEX[0] || tiers.REASONING[0]; const updatedLitellmParams = { ...modelData.litellm_params, - complexity_router_config: { - tiers, - classifier_type, - ...(classifier_type === "llm" ? { classifier_llm_config } : {}), - }, + complexity_router_config: buildUpdatedComplexityRouterConfig( + modelData.litellm_params?.complexity_router_config, + complexityRouterConfig, + customTechnicalKeywords, + ), complexity_router_default_model: defaultModel, }; const updatedModelInfo = { @@ -264,6 +333,8 @@ const EditAutoRouterModal: React.FC = ({ onChange={(config) => { setComplexityRouterConfig(config); }} + customTechnicalKeywords={customTechnicalKeywords} + onCustomTechnicalKeywordsChange={setCustomTechnicalKeywords} />
) : ( 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/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index f0e27607ffb..a57676d07b8 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -91,6 +91,16 @@ vi.mock("@/components/team/member_permissions", () => ({ default: vi.fn(() =>
Member Permissions
), })); +vi.mock("@/components/common_components/ModelAliasManager", () => ({ + default: vi.fn(({ initialModelAliases, onAliasUpdate }) => ( +
+
{JSON.stringify(initialModelAliases)}
+ + +
+ )), +})); + vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ useAccessGroups: vi.fn().mockReturnValue({ data: [ @@ -858,4 +868,160 @@ describe("TeamInfoView", () => { }); }); }); + + describe("model aliases", () => { + const openSettingsEditor = async (user: ReturnType) => { + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /edit settings/i })); + + await waitFor(() => { + expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); + }); + }; + + it("should render existing model aliases in the read-only settings view", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + litellm_model_table: { model_aliases: { "my-smart-model": "gpt-4", "my-fast-model": "gpt-3.5-turbo" } }, + }), + ); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByText("Team Settings")).toBeInTheDocument(); + }); + + expect(screen.getByText("Model Aliases")).toBeInTheDocument(); + expect(screen.getByText("my-smart-model")).toBeInTheDocument(); + expect(screen.getByText("gpt-4")).toBeInTheDocument(); + expect(screen.getByText("my-fast-model")).toBeInTheDocument(); + expect(screen.getByText("gpt-3.5-turbo")).toBeInTheDocument(); + }); + + it("should show an empty state when the team has no model aliases", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ litellm_model_table: null })); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByText("Team Settings")).toBeInTheDocument(); + }); + + expect(screen.getByText("No model aliases configured")).toBeInTheDocument(); + }); + + it("should seed the alias editor from existing team aliases", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + models: ["gpt-4"], + litellm_model_table: { model_aliases: { "my-smart-model": "gpt-4" } }, + }), + ); + + renderWithProviders(); + + await openSettingsEditor(user); + + expect(screen.getByTestId("alias-editor-initial")).toHaveTextContent( + JSON.stringify({ "my-smart-model": "gpt-4" }), + ); + }); + + it("should pass model_aliases to teamUpdateCall when aliases are added", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] })); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + + await user.click(screen.getByRole("button", { name: "Set Alias" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + team_id: "123", + model_aliases: { "gpt-4o": "gpt-4" }, + }), + ); + }); + }); + + it("should send an empty model_aliases map to clear existing aliases", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + models: ["gpt-4"], + litellm_model_table: { model_aliases: { "my-smart-model": "gpt-4" } }, + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + + await user.click(screen.getByRole("button", { name: "Clear Aliases" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + const payload = vi.mocked(networking.teamUpdateCall).mock.calls[0][1] as Record; + expect(payload.model_aliases).toEqual({}); + }); + + it("should not include model_aliases when the team has none and the editor is untouched", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ models: ["gpt-4"], litellm_model_table: null }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + const payload = vi.mocked(networking.teamUpdateCall).mock.calls[0][1] as Record; + expect(payload).not.toHaveProperty("model_aliases"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 3b80c598d34..74f201b40c6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -34,6 +34,7 @@ import { CheckIcon, CopyIcon } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; +import ModelAliasManager from "../common_components/ModelAliasManager"; import AgentSelector from "../agent_management/AgentSelector"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; import DurationSelect from "../common_components/DurationSelect"; @@ -106,7 +107,7 @@ export interface TeamData { budget_reset_at: string | null; model_id: string | null; litellm_model_table: { - model_aliases: Record; + model_aliases: Record | null; } | null; created_at: string; access_group_ids?: string[]; @@ -206,6 +207,7 @@ const TeamInfoView: React.FC = ({ const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [isDeleting, setIsDeleting] = useState(false); const [isTeamSaving, setIsTeamSaving] = useState(false); + const [teamModelAliases, setTeamModelAliases] = useState>({}); const routerSettingsRef = React.useRef(null); const [organization, setOrganization] = useState(null); const { userRole, userId } = useAuthorized(); @@ -619,6 +621,11 @@ const TeamInfoView: React.FC = ({ updateData.default_team_member_models = values.default_team_member_models; } + const previousModelAliases = info.litellm_model_table?.model_aliases ?? {}; + if (Object.keys(teamModelAliases).length > 0 || Object.keys(previousModelAliases).length > 0) { + updateData.model_aliases = teamModelAliases; + } + // Handle router_settings - read fresh values from DOM at save time. const currentRouterSettings = routerSettingsRef.current?.getValue(); if (currentRouterSettings?.router_settings) { @@ -905,7 +912,13 @@ const TeamInfoView: React.FC = ({
Team Settings {canEditTeam && !isEditing && ( - )} @@ -1026,6 +1039,24 @@ const TeamInfoView: React.FC = ({ /> + + Model Aliases{" "} + + + + + } + > + + + @@ -1520,6 +1551,26 @@ const TeamInfoView: React.FC = ({
)} +
+ Model Aliases + {(() => { + const aliasEntries = Object.entries(info.litellm_model_table?.model_aliases ?? {}); + if (aliasEntries.length === 0) { + return
No model aliases configured
; + } + return ( +
+ {aliasEntries.map(([alias, target]) => ( +
+ {alias} + {" -> "} + {target} +
+ ))} +
+ ); + })()} +
Rate Limits
TPM: {info.tpm_limit || "Unlimited"}
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 ( +