diff --git a/.circleci/config.yml b/.circleci/config.yml index cc485aa0595..e8a8483781b 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2744,84 +2744,6 @@ jobs: file: ./coverage.xml flags: circleci - ui_build: - docker: - - image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad - auth: - username: ${DOCKERHUB_USERNAME} - password: ${DOCKERHUB_PASSWORD} - resource_class: medium+ - working_directory: ~/project - steps: - - checkout - - skip_if_unrelated_changes: - category: client - - setup_google_dns - - restore_cache: - keys: - - ui-build-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }} - - ui-build-deps-v1- - - restore_cache: - keys: - - ui-nextjs-cache-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }} - - ui-nextjs-cache-v1- - - run: - name: Install dependencies - command: | - cd ui/litellm-dashboard - npm ci - - save_cache: - key: ui-build-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }} - paths: - - ui/litellm-dashboard/node_modules - - run: - name: Build UI - command: | - cd ui/litellm-dashboard - source ./build_ui.sh - - save_cache: - key: ui-nextjs-cache-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }} - paths: - - ui/litellm-dashboard/.next/cache - - persist_to_workspace: - root: . - paths: - - litellm/proxy/_experimental/out - - ui_unit_tests: - docker: - - image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad - auth: - username: ${DOCKERHUB_USERNAME} - password: ${DOCKERHUB_PASSWORD} - resource_class: xlarge - working_directory: ~/project - steps: - - checkout - - skip_if_unrelated_changes: - category: client - - setup_google_dns - - restore_cache: - keys: - - ui-unit-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }} - - ui-unit-deps-v1- - - run: - name: Install dependencies - command: | - cd ui/litellm-dashboard - npm ci - - save_cache: - key: ui-unit-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }} - paths: - - ui/litellm-dashboard/node_modules - - run: - name: Run UI unit tests (Vitest) - command: | - cd ui/litellm-dashboard - - CI=true npm run test -- --run \ - --pool forks --poolOptions.forks.maxForks=6 - e2e_ui_testing: docker: - image: cimg/python:3.12-browsers@sha256:b432899af01c9a311bf74f4f22e9ada2e5306d4b1b4383f8d29e1228a5844ef2 @@ -3181,12 +3103,6 @@ workflows: filters: *main_branches - litellm_router_unit_testing: filters: *main_branches - - ui_build: - filters: *main_branches - - ui_unit_tests: - requires: - - ui_build - filters: *main_branches - auth_ui_unit_tests: filters: *main_branches - proxy_behavior_tests: diff --git a/.github/workflows/test-unit-proxy-legacy.yml b/.github/workflows/test-unit-proxy-legacy.yml deleted file mode 100644 index e8ca36fb30d..00000000000 --- a/.github/workflows/test-unit-proxy-legacy.yml +++ /dev/null @@ -1,106 +0,0 @@ -name: "Unit Tests: Proxy Legacy Tests" - -on: - pull_request: - branches: - - main - - litellm_internal_staging - - litellm_oss_staging - - "litellm_**" - push: - branches: - - main - - litellm_internal_staging - -permissions: - contents: read - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} - cancel-in-progress: ${{ github.event_name == 'pull_request' }} - -jobs: - test: - runs-on: ubuntu-latest - timeout-minutes: 20 - strategy: - fail-fast: false - matrix: - test-group: - - name: "auth-and-jwt" - path: "tests/proxy_unit_tests/test_[a-j]*.py" - - name: "key-generation" - path: "tests/proxy_unit_tests/test_[k-o]*.py" - - name: "proxy-config" - path: "tests/proxy_unit_tests/test_prisma*.py tests/proxy_unit_tests/test_prompt*.py tests/proxy_unit_tests/test_proxy_[c-r]*.py" - - name: "proxy-server" - path: "tests/proxy_unit_tests/test_proxy_server.py" - - name: "proxy-server-extras" - path: "tests/proxy_unit_tests/test_proxy_server_*.py tests/proxy_unit_tests/test_proxy_setting_guardrails.py" - - name: "proxy-utils" - path: "tests/proxy_unit_tests/test_proxy_utils.py" - - name: "proxy-token-counter" - path: "tests/proxy_unit_tests/test_proxy_token_counter.py" - - name: "proxy-response-and-misc" - path: "tests/proxy_unit_tests/test_[r-t]*.py" - - name: "proxy-user-auth-and-spend" - path: "tests/proxy_unit_tests/test_[u-z]*.py" - - name: ${{ matrix.test-group.name }} - - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Detect backend-relevant changes - id: changes - uses: ./.github/actions/detect-backend-changes - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Cache uv dependencies - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 - with: - path: | - ~/.cache/uv - .venv - key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }} - restore-keys: | - ${{ runner.os }}-uv- - - - name: Install dependencies - if: steps.changes.outputs.decision != 'skip' - run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - - - name: Cache Prisma binaries - if: steps.changes.outputs.decision != 'skip' - uses: ./.github/actions/cache-prisma-binaries - - - name: Generate Prisma client - if: steps.changes.outputs.decision != 'skip' - run: | - uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - - name: Run tests - ${{ matrix.test-group.name }} - if: steps.changes.outputs.decision != 'skip' - env: - TEST_PATH: ${{ matrix.test-group.path }} - run: | - uv run --no-sync pytest ${TEST_PATH} \ - --tb=short -vv \ - --maxfail=10 \ - -n 2 \ - --reruns 1 \ - --reruns-delay 1 \ - --dist=loadscope \ - --durations=20 diff --git a/litellm/constants.py b/litellm/constants.py index 554165f5d39..6449834d6a4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1491,6 +1491,9 @@ SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEA SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float( os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5) ) +SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300")) +SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30")) +SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000")) TOOL_SPEND_TOP_TOOLS: Final = 100 SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) @@ -1742,6 +1745,9 @@ PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900 # Furthest back the catch-up pass looks for unpriced PTU days when a deployment # declares no ptu_effective_from, bounding the scan for an open-ended window. PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90 +# Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide +# expiry cannot produce an alert too large for the channel delivering it. +PTU_LAPSED_ALERT_LIMIT: Final[int] = 10 # Slack allowed when deciding a sentinel row is stale. The row's updated_at and the # run's cutoff are stamped by different hosts, so clock skew between them must not let # one run delete a charge another just wrote. A stale row is hours old and a concurrent diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index db253b1517d..8720f561e14 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -2,8 +2,9 @@ # On success, logs events to Langfuse import os import traceback -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast from packaging.version import Version @@ -30,6 +31,7 @@ from litellm.types.utils import ( ImageResponse, ModelResponse, RerankResponse, + StandardLoggingMetadata, StandardLoggingPayload, StandardLoggingPromptManagementMetadata, TextCompletionResponse, @@ -46,6 +48,11 @@ else: Langfuse = Any +_DENIED_STEERING_KEYS: Final = frozenset({"headers", "endpoint", "caching_groups", "previous_models"}) +_NO_METADATA: Final[Mapping[str, Any]] = MappingProxyType({}) +_REDACTED_PROXY_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "cookie", "referer"}) + + def _extract_cache_read_input_tokens(usage_obj) -> int: """ Extract cache_read_input_tokens from usage object. @@ -512,16 +519,14 @@ class LangFuseLogger: else [] ) - if standard_logging_object is None: - end_user_id = None - prompt_management_metadata: StandardLoggingPromptManagementMetadata | None = None - else: - end_user_id = standard_logging_object["metadata"].get("user_api_key_end_user_id", None) - - prompt_management_metadata = cast( - StandardLoggingPromptManagementMetadata | None, - standard_logging_object["metadata"].get("prompt_management_metadata", None), - ) + allowlisted_metadata: Final[StandardLoggingMetadata | dict[str, Any]] = ( + standard_logging_object["metadata"] if standard_logging_object is not None else _NO_METADATA + ) + end_user_id: Final = allowlisted_metadata.get("user_api_key_end_user_id", None) + prompt_management_metadata: Final[StandardLoggingPromptManagementMetadata | None] = cast( + StandardLoggingPromptManagementMetadata | None, + allowlisted_metadata.get("prompt_management_metadata", None), + ) # Clean Metadata before logging - never log raw metadata # the raw metadata can contain circular references which leads to infinite recursion @@ -540,12 +545,7 @@ class LangFuseLogger: tags.append(f"{key}:{value}") # clean litellm metadata before logging - if key in [ - "headers", - "endpoint", - "caching_groups", - "previous_models", - ]: + if key in _DENIED_STEERING_KEYS: continue else: clean_metadata[key] = value @@ -630,19 +630,18 @@ class LangFuseLogger: trace_params["output"] = output if not mask_output else "redacted-by-litellm" if debug is True or (isinstance(debug, str) and debug.lower() == "true"): - if "metadata" in trace_params: - # log the raw_metadata in the trace - trace_params["metadata"]["metadata_passed_to_litellm"] = metadata - else: - trace_params["metadata"] = {"metadata_passed_to_litellm": metadata} + debug_metadata: Final = { + key: value for key, value in metadata.items() if isinstance(value, (str, int, float, bool)) + } + trace_params["metadata"] = { + **(trace_params.get("metadata") or _NO_METADATA), + "metadata_passed_to_litellm": debug_metadata, + } cost: Final = kwargs.get("response_cost", None) verbose_logger.debug("trace: %s", cost) - clean_metadata["litellm_response_cost"] = cost - if standard_logging_object is not None: - hidden_params: Final = standard_logging_object.get("hidden_params", {}) - clean_metadata["hidden_params"] = filter_exceptions_from_params(hidden_params) + hidden_params: Final = standard_logging_object.get("hidden_params") if standard_logging_object else None if ( litellm.langfuse_default_tags is not None @@ -654,22 +653,24 @@ class LangFuseLogger: tags.append(f"proxy_base_url:{proxy_base_url}") api_base: Final = litellm_params.get("api_base", None) - if api_base: - clean_metadata["api_base"] = api_base - vertex_location: Final = kwargs.get("vertex_location", None) - if vertex_location: - clean_metadata["vertex_location"] = vertex_location - aws_region_name: Final = kwargs.get("aws_region_name", None) - if aws_region_name: - clean_metadata["aws_region_name"] = aws_region_name + + candidate_enrichments: Final = ( + ("litellm_response_cost", cost, True), + ("hidden_params", filter_exceptions_from_params(hidden_params), hidden_params is not None), + ("api_base", api_base, bool(api_base)), + ("vertex_location", vertex_location, bool(vertex_location)), + ("aws_region_name", aws_region_name, bool(aws_region_name)), + ("cache_hit", kwargs.get("cache_hit") or False, self._supports_tags() and "cache_hit" in kwargs), + ) + enrichments: Final[Mapping[str, Any]] = { + key: value for key, value, include in candidate_enrichments if include + } if self._supports_tags(): - if "cache_hit" in kwargs: - if kwargs["cache_hit"] is None: - kwargs["cache_hit"] = False - clean_metadata["cache_hit"] = kwargs["cache_hit"] + if "cache_hit" in kwargs and kwargs["cache_hit"] is None: + kwargs["cache_hit"] = False # rebind-ok: pre-existing normalization other integrations rely on if existing_trace_id is None: trace_params.update({"tags": tags}) @@ -682,13 +683,13 @@ class LangFuseLogger: if headers: for key, value in headers.items(): # these headers can leak our API keys and/or JWT tokens - if key.lower() not in ["authorization", "cookie", "referer"]: + if key.lower() not in _REDACTED_PROXY_HEADERS: clean_headers[key] = value trace: Final[StatefulTraceClient] = self.Langfuse.trace(**trace_params) # Log provider specific information as a span - log_provider_specific_information_as_span(trace, clean_metadata) + log_provider_specific_information_as_span(trace, enrichments) # Log guardrail information as a span self._log_guardrail_information_as_span( @@ -761,7 +762,10 @@ class LangFuseLogger: "output": output if not mask_output else "redacted-by-litellm", "usage": usage, "usage_details": usage_details, - "metadata": log_requester_metadata(clean_metadata), + "metadata": { + **log_requester_metadata(redact_user_api_key_info(metadata=allowlisted_metadata)), + **enrichments, + }, "level": level, "version": clean_metadata.pop("version", None), } @@ -1058,7 +1062,7 @@ def _add_prompt_to_generation_params( def log_provider_specific_information_as_span( trace, - clean_metadata, + clean_metadata: Mapping[str, Any], ): """ Logs provider-specific information as spans. @@ -1098,7 +1102,7 @@ def log_provider_specific_information_as_span( ) -def log_requester_metadata(clean_metadata: dict): +def log_requester_metadata(clean_metadata: Mapping[str, Any]): returned_metadata: Final = {} requester_metadata: Final = clean_metadata.get("requester_metadata") or {} for k, v in clean_metadata.items(): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index bad8e93e0c5..d23466938f2 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -34,12 +34,16 @@ class ExceptionCheckers: """ @staticmethod - def is_error_str_rate_limit(error_str: str) -> bool: + def is_error_str_rate_limit(error_str: str, status_code: int | None = None) -> bool: """ Check if an error string indicates a rate limit error. Args: error_str: The error string to check + status_code: The HTTP status the provider returned, when known. Gates only the + bare-number branch: providers echo the request back in validation errors and + 429 is an ordinary token id, so an echoed prompt can put a standalone 429 in + the body of a 400. The phrase branches stay ungated (#11455). Returns: True if the error indicates a rate limit, False otherwise @@ -47,8 +51,9 @@ class ExceptionCheckers: if not isinstance(error_str, str): return False - # Only treat 429 as a rate limit signal when it appears as a standalone token - if re.search(r"\b429\b", error_str): + # A standalone 429 counts unless the provider's own status says otherwise. The + # status is read off an arbitrary exception, so a non-integer means "unknown". + if re.search(r"\b429\b", error_str) and (not isinstance(status_code, int) or status_code == 429): return True _error_str_lower: Final = error_str.lower() @@ -280,7 +285,9 @@ def _map_openai_exception( else: exception_provider = custom_llm_provider[0].upper() + custom_llm_provider[1:] + "Exception" - if ExceptionCheckers.is_error_str_rate_limit(error_str): + if ExceptionCheckers.is_error_str_rate_limit( + error_str, status_code=getattr(original_exception, "status_code", None) + ): raise RateLimitError( message=f"RateLimitError: {exception_provider} - {message}", model=model, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 9aa5a4f465f..b444c77d718 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -5,6 +5,7 @@ This file contains common utils for anthropic calls. import copy import re from collections.abc import Mapping, Sequence +from datetime import datetime, timezone from types import MappingProxyType from typing import Any, Final, Literal @@ -12,6 +13,7 @@ import httpx from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm +from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, ) @@ -28,6 +30,7 @@ from litellm.types.llms.anthropic import ( AnthropicMcpServerTool, ) from litellm.types.llms.openai import AllMessageValues +from litellm.types.proxy.model_listing import ModelInfoResponse _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$") _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$") @@ -1221,3 +1224,39 @@ def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: additional_headers: Final = {**llm_response_headers, **openai_headers} return additional_headers + + +def _anthropic_model_entry(model: ModelInfoResponse, created_at: str) -> Mapping[str, object]: + token_limits: Final = ( + ("max_input_tokens", model.get("max_input_tokens")), + ("max_tokens", model.get("max_output_tokens")), + ) + return { # mutable-ok: JSON response body, serialized by the route and never mutated + "type": "model", + "id": model["id"], + "display_name": model["id"], + "created_at": created_at, + **{name: limit for name, limit in token_limits if limit is not None}, # mutable-ok: merged into the body above + } + + +def create_anthropic_model_list_response(models: Sequence[ModelInfoResponse]) -> Mapping[str, object]: + """Build the Anthropic-native /v1/models envelope. + + Clients that send an anthropic-version header parse the Anthropic Models API + shape (type/display_name/created_at plus has_more/first_id/last_id) and filter + the list themselves, so every model is returned here. The token limits carry + over from the OpenAI-shaped listing, named as the Messages API names them + """ + created_at: Final = ( + datetime.fromtimestamp(DEFAULT_MODEL_CREATED_AT_TIME, tz=timezone.utc).isoformat().replace("+00:00", "Z") + ) + data: Final = [ # mutable-ok: JSON response body, serialized by the route and never mutated + _anthropic_model_entry(model, created_at) for model in models + ] + return { # mutable-ok: JSON response body, serialized by the route and never mutated + "data": data, + "has_more": False, + "first_id": models[0]["id"] if models else None, + "last_id": models[-1]["id"] if models else None, + } diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 647ccc33a44..b8b07af59c6 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -19,9 +19,9 @@ from litellm.llms.bedrock.common_utils import ( convert_bedrock_invoke_output_format_to_inline_schema, get_anthropic_beta_from_headers, normalize_bedrock_opus_output_config_effort, + normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, pop_bedrock_invoke_output_config_format, - remove_custom_field_from_tools, ) from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER from litellm.types.llms.openai import AllMessageValues @@ -243,8 +243,8 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): if "anthropic_version" not in anthropic_request: anthropic_request["anthropic_version"] = self.anthropic_version - # Remove `custom` field from tools (Bedrock doesn't support it) - remove_custom_field_from_tools(anthropic_request) + # Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it) + normalize_custom_field_on_tools(anthropic_request) normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_request) return anthropic_request diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 48bc60a07e5..4ad20772ed0 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -176,13 +176,14 @@ def convert_bedrock_invoke_output_format_to_inline_schema( request_body["messages"] = new_messages -def remove_custom_field_from_tools(request_body: dict) -> None: +def normalize_custom_field_on_tools(request_body: dict) -> None: """ - Remove ``custom`` field from each tool in the request body. + Drop the ``custom`` field from each tool, first hoisting a boolean + ``custom.defer_loading`` onto the top-level ``defer_loading`` flag that + Bedrock and Anthropic actually document, unless the tool already carries one. - Claude Code (v2.1.69+) sends ``custom: {defer_loading: true}`` on tool - definitions, which Anthropic's API accepts but Bedrock rejects with - ``"Extra inputs are not permitted"``. + Claude Code (v2.1.69+) is reported to send ``custom: {defer_loading: true}`` on + tool definitions, which Bedrock rejects with ``"Extra inputs are not permitted"``. Args: request_body: The request dictionary to modify in-place. @@ -193,8 +194,14 @@ def remove_custom_field_from_tools(request_body: dict) -> None: if not tools or not isinstance(tools, list): return for tool in tools: - if isinstance(tool, dict): - tool.pop("custom", None) + if not isinstance(tool, dict): + continue + custom: dict[str, object] | None = tool.pop("custom", None) + if not isinstance(custom, dict) or "defer_loading" in tool: + continue + deferred: object = custom.get("defer_loading") + if isinstance(deferred, bool): + tool["defer_loading"] = deferred def normalize_json_schema_custom_types_to_object(schema: dict) -> None: diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 8d039d95bb1..372cf110f7c 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -33,9 +33,9 @@ from litellm.llms.bedrock.common_utils import ( get_anthropic_beta_from_headers, is_claude_4_5_on_bedrock, normalize_bedrock_opus_output_config_effort, + normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, pop_bedrock_invoke_output_config_format, - remove_custom_field_from_tools, ) from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, @@ -749,11 +749,9 @@ class AmazonAnthropicClaudeMessagesConfig( model, ) - # 5b. Remove `custom` field from tools (Bedrock doesn't support it) - # Claude Code sends `custom: {defer_loading: true}` on tool definitions, - # which causes Bedrock to reject the request with "Extra inputs are not permitted" + # 5b. Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it) # Ref: https://github.com/BerriAI/litellm/issues/22847 - remove_custom_field_from_tools(anthropic_messages_request) + normalize_custom_field_on_tools(anthropic_messages_request) normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_messages_request) ensure_bedrock_anthropic_messages_tool_names(anthropic_messages_request) diff --git a/litellm/llms/openai/cost_calculation.py b/litellm/llms/openai/cost_calculation.py index eafabdb880d..0352d246c09 100644 --- a/litellm/llms/openai/cost_calculation.py +++ b/litellm/llms/openai/cost_calculation.py @@ -109,15 +109,16 @@ def cost_per_second(model: str, custom_llm_provider: str | None, duration: float prompt_cost = 0.0 completion_cost = 0.0 ## Speech / Audio cost calculation - if "output_cost_per_second" in model_info and model_info["output_cost_per_second"] is not None: + output_cost_per_second: Final = model_info.get("output_cost_per_second") + if output_cost_per_second is not None and output_cost_per_second > 0: verbose_logger.debug( "For model=%s - output_cost_per_second: %s; duration: %s", model, - model_info.get("output_cost_per_second"), + output_cost_per_second, duration, ) ## COST PER SECOND ## - completion_cost = model_info["output_cost_per_second"] * duration + completion_cost = output_cost_per_second * duration elif "input_cost_per_second" in model_info and model_info["input_cost_per_second"] is not None: verbose_logger.debug( "For model=%s - input_cost_per_second: %s; duration: %s", diff --git a/litellm/main.py b/litellm/main.py index 16eff5a0f3e..04ae410db6f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5616,7 +5616,12 @@ def completion( elif custom_llm_provider == "hosted_vllm": response = _complete_hosted_vllm(_dispatch_ctx) elif ( - model in litellm.open_ai_chat_completion_models + # A known OpenAI model name only decides the route when nothing else + # resolved a provider. get_llm_provider() already maps these names to + # "openai", so a different value here was asked for explicitly (or came + # from a register_model entry), and the provider config built for it + # would be handed to the OpenAI handler. + (model in litellm.open_ai_chat_completion_models and custom_llm_provider in (None, "openai")) or custom_llm_provider == "custom_openai" or custom_llm_provider == "deepinfra" or custom_llm_provider == "perplexity" 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 554b6ea952e..95a3806e8ad 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 @@ -1190,7 +1190,7 @@ class MCPRequestHandler: DEPRECATED: This method is deprecated in favor of server-specific auth headers using the format x-mcp-{{server_alias}}-{{header_name}} instead. """ - mcp_client_side_auth_header_name: Final[str] = MCPRequestHandler._get_mcp_client_side_auth_header_name() + mcp_client_side_auth_header_name: Final[str] = MCPRequestHandler.get_mcp_client_side_auth_header_name() auth_header: Final = headers.get(mcp_client_side_auth_header_name) if auth_header: verbose_logger.warning( @@ -1265,7 +1265,7 @@ class MCPRequestHandler: return oauth2_headers @staticmethod - def _get_mcp_client_side_auth_header_name() -> str: + def get_mcp_client_side_auth_header_name() -> str: """ Get the header name used to pass the MCP auth header to the MCP server diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c8ff6e262d2..a1adda2bc95 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -118,6 +118,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( is_short_mcp_tool_prefix_enabled, iter_known_server_prefixes, iter_known_tool_name_spellings, + logging_safe_mcp_headers, match_known_server_prefix, match_known_tool_name, merge_mcp_headers, @@ -4603,6 +4604,7 @@ class MCPServerManager: ), "user_api_key_hash": (getattr(user_api_key_auth, "api_key_hash", None) if user_api_key_auth else None), "incoming_bearer_token": incoming_bearer_token, + "headers": logging_safe_mcp_headers(raw_headers), } # Create MCP request object for processing diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index a9a3367cd93..125dc3d773d 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -1042,100 +1042,15 @@ def _build_sampling_request( raw_headers: dict[str, str] | None = None, client_ip: str | None = None, ) -> "Request": - """Build a synthetic FastAPI Request for sampling sub-calls. + """The synthetic FastAPI Request for sampling sub-calls, carrying the original + MCP connection's headers and client IP.""" + from litellm.proxy._experimental.mcp_server.utils import build_synthetic_mcp_request - Converts the original MCP connection's HTTP headers into ASGI - scope format so that ``add_litellm_data_to_request`` can apply - header-dependent guardrails, tag-based routing, trace correlation, - and ``forward_llm_provider_auth_headers``. - - Key fields populated: - - **headers**: All original HTTP headers are forwarded (except - hop-by-hop: content-length, transfer-encoding). This ensures - ``traceparent``, ``authorization``, ``user-agent``, and - ``x-litellm-api-key`` are visible to pre-call utils. - - **client**: The ASGI ``(host, port)`` tuple so that - ``request.client.host`` returns the real client IP for - IP-based routing and guardrails. - - **server**: Derived from the running proxy's ``server_host`` - / ``server_port`` when available, avoiding the misleading - ``127.0.0.1:0`` placeholder. - - **x-forwarded-for**: Injected from ``client_ip`` if the - original headers don't already carry it, as a fallback for - IP attribution. - """ - from fastapi import Request - - # --- Build ASGI headers --- - _scope_headers: Final[list[tuple[bytes, bytes]]] = [(b"content-type", b"application/json")] - # Hop-by-hop headers that must NOT be forwarded into the - # synthetic request (they describe the original HTTP framing, - # not the logical request). - _HOP_BY_HOP: Final = frozenset( - { - "content-length", - "transfer-encoding", - "connection", - "keep-alive", - "upgrade", - "te", - "trailer", - } + return build_synthetic_mcp_request( + path="/mcp/sampling/createMessage", + raw_headers=raw_headers, + client_ip=client_ip, ) - if raw_headers: - for hdr_name, hdr_value in raw_headers.items(): - _key = hdr_name.lower() - # Skip content-type (already set), x-forwarded-for (use resolved - # client_ip instead to prevent spoofing), and hop-by-hop headers - if _key in {"content-type", "x-forwarded-for"} or _key in _HOP_BY_HOP: - continue - _scope_headers.append( - ( - _key.encode("latin-1", errors="replace"), - hdr_value.encode("utf-8"), - ) - ) - - # Inject x-forwarded-for from captured client_ip if the - # original headers don't already carry it - if client_ip and not any(h[0] == b"x-forwarded-for" for h in _scope_headers): - _scope_headers.append((b"x-forwarded-for", client_ip.encode("utf-8"))) - - # --- Derive server (host, port) from the running proxy --- - _server_host = "127.0.0.1" - _server_port = 4000 # LiteLLM default - try: - from litellm.proxy import proxy_server - - _proxy_host: Final[str | None] = getattr(proxy_server, "server_host", None) - _proxy_port: Final[str | int | None] = getattr(proxy_server, "server_port", None) - - if _proxy_host: - _server_host = str(_proxy_host) - if _proxy_port: - _server_port = int(_proxy_port) - except (ImportError, AttributeError, TypeError, ValueError): - pass - - # --- Build ASGI client tuple for request.client.host --- - _client_tuple = None - if client_ip: - _client_tuple = (client_ip, 0) - - scope: Final[dict[str, object]] = { - "type": "http", - "method": "POST", - "path": "/mcp/sampling/createMessage", - "scheme": "http", - "server": (_server_host, _server_port), - "query_string": b"", - "root_path": "", - "headers": _scope_headers, - } - if _client_tuple is not None: - scope["client"] = _client_tuple - - return Request(scope=scope) async def _build_completion_kwargs( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 49a1f1314f0..f237529b319 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -58,9 +58,11 @@ from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_VERSION, MCPMissingUserEnvVarsError, add_server_prefix_to_name, + build_synthetic_mcp_request, extract_mcp_tool_result_error_message, get_server_prefix, iter_known_server_prefixes, + logging_safe_mcp_headers, match_known_tool_name, ) from litellm.proxy._types import ( @@ -860,11 +862,11 @@ if MCP_AVAILABLE: name: str, arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, ) -> LiteLLMLoggingObj | None: """Run the pre-call pipeline (guardrails + logging setup) for a virtual mcp_tool_call so the SSE path spend-logs like the REST path.""" - from fastapi import Request - from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -874,13 +876,10 @@ if MCP_AVAILABLE: proxy_logging_obj, ) - request: Final = Request( - scope={ - "type": "http", - "method": "POST", - "path": "/mcp/tools/call", - "headers": [(b"content-type", b"application/json")], - } + request: Final = build_synthetic_mcp_request( + path="/mcp/tools/call", + raw_headers=raw_headers, + client_ip=client_ip, ) _, virtual_logging_obj = await ProxyBaseLLMRequestProcessing( data={"name": name, "arguments": arguments} @@ -952,7 +951,11 @@ if MCP_AVAILABLE: assert user_api_key_auth is not None # guaranteed by the flag check above virtual_logging_obj: Final = await _build_virtual_call_logging_obj( - name=name, arguments=args, user_api_key_auth=user_api_key_auth + name=name, + arguments=args, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, ) return await handle_mcp_tool_call( tool_name=args.get("tool_name", ""), @@ -979,7 +982,6 @@ if MCP_AVAILABLE: Raises: HTTPException: If tool not found or arguments missing """ - from fastapi import Request from mcp.server.lowlevel.server import request_ctx from mcp.types import CallToolResult @@ -1041,13 +1043,10 @@ if MCP_AVAILABLE: body_data["litellm_trace_id"] = chain_id body_data["litellm_session_id"] = chain_id - request: Final = Request( - scope={ - "type": "http", - "method": "POST", - "path": "/mcp/tools/call", - "headers": [(b"content-type", b"application/json")], - } + request: Final = build_synthetic_mcp_request( + path="/mcp/tools/call", + raw_headers=raw_headers, + client_ip=_client_ip, ) if user_api_key_auth is not None: data = await add_litellm_data_to_request( @@ -1905,6 +1904,7 @@ if MCP_AVAILABLE: "litellm_trace_id": effective_litellm_trace_id, "metadata": { "spend_logs_metadata": spend_logs_metadata, + "headers": logging_safe_mcp_headers(raw_headers), **({"tags": request_tags} if request_tags else {}), }, # Provide a small input payload for standard logging diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 5534b8b5a3c..4cf84dd0725 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -7,6 +7,7 @@ import importlib import json import os import re +import typing from collections.abc import Iterable, Iterator, Mapping, MutableMapping, MutableSequence from collections.abc import Set as AbstractSet from typing import Any, Final, Protocol @@ -14,6 +15,9 @@ from urllib.parse import quote from litellm.types.mcp_server.mcp_server_manager import MCPServer +if typing.TYPE_CHECKING: + from fastapi import Request + class _McpServerLike(Protocol): @property @@ -862,3 +866,146 @@ def set_mcp_tool_result_structured_content(result: object, value: object) -> boo return True except (AttributeError, TypeError, ValueError): return False + + +_HOP_BY_HOP_HEADERS: Final = frozenset( + { + "content-length", + "transfer-encoding", + "connection", + "keep-alive", + "upgrade", + "te", + "trailer", + } +) + +_SYNTHETIC_REQUEST_EXCLUDED_HEADERS: Final = _HOP_BY_HOP_HEADERS | frozenset({"content-type", "x-forwarded-for"}) + +_SYNTHETIC_REQUEST_SERVER: Final = ("127.0.0.1", 4000) + +_MCP_SERVER_AUTH_HEADER_PREFIX: Final = "x-mcp-" + + +def _custom_litellm_key_header_name() -> str | None: + """``general_settings.litellm_key_header_name``, the deployment's custom header name for + the proxy virtual key, so it is stripped from observability copies like the standard ones.""" + try: + from litellm.proxy.proxy_server import general_settings + except ImportError: + return None + return general_settings.get("litellm_key_header_name") if general_settings else None + + +def _mcp_client_side_auth_header_name() -> str: + """The header name the client passes the upstream MCP credential in, falling back to the + default when ``general_settings`` is unavailable (the SDK, outside a running proxy).""" + from .auth.user_api_key_auth_mcp import MCPRequestHandler + + try: + return MCPRequestHandler.get_mcp_client_side_auth_header_name() + except ImportError: + return MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME + + +def _upstream_credential_headers(header_names: Iterable[str]) -> frozenset[str]: + """Lowercased names of the headers in ``header_names`` that carry an upstream MCP + credential rather than request context: the configured client side auth header and + the per-server ``x-mcp-{alias}-{header}`` family. ``clean_headers`` only knows the + credential headers of the chat completions path, so these are dropped on top of it. + """ + from .auth.user_api_key_auth_mcp import MCPRequestHandler + + non_credential: Final = frozenset( + { + MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME.lower(), + MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME.lower(), + } + ) + client_side_auth: Final = _mcp_client_side_auth_header_name().lower() + return frozenset( + name + for name in (raw_name.lower() for raw_name in header_names) + if name == client_side_auth or (name.startswith(_MCP_SERVER_AUTH_HEADER_PREFIX) and name not in non_credential) + ) + + +def build_synthetic_mcp_request( + *, + path: str, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, +) -> "Request": + """A synthetic FastAPI ``Request`` carrying the MCP connection's HTTP headers. + + The MCP protocol transports do not hand a per-call ``Request`` to the tool + handlers, so one is reconstructed from the connection's ``raw_headers``. That + lets ``add_litellm_data_to_request`` derive ``metadata.headers``, + ``proxy_server_request``, header-based tags, guardrails and trace correlation + exactly as on the chat completions path. Hop-by-hop headers describe the + original HTTP framing rather than the logical request, so they are dropped, and + ``x-forwarded-for`` comes from the resolved ``client_ip`` to avoid spoofing. Upstream + MCP credentials and the deployment's proxy key header, including a custom + ``litellm_key_header_name``, are dropped so they cannot reach a callback or a guardrail + through the derived metadata even when a caller omits ``general_settings``. + """ + from fastapi import Request + + custom_key_header: Final = _custom_litellm_key_header_name() + excluded: Final = ( + _SYNTHETIC_REQUEST_EXCLUDED_HEADERS + | _upstream_credential_headers(raw_headers.keys() if raw_headers else ()) + | (frozenset({custom_key_header.lower()}) if custom_key_header else frozenset()) + ) + forwarded: Final = tuple( + ( + name.lower().encode("latin-1", errors="replace"), + value.encode("utf-8", errors="replace"), + ) + for name, value in (raw_headers.items() if raw_headers else ()) + if name.lower() not in excluded + ) + xff: Final = ((b"x-forwarded-for", client_ip.encode("utf-8")),) if client_ip else () + return Request( + scope={ + "type": "http", + "method": "POST", + "path": path, + "scheme": "http", + "server": _SYNTHETIC_REQUEST_SERVER, + "query_string": b"", + "root_path": "", + "headers": ((b"content-type", b"application/json"), *forwarded, *xff), + **({"client": (client_ip, 0)} if client_ip else {}), + } + ) + + +def logging_safe_mcp_headers(raw_headers: Mapping[str, str] | None) -> Mapping[str, str]: + """The MCP request's client headers, sanitized the way the chat completions path + sanitizes them before they reach a logging callback or a guardrail: proxy key + headers stripped, including the custom key header name the deployment configured, + upstream MCP credentials dropped, and credential-bearing values masked. + + Client-controlled behaviour flags (``litellm-disable-message-redaction``) are dropped + too: these headers are read back out of the metadata to change proxy behaviour, so + leaving one in place would let any MCP client turn off the redaction an admin + configured. This path carries no key or team object to authorize an opt-out with, so + it always strips them.""" + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import ( + UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS, + clean_headers, + redact_credential_headers, + ) + + excluded: Final = ( + _upstream_credential_headers(raw_headers.keys() if raw_headers else ()) + | UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS + ) + cleaned: Final = clean_headers( + Headers(raw_headers), + litellm_key_header_name=_custom_litellm_key_header_name(), + ) + return redact_credential_headers({name: value for name, value in cleaned.items() if name.lower() not in excluded}) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6c01575e271..385f39e02a4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2514,6 +2514,22 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False.", ) + maximum_spend_logs_cleanup_batch_size: int | None = Field( + None, + description="Rows deleted per DELETE statement by the spend log cleanup job. Defaults to 1000.", + ) + maximum_spend_logs_cleanup_max_batches: int | None = Field( + None, + description="Maximum DELETE statements the spend log cleanup job issues per table per run. Defaults to 500.", + ) + maximum_spend_logs_cleanup_run_budget: str | None = Field( + None, + description="Wall-clock budget for one spend log cleanup run (e.g. '5m'), shared across every table it prunes. A run that hits the budget stops and the next run resumes from where it left off. Defaults to '5m'.", + ) + maximum_spend_logs_cleanup_batch_timeout: str | None = Field( + None, + description="Postgres statement_timeout and lock_timeout applied to each spend log cleanup delete batch (e.g. '30s'), so cleanup cannot hold row locks or a connection indefinitely. Defaults to '30s'.", + ) mcp_internal_ip_ranges: list[str] | None = Field( None, description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 0239368ad0e..51050e62494 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -14,6 +14,7 @@ import math import re import time from collections.abc import Iterator, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast from fastapi import HTTPException, Request, status @@ -376,6 +377,16 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None zero_cost_cache[model_name] = False return False + if _has_ptu_flat_cost(model_name, llm_router): + verbose_proxy_logger.debug( + "Model %s prices reserved PTU capacity as a flat cost, so its zero per-token " + "rate is not a free model (enforce budget)", + safe_name, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + verbose_proxy_logger.debug( "Model %s has zero cost explicitly configured (input: %s, output: %s)", safe_name, @@ -394,6 +405,24 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None return True +_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool: + """Whether any deployment in the model group bills reserved PTU capacity as a flat cost. + + Such a deployment carries an explicit zero per-token price so the flat cost is not charged + twice, which otherwise reads here as a free model and waives every budget check for it. + """ + for deployment in llm_router.model_list: + if deployment.get("model_name") != model: + continue + model_info = deployment.get("model_info") or _NO_MODEL_INFO + if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None: + return True + return False + + def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool: """ Check if any deployment in the model group has cost fields explicitly diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 1869328c039..60a03689804 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,7 +1,10 @@ import copy import os from collections.abc import Callable, Iterable -from typing import TYPE_CHECKING, Any, Final, Optional +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias + +from typing_extensions import assert_never import litellm from litellm import get_secret @@ -50,6 +53,66 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +@dataclass(frozen=True, slots=True) +class _CallbackResolvedToClass: + entry: str + loaded: type + tag: Literal["resolved_to_class"] = "resolved_to_class" + + +@dataclass(frozen=True, slots=True) +class _CallbackNotDispatchable: + entry: str + loaded: object + tag: Literal["not_dispatchable"] = "not_dispatchable" + + +_CallbackLoadError: TypeAlias = _CallbackResolvedToClass | _CallbackNotDispatchable + + +def _classify_loaded_callback(entry: str, loaded: object) -> CustomLogger | Callable[..., object] | _CallbackLoadError: + """ + Decide whether what a ``litellm_settings.callbacks`` dotted path resolved to can be dispatched. + + A dotted path only ever runs as a ``CustomLogger`` instance or as a callback function. Anything + else (most commonly a class instead of an instance) used to load without complaint and then be + skipped on every request, with no log line and no error. + """ + if isinstance(loaded, CustomLogger) or (callable(loaded) and not isinstance(loaded, type)): + return loaded + if isinstance(loaded, type): + return _CallbackResolvedToClass(entry=entry, loaded=loaded) + return _CallbackNotDispatchable(entry=entry, loaded=loaded) + + +def _raise_callback_load_error(error: _CallbackLoadError) -> NoReturn: + """The one edge that raises: map a load error onto config load's failure contract.""" + match error: + case _CallbackResolvedToClass(): + module_path: Final = error.entry.rsplit(".", 1)[0] if "." in error.entry else error.entry + raise ValueError( + f"litellm_settings.callbacks entry '{error.entry}' resolved to the class " + f"{error.loaded.__module__}.{error.loaded.__qualname__}, which is neither a " + "CustomLogger instance nor a callable, so the proxy would never run it." + f" Point it at an instance instead, e.g. add `proxy_handler_instance = {error.loaded.__name__}()` to " + f'{module_path} and set `callbacks: ["{module_path}.proxy_handler_instance"]`.' + ) + case _CallbackNotDispatchable(): + raise ValueError( + f"litellm_settings.callbacks entry '{error.entry}' resolved to " + f"{type(error.loaded).__name__} {error.loaded!r}, which is neither a " + "CustomLogger instance nor a callable, so the proxy would never run it." + ) + assert_never(error) + + +def _loaded_callback_or_raise(entry: str, loaded: object) -> CustomLogger | Callable[..., object]: + resolved: Final = _classify_loaded_callback(entry=entry, loaded=loaded) + if isinstance(resolved, _CallbackResolvedToClass | _CallbackNotDispatchable): + _raise_callback_load_error(resolved) + return resolved + + def initialize_callbacks_on_proxy( value: Any, premium_user: bool, @@ -305,9 +368,12 @@ def initialize_callbacks_on_proxy( "%s attempting to import custom calback=%s %s", blue_color_code, callback, reset_color_code ) imported_list.append( - get_instance_fn( - value=callback, - config_file_path=config_file_path, + _loaded_callback_or_raise( + entry=callback, + loaded=get_instance_fn( + value=callback, + config_file_path=config_file_path, + ), ) ) if isinstance(litellm.callbacks, list): @@ -321,9 +387,12 @@ def initialize_callbacks_on_proxy( PrometheusLogger._mount_metrics_endpoint() else: litellm.callbacks = [ - get_instance_fn( - value=value, - config_file_path=config_file_path, + _loaded_callback_or_raise( + entry=value, + loaded=get_instance_fn( + value=value, + config_file_path=config_file_path, + ), ) ] verbose_proxy_logger.debug("%s Initialized Callbacks - %s %s", blue_color_code, litellm.callbacks, reset_color_code) diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 9f01c719a5f..d19023862cb 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -1,22 +1,60 @@ import asyncio +import time +from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Final +from typing import Final, Literal, TypeAlias + +from pydantic import BaseModel, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache from litellm.constants import ( SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS, SPEND_LOG_CLEANUP_BATCH_SIZE, + SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS, SPEND_LOG_CLEANUP_JOB_NAME, SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES, + SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP, + SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS, SPEND_LOG_RUN_LOOPS, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import ( + RunOutcome, + SpendLogCleanupMetrics, +) from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import ( + RemainingTimeoutMs, SpendLogsPartitionManager, ) from litellm.proxy.utils import PrismaClient +StopReason: TypeAlias = Literal["exhausted", "budget_exhausted", "batch_cap_reached", "aborted"] + + +@dataclass(frozen=True, slots=True) +class TableCleanupResult: + """Outcome of pruning one table, so the caller can report why a run ended.""" + + rows_deleted: int + stop_reason: StopReason + + +class _RemainingRow(BaseModel): + """One row of the capped outstanding-rows probe, validated out of prisma's untyped result.""" + + remaining: int + + +_REMAINING_ROWS: Final = TypeAdapter(list[_RemainingRow]) + +SPEND_LOG_CLEANUP_BOUND_SETTINGS: Final = ( + "maximum_spend_logs_cleanup_batch_size", + "maximum_spend_logs_cleanup_max_batches", + "maximum_spend_logs_cleanup_run_budget", + "maximum_spend_logs_cleanup_batch_timeout", +) + class SpendLogCleanup: """ @@ -26,6 +64,24 @@ class SpendLogCleanup: dropping whole partitions (instant, frees disk immediately). Otherwise it falls back to deleting logs in batches. Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments. + + Every run is bounded so it can never monopolise the database: a wall-clock + budget shared across all tables, a per-table batch cap, and a Postgres + statement/lock timeout on every statement the job issues, deletes and the + outstanding-rows probe alike. A run that hits a bound stops cleanly and the + next run resumes from where it left off, because the cutoff is recomputed + and deleted rows are gone. + + The budget is a hard wall clock, not an advisory one. Every statement this + job issues, deletes, the outstanding-rows probe and partition DDL alike, is + issued with a timeout clamped to the budget that is still left, so one + started just under the deadline is cancelled by Postgres at the deadline + rather than running a further batch timeout past it. No statement is issued + at all once the budget is spent, which is why the probe is skipped on that + path. Partition DDL additionally carries a lock_timeout, because it takes an + ACCESS EXCLUSIVE lock and would otherwise queue behind a long-running reader + for as long as that reader lives; a partition this run cannot get is left + for the next one. """ def __init__( @@ -34,17 +90,88 @@ class SpendLogCleanup: redis_cache: RedisCache | None = None, partition_manager: SpendLogsPartitionManager | None = None, ): - self.batch_size = SPEND_LOG_CLEANUP_BATCH_SIZE self.retention_seconds: int | None = None self.partition_manager = partition_manager or SpendLogsPartitionManager() from litellm.proxy.proxy_server import general_settings as default_settings self.general_settings = general_settings or default_settings + self._refresh_bounds() from litellm.proxy.proxy_server import proxy_logging_obj pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager self.pod_lock_manager = pod_lock_manager - verbose_proxy_logger.info("SpendLogCleanup initialized with batch size: %s", self.batch_size) + verbose_proxy_logger.info( + "SpendLogCleanup initialized: batch_size=%s max_batches=%s run_budget=%ss batch_timeout=%ss", + self.batch_size, + self.max_batches, + self.run_budget_seconds, + self.batch_timeout_seconds, + ) + + def _refresh_bounds(self) -> None: + """ + Re-read every bound in SPEND_LOG_CLEANUP_BOUND_SETTINGS from settings. + + The scheduler holds one long-lived instance, so a bound captured at + construction would never reflect a dashboard change. general_settings is + the same dict the periodic config reload mutates in place, so reading it + per run is what makes these knobs live. Every bound falls back to its + shipped default, so clearing a field restores that default. + """ + self.batch_size: int = self._positive_int_setting( + "maximum_spend_logs_cleanup_batch_size", SPEND_LOG_CLEANUP_BATCH_SIZE + ) + self.max_batches: int = self._positive_int_setting( + "maximum_spend_logs_cleanup_max_batches", SPEND_LOG_RUN_LOOPS + ) + self.run_budget_seconds: float = self._duration_setting( + "maximum_spend_logs_cleanup_run_budget", SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS + ) + self.batch_timeout_seconds: float = self._duration_setting( + "maximum_spend_logs_cleanup_batch_timeout", SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS + ) + + def _positive_int_setting(self, setting_name: str, default: int) -> int: + """ + Read a positive-integer knob, falling back to the default when unset or unusable. + """ + raw: Final = self.general_settings.get(setting_name) + if raw is None: + return default + try: + parsed: Final = int(raw) + except (TypeError, ValueError): + verbose_proxy_logger.warning("Invalid %s value: %s, using default %s", setting_name, raw, default) + return default + if parsed <= 0: + verbose_proxy_logger.warning("%s must be positive, got %s, using default %s", setting_name, parsed, default) + return default + return parsed + + def _duration_setting(self, setting_name: str, default_seconds: float) -> float: + """ + Read a duration knob (e.g. '5m'), falling back to the default when unset or unusable. + + The knob must never be able to remove the bound it exists to enforce, so + anything the parser rejects (including the non-finite spellings 'inf' and + 'nan') and anything non-positive falls back rather than being honoured. + """ + raw: Final = self.general_settings.get(setting_name) + if raw is None: + return default_seconds + try: + parsed: Final = float(duration_in_seconds(str(raw))) + except (ValueError, TypeError) as e: + verbose_proxy_logger.warning( + "Invalid %s value: %s (%s), using default %ss", setting_name, raw, e, default_seconds + ) + return default_seconds + if parsed <= 0: + verbose_proxy_logger.warning( + "%s must be a positive duration, got %s, using default %ss", setting_name, raw, default_seconds + ) + return default_seconds + return parsed def _retention_seconds_for(self, setting_name: str) -> int | None: """ @@ -78,6 +205,91 @@ class SpendLogCleanup: self.retention_seconds = self._retention_seconds_for("maximum_spend_logs_retention_period") return self.retention_seconds is not None + def _timeout_ms(self, deadline: float) -> int: + """ + The per-statement bound in milliseconds: the batch timeout, or whatever + is left of the run budget, whichever is smaller. + + Clamping to the remaining budget is what makes the budget a real + wall-clock bound rather than an advisory one. Postgres offers no "stop + at time T", only a per-statement duration, so a statement issued just + under the deadline would otherwise run a full batch timeout past it, and + with several tables those overruns stack. + + Interpolating this into SQL is safe by construction: an int cannot carry + SQL, and SET does not accept a bind parameter. + """ + remaining_ms: Final = int((deadline - time.monotonic()) * 1000) + return max(1, min(int(self.batch_timeout_seconds * 1000), remaining_ms)) + + def _remaining_timeout_ms(self, deadline: float) -> RemainingTimeoutMs: + """ + The per-statement bound for work this job delegates, as a callable. + + Partition maintenance issues one statement per partition, so handing it a + number would bound each statement by the budget that was left before the + FIRST one and never by what remains. Re-evaluating per statement is what + makes the loop itself bounded, and None tells the callee to stop rather + than issue a statement it has no budget for. + """ + + def remaining() -> int | None: + return None if time.monotonic() >= deadline else self._timeout_ms(deadline) + + return remaining + + async def _execute_delete_batch( + self, prisma_client: PrismaClient, delete_sql: str, cutoff_date: datetime, deadline: float + ) -> int | None: + """ + Run one delete batch under a Postgres statement and lock timeout. + + The timeouts are what actually bound the work: a Prisma transaction + timeout cannot interrupt a statement that is already executing, so + without these a single batch blocked behind a lock would hold its + connection, and the row locks it already took, indefinitely. SET LOCAL + scopes both to this transaction so the pooled connection is unaffected. + + Returns the row count, or None when the driver returned something that + is not a row count. That is a contract violation rather than a transient + fault, so the caller stops instead of retrying. + """ + timeout_ms: Final = self._timeout_ms(deadline) + async with prisma_client.db.tx() as tx: + await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}") + await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}") + deleted_result: Final = await tx.execute_raw(delete_sql, cutoff_date, self.batch_size) + return deleted_result if isinstance(deleted_result, int) else None + + async def _count_remaining( + self, prisma_client: PrismaClient, cutoff_date: datetime, table_name: str, time_column: str, deadline: float + ) -> int | None: + """ + Count expired rows still outstanding, stopping at a cap. + + An uncapped COUNT(*) over an expired backlog would itself be the kind of + long scan this job exists to avoid, so the probe reads at most + SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP index entries. A result equal to + the cap means "at least this many". + """ + count_sql: Final = f""" + SELECT count(*)::int AS remaining FROM ( + SELECT 1 FROM "{table_name}" + WHERE "{time_column}" < $1::timestamptz + LIMIT $2 + ) capped + """ + try: + async with prisma_client.db.tx() as tx: + await tx.execute_raw(f"SET LOCAL statement_timeout = {self._timeout_ms(deadline)}") + rows: Final = _REMAINING_ROWS.validate_python( + await tx.query_raw(count_sql, cutoff_date, SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP) + ) + except Exception as e: # noqa: BLE001 - an observability probe must never fail the cleanup run + verbose_proxy_logger.warning("Could not count remaining %s rows: %s", table_name, e) + return None + return rows[0].remaining if rows else None + async def _delete_old_rows_batched( self, prisma_client: PrismaClient, @@ -85,10 +297,14 @@ class SpendLogCleanup: table_name: str, key_columns: tuple[str, ...], time_column: str, - ) -> int: + deadline: float, + ) -> TableCleanupResult: """ - Helper method to delete a table's rows older than the cutoff in batches. - Returns the total number of rows deleted. + Delete a table's rows older than the cutoff in batches. + + Stops at whichever bound is reached first: the backlog running out, the + shared wall-clock deadline, the per-table batch cap, or too many + consecutive batch failures. """ key_list: Final = ", ".join(f'"{col}"' for col in key_columns) delete_sql: Final = f""" @@ -103,23 +319,46 @@ class SpendLogCleanup: run_count = 0 consecutive_failures = 0 while True: - if run_count > SPEND_LOG_RUN_LOOPS: + if time.monotonic() >= deadline: + verbose_proxy_logger.info( + "Run budget exhausted during %s cleanup after %d rows; the next run resumes from here", + table_name, + total_deleted, + ) + return await self._finish_table( + prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline + ) + if run_count >= self.max_batches: verbose_proxy_logger.info( "Max batches reached for %s cleanup, remaining rows will be deleted in next run", table_name ) - break - # Step 1: Find rows and delete them in one go without fetching to application - # Delete in batches, limited by self.batch_size - try: - deleted_result = await prisma_client.db.execute_raw( - delete_sql, - cutoff_date, - self.batch_size, + return await self._finish_table( + prisma_client, cutoff_date, table_name, time_column, total_deleted, "batch_cap_reached", deadline ) + # Find rows and delete them in one go without fetching to application + batch_started_at = time.monotonic() + try: + batch_result = await self._execute_delete_batch(prisma_client, delete_sql, cutoff_date, deadline) except Exception as batch_exc: + if time.monotonic() >= deadline: + # The statement timeout was clamped to the budget that was + # left, so this batch was cancelled by the deadline itself. + # That is the bound working, not a database fault, and + # counting it would both inflate the failure metric and push + # every budget-exhausted run toward the abort threshold. + verbose_proxy_logger.info( + "Run budget exhausted mid-batch during %s cleanup after %d rows; " + "the next run resumes from here", + table_name, + total_deleted, + ) + return await self._finish_table( + prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline + ) # A single batch failure (e.g. Prisma/DB timeout) must not abort # the whole run — subsequent batches may still succeed. consecutive_failures += 1 + SpendLogCleanupMetrics.record_batch_failure(table_name) verbose_proxy_logger.exception( "%s cleanup batch failed " "(run_count=%d, consecutive_failures=%d, batch_size=%d, " @@ -140,28 +379,31 @@ class SpendLogCleanup: consecutive_failures, total_deleted, ) - break + return await self._finish_table( + prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline + ) await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS) continue - consecutive_failures = 0 - - deleted_count = 0 - if isinstance(deleted_result, int): - deleted_count = deleted_result - else: + if batch_result is None: verbose_proxy_logger.error( - "Unexpected execute_raw return type for %s cleanup: %s; aborting cleanup to avoid infinite loop", + "Unexpected execute_raw return type for %s cleanup; aborting cleanup to avoid infinite loop", table_name, - type(deleted_result), ) - break + return await self._finish_table( + prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline + ) + consecutive_failures = 0 + deleted_count = batch_result + SpendLogCleanupMetrics.record_batch(table_name, deleted_count, time.monotonic() - batch_started_at) verbose_proxy_logger.info("Deleted %s %s rows in this batch", deleted_count, table_name) if deleted_count == 0: verbose_proxy_logger.info("No more %s rows to delete. Total deleted: %s", table_name, total_deleted) - break + return await self._finish_table( + prisma_client, cutoff_date, table_name, time_column, total_deleted, "exhausted", deadline + ) total_deleted += deleted_count run_count += 1 @@ -169,18 +411,49 @@ class SpendLogCleanup: # Add a small sleep to prevent overwhelming the database await asyncio.sleep(0.1) - return total_deleted + async def _finish_table( + self, + prisma_client: PrismaClient, + cutoff_date: datetime, + table_name: str, + time_column: str, + rows_deleted: int, + stop_reason: StopReason, + deadline: float, + ) -> TableCleanupResult: + """ + Publish how much of this table is still outstanding, then report the run's result. - async def _delete_old_logs(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int: + The probe is skipped once the budget is spent. It is the one piece of + work that would otherwise be ISSUED after the deadline, and every table + exits through here, including the ones a spent run never started, so + keeping it would put one more statement per table past the bound. A run + that ends this way already reports "budget_exhausted", which tells an + operator the backlog was not drained; the gauge simply keeps its value + from the last run that finished inside its budget. + """ + if time.monotonic() >= deadline: + return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason) + remaining: Final = await self._count_remaining(prisma_client, cutoff_date, table_name, time_column, deadline) + if remaining is not None: + SpendLogCleanupMetrics.set_rows_remaining(table_name, remaining) + return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason) + + async def _delete_old_logs( + self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float + ) -> TableCleanupResult: return await self._delete_old_rows_batched( prisma_client, cutoff_date, table_name="LiteLLM_SpendLogs", key_columns=("request_id", "startTime"), time_column="startTime", + deadline=deadline, ) - async def _delete_old_tool_index_rows(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int: + async def _delete_old_tool_index_rows( + self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float + ) -> TableCleanupResult: # SpendLogToolIndex rows are derived from spend logs, so they expire on the # same cutoff; rows older than retention point at already-deleted logs. return await self._delete_old_rows_batched( @@ -189,17 +462,87 @@ class SpendLogCleanup: table_name="LiteLLM_SpendLogToolIndex", key_columns=("request_id", "tool_name"), time_column="start_time", + deadline=deadline, ) - async def _delete_old_autorouter_session_rows(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int: + async def _delete_old_autorouter_session_rows( + self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float + ) -> TableCleanupResult: return await self._delete_old_rows_batched( prisma_client, cutoff_date, table_name="LiteLLM_AutoRouterSession", key_columns=("api_key", "session_id", "router_name"), time_column="last_turn_at", + deadline=deadline, ) + async def _clean_spend_log_tables( + self, prisma_client: PrismaClient, deadline: float + ) -> tuple[TableCleanupResult, ...]: + """ + Prune the spend logs and the tool index rows derived from them. + + When the table is range-partitioned, whole expired partitions are dropped + first because that reclaims disk immediately. Expired rows can still sit in + the DEFAULT partition (backfill, coverage gaps) or in a partition that spans + the cutoff, so retention still deletes those stragglers row-wise. + """ + cutoff_date: Final = datetime.now(timezone.utc) - timedelta(seconds=float(self.retention_seconds or 0)) + verbose_proxy_logger.info("Removing logs older than %s", cutoff_date.isoformat()) + + # Partition maintenance is DDL taking an ACCESS EXCLUSIVE lock, so it is + # only STARTED while the run still has budget, and each statement carries + # the same timeouts the batches do. Without those, a DROP would queue + # behind any long-running reader for as long as that reader lives, which + # is the one way this job could still outlast its budget without bound. + remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline) + if time.monotonic() >= deadline: + verbose_proxy_logger.info("Run budget already spent, skipping partition maintenance this run") + elif self.general_settings.get( + "use_spend_logs_partitioning", False + ) and await self.partition_manager.is_partitioned(prisma_client, remaining_timeout_ms): + await self.partition_manager.ensure_partitions(prisma_client, remaining_timeout_ms) + dropped: Final = await self.partition_manager.drop_partitions_older_than( + prisma_client, cutoff_date, remaining_timeout_ms + ) + verbose_proxy_logger.info("Dropped %d expired spend-log partitions: %s", len(dropped), dropped) + + logs_result: Final = await self._delete_old_logs(prisma_client, cutoff_date, deadline) + verbose_proxy_logger.info("Deleted %s logs", logs_result.rows_deleted) + + index_result: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date, deadline) + verbose_proxy_logger.info("Deleted %s expired tool index rows", index_result.rows_deleted) + return (logs_result, index_result) + + async def _clean_session_rollup( + self, prisma_client: PrismaClient, retention_seconds: int, deadline: float + ) -> tuple[TableCleanupResult, ...]: + """ + Prune auto-router session rollup rows, which carry their own retention horizon. + """ + session_cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=float(retention_seconds)) + sessions_result: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff, deadline) + verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted) + return (sessions_result,) + + @staticmethod + def _run_outcome(results: tuple[TableCleanupResult, ...]) -> RunOutcome: + """ + Report the most operationally significant reason the run stopped. + + A bound that was hit matters more than a table that simply ran dry, so + those win over "completed", and an abort wins over everything. + """ + reasons: Final = frozenset(result.stop_reason for result in results) + if "aborted" in reasons: + return "aborted" + if "budget_exhausted" in reasons: + return "budget_exhausted" + if "batch_cap_reached" in reasons: + return "batch_cap_reached" + return "completed" + async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None: """ Main cleanup function. Deletes old spend logs in batches. @@ -209,16 +552,19 @@ class SpendLogCleanup: lock_acquired = False try: verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now()) + self._refresh_bounds() delete_spend_logs: Final = self._should_delete_spend_logs() autorouter_retention_seconds: Final = self._retention_seconds_for( "maximum_autorouter_session_retention_period" ) if not delete_spend_logs and autorouter_retention_seconds is None: + SpendLogCleanupMetrics.record_run("skipped_disabled") return if delete_spend_logs and self.retention_seconds is None: verbose_proxy_logger.error("Retention seconds is None, cannot proceed with cleanup") + SpendLogCleanupMetrics.record_run("skipped_disabled") return # If we have a pod lock manager, try to acquire the lock @@ -235,43 +581,23 @@ class SpendLogCleanup: if not lock_acquired: verbose_proxy_logger.info("Another pod is already running cleanup") + SpendLogCleanupMetrics.record_run("skipped_locked") return - if delete_spend_logs and self.retention_seconds is not None: - cutoff_date: Final = datetime.now(timezone.utc) - timedelta(seconds=float(self.retention_seconds)) - verbose_proxy_logger.info("Removing logs older than %s", cutoff_date.isoformat()) + deadline: Final = time.monotonic() + self.run_budget_seconds - if self.general_settings.get( - "use_spend_logs_partitioning", False - ) and await self.partition_manager.is_partitioned(prisma_client): - await self.partition_manager.ensure_partitions(prisma_client) - dropped: Final = await self.partition_manager.drop_partitions_older_than(prisma_client, cutoff_date) - verbose_proxy_logger.info( - "Dropped %d expired spend-log partitions: %s", - len(dropped), - dropped, - ) - # DROP only reclaims whole expired partitions. Expired rows can - # still sit in the DEFAULT partition (backfill, coverage gaps) - # or in a partition that spans the cutoff, so retention must - # also delete those stragglers row-wise. - total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) - verbose_proxy_logger.info( - "Deleted %s expired logs not covered by dropped partitions", total_deleted - ) - else: - total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) - verbose_proxy_logger.info("Deleted %s logs", total_deleted) + spend_log_results: Final = ( + await self._clean_spend_log_tables(prisma_client, deadline) + if delete_spend_logs and self.retention_seconds is not None + else () + ) + session_results: Final = ( + await self._clean_session_rollup(prisma_client, autorouter_retention_seconds, deadline) + if autorouter_retention_seconds is not None + else () + ) - index_deleted: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date) - verbose_proxy_logger.info("Deleted %s expired tool index rows", index_deleted) - - if autorouter_retention_seconds is not None: - session_cutoff: Final = datetime.now(timezone.utc) - timedelta( - seconds=float(autorouter_retention_seconds) - ) - sessions_deleted: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff) - verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_deleted) + SpendLogCleanupMetrics.record_run(self._run_outcome(spend_log_results + session_results)) except Exception as e: # .exception() captures the traceback; str(e) alone on a Prisma/DB @@ -281,6 +607,7 @@ class SpendLogCleanup: type(e).__name__, e, ) + SpendLogCleanupMetrics.record_run("aborted") return # Return after error handling finally: # Only release the lock if it was actually acquired diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py new file mode 100644 index 00000000000..340aeab938c --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py @@ -0,0 +1,122 @@ +""" +Prometheus metrics for the spend-log retention cleanup job. + +The job runs in the background on a single elected pod, so its cost is invisible +from request-path metrics. These instruments make a run's database footprint +observable: how much it deleted, how long each batch took, how much work is +still outstanding, and why a run stopped. + +``prometheus_client`` is an optional dependency, so every recorder degrades to a +no-op when it is absent. +""" + +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from litellm._logging import verbose_proxy_logger + +if TYPE_CHECKING: + # aliased so the annotations below cannot be mistaken for collections.Counter + from prometheus_client import Counter as PrometheusCounter + from prometheus_client import Gauge as PrometheusGauge + from prometheus_client import Histogram as PrometheusHistogram + +RunOutcome: TypeAlias = Literal[ + "completed", + "budget_exhausted", + "batch_cap_reached", + "skipped_locked", + "skipped_disabled", + "aborted", +] + +_BATCH_DURATION_BUCKETS: Final = (0.005, 0.025, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0) +_TABLE_LABEL: Final = ("table",) +_OUTCOME_LABEL: Final = ("outcome",) + + +class SpendLogCleanupMetrics: + """ + Lazily-registered Prometheus instruments for the retention cleanup job. + + Registration is deferred to first use so that importing this module never + touches the Prometheus registry, which keeps it safe to import from the + proxy regardless of whether Prometheus is a configured callback. + """ + + _initialized: bool = False + rows_deleted: "PrometheusCounter | None" = None + batch_duration: "PrometheusHistogram | None" = None + rows_remaining: "PrometheusGauge | None" = None + batch_failures: "PrometheusCounter | None" = None + runs: "PrometheusCounter | None" = None + + @classmethod + def _ensure_initialized(cls) -> None: + if cls._initialized: + return + cls._initialized = True + try: + # prometheus_client is an optional extra, so it is resolved here rather + # than at module import: this module is reachable from proxy startup + # regardless of whether Prometheus is a configured callback. + from prometheus_client import Counter, Gauge, Histogram + + cls.rows_deleted = Counter( + "litellm_spend_log_cleanup_rows_deleted_total", + "Rows deleted by the spend-log retention cleanup job", + labelnames=_TABLE_LABEL, + ) + cls.batch_duration = Histogram( + "litellm_spend_log_cleanup_batch_duration_seconds", + "Wall-clock duration of one retention cleanup delete batch", + labelnames=_TABLE_LABEL, + buckets=_BATCH_DURATION_BUCKETS, + ) + cls.rows_remaining = Gauge( + "litellm_spend_log_cleanup_rows_remaining", + "Expired rows still awaiting deletion, counted only up to " + "SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP so the probe itself cannot scan a " + "large table; a value equal to that cap means at least that many remain", + labelnames=_TABLE_LABEL, + multiprocess_mode="livemax", + ) + cls.batch_failures = Counter( + "litellm_spend_log_cleanup_batch_failures_total", + "Retention cleanup delete batches that raised", + labelnames=_TABLE_LABEL, + ) + cls.runs = Counter( + "litellm_spend_log_cleanup_runs_total", + "Retention cleanup runs, labelled by why the run ended", + labelnames=_OUTCOME_LABEL, + ) + except Exception as e: # noqa: BLE001 - a metrics problem must never fail the cleanup run + # Covers the extra being absent, a duplicate registration (repeated + # imports under a test runner), and registry misconfiguration alike. + verbose_proxy_logger.warning("Could not register spend-log cleanup metrics: %s", e) + + @classmethod + def record_batch(cls, table_name: str, rows_deleted: int, duration_seconds: float) -> None: + cls._ensure_initialized() + if cls.rows_deleted is not None: + cls.rows_deleted.labels(table=table_name).inc(rows_deleted) + if cls.batch_duration is not None: + cls.batch_duration.labels(table=table_name).observe(duration_seconds) + + @classmethod + def record_batch_failure(cls, table_name: str) -> None: + cls._ensure_initialized() + if cls.batch_failures is not None: + cls.batch_failures.labels(table=table_name).inc() + + @classmethod + def set_rows_remaining(cls, table_name: str, remaining: int) -> None: + cls._ensure_initialized() + if cls.rows_remaining is not None: + cls.rows_remaining.labels(table=table_name).set(remaining) + + @classmethod + def record_run(cls, outcome: RunOutcome) -> None: + cls._ensure_initialized() + if cls.runs is not None: + cls.runs.labels(outcome=outcome).inc() diff --git a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py index df17721d8e5..221c142d9d3 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py +++ b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py @@ -14,8 +14,9 @@ keeps the batched-DELETE path, so existing deployments are untouched. """ import re +from collections.abc import Callable from datetime import date, datetime, timedelta, timezone -from typing import Final +from typing import TYPE_CHECKING, Final, TypeAlias from litellm._logging import verbose_proxy_logger from litellm.constants import ( @@ -23,8 +24,23 @@ from litellm.constants import ( SPEND_LOG_PARTITION_PRECREATE_AHEAD, ) +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs" +RemainingTimeoutMs: TypeAlias = Callable[[], "int | None"] +""" +The per-statement bound in milliseconds, or None once the caller's budget is +spent. + +Injected rather than passed as a number so it is re-evaluated before EVERY +statement: a value read once at entry would let a loop issue N statements each +bounded by the budget that was left before the first of them, which is not a +bound on the loop at all. The caller owns the policy; this module only asks how +much time it may still use. +""" + PartitionInterval = str # "day" | "week" | "month" VALID_PARTITION_INTERVALS: Final = {"day", "week", "month"} @@ -116,21 +132,26 @@ class SpendLogsPartitionManager: self.interval = interval self.precreate_ahead = precreate_ahead - async def is_partitioned(self, prisma_client) -> bool: + async def is_partitioned(self, prisma_client: "PrismaClient", remaining_timeout_ms: RemainingTimeoutMs) -> bool: + budget_ms: Final = remaining_timeout_ms() + if budget_ms is None: + return False try: - rows: Final = await prisma_client.db.query_raw( - """ - SELECT EXISTS ( - SELECT 1 - FROM pg_partitioned_table pt - JOIN pg_class c ON c.oid = pt.partrelid - JOIN pg_namespace n ON n.oid = c.relnamespace - WHERE c.relname = $1 - AND n.nspname = current_schema() - ) AS partitioned - """, - SPEND_LOGS_TABLE, - ) + async with prisma_client.db.tx() as tx: + await tx.execute_raw(f"SET LOCAL statement_timeout = {budget_ms}") + rows: Final = await tx.query_raw( + """ + SELECT EXISTS ( + SELECT 1 + FROM pg_partitioned_table pt + JOIN pg_class c ON c.oid = pt.partrelid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE c.relname = $1 + AND n.nspname = current_schema() + ) AS partitioned + """, + SPEND_LOGS_TABLE, + ) except Exception as e: verbose_proxy_logger.warning( "Could not determine if %s is partitioned, assuming it is not: %s", @@ -140,7 +161,25 @@ class SpendLogsPartitionManager: return False return bool(rows and rows[0].get("partitioned")) - async def ensure_partitions(self, prisma_client) -> list[str]: + @staticmethod + async def _execute_bounded_ddl(prisma_client: "PrismaClient", statement: str, timeout_ms: int) -> None: + """ + Run one DDL statement under a Postgres statement and lock timeout. + + Partition DDL takes an ACCESS EXCLUSIVE lock, so an unbounded statement + queues behind any long-running reader for as long as that reader lives, + and the caller's run budget cannot cut it short. lock_timeout bounds the + wait for the lock and statement_timeout bounds the work itself, so a + partition this run cannot get is simply left for the next one. + """ + async with prisma_client.db.tx() as tx: + await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}") + await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}") + await tx.execute_raw(statement) + + async def ensure_partitions( + self, prisma_client: "PrismaClient", remaining_timeout_ms: RemainingTimeoutMs + ) -> list[str]: """ Ensure the current and upcoming partitions exist, returning the names now present. CREATE TABLE IF NOT EXISTS is a no-op for partitions that @@ -150,42 +189,61 @@ class SpendLogsPartitionManager: for name, lower, upper in upcoming_partitions( datetime.now(timezone.utc).date(), self.interval, self.precreate_ahead ): + budget_ms = remaining_timeout_ms() + if budget_ms is None: + verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run") + break try: - await prisma_client.db.execute_raw( + await self._execute_bounded_ddl( + prisma_client, f'CREATE TABLE IF NOT EXISTS "{name}" ' f'PARTITION OF "{SPEND_LOGS_TABLE}" ' - f"FOR VALUES FROM ('{lower.isoformat()}') TO ('{upper.isoformat()}')" + f"FOR VALUES FROM ('{lower.isoformat()}') TO ('{upper.isoformat()}')", + budget_ms, ) ensured.append(name) except Exception as e: verbose_proxy_logger.warning("Failed to ensure spend-log partition %s: %s", name, e) return ensured - async def _list_partitions(self, prisma_client) -> list[tuple[str, datetime | None]]: - rows: Final = await prisma_client.db.query_raw( - """ - SELECT c.relname AS name, - pg_get_expr(c.relpartbound, c.oid) AS bound - FROM pg_inherits i - JOIN pg_class c ON c.oid = i.inhrelid - JOIN pg_class p ON p.oid = i.inhparent - JOIN pg_namespace n ON n.oid = p.relnamespace - WHERE p.relname = $1 - AND n.nspname = current_schema() - """, - SPEND_LOGS_TABLE, - ) + async def _list_partitions( + self, prisma_client: "PrismaClient", timeout_ms: int + ) -> list[tuple[str, datetime | None]]: + async with prisma_client.db.tx() as tx: + await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}") + rows: Final = await tx.query_raw( + """ + SELECT c.relname AS name, + pg_get_expr(c.relpartbound, c.oid) AS bound + FROM pg_inherits i + JOIN pg_class c ON c.oid = i.inhrelid + JOIN pg_class p ON p.oid = i.inhparent + JOIN pg_namespace n ON n.oid = p.relnamespace + WHERE p.relname = $1 + AND n.nspname = current_schema() + """, + SPEND_LOGS_TABLE, + ) return [(row["name"], parse_partition_upper_bound(row.get("bound") or "")) for row in rows] - async def drop_partitions_older_than(self, prisma_client, cutoff: datetime) -> list[str]: + async def drop_partitions_older_than( + self, prisma_client: "PrismaClient", cutoff: datetime, remaining_timeout_ms: RemainingTimeoutMs + ) -> list[str]: """DROP every partition whose whole range is older than `cutoff`.""" + list_budget_ms: Final = remaining_timeout_ms() + if list_budget_ms is None: + return [] cutoff_naive: Final = cutoff.astimezone(timezone.utc).replace(tzinfo=None) - partitions: Final = await self._list_partitions(prisma_client) + partitions: Final = await self._list_partitions(prisma_client, list_budget_ms) to_drop: Final = select_partitions_to_drop(partitions, cutoff_naive) dropped: Final[list[str]] = [] for name in to_drop: + budget_ms = remaining_timeout_ms() + if budget_ms is None: + verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run") + break try: - await prisma_client.db.execute_raw(f'DROP TABLE IF EXISTS "{name}"') + await self._execute_bounded_ddl(prisma_client, f'DROP TABLE IF EXISTS "{name}"', budget_ms) dropped.append(name) except Exception as e: verbose_proxy_logger.warning("Failed to drop spend-log partition %s: %s", name, e) diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 8287ef1addf..5aeb52be535 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -103,6 +103,17 @@ class RoutingPrismaWrapper: def reader(self) -> PrismaWrapper: return self._reader + @property + def read_target(self) -> PrismaWrapper: + """The wrapper `_TOP_LEVEL_READ_METHODS` dispatch to right now. + + Callers that need to reason about the engine a read actually ran on + (e.g. recovering from prepared statements that went stale on it) must + consult this rather than `writer`, and `__getattr__` routes through it + so the two cannot drift apart. + """ + return self._writer if self._reader_unavailable else self._reader + @property def reader_unavailable(self) -> bool: return self._reader_unavailable @@ -254,8 +265,7 @@ class RoutingPrismaWrapper: def __getattr__(self, name: str) -> Any: if name in _TOP_LEVEL_READ_METHODS: - target: Final = self._writer if self._reader_unavailable else self._reader - return getattr(target, name) + return getattr(self.read_target, name) writer_attr: Final = getattr(self._writer, name) # Per-model action accessors are non-callable instances that expose # both `find_many` and `create`. Methods like execute_raw / batch_ / diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 251ed1feb10..0a5626ba0a7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -274,7 +274,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) -_UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset( +UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset( { "litellm-disable-message-redaction", } @@ -355,7 +355,7 @@ def _strip_untrusted_request_header_controls( return for header_name in list(headers.keys()): - if isinstance(header_name, str) and header_name.lower() in _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: + if isinstance(header_name, str) and header_name.lower() in UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: if allow_client_message_redaction_opt_out: continue headers.pop(header_name, None) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index c00c2a5ba4c..2271501d480 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -14,11 +14,11 @@ from litellm.proxy.auth.auth_checks import ( _cache_access_object, _cache_key_object, _cache_team_object, - _delete_cache_access_object, _get_team_object_from_cache, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_cache from litellm.proxy.utils import get_prisma_client_or_throw from litellm.repositories.table_repositories import AccessGroupRepository from litellm.types.access_group import ( @@ -146,22 +146,6 @@ async def _cache_access_group_record(record: _AccessGroupRecord) -> None: ) -async def _invalidate_cache_access_group(access_group_id: str) -> None: - """ - Invalidate (delete) an access group entry from both in-memory and Redis caches. - - Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server - to avoid circular imports, following the same pattern as key_management_endpoints. - """ - from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache - - await _delete_cache_access_object( - access_group_id=access_group_id, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - # --------------------------------------------------------------------------- # DB sync helpers (called inside a Prisma transaction) # --------------------------------------------------------------------------- @@ -595,7 +579,7 @@ async def delete_access_group( from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache - await _invalidate_cache_access_group(access_group_id) + await invalidate_access_group_cache(access_group_id) await _patch_team_caches_remove_access_group( affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 116dea464ff..7e190e8b19d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -88,6 +88,11 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) +from litellm.proxy.management_helpers.access_group_key_sync import ( + sync_key_access_group_membership, + sync_key_regeneration_access_group_membership, + sync_key_update_access_group_membership, +) from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, @@ -2347,6 +2352,17 @@ async def _process_single_key_update( proxy_logging_obj=proxy_logging_obj, ) + # After the key's own cache entry is dropped, so a failure here cannot leave the key + # authenticating against the access groups it just lost. + await sync_key_update_access_group_membership( + prisma_client=prisma_client, + key_token=_hash_token_if_needed( + _resolve_token_to_update(data=update_key_request, existing_key_row=existing_key_row) + ), + data=update_key_request, + existing_key_row=existing_key_row, + ) + # Trigger async hook asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( @@ -2828,6 +2844,15 @@ async def update_key_fn( proxy_logging_obj=proxy_logging_obj, ) + # After the key's own cache entry is dropped, so a failure here cannot leave the key + # authenticating against the access groups it just lost. + await sync_key_update_access_group_membership( + prisma_client=prisma_client, + key_token=_hash_token_if_needed(key), + data=data, + existing_key_row=existing_key_row, + ) + if data.spend is not None: from litellm.proxy.proxy_server import spend_counter_cache @@ -3771,7 +3796,7 @@ async def generate_key_helper_fn( auto_rotate: bool | None = None, rotation_interval: str | None = None, router_settings: dict | None = None, - access_group_ids: list | None = None, + access_group_ids: list[str] | None = None, budget_limits: list | None = None, # multiple concurrent budget windows ): from litellm.proxy.proxy_server import premium_user, prisma_client @@ -3979,6 +4004,14 @@ async def generate_key_helper_fn( create_key_response: Final = await prisma_client.insert_data(data=key_data, table_name="key") key_data["token_id"] = getattr(create_key_response, "token", None) + created_token_hash: Final = getattr(create_key_response, "token", None) + if isinstance(created_token_hash, str): + await sync_key_access_group_membership( + prisma_client=prisma_client, + key_token=created_token_hash, + previous_access_group_ids=None, + updated_access_group_ids=access_group_ids, + ) key_data["litellm_budget_table"] = getattr(create_key_response, "litellm_budget_table", None) key_data["created_at"] = getattr(create_key_response, "created_at", None) key_data["updated_at"] = getattr(create_key_response, "updated_at", None) @@ -4196,6 +4229,7 @@ async def delete_verification_tokens( deleted_tokens = [key.token for key in authorized_keys] if len(deleted_tokens) != len(tokens): failed_tokens = [token for token in tokens if token not in deleted_tokens] + else: raise Exception("DB not connected. prisma_client is None") except Exception as e: @@ -4211,6 +4245,16 @@ async def delete_verification_tokens( hashed_token = hash_token(cast(str, key)) user_api_key_cache.delete_cache(hashed_token) + # After credential invalidation, so a failure here can never keep a deleted key alive. + for deleted_key in authorized_keys: + if deleted_key.token is not None: + await sync_key_access_group_membership( + prisma_client=prisma_client, + key_token=deleted_key.token, + previous_access_group_ids=deleted_key.access_group_ids, + updated_access_group_ids=None, + ) + return { "deleted_keys": deleted_tokens, "failed_tokens": failed_tokens, @@ -4726,6 +4770,15 @@ async def _execute_virtual_key_regeneration( proxy_logging_obj=proxy_logging_obj, ) + # After credential invalidation, so a failure here can never keep the old key alive. + await sync_key_regeneration_access_group_membership( + prisma_client=prisma_client, + previous_key_token=hashed_api_key, + new_key_token=new_token_hash, + data=data, + existing_key_row=key_in_db, + ) + response: Final = GenerateKeyResponse.model_validate(updated_token_dict) asyncio.create_task( KeyManagementEventHooks.async_key_rotated_hook( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 4bc72f554b1..912e18150b3 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -89,6 +89,7 @@ from litellm.types.router import ( ModelInfo, updateDeployment, ) +from litellm.types.utils import CustomPricingLiteLLMParams from litellm.utils import get_utc_datetime router: Final = APIRouter() @@ -242,6 +243,7 @@ def _raise_on_strategy_router_write_violation( _PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to") +_PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"}) def _explicitly_cleared_ptu_fields(model_info: ModelInfo | None) -> frozenset[str]: @@ -265,9 +267,10 @@ def _merged_ptu_model_info(*, db_model: Deployment, patch_data: updateDeployment A PTU invariant holds over the deployment as it will exist, not over whichever subset of fields a caller happened to send. """ - empty: Final[Mapping[str, object]] = MappingProxyType({}) - stored: Final = db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else empty - incoming: Final = patch_data.model_info.model_dump(exclude_none=True) if patch_data.model_info else empty + stored: Final = db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else _EMPTY_MODEL_INFO + incoming: Final = ( + patch_data.model_info.model_dump(exclude_none=True) if patch_data.model_info else _EMPTY_MODEL_INFO + ) cleared: Final = _explicitly_cleared_ptu_fields(patch_data.model_info) return MappingProxyType({k: v for k, v in {**stored, **incoming}.items() if k not in cleared}) @@ -339,6 +342,140 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: ) +# The six mirrored pricing fields plus the three remaining fields +# Router._inherit_builtin_cache_pricing back-fills from the public cost map. An unset field is +# what that back-fill targets, so a field left out here is one a PTU deployment still bills. +_PTU_ZEROED_PRICING_FIELDS: Final = SPECIAL_MODEL_INFO_PARAMS + ( + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost_above_200k_tokens", +) +_PTU_ZEROED_PRICING: Final[Mapping[str, float]] = MappingProxyType(dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0)) +_NO_PRICING_OVERRIDE: Final[Mapping[str, float]] = MappingProxyType({}) +_EMPTY_MODEL_INFO: Final[Mapping[str, object]] = _NO_PRICING_OVERRIDE +# Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges +# (an embedding's output_vector_size, the regional uplift multipliers), and zeroing one of +# those would destroy the deployment's configuration rather than stop a charge. +_CUSTOM_PRICING_FIELDS: Final = frozenset(f for f in CustomPricingLiteLLMParams.model_fields if "cost" in f) + + +def _is_nonzero_price(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value != 0 + + +def _is_zero_price(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value == 0 + + +def _raise_if_ptu_deployment_is_priced(*, model_info: Mapping[str, object], supplied: Mapping[str, object]) -> None: + """Refuse a rate the caller supplies for a deployment that bills reserved capacity. + + Separate from the zeroing so the team-model path can run it before it touches the team, whose + ACL write autocommits: a refusal raised after it would leave the team changed and the + deployment row never written. + """ + if not is_ptu_cost_attribution_enabled(): + return + if model_info.get("ptu_count") is None or model_info.get("cost_per_ptu_per_hour") is None: + return + priced: Final = tuple(sorted(field for field in _CUSTOM_PRICING_FIELDS if _is_nonzero_price(supplied.get(field)))) + if not priced: + return + raise HTTPException( + status_code=400, + detail=( + f"A PTU deployment bills by reserved capacity, so {', '.join(priced)} cannot be charged on " + "top of it. Send 0 or no value, or remove ptu_count and cost_per_ptu_per_hour to bill per token." + ), + ) + + +def _ptu_zeroed_pricing( + *, + model_info: Mapping[str, object], + litellm_params: Mapping[str, object], + supplied: Mapping[str, object], +) -> Mapping[str, float]: + """The pricing a PTU deployment must carry, empty unless one is being stored. + + Reserved capacity is already billed by the flat cost the rollup writes, so charging the + traffic it serves bills the same tokens twice. Left unset the rate falls back to the public + cost map, which makes the double charge the default rather than an opt-in. + + Only a price the caller supplies is refused. A non-zero price already on the row is zeroed + instead, so a deployment priced through a path this rule does not cover heals on its next + save rather than rejecting every later edit of a field that has nothing to do with pricing. + + ``supplied`` is the caller's litellm_params alone, because that is the blob a price is + authored on. model_info's copy is written by the server, both by the mirror in + ``Deployment.__init__`` and by the cost-map defaults /model/info fills in, so a client that + round-trips a model_info blob sends back prices it never chose. + """ + if not is_ptu_cost_attribution_enabled(): + return _NO_PRICING_OVERRIDE + if model_info.get("ptu_count") is None or model_info.get("cost_per_ptu_per_hour") is None: + return _NO_PRICING_OVERRIDE + _raise_if_ptu_deployment_is_priced(model_info=model_info, supplied=supplied) + stored: Final = frozenset( + field + for field in _CUSTOM_PRICING_FIELDS + if _is_nonzero_price(model_info.get(field)) or _is_nonzero_price(litellm_params.get(field)) + ) + if not stored: + return _PTU_ZEROED_PRICING + return MappingProxyType({**_PTU_ZEROED_PRICING, **dict.fromkeys(stored, 0.0)}) + + +def _ptu_pricing_delta( + *, + stored_model_info: Mapping[str, object], + model_info: Mapping[str, object], + litellm_params: Mapping[str, object], + patch: updateDeployment, +) -> tuple[Mapping[str, float], frozenset[str]]: + """The pricing a patch must write into both blobs, and the pricing it must drop from them. + + A patch that takes the deployment off PTU takes the zeroed pricing with it, since the zeros + exist only to stop the double charge. Left behind they would serve the deployment for free. + Reading the stored row rather than the patch alone keeps that release off a deployment that + never carried PTU config, whose zero price is a rate its operator chose. A zero the patch + itself carries is released with the rest, because the dashboard echoes the whole stored + blob on every save, so a supplied zero cannot be told apart from the one this rule wrote. + + The release spans every field the zeroing could have written, not just the mirrored ones, or + a rate zeroed on the way in (per-second, per-character tiers) would bill nothing forever. + """ + supplied: Final = patch.litellm_params.model_dump(exclude_none=True) if patch.litellm_params else _EMPTY_MODEL_INFO + zeroed: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=supplied) + if zeroed: + return zeroed, frozenset() + was_ptu: Final = any(stored_model_info.get(field) is not None for field in _PTU_PRICED_PAIR) + if not was_ptu or not _explicitly_cleared_ptu_fields(patch.model_info) & _PTU_PRICED_PAIR: + return _NO_PRICING_OVERRIDE, frozenset() + return _NO_PRICING_OVERRIDE, frozenset( + field + for field in _CUSTOM_PRICING_FIELDS.union(_PTU_ZEROED_PRICING_FIELDS) + if _is_zero_price(model_info.get(field)) or _is_zero_price(litellm_params.get(field)) + ) + + +def _ptu_priced_deployment(model_params: Deployment) -> Deployment: + """``model_params`` with PTU pricing applied, or itself when it configures no PTU.""" + model_info: Final = model_params.model_info.model_dump(exclude_none=True) + litellm_params: Final = model_params.litellm_params.model_dump(exclude_none=True) + override: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=litellm_params) + if not override: + return model_params + return model_params.model_copy( + update=MappingProxyType( + { + "litellm_params": model_params.litellm_params.model_copy(update=override), + "model_info": model_params.model_info.model_copy(update=override), + } + ) + ) + + def _parse_ptu_datetime(value: object) -> datetime.datetime | None: """``value`` as a datetime, parsing an ISO string, else None.""" if isinstance(value, datetime.datetime): @@ -404,6 +541,19 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr merged_model_info.pop(field, None) _validate_ptu_model_info(merged_model_info) + ptu_pricing, ptu_released = _ptu_pricing_delta( + stored_model_info=db_model.model_info.model_dump(exclude_none=True) + if db_model.model_info + else _EMPTY_MODEL_INFO, + model_info=merged_model_info, + litellm_params=merged_litellm_params, + patch=updated_patch, + ) + merged_model_info.update(ptu_pricing) + merged_litellm_params.update(ptu_pricing) + for field in ptu_released: + merged_model_info.pop(field, None) + merged_litellm_params.pop(field, None) # convert to prisma compatible format @@ -863,6 +1013,12 @@ async def _update_team_model_in_db( if patch_data.model_info is not None: _raise_if_ptu_cost_attribution_disabled(patch_data.model_info.model_dump(exclude_none=True)) _validate_ptu_model_info(_merged_ptu_model_info(db_model=db_model, patch_data=patch_data)) + _raise_if_ptu_deployment_is_priced( + model_info=_merged_ptu_model_info(db_model=db_model, patch_data=patch_data), + supplied=( + patch_data.litellm_params.model_dump(exclude_none=True) if patch_data.litellm_params else _EMPTY_MODEL_INFO + ), + ) patch_team_id: Final = patch_data.model_info.team_id if patch_data.model_info else None @@ -1589,6 +1745,7 @@ async def add_new_model( incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True) _raise_if_ptu_cost_attribution_disabled(incoming_model_info) _validate_ptu_model_info(incoming_model_info) + priced_model_params: Final = _ptu_priced_deployment(model_params) if store_model_in_db is True: """ @@ -1602,13 +1759,13 @@ async def add_new_model( _original_litellm_model_name: Final = model_params.model_name if model_params.model_info.team_id is None: model_response = await _add_model_to_db( - model_params=model_params, + model_params=priced_model_params, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) else: model_response = await _add_team_model_to_db( - model_params=model_params, + model_params=priced_model_params, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) @@ -1620,9 +1777,9 @@ async def add_new_model( if "slack" in _alerting: # send notification - new model added await proxy_logging_obj.slack_alerting_instance.model_added_alert( - model_name=model_params.model_name, + model_name=priced_model_params.model_name, litellm_model_name=_original_litellm_model_name, - passed_model_info=model_params.model_info, + passed_model_info=priced_model_params.model_info, ) except Exception as e: verbose_proxy_logger.exception("Exception in add_new_model: %s", e) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b72b5d218b2..3d7f0808fb9 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -15,6 +15,7 @@ import math import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone +from types import MappingProxyType from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast import fastapi @@ -106,6 +107,12 @@ from litellm.proxy.management_endpoints.organization_endpoints import ( from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, ) +from litellm.proxy.management_helpers.access_group_team_sync import ( + AccessGroupSyncTx, + invalidate_access_group_caches, + reconcile_team_access_group_membership, + sync_team_access_group_membership, +) from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, enforce_all_proxy_mcp_servers_grant_is_admin_only, @@ -315,10 +322,17 @@ class _TeamIdInFilter(TypedDict, total=False): team_id: Mapping[str, Sequence[str]] +class _TeamCreateTx(AccessGroupSyncTx, Protocol): + @property + def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ... + + _STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """ UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams) """ +_INCLUDE_MODEL_TABLE: Final = MappingProxyType({"litellm_model_table": True}) + def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable) @@ -1511,10 +1525,15 @@ async def new_team( complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict) team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict - team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create( - data=team_creation_data, - include={"litellm_model_table": True}, - ) + tx: _TeamCreateTx + async with prisma_client.db.tx() as tx: + team_row: Final[LiteLLM_TeamTable] = await tx.litellm_teamtable.create( + data=team_creation_data, + include=_INCLUDE_MODEL_TABLE, + ) + affected_access_groups: Final = await reconcile_team_access_group_membership(tx, team_row.team_id) + + await invalidate_access_group_caches(affected_access_groups) ## ADD TEAM ID TO USER TABLE ## team_member_add_request: Final = TeamMemberAddRequest( @@ -2217,6 +2236,7 @@ async def update_team( ) verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id) + await sync_team_access_group_membership(prisma_client=prisma_client, team_id=team_row.team_id) await _refresh_cached_team( team_row=team_row, user_api_key_cache=user_api_key_cache, @@ -3850,6 +3870,9 @@ async def delete_team( # keeping the first one means a failure here still leaves a team the admin can retry deleting. await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client) + for deleted_team in team_rows: + await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id) + return deleted_teams diff --git a/litellm/proxy/management_helpers/access_group_key_sync.py b/litellm/proxy/management_helpers/access_group_key_sync.py new file mode 100644 index 00000000000..5d43cb29978 --- /dev/null +++ b/litellm/proxy/management_helpers/access_group_key_sync.py @@ -0,0 +1,173 @@ +""" +Reverse sync for the key side of the key <-> access group relationship. + +`litellm_accessgrouptable.assigned_key_ids` and `litellm_verificationtoken.access_group_ids` +are the two halves of one relationship and BOTH are read: the access group's +attached-keys view reads the former, and so does the grant check in +`auth_checks.get_authorized_resources_from_key_access_groups`, which authorizes a +key only when the group lists the key's token (or the key's team). The access-group +endpoints maintain both halves already; this module is what the key write paths call +so an edit from that side is mirrored back. + +Every write is a single guarded statement rather than a read-modify-write. Prisma has no +atomic scalar-list removal (see `TeamRepository.remove_member`), and the read-modify-write +it otherwise forces is not safe here: a lost update would put an already revoked token back +into a group and restore its grants, or drop a grant an admin just made. The guards also +make each statement idempotent, so a retry cannot duplicate an entry. Each statement covers +every group the request touches at once, so the size of the caller's id list does not turn +into a matching number of round trips, and returns the ids it actually moved so only those +groups are dropped from cache. + +It deliberately lives outside `access_group_endpoints`, which is a lazily +registered feature router (see `_lazy_features.LAZY_FEATURES`). Importing that +module eagerly from `key_management_endpoints` would put it in `sys.modules` +without its router ever being included, which drops its routes from the OpenAPI +schema. +""" + +from collections.abc import Sequence +from typing import Final, Protocol + +from pydantic import BaseModel + +from litellm.proxy._types import ( + LiteLLM_VerificationToken, + RegenerateKeyRequest, + UpdateKeyRequest, +) +from litellm.proxy.auth.auth_checks import ( + _delete_cache_access_object, # pyright: ignore[reportPrivateUsage] # the access-group endpoints reach for this same cache primitive +) +from litellm.repositories.table_repositories import AccessGroupRepository + + +class _MovedGroupRow(BaseModel): + access_group_id: str + + +class _RawExecutor(Protocol): + async def query_raw(self, query: str, *args: str | Sequence[str]) -> Sequence[object]: ... + + +_ATTACH_KEY_SQL: Final = ( + 'UPDATE "LiteLLM_AccessGroupTable" ' + 'SET "assigned_key_ids" = array_append("assigned_key_ids", $1) ' + 'WHERE "access_group_id" = ANY($2::text[]) AND NOT ($1 = ANY("assigned_key_ids")) ' + 'RETURNING "access_group_id"' +) + +_DETACH_KEY_SQL: Final = ( + 'UPDATE "LiteLLM_AccessGroupTable" ' + 'SET "assigned_key_ids" = array_remove("assigned_key_ids", $1) ' + 'WHERE "access_group_id" = ANY($2::text[]) AND $1 = ANY("assigned_key_ids") ' + 'RETURNING "access_group_id"' +) + +_REPOINT_KEY_SQL: Final = ( + 'UPDATE "LiteLLM_AccessGroupTable" ' + 'SET "assigned_key_ids" = array_append(array_remove(array_remove("assigned_key_ids", $1), $2), $2) ' + 'WHERE $1 = ANY("assigned_key_ids") ' + 'RETURNING "access_group_id"' +) + + +def _raw_executor(prisma_client: object) -> _RawExecutor: + """Narrow the untyped Prisma client down to the raw-query call this module makes.""" + return AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + + +async def _invalidate_access_group_cache(access_group_id: str) -> None: + """ + Drop an access group entry from both the in-memory and Redis caches. + + Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server + to avoid circular imports, following the same pattern as key_management_endpoints. + """ + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + + await _delete_cache_access_object( + access_group_id=access_group_id, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def _invalidate_moved_groups(moved_rows: Sequence[object]) -> None: + for row in moved_rows: + await _invalidate_access_group_cache(_MovedGroupRow.model_validate(row).access_group_id) + + +async def _write_membership(prisma_client: object, sql: str, access_group_ids: frozenset[str], key_token: str) -> None: + """Run one guarded membership statement for every listed group, dropping the cache of those it moved.""" + if not access_group_ids: + return + await _invalidate_moved_groups( + await _raw_executor(prisma_client).query_raw(sql, key_token, sorted(access_group_ids)) + ) + + +async def sync_key_access_group_membership( + prisma_client: object, + key_token: str, + previous_access_group_ids: Sequence[str] | None, + updated_access_group_ids: Sequence[str] | None, +) -> None: + """Mirror a key-side change to `access_group_ids` onto each access group's `assigned_key_ids`.""" + previous: Final = frozenset(previous_access_group_ids or ()) + updated: Final = frozenset(updated_access_group_ids or ()) + + await _write_membership(prisma_client, _ATTACH_KEY_SQL, updated - previous, key_token) + await _write_membership(prisma_client, _DETACH_KEY_SQL, previous - updated, key_token) + + +async def sync_key_update_access_group_membership( + prisma_client: object, + key_token: str, + data: UpdateKeyRequest | RegenerateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, +) -> None: + """ + Mirror a key UPDATE onto the group side, honouring `exclude_unset` semantics. + + The key row is written from `model_dump(exclude_unset=True)`, so a request that never + mentions `access_group_ids` leaves the key's own list alone and must leave the group's + copy alone too. Reading the attribute instead of `model_fields_set` would see None on + every unrelated edit and withdraw the token from every group it belongs to. + """ + if "access_group_ids" not in data.model_fields_set: + return + await sync_key_access_group_membership( + prisma_client=prisma_client, + key_token=key_token, + previous_access_group_ids=existing_key_row.access_group_ids, + updated_access_group_ids=data.access_group_ids, + ) + + +async def sync_key_regeneration_access_group_membership( + prisma_client: object, + previous_key_token: str, + new_key_token: str, + data: RegenerateKeyRequest | None, + existing_key_row: LiteLLM_VerificationToken, +) -> None: + """ + Re-point every group's copy from the old token to the regenerated one. + + Regeneration replaces the token, which is the identity `assigned_key_ids` stores, so + leaving the old hash behind both points the group at a row that no longer exists and + denies the regenerated key the group's grants. The swap is driven by the groups that + hold the old token when the statement runs, not by the key row read earlier, so a group + edited in between is neither resurrected nor skipped. Removing the new token before + appending it keeps a re-run from duplicating it. + """ + await _invalidate_moved_groups( + await _raw_executor(prisma_client).query_raw(_REPOINT_KEY_SQL, previous_key_token, new_key_token) + ) + if data is not None: + await sync_key_update_access_group_membership( + prisma_client=prisma_client, + key_token=new_key_token, + data=data, + existing_key_row=existing_key_row, + ) diff --git a/litellm/proxy/management_helpers/access_group_team_sync.py b/litellm/proxy/management_helpers/access_group_team_sync.py new file mode 100644 index 00000000000..55c0346e375 --- /dev/null +++ b/litellm/proxy/management_helpers/access_group_team_sync.py @@ -0,0 +1,155 @@ +""" +Reverse sync for the team side of the team <-> access group relationship. + +`litellm_accessgrouptable.assigned_team_ids` and `litellm_teamtable.access_group_ids` +are two copies of the same relationship, and both are read: the access group's +attached-teams view reads the former, and so does the key-side grant check in +`auth_checks.get_authorized_resources_from_key_access_groups`. The access-group +endpoints maintain both copies already; this module is what the team write paths +call so an edit from that side is mirrored back. + +It deliberately lives outside `access_group_endpoints`, which is a lazily +registered feature router (see `_lazy_features.LAZY_FEATURES`). Importing that +module eagerly from `team_endpoints` would put it in `sys.modules` without its +router ever being included, which drops its routes from the OpenAPI schema. +""" + +import asyncio +from collections.abc import Mapping, Sequence +from typing import Final, Protocol + +from pydantic import BaseModel, TypeAdapter + +from litellm.proxy.auth.auth_checks import _delete_cache_access_object + +# hashtext collisions only cost two unrelated teams a little serialization, and the +# lock is never taken by the access-group endpoints, so it cannot join their +# access-group-then-team lock order to form a cycle. +_LOCK_TEAM_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked" + +_READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1' + +# The groups the team is on either side of the reconcile, so the cache step is driven by +# desired state rather than by which rows this attempt happened to change. A retry after a +# failed invalidation finds the same set even though its statements are already no-ops. +_AFFECTED_SQL: Final = """ +SELECT access_group_id FROM "LiteLLM_AccessGroupTable" +WHERE access_group_id = ANY($2::TEXT[]) + OR $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[])) +""" + +_ATTACH_SQL: Final = """ +UPDATE "LiteLLM_AccessGroupTable" +SET assigned_team_ids = array_append(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]), $1) +WHERE access_group_id = ANY($2::TEXT[]) + AND NOT ($1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]))) +RETURNING access_group_id +""" + +_DETACH_SQL: Final = """ +UPDATE "LiteLLM_AccessGroupTable" +SET assigned_team_ids = array_remove(assigned_team_ids, $1) +WHERE $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[])) + AND NOT (access_group_id = ANY($2::TEXT[])) +RETURNING access_group_id +""" + + +class _AffectedGroup(BaseModel): + access_group_id: str + + +class _TeamGroups(BaseModel): + access_group_ids: tuple[str, ...] | None = None + + +_AffectedGroups: Final = TypeAdapter(tuple[_AffectedGroup, ...]) +_TeamRows: Final = TypeAdapter(tuple[_TeamGroups, ...]) + + +class AccessGroupSyncTx(Protocol): + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... + + +class _Transaction(Protocol): + async def __aenter__(self) -> AccessGroupSyncTx: ... + + async def __aexit__(self, *exc_info: object) -> None: ... + + +class _PrismaDb(Protocol): + def tx(self) -> _Transaction: ... + + +class _PrismaClient(Protocol): + @property + def db(self) -> _PrismaDb: ... + + +async def invalidate_access_group_cache(access_group_id: str) -> None: + """ + Drop an access group entry from both the in-memory and Redis caches. + + Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server + to avoid circular imports, following the same pattern as key_management_endpoints. + """ + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + + await _delete_cache_access_object( + access_group_id=access_group_id, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def invalidate_access_group_caches(access_group_ids: Sequence[str]) -> None: + """ + Drop every given access group from the caches, then raise if any drop failed. + + Every entry is attempted even when one raises, so a single unreachable cache cannot + leave the rest of the reconciled groups serving a grant the admin revoked. + """ + outcomes: Final = await asyncio.gather( + *(invalidate_access_group_cache(access_group_id) for access_group_id in access_group_ids), + return_exceptions=True, + ) + for outcome in outcomes: + if isinstance(outcome, BaseException): + raise outcome + + +async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id: str) -> tuple[str, ...]: + """ + Reconcile every access group's `assigned_team_ids` against the team's own + `access_group_ids`, and return the groups whose cache the caller has to drop once the + transaction commits. + + Call this inside the transaction that writes the team row, or after that row is + written or deleted: a team with no row reconciles to an empty set, which detaches it + from every group. + + The team row is read here rather than passed in, under an advisory lock held for the + rest of the transaction. That is what makes concurrent writes to the same team + converge, since each mirror reconciles against the row as the transaction sees it + instead of against the snapshot its own caller happened to see. It also means a retry + heals a sync that failed partway, where a before/after delta would compute nothing. + + Both mirror statements are set-based and mutate the array inside the statement, so a + concurrent write for a different team cannot be lost the way a read-modify-write of + the whole array can, and the pair commits together or not at all. + """ + await tx.query_raw(_LOCK_TEAM_SQL, team_id) + team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id)) + desired: Final = (team_rows[0].access_group_ids or ()) if team_rows else () + affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired)) + await tx.query_raw(_ATTACH_SQL, team_id, desired) + await tx.query_raw(_DETACH_SQL, team_id, desired) + return tuple(group.access_group_id for group in affected) + + +async def sync_team_access_group_membership(prisma_client: _PrismaClient, team_id: str) -> None: + """Reconcile the mirror for an already committed team write, in its own transaction.""" + async with prisma_client.db.tx() as tx: + affected: Final = await reconcile_team_access_group_membership(tx, team_id) + + await invalidate_access_group_caches(affected) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py index 132af097a55..597c2e742b3 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py @@ -82,7 +82,7 @@ class CoherePassthroughLoggingHandler(BasePassthroughLoggingHandler): Handle Cohere passthrough logging with route detection and cost tracking. """ # Check if this is an embed endpoint - if "/v1/embed" in url_route: + if "/v1/embed" in url_route and "/v1/embeddings" not in url_route: model: Final = request_body.get("model", response_body.get("model", "")) try: cohere_embed_config: Final = CohereEmbeddingConfig() diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 65ebc2728c6..1c8bce28454 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -31,8 +31,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( EndpointType, PassthroughStandardLoggingPayload, ) -from litellm.types.utils import ImageResponse, LlmProviders, PassthroughCallTypes -from litellm.utils import ModelResponse, TextCompletionResponse +from litellm.types.utils import EmbeddingResponse, ImageResponse, LlmProviders, PassthroughCallTypes +from litellm.utils import ModelResponse, TextCompletionResponse, convert_to_model_response_object # Hostnames that route to OpenAI-compatible APIs. # @@ -143,6 +143,14 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): "/v1/responses" in parsed_url.path or "/responses" in parsed_url.path ) + @staticmethod + def is_openai_embeddings_route(url_route: str) -> bool: + """Check if the URL route is an OpenAI embeddings endpoint.""" + if not url_route: + return False + parsed_url: Final = urlparse(url_route) + return _is_openai_compatible_host(parsed_url.hostname) and "/v1/embeddings" in parsed_url.path + def _get_user_from_metadata( self, passthrough_logging_payload: PassthroughStandardLoggingPayload, @@ -271,22 +279,21 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): **kwargs, ) -> PassThroughEndpointLoggingTypedDict: """ - Handle OpenAI passthrough logging with cost tracking for chat completions, image generation, image editing, and responses API. + Handle OpenAI passthrough logging with cost tracking for chat completions, + embeddings, image generation, image editing, and responses API. """ - # Check if this is a supported endpoint for cost tracking is_chat_completions: Final = OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route) + is_embeddings: Final = OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(url_route) is_image_generation: Final = OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route) is_image_editing: Final = OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) is_responses: Final = OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) - if not (is_chat_completions or is_image_generation or is_image_editing or is_responses): - # For unsupported endpoints, return None to let the system fall back to generic behavior + if not (is_chat_completions or is_embeddings or is_image_generation or is_image_editing or is_responses): return { "result": None, "kwargs": kwargs, } - # Extract model from request or response model: Final = request_body.get("model", response_body.get("model", "")) if not model: verbose_proxy_logger.warning("No model found in request or response for OpenAI passthrough cost tracking") @@ -307,7 +314,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): try: response_cost = 0.0 litellm_model_response: ( - ModelResponse | TextCompletionResponse | ImageResponse | ResponsesAPIResponse | None + ModelResponse | TextCompletionResponse | EmbeddingResponse | ImageResponse | ResponsesAPIResponse | None ) = None handler_instance: Final = OpenAIPassthroughLoggingHandler() @@ -338,6 +345,19 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): model=model, custom_llm_provider=custom_llm_provider, ) + elif is_embeddings: + litellm_model_response = convert_to_model_response_object( + response_object=response_body, + model_response_object=EmbeddingResponse(), + response_type="embedding", + ) + response_cost = litellm.completion_cost( + completion_response=litellm_model_response, + model=model, + custom_llm_provider=custom_llm_provider, + call_type="aembedding", + ) + litellm_model_response._hidden_params["response_cost"] = response_cost elif is_image_generation: # Handle image generation cost calculation response_cost = OpenAIPassthroughLoggingHandler._calculate_image_generation_cost( @@ -432,9 +452,13 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): endpoint_type: Final = ( "chat_completions" if is_chat_completions + else "embeddings" + if is_embeddings else "image_generation" if is_image_generation else "image_editing" + if is_image_editing + else "responses" ) verbose_proxy_logger.debug( f"OpenAI passthrough cost tracking - Endpoint: {endpoint_type}, Model: {model}, Cost: ${response_cost:.6f}" diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 3dcc8257b82..34286b203c7 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -349,10 +349,14 @@ class PassThroughEndpointLogging: return True return False - def is_cohere_route(self, url_route: str): + def is_cohere_route(self, url_route: str) -> bool: for route in self.TRACKED_COHERE_ROUTES: - if route in url_route: - return True + if route not in url_route: + continue + if route == "/v1/embed" and "/v1/embeddings" in url_route: + continue + return True + return False def is_assemblyai_route(self, url_route: str): parsed_url: Final = urlparse(url_route) @@ -429,6 +433,7 @@ class PassThroughEndpointLogging: return ( OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route) + or OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(url_route) or OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route) or OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) or OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 968c645ffcd..359187f81cb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -367,7 +367,10 @@ from litellm.proxy.config_resolvers.alerting import ( ) from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router -from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup +from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( + SPEND_LOG_CLEANUP_BOUND_SETTINGS, + SpendLogCleanup, +) from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, @@ -4079,6 +4082,7 @@ class ProxyConfig: # precedence over stale DB-cached values for these specific keys # during periodic config reloads (_update_general_settings). self._yaml_general_settings_keys: set[str] = set() # mutable-ok: populated once at startup, read-only thereafter # fmt: skip + self._yaml_spend_log_cleanup_bounds: dict[str, object] = {} # mutable-ok: snapshot of YAML bounds at load time # fmt: skip def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -5015,6 +5019,12 @@ class ProxyConfig: # These keys take precedence over DB-cached values during periodic # reloads (see _update_general_settings). self._yaml_general_settings_keys = set(general_settings.keys()) # mutable-ok: snapshot of YAML keys at load time # fmt: skip + # The VALUES matter for the cleanup bounds, not just which keys were + # set: clearing one from the dashboard has to fall back to what the + # YAML declared, and a set of names cannot answer that. + self._yaml_spend_log_cleanup_bounds = { # mutable-ok: snapshot of YAML bounds at load time # fmt: skip + key: general_settings[key] for key in SPEND_LOG_CLEANUP_BOUND_SETTINGS if key in general_settings + } ### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ### key_management_settings: Final = general_settings.get("key_management_settings", None) @@ -6299,6 +6309,18 @@ class ProxyConfig: if old_session_value != new_session_value: await self._reschedule_spend_log_cleanup_job() + ## SPEND LOG CLEANUP BOUNDS ## + # The dashboard writes these straight to the DB, so without copying them + # here the running cleanup job never sees them. A key the DB no longer + # carries was cleared from the dashboard, and falls back to whatever + # config.yaml declared, or to None (the shipped default) when it declared + # nothing. Leaving the deleted DB value in memory would keep enforcing the + # bound the operator just removed. + for cleanup_key in SPEND_LOG_CLEANUP_BOUND_SETTINGS: + general_settings[cleanup_key] = _general_settings.get( + cleanup_key, self._yaml_spend_log_cleanup_bounds.get(cleanup_key) + ) + for key in ( "user_url_allowed_hosts", "user_url_validation", @@ -9471,6 +9493,7 @@ class ProxyStartupEvent: "/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"] ) # if project requires model list async def model_list( + request: Request = None, # pyright: ignore[reportArgumentType] # FastAPI always injects the Request; the None default only serves direct in-process callers user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), return_wildcard_routes: bool | None = False, team_id: str | None = None, @@ -9507,6 +9530,9 @@ async def model_list( settings: Final = cast(dict[str, object], general_settings) # any-ok: legacy settings + from litellm.llms.anthropic.common_utils import ( + create_anthropic_model_list_response, + ) from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_privileges, ) @@ -9514,6 +9540,12 @@ async def model_list( create_model_info_response, get_available_models_for_user, ) + from litellm.types.proxy.model_listing import ModelInfoResponse + + http_request: Final = cast(Request | None, request) # cast-ok: in-process callers pass no request + wants_anthropic_format: Final = ( + http_request is not None and http_request.headers.get("anthropic-version") is not None + ) # Validate scope parameter if provided if scope is not None and scope != "expand": @@ -9597,6 +9629,10 @@ async def model_list( model_info["id"] = response_id model_data.append(model_info) + if wants_anthropic_format: + admin_listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above + return create_anthropic_model_list_response(admin_listing) + return dict( data=model_data, object="list", @@ -9637,6 +9673,10 @@ async def model_list( model_info["id"] = response_id model_data.append(model_info) + if wants_anthropic_format: + listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above + return create_anthropic_model_list_response(listing) + return dict( data=model_data, object="list", @@ -15529,6 +15569,10 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "store_model_in_db": "Boolean", "store_prompts_in_spend_logs": "Boolean", "maximum_spend_logs_retention_period": "String", + "maximum_spend_logs_cleanup_batch_size": "Integer", + "maximum_spend_logs_cleanup_max_batches": "Integer", + "maximum_spend_logs_cleanup_run_budget": "String", + "maximum_spend_logs_cleanup_batch_timeout": "String", "mcp_internal_ip_ranges": "List", "mcp_trusted_proxy_ranges": "List", "mcp_xff_num_trusted_hops": "Integer", diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index fcc6aac1c14..e24e5b21583 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2062,6 +2062,44 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "NVIDIA_RIVA", + "provider_display_name": "Nvidia Riva", + "litellm_provider": "nvidia_riva", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "grpc.nvcf.nvidia.com:443", + "tooltip": "host:port of the Riva gRPC endpoint. Use grpc.nvcf.nvidia.com:443 for NVCF-hosted Riva, or your own host (e.g. localhost:50051) when self-hosting. Riva has no public default, so this is required.", + "required": true, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": "nvapi-...", + "tooltip": "Sent as gRPC authorization metadata. Required for NVCF-hosted Riva, optional for self-hosted deployments without auth.", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "nvcf_function_id", + "label": "NVCF Function ID", + "placeholder": "1598d209-5e27-4d3c-8079-4751568b1081", + "tooltip": "NVCF function id of the hosted Riva model. Setting it turns on TLS and the function-id gRPC metadata. Leave empty for self-hosted Riva.", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "nvidia_riva/nvidia/parakeet-ctc-1_1b-asr" + }, { "provider": "Ollama", "provider_display_name": "Ollama", diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 029648f7901..efdbda47fdc 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger from litellm.constants import ( + PTU_LAPSED_ALERT_LIMIT, PTU_PRUNE_SKEW_GRACE_SECONDS, PTU_ROLLUP_JOB_ID, PTU_ROLLUP_LOCK_TTL_SECONDS, @@ -45,6 +46,7 @@ class RollupResult: models_processed: int rows_written: int rows_failed: int = 0 + lapsed: tuple[str, ...] = () @dataclass(frozen=True, slots=True) @@ -387,6 +389,34 @@ async def run_ptu_flat_cost_rollup( models_processed=len(ptu_models), rows_written=rows_written, rows_failed=rows_failed, + lapsed=_lapsed_models(ptu_models, run_started), + ) + + +def _slack_safe(model_name: str) -> str: + """``model_name`` with the characters Slack reads as markup escaped. + + A model name is operator-supplied and this alert is delivered to an operator channel, so an + unescaped name could post a channel-wide mention or a disguised link. + """ + return model_name.replace("&", "&").replace("<", "<").replace(">", ">") + + +def _lapsed_models(ptu_models: tuple[PTUModel, ...], now: datetime) -> tuple[str, ...]: + """PTU deployments whose window has closed, newest bound first. + + The provider bills reserved capacity until the deployment is deleted, so a closed window + stops this attribution without stopping the charge. The deployment is left alone: the + window is what the operator asked to be attributed, and per-token pricing would invent a + charge the provider does not make for reserved capacity. + """ + return tuple( + _slack_safe(model.model_name) + for model in sorted( + (m for m in ptu_models if m.effective_to is not None and m.effective_to <= now), + key=lambda m: m.effective_to, + reverse=True, + ) ) @@ -585,6 +615,14 @@ async def _run_and_alert( f"{result.rows_written + result.rows_failed} team charges failed to write. Those teams show no PTU " f"cost for that date until the rollup is rerun for it.", ) + if result.lapsed: + await _deliver_alert( + alert, + f"PTU flat-cost attribution has stopped for {len(result.lapsed)} deployment(s) whose effective " + f"window has closed: {', '.join(result.lapsed[:PTU_LAPSED_ALERT_LIMIT])}. Reserved capacity is billed " + "until the deployment is deleted, so a deployment still serving traffic is still being charged for " + "by the provider with nothing attributing it here. Extend the window, or retire the deployment.", + ) if target_date is None: await _backfill_and_alert(prisma_client, alert=alert) return result diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ef0376ade84..ec1aa262736 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -16,6 +16,7 @@ from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, TypeVar, Union, cast, overload from litellm import _custom_logger_compatible_callbacks_literal @@ -679,6 +680,7 @@ class ProxyLogging: # (e.g. MCPJWTSigner) to independently verify the caller's identity # before re-signing an outbound token (FR-5 verify+re-sign). "incoming_bearer_token": kwargs.get("incoming_bearer_token"), + "metadata": {"headers": kwargs.get("headers") or {}}, } return synthetic_data @@ -3005,6 +3007,62 @@ async def prefetch_config_params(prisma_client: "PrismaClient | None", param_nam ) +class _ForcedRecreateDeclined(Exception): + """A forced recreate was declined by the engine-generation guard. + + Distinct from a reconnect *failure*: the machinery worked, it just found + that another path had already replaced the writer, so it left the engines + alone. The caller's engine may still be poisoned, so the cycle must not + report success, but it must not count as a failure either, or the record + of what could not be repaired would gate the retry that recovers. + """ + + +@dataclass(frozen=True, slots=True) +class _StaleReadEngine: + """The read engine a query observed, identified rather than only counted. + + `PrismaClient.read_db` resolves to the reader while it is available and to + the writer once it is not, and the two carry independent generation + counters that both start at zero and advance on the same reconnect + cadence. A bare generation compared across that switch would silently pit + one engine's counter against another's, so the wrapper is carried with the + number and a switch counts as the engine having moved. + + Holding the wrapper itself rather than its `id()` is load-bearing, not + incidental: the strong reference keeps the wrapper alive, so its address + cannot be recycled under a stored observation and match an unrelated + engine later. It is only free because writer and reader both live as long + as the client does; a replaceable reader would make this a retention leak. + """ + + wrapper: PrismaWrapper + generation: int + + @classmethod + def observe(cls, wrapper: PrismaWrapper) -> "_StaleReadEngine": + return cls(wrapper=wrapper, generation=wrapper.engine_generation) + + def is_still_live(self, current: PrismaWrapper) -> bool: + """Whether this exact engine is still serving reads, unreplaced. + + A True answer must never be the only thing standing between a poisoned + engine and its repair. The generation moves only after a replacement + connects, and a recreate whose connect raises leaves it unmoved until + some later recreate succeeds, so this can report an engine as live + after it has stopped working. What bounds that is the failed-repair + record in `_cooldown_applies`, written by a repair attempt that fails + rather than by whatever broke the engine: the two need not be the same + recreate, since the synchronous token-refresh fallback in + `PrismaWrapper.__getattr__` recreates outside the reconnect machinery + and records nothing. The record is written only for callers that named + an engine, and it collapses the rest of the burst for up to one + cooldown window rather than guaranteeing a repair, since the cooldown + conjunct underneath it still expires and lets a later caller retry. + """ + return self.wrapper is current and self.generation == current.engine_generation + + class PrismaClient: spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() @@ -3152,6 +3210,14 @@ class PrismaClient: float(os.getenv("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", "0.1")), ) self._consecutive_reconnect_failures: int = 0 + # Last generation of each read engine whose repair was attempted and + # failed. Scoped to the engine rather than counted globally so an + # unrelated reconnect failure cannot suppress a stale reader's + # recovery, and keyed per wrapper rather than held in one slot so a + # writer failure cannot evict the reader's record and hand the waiver + # back to a caller whose engine is still unrepaired. Bounded at two + # entries: a client has one writer and at most one reader. + self._failed_recreate_generations: Mapping[PrismaWrapper, int] = MappingProxyType({}) self._reconnect_escalation_threshold: int = max(1, int(os.getenv("PRISMA_RECONNECT_ESCALATION_THRESHOLD", "3"))) self._engine_pidfd: int = -1 self._engine_pid: int = 0 @@ -3167,6 +3233,19 @@ class PrismaClient: return self.db.writer return self.db + @property + def read_db(self) -> PrismaWrapper: + """Underlying wrapper that top-level reads are dispatched to. + + Identical to `writer_db` without a read replica. With one configured + it is the reader, which is the engine `query_first` actually runs on, + so anything reasoning about the state of the connection that served a + read has to consult this rather than the writer. + """ + if isinstance(self.db, RoutingPrismaWrapper): + return self.db.read_target + return self.db + def tx(self) -> "TransactionManager": """Open an interactive transaction on the writer. @@ -3390,18 +3469,30 @@ class PrismaClient: `attempt_db_reconnect`, which is singleflight: when a schema change poisons every pooled connection at once, the first cached-plan error recreates the client and the concurrent waiters reuse that single - recreate instead of racing to kill each other's fresh engine. We then - retry the identical query exactly once. + recreate instead of racing to kill each other's fresh engine. We pass + `force_recreate` so the reconnect skips its `SELECT 1` liveness probe: + the connection is healthy here, it is the prepared statements on it + that are stale, so a passing probe would otherwise skip the recreate + and leave the retry to hit the same error. We then retry the identical + query exactly once. The retry reuses the original query byte-for-byte. Mutating the SQL (e.g. injecting a unique comment) would defeat PostgreSQL's plan cache, forcing a fresh plan on every request and pegging the database CPU. - If the reconnect is skipped because a recent reconnect is still within - its cooldown, the retry runs against the same connection and may fail - again; the get_data backoff decorator re-runs the lookup and a later - attempt reconnects once the cooldown elapses. + The reconnect cooldown must not gate the engine this query itself saw + as stale, or a migration landing within the cooldown of an earlier + reconnect leaves auth failing until it elapses. The engine observed + before the query names it, so the reconnect bypasses the cooldown only + while that same engine is still the live one. + + It is observed from `read_db`, not `writer_db`: `query_first` is a + top-level read, so with a read replica configured it runs on the reader + and it is the reader's prepared statements that went stale. Naming the + writer here would let an unrelated writer reconnect re-arm the cooldown + while the reader stayed poisoned. """ + stale_read_engine: Final = _StaleReadEngine.observe(self.read_db) try: return await self.db.query_first(sql_query, *args) except Exception as e: @@ -3413,7 +3504,11 @@ class PrismaClient: "query. This may occur during rolling deployments when schema " "changes are applied." ) - await self.attempt_db_reconnect(reason="postgres_cached_plan_error") + await self.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=stale_read_engine, + ) return await self.db.query_first(sql_query, *args) @backoff.on_exception( @@ -4696,7 +4791,11 @@ class PrismaClient: self._cleanup_engine_watcher() asyncio.create_task(self._start_engine_watcher()) - async def _run_reconnect_cycle(self, timeout_seconds: float | None = None) -> None: + async def _run_reconnect_cycle( + self, + timeout_seconds: float | None = None, + force_recreate: bool = False, + ) -> None: """ Run a reconnect cycle with a single overall timeout budget. @@ -4707,6 +4806,11 @@ class PrismaClient: the client via the non-blocking kill-then-construct flow rather than calling disconnect(), which blocks the event loop on the synchronous subprocess.Popen.wait() inside prisma-client-py (see issue #26191). + + `force_recreate` skips the direct path's liveness probe, for callers + whose failure lives in the session state rather than the connection + (stale prepared statements after a schema change): a reachable writer + proves nothing about those, so the probe must not veto the recreate. """ effective_timeout: Final = ( timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds @@ -4746,8 +4850,29 @@ class PrismaClient: # direct path there is no SELECT 1 probe here, so the generation # guard is the only thing standing between a crash-reconnect and # a refresh that raced it. - await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation) + recreated: Final = await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation) await self._start_engine_watcher() + # Same contract as the direct path below: a forced caller asked + # for its engine to be replaced, so a decline is not a success. + # Reachable here because the escalation threshold flips + # `_engine_confirmed_dead`, which routes the next cycle, forced + # callers included, down this branch. + if force_recreate is True and recreated is False: + # Clear the dead-engine flag first, restoring the policy the + # non-forced path already has: a decline does not raise for + # it, so it falls through to the clear below. Only the + # forced branch would strand the flag, and stranding it + # routes the next cycle back down this probe-free branch, + # where the refreshed generation now matches and the + # recreate kills the healthy engine a refresh just spawned + # (#29176). This has to stay AFTER `_start_engine_watcher` + # above: clearing the flag while the watcher is still torn + # down would be worse than either alone. + self._engine_confirmed_dead = False + raise _ForcedRecreateDeclined( + "Forced Prisma recreate declined by the generation guard; " + "the engine that failed was not replaced" + ) await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) # Only clear the "dead engine" flag after the heavy reconnect @@ -4772,44 +4897,106 @@ class PrismaClient: # detect a refresh that landed since cycle entry and skip the # redundant restart. writer: Final = self.writer_db - try: - await writer.query_raw("SELECT 1") - verbose_proxy_logger.info( - "Writer healthy on probe; skipping recreate (engine " - "likely already replaced by a token refresh)." - ) - if isinstance(self.db, RoutingPrismaWrapper): - self.db.mark_writer_recovered() - await self._start_engine_watcher() - return - except Exception as probe_err: - verbose_proxy_logger.warning( - "Writer probe failed (%s); recreating Prisma client.", - probe_err, - ) + if force_recreate is False: + try: + await writer.query_raw("SELECT 1") + verbose_proxy_logger.info( + "Writer healthy on probe; skipping recreate (engine " + "likely already replaced by a token refresh)." + ) + if isinstance(self.db, RoutingPrismaWrapper): + self.db.mark_writer_recovered() + await self._start_engine_watcher() + return + except Exception as probe_err: + verbose_proxy_logger.warning( + "Writer probe failed (%s); recreating Prisma client.", + probe_err, + ) # Fresh Prisma client + new engine subprocess. The previous # "lightweight" path called `disconnect()` which blocks the # event loop on `subprocess.Popen.wait()`; since that call # ends up killing the engine anyway, we do it non-blockingly # via `_kill_engine_process` inside `recreate_prisma_client`. self._cleanup_engine_watcher() - await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation) + recreated: Final = await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation) await self._start_engine_watcher() # Smoke-test the writer specifically; query_raw on the routing # wrapper sends to the reader, which would not validate the - # newly-recreated writer engine. + # newly-recreated writer engine. The reader is left to the + # caller's own retried query, a stronger check than SELECT 1, + # and a reader that fails to come back sets `_reader_unavailable` + # so reads fall through to the writer just recreated here. await self.writer_db.query_raw("SELECT 1") + # A recreate can decline: the optimistic-lock guard no-ops when + # the writer generation moved since cycle entry, and the routing + # wrapper then leaves the reader untouched as well. Callers that + # merely suspect a transport blip are happy either way, but a + # forced caller asked for this engine to be replaced because its + # session state is poisoned, and it was not. Do not report that + # as a success: it would reset the consecutive-failure count and + # log a repair that never happened. + if force_recreate is True and recreated is False: + raise _ForcedRecreateDeclined( + "Forced Prisma recreate declined by the generation guard; " + "the engine that failed was not replaced" + ) await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) + def _cooldown_applies(self, stale_read_engine: "_StaleReadEngine | None") -> bool: + """ + Whether the reconnect cooldown should still gate this caller. + + The cooldown collapses a burst of callers onto one recreate, so it + keeps gating a caller whose named engine has already been replaced: + that recreate is the one it was waiting for. While that engine is still + the live one the damage is still being served, so deferring to an + unrelated reconnect's cooldown would leave it broken until the cooldown + elapses. + + A named engine always describes the one that served the failing read + (see `_query_first_with_cached_plan_fallback`), so it is compared + against `read_db`, identity included: `read_db` can resolve to a + different wrapper than it did at observation time. + + The waiver is withdrawn once a repair of this same engine has been + tried and failed. A failed recreate leaves the generation where it was, + so without this every queued caller would still see its own engine live + and run its own full recreate serially instead of collapsing onto one + attempt, which is what the cooldown is for. The record is scoped to the + engine rather than to a global failure count: an unrelated reconnect + failing somewhere else says nothing about whether this engine can be + repaired, and gating on it would suppress the recovery this method + exists to allow. + + The record is never cleared, and does not need to be. Generations are + monotonic per wrapper, so once the engine is repaired every later + caller names a higher one and the entry can never match again. And this + method is only ever the first half of the gate: the cooldown window + itself still expires, so an engine that can never be repaired degrades + to the plain cooldown rather than being suppressed forever. + """ + if stale_read_engine is None: + return True + if self._failed_recreate_generations.get(stale_read_engine.wrapper) == stale_read_engine.generation: + return True + return not stale_read_engine.is_still_live(self.read_db) + async def _attempt_reconnect_inside_lock( self, force: bool, reason: str, timeout_seconds: float | None, + force_recreate: bool = False, + stale_read_engine: "_StaleReadEngine | None" = None, ) -> bool: now: Final = time.time() - if force is False and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds: + if ( + force is False + and self._cooldown_applies(stale_read_engine) + and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds + ): verbose_proxy_logger.debug( "Skipping DB reconnect attempt inside lock due to cooldown. reason=%s", reason, @@ -4833,12 +5020,43 @@ class PrismaClient: reconnect_succeeded = False try: - await self._run_reconnect_cycle(timeout_seconds=timeout_seconds) + await self._run_reconnect_cycle(timeout_seconds=timeout_seconds, force_recreate=force_recreate) reconnect_succeeded = True self._consecutive_reconnect_failures = 0 verbose_proxy_logger.info("Prisma DB reconnect succeeded. reason=%s", reason) + except _ForcedRecreateDeclined as declined: + # A decline is raised only when the recreate returns False, which + # happens only at the generation guard, and the generation moves + # only after a replacement has connected. So a decline is proof + # that a replacement SUCCEEDED, and zeroing a consecutive-failure + # count on that proof is right by definition rather than by + # analogy to what a reported success used to do. Note what it + # proves is that the WRITER was replaced, not that this caller's + # engine was repaired: on a read replica the reader can still be + # poisoned, since the wrapper returns before touching it. Leaving + # the count at the threshold would let the escalation check above + # re-arm the dead-engine flag on the very next attempt and send a + # healthy replacement back down the probe-free heavy path. + self._consecutive_reconnect_failures = 0 + verbose_proxy_logger.warning("Prisma DB reconnect declined. reason=%s detail=%s", reason, declined) except Exception as reconnect_err: self._consecutive_reconnect_failures += 1 + # Remember WHICH engine could not be repaired, so the rest of this + # caller's burst collapses onto the cooldown instead of each + # retrying the recreate that just failed. Recorded only for a + # caller that named a generation: a watchdog or transport-error + # reconnect failing here is unrelated to any stale read engine and + # must not suppress its waiver. + if stale_read_engine is not None: + # Key off the wrapper the CALLER named, never a freshly resolved + # `read_db`. A failed reader recreate is itself what marks the + # reader unavailable, so re-resolving here would file the + # reader's failure under the writer: the poisoned reader would + # lose its record and the healthy writer would gain a spurious + # one, wrong in both directions at once. + self._failed_recreate_generations = MappingProxyType( + {**self._failed_recreate_generations, stale_read_engine.wrapper: stale_read_engine.generation} + ) verbose_proxy_logger.error( "Prisma DB reconnect failed (%d consecutive). reason=%s error=%s", self._consecutive_reconnect_failures, @@ -4856,15 +5074,35 @@ class PrismaClient: force: bool = False, timeout_seconds: float | None = None, lock_timeout_seconds: float | None = None, + force_recreate: bool = False, + stale_read_engine: "_StaleReadEngine | None" = None, ) -> bool: """ Attempt to reconnect the Prisma client in a singleflight manner. + `force` bypasses the cooldown unconditionally; `force_recreate` + bypasses the liveness probe that would otherwise skip recreating a + reachable engine; `stale_read_engine` bypasses the cooldown only while + the engine that produced the caller's failure is still the live one + (see `_cooldown_applies`). + + A `force_recreate` caller can also get False for a third reason: the + generation guard declined because another path had already replaced + the engine, which is a successful outcome reported as False. Callers + that branch on the return value (`exception_handler` raises on False, + `auth_checks` retries only on True) would misread that as a dead end, + and are safe today only because neither passes `force_recreate`. Do + not add it to one of them without revisiting how it reads the result. + Returns: bool: True if reconnection succeeded, else False. """ now: Final = time.time() - if force is False and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds: + if ( + force is False + and self._cooldown_applies(stale_read_engine) + and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds + ): verbose_proxy_logger.debug( "Skipping DB reconnect attempt due to cooldown. reason=%s", reason, @@ -4873,7 +5111,9 @@ class PrismaClient: if lock_timeout_seconds is None: async with self._db_reconnect_lock: - return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) + return await self._attempt_reconnect_inside_lock( + force, reason, timeout_seconds, force_recreate, stale_read_engine + ) lock_acquired_by_timeout_task = False @@ -4922,7 +5162,9 @@ class PrismaClient: return False try: - return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) + return await self._attempt_reconnect_inside_lock( + force, reason, timeout_seconds, force_recreate, stale_read_engine + ) finally: self._db_reconnect_lock.release() diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index c6e17502e5d..56818717c09 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.utils import ( + logging_safe_mcp_headers, split_server_prefix_from_name, strip_known_server_prefix, ) @@ -653,6 +654,7 @@ class LiteLLM_Proxy_MCP_Handler: tool_results: Final[list[MCPToolResult]] = [] tool_call_id: str | None = None rules_obj: Final = Rules() + logging_safe_headers: Final = logging_safe_mcp_headers(raw_headers) for tool_call in tool_calls: logging_request_data: dict[str, object] = {} tool_name: str | None = None @@ -697,6 +699,7 @@ class LiteLLM_Proxy_MCP_Handler: "tool_call_id": tool_call_id, "tool_name": sanitized_tool_name, "server_name": server_name, + "headers": logging_safe_headers, } logging_request_data = { "model": f"MCP: {tool_name}", @@ -708,7 +711,7 @@ class LiteLLM_Proxy_MCP_Handler: "proxy_server_request": { "url": "/mcp/tools/call", "method": "POST", - "headers": {}, + "headers": logging_safe_headers, "body": { "name": sanitized_tool_name, "arguments": parsed_arguments, diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index afe55603466..82498ec10cd 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -55,11 +55,13 @@ else merge_base=$(git merge-base origin/litellm_internal_staging HEAD 2>/dev/null) || { echo "check: cannot resolve the merge base with origin/litellm_internal_staging." >&2 echo " Fix: git fetch origin litellm_internal_staging" >&2 + echo "check: FAIL" exit 1 } scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMRD "$merge_base")" "$untracked" | sed '/^$/d' | sort -u) if [ -z "$scope" ]; then echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)" + echo "check: PASS" exit 0 fi echo "check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging:" @@ -281,4 +283,30 @@ if [ -n "${gen_pid:-}" ]; then cat "$gen_log"; rm -f "$gen_log" fi +summary_item() { + local check_name=$1 triggered=$2 skip_reason=$3 + if [ -n "$triggered" ]; then + echo " ran: $check_name" + else + echo " skipped: $check_name ($skip_reason)" + fi +} + +echo "check: summary" +summary_item "Python lint (make lint)" "$litellm_py_files" "no litellm/ Python files in scope" +summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_files" "no tests/e2e Python files in scope" +summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope" +summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope" + +if [ -z "$litellm_py_files$e2e_py_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then + echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2 + printf '%s\n' "$scope" | sed 's/^/ /' >&2 + echo " A pass here is a no-op, not a lint verdict." >&2 +fi + +if [ "$status" -eq 0 ]; then + echo "check: PASS" +else + echo "check: FAIL" +fi exit $status diff --git a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts index 07a75dc007d..b8424b06115 100644 --- a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts @@ -19,12 +19,11 @@ test.describe("Internal User", () => { // Open the team dropdown — seeded internal user is a member of // e2e-team-crud and e2e-team-org, so we expect at least the CRUD alias. - const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" }); + const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); - await expect(page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ALIAS).first()).toBeVisible({ - timeout: 5_000, - }); + const dropdown = page.locator('[data-slot="combobox-content"]:visible'); + await expect(dropdown.getByText(E2E_TEAM_CRUD_ALIAS).first()).toBeVisible({ timeout: 5_000 }); }); test("Team info page omits the Settings tab for non-admin members", async ({ page }) => { diff --git a/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts index 1b048198456..c44305187f1 100644 --- a/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts @@ -27,18 +27,18 @@ test.describe("Internal User with no team memberships", () => { await page.getByRole("button", { name: /Create New Key/i }).click(); await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); - const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" }); + const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); - const dropdown = page.locator(".ant-select-dropdown:visible").first(); + const dropdown = page.locator('[data-slot="combobox-content"]:visible').first(); await expect(dropdown).toBeVisible({ timeout: 5_000 }); // Wait for the settled-empty state, not a transient one. The dropdown shows - // a spinner while teams load and only swaps in "No teams found" once the - // request resolves with nothing (team_dropdown.tsx renders the spinner when - // isLoading and this copy otherwise). Asserting on it means a regression - // where teams DO load for this user fails here instead of racing a one-shot - // count() against an in-flight request. + // "Loading teams…" while teams load and only swaps in "No teams found" once + // the request resolves with nothing (team_dropdown.tsx passes both copies to + // PaginatedSearchSelect). Asserting on it means a regression where teams DO + // load for this user fails here instead of racing a one-shot count() against + // an in-flight request. await expect(dropdown.getByText("No teams found")).toBeVisible({ timeout: 10_000 }); await expect(dropdown.getByRole("option")).toHaveCount(0); }); diff --git a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts index 7d5058a8140..68319154554 100644 --- a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts @@ -18,10 +18,10 @@ test.describe("Internal User with team memberships", () => { await page.getByRole("button", { name: /Create New Key/i }).click(); await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); - const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" }); + const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); - const dropdown = page.locator(".ant-select-dropdown:visible").first(); + const dropdown = page.locator('[data-slot="combobox-content"]:visible').first(); await expect(dropdown).toBeVisible({ timeout: 5_000 }); // Both seeded memberships render, and nothing else does — proving the diff --git a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts index 461d9dfd9f8..1b11ea69f97 100644 --- a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -328,11 +328,11 @@ test.describe("Add Model", () => { const teamByokRow = page.locator(".ant-form-item", { hasText: "Team-BYOK Model" }); await teamByokRow.getByRole("switch").click(); - // TeamDropdown's options carry custom markup and no role="option", so match by text. - const teamDropdown = page.getByTestId("team-dropdown"); + // TeamDropdown options show the alias above the team id, so match on the id line by text. + const teamDropdown = page.getByTestId("team-dropdown").getByRole("combobox"); await expect(teamDropdown).toBeVisible({ timeout: 5_000 }); await teamDropdown.click(); - const teamOption = page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ID).first(); + const teamOption = page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ID).first(); await expect(teamOption).toBeVisible({ timeout: 5_000 }); await teamOption.click(); diff --git a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts index 1ff3ef274b6..d9b0f959c9f 100644 --- a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts @@ -40,11 +40,11 @@ test.describe("Proxy Admin - Keys", () => { const keyName = `e2e-admin-key-${Date.now()}`; await page.getByTestId("base-input").fill(keyName); - // Select team — the team dropdown has placeholder "Search or select a team" - const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" }); + // Select team + const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); - await page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ALIAS).first().click(); + await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ALIAS).first().click(); // Select models await page.locator(".ant-select-selection-overflow").click(); @@ -157,7 +157,7 @@ test.describe("Proxy Admin - Keys", () => { await page.getByRole("button", { name: "More key actions" }).click(); await page.getByRole("menuitem", { name: "Delete Key" }).click(); - const modal = page.locator(".ant-modal:visible"); + const modal = page.getByRole("dialog", { name: "Delete Key" }); await expect(modal).toBeVisible({ timeout: 5_000 }); await modal.locator("input").fill(E2E_DELETE_KEY_ALIAS); diff --git a/tests/e2e/ui/tests/proxy-admin/teams.spec.ts b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts index 92d22f11f4d..d7c8eb6237e 100644 --- a/tests/e2e/ui/tests/proxy-admin/teams.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts @@ -47,10 +47,10 @@ test.describe("Proxy Admin - Teams", () => { // Fill Team Name — the input has id="team_alias" await dialog.locator("#team_alias").fill(uniqueAlias); - // Select models — the models multi-select is inside the modal - // Click to open dropdown, select "All Proxy Models" - await dialog.locator(".ant-select-selection-overflow").first().click(); - await page.locator(".ant-select-dropdown:visible").getByText("All Proxy Models").click(); + // Select models — the models multi-select is inside the modal. Its popup is + // portaled to the body, so scope the option lookup to the page, not the dialog. + await dialog.getByTestId("create-team-models-select").getByRole("combobox").click(); + await page.getByRole("option", { name: "All Proxy Models", exact: true }).click(); await page.keyboard.press("Escape"); // Submit — click the submit button inside the dialog (not the header button) @@ -129,7 +129,7 @@ test.describe("Proxy Admin - Teams", () => { await teamRow.locator('[data-testid^="team-actions-"]').click(); await page.getByTestId("team-action-delete").click(); - const modal = page.locator(".ant-modal:visible"); + const modal = page.getByRole("dialog", { name: "Delete Team?" }); await expect(modal).toBeVisible({ timeout: 5_000 }); await modal.locator("input").fill(E2E_TEAM_DELETE_ALIAS); await modal.getByRole("button", { name: /Force Delete|Delete/i }).click(); @@ -191,11 +191,11 @@ test.describe("Proxy Admin - Teams", () => { const modelsSelect = page.locator("[data-testid='models-select']"); await expect(modelsSelect).toBeVisible({ timeout: 10_000 }); - const anthropicTag = modelsSelect - .locator(".ant-select-selection-item") + const anthropicChip = modelsSelect + .locator('[data-slot="combobox-chip"]') .filter({ hasText: "fake-anthropic-claude" }); - await expect(anthropicTag).toBeVisible({ timeout: 5_000 }); - await anthropicTag.locator(".ant-select-selection-item-remove").click(); + await expect(anthropicChip).toBeVisible({ timeout: 5_000 }); + await anthropicChip.locator('[data-slot="combobox-chip-remove"]').click(); await page.getByRole("button", { name: "Save Changes" }).click(); diff --git a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts index be4526f7089..d71d5e6c0fe 100644 --- a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts +++ b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts @@ -105,7 +105,7 @@ test.describe("Team Admin", () => { await expect(row).toBeVisible({ timeout: 10_000 }); await row.getByTestId("delete-member").click(); - const modal = page.locator(".ant-modal:visible"); + const modal = page.getByRole("dialog", { name: "Delete Team Member" }); await expect(modal).toBeVisible({ timeout: 5_000 }); const remove = await captureRequestBody(page, { method: "POST", urlIncludes: "/team/member_delete" }, async () => { @@ -139,10 +139,10 @@ test.describe("Team Admin", () => { await page.getByTestId("base-input").fill(keyName); // Team selector — same locator pattern as the proxy-admin keys test. - const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" }); + const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); - await page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ALIAS).first().click(); + await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ALIAS).first().click(); // Models — pick "All Team Models" await page.locator(".ant-select-selection-overflow").click(); diff --git a/tests/proxy_admin_ui_tests/test_access_group_team_sync.py b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py new file mode 100644 index 00000000000..629d77f20fc --- /dev/null +++ b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py @@ -0,0 +1,230 @@ +""" +Real-Postgres coverage for the team -> access group mirror. + +`sync_team_access_group_membership` reconciles `assigned_team_ids` with two raw +statements, and a mocked prisma cannot tell whether that SQL is right: a fake has to +reimplement the array semantics in Python, so it passes no matter what the SQL says. +These tests run the statements against the same Postgres CI seeds for the admin UI +suite, which is the only place a `NOT (... = ANY(...))` guard going missing shows up. +""" + +import asyncio +import os +import sys +from contextlib import asynccontextmanager +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.proxy.management_helpers.access_group_team_sync import ( + reconcile_team_access_group_membership, + sync_team_access_group_membership, +) + +TEAM = "ags-team-a" +OTHER_TEAM = "ags-team-b" +GROUPS = ("ags-group-1", "ags-group-2", "ags-group-3") +_DELETE_SEEDED = 'DELETE FROM "LiteLLM_AccessGroupTable" WHERE access_group_id = ANY($1::TEXT[])' +_DELETE_TEAMS = 'DELETE FROM "LiteLLM_TeamTable" WHERE team_id = ANY($1::TEXT[])' + + +@asynccontextmanager +async def _clean_db(): + """Connects inside the running test's loop. An async fixture would be torn up on a + different loop than the test body, which prisma's engine lock refuses outright.""" + from prisma import Prisma + + if not os.getenv("DATABASE_URL"): + pytest.fail("DATABASE_URL is required; these tests must not silently skip") + + db = Prisma() + await db.connect() + try: + await db.execute_raw(_DELETE_SEEDED, list(GROUPS)) + await db.execute_raw(_DELETE_TEAMS, [TEAM, OTHER_TEAM]) + yield db + finally: + await db.execute_raw(_DELETE_SEEDED, list(GROUPS)) + await db.execute_raw(_DELETE_TEAMS, [TEAM, OTHER_TEAM]) + await db.disconnect() + + +async def _seed(db, assignments): + for group_id, team_ids in assignments.items(): + await db.litellm_accessgrouptable.create( + data={ + "access_group_id": group_id, + "access_group_name": group_id, + "assigned_team_ids": team_ids, + } + ) + + +async def _read(db): + rows = await db.query_raw( + 'SELECT access_group_id, assigned_team_ids FROM "LiteLLM_AccessGroupTable" ' + "WHERE access_group_id = ANY($1::TEXT[])", + list(GROUPS), + ) + return {row["access_group_id"]: sorted(row["assigned_team_ids"] or []) for row in rows} + + +async def _set_team_groups(db, team_id, access_group_ids): + """The mirror reads the committed team row, so the desired state is written there.""" + if access_group_ids is None: + await db.execute_raw(_DELETE_TEAMS, [team_id]) + return + await db.litellm_teamtable.upsert( + where={"team_id": team_id}, + data={ + "create": {"team_id": team_id, "access_group_ids": list(access_group_ids)}, + "update": {"access_group_ids": list(access_group_ids)}, + }, + ) + + +async def _sync(db, team_id, access_group_ids): + await _set_team_groups(db, team_id, access_group_ids) + with patch( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + new_callable=AsyncMock, + ) as invalidate: + await sync_team_access_group_membership(prisma_client=SimpleNamespace(db=db), team_id=team_id) + return {call.args[0] for call in invalidate.call_args_list} + + +@pytest.mark.asyncio +async def test_reconcile_attaches_and_detaches_without_touching_other_teams(): + """The detach must be scoped to groups the team dropped. Losing that scope would + strip the team from the very groups it just kept, silently revoking live grants.""" + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [TEAM, OTHER_TEAM], GROUPS[1]: [TEAM], GROUPS[2]: [OTHER_TEAM]}) + + invalidated = await _sync(db, TEAM, [GROUPS[1], GROUPS[2]]) + + assert await _read(db) == { + GROUPS[0]: [OTHER_TEAM], + GROUPS[1]: [TEAM], + GROUPS[2]: sorted([TEAM, OTHER_TEAM]), + } + assert invalidated == {GROUPS[0], GROUPS[1], GROUPS[2]} + + +@pytest.mark.asyncio +async def test_reconcile_is_idempotent_so_a_retry_heals_rather_than_duplicates(): + """Reconciling to the same desired state twice must leave the rows alone and still name + the team's groups for the cache step, so a retry after a failed cache drop reaches them. + A delta-based mirror would instead go quiet once the rows match, leaving the caches + serving a grant the admin already revoked.""" + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [], GROUPS[1]: [TEAM], GROUPS[2]: []}) + + first = await _sync(db, TEAM, [GROUPS[0], GROUPS[1]]) + after_first = await _read(db) + second = await _sync(db, TEAM, [GROUPS[0], GROUPS[1]]) + + assert after_first == {GROUPS[0]: [TEAM], GROUPS[1]: [TEAM], GROUPS[2]: []} + assert await _read(db) == after_first + assert first == {GROUPS[0], GROUPS[1]} + assert second == first + + +@pytest.mark.asyncio +async def test_reconcile_handles_a_null_array_column(): + """`assigned_team_ids` is nullable in Postgres. Without COALESCE both statements + evaluate their guard to NULL, skip the row, and the grant silently never syncs.""" + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [], GROUPS[1]: []}) + await db.execute_raw( + 'UPDATE "LiteLLM_AccessGroupTable" SET assigned_team_ids = NULL WHERE access_group_id = $1', + GROUPS[0], + ) + + await _sync(db, TEAM, [GROUPS[0]]) + + assert await _read(db) == {GROUPS[0]: [TEAM], GROUPS[1]: []} + + +@pytest.mark.asyncio +async def test_passing_none_detaches_the_team_from_every_group(): + """Team deletion. A group the deleted row never listed must still let the team go, + otherwise the id dangles under Attached Teams and grants again if it is reused.""" + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [TEAM, OTHER_TEAM], GROUPS[1]: [TEAM], GROUPS[2]: [OTHER_TEAM]}) + + invalidated = await _sync(db, TEAM, None) + + assert await _read(db) == {GROUPS[0]: [OTHER_TEAM], GROUPS[1]: [], GROUPS[2]: [OTHER_TEAM]} + assert invalidated == {GROUPS[0], GROUPS[1]} + + +@pytest.mark.asyncio +async def test_a_failed_mirror_takes_the_new_team_row_with_it(): + """`/team/new` inserts the team and mirrors it in one transaction. Mirroring in a + transaction of its own instead leaves a committed team whose groups never learned about + it, and the retry with that same team id comes back as a duplicate.""" + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]}) + + with pytest.raises(RuntimeError): + async with db.tx() as tx: + await tx.litellm_teamtable.create(data={"team_id": TEAM, "access_group_ids": [GROUPS[0]]}) + await reconcile_team_access_group_membership(tx, TEAM) + raise RuntimeError("the cache handoff blew up") + + assert await _read(db) == {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]} + assert await db.litellm_teamtable.find_unique(where={"team_id": TEAM}) is None + + +@pytest.mark.asyncio +async def test_a_concurrent_writer_cannot_replay_a_stale_team_row_over_a_newer_one(): + """ + Two writers edit one team at once. Whichever team row commits last is the admin's + final intent and the mirror must match it, so the mirror has to hold the team's + advisory lock across its read and its writes. + + A second connection holds that lock and changes the team underneath, which pins the + interleaving instead of hoping a sleep lands in the gap. With the lock the sync waits + and then reads the new row. Without it the sync reads the old row and writes a group + the admin already moved off, which keeps granting to that team. + """ + from prisma import Prisma + + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [], GROUPS[1]: []}) + await _sync(db, TEAM, [GROUPS[0]]) + assert await _read(db) == {GROUPS[0]: [TEAM], GROUPS[1]: []} + + blocker = Prisma() + await blocker.connect() + sync_started = asyncio.Event() + + async def competing_sync(): + sync_started.set() + with patch( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + new_callable=AsyncMock, + ): + await sync_team_access_group_membership(prisma_client=SimpleNamespace(db=db), team_id=TEAM) + + try: + async with blocker.tx(timeout=timedelta(seconds=30)) as held: + await held.query_raw("SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked", TEAM) + task = asyncio.create_task(competing_sync()) + await sync_started.wait() + await asyncio.sleep(0.2) + assert not task.done(), "the mirror did not wait on the team's advisory lock" + await held.execute_raw( + 'UPDATE "LiteLLM_TeamTable" SET access_group_ids = $1 WHERE team_id = $2', + [GROUPS[1]], + TEAM, + ) + await asyncio.wait_for(task, timeout=30) + finally: + await blocker.disconnect() + + assert await _read(db) == {GROUPS[0]: [], GROUPS[1]: [TEAM]} diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index c83a3fa2b73..de04a65c310 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -314,7 +314,7 @@ class TestLangfuseUsageDetails(unittest.TestCase): "litellm_params": {"metadata": {}}, "optional_params": {}, "litellm_call_id": "test-call-id-null-usage", - "standard_logging_object": None, + "standard_logging_object": self._build_standard_logging_payload(), "response_cost": 0.0, } @@ -382,16 +382,14 @@ class TestLangfuseUsageDetails(unittest.TestCase): "model_id": "model-123", "model_group": "openai", "api_base": "https://api.openai.com", + # only real StandardLoggingMetadata fields: session_id, trace_name, + # headers and friends are request-metadata keys the allowlist drops, + # so a payload carrying them cannot occur in production "metadata": { "user_api_key_end_user_id": None, "prompt_management_metadata": None, - "session_id": None, - "trace_name": None, - "trace_version": None, - "headers": None, - "endpoint": None, - "caching_groups": None, - "previous_models": None, + "user_api_key_hash": "hashed-key", + "user_api_key_alias": "canary-alias", }, "hidden_params": {}, "request_tags": [], @@ -503,14 +501,251 @@ class TestLangfuseUsageDetails(unittest.TestCase): # litellm_trace_id should be preferred over litellm_call_id assert self.last_trace_kwargs.get("id") == "trace-id-from-kwargs" - def test_log_langfuse_v2_uses_litellm_trace_id_when_standard_logging_object_none( - self, - ): + CANARY = "sk-lf-canary-SECRET-d4e5f6" + + def _canary_request_metadata(self): + """Raw request metadata shaped like the proxy builds it, credentials included.""" + from litellm.proxy._types import UserAPIKeyAuth + + team_logging = [ + { + "callback_name": "langfuse", + "callback_vars": {"langfuse_secret_key": self.CANARY}, + } + ] + return { + "user_api_key_auth": UserAPIKeyAuth( + api_key="hashed-key", + team_metadata={"logging": team_logging}, + ), + "user_api_key_team_metadata": {"logging": team_logging}, + "user_api_key_metadata": {"secret_manager_settings": {"vault_token": self.CANARY}}, + "session_id": "canary-session", + "trace_name": "canary-trace", + "first_custom": "keep-first", + "second_custom": "keep-second", + "endpoint": "/v1/chat/completions", + "headers": {"authorization": f"Bearer {self.CANARY}"}, + } + + def _emitted_payload_text(self): + """Every blob this logger handed to the langfuse SDK, as one searchable string.""" + import json + + blobs = [self.last_trace_kwargs] + if self.mock_langfuse_trace.generation.call_args is not None: + blobs.append(self.mock_langfuse_trace.generation.call_args.kwargs) + blobs.extend(call.kwargs for call in self.mock_langfuse_trace.span.call_args_list) + return json.dumps(blobs, default=repr) + + def _drive_with_canary(self, extra_metadata=None, hidden_params=None): + metadata = {**self._canary_request_metadata(), **(extra_metadata or {})} + payload = self._build_standard_logging_payload(trace_id="canary-trace-id") + if hidden_params is not None: + payload["hidden_params"] = hidden_params + kwargs = {**self._build_langfuse_kwargs(payload), "response_cost": 0.25} + self.last_trace_kwargs = {} + self.mock_langfuse_trace.generation.reset_mock() + self.mock_langfuse_trace.span.reset_mock() + + with patch( + "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", + side_effect=lambda generation_params, **kw: generation_params, + create=True, + ): + self.logger._log_langfuse_v2( + user_id="user-1", + metadata=metadata, + litellm_params={"metadata": metadata}, + output=None, + start_time=datetime.datetime(2024, 1, 1, 12, 0, 0), + end_time=datetime.datetime(2024, 1, 1, 12, 0, 1), + kwargs=kwargs, + optional_params={}, + input=None, + response_obj=None, + level="INFO", + litellm_call_id="canary-call-id", + ) + return self.mock_langfuse_trace.generation.call_args.kwargs["metadata"] + + def test_team_callback_credentials_never_reach_langfuse(self): """ - When standard_logging_object is None (failure case where - get_standard_logging_object_payload threw), litellm_trace_id from kwargs - should be used as the Langfuse trace_id. This matches the DB Session ID. + Regression for the credential leak: request metadata carries the whole + UserAPIKeyAuth object, whose team_metadata holds the customer's own langfuse + keys. The emitted blob is sourced from StandardLoggingPayload, so none of the + three credential carriers can ride along. """ + generation_metadata = self._drive_with_canary() + + assert self.CANARY not in self._emitted_payload_text() + for leaked_key in ( + "user_api_key_auth", + "user_api_key_team_metadata", + "user_api_key_metadata", + ): + assert leaked_key not in generation_metadata + + def test_debug_langfuse_dump_carries_no_credentials(self): + """ + debug_langfuse dumps request metadata into the trace as a second emit site. + It must be sourced from the allowlisted payload too. + """ + self._drive_with_canary(extra_metadata={"debug_langfuse": True}) + + dumped = self.last_trace_kwargs["metadata"]["metadata_passed_to_litellm"] + assert "user_api_key_auth" not in dumped + assert self.CANARY not in self._emitted_payload_text() + + def test_raw_request_metadata_reaches_the_emitted_blob_through_no_key(self): + """ + The emitted blob is the allowlist plus litellm enrichments, nothing else. + Nothing from raw request metadata is copied across, whatever its type, which + is what makes the credential exclusion structural rather than a filter that + has to be kept correct. Proxy callers keep their own metadata under the + allowlisted requester_metadata key. + """ + generation_metadata = self._drive_with_canary() + + for caller_key in ("first_custom", "second_custom", "session_id", "trace_name"): + assert caller_key not in generation_metadata + + def test_provider_specific_span_receives_the_emitted_blob(self): + """ + The provider span reads hidden_params, which is an enrichment on the emitted + blob rather than a key of request metadata. Handing it the steering dict + instead would silently stop emitting vertex grounding spans. + """ + self._drive_with_canary(hidden_params={"vertex_ai_grounding_metadata": ["ground-a", "ground-b"]}) + + span_inputs = [call.kwargs.get("input") for call in self.mock_langfuse_trace.span.call_args_list] + assert span_inputs == ["ground-a", "ground-b"] + assert self.CANARY not in self._emitted_payload_text() + + def test_caller_cannot_spoof_an_allowlisted_identity_field(self): + """ + Request metadata never reaches the blob, so a caller naming user_api_key_alias + cannot have their value emitted in place of the proxy-resolved one. + """ + generation_metadata = self._drive_with_canary( + extra_metadata={"user_api_key_alias": "spoofed-by-caller"} + ) + + assert generation_metadata["user_api_key_alias"] == "canary-alias" + + def test_caller_nested_metadata_cannot_erase_a_litellm_enrichment(self): + """ + log_requester_metadata drops any top-level key whose name also appears inside + requester_metadata. Sourcing the blob from the allowlist populates that nested + dict for real, so a caller naming a key litellm_response_cost would otherwise + blank out the cost litellm computed. Enrichments are layered after the dedupe. + """ + payload = self._build_standard_logging_payload(trace_id="canary-trace-id") + payload["metadata"]["requester_metadata"] = {"litellm_response_cost": "caller-value", "api_base": "caller"} + kwargs = {**self._build_langfuse_kwargs(payload), "response_cost": 0.25} + metadata = self._canary_request_metadata() + self.mock_langfuse_trace.generation.reset_mock() + + with patch( + "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", + side_effect=lambda generation_params, **kw: generation_params, + create=True, + ): + self.logger._log_langfuse_v2( + user_id="user-1", + metadata=metadata, + litellm_params={"metadata": metadata, "api_base": "https://real-api-base"}, + output=None, + start_time=datetime.datetime(2024, 1, 1, 12, 0, 0), + end_time=datetime.datetime(2024, 1, 1, 12, 0, 1), + kwargs=kwargs, + optional_params={}, + input=None, + response_obj=None, + level="INFO", + litellm_call_id="canary-call-id", + ) + + generation_metadata = self.mock_langfuse_trace.generation.call_args.kwargs["metadata"] + assert generation_metadata["litellm_response_cost"] == 0.25 + assert generation_metadata["api_base"] == "https://real-api-base" + + def test_denied_steering_keys_and_enrichments(self): + """ + endpoint is a plain string, so without the deny-list it would ride the + string re-injection straight into the emitted blob. The enrichments are + litellm-computed and must survive the move off clean_metadata. + """ + generation_metadata = self._drive_with_canary() + + assert "endpoint" not in generation_metadata + assert "headers" not in generation_metadata + assert generation_metadata["litellm_response_cost"] == 0.25 + assert "hidden_params" in generation_metadata + + def test_cache_hit_is_normalized_on_the_shared_kwargs(self): + """ + kwargs here is the shared model_call_details dict. Callbacks that run after + langfuse read cache_hit off it and copy it into their own payloads, so + dropping the None to False normalization records None for datadog, logfire, + generic_api and spend tracking. + """ + metadata = self._canary_request_metadata() + payload = self._build_standard_logging_payload(trace_id="canary-trace-id") + kwargs = {**self._build_langfuse_kwargs(payload), "cache_hit": None} + + with patch( + "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", + side_effect=lambda generation_params, **kw: generation_params, + create=True, + ): + self.logger._log_langfuse_v2( + user_id="user-1", + metadata=metadata, + litellm_params={"metadata": metadata}, + output=None, + start_time=datetime.datetime(2024, 1, 1, 12, 0, 0), + end_time=datetime.datetime(2024, 1, 1, 12, 0, 1), + kwargs=kwargs, + optional_params={}, + input=None, + response_obj=None, + level="INFO", + litellm_call_id="canary-call-id", + ) + + assert kwargs["cache_hit"] is False + + def test_redact_user_api_key_info_still_strips_the_emitted_blob(self): + """ + The flag used to act on the raw-derived blob. That blob is now sourced from + StandardLoggingPayload, which is where the user_api_key_* fields live, so the + redaction has to run on the assembled payload or the flag silently stops working. + """ + with patch.object(litellm, "redact_user_api_key_info", True): + generation_metadata = self._drive_with_canary() + + assert not [key for key in generation_metadata if key.startswith("user_api_key")] + + def test_steering_keys_still_read_from_raw_metadata(self): + """ + Only the emitted payload moves to StandardLoggingPayload. The control fields + keep reading raw metadata, which is what Braintrust's migration got wrong. + """ + self._drive_with_canary() + + assert self.last_trace_kwargs.get("session_id") == "canary-session" + assert self.last_trace_kwargs.get("name") == "canary-trace" + + def test_failure_trace_survives_a_missing_standard_logging_object(self): + """ + get_standard_logging_object_payload is fail-open and returns None on any + exception, which is exactly the failed-request case Langfuse most needs to + show. The trace is still emitted with the litellm_trace_id fallback, and the + blob degrades to caller strings plus enrichments rather than falling back to + raw metadata, which would ship the UserAPIKeyAuth object. + """ + metadata = self._canary_request_metadata() kwargs = { "standard_logging_object": None, "model": "gpt-4", @@ -520,16 +755,17 @@ class TestLangfuseUsageDetails(unittest.TestCase): "litellm_trace_id": "trace-id-failure", } self.last_trace_kwargs = {} + self.mock_langfuse_trace.generation.reset_mock() with patch( "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", side_effect=lambda generation_params, **kwargs: generation_params, create=True, ): - self.logger._log_langfuse_v2( + trace_id, _ = self.logger._log_langfuse_v2( user_id="user-1", - metadata={}, - litellm_params={"metadata": {}}, + metadata=metadata, + litellm_params={"metadata": metadata}, output=None, start_time=datetime.datetime.utcnow(), end_time=datetime.datetime.utcnow(), @@ -541,8 +777,18 @@ class TestLangfuseUsageDetails(unittest.TestCase): litellm_call_id="call-id-different", ) - # Must use litellm_trace_id, not litellm_call_id + import json + + assert trace_id == "trace-id-failure" assert self.last_trace_kwargs.get("id") == "trace-id-failure" + generation_metadata = self.mock_langfuse_trace.generation.call_args.kwargs["metadata"] + assert "user_api_key_auth" not in generation_metadata + assert self.CANARY not in self._emitted_payload_text() + assert "first_custom" not in generation_metadata + # hidden_params comes off the payload, so it is omitted rather than emitted + # as an unserializable placeholder + assert "hidden_params" not in generation_metadata + json.dumps(generation_metadata) def test_log_langfuse_v2_session_id_passed_as_trace_session_id(self): """ diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 1fcee1b1c42..d5676aaf288 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -133,6 +133,40 @@ class TestExceptionCheckers: result = ExceptionCheckers.is_error_str_rate_limit(error_str) assert result is True + def test_bare_429_in_body_is_ignored_when_status_code_says_otherwise(self): + """A 429 echoed back inside a 400's body is not a rate limit. + + Word boundaries don't help: 429 is an ordinary token id (" that" in several + tokenisers), so an echoed prompt_token_ids array reads as a standalone 429. + """ + error_str = ( + '{"error":{"message":"`tools` must not be an empty array",' + '"type":"invalid_request_error"},' + '"prompt_token_ids":[9906,429,1234]}' + ) + assert ExceptionCheckers.is_error_str_rate_limit(error_str, status_code=400) is False + + def test_bare_429_still_detected_without_a_status_code(self): + """With no status available, a standalone 429 still counts (unchanged behaviour).""" + + assert ExceptionCheckers.is_error_str_rate_limit("HTTP 429 Too Many Requests") is True + assert ExceptionCheckers.is_error_str_rate_limit("HTTP 429 Too Many Requests", status_code=None) is True + assert ExceptionCheckers.is_error_str_rate_limit("HTTP 429 Too Many Requests", status_code=429) is True + + def test_non_integer_status_code_does_not_suppress_bare_429(self): + """A non-integer status counts as unknown, not as a contradiction.""" + + assert ExceptionCheckers.is_error_str_rate_limit("HTTP 429 Too Many Requests", status_code="not-an-int") is True + + def test_rate_limit_phrase_is_honoured_under_a_non_429_status(self): + """Phrase matching stays ungated: some providers report a real rate limit in + the text under a non-429 status (#11455).""" + + assert ( + ExceptionCheckers.is_error_str_rate_limit("FireworksException - rate limit exceeded", status_code=400) + is True + ) + def test_is_azure_content_policy_violation_error_with_policy_violation_text(self): """Test detection of Azure content policy violation with explicit policy violation text""" @@ -300,6 +334,54 @@ def test_lemonade_context_window_error_mapping(): assert excinfo.value.model == model +def test_openai_compatible_400_with_bare_429_in_body_maps_to_bad_request(): + """A provider 400 whose echoed body contains a 429 must stay a 400. + + ``is_error_str_rate_limit`` runs before the status-code branch for + openai-compatible providers, so a validation error echoing the request back came + out as RateLimitError, which tells the caller to retry a request that cannot + succeed and books the failure against provider throttling. + """ + error_message = ( + '{"error":{"message":"`tools` must not be an empty array",' + '"type":"invalid_request_error","code":400},' + '"prompt_token_ids":[9906,429,1234]}' + ) + original_exception = OpenAIError( + status_code=400, + message=error_message, + headers={}, + ) + + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="deepseek-ai/DeepSeek-V3", + original_exception=original_exception, + custom_llm_provider="deepinfra", + ) + + assert excinfo.value.status_code == 400 + assert excinfo.value.llm_provider == "deepinfra" + + +def test_openai_compatible_429_still_maps_to_rate_limit(): + """A real 429 still maps to RateLimitError.""" + original_exception = OpenAIError( + status_code=429, + message='{"error":{"message":"Too Many Requests","type":"rate_limit_error"}}', + headers={}, + ) + + with pytest.raises(litellm.RateLimitError) as excinfo: + exception_type( + model="deepseek-ai/DeepSeek-V3", + original_exception=original_exception, + custom_llm_provider="deepinfra", + ) + + assert excinfo.value.status_code == 429 + + @pytest.mark.parametrize( "error_message", [ diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 9df72108332..431030bcf2e 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -2028,3 +2028,82 @@ class TestCapabilityProbeUsesCallerProvider: AnthropicModelInfo._is_adaptive_thinking_model("claude-opus-4-8", "anthropic") is True ) +def test_create_anthropic_model_list_response_shape(): + from litellm.llms.anthropic.common_utils import ( + create_anthropic_model_list_response, + ) + + response = create_anthropic_model_list_response( + [ + {"id": "claude-opus-4-6", "object": "model", "created": 0, "owned_by": "openai"}, + {"id": "gpt-4o", "object": "model", "created": 0, "owned_by": "openai"}, + {"id": "claude-haiku-4-5", "object": "model", "created": 0, "owned_by": "openai"}, + ] + ) + + assert "object" not in response + assert response["has_more"] is False + assert response["first_id"] == "claude-opus-4-6" + assert response["last_id"] == "claude-haiku-4-5" + assert [m["id"] for m in response["data"]] == [ + "claude-opus-4-6", + "gpt-4o", + "claude-haiku-4-5", + ] + for entry in response["data"]: + assert entry["type"] == "model" + assert entry["display_name"] == entry["id"] + # ISO 8601 with a Z suffix, as the Anthropic Models API returns. + assert entry["created_at"].endswith("Z") + assert "+00:00" not in entry["created_at"] + assert "max_input_tokens" not in entry + assert "max_tokens" not in entry + + +def test_create_anthropic_model_list_response_carries_token_limits(): + from litellm.llms.anthropic.common_utils import ( + create_anthropic_model_list_response, + ) + + response = create_anthropic_model_list_response( + [ + { + "id": "claude-opus-4-6", + "object": "model", + "created": 0, + "owned_by": "openai", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + }, + { + "id": "input-only", + "object": "model", + "created": 0, + "owned_by": "openai", + "max_input_tokens": 8192, + }, + {"id": "unknown-limits", "object": "model", "created": 0, "owned_by": "openai"}, + ] + ) + + opus, input_only, unknown = response["data"] + assert opus["max_input_tokens"] == 200000 + assert opus["max_tokens"] == 64000 + assert "max_output_tokens" not in opus + assert input_only["max_input_tokens"] == 8192 + assert "max_tokens" not in input_only + assert "max_input_tokens" not in unknown + assert "max_tokens" not in unknown + + +def test_create_anthropic_model_list_response_empty(): + from litellm.llms.anthropic.common_utils import ( + create_anthropic_model_list_response, + ) + + response = create_anthropic_model_list_response([]) + + assert response["data"] == [] + assert response["has_more"] is False + assert response["first_id"] is None + assert response["last_id"] is None \ No newline at end of file diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 02bd3535c0a..fd66667af64 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -16,8 +16,8 @@ sys.path.insert(0, os.path.abspath("../../../../../..")) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.bedrock.common_utils import ( ensure_bedrock_anthropic_messages_tool_names, + normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, - remove_custom_field_from_tools, ) from litellm.constants import ( BEDROCK_MIN_THINKING_BUDGET_TOKENS, @@ -353,12 +353,13 @@ def test_remove_ttl_from_cache_control(): assert request5 == {} -def test_remove_custom_field_from_tools(): +def test_normalize_custom_field_on_tools(): """ - Ensure the `custom` field is stripped from every tool definition. + Ensure the `custom` field is stripped from every tool definition, and that a + boolean `custom.defer_loading` is hoisted onto the top-level `defer_loading` + flag Bedrock documents instead of being dropped with the wrapper. - Claude Code v2.1.69+ sends `custom: {defer_loading: true}` on tool - objects. Bedrock does not accept this extra field and returns + Bedrock does not accept a `custom` object on a tool and returns "Extra inputs are not permitted". Ref: https://github.com/BerriAI/litellm/issues/22847 @@ -381,29 +382,94 @@ def test_remove_custom_field_from_tools(): ] } - remove_custom_field_from_tools(request) + normalize_custom_field_on_tools(request) for tool in request["tools"]: assert "custom" not in tool, f"Tool {tool['name']} still has 'custom' field" # Other fields should be preserved assert request["tools"][0]["name"] == "Read" assert request["tools"][1]["name"] == "Write" + # `custom.defer_loading` is hoisted; the tool that never carried it is untouched + assert request["tools"][0]["defer_loading"] is True + assert "defer_loading" not in request["tools"][1] # Case 2: request without tools key (should not raise error) request2 = {"messages": [{"role": "user", "content": "hi"}]} - remove_custom_field_from_tools(request2) + normalize_custom_field_on_tools(request2) assert "tools" not in request2 # Case 3: empty tools list (should not raise error) request3 = {"tools": []} - remove_custom_field_from_tools(request3) + normalize_custom_field_on_tools(request3) assert request3["tools"] == [] # Case 4: tools with None value (should not raise error) request4 = {"tools": None} - remove_custom_field_from_tools(request4) + normalize_custom_field_on_tools(request4) assert request4["tools"] is None + # Case 5: an explicit top-level flag wins over a conflicting wrapped one + request5 = { + "tools": [ + {"name": "Read", "defer_loading": False, "custom": {"defer_loading": True}} + ] + } + normalize_custom_field_on_tools(request5) + assert request5["tools"][0] == {"name": "Read", "defer_loading": False} + + # Case 6: a non-boolean `custom.defer_loading` is dropped, never forwarded + for junk in ("true", 1, None, {"nested": True}): + request6 = {"tools": [{"name": "Read", "custom": {"defer_loading": junk}}]} + normalize_custom_field_on_tools(request6) + assert request6["tools"][0] == {"name": "Read"}, f"leaked defer_loading={junk!r}" + + # Case 7: a `custom` that is not a dict is dropped without raising + request7 = { + "tools": [ + {"name": "Read", "custom": "defer_loading"}, + {"name": "Write", "custom": None}, + ] + } + normalize_custom_field_on_tools(request7) + assert request7["tools"] == [{"name": "Read"}, {"name": "Write"}] + + +@pytest.mark.parametrize( + "deferred_marker", [{"custom": {"defer_loading": True}}, {"defer_loading": True}] +) +def test_bedrock_invoke_messages_transform_emits_top_level_defer_loading( + deferred_marker, +): + """A deferred tool must reach Bedrock as top-level ``defer_loading``, whether the + client wrapped the flag in ``custom`` or sent it top-level, and the outbound body + must still carry the Bedrock tool-search beta.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_optional_request_params={ + "max_tokens": 128, + "stream": False, + "betas": ["advanced-tool-use-2025-11-20"], + "tools": [ + { + "name": "Read", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {}}, + **deferred_marker, + }, + {"type": "tool_search_tool_regex_20251119", "name": "tool_search"}, + ], + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert result["tools"][0]["defer_loading"] is True + assert "custom" not in result["tools"][0] + assert result["anthropic_beta"] == ["tool-search-tool-2025-10-19"] + def test_normalize_tool_input_schema_types_for_bedrock_invoke(): """ diff --git a/tests/test_litellm/llms/openai/test_cost_calculation.py b/tests/test_litellm/llms/openai/test_cost_calculation.py new file mode 100644 index 00000000000..9b6aec1966c --- /dev/null +++ b/tests/test_litellm/llms/openai/test_cost_calculation.py @@ -0,0 +1,83 @@ +"""Tests for per-second transcription cost calculation.""" + +import pytest + +import litellm +from litellm.llms.openai.cost_calculation import cost_per_second + + +def _register_stt(name: str, **pricing: float) -> None: + litellm.register_model( + { + name: { + "mode": "audio_transcription", + "litellm_provider": "openai", + **pricing, + } + }, + persist_across_reloads=False, + ) + + +def test_input_rate_bills_when_output_rate_is_zero(): + """A declared-but-zero output rate must not suppress the real input rate.""" + _register_stt( + "test-stt-zero-output", + input_cost_per_second=5e-05, + output_cost_per_second=0.0, + ) + + prompt_cost, completion_cost = cost_per_second( + model="test-stt-zero-output", custom_llm_provider="openai", duration=300.0 + ) + + assert prompt_cost == pytest.approx(0.015) + assert completion_cost == 0.0 + + +def test_output_rate_takes_precedence_when_both_are_billable(): + """Entries duplicating one rate into both fields must not be billed twice.""" + _register_stt( + "test-stt-both-rates", + input_cost_per_second=1e-04, + output_cost_per_second=1e-04, + ) + + prompt_cost, completion_cost = cost_per_second( + model="test-stt-both-rates", custom_llm_provider="openai", duration=10.0 + ) + + assert prompt_cost + completion_cost == pytest.approx(1e-03) + + +def test_output_rate_alone_still_bills(): + _register_stt("test-stt-output-only", output_cost_per_second=3e-05) + + prompt_cost, completion_cost = cost_per_second( + model="test-stt-output-only", custom_llm_provider="openai", duration=60.0 + ) + + assert prompt_cost == 0.0 + assert completion_cost == pytest.approx(1.8e-03) + + +@pytest.mark.parametrize( + "model, provider", + [ + ("deepgram/nova-3", "deepgram"), + ("groq/whisper-large-v3", "groq"), + ("elevenlabs/scribe_v1", "elevenlabs"), + ("assemblyai/best", "assemblyai"), + ("whisper-1", "openai"), + ], +) +def test_shipped_per_second_models_bill_a_non_zero_cost(model, provider): + prompt_cost, completion_cost = cost_per_second(model=model, custom_llm_provider=provider, duration=60.0) + + assert prompt_cost + completion_cost > 0.0 + + +def test_whisper_bills_its_documented_rate_once(): + prompt_cost, completion_cost = cost_per_second(model="whisper-1", custom_llm_provider="openai", duration=30.0) + + assert prompt_cost + completion_cost == pytest.approx(0.003) 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 0b95a497882..52dc91ce24d 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 @@ -2874,7 +2874,7 @@ class TestMCPCustomHeaderName: mock_general_settings.get.return_value = general_setting # Call the method - result = MCPRequestHandler._get_mcp_client_side_auth_header_name() + result = MCPRequestHandler.get_mcp_client_side_auth_header_name() # Assert the result assert result == expected_header_name @@ -2938,7 +2938,7 @@ class TestMCPCustomHeaderName: # Mock the header name method with patch.object( MCPRequestHandler, - "_get_mcp_client_side_auth_header_name", + "get_mcp_client_side_auth_header_name", return_value=custom_header_name, ): # Create headers from the test data @@ -2963,7 +2963,7 @@ class TestMCPCustomHeaderName: # Mock the custom header name with patch.object( MCPRequestHandler, - "_get_mcp_client_side_auth_header_name", + "get_mcp_client_side_auth_header_name", return_value="custom-auth-header", ): # Create ASGI scope with custom header diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index b56a12db5b1..4081681daef 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1196,3 +1196,45 @@ class TestOpenApiResolvedUpstreamAuth: ) assert resolved is None lookup.assert_not_awaited() + + +class TestPreCallToolCheckExposesClientHeaders: + """The pre_mcp_call guardrail payload must carry the caller's sanitized HTTP headers.""" + + @pytest.mark.asyncio + async def test_sanitized_client_headers_reach_the_guardrail_payload(self): + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="test_server", + server_name="test_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + captured: Dict[str, Any] = {} + + def capture(request_obj, kwargs): + captured.update(kwargs) + return {"model": "fake"} + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(return_value=MagicMock()) + proxy_logging._convert_mcp_to_llm_format = MagicMock(side_effect=capture) + proxy_logging.pre_call_hook = AsyncMock(return_value=None) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object(manager, "check_tool_permission_for_key_team", new_callable=AsyncMock): + with patch.object(manager, "validate_allowed_params"): + await manager.pre_call_tool_check( + name="test_tool", + arguments={"key": "val"}, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + raw_headers={"x-nuid": "nuid-1", "x-litellm-api-key": "sk-proxy"}, + ) + + assert captured["headers"] == {"x-nuid": "nuid-1"} 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 850d01c6e34..7df83065865 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 @@ -77,7 +77,7 @@ async def test_mcp_server_tool_call_body_contains_request_data(): # Mock the add_litellm_data_to_request function to capture the data captured_data = {} - async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config): + async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs): captured_data.update(data) # Simulate the proxy_server_request creation captured_data["proxy_server_request"] = { @@ -116,6 +116,107 @@ async def test_mcp_server_tool_call_body_contains_request_data(): assert body["arguments"] == tool_arguments +@pytest.mark.asyncio +async def test_mcp_server_tool_call_forwards_client_headers_to_logging(): + """The MCP protocol path must hand the connection's client headers to the pre-call + pipeline, so logging callbacks and guardrails see them the way the REST path does.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + mcp_server_tool_call, + set_auth_context, + ) + except ImportError: + pytest.skip("MCP server not available") + + set_auth_context( + UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + raw_headers={ + "x-nuid": "nuid-1", + "x-app-id": "app-1", + "content-length": "42", + "x-forwarded-for": "9.9.9.9", + }, + client_ip="1.2.3.4", + ) + + captured_headers = {} + + async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs): + captured_headers.update(request.headers) + return data + + async def mock_call_mcp_tool(*args, **kwargs): + return [{"type": "text", "text": "mocked response"}] + + with patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ): + with patch( + "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + mock_call_mcp_tool, + ): + with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): + await mcp_server_tool_call("test_tool", {"param": "value"}) + + assert captured_headers.get("x-nuid") == "nuid-1" + assert captured_headers.get("x-app-id") == "app-1" + assert "content-length" not in captured_headers + assert captured_headers.get("x-forwarded-for") == "1.2.3.4" + + +@pytest.mark.asyncio +async def test_mcp_server_tool_call_strips_custom_litellm_key_header(): + """The deployment can rename the proxy key header via general_settings.litellm_key_header_name. + The pre-call pipeline only knows that name if it is passed in, so without it the virtual key + reaches metadata.headers and proxy_server_request.headers in plaintext.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + mcp_server_tool_call, + set_auth_context, + ) + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + except ImportError: + pytest.skip("MCP server not available") + + set_auth_context( + UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + raw_headers={"x-company-key": "sk-proxy-secret", "x-nuid": "nuid-1"}, + client_ip="1.2.3.4", + ) + + captured_data = {} + + async def capturing_add_litellm_data_to_request(**kwargs): + data = await add_litellm_data_to_request(**kwargs) + captured_data.update(data) + return data + + async def mock_call_mcp_tool(*args, **kwargs): + return [{"type": "text", "text": "mocked response"}] + + with patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request", + capturing_add_litellm_data_to_request, + ): + with patch( + "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + mock_call_mcp_tool, + ): + with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): + with patch.dict( + "litellm.proxy.proxy_server.general_settings", + {"litellm_key_header_name": "x-company-key"}, + clear=False, + ): + await mcp_server_tool_call("test_tool", {"param": "value"}) + + metadata_headers = captured_data["metadata"]["headers"] + assert metadata_headers.get("x-nuid") == "nuid-1" + assert "x-company-key" not in metadata_headers + assert "x-company-key" not in captured_data["proxy_server_request"]["headers"] + + @pytest.mark.asyncio async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(): """The MCP session manager serializes handler exceptions as JSON-RPC errors, so a mid-session @@ -133,7 +234,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(): set_auth_context(UserAPIKeyAuth(api_key="test_key", user_id="test_user")) - async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config): + async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs): return data async def mock_call_mcp_tool(*args, **kwargs): @@ -1245,7 +1346,7 @@ async def test_mcp_server_tool_call_body_with_none_arguments(): # Mock the add_litellm_data_to_request function to capture the data captured_data = {} - async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config): + async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs): captured_data.update(data) captured_data["proxy_server_request"] = { "url": str(request.url), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py index 73fdee9cde3..00ed4e91efa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py @@ -1,7 +1,11 @@ +from unittest.mock import patch + import pytest from fastapi import HTTPException from litellm.proxy._experimental.mcp_server.utils import ( + build_synthetic_mcp_request, + logging_safe_mcp_headers, validate_and_normalize_mcp_server_payload, validate_tool_display_names, ) @@ -47,3 +51,99 @@ class TestValidateAndNormalizeMcpServerPayload: tool_name_to_display_name={"read_wiki_structure": "browse_repo_docs"}, ) validate_and_normalize_mcp_server_payload(payload) + + +class TestLoggingSafeMcpHeaders: + def test_returns_empty_for_missing_headers(self): + assert logging_safe_mcp_headers(None) == {} + assert logging_safe_mcp_headers({}) == {} + + def test_exposes_custom_headers_and_masks_credentials(self): + safe = logging_safe_mcp_headers( + { + "x-nuid": "nuid-1", + "x-app-id": "app-1", + "x-litellm-api-key": "sk-proxy", + "cookie": "session=secret", + } + ) + assert safe == { + "x-nuid": "nuid-1", + "x-app-id": "app-1", + "cookie": "***REDACTED***", + } + + def test_strips_custom_litellm_key_header(self): + """general_settings.litellm_key_header_name carries the proxy virtual key, so it must + never reach a callback or a guardrail even though clean_headers cannot know its name.""" + with patch.dict( + "litellm.proxy.proxy_server.general_settings", + {"litellm_key_header_name": "x-company-key"}, + clear=False, + ): + safe = logging_safe_mcp_headers({"x-company-key": "sk-proxy", "x-nuid": "nuid-1"}) + + assert safe == {"x-nuid": "nuid-1"} + + def test_strips_client_controlled_redaction_opt_out(self): + """litellm-disable-message-redaction is read back out of the logged metadata to turn off + redaction, so leaving it in place lets any MCP client undo what an admin configured.""" + safe = logging_safe_mcp_headers({"litellm-disable-message-redaction": "true", "x-nuid": "nuid-1"}) + + assert safe == {"x-nuid": "nuid-1"} + + def test_strips_upstream_mcp_credentials(self): + safe = logging_safe_mcp_headers( + { + "x-mcp-auth": "Bearer upstream", + "x-mcp-github-authorization": "Bearer gh_token", + "x-mcp-zapier-x-api-key": "zapier-key", + "x-nuid": "nuid-1", + } + ) + + assert safe == {"x-nuid": "nuid-1"} + + def test_strips_custom_mcp_client_side_auth_header(self): + with patch.dict( + "litellm.proxy.proxy_server.general_settings", + {"mcp_client_side_auth_header_name": "x-upstream-token"}, + clear=False, + ): + safe = logging_safe_mcp_headers({"x-upstream-token": "Bearer upstream", "x-nuid": "nuid-1"}) + + assert safe == {"x-nuid": "nuid-1"} + + +class TestBuildSyntheticMcpRequest: + def test_forwards_client_headers_without_upstream_credentials(self): + """The synthetic request feeds add_litellm_data_to_request, which derives + metadata.headers, so upstream MCP credentials must not ride along.""" + request = build_synthetic_mcp_request( + path="/mcp/tools/call", + raw_headers={ + "x-nuid": "nuid-1", + "x-mcp-auth": "Bearer upstream", + "x-mcp-github-authorization": "Bearer gh_token", + }, + ) + + assert request.headers.get("x-nuid") == "nuid-1" + assert "x-mcp-auth" not in request.headers + assert "x-mcp-github-authorization" not in request.headers + + def test_drops_custom_litellm_key_header(self): + """Callers such as the sampling flow build metadata off this request, so the + deployment's custom proxy key header must never be forwarded on it.""" + with patch.dict( + "litellm.proxy.proxy_server.general_settings", + {"litellm_key_header_name": "x-company-key"}, + clear=False, + ): + request = build_synthetic_mcp_request( + path="/mcp/sampling/createMessage", + raw_headers={"x-company-key": "sk-proxy-secret", "x-nuid": "nuid-1"}, + ) + + assert request.headers.get("x-nuid") == "nuid-1" + assert "x-company-key" not in request.headers diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index bfd4ffe1593..0c73c3fcf22 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -21,6 +21,10 @@ from litellm.proxy.common_utils.callback_utils import ( strip_callback_config, ) import litellm +from litellm.caching.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging from unittest.mock import patch from litellm.proxy.common_utils.callback_utils import process_callback @@ -491,3 +495,163 @@ def test_strip_callback_config_drops_credential_bearing_slots(): @pytest.mark.parametrize("value", [None, "not-a-dict", 42]) def test_strip_callback_config_passes_through_non_dicts(value): assert strip_callback_config(value) is value + + +# --------------------------------------------------------------------------- +# initialize_callbacks_on_proxy: dotted-path entries must resolve to something +# the request path can actually dispatch +# --------------------------------------------------------------------------- + +_PROBE_MODULE_NAME = "custom_callback_probe" + +_PROBE_MODULE_SOURCE = ''' +from litellm.integrations.custom_logger import CustomLogger + + +class FloorMaxTokens(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + data["max_tokens"] = 16 + return data + + +class NotALogger: + pass + + +def log_event_fn(kwargs, response_obj, start_time, end_time): + return None + + +NOT_A_CALLBACK = "some-plain-string" + +proxy_handler_instance = FloorMaxTokens() +''' + + +@pytest.fixture +def probe_config_path(tmp_path): + """Write a callback module next to a config.yaml, the layout get_instance_fn's file + branch expects, and restore every global the load + dispatch path touches. + + ``ProxyLogging._callback_capabilities_cache`` is keyed on the id()s of the + litellm.callbacks members, so an entry left behind here can be read back by an + unrelated test whose (len, ids) signature happens to collide. + """ + (tmp_path / f"{_PROBE_MODULE_NAME}.py").write_text(_PROBE_MODULE_SOURCE) + + original_callbacks = ( + list(litellm.callbacks) if isinstance(litellm.callbacks, list) else litellm.callbacks + ) + litellm.callbacks = [] + ProxyLogging._callback_capabilities_cache.clear() + try: + yield str(tmp_path / "config.yaml") + finally: + litellm.callbacks = original_callbacks + ProxyLogging._callback_capabilities_cache.clear() + + +def _load_callbacks(value, config_file_path): + initialize_callbacks_on_proxy( + value=value, + premium_user=False, + config_file_path=config_file_path, + litellm_settings={}, + callback_specific_params={}, + ) + + +def test_initialize_callbacks_on_proxy_rejects_class_valued_entry(probe_config_path): + """A class path loads an object that fails the `isinstance(_callback, CustomLogger)` + dispatch gate in ProxyLogging.pre_call_hook, so the proxy used to boot clean and + silently never run the hook. Config load must fail instead.""" + entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" + + with pytest.raises(ValueError) as exc_info: + _load_callbacks([entry], probe_config_path) + + message = str(exc_info.value) + assert entry in message + assert "the class" in message + assert "FloorMaxTokens" in message + assert f"{_PROBE_MODULE_NAME}.proxy_handler_instance" in message + assert litellm.callbacks == [] + + +@pytest.mark.parametrize( + "attribute, expected_fragment", + [ + ("NotALogger", "the class"), + ("NOT_A_CALLBACK", "str 'some-plain-string'"), + ], +) +def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values( + probe_config_path, attribute, expected_fragment +): + entry = f"{_PROBE_MODULE_NAME}.{attribute}" + + with pytest.raises(ValueError) as exc_info: + _load_callbacks([entry], probe_config_path) + + message = str(exc_info.value) + assert entry in message + assert expected_fragment in message + assert litellm.callbacks == [] + + +def test_initialize_callbacks_on_proxy_rejects_class_valued_non_list_value(probe_config_path): + entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" + + with pytest.raises(ValueError) as exc_info: + _load_callbacks(entry, probe_config_path) + + assert entry in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_initialize_callbacks_on_proxy_instance_entry_runs_pre_call_hook(probe_config_path): + """Positive control: the supported shape must still load AND still run. Drives the + real ProxyLogging.pre_call_hook, which is where a class-valued entry goes silent.""" + _load_callbacks([f"{_PROBE_MODULE_NAME}.proxy_handler_instance"], probe_config_path) + + assert len(litellm.callbacks) == 1 + assert isinstance(litellm.callbacks[0], CustomLogger) + + ProxyLogging._callback_capabilities_cache.clear() + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + data = await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-probe"), + data={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 1, + "metadata": {}, + }, + call_type="acompletion", + ) + + assert data["max_tokens"] == 16 + + +def test_initialize_callbacks_on_proxy_keeps_known_string_callback(probe_config_path): + """Non-narrowing control: a known callback name never reaches get_instance_fn and + stays a plain string in litellm.callbacks.""" + _load_callbacks(["langfuse"], probe_config_path) + + assert litellm.callbacks == ["langfuse"] + + +def test_initialize_callbacks_on_proxy_accepts_plain_function_callback(probe_config_path): + """Non-narrowing control: litellm.callbacks is typed + `Callable | | CustomLogger`, so a dotted path resolving to a plain + function is a supported shape and must keep loading.""" + _load_callbacks([f"{_PROBE_MODULE_NAME}.log_event_fn"], probe_config_path) + + assert [getattr(cb, "__name__", None) for cb in litellm.callbacks] == ["log_event_fn"] + + +def test_initialize_callbacks_on_proxy_accepts_instance_non_list_value(probe_config_path): + _load_callbacks(f"{_PROBE_MODULE_NAME}.proxy_handler_instance", probe_config_path) + + assert len(litellm.callbacks) == 1 + assert isinstance(litellm.callbacks[0], CustomLogger) diff --git a/tests/test_litellm/proxy/conftest.py b/tests/test_litellm/proxy/conftest.py index 230ccaf5fd4..61752997f0f 100644 --- a/tests/test_litellm/proxy/conftest.py +++ b/tests/test_litellm/proxy/conftest.py @@ -43,32 +43,51 @@ def disconnected_prisma() -> DisconnectedPrisma: return DisconnectedPrisma() -@pytest.fixture(autouse=True) -def _isolate_proxy_module_globals(): - """ - Snapshot and restore module-level globals on litellm.proxy.proxy_server - that tests sometimes mutate via raw setattr (not monkeypatch). +_MODULE_GLOBAL_MISSING = object() +_proxy_module_globals_snapshot = pytest.StashKey[Dict[str, object]]() - Without this, a leaked value — e.g. master_key set by a sibling test — + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_setup(item): + """ + Snapshot module-level globals on litellm.proxy.proxy_server before any + fixture runs, and restore them in pytest_runtest_teardown after every + fixture finalizer has run. + + Without this, a leaked value (e.g. master_key set by a sibling test) flips the auth short-circuit in user_api_key_auth and causes unrelated tests in the same xdist worker to return 401 instead of 200. + + This must be a hook pair, not an autouse fixture: an autouse fixture in + the root conftest requests monkeypatch, so monkeypatch's undo stack + unwinds after every other fixture finalizer. A test that monkeypatches a + global while a fixture has it patched records the fixture's mock as the + "original", and monkeypatch.undo re-plants that mock after all restores + have run, poisoning the global for the rest of the xdist worker. """ from litellm.proxy import proxy_server - sentinel = object() - snapshot = { - name: getattr(proxy_server, name, sentinel) + item.stash[_proxy_module_globals_snapshot] = { + name: getattr(proxy_server, name, _MODULE_GLOBAL_MISSING) for name in _PROXY_MODULE_GLOBALS_TO_ISOLATE } - try: - yield - finally: - for name, value in snapshot.items(): - if value is sentinel: - if hasattr(proxy_server, name): - delattr(proxy_server, name) - else: - setattr(proxy_server, name, value) + yield + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_teardown(item, nextitem): + yield + snapshot = item.stash.get(_proxy_module_globals_snapshot, None) + if snapshot is None: + return + from litellm.proxy import proxy_server + + for name, value in snapshot.items(): + if value is _MODULE_GLOBAL_MISSING: + if hasattr(proxy_server, name): + delattr(proxy_server, name) + else: + setattr(proxy_server, name, value) @pytest.fixture(autouse=True) diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py index 289de707387..e949afce57b 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py @@ -3,6 +3,7 @@ Tests for SpendLogsPartitionManager: partition naming/bounds math, retention selection, the non-partitioned no-op safety path, and the drop/ensure SQL flow. """ +from contextlib import asynccontextmanager from datetime import date, datetime, timezone from unittest.mock import AsyncMock, MagicMock @@ -19,6 +20,46 @@ from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import ( ) +DDL_TIMEOUT_MS = 30000 + + +def _budget(ms: "int | None" = DDL_TIMEOUT_MS): + """The injected per-statement bound: a callable re-read before each statement.""" + return lambda: ms + + +def _wire_tx(db) -> list[str]: + """ + Model the prisma seam the partition DDL uses. + + Every statement this manager issues, DDL and catalog query alike, runs inside + db.tx() so it can carry SET LOCAL timeouts. Those SET LOCAL statements are + collected in the returned list rather than forwarded, so assertions on + db.execute_raw and db.query_raw still see only the real statements. + """ + session_settings: list[str] = [] + + @asynccontextmanager + async def _tx(): + tx = MagicMock() + + async def _execute_raw(sql, *args): + if sql.lstrip().upper().startswith("SET LOCAL"): + session_settings.append(sql.strip()) + return 0 + return await db.execute_raw(sql, *args) + + async def _query_raw(sql, *args): + return await db.query_raw(sql, *args) + + tx.execute_raw = _execute_raw + tx.query_raw = _query_raw + yield tx + + db.tx = _tx + return session_settings + + def test_period_start_per_interval(): d = date(2026, 6, 3) # a Wednesday assert period_start(d, "day") == date(2026, 6, 3) @@ -78,11 +119,13 @@ async def test_is_partitioned_true_and_false(): client_true = MagicMock() client_true.db.query_raw = AsyncMock(return_value=[{"partitioned": True}]) - assert await mgr.is_partitioned(client_true) is True + _wire_tx(client_true.db) + assert await mgr.is_partitioned(client_true, _budget()) is True client_false = MagicMock() client_false.db.query_raw = AsyncMock(return_value=[{"partitioned": False}]) - assert await mgr.is_partitioned(client_false) is False + _wire_tx(client_false.db) + assert await mgr.is_partitioned(client_false, _budget()) is False @pytest.mark.asyncio @@ -94,13 +137,14 @@ async def test_catalog_queries_are_scoped_to_current_schema(): mgr = SpendLogsPartitionManager() client = MagicMock() client.db.query_raw = AsyncMock(return_value=[]) + _wire_tx(client.db) - await mgr.is_partitioned(client) + await mgr.is_partitioned(client, _budget()) is_partitioned_sql = client.db.query_raw.call_args.args[0] assert "pg_namespace" in is_partitioned_sql assert "current_schema()" in is_partitioned_sql - await mgr._list_partitions(client) + await mgr._list_partitions(client, DDL_TIMEOUT_MS) list_sql = client.db.query_raw.call_args.args[0] assert "pg_namespace" in list_sql assert "current_schema()" in list_sql @@ -112,7 +156,10 @@ async def test_is_partitioned_swallows_errors_and_returns_false(): mgr = SpendLogsPartitionManager() client = MagicMock() client.db.query_raw = AsyncMock(side_effect=Exception("db down")) - assert await mgr.is_partitioned(client) is False + # Wire the real seam: without it the async with itself raises, and the test + # would pass on the wrong exception. + _wire_tx(client.db) + assert await mgr.is_partitioned(client, _budget()) is False @pytest.mark.asyncio @@ -133,9 +180,10 @@ async def test_drop_partitions_older_than_drops_expired_only(): ] ) client.db.execute_raw = AsyncMock(return_value=0) + _wire_tx(client.db) cutoff = datetime(2026, 6, 5, 0, 0, 0, tzinfo=timezone.utc) - dropped = await mgr.drop_partitions_older_than(client, cutoff) + dropped = await mgr.drop_partitions_older_than(client, cutoff, _budget()) assert dropped == ["LiteLLM_SpendLogs_p20260601"] executed = " ".join(call.args[0] for call in client.db.execute_raw.call_args_list) @@ -149,8 +197,9 @@ async def test_ensure_partitions_issues_create_for_each_period(): mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2) client = MagicMock() client.db.execute_raw = AsyncMock(return_value=0) + _wire_tx(client.db) - created = await mgr.ensure_partitions(client) + created = await mgr.ensure_partitions(client, _budget()) assert len(created) == 3 # current + 2 ahead assert client.db.execute_raw.await_count == 3 @@ -159,6 +208,105 @@ async def test_ensure_partitions_issues_create_for_each_period(): assert "CREATE TABLE IF NOT EXISTS" in first_sql +@pytest.mark.asyncio +async def test_partition_ddl_carries_a_statement_and_lock_timeout(): + """ + Partition DDL takes an ACCESS EXCLUSIVE lock, so an unbounded DROP queues + behind any long-running reader for as long as that reader lives. That is the + one path by which cleanup could outlast its run budget without bound, and + lock_timeout is what bounds the wait rather than only the work. + """ + mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=0) + client = MagicMock() + client.db.execute_raw = AsyncMock(return_value=0) + client.db.query_raw = AsyncMock( + return_value=[ + { + "name": "LiteLLM_SpendLogs_p20260601", + "bound": "FOR VALUES FROM ('2026-06-01 00:00:00') TO ('2026-06-02 00:00:00')", + } + ] + ) + session_settings = _wire_tx(client.db) + + await mgr.ensure_partitions(client, _budget(7000)) + await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), _budget(7000)) + + # Three statements were issued: the CREATE, the catalog list the drop needs, + # and the DROP. All three carry a statement timeout; only the two that take + # a lock also carry a lock timeout, since the catalog read takes none. + assert session_settings.count("SET LOCAL statement_timeout = 7000") == 3 + assert session_settings.count("SET LOCAL lock_timeout = 7000") == 2 + + +@pytest.mark.asyncio +async def test_catalog_queries_carry_a_statement_timeout(): + """ + Bounding only the DDL leaves the two catalog lookups as statements this job + issues with no bound at all, so a run could still outlast its budget waiting + on one. Every statement the manager issues carries the caller's timeout. + """ + mgr = SpendLogsPartitionManager() + client = MagicMock() + client.db.query_raw = AsyncMock(return_value=[]) + session_settings = _wire_tx(client.db) + + await mgr.is_partitioned(client, _budget(4000)) + assert session_settings == ["SET LOCAL statement_timeout = 4000"], ( + f"is_partitioned issued no statement timeout: {session_settings}" + ) + + session_settings.clear() + await mgr._list_partitions(client, 4000) + assert session_settings == ["SET LOCAL statement_timeout = 4000"], ( + f"_list_partitions issued no statement timeout: {session_settings}" + ) + + +@pytest.mark.asyncio +async def test_partition_loops_stop_when_the_budget_runs_out_mid_way(): + """ + Each loop issues one statement per partition, so a bound read once at entry + would let N statements each run for the budget that was left before the + first of them. The bound is re-read per statement and the loop stops. + """ + mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=4) + client = MagicMock() + client.db.execute_raw = AsyncMock(return_value=0) + _wire_tx(client.db) + + # Budget for two statements, then spent. + calls = {"n": 0} + + def budget() -> "int | None": + calls["n"] += 1 + return 5000 if calls["n"] <= 2 else None + + created = await mgr.ensure_partitions(client, budget) + + assert len(created) == 2, f"the loop ran past its budget and created {len(created)}" + assert client.db.execute_raw.await_count == 2 + + +@pytest.mark.asyncio +async def test_partition_maintenance_issues_nothing_when_the_budget_is_already_spent(): + """A run with no budget left must not issue even the catalog lookups.""" + mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2) + client = MagicMock() + client.db.execute_raw = AsyncMock(return_value=0) + client.db.query_raw = AsyncMock(return_value=[]) + _wire_tx(client.db) + + spent = _budget(None) + + assert await mgr.is_partitioned(client, spent) is False + assert await mgr.ensure_partitions(client, spent) == [] + assert await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), spent) == [] + + client.db.execute_raw.assert_not_awaited() + client.db.query_raw.assert_not_awaited() + + def test_unsupported_interval_raises(): with pytest.raises(ValueError): period_start(date(2026, 6, 1), "year") @@ -178,8 +326,9 @@ async def test_ensure_partitions_continues_when_one_create_fails(): mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2) client = MagicMock() client.db.execute_raw = AsyncMock(side_effect=[0, Exception("overlap"), 0]) + _wire_tx(client.db) - created = await mgr.ensure_partitions(client) + created = await mgr.ensure_partitions(client, _budget()) # the failed partition is skipped, the others still created assert len(created) == 2 @@ -202,8 +351,9 @@ async def test_invalid_interval_does_not_abort_ensure_partitions(): mgr = SpendLogsPartitionManager(interval="fortnight", precreate_ahead=1) client = MagicMock() client.db.execute_raw = AsyncMock(return_value=0) + _wire_tx(client.db) - created = await mgr.ensure_partitions(client) + created = await mgr.ensure_partitions(client, _budget()) assert len(created) == 2 # current + 1 ahead, day-based fallback @@ -225,9 +375,10 @@ async def test_drop_partitions_continues_when_one_drop_fails(): ] ) client.db.execute_raw = AsyncMock(side_effect=[Exception("locked"), 0]) + _wire_tx(client.db) cutoff = datetime(2026, 6, 10, 0, 0, 0, tzinfo=timezone.utc) - dropped = await mgr.drop_partitions_older_than(client, cutoff) + dropped = await mgr.drop_partitions_older_than(client, cutoff, _budget()) # both were eligible; the first drop failed so only the second is reported assert dropped == ["LiteLLM_SpendLogs_p20260602"] 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 0ebdc07d282..bdf09a95e4b 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 @@ -791,6 +791,7 @@ async def test_generate_key_helper_fn_with_access_group_ids(monkeypatch): mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( return_value=MagicMock(object_permission_id=None) ) + mock_prisma_client.db.query_raw = AsyncMock(return_value=[]) captured_key_data = {} @@ -15703,3 +15704,743 @@ async def test_key_generate_omitted_budget_duration_still_filled_by_upperbound(m assert key_row["budget_duration"] == "30d" assert key_row["budget_reset_at"] is not None +from litellm.proxy.management_helpers.access_group_key_sync import ( + _ATTACH_KEY_SQL, + _DETACH_KEY_SQL, + _REPOINT_KEY_SQL, +) + +ACCESS_GROUP_SYNC_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" + + +def _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups): + """ + Back the access group table with an in-memory dict so the sync's writes are observable. + + The sync writes through guarded set-based SQL statements, so this emulates exactly what + Postgres does with them, including the guards that make each one idempotent and the + `RETURNING` clause that reports which groups actually moved. + """ + + def _repoint(previous_token, new_token): + moved = [ + group_id + for group_id, stored in access_groups.items() + if previous_token in stored["assigned_key_ids"] + ] + for group_id in moved: + current = access_groups[group_id]["assigned_key_ids"] + access_groups[group_id]["assigned_key_ids"] = [ + *(t for t in current if t not in (previous_token, new_token)), + new_token, + ] + return moved + + def _attach(key_token, access_group_ids): + moved = [ + group_id + for group_id in access_group_ids + if group_id in access_groups + and key_token not in access_groups[group_id]["assigned_key_ids"] + ] + for group_id in moved: + stored = access_groups[group_id] + stored["assigned_key_ids"] = [*stored["assigned_key_ids"], key_token] + return moved + + def _detach(key_token, access_group_ids): + moved = [ + group_id + for group_id in access_group_ids + if group_id in access_groups + and key_token in access_groups[group_id]["assigned_key_ids"] + ] + for group_id in moved: + stored = access_groups[group_id] + stored["assigned_key_ids"] = [ + t for t in stored["assigned_key_ids"] if t != key_token + ] + return moved + + async def _query_raw(query, *args): + if query == _REPOINT_KEY_SQL: + moved = _repoint(*args) + elif query == _ATTACH_KEY_SQL: + moved = _attach(*args) + else: + assert query == _DETACH_KEY_SQL, f"unexpected statement: {query}" + moved = _detach(*args) + return [{"access_group_id": group_id} for group_id in moved] + + raw_mock = AsyncMock(side_effect=_query_raw) + mock_prisma_client.db.query_raw = raw_mock + return raw_mock + + +async def _authorized_models_for_key(access_groups, token, key_access_group_ids): + """Run the real auth-time reader against the post-sync access group rows.""" + from litellm.proxy._types import LiteLLM_AccessGroupTable, LiteLLM_TeamTable + from litellm.proxy.auth.auth_checks import ( + get_authorized_resources_from_key_access_groups, + ) + + async def _get_access_object(*, access_group_id, **_kwargs): + stored = access_groups[access_group_id] + return LiteLLM_AccessGroupTable( + access_group_id=access_group_id, + access_group_name=access_group_id, + access_model_names=list(stored["access_model_names"]), + assigned_team_ids=[], + assigned_key_ids=list(stored["assigned_key_ids"]), + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + side_effect=_get_access_object, + ), + ): + return await get_authorized_resources_from_key_access_groups( + valid_token=UserAPIKeyAuth( + token=token, + models=[], + team_id="team-a", + access_group_ids=list(key_access_group_ids), + ), + team_object=LiteLLM_TeamTable(team_id="team-a", models=[]), + resource_field="access_model_names", + ) + + +@pytest.mark.asyncio +async def test_update_key_syncs_access_group_assigned_key_ids_in_both_directions( + monkeypatch, +): + """ + A key-side edit of `access_group_ids` must be mirrored onto every affected access + group's `assigned_key_ids`, in one operation, in both directions. + + `assigned_key_ids` is not display-only. `get_authorized_resources_from_key_access_groups` + reads it as an authorization input and authorizes only when the group lists the key's + token, so a group the key just added must start granting its resources and a group the + key dropped must stop. A single-direction assertion would pass against a fix that only + ever adds (or only ever removes), so this covers add, remove, untouched, and the + authorization consequence of each. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + key_in_db = LiteLLM_VerificationToken( + token=ACCESS_GROUP_SYNC_TOKEN, + user_id="test-user", + access_group_ids=["ag-drop", "ag-keep"], + ) + access_groups = { + "ag-drop": { + "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN], + "access_model_names": ["dropped-model"], + }, + "ag-keep": { + "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN], + "access_model_names": ["kept-model"], + }, + "ag-add": {"assigned_key_ids": [], "access_model_names": ["added-model"]}, + } + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_in_db + ) + mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock( + return_value=None + ) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {}}) + raw_mock = _access_group_table_mocks( + monkeypatch, mock_prisma_client, access_groups + ) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache", + new_callable=AsyncMock, + ) as invalidate_cache, + ): + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest( + key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=["ag-keep", "ag-add"] + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ), + litellm_changed_by=None, + ) + + assert access_groups["ag-drop"]["assigned_key_ids"] == [] + assert access_groups["ag-add"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN] + assert access_groups["ag-keep"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN] + + # Both halves go out as single guarded statements. A read-modify-write here lets two + # admins editing one group lose each other's change: an attach can vanish, and a detach + # can put an already revoked token back and restore its grants. + assert sorted(call.args for call in raw_mock.call_args_list) == sorted( + [ + (_ATTACH_KEY_SQL, ACCESS_GROUP_SYNC_TOKEN, ["ag-add"]), + (_DETACH_KEY_SQL, ACCESS_GROUP_SYNC_TOKEN, ["ag-drop"]), + ] + ) + assert {call.args[0] for call in invalidate_cache.call_args_list} == { + "ag-drop", + "ag-add", + } + + authorized_models = await _authorized_models_for_key( + access_groups, + ACCESS_GROUP_SYNC_TOKEN, + ["ag-drop", "ag-keep", "ag-add"], + ) + assert sorted(authorized_models) == ["added-model", "kept-model"] + + +@pytest.mark.asyncio +async def test_update_key_leaves_access_groups_alone_when_field_is_unset(monkeypatch): + """ + An update that never mentions `access_group_ids` must not touch the group rows. + + `prepare_key_update_data` writes from `model_dump(exclude_unset=True)`, so an omitted + field leaves the key row's own list intact. Reading the request attribute instead of + its `model_fields_set` would see None and wipe every group's copy of the token on any + unrelated edit, e.g. a max_budget change. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + key_in_db = LiteLLM_VerificationToken( + token=ACCESS_GROUP_SYNC_TOKEN, + user_id="test-user", + access_group_ids=["ag-keep"], + ) + access_groups = { + "ag-keep": { + "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN], + "access_model_names": ["kept-model"], + }, + } + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_in_db + ) + mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock( + return_value=None + ) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {}}) + raw_mock = _access_group_table_mocks( + monkeypatch, mock_prisma_client, access_groups + ) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ): + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=ACCESS_GROUP_SYNC_TOKEN, max_budget=50.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ), + litellm_changed_by=None, + ) + + raw_mock.assert_not_called() + assert access_groups["ag-keep"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN] + assert await _authorized_models_for_key( + access_groups, ACCESS_GROUP_SYNC_TOKEN, ["ag-keep"] + ) == ["kept-model"] + + +@pytest.mark.asyncio +async def test_bulk_update_keys_syncs_access_group_assigned_key_ids(monkeypatch): + """ + /key/bulk_update and /team/keys/bulk_update reach the DB through + `_process_single_key_update`, which is a separate write path from /key/update's own + inline one. Both have to maintain the group's copy or a bulk attach grants nothing. + """ + key_in_db = LiteLLM_VerificationToken( + token=ACCESS_GROUP_SYNC_TOKEN, + user_id="test-user", + access_group_ids=["ag-drop"], + ) + access_groups = { + "ag-drop": { + "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN], + "access_model_names": ["dropped-model"], + }, + "ag-add": {"assigned_key_ids": [], "access_model_names": ["added-model"]}, + } + + mock_prisma_client = AsyncMock() + mock_prisma_client.update_data = AsyncMock(return_value={"data": {}}) + _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", + new_callable=AsyncMock, + ), + ): + await _process_single_key_update( + update_key_request=UpdateKeyRequest( + key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=["ag-add"] + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ), + litellm_changed_by=None, + prisma_client=mock_prisma_client, + user_api_key_cache=AsyncMock(), + proxy_logging_obj=MagicMock(), + llm_router=None, + existing_key_row=key_in_db, + ) + + assert access_groups["ag-drop"]["assigned_key_ids"] == [] + assert access_groups["ag-add"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN] + assert await _authorized_models_for_key( + access_groups, ACCESS_GROUP_SYNC_TOKEN, ["ag-drop", "ag-add"] + ) == ["added-model"] + + +@pytest.mark.asyncio +async def test_delete_key_withdraws_token_from_its_access_groups(monkeypatch): + """ + Deleting a key must withdraw its token from every group that lists it. + + Without the withdrawal the group keeps a token that no longer resolves to a row, so + the access group page lists a key that does not exist and the list grows without bound. + """ + key_in_db = LiteLLM_VerificationToken( + token=ACCESS_GROUP_SYNC_TOKEN, + user_id="test-user", + access_group_ids=["ag-keep"], + ) + access_groups = { + "ag-keep": { + "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN, "other-key"], + "access_model_names": ["kept-model"], + }, + } + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[key_in_db] + ) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_keys": 1}) + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) + + mock_cache = MagicMock() + mock_cache.delete_cache = MagicMock() + + with patch( + "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache", + new_callable=AsyncMock, + ): + await delete_verification_tokens( + tokens=[ACCESS_GROUP_SYNC_TOKEN], + user_api_key_cache=mock_cache, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ), + litellm_changed_by="admin-user", + ) + + assert access_groups["ag-keep"]["assigned_key_ids"] == ["other-key"] + + +@pytest.mark.asyncio +async def test_generate_key_records_token_in_its_access_groups(monkeypatch): + """ + /key/generate with `access_group_ids` must record the new token on the group side. + + The key row's own list alone does not authorize: the group has to list the token back + or `get_authorized_resources_from_key_access_groups` contributes nothing, so a key + created against a group silently gets none of its models. + """ + access_groups = { + "ag-add": {"assigned_key_ids": [], "access_model_names": ["added-model"]}, + } + + created_key = MagicMock() + created_key.token = ACCESS_GROUP_SYNC_TOKEN + created_key.litellm_budget_table = None + created_key.created_at = None + created_key.updated_at = None + + mock_prisma_client = AsyncMock() + mock_prisma_client.insert_data = AsyncMock(return_value=created_key) + _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.store_audit_logs", False) + + with patch( + "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache", + new_callable=AsyncMock, + ): + await generate_key_helper_fn( + request_type="key", + access_group_ids=["ag-add"], + table_name="key", + user_id="test-user", + ) + + assert access_groups["ag-add"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN] + assert await _authorized_models_for_key( + access_groups, ACCESS_GROUP_SYNC_TOKEN, ["ag-add"] + ) == ["added-model"] + + +@pytest.mark.asyncio +async def test_regenerate_key_repoints_access_group_assigned_key_ids(monkeypatch): + """ + Regeneration replaces the key's token, which is the identity `assigned_key_ids` stores. + + Leaving the old hash behind points the group at a token that no longer exists AND + denies the regenerated key the group's grants, so the group's copy has to be + re-pointed from the old hash to the new one in the same operation. + """ + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + from litellm.proxy.utils import hash_token + + new_token_hash = hash_token("sk-newtoken1234ab12") + existing_key = LiteLLM_VerificationToken( + token="abc123", + user_id="user-1", + models=["gpt-4"], + access_group_ids=["ag-keep"], + ) + access_groups = { + "ag-keep": { + "assigned_key_ids": ["abc123"], + "access_model_names": ["kept-model"], + }, + } + + mock_prisma_client = _make_regenerate_mock_prisma() + _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups) + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache", + new_callable=AsyncMock, + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=RegenerateKeyRequest(), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert access_groups["ag-keep"]["assigned_key_ids"] == [new_token_hash] + assert await _authorized_models_for_key( + access_groups, new_token_hash, ["ag-keep"] + ) == ["kept-model"] + assert ( + await _authorized_models_for_key(access_groups, "abc123", ["ag-keep"]) == [] + ) + + +@pytest.mark.asyncio +async def test_key_write_paths_revoke_the_key_cache_before_syncing_access_groups( + monkeypatch, +): + """ + Credential invalidation must not sit behind the group sync on any key write path. + + The cached auth object still carries the key's old `access_group_ids`, so if the sync + raises first, the request fails with the key still authenticating against groups it + just lost, until that entry expires. Ordering it last means a failed sync degrades to + the stale listing this PR fixes rather than to a stale grant. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + order = [] + + key_in_db = LiteLLM_VerificationToken( + token=ACCESS_GROUP_SYNC_TOKEN, + user_id="test-user", + access_group_ids=["ag-drop"], + ) + access_groups = { + "ag-drop": { + "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN], + "access_model_names": ["dropped-model"], + }, + } + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_in_db + ) + mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock( + return_value=None + ) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {}}) + _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups) + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=lambda *a, **k: order.append("sync") or [] + ) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + side_effect=lambda **kwargs: order.append("revoke_key_cache"), + ), + patch( + "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache", + new_callable=AsyncMock, + ), + ): + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=[]), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ), + litellm_changed_by=None, + ) + + assert order == ["revoke_key_cache", "sync"] + + +@pytest.mark.asyncio +async def test_update_key_syncs_many_access_groups_in_one_statement_per_direction( + monkeypatch, +): + """ + The number of groups on a request must not become a matching number of round trips. + + Anyone allowed to assign access groups picks the size of `access_group_ids`, so a + per-group statement lets one /key/update hold a connection for hundreds of sequential + writes. Both halves are set-based, so the cost is two statements no matter the size. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + dropped = [f"ag-drop-{i}" for i in range(60)] + added = [f"ag-add-{i}" for i in range(60)] + key_in_db = LiteLLM_VerificationToken( + token=ACCESS_GROUP_SYNC_TOKEN, + user_id="test-user", + access_group_ids=dropped, + ) + access_groups = { + **{ + group_id: { + "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN], + "access_model_names": [f"{group_id}-model"], + } + for group_id in dropped + }, + **{ + group_id: {"assigned_key_ids": [], "access_model_names": [f"{group_id}-model"]} + for group_id in added + }, + } + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_in_db + ) + mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock( + return_value=None + ) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {}}) + raw_mock = _access_group_table_mocks( + monkeypatch, mock_prisma_client, access_groups + ) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache", + new_callable=AsyncMock, + ), + ): + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest( + key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=added + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ), + litellm_changed_by=None, + ) + + assert [call.args[0] for call in raw_mock.call_args_list] == [ + _ATTACH_KEY_SQL, + _DETACH_KEY_SQL, + ] + assert sorted(raw_mock.call_args_list[0].args[2]) == sorted(added) + assert sorted(raw_mock.call_args_list[1].args[2]) == sorted(dropped) + assert all( + access_groups[group_id]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN] + for group_id in added + ) + assert all(access_groups[group_id]["assigned_key_ids"] == [] for group_id in dropped) + + +@pytest.mark.asyncio +async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( + monkeypatch, +): + """ + Regeneration must move whatever the groups hold when it writes, not the key row's list. + + That list is read before the new token exists, so replaying it re-adds the key to a + group an admin revoked in between and leaves the dead hash in a group an admin attached + in between, which silently restores one grant and drops another. + """ + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + from litellm.proxy.utils import hash_token + + new_token_hash = hash_token("sk-newtoken1234ab12") + existing_key = LiteLLM_VerificationToken( + token="abc123", + user_id="user-1", + models=["gpt-4"], + access_group_ids=["ag-revoked-since"], + ) + access_groups = { + "ag-revoked-since": { + "assigned_key_ids": [], + "access_model_names": ["revoked-model"], + }, + "ag-attached-since": { + "assigned_key_ids": ["abc123"], + "access_model_names": ["attached-model"], + }, + } + + mock_prisma_client = _make_regenerate_mock_prisma() + _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups) + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache", + new_callable=AsyncMock, + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=RegenerateKeyRequest(), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert access_groups["ag-revoked-since"]["assigned_key_ids"] == [] + assert access_groups["ag-attached-since"]["assigned_key_ids"] == [new_token_hash] + assert await _authorized_models_for_key( + access_groups, new_token_hash, ["ag-revoked-since", "ag-attached-since"] + ) == ["attached-model"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index b6c2fd4c8e3..84dee5b05c5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1688,9 +1688,15 @@ class TestTemporaryMCPSessionEndpoints: expires_at=datetime.utcnow() - timedelta(seconds=30), ) cache = {"expired": expired_entry} - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers", - cache, + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers", + cache, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=None, + ), ): result = await get_cached_temporary_mcp_server("expired") @@ -2274,6 +2280,10 @@ class TestTemporaryMCPSessionEndpoints: "litellm.proxy.management_endpoints.mcp_management_endpoints._cache_temporary_mcp_server_in_redis", AsyncMock(), ) as redis_cache_mock, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=None, + ), ): response = await add_session_mcp_server( payload=payload, @@ -3419,6 +3429,10 @@ class TestTemporaryMCPSessionEndpoints: "litellm.proxy.management_endpoints.mcp_management_endpoints.decrypt_value_helper", return_value=serialized, ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=None, + ), ): result = await get_cached_temporary_mcp_server("from-redis") finally: diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index 178e9379ea7..6e670e48b6a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -4,6 +4,7 @@ import datetime import json from contextlib import ExitStack from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import patch as patch_ctx import pytest from fastapi import HTTPException @@ -14,15 +15,28 @@ from litellm.proxy._types import ( ReconcileOutcome, UserAPIKeyAuth, ) +from litellm.proxy.auth.auth_checks import _is_model_cost_zero from litellm.proxy.management_endpoints.model_management_endpoints import ( + _PTU_ZEROED_PRICING_FIELDS, _merged_ptu_model_info, + _update_team_model_in_db, + _ptu_priced_deployment, + _ptu_zeroed_pricing, _raise_if_ptu_cost_attribution_disabled, _validate_ptu_model_info, add_new_model, update_db_model, ) from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR -from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment +from litellm.router import Router +from litellm.types.router import ( + SPECIAL_MODEL_INFO_PARAMS, + Deployment, + LiteLLM_Params, + ModelInfo, + updateDeployment, + updateLiteLLMParams, +) def test_model_info_accepts_valid_ptu_fields(): @@ -717,3 +731,390 @@ class TestAddNewModelPtuGate: assert result.model_id == "ptu-gate-model" add_team_model_to_db.assert_called_once() + + + +class TestPtuDeploymentsAreNotBilledPerToken: + """Reserved capacity is billed by the flat cost the rollup writes, so a PTU deployment must + not also bill the traffic that capacity serves.""" + + PTU = {"ptu_count": 15, "cost_per_ptu_per_hour": 2.0} + + @pytest.fixture(autouse=True) + def _flag_on(self, monkeypatch): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + # update_db_model encrypts every litellm_params value it is handed, and the salt falls + # back to the master key the proxy sets at boot, which no unit test has. + monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-key") + + @staticmethod + def _zeroed(model_info=None, litellm_params=None, supplied=None): + return _ptu_zeroed_pricing( + model_info=model_info if model_info is not None else {}, + litellm_params=litellm_params if litellm_params is not None else {}, + supplied=supplied if supplied is not None else {}, + ) + + def test_a_deployment_without_ptu_config_keeps_its_pricing(self): + assert self._zeroed(model_info={"team_id": "t"}, litellm_params={"input_cost_per_token": 5e-07}) == {} + + def test_a_half_set_pair_is_not_treated_as_ptu(self): + assert self._zeroed(model_info={"ptu_count": 15}) == {} + + def test_every_field_the_cost_map_could_fill_is_zeroed(self): + assert self._zeroed(model_info=self.PTU) == dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0) + + def test_nothing_is_zeroed_while_the_feature_is_disabled(self, monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + assert self._zeroed(model_info=self.PTU) == {} + + @pytest.mark.parametrize("field", ["input_cost_per_token", "cache_read_input_token_cost", "input_cost_per_second"]) + def test_a_price_the_caller_supplies_is_refused(self, field): + """Every custom-pricing field, not only the mirrored ones: per-second pricing bills a + PTU deployment just as surely as per-token pricing does.""" + with pytest.raises(HTTPException) as exc: + self._zeroed(model_info=self.PTU, supplied={field: 5e-07}) + assert exc.value.status_code == 400 + assert field in str(exc.value.detail) + + def test_a_price_the_caller_supplies_as_zero_is_accepted(self): + assert self._zeroed(model_info={**self.PTU, "input_cost_per_token": 0}, supplied={"input_cost_per_token": 0})[ + "input_cost_per_token" + ] == 0 + + def test_a_price_already_on_the_row_is_zeroed_rather_than_refused(self): + """A row priced through a path this rule does not cover must heal on its next save. The + alternative refuses every later edit of a field that has nothing to do with pricing.""" + zeroed = self._zeroed(model_info={**self.PTU, "input_cost_per_second": 3.0}, litellm_params={}) + assert zeroed["input_cost_per_second"] == 0 + assert zeroed["input_cost_per_token"] == 0 + + @pytest.mark.asyncio + async def test_a_refused_price_does_not_leave_the_team_changed(self): + """The team ACL write autocommits, so the refusal has to run before it. Otherwise a + rejected edit grants the team a model whose settings were never saved.""" + db_model = Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo( + id="dep-0", + team_id="team-1", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ), + ) + patch = updateDeployment( + litellm_params=updateLiteLLMParams(model="openai/gpt-4o", input_cost_per_token=5e-07), + model_info=ModelInfo(id="dep-0", team_id="team-2"), + ) + endpoints = "litellm.proxy.management_endpoints.model_management_endpoints" + setup_new = AsyncMock() + update_existing = AsyncMock() + with ExitStack() as stack: + stack.enter_context( + patch_ctx(f"{endpoints}.ModelManagementAuthChecks.allow_team_model_action", AsyncMock(return_value=True)) + ) + stack.enter_context(patch_ctx(f"{endpoints}._setup_new_team_model_assignment", setup_new)) + stack.enter_context(patch_ctx(f"{endpoints}._update_existing_team_model_assignment", update_existing)) + stack.enter_context(patch_ctx("litellm.proxy.proxy_server.premium_user", True)) + with pytest.raises(HTTPException) as exc: + await _update_team_model_in_db( + db_model=db_model, + patch_data=patch, + user_api_key_dict=UserAPIKeyAuth(user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN), + prisma_client=MagicMock(), + ) + + assert exc.value.status_code == 400 + setup_new.assert_not_called() + update_existing.assert_not_called() + + def test_a_setting_that_is_not_a_charge_is_left_alone(self): + """CustomPricingLiteLLMParams also carries an embedding's output vector size and the + regional uplift multipliers. Zeroing one of those destroys the deployment's config, and + refusing it answers with a message calling a setting a charge.""" + priced = _ptu_priced_deployment( + Deployment( + model_name="embeddings", + litellm_params=LiteLLM_Params( + model="azure/text-embedding-3-large", + output_vector_size=1536, + regional_processing_uplift_multiplier_eu=1.15, + ), + model_info=ModelInfo( + id="dep-emb", + team_id="team-1", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ), + ) + ) + assert priced.litellm_params.get("output_vector_size") == 1536 + assert priced.litellm_params.get("regional_processing_uplift_multiplier_eu") == 1.15 + assert priced.litellm_params.get("input_cost_per_token") == 0 + + def test_removing_ptu_config_releases_every_rate_it_zeroed(self): + """The zeroing covers any stored rate, so a release that only spans the mirrored fields + leaves a per-second deployment billing nothing for that dimension forever.""" + on = update_db_model( + db_model=Deployment( + model_name="audio", + litellm_params=LiteLLM_Params(model="azure/whisper", input_cost_per_second=0.006), + model_info=ModelInfo(id="dep-audio", team_id="t"), + ), + updated_patch=updateDeployment( + model_info=ModelInfo( + id="dep-audio", + team_id="t", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ) + ), + ) + assert json.loads(on["litellm_params"])["input_cost_per_second"] == 0 + + off = update_db_model( + db_model=Deployment( + model_name="audio", + litellm_params=LiteLLM_Params(**json.loads(on["litellm_params"])), + model_info=ModelInfo(**json.loads(on["model_info"])), + ), + updated_patch=updateDeployment( + model_info=ModelInfo(id="dep-audio", ptu_count=None, cost_per_ptu_per_hour=None) + ), + ) + assert "input_cost_per_second" not in json.loads(off["litellm_params"]) + + @pytest.mark.parametrize( + "backend", ["azure/gpt-4o", "anthropic/claude-sonnet-4-5", "bedrock/anthropic.claude-sonnet-4-20250514-v1:0"] + ) + def test_the_cost_map_contributes_no_price_to_a_priced_ptu_deployment(self, backend): + """The acceptance criterion, read off the entry the router registers for the deployment. + + Zeroing only the per-token pair leaves the cache-tier fields unset, which is exactly what + Router._inherit_builtin_cache_pricing back-fills from the public cost map, so a cached + prompt would still be billed at the public rate.""" + priced = _ptu_priced_deployment( + Deployment( + model_name="ptu-deployment", + litellm_params=LiteLLM_Params(model=backend, api_key="fake-key"), + model_info=ModelInfo( + id="dep-ptu", + team_id="team-1", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ), + ) + ) + registered = Router._deployment_model_cost_payload(priced) + charged = {k: v for k, v in registered.items() if "cost" in k and k != "cost_per_ptu_per_hour" and v} + assert charged == {} + + def test_the_zeroed_pricing_does_not_waive_budget_enforcement(self): + """A zero price otherwise tells auth the model is free and skips every budget check.""" + priced = _ptu_priced_deployment( + Deployment( + model_name="model_name_team-1_dep-ptu", + litellm_params=LiteLLM_Params(model="gemini/gemini-2.5-flash", api_key="fake-key"), + model_info=ModelInfo( + id="dep-ptu", + team_id="team-1", + team_public_model_name="ptu-model", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ), + ) + ) + router = Router(model_list=[priced.to_json(exclude_none=True)]) + assert _is_model_cost_zero(model="model_name_team-1_dep-ptu", llm_router=router) is False + assert _is_model_cost_zero(model="ptu-model", llm_router=router) is False + + def test_an_unrelated_patch_heals_a_deployment_stored_before_this_rule(self): + """Both blobs, because litellm_params wins over model_info wherever the two are merged.""" + written = update_db_model( + db_model=_deployment_with_stored_ptu(), + updated_patch=updateDeployment(model_name="gpt-4o-renamed"), + ) + for blob in ("model_info", "litellm_params"): + stored = json.loads(written[blob]) + assert all(stored[field] == 0 for field in _PTU_ZEROED_PRICING_FIELDS), blob + + def test_an_unrelated_patch_of_a_ptu_row_that_carries_a_price_is_not_refused(self): + """The pause toggle and the credential-rotation modal send no pricing at all. Refusing + them because the stored row is mispriced blocks flows that cannot fix it.""" + priced_ptu = Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", input_cost_per_token=5e-07), + model_info=ModelInfo( + id="dep-0", + team_id="t", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ), + ) + written = update_db_model(db_model=priced_ptu, updated_patch=updateDeployment(model_name="renamed")) + assert written["model_name"] == "renamed" + assert json.loads(written["litellm_params"])["input_cost_per_token"] == 0 + + def test_removing_ptu_config_hands_per_token_billing_back(self): + """Left behind, the zeros this rule wrote would serve the deployment for free forever.""" + zeros = dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0) + written = update_db_model( + db_model=Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", **zeros), + model_info=ModelInfo( + id="dep-0", + team_id="t", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + **zeros, + ), + ), + updated_patch=updateDeployment( + model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None) + ), + ) + for blob in ("model_info", "litellm_params"): + stored = json.loads(written[blob]) + assert not any(field in stored for field in _PTU_ZEROED_PRICING_FIELDS), blob + + def test_the_dashboard_clear_releases_the_zeros_it_echoes_back(self): + """The edit form re-sends the whole stored model_info on every save, so the clearing + patch carries the zeros this rule wrote. Treating those as a rate the operator chose + left the deployment serving free and reading as a free model to the budget checks.""" + zeros = dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0) + written = update_db_model( + db_model=Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", **zeros), + model_info=ModelInfo( + id="dep-0", + team_id="t", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + **zeros, + ), + ), + updated_patch=updateDeployment( + model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None, **zeros) + ), + ) + stored = json.loads(written["model_info"]) + assert not any(field in stored for field in _PTU_ZEROED_PRICING_FIELDS) + + def test_a_deployment_that_never_had_ptu_keeps_a_price_its_operator_set_to_zero(self): + """The dashboard sends both PTU keys as null on every save while the feature is on, so a + release keyed on the patch alone would strip a deliberate zero rate from any model.""" + free = Deployment( + model_name="free-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", input_cost_per_token=0.0), + model_info=ModelInfo(id="dep-free", team_id="t", input_cost_per_token=0.0), + ) + written = update_db_model( + db_model=free, + updated_patch=updateDeployment( + model_info=ModelInfo(id="dep-free", ptu_count=None, cost_per_ptu_per_hour=None) + ), + ) + for blob in ("model_info", "litellm_params"): + assert json.loads(written[blob])["input_cost_per_token"] == 0, blob + + def test_a_patch_pricing_a_ptu_deployment_is_refused(self): + with pytest.raises(HTTPException) as exc: + update_db_model( + db_model=_deployment_with_stored_ptu(), + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(model="openai/gpt-4o", input_cost_per_token=5e-07) + ), + ) + assert exc.value.status_code == 400 + + def test_a_price_the_client_only_echoes_back_is_not_read_as_an_attempt_to_charge(self): + """/model/info fills missing rates from the public cost map and the edit form re-sends the + whole blob, so a model_info price is one the server wrote. Reading it as the operator's + refused every attempt to put an existing deployment on PTU from the dashboard.""" + written = update_db_model( + db_model=_deployment_without_ptu(), + updated_patch=updateDeployment( + model_info=ModelInfo( + id="dep-0", + team_id="t", + input_cost_per_token=3e-07, + output_cost_per_token=2.5e-06, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ) + ), + ) + stored = json.loads(written["model_info"]) + assert stored["ptu_count"] == 15 + assert stored["input_cost_per_token"] == 0 + assert stored["output_cost_per_token"] == 0 + + def test_adding_ptu_config_to_an_already_priced_deployment_is_refused(self): + priced = Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo(id="dep-0", team_id="t"), + ) + with pytest.raises(HTTPException) as exc: + update_db_model( + db_model=priced, + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(model="openai/gpt-4o", input_cost_per_token=5e-07), + model_info=ModelInfo( + id="dep-0", + team_id="t", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ), + ), + ) + assert exc.value.status_code == 400 + + def test_a_deployment_without_ptu_config_keeps_its_pricing_through_a_patch(self): + priced = Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo(id="dep-0", team_id="t", input_cost_per_token=5e-07), + ) + stored = json.loads( + update_db_model(db_model=priced, updated_patch=updateDeployment(model_name="renamed"))["model_info"] + ) + assert stored["input_cost_per_token"] == 5e-07 + + @pytest.mark.asyncio + async def test_model_new_stores_zero_pricing_on_both_blobs(self): + (_, add_team_model_to_db), patches = TestAddNewModelPtuGate._patched_proxy("ptu-priced-model") + admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with ExitStack() as stack: + for active_patch in patches: + stack.enter_context(active_patch) + await add_new_model( + model_params=TestAddNewModelPtuGate._ptu_deployment("ptu-priced-model"), + user_api_key_dict=admin, + ) + + written = add_team_model_to_db.call_args.kwargs["model_params"] + assert all(getattr(written.model_info, field, None) == 0 for field in SPECIAL_MODEL_INFO_PARAMS) + assert all(written.litellm_params.get(field) == 0 for field in _PTU_ZEROED_PRICING_FIELDS) + + @pytest.mark.asyncio + async def test_model_new_refuses_a_priced_ptu_deployment(self): + (_, add_team_model_to_db), patches = TestAddNewModelPtuGate._patched_proxy("ptu-priced-model") + admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + base = TestAddNewModelPtuGate._ptu_deployment("ptu-priced-model") + deployment = base.model_copy( + update={"litellm_params": base.litellm_params.model_copy(update={"input_cost_per_token": 5e-07})} + ) + + with ExitStack() as stack: + for active_patch in patches: + stack.enter_context(active_patch) + with pytest.raises(Exception) as exc: + await add_new_model(model_params=deployment, user_api_key_dict=admin) + + assert "input_cost_per_token" in str(exc.value) + add_team_model_to_db.assert_not_called() 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 f038af4c7a1..c6960ecda5a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2,7 +2,9 @@ import asyncio import json import os import sys +from contextlib import asynccontextmanager from datetime import datetime, timezone +from types import SimpleNamespace from typing import Optional, cast from unittest.mock import AsyncMock, MagicMock, call, patch @@ -68,6 +70,21 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( # Setup TestClient client = TestClient(app) + +def _wire_team_create_tx(prisma_client): + """`/team/new` inserts the team and mirrors it onto the access groups in one transaction, + so a mocked client has to hand its team table back out of `db.tx()`.""" + + @asynccontextmanager + async def _tx(): + yield SimpleNamespace( + litellm_teamtable=prisma_client.db.litellm_teamtable, + query_raw=AsyncMock(return_value=[]), + ) + + prisma_client.db.tx = lambda *_args, **_kwargs: _tx() + + # Mock prisma_client mock_prisma_client = MagicMock() # Set up async mock for db operations @@ -400,6 +417,7 @@ async def test_new_team_rejects_a_duration_that_never_advances( mock_team_create = AsyncMock() mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) with pytest.raises(ProxyException) as exc_info: await new_team( @@ -481,6 +499,7 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): mock_team_count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = mock_team_count mock_db_client.db.litellm_teamtable.update = AsyncMock( return_value=team_create_result @@ -570,6 +589,7 @@ async def test_new_team_with_mcp_tool_permissions(mock_db_client, mock_admin_aut mock_db_client.db.litellm_teamtable.create = AsyncMock( return_value=team_create_result ) + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable.update = AsyncMock( return_value=team_create_result @@ -663,6 +683,7 @@ async def test_new_team_disable_auto_add_proxy_admin_flag( mock_db_client.db.litellm_teamtable.create = AsyncMock( return_value=team_create_result ) + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_usertable = MagicMock() mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) @@ -4430,6 +4451,7 @@ async def test_new_team_max_budget_within_user_limit(): mock_prisma.db.litellm_teamtable.create = AsyncMock( return_value=mock_created_team ) + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_created_team ) @@ -4573,6 +4595,7 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit(): mock_prisma.db.litellm_teamtable.create = AsyncMock( return_value=mock_created_team ) + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_created_team ) @@ -4721,6 +4744,7 @@ async def test_new_team_org_scoped_models_bypasses_user_limit(): mock_prisma.db.litellm_teamtable.create = AsyncMock( return_value=mock_created_team ) + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_created_team ) @@ -6567,6 +6591,7 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit(): mock_created_team.rpm_limit = 1000 mock_created_team.metadata = None mock_created_team.members_with_roles = [] + mock_created_team.access_group_ids = None mock_created_team.model_dump.return_value = { "team_id": "new-bypass-team-id", "team_alias": "org-bypass-test-team", @@ -6578,6 +6603,7 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit(): mock_prisma.db.litellm_teamtable.create = AsyncMock( return_value=mock_created_team ) + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_created_team ) @@ -6856,6 +6882,7 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(): mock_updated_team.team_id = "org-team-update-bypass-123" mock_updated_team.tpm_limit = 10000 mock_updated_team.rpm_limit = 1000 + mock_updated_team.access_group_ids = None mock_updated_team.model_dump.return_value = { "team_id": "org-team-update-bypass-123", "tpm_limit": 10000, @@ -7009,6 +7036,7 @@ async def test_update_team_guardrails_with_org_id(): "guardrails": ["aporia-pre-call", "aporia-post-call"] } mock_updated_team.litellm_model_table = None + mock_updated_team.access_group_ids = None mock_updated_team.model_dump.return_value = { "team_id": "team-guardrails-123", "organization_id": "test-org-guardrails", @@ -7937,6 +7965,7 @@ async def test_new_team_soft_budget_validation( mock_prisma.db.litellm_teamtable.create = AsyncMock( return_value=mock_created_team ) + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_created_team ) @@ -8236,6 +8265,7 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth): mock_team_count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = mock_team_count mock_db_client.db.litellm_teamtable.update = AsyncMock( return_value=team_create_result @@ -9626,6 +9656,7 @@ async def test_new_team_encrypts_callback_vars( team_create_result.model_dump.return_value = {"team_id": "team-456"} mock_team_create = AsyncMock(return_value=team_create_result) mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable.update = AsyncMock( return_value=team_create_result @@ -10786,6 +10817,7 @@ async def test_new_team_validator_runs_without_metadata_and_rejection_blocks_cre ): mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.db.litellm_teamtable.create = AsyncMock() + _wire_team_create_tx(mock_prisma) mock_license.is_team_count_over_limit.return_value = False with pytest.raises(ProxyException) as exc_info: @@ -10820,6 +10852,7 @@ async def test_new_team_validator_accept_proceeds_to_create(mock_db_client, mock team_create_result.model_dump.return_value = {"team_id": "team-accept-1"} mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=team_create_result) + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) mock_db_client.db.litellm_usertable = MagicMock() @@ -10859,6 +10892,7 @@ async def test_new_team_rejection_precedes_model_alias_write(): ): mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.db.litellm_teamtable.create = AsyncMock() + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1")) mock_license.is_team_count_over_limit.return_value = False @@ -11626,6 +11660,7 @@ def _wire_new_team_prisma(mock_db_client): mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=created_team) + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=created_team) mock_db_client.db.litellm_usertable = MagicMock() mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) @@ -11714,3 +11749,342 @@ async def test_new_team_explicit_null_max_budget_still_takes_configured_default( team_data = mock_team_create.call_args.kwargs["data"] assert team_data.get("max_budget") == 100.0 + + +class _FakeMirrorDb: + """Stands in for prisma inside the access-group mirror. + + Dispatches on the statement so a change to the SQL's shape is visible here, but it + cannot validate the SQL itself: it reimplements the array semantics in Python, so it + passes whatever the statement says. Correctness of the SQL is pinned against a real + Postgres in tests/proxy_admin_ui_tests/test_access_group_team_sync.py. + """ + + def __init__(self, access_groups, teams, plain_lists=False): + self._access_groups = access_groups + self._teams = teams + self._plain_lists = plain_lists + self.transactions = [] + + def _team_ids(self, group_id): + stored = self._access_groups[group_id] + return stored if self._plain_lists else stored["assigned_team_ids"] + + async def _query_raw(self, sql, *args): + assert self._open, "mirror statement ran outside a transaction" + if "pg_advisory_xact_lock" in sql: + self.transactions[-1].append("lock") + return [{"locked": False}] + if "LiteLLM_TeamTable" in sql: + self.transactions[-1].append("read") + team_id = args[0] + if team_id not in self._teams: + return [] + return [{"access_group_ids": list(self._teams[team_id])}] + + team_id, desired = args + if sql.lstrip().startswith("SELECT"): + self.transactions[-1].append("affected") + affected = [g for g in self._access_groups if g in desired or team_id in self._team_ids(g)] + return [{"access_group_id": group_id} for group_id in affected] + + if "array_append" in sql: + self.transactions[-1].append("attach") + changed = [ + g for g in desired if g in self._access_groups and team_id not in self._team_ids(g) + ] + for group_id in changed: + self._team_ids(group_id).append(team_id) + else: + self.transactions[-1].append("detach") + changed = [ + g for g in self._access_groups if team_id in self._team_ids(g) and g not in desired + ] + for group_id in changed: + self._team_ids(group_id).remove(team_id) + return [{"access_group_id": group_id} for group_id in changed] + + async def _create_team(self, data, include=None): + self.transactions[-1].append("create") + team_id = data["team_id"] + self._teams[team_id] = list(data.get("access_group_ids") or ()) + return SimpleNamespace( + team_id=team_id, + access_group_ids=list(self._teams[team_id]), + model_dump=lambda: {"team_id": team_id}, + ) + + def tx(self, *_args, **_kwargs): + outer = self + + class _Tx: + async def __aenter__(self): + outer.transactions.append([]) + outer._open = True + return SimpleNamespace( + query_raw=outer._query_raw, + litellm_teamtable=SimpleNamespace(create=outer._create_team), + ) + + async def __aexit__(self, *_exc_info): + outer._open = False + return None + + return _Tx() + + _open = False + + +@pytest.mark.asyncio +async def test_update_team_syncs_access_group_assigned_team_ids_in_both_directions(): + """ + A team-side edit of `access_group_ids` must be mirrored onto every affected access + group's `assigned_team_ids`, in one transaction, in both directions. + + `assigned_team_ids` is not display-only. `get_authorized_resources_from_key_access_groups` + reads it as an authorization input, so a group the team dropped must stop granting its + resources to keys on that team, and a group the team added must start granting them. + A single-direction assertion would pass against a fix that only ever removes (or only + ever adds), so this covers add, remove, untouched, and the authorization consequence. + """ + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.auth.auth_checks import ( + get_authorized_resources_from_key_access_groups, + ) + + access_groups = { + "ag-drop": {"assigned_team_ids": ["team-a"], "access_model_names": ["dropped-model"]}, + "ag-keep": {"assigned_team_ids": ["team-a"], "access_model_names": ["kept-model"]}, + "ag-add": {"assigned_team_ids": [], "access_model_names": ["added-model"]}, + "ag-other-team": {"assigned_team_ids": ["team-b"], "access_model_names": ["other-model"]}, + } + committed_team_groups = ["ag-keep", "ag-add"] + fake_db = _FakeMirrorDb(access_groups, {"team-a": committed_team_groups}) + + existing_team = MagicMock() + existing_team.access_group_ids = ["ag-drop", "ag-keep"] + existing_team.metadata = {} + existing_team.max_budget = None + existing_team.organization_id = None + existing_team.team_alias = "team-a" + existing_team.model_dump.return_value = {"team_id": "team-a", "team_alias": "team-a"} + + updated_team = MagicMock() + updated_team.team_id = "team-a" + updated_team.access_group_ids = committed_team_groups + updated_team.model_dump.return_value = {"team_id": "team-a"} + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as prisma, + patch("litellm.proxy.proxy_server.llm_router"), + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.management_endpoints.team_endpoints._refresh_cached_team"), + patch( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + new_callable=AsyncMock, + ) as invalidate_cache, + ): + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) + prisma.db.tx = fake_db.tx + prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + + await update_team( + data=UpdateTeamRequest(team_id="team-a", access_group_ids=committed_team_groups), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + assert access_groups["ag-drop"]["assigned_team_ids"] == [] + assert access_groups["ag-add"]["assigned_team_ids"] == ["team-a"] + assert access_groups["ag-keep"]["assigned_team_ids"] == ["team-a"] + assert access_groups["ag-other-team"]["assigned_team_ids"] == ["team-b"] + + assert fake_db.transactions == [["lock", "read", "affected", "attach", "detach"]] + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-drop", "ag-keep", "ag-add"} + + async def _get_access_object(*, access_group_id, **_kwargs): + stored = access_groups[access_group_id] + return LiteLLM_AccessGroupTable( + access_group_id=access_group_id, + access_group_name=access_group_id, + access_model_names=list(stored["access_model_names"]), + assigned_team_ids=list(stored["assigned_team_ids"]), + assigned_key_ids=[], + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + side_effect=_get_access_object, + ), + ): + authorized_models = await get_authorized_resources_from_key_access_groups( + valid_token=UserAPIKeyAuth( + token="sk-hash", + models=[], + team_id="team-a", + access_group_ids=["ag-drop", "ag-keep", "ag-add"], + ), + team_object=LiteLLM_TeamTable(team_id="team-a", models=[]), + resource_field="access_model_names", + ) + + assert sorted(authorized_models) == ["added-model", "kept-model"] + + +@pytest.mark.asyncio +async def test_sync_reads_the_committed_team_row_rather_than_the_callers_snapshot(): + """ + The mirror takes no desired-state argument on purpose. It locks the team and reads + the row as committed, so two concurrent writers for one team converge on the row the + last one committed instead of each replaying its own stale snapshot. Reconciling also + means a retry heals a half-applied sync, where a before/after delta computes nothing. + + The same holds for the cache step: the groups to drop come from the reconciled set, + not from the rows this attempt happened to change, so a retry after an unreachable + cache still drops the entries even though its statements are now no-ops. + + A team with no row at all is deletion, and must detach from every group. + """ + from litellm.proxy.management_helpers.access_group_team_sync import ( + sync_team_access_group_membership, + ) + + access_groups = {"ag-1": ["team-a", "team-b"], "ag-2": ["team-a"], "ag-3": []} + teams = {"team-a": ["ag-2", "ag-3"]} + fake_db = _FakeMirrorDb(access_groups, teams, plain_lists=True) + prisma_client = SimpleNamespace(db=SimpleNamespace(tx=fake_db.tx)) + + with patch( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + new_callable=AsyncMock, + side_effect=[ConnectionError("redis unreachable"), None, None], + ) as invalidate_cache: + with pytest.raises(ConnectionError): + await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a") + assert access_groups == {"ag-1": ["team-b"], "ag-2": ["team-a"], "ag-3": ["team-a"]} + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-1", "ag-2", "ag-3"} + + invalidate_cache.reset_mock() + invalidate_cache.side_effect = None + await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a") + assert access_groups == {"ag-1": ["team-b"], "ag-2": ["team-a"], "ag-3": ["team-a"]} + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-2", "ag-3"} + + invalidate_cache.reset_mock() + del teams["team-a"] + await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a") + assert access_groups == {"ag-1": ["team-b"], "ag-2": [], "ag-3": []} + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-2", "ag-3"} + + assert fake_db.transactions == [["lock", "read", "affected", "attach", "detach"]] * 3 + + +@pytest.mark.asyncio +async def test_new_team_and_delete_team_both_drive_the_mirror(): + """Every writer of `team.access_group_ids` has to reach the mirror, not just update. + These pin the wiring on the other two paths; the mirror's own behavior is covered above. + + Creation has to insert the team row and mirror it in one transaction. With the mirror + in a transaction of its own, a sync that fails leaves a committed team whose groups + never learned about it, and the retry is rejected as a duplicate team id.""" + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import DeleteTeamRequest, NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import delete_team, new_team + + access_groups = {"ag-1": [], "ag-2": []} + fake_db = _FakeMirrorDb(access_groups, {}, plain_lists=True) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as prisma, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch("litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", new_callable=AsyncMock), + patch( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + new_callable=AsyncMock, + ) as invalidate_cache, + ): + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + prisma.db.tx = fake_db.tx + prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + prisma.get_data = AsyncMock(return_value=None) + + await new_team( + data=NewTeamRequest(team_id="team-new", team_alias="new", access_group_ids=["ag-1"]), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + assert access_groups == {"ag-1": ["team-new"], "ag-2": []} + assert fake_db.transactions == [["create", "lock", "read", "affected", "attach", "detach"]] + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-1"} + + team_row = LiteLLM_TeamTable(team_id="team-gone", models=[], access_group_ids=["ag-1"]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as prisma, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.management_endpoints.team_endpoints._persist_deleted_team_records", new_callable=AsyncMock), + patch("litellm.proxy.management_endpoints.team_endpoints._verify_team_access", new_callable=AsyncMock), + patch( + "litellm.proxy.management_endpoints.team_endpoints.sync_team_access_group_membership", + new_callable=AsyncMock, + ) as sync, + ): + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + prisma.delete_data = AsyncMock(return_value=[team_row]) + prisma.db.execute_raw = AsyncMock(return_value=0) + prisma.db.litellm_teammembership.delete_many = AsyncMock(return_value=0) + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-gone"]), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + assert sync.await_args_list[0].kwargs["team_id"] == "team-gone" + + +@pytest.mark.asyncio +async def test_invalidate_access_group_cache_deletes_the_cached_object(): + """The mirror's cache step is what stops a revoked group granting from cache until TTL, + so pin that it actually reaches the delete rather than only being called.""" + from litellm.proxy.management_helpers.access_group_team_sync import ( + invalidate_access_group_cache, + ) + + cache, logging_obj = MagicMock(), MagicMock() + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", logging_obj), + patch( + "litellm.proxy.management_helpers.access_group_team_sync._delete_cache_access_object", + new_callable=AsyncMock, + ) as delete_cached, + ): + await invalidate_access_group_cache("ag-1") + + assert delete_cached.await_args.kwargs == { + "access_group_id": "ag-1", + "user_api_key_cache": cache, + "proxy_logging_obj": logging_obj, + } 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 da805b864fc..b83b862d6b8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2,6 +2,7 @@ import asyncio import json import os import sys +from contextlib import asynccontextmanager from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -37,6 +38,20 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( ) +def _wire_team_create_tx(prisma_client): + """`/team/new` inserts the team and mirrors it onto the access groups in one transaction, + so a mocked client has to hand its team table back out of `db.tx()`.""" + + @asynccontextmanager + async def _tx(): + yield SimpleNamespace( + litellm_teamtable=prisma_client.db.litellm_teamtable, + query_raw=AsyncMock(return_value=[]), + ) + + prisma_client.db.tx = lambda *_args, **_kwargs: _tx() + + def test_microsoft_sso_handler_openid_from_response_user_principal_name(): # Arrange # Create a mock response similar to what Microsoft SSO would return @@ -577,6 +592,7 @@ async def test_default_team_params(team_params): mock_prisma = MagicMock() mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None) mock_prisma.db.litellm_teamtable.create = AsyncMock() + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) mock_prisma.jsonify_team_object = MagicMock(side_effect=mock_jsonify_team_object) @@ -624,6 +640,7 @@ async def test_default_team_params_organization_id_reaches_sso_created_team(team mock_prisma = MagicMock() mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None) mock_prisma.db.litellm_teamtable.create = AsyncMock() + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) @@ -671,6 +688,7 @@ async def test_create_team_without_default_params(): mock_prisma = MagicMock() mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None) mock_prisma.db.litellm_teamtable.create = AsyncMock() + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) mock_prisma.jsonify_team_object = MagicMock(side_effect=mock_jsonify_team_object) diff --git a/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py b/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py new file mode 100644 index 00000000000..eb11292cf42 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py @@ -0,0 +1,39 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.management_helpers.access_group_team_sync import ( + invalidate_access_group_caches, +) + + +@pytest.mark.asyncio +async def test_one_unreachable_cache_does_not_skip_the_other_groups(monkeypatch): + """ + `assigned_team_ids` is an authorization input, so a group whose cache still holds the + revoked grant keeps serving it until the entry is dropped. + + A sequential loop would stop at the first failing group and leave the groups behind it + serving stale grants, and swallowing the failure would report success to the admin for + a revoke that never took effect. Every group has to be attempted, and the endpoint has + to fail so the caller can retry. + """ + attempted: list[str] = [] + + async def _invalidate(access_group_id: str) -> None: + attempted.append(access_group_id) + if access_group_id == "ag-redis-down": + raise ConnectionError("redis unreachable") + + monkeypatch.setattr( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + _invalidate, + ) + + with pytest.raises(ConnectionError): + await invalidate_access_group_caches(("ag-redis-down", "ag-2", "ag-3")) + + assert attempted == ["ag-redis-down", "ag-2", "ag-3"] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py index 887aedaf0aa..6d7011fe10c 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py @@ -7,9 +7,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest -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 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.cohere_passthrough_logging_handler import ( @@ -69,12 +67,8 @@ class TestCoherePassthroughLoggingHandler: ) @patch("litellm.completion_cost") - @patch( - "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" - ) - @patch( - "litellm.llms.cohere.embed.v1_transformation.CohereEmbeddingConfig._transform_response" - ) + @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") + @patch("litellm.llms.cohere.embed.v1_transformation.CohereEmbeddingConfig._transform_response") def test_cohere_embed_passthrough_cost_tracking( self, mock_transform_response, mock_get_standard_logging, mock_completion_cost ): @@ -92,9 +86,7 @@ class TestCoherePassthroughLoggingHandler: mock_embedding_response.object = "list" from litellm.types.utils import Usage - mock_embedding_response.usage = Usage( - prompt_tokens=3, completion_tokens=0, total_tokens=3 - ) + mock_embedding_response.usage = Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3) mock_transform_response.return_value = mock_embedding_response mock_completion_cost.return_value = 3.6e-07 # Expected cost for embed-v4.0 @@ -151,6 +143,38 @@ class TestCoherePassthroughLoggingHandler: assert hasattr(result["result"], "model") assert result["result"].model == "embed-english-v3.0" + @patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler.BasePassthroughLoggingHandler.passthrough_chat_handler" + ) + @patch("litellm.completion_cost") + def test_openai_embeddings_route_does_not_use_cohere_embed_path(self, mock_completion_cost, mock_chat_handler): + mock_chat_handler.return_value = {"result": None, "kwargs": {}} + response_body = { + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}], + "usage": {"prompt_tokens": 6, "total_tokens": 6}, + } + result = self.handler.cohere_passthrough_handler( + httpx_response=self._create_mock_httpx_response(response_body), + response_body=response_body, + logging_obj=self._create_mock_logging_obj(), + url_route="https://api.openai.com/v1/embeddings", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": "text-embedding-3-small", "input": "PROOF_SENTINEL_TEXT"}, + passthrough_logging_payload=PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/embeddings", + request_body={"model": "text-embedding-3-small", "input": "PROOF_SENTINEL_TEXT"}, + request_method="POST", + ), + ) + mock_completion_cost.assert_not_called() + mock_chat_handler.assert_called_once() + assert result == {"result": None, "kwargs": {}} + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 05051ab3745..664015003e4 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -8,9 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -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 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( @@ -70,9 +68,7 @@ class TestOpenAIPassthroughLoggingHandler: mock_response.headers = {"content-type": "application/json"} return mock_response - def _create_passthrough_logging_payload( - self, user: str = "test_user" - ) -> PassthroughStandardLoggingPayload: + def _create_passthrough_logging_payload(self, user: str = "test_user") -> PassthroughStandardLoggingPayload: """Create a mock passthrough logging payload""" return PassthroughStandardLoggingPayload( url="https://api.openai.com/v1/chat/completions", @@ -113,9 +109,7 @@ class TestOpenAIPassthroughLoggingHandler: # Negative cases assert ( - OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route( - "https://api.openai.com/v1/models" - ) + OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/models") == False ) assert ( @@ -125,15 +119,10 @@ class TestOpenAIPassthroughLoggingHandler: == False ) assert ( - OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route( - "https://api.anthropic.com/v1/messages" - ) - == False - ) - assert ( - OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("") + OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.anthropic.com/v1/messages") == False ) + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("") == False def test_is_openai_image_generation_route(self): """Test OpenAI image generation route detection""" @@ -159,9 +148,7 @@ class TestOpenAIPassthroughLoggingHandler: == False ) assert ( - OpenAIPassthroughLoggingHandler.is_openai_image_generation_route( - "https://api.openai.com/v1/images/edits" - ) + OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://api.openai.com/v1/images/edits") == False ) assert ( @@ -170,32 +157,23 @@ class TestOpenAIPassthroughLoggingHandler: ) == False ) - assert ( - OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("") - == False - ) + assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("") == False def test_is_openai_image_editing_route(self): """Test OpenAI image editing route detection""" # Positive cases assert ( - OpenAIPassthroughLoggingHandler.is_openai_image_editing_route( - "https://api.openai.com/v1/images/edits" - ) + OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/images/edits") == True ) assert ( - OpenAIPassthroughLoggingHandler.is_openai_image_editing_route( - "https://openai.azure.com/v1/images/edits" - ) + OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://openai.azure.com/v1/images/edits") == True ) # Negative cases assert ( - OpenAIPassthroughLoggingHandler.is_openai_image_editing_route( - "https://api.openai.com/v1/chat/completions" - ) + OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/chat/completions") == False ) assert ( @@ -210,118 +188,91 @@ class TestOpenAIPassthroughLoggingHandler: ) == False ) - assert ( - OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("") == False - ) + assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("") == False def test_is_openai_responses_route(self): """Test OpenAI responses API route detection""" # Positive cases + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/responses") == True assert ( - OpenAIPassthroughLoggingHandler.is_openai_responses_route( - "https://api.openai.com/v1/responses" - ) - == True - ) - assert ( - OpenAIPassthroughLoggingHandler.is_openai_responses_route( - "https://openai.azure.com/v1/responses" - ) - == True - ) - assert ( - OpenAIPassthroughLoggingHandler.is_openai_responses_route( - "https://api.openai.com/responses" - ) - == True + OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://openai.azure.com/v1/responses") == True ) + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/responses") == True # Negative cases assert ( - OpenAIPassthroughLoggingHandler.is_openai_responses_route( - "https://api.openai.com/v1/chat/completions" - ) + OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/chat/completions") == False ) assert ( - OpenAIPassthroughLoggingHandler.is_openai_responses_route( - "https://api.openai.com/v1/images/generations" - ) + OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/images/generations") == False ) assert ( - OpenAIPassthroughLoggingHandler.is_openai_responses_route( - "http://localhost:4000/openai/v1/responses" - ) + OpenAIPassthroughLoggingHandler.is_openai_responses_route("http://localhost:4000/openai/v1/responses") == False ) assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("") == False + def test_is_openai_embeddings_route(self): + assert ( + OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("https://api.openai.com/v1/embeddings") is True + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("https://openai.azure.com/v1/embeddings") is True + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_embeddings_route( + "https://my-resource.cognitiveservices.azure.com/v1/embeddings" + ) + is True + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_embeddings_route( + "https://my-resource.openai.azure.com/openai/deployments/text-embedding-3-small/embeddings" + ) + is False + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("https://api.openai.com/v1/chat/completions") + is False + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_embeddings_route( + "http://localhost:4000/openai_passthrough/v1/embeddings" + ) + is False + ) + assert OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("") is False + def test_is_openai_route_recognizes_cognitiveservices_azure_com(self): """Azure OpenAI resources created via the newer "Azure AI Foundry" / Cognitive Services pathway live on `*.cognitiveservices.azure.com` - subdomains rather than the older `openai.azure.com`. All four + subdomains rather than the older `openai.azure.com`. The is_openai_*_route methods must recognize both Azure subdomains so cost tracking applies regardless of which Azure naming the user's resource happens to be on. """ - cognitive_chat = ( - "https://my-resource.cognitiveservices.azure.com/v1/chat/completions" - ) - cognitive_images_gen = ( - "https://my-resource.cognitiveservices.azure.com/v1/images/generations" - ) - cognitive_images_edit = ( - "https://my-resource.cognitiveservices.azure.com/v1/images/edits" - ) - cognitive_responses = ( - "https://my-resource.cognitiveservices.azure.com/v1/responses" - ) + cognitive_chat = "https://my-resource.cognitiveservices.azure.com/v1/chat/completions" + cognitive_images_gen = "https://my-resource.cognitiveservices.azure.com/v1/images/generations" + cognitive_images_edit = "https://my-resource.cognitiveservices.azure.com/v1/images/edits" + cognitive_responses = "https://my-resource.cognitiveservices.azure.com/v1/responses" + cognitive_embeddings = "https://my-resource.cognitiveservices.azure.com/v1/embeddings" - assert ( - OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route( - cognitive_chat - ) - is True - ) - assert ( - OpenAIPassthroughLoggingHandler.is_openai_image_generation_route( - cognitive_images_gen - ) - is True - ) - assert ( - OpenAIPassthroughLoggingHandler.is_openai_image_editing_route( - cognitive_images_edit - ) - is True - ) - assert ( - OpenAIPassthroughLoggingHandler.is_openai_responses_route( - cognitive_responses - ) - is True - ) + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(cognitive_chat) is True + assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(cognitive_images_gen) is True + assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(cognitive_images_edit) is True + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_responses) is True + assert OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(cognitive_embeddings) is True # Cross-route negatives still hold for cognitiveservices hosts. - assert ( - OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route( - cognitive_responses - ) - is False - ) - assert ( - OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_chat) - is False - ) + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(cognitive_responses) is False + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_chat) is False + assert OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(cognitive_chat) is False @patch("litellm.completion_cost") - @patch( - "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" - ) - def test_openai_passthrough_handler_success( - self, mock_get_standard_logging, mock_completion_cost - ): + @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") + def test_openai_passthrough_handler_success(self, mock_get_standard_logging, mock_completion_cost): """Test successful cost tracking for OpenAI chat completions""" # Arrange mock_completion_cost.return_value = 0.000045 @@ -370,9 +321,7 @@ class TestOpenAIPassthroughLoggingHandler: assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai" @patch("litellm.completion_cost") - def test_openai_passthrough_handler_non_chat_completions( - self, mock_completion_cost - ): + def test_openai_passthrough_handler_non_chat_completions(self, mock_completion_cost): """Test that non-chat-completions routes fall back to base handler""" # Arrange mock_httpx_response = self._create_mock_httpx_response() @@ -406,12 +355,8 @@ class TestOpenAIPassthroughLoggingHandler: # The important thing is that our specific OpenAI handler logic didn't run @patch("litellm.completion_cost") - @patch( - "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" - ) - def test_openai_passthrough_handler_with_user_tracking( - self, mock_get_standard_logging, mock_completion_cost - ): + @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") + def test_openai_passthrough_handler_with_user_tracking(self, mock_get_standard_logging, mock_completion_cost): """Test cost tracking with user information""" # Arrange mock_completion_cost.return_value = 0.000123 @@ -464,15 +409,10 @@ class TestOpenAIPassthroughLoggingHandler: assert "litellm_params" in result["kwargs"] assert "proxy_server_request" in result["kwargs"]["litellm_params"] assert "body" in result["kwargs"]["litellm_params"]["proxy_server_request"] - assert ( - result["kwargs"]["litellm_params"]["proxy_server_request"]["body"]["user"] - == "test_user_123" - ) + assert result["kwargs"]["litellm_params"]["proxy_server_request"]["body"]["user"] == "test_user_123" @patch("litellm.completion_cost") - def test_openai_passthrough_handler_cost_calculation_error( - self, mock_completion_cost - ): + def test_openai_passthrough_handler_cost_calculation_error(self, mock_completion_cost): """Test error handling in cost calculation""" # Arrange mock_completion_cost.side_effect = Exception("Cost calculation failed") @@ -521,9 +461,7 @@ class TestOpenAIPassthroughLoggingHandler: @patch(f"{OpenAIPassthroughLoggingHandler.__module__}.get_standard_logging_object_payload") @patch("litellm.completion_cost", return_value=3.3e-06) - def test_streaming_responses_cost_uses_completed_response( - self, mock_completion_cost, mock_get_standard_logging - ): + def test_streaming_responses_cost_uses_completed_response(self, mock_completion_cost, mock_get_standard_logging): response_id = "resp_PROOFSENTINEL0123456789abcdef" completed_event = { "type": "response.completed", @@ -796,12 +734,8 @@ class TestOpenAIPassthroughLoggingHandler: mock_completion_cost.assert_not_called() @patch("litellm.completion_cost") - @patch( - "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" - ) - def test_different_models_cost_tracking( - self, mock_get_standard_logging, mock_completion_cost - ): + @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") + def test_different_models_cost_tracking(self, mock_get_standard_logging, mock_completion_cost): """Test cost tracking for different OpenAI models""" # Arrange mock_get_standard_logging.return_value = {"test": "logging_payload"} @@ -868,12 +802,8 @@ class TestOpenAIPassthroughLoggingHandler: assert handler.get_provider_config("gpt-4o") is not None @patch("litellm.completion_cost") - @patch( - "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" - ) - def test_azure_passthrough_tags_metadata_model_provider( - self, mock_get_standard_logging, mock_completion_cost - ): + @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") + def test_azure_passthrough_tags_metadata_model_provider(self, mock_get_standard_logging, mock_completion_cost): """Test that tags, metadata, model, and custom_llm_provider are preserved for Azure passthrough in UI""" # Arrange mock_completion_cost.return_value = 0.000045 @@ -929,9 +859,7 @@ class TestOpenAIPassthroughLoggingHandler: # Verify model and custom_llm_provider are set correctly assert result["kwargs"]["model"] == "gpt-4o" - assert ( - result["kwargs"]["custom_llm_provider"] == "azure" - ) # Should preserve Azure, not default to "openai" + assert result["kwargs"]["custom_llm_provider"] == "azure" # Should preserve Azure, not default to "openai" assert result["kwargs"]["response_cost"] == 0.000045 # Verify metadata tags are preserved in litellm_params @@ -955,12 +883,8 @@ class TestOpenAIPassthroughLoggingHandler: assert call_args[1]["custom_llm_provider"] == "azure" @patch("litellm.completion_cost") - @patch( - "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" - ) - @patch( - "litellm.llms.openai.responses.transformation.OpenAIResponsesAPIConfig.transform_response_api_response" - ) + @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") + @patch("litellm.llms.openai.responses.transformation.OpenAIResponsesAPIConfig.transform_response_api_response") def test_responses_api_cost_tracking( self, mock_transform_responses, @@ -1052,9 +976,7 @@ class TestOpenAIPassthroughLoggingHandler: assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai" @patch("litellm.completion_cost") - @patch( - "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" - ) + @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") def test_responses_api_uses_responses_transformer_not_chat_completions( self, mock_get_standard_logging, mock_completion_cost ): @@ -1185,9 +1107,7 @@ class TestOpenAIPassthroughIntegration: mock_response.headers = {"content-type": "application/json"} return mock_response - def _create_passthrough_logging_payload( - self, user: str = "test_user" - ) -> PassthroughStandardLoggingPayload: + def _create_passthrough_logging_payload(self, user: str = "test_user") -> PassthroughStandardLoggingPayload: """Create a mock passthrough logging payload""" return PassthroughStandardLoggingPayload( url="https://api.openai.com/v1/chat/completions", @@ -1201,59 +1121,32 @@ class TestOpenAIPassthroughIntegration: def test_is_openai_route_detection(self): """Test OpenAI route detection in the main success handler""" # Positive cases - assert ( - self.handler.is_openai_route("https://api.openai.com/v1/chat/completions") - == True - ) - assert ( - self.handler.is_openai_route("https://openai.azure.com/v1/chat/completions") - == True - ) + assert self.handler.is_openai_route("https://api.openai.com/v1/chat/completions") == True + assert self.handler.is_openai_route("https://openai.azure.com/v1/chat/completions") == True assert self.handler.is_openai_route("https://api.openai.com/v1/models") == True # Azure OpenAI on the shared Cognitive Services domain, identified by an # OpenAI-style path segment. assert ( - self.handler.is_openai_route( - "https://my-resource.cognitiveservices.azure.com/v1/chat/completions" - ) - == True + self.handler.is_openai_route("https://my-resource.cognitiveservices.azure.com/v1/chat/completions") == True ) # Negative cases - assert ( - self.handler.is_openai_route( - "http://localhost:4000/openai/v1/chat/completions" - ) - == False - ) - assert ( - self.handler.is_openai_route("https://api.anthropic.com/v1/messages") - == False - ) - assert ( - self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript") - == False - ) + assert self.handler.is_openai_route("http://localhost:4000/openai/v1/chat/completions") == False + assert self.handler.is_openai_route("https://api.anthropic.com/v1/messages") == False + assert self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript") == False # Non-OpenAI Azure Cognitive Services share the `cognitiveservices.azure.com` # domain but must NOT be classified as OpenAI routes (no OpenAI path segment). assert ( - self.handler.is_openai_route( - "https://my-resource.cognitiveservices.azure.com/speechtotext/v3.1/recognize" - ) + self.handler.is_openai_route("https://my-resource.cognitiveservices.azure.com/speechtotext/v3.1/recognize") == False ) assert ( - self.handler.is_openai_route( - "https://my-resource.cognitiveservices.azure.com/vision/v3.2/analyze" - ) - == False + self.handler.is_openai_route("https://my-resource.cognitiveservices.azure.com/vision/v3.2/analyze") == False ) # A look-alike domain that merely contains an OpenAI host as a substring # must be rejected by the suffix-based hostname match. assert ( - self.handler.is_openai_route( - "https://cognitiveservices.azure.com.attacker.example/v1/chat/completions" - ) + self.handler.is_openai_route("https://cognitiveservices.azure.com.attacker.example/v1/chat/completions") == False ) assert self.handler.is_openai_route("") == False @@ -1274,52 +1167,188 @@ class TestOpenAIPassthroughIntegration: remove Responses from the OR-chain without a test failure. """ # Responses must be supported on api.openai.com and openai.azure.com. - assert ( - self.handler._is_supported_openai_endpoint( - "https://api.openai.com/v1/responses" - ) - is True - ) - assert ( - self.handler._is_supported_openai_endpoint( - "https://openai.azure.com/v1/responses" - ) - is True - ) + assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/responses") is True + assert self.handler._is_supported_openai_endpoint("https://openai.azure.com/v1/responses") is True # The other supported endpoints stay supported (no regression). - assert ( - self.handler._is_supported_openai_endpoint( - "https://api.openai.com/v1/chat/completions" - ) - is True - ) - assert ( - self.handler._is_supported_openai_endpoint( - "https://api.openai.com/v1/images/generations" - ) - is True - ) - assert ( - self.handler._is_supported_openai_endpoint( - "https://api.openai.com/v1/images/edits" - ) - is True - ) + assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/chat/completions") is True + assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/images/generations") is True + assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/images/edits") is True # Unsupported OpenAI endpoints (e.g. /v1/models) still return False. + assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/models") is False assert ( self.handler._is_supported_openai_endpoint( - "https://api.openai.com/v1/models" + "https://my-resource.openai.azure.com/openai/deployments/text-embedding-3-small/embeddings" ) is False ) + def test_is_supported_openai_endpoint_includes_embeddings(self): + assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/embeddings") is True + assert self.handler._is_supported_openai_endpoint("https://openai.azure.com/v1/embeddings") is True + + def test_is_cohere_route_does_not_match_openai_embeddings(self): + assert self.handler.is_cohere_route("https://api.cohere.com/v1/embed") is True + assert self.handler.is_cohere_route("https://api.cohere.com/v2/chat") is True + assert self.handler.is_cohere_route("https://api.openai.com/v1/embeddings") is False + assert self.handler.is_cohere_route("https://api.cohere.com/v1/rerank") is False + assert self.handler.is_cohere_route("http://localhost:4000/openai_passthrough/v1/embeddings") is False + + @patch("litellm.completion_cost") + @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") + def test_openai_passthrough_handler_embeddings_sets_response_cost( + self, mock_get_standard_logging, mock_completion_cost + ): + mock_completion_cost.return_value = 2.8e-07 + mock_get_standard_logging.return_value = {"test": "logging_payload"} + + response_body = { + "object": "list", + "model": "text-embedding-3-small", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": [0.1, 0.2], + } + ], + "usage": {"prompt_tokens": 14, "total_tokens": 14}, + } + mock_httpx_response = self._create_mock_httpx_response(response_body) + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/embeddings", + request_body={ + "model": "text-embedding-3-small", + "input": "PROOF_SENTINEL_TEXT", + }, + request_method="POST", + ) + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "litellm_params": {}, + } + + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=response_body, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/embeddings", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={ + "model": "text-embedding-3-small", + "input": "PROOF_SENTINEL_TEXT", + }, + **kwargs, + ) + + assert result["result"] is not None + assert result["kwargs"]["response_cost"] == 2.8e-07 + assert result["kwargs"]["model"] == "text-embedding-3-small" + assert result["kwargs"]["custom_llm_provider"] == "openai" + assert result["result"]._hidden_params["response_cost"] == 2.8e-07 + mock_completion_cost.assert_called_once() + assert mock_completion_cost.call_args.kwargs["call_type"] == "aembedding" + assert mock_logging_obj.model_call_details["response_cost"] == 2.8e-07 + + @patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.passthrough_chat_handler" + ) + @patch("litellm.completion_cost") + def test_openai_passthrough_handler_embeddings_without_model_falls_back( + self, mock_completion_cost, mock_chat_handler + ): + mock_chat_handler.return_value = {"result": None, "kwargs": {}} + response_body = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=self._create_mock_httpx_response(response_body), + response_body=response_body, + logging_obj=self._create_mock_logging_obj(), + url_route="https://api.openai.com/v1/embeddings", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"input": "PROOF_SENTINEL_TEXT"}, + passthrough_logging_payload=PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/embeddings", + request_body={"input": "PROOF_SENTINEL_TEXT"}, + request_method="POST", + ), + ) + mock_completion_cost.assert_not_called() + mock_chat_handler.assert_called_once() + assert result == {"result": None, "kwargs": {}} + @patch( "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler" ) @pytest.mark.asyncio - async def test_success_handler_dispatches_responses_api_to_openai_handler( - self, mock_openai_handler - ): + async def test_success_handler_dispatches_embeddings_to_openai_handler(self, mock_openai_handler): + mock_openai_handler.return_value = { + "result": {"object": "list"}, + "kwargs": { + "response_cost": 2.8e-07, + "model": "text-embedding-3-small", + "custom_llm_provider": "openai", + }, + } + + mock_httpx_response = MagicMock(spec=httpx.Response) + mock_httpx_response.text = ( + '{"object":"list","model":"text-embedding-3-small",' + '"data":[{"object":"embedding","index":0,"embedding":[0.1]}],' + '"usage":{"prompt_tokens":14,"total_tokens":14}}' + ) + + mock_logging_obj = AsyncMock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.async_success_handler = AsyncMock() + + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/embeddings", + request_body={ + "model": "text-embedding-3-small", + "input": "PROOF_SENTINEL_TEXT", + }, + request_method="POST", + ) + + await self.handler.pass_through_async_success_handler( + httpx_response=mock_httpx_response, + response_body={ + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}], + "usage": {"prompt_tokens": 14, "total_tokens": 14}, + }, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/embeddings", + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={ + "model": "text-embedding-3-small", + "input": "PROOF_SENTINEL_TEXT", + }, + passthrough_logging_payload=passthrough_payload, + ) + + mock_openai_handler.assert_called_once() + assert mock_openai_handler.call_args.kwargs["url_route"] == "https://api.openai.com/v1/embeddings" + + @patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler" + ) + @pytest.mark.asyncio + async def test_success_handler_dispatches_responses_api_to_openai_handler(self, mock_openai_handler): """End-to-end dispatch test for the Responses API path. Pre-fix: `_is_supported_openai_endpoint` returned False for @@ -1395,9 +1424,7 @@ class TestOpenAIPassthroughIntegration: } mock_httpx_response = MagicMock(spec=httpx.Response) - mock_httpx_response.text = ( - '{"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}' - ) + mock_httpx_response.text = '{"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}' mock_logging_obj = AsyncMock() mock_logging_obj.model_call_details = {} @@ -1590,14 +1617,10 @@ class TestOpenAIPassthroughIntegration: # Test the _response_cost_calculator method calculated_cost = logging_obj._response_cost_calculator(result=image_response) - assert ( - calculated_cost == test_cost - ), f"Expected {test_cost}, got {calculated_cost}" + assert calculated_cost == test_cost, f"Expected {test_cost}, got {calculated_cost}" @patch("litellm.cost_calculator.default_image_cost_calculator") - def test_openai_passthrough_handler_image_generation( - self, mock_image_cost_calculator - ): + def test_openai_passthrough_handler_image_generation(self, mock_image_cost_calculator): """Test successful cost tracking for OpenAI image generation""" # Arrange mock_image_cost_calculator.return_value = 0.040 diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_models.py b/tests/test_litellm/proxy/proxy_server/test_routes_models.py index 381835fbc14..f18c5998b8c 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_models.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_models.py @@ -99,6 +99,62 @@ def test_get_models_happy_path(client, auth_as, patched_models, path): } +@pytest.mark.parametrize("path", ["/v1/models", "/models"]) +def test_get_models_anthropic_format_when_header_present( + client, auth_as, patched_models, path +): + """Pins: ``GET /v1/models`` returns the Anthropic-native models shape when + the caller sends an ``anthropic-version`` header (Claude Code gateway + discovery), while the default OpenAI shape is unchanged without it.""" + with auth_as(): + response = client.get(path, headers={"anthropic-version": "2023-06-01"}) + assert response.status_code == 200 + body = response.json() + assert "object" not in body + assert body["has_more"] is False + assert body["first_id"] == "gpt-4" + assert body["last_id"] == "claude-sonnet" + assert [m["id"] for m in body["data"]] == ["gpt-4", "claude-sonnet"] + for entry in body["data"]: + assert entry["type"] == "model" + assert entry["display_name"] == entry["id"] + assert entry["created_at"].endswith("Z") + + +@pytest.mark.parametrize("path", ["/v1/models", "/models"]) +def test_anthropic_format_exposes_token_limits( + client, auth_as, patched_models, monkeypatch, path +): + """Claude Code sizes requests off the listing, so the Anthropic-native entries + carry the same token limits the OpenAI listing resolves, with the output budget + named max_tokens as the Messages API names it.""" + from litellm.proxy import utils as proxy_utils + + def _create_model_info_response(model_id, provider="openai", **kwargs): + if model_id != "claude-sonnet": + return _stub_model_info_response(model_id=model_id, provider=provider) + return { + **_stub_model_info_response(model_id=model_id, provider=provider), + "max_input_tokens": 200000, + "max_output_tokens": 64000, + } + + monkeypatch.setattr( + proxy_utils, "create_model_info_response", _create_model_info_response + ) + + with auth_as(): + response = client.get(path, headers={"anthropic-version": "2023-06-01"}) + + assert response.status_code == 200 + gpt_4, claude = response.json()["data"] + assert claude["max_input_tokens"] == 200000 + assert claude["max_tokens"] == 64000 + assert "max_output_tokens" not in claude + assert "max_input_tokens" not in gpt_4 + assert "max_tokens" not in gpt_4 + + @pytest.mark.parametrize("path", ["/v1/models", "/models"]) def test_get_models_invalid_scope_returns_400(client, auth_as, patched_models, path): """Pins: ``GET /v1/models``, ``GET /models`` (error path: invalid scope).""" @@ -130,3 +186,50 @@ def test_get_model_by_id_not_found(client, auth_as, patched_models, path): response = client.get(path) assert response.status_code == 404 assert "not found" in response.text.lower() + + +@pytest.mark.parametrize("params", [{}, {"scope": "expand"}]) +def test_anthropic_format_returns_public_team_model_name( + client, auth_as, patched_models, monkeypatch, params +): + """Regression: the Anthropic-native listing must go through the same team + name translation as the OpenAI listing, so a caller never sees the internal + ``model_name_{team_id}_{uuid}`` routing key.""" + from litellm.proxy import utils as proxy_utils + from litellm.proxy.auth import model_checks + + internal_name = "model_name_team-1_c0ffee" + + patched_models.get_model_list = MagicMock( + return_value=[ + { + "model_name": internal_name, + "model_info": { + "team_id": "team-1", + "team_public_model_name": "gpt-4-team", + }, + } + ] + ) + patched_models.get_model_names = MagicMock(return_value=[internal_name]) + + async def _fake_get_available_models_for_user(**kwargs): + return [internal_name] + + monkeypatch.setattr( + proxy_utils, + "get_available_models_for_user", + _fake_get_available_models_for_user, + ) + monkeypatch.setattr( + model_checks, "get_complete_model_list", lambda **kwargs: [internal_name] + ) + + with auth_as(): + response = client.get( + "/v1/models", params=params, headers={"anthropic-version": "2023-06-01"} + ) + + assert response.status_code == 200 + assert [m["id"] for m in response.json()["data"]] == ["gpt-4-team"] + assert internal_name not in response.text diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 88dc07e741b..e99bdfb5c35 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -243,6 +243,34 @@ def test_bedrock_mantle_provider_fields(): assert fields_by_key["api_base"]["field_type"] == "text" +def test_nvidia_riva_provider_fields(): + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + response = test_client.get("/public/providers/fields") + assert response.status_code == 200 + providers = response.json() + + riva = next((p for p in providers if p["provider"] == "NVIDIA_RIVA"), None) + assert riva is not None, "NVIDIA Riva provider entry not found" + + assert riva["provider_display_name"] == "Nvidia Riva" + assert riva["litellm_provider"] == LlmProviders.NVIDIA_RIVA.value + assert riva["default_model_placeholder"].startswith("nvidia_riva/") + + fields_by_key = {f["key"]: f for f in riva["credential_fields"]} + + assert fields_by_key["api_base"]["required"] is True + assert fields_by_key["api_base"]["field_type"] == "text" + + assert fields_by_key["api_key"]["required"] is False + assert fields_by_key["api_key"]["field_type"] == "password" + + assert "nvcf_function_id" in fields_by_key + assert fields_by_key["nvcf_function_id"]["required"] is False + + def test_google_ai_studio_provider_fields_expose_api_base(): """The Google AI Studio (gemini) credential form must let admins set a custom api_base so they can point at a Gemini-compatible gateway (e.g. a self-hosted diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py index d17f6293cc3..736fc13d137 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py @@ -657,6 +657,87 @@ async def test_scheduled_rollup_stays_quiet_when_every_charge_landed(): alert.assert_not_awaited() +@pytest.mark.asyncio +async def test_scheduled_rollup_alerts_once_a_ptu_window_has_closed(): + """Reserved capacity is billed until the deployment is deleted, so a closed window stops + the attribution without stopping the charge. Nobody notices unless it is escalated.""" + ptu = { + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "team_id": "t", + "ptu_effective_from": "2020-01-01T00:00:00Z", + "ptu_effective_to": "2020-02-01T00:00:00Z", + } + prisma, _ = _prisma_with_models([_model_row(model_id="dep-lapsed", model_info=ptu)]) + alert = AsyncMock() + + result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert) + + assert result.lapsed == ("gpt-4o-mini-ptu",) + alert.assert_awaited_once() + message = alert.await_args.args[0] + assert "window has closed" in message + assert "gpt-4o-mini-ptu" in message + + +@pytest.mark.asyncio +async def test_a_model_name_cannot_smuggle_slack_markup_into_the_alert(): + """The alert lands in an operator channel and a model name is operator-supplied, so an + unescaped name could post a channel-wide mention.""" + ptu = { + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "team_id": "t", + "ptu_effective_from": "2020-01-01T00:00:00Z", + "ptu_effective_to": "2020-02-01T00:00:00Z", + } + row = _model_row(model_id="dep-x", model_name=" & ", model_info=ptu) + prisma, _ = _prisma_with_models([row]) + alert = AsyncMock() + + await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert) + + message = alert.await_args.args[0] + assert "" not in message + assert "<!channel>" in message + + +@pytest.mark.asyncio +async def test_an_open_ptu_window_raises_no_lapsed_alert(): + ptu = { + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "team_id": "t", + "ptu_effective_from": "2020-01-01T00:00:00Z", + "ptu_effective_to": "2999-01-01T00:00:00Z", + } + prisma, _ = _prisma_with_models([_model_row(model_id="dep-open", model_info=ptu)]) + alert = AsyncMock() + + result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert) + + assert result.lapsed == () + alert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_an_open_ended_ptu_window_raises_no_lapsed_alert(): + """No end bound means the operator never asked the attribution to stop.""" + ptu = { + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "team_id": "t", + "ptu_effective_from": "2020-01-01T00:00:00Z", + } + prisma, _ = _prisma_with_models([_model_row(model_id="dep-forever", model_info=ptu)]) + alert = AsyncMock() + + result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert) + + assert result.lapsed == () + alert.assert_not_awaited() + + @pytest.mark.asyncio async def test_a_broken_alert_channel_does_not_fail_the_rollup(): rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})] diff --git a/tests/test_litellm/proxy/test_conftest.py b/tests/test_litellm/proxy/test_conftest.py new file mode 100644 index 00000000000..6df692a67c9 --- /dev/null +++ b/tests/test_litellm/proxy/test_conftest.py @@ -0,0 +1,31 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +@pytest.fixture +def fixture_planted_prisma_mock(): + with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()): + yield + + +def test_monkeypatch_over_fixture_patched_prisma_client( + fixture_planted_prisma_mock, monkeypatch +): + """ + Mirrors the flake in test_team_endpoints.py: an autouse fixture patches + prisma_client, the test monkeypatches the same global, and monkeypatch + records the fixture's MagicMock as the value to restore. Its undo runs + after every other finalizer, so without hook-level isolation the mock + leaks and every later no-database test on the worker fails awaiting it. + """ + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock()) + assert isinstance(proxy_server.prisma_client, AsyncMock) + + +def test_prisma_client_did_not_leak_from_previous_test(): + import litellm.proxy.proxy_server as proxy_server + + assert not isinstance(proxy_server.prisma_client, MagicMock) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 57b2c874962..918d39646b0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -6872,6 +6872,91 @@ async def test_update_general_settings_propagates_apply_user_budget_to_team_keys assert ps.general_settings["apply_user_budget_to_team_keys"] is True +@pytest.mark.asyncio +async def test_update_general_settings_propagates_spend_log_cleanup_bounds(): + """The dashboard writes the cleanup bounds straight to the DB config, so + without runtime propagation the scheduled job never sees them and the knobs + do nothing until the process restarts.""" + from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( + SPEND_LOG_CLEANUP_BOUND_SETTINGS, + ) + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + db_settings = { + "maximum_spend_logs_cleanup_batch_size": 2000, + "maximum_spend_logs_cleanup_max_batches": 250, + "maximum_spend_logs_cleanup_run_budget": "90s", + "maximum_spend_logs_cleanup_batch_timeout": "10s", + } + assert set(db_settings) == set(SPEND_LOG_CLEANUP_BOUND_SETTINGS) + + with patch("litellm.proxy.proxy_server.general_settings", {}): + await proxy_config._update_general_settings(db_general_settings=db_settings) + + import litellm.proxy.proxy_server as ps + + assert {key: ps.general_settings.get(key) for key in db_settings} == db_settings + + +@pytest.mark.asyncio +async def test_update_general_settings_clears_a_spend_log_cleanup_bound_dropped_from_the_db(): + """Blanking the field in the dashboard deletes the key outright, so leaving + the last value in memory would keep a bound the operator just removed.""" + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + with patch( + "litellm.proxy.proxy_server.general_settings", + {"maximum_spend_logs_cleanup_run_budget": "90s", "maximum_spend_logs_cleanup_batch_timeout": "10s"}, + ): + await proxy_config._update_general_settings( + db_general_settings={"maximum_spend_logs_cleanup_batch_timeout": "10s"} + ) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["maximum_spend_logs_cleanup_run_budget"] is None + assert ps.general_settings["maximum_spend_logs_cleanup_batch_timeout"] == "10s" + + +@pytest.mark.asyncio +async def test_update_general_settings_keeps_a_yaml_set_spend_log_cleanup_bound(): + """A YAML-set bound never appears in the DB object, so treating its absence + as a dashboard clear would discard the deployed config on every reload.""" + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._yaml_spend_log_cleanup_bounds = {"maximum_spend_logs_cleanup_run_budget": "90s"} + + with patch("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_cleanup_run_budget": "90s"}): + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": True}) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["maximum_spend_logs_cleanup_run_budget"] == "90s" + + +@pytest.mark.asyncio +async def test_update_general_settings_clearing_a_db_override_falls_back_to_the_yaml_bound(): + """Clearing a dashboard override of a YAML-declared bound must restore the + YAML value. Leaving the deleted override in memory would keep enforcing the + bound the operator just removed, until the process restarted.""" + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._yaml_spend_log_cleanup_bounds = {"maximum_spend_logs_cleanup_run_budget": "90s"} + + # Memory currently holds the dashboard override, and the DB no longer carries it. + with patch("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_cleanup_run_budget": "30s"}): + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": True}) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["maximum_spend_logs_cleanup_run_budget"] == "90s" + + @pytest.mark.asyncio async def test_update_general_settings_apply_user_budget_to_team_keys_yaml_wins(): """A DB value must not silently override an explicit YAML setting on reload.""" diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 03eef14dacb..87fbdd4c933 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -2,12 +2,66 @@ Test cases for spend log cleanup functionality """ +import asyncio +import math +import time +from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock import pytest -from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup +from litellm.constants import ( + SPEND_LOG_CLEANUP_BATCH_SIZE, + SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP, + SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS, +) +from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( + SPEND_LOG_CLEANUP_BOUND_SETTINGS, + SpendLogCleanup, + TableCleanupResult, +) +from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import ( + SpendLogCleanupMetrics, +) + + +def _far_deadline() -> float: + """A run deadline far enough out that only the other bounds can stop a batch loop.""" + return time.monotonic() + 3600 + + +def _wire_tx(db): + """ + Model the prisma seam the cleanup job actually uses. + + Every statement the job issues runs inside db.tx() so it can carry a SET + LOCAL statement_timeout. Batch and probe statements are forwarded to + db.execute_raw and db.query_raw, which is what tests configure and assert + on, while the SET LOCAL statements are answered here so they neither consume + a side_effect entry nor show up in the recorded call list. Lookup is + deferred to call time so this can be wired before a test assigns its own + execute_raw. + """ + + @asynccontextmanager + async def _tx(): + tx = MagicMock() + + async def _execute_raw(sql, *args): + if sql.lstrip().upper().startswith("SET LOCAL"): + return 0 + return await db.execute_raw(sql, *args) + + async def _query_raw(sql, *args): + return await db.query_raw(sql, *args) + + tx.execute_raw = _execute_raw + tx.query_raw = _query_raw + yield tx + + db.tx = _tx + db.query_raw = AsyncMock(return_value=[{"remaining": 0}]) def test_spend_log_cleanup_cron_scheduling(): @@ -49,6 +103,7 @@ def test_spend_log_cleanup_cron_scheduler_integration(): # Mock scheduler mock_scheduler = MagicMock() mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_cleanup_instance = MagicMock() # Test Case 1: Cron-based scheduling @@ -155,7 +210,9 @@ async def test_cleanup_old_spend_logs_batch_deletion(): # Setup Prisma client mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_db = MagicMock() + _wire_tx(mock_db) # Mock execute_raw to return deleted counts (3 spend-log batches, then the # tool-index cleanup's first batch returning 0) @@ -207,7 +264,9 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): """ # Setup Prisma client mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_db = MagicMock() + _wire_tx(mock_db) mock_db.execute_raw = AsyncMock(return_value=0) mock_prisma_client.db = mock_db @@ -244,6 +303,7 @@ async def test_cleanup_drops_partitions_when_enabled_and_partitioned(): from unittest.mock import AsyncMock, MagicMock mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_prisma_client.db.execute_raw = AsyncMock(return_value=0) partition_manager = MagicMock() @@ -285,6 +345,7 @@ async def test_cleanup_uses_delete_when_partitioning_not_enabled(): from unittest.mock import AsyncMock, MagicMock mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0]) partition_manager = MagicMock() @@ -316,6 +377,7 @@ async def test_cleanup_uses_delete_when_not_partitioned(): from unittest.mock import AsyncMock, MagicMock mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0]) partition_manager = MagicMock() @@ -346,6 +408,7 @@ async def test_cleanup_old_spend_logs_no_retention_period(): Test that no logs are deleted when no retention period is set """ mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_prisma_client.db.execute_raw = AsyncMock() cleaner = SpendLogCleanup(general_settings={}) # no retention @@ -361,6 +424,7 @@ async def test_lock_not_released_when_not_acquired(): before the lock is ever acquired. """ mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_prisma_client.db.execute_raw = AsyncMock() mock_redis_cache = MagicMock() @@ -418,7 +482,9 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return(): """should abort deletion loop immediately when execute_raw returns a non-int (e.g. None or dict), preventing an infinite loop.""" mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_db = MagicMock() + _wire_tx(mock_db) mock_db.execute_raw = AsyncMock(return_value=None) mock_prisma_client.db = mock_db @@ -427,17 +493,19 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return(): ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) - total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date) + result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) assert mock_db.execute_raw.call_count == 1 - assert total_deleted == 0 + assert result.rows_deleted == 0 @pytest.mark.asyncio async def test_delete_old_logs_continues_on_valid_int_return(): """should continue deletion loop across batches when execute_raw returns valid int counts.""" mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_db = MagicMock() + _wire_tx(mock_db) mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0]) mock_prisma_client.db = mock_db @@ -446,35 +514,37 @@ async def test_delete_old_logs_continues_on_valid_int_return(): ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) - total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date) + result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) assert mock_db.execute_raw.call_count == 3 - assert total_deleted == 800 + assert result.rows_deleted == 800 @pytest.mark.asyncio -async def test_delete_old_rows_stops_at_max_batches(monkeypatch): - """The run-loop backstop must halt a cleanup that keeps finding rows, so a - huge backlog is spread across scheduled runs instead of one unbounded loop.""" - import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - - monkeypatch.setattr(cleanup_module, "SPEND_LOG_RUN_LOOPS", 2) - +async def test_delete_old_rows_stops_at_max_batches(): + """The batch cap must halt a cleanup that keeps finding rows, so a huge + backlog is spread across scheduled runs instead of one unbounded loop, and + the operator-facing knob must mean exactly the number of statements it names.""" mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_db = MagicMock() + _wire_tx(mock_db) mock_db.execute_raw = AsyncMock(return_value=1000) mock_prisma_client.db = mock_db cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "maximum_spend_logs_cleanup_max_batches": 2, + } ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) - total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date) + result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) - # run_count exceeds the cap only after 3 full batches (0, 1, 2) - assert mock_db.execute_raw.call_count == 3 - assert total_deleted == 3000 + assert mock_db.execute_raw.call_count == 2 + assert result.rows_deleted == 2000 + assert result.stop_reason == "batch_cap_reached" @pytest.mark.asyncio @@ -482,7 +552,9 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key(): """Tool index rows are derived from spend logs and expire on the same cutoff; the delete must match on the table's composite primary key.""" mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_db = MagicMock() + _wire_tx(mock_db) mock_db.execute_raw = AsyncMock(side_effect=[5, 0]) mock_prisma_client.db = mock_db @@ -491,9 +563,9 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key(): ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) - total_deleted = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date) + result = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date, _far_deadline()) - assert total_deleted == 5 + assert result.rows_deleted == 5 delete_sql = mock_db.execute_raw.call_args_list[0][0][0] assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in delete_sql assert 'WHERE ("request_id", "tool_name") IN' in delete_sql @@ -513,7 +585,9 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) ) mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_db = MagicMock() + _wire_tx(mock_db) # batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed, # batch 5 returns 0 → loop exits naturally. mock_db.execute_raw = AsyncMock( @@ -526,11 +600,11 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) - total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date) + result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) # All 5 batches should have been attempted; 100 + 200 + 50 = 350 deleted. assert mock_db.execute_raw.call_count == 5 - assert total_deleted == 350 + assert result.rows_deleted == 350 @pytest.mark.asyncio @@ -548,7 +622,9 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch): ) mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_db = MagicMock() + _wire_tx(mock_db) # Every batch raises — must abort after exactly 3 attempts, not loop forever. mock_db.execute_raw = AsyncMock( side_effect=ConnectionError("simulated persistent DB outage") @@ -560,10 +636,10 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch): ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) - total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date) + result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) assert mock_db.execute_raw.call_count == 3 - assert total_deleted == 0 + assert result.rows_deleted == 0 @pytest.mark.asyncio @@ -580,7 +656,9 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc ) mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_db = MagicMock() + _wire_tx(mock_db) # Pattern: fail, fail, success (resets counter), fail, fail, success, done. # Without reset, three of these would trip abort; with reset, they don't. mock_db.execute_raw = AsyncMock( @@ -601,10 +679,10 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) - total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date) + result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) assert mock_db.execute_raw.call_count == 7 - assert total_deleted == 150 + assert result.rows_deleted == 150 @pytest.mark.asyncio @@ -617,6 +695,7 @@ async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch): monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger) mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) # Force the outer try/except to fire by making _should_delete_spend_logs raise. cleaner = cleanup_module.SpendLogCleanup( general_settings={"maximum_spend_logs_retention_period": "7d"} @@ -653,7 +732,9 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch ) mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) mock_db = MagicMock() + _wire_tx(mock_db) mock_db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down")) mock_prisma_client.db = mock_db @@ -698,6 +779,7 @@ def _mock_prisma_for_retention(side_effect: list) -> "MagicMock": from unittest.mock import AsyncMock, MagicMock client = MagicMock() + _wire_tx(client.db) client.db.execute_raw = AsyncMock(side_effect=side_effect) return client @@ -753,3 +835,536 @@ async def test_no_retention_keys_means_no_cleanup_at_all(): cleaner.pod_lock_manager = None await cleaner.cleanup_old_spend_logs(client) assert client.db.execute_raw.await_count == 0 + + +@pytest.mark.asyncio +async def test_run_budget_stops_the_loop_and_leaves_the_backlog_for_the_next_run(): + """ + The wall-clock budget is the bound that keeps a large backlog from turning + into one multi-hour run. With rows always available, the loop must stop on + the deadline rather than on the batch cap, and must report that reason so + operators can tell a budgeted stop from a drained table. + """ + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock(return_value=1000) + mock_prisma_client.db = mock_db + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + # Comfortably more batches than a sub-second budget can reach (each + # batch sleeps 0.1s), but small enough that a broken deadline fails + # this test in seconds instead of hanging it + "maximum_spend_logs_cleanup_max_batches": 50, + } + ) + + cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) + started_at = time.monotonic() + result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, time.monotonic() + 0.25) + elapsed = time.monotonic() - started_at + + assert result.stop_reason == "budget_exhausted" + assert elapsed < 3, f"budgeted run overran its deadline: {elapsed}s" + assert mock_db.execute_raw.call_count < 50 + assert result.rows_deleted > 0 + + +@pytest.mark.asyncio +async def test_run_budget_is_shared_across_tables_not_granted_per_table(): + """ + A per-table budget would let a run take N times the configured bound. The + deadline is computed once per run, so once it is spent on the first table + the later tables must stop immediately rather than each getting a fresh one. + """ + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock(return_value=1000) + mock_prisma_client.db = mock_db + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "maximum_autorouter_session_retention_period": "365d", + # Comfortably more batches than a sub-second budget can reach (each + # batch sleeps 0.1s), but small enough that a broken deadline fails + # this test in seconds instead of hanging it + "maximum_spend_logs_cleanup_max_batches": 50, + "maximum_spend_logs_cleanup_run_budget": "1s", + } + ) + cleaner.pod_lock_manager = None + + started_at = time.monotonic() + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + elapsed = time.monotonic() - started_at + + # three tables are eligible; a per-table budget would push this past 3s + assert elapsed < 2.5, f"budget was granted per table, not per run: {elapsed}s" + tables_touched = {call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list} + assert "LiteLLM_SpendLogs" in tables_touched + + +@pytest.mark.asyncio +async def test_each_batch_carries_a_statement_and_lock_timeout(): + """ + A Prisma transaction timeout cannot interrupt a statement already running, + so the Postgres statement_timeout and lock_timeout are the only things + stopping one batch from holding row locks and a pooled connection + indefinitely. Both must be set, inside the batch's own transaction, and + scoped with SET LOCAL so the pooled connection is left unchanged. + """ + recorded: list[str] = [] + + mock_prisma_client = MagicMock() + mock_db = MagicMock() + + @asynccontextmanager + async def _tx(): + tx = MagicMock() + + async def _execute_raw(sql, *args): + recorded.append(sql.strip()) + return 0 + + tx.execute_raw = _execute_raw + yield tx + + mock_db.tx = _tx + mock_db.query_raw = AsyncMock(return_value=[{"remaining": 0}]) + mock_prisma_client.db = mock_db + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "maximum_spend_logs_cleanup_batch_timeout": "12s", + } + ) + + await cleaner._delete_old_logs( + mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() + ) + + assert "SET LOCAL statement_timeout = 12000" in recorded + assert "SET LOCAL lock_timeout = 12000" in recorded + # the timeouts must precede the delete they are meant to bound + assert recorded.index("SET LOCAL statement_timeout = 12000") < next( + i for i, sql in enumerate(recorded) if sql.startswith("DELETE") + ) + + +@pytest.mark.parametrize( + "setting_value", + ["inf", "-inf", "nan", "1e400", "0s", "-5m", "not-a-duration"], +) +def test_a_non_finite_or_non_positive_budget_falls_back_to_the_default(setting_value): + """ + The knob must not be able to remove the bound it exists to enforce. + + 'inf', 'nan' and '1e400' are the spellings that would turn the deadline + into no deadline at all, and '0s' and '-5m' would make every run stop before + deleting anything. All of them must land on the default rather than being + honoured, and the resulting budget must be usable arithmetic. + """ + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "maximum_spend_logs_cleanup_run_budget": setting_value, + } + ) + + assert cleaner.run_budget_seconds == SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS + assert math.isfinite(cleaner.run_budget_seconds) + assert cleaner.run_budget_seconds > 0 + + +@pytest.mark.parametrize("setting_value", [0, -1, "abc", "", 2.9]) +def test_a_bad_batch_size_falls_back_to_the_default(setting_value): + """A zero or negative batch size would make every DELETE a no-op and the + loop spin, so unusable values must fall back rather than be honoured.""" + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "maximum_spend_logs_cleanup_batch_size": setting_value, + } + ) + + assert cleaner.batch_size >= 1 + + +def test_operator_knobs_override_the_env_defaults(): + """The knobs are meant to be reachable from general_settings (and therefore + from the admin UI), not only from environment variables.""" + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "maximum_spend_logs_cleanup_batch_size": 250, + "maximum_spend_logs_cleanup_max_batches": 7, + "maximum_spend_logs_cleanup_run_budget": "90s", + "maximum_spend_logs_cleanup_batch_timeout": "2m", + } + ) + + assert cleaner.batch_size == 250 + assert cleaner.max_batches == 7 + assert cleaner.run_budget_seconds == 90 + assert cleaner.batch_timeout_seconds == 120 + + +_BOUND_SETTING_CASES = ( + ("maximum_spend_logs_cleanup_batch_size", 137, "batch_size", 137), + ("maximum_spend_logs_cleanup_max_batches", 9, "max_batches", 9), + ("maximum_spend_logs_cleanup_run_budget", "45s", "run_budget_seconds", 45.0), + ("maximum_spend_logs_cleanup_batch_timeout", "8s", "batch_timeout_seconds", 8.0), +) + + +@pytest.mark.parametrize("setting_name, setting_value, attribute, expected", _BOUND_SETTING_CASES) +@pytest.mark.asyncio +async def test_a_bound_changed_after_construction_reaches_the_next_run( + setting_name, setting_value, attribute, expected +): + """The scheduler holds one long-lived instance and the config reload mutates + general_settings in place, so a bound captured at construction would leave + every dashboard change inert until the process restarts.""" + settings = {"maximum_spend_logs_retention_period": "7d"} + cleaner = SpendLogCleanup(general_settings=settings) + cleaner.pod_lock_manager = None + assert getattr(cleaner, attribute) != expected + + settings[setting_name] = setting_value + + await cleaner.cleanup_old_spend_logs(_mock_prisma_for_retention([0, 0])) + + assert getattr(cleaner, attribute) == expected + + +@pytest.mark.parametrize("cleared_to_none", [True, False]) +@pytest.mark.asyncio +async def test_a_bound_cleared_after_construction_falls_back_to_its_default(cleared_to_none): + """Blanking the field in the dashboard has to restore the shipped default + rather than leave the operator's old bound in force, whether the reload + spells the clear as an explicit None or as an absent key.""" + settings = {"maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_batch_size": 137} + cleaner = SpendLogCleanup(general_settings=settings) + cleaner.pod_lock_manager = None + assert cleaner.batch_size == 137 + + if cleared_to_none: + settings["maximum_spend_logs_cleanup_batch_size"] = None + else: + del settings["maximum_spend_logs_cleanup_batch_size"] + + await cleaner.cleanup_old_spend_logs(_mock_prisma_for_retention([0, 0])) + + assert cleaner.batch_size == SPEND_LOG_CLEANUP_BATCH_SIZE + + +def test_every_declared_bound_setting_is_covered_by_a_live_reread_case(): + """A bound added to the declared set without a live-reread case would be + propagated by the proxy and then ignored by the running job.""" + assert {case[0] for case in _BOUND_SETTING_CASES} == set(SPEND_LOG_CLEANUP_BOUND_SETTINGS) + + +@pytest.mark.asyncio +async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table(): + """The remaining-eligible-rows metric must never itself become the long + scan this job exists to avoid, so its probe carries a LIMIT.""" + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock(return_value=0) + mock_prisma_client.db = mock_db + + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + + await cleaner._delete_old_logs( + mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() + ) + + count_sql = mock_db.query_raw.call_args[0][0] + assert "count(*)" in count_sql + assert "LIMIT $2" in count_sql + assert mock_db.query_raw.call_args[0][2] == SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP + + +@pytest.mark.asyncio +async def test_a_run_skipped_because_another_pod_holds_the_lock_is_reported(): + """Operators need to tell "nothing to do" apart from "someone else is doing + it", so a lock-skipped run is recorded under its own outcome.""" + recorded: list[str] = [] + original_record_run = SpendLogCleanupMetrics.record_run + + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = MagicMock() + cleaner.pod_lock_manager.redis_cache = MagicMock() + cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=False) + cleaner.pod_lock_manager.release_lock = AsyncMock() + + SpendLogCleanupMetrics.record_run = classmethod(lambda cls, outcome: recorded.append(outcome)) + try: + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + finally: + SpendLogCleanupMetrics.record_run = original_record_run + + assert recorded == ["skipped_locked"] + cleaner.pod_lock_manager.release_lock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_the_outstanding_rows_probe_carries_a_statement_timeout(): + """ + The probe is a statement like any other, so if it were issued bare a slow one + would hold a connection past the budget the job advertises, which is exactly + what the bounds exist to prevent. With budget to spare it carries the same + per-statement timeout the delete batches do. + """ + recorded: list[str] = [] + + mock_prisma_client = MagicMock() + mock_db = MagicMock() + + @asynccontextmanager + async def _tx(): + tx = MagicMock() + + async def _execute_raw(sql, *args): + recorded.append(sql.strip()) + return 0 + + async def _query_raw(sql, *args): + recorded.append(sql.strip()) + return [{"remaining": 7}] + + tx.execute_raw = _execute_raw + tx.query_raw = _query_raw + yield tx + + mock_db.tx = _tx + mock_prisma_client.db = mock_db + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "maximum_spend_logs_cleanup_batch_timeout": "8s", + } + ) + + remaining = await cleaner._count_remaining( + mock_prisma_client, + datetime.now(timezone.utc) - timedelta(days=7), + "LiteLLM_SpendLogs", + "startTime", + _far_deadline(), + ) + + assert remaining == 7 + count_index = next(i for i, sql in enumerate(recorded) if sql.startswith("SELECT count(*)")) + assert "SET LOCAL statement_timeout = 8000" in recorded[:count_index], ( + f"the probe ran without a statement timeout: {recorded}" + ) + + +@pytest.mark.asyncio +async def test_a_statement_timeout_is_clamped_to_the_budget_that_is_left(): + """ + Postgres has no 'stop at time T', only a per-statement duration, so a batch + issued just under the deadline would run a whole batch timeout past it and + the run budget would be advisory. Clamping the timeout to the remaining + budget is what makes the budget a real wall clock. + """ + recorded: list[str] = [] + client = MagicMock() + + @asynccontextmanager + async def _tx(): + tx = MagicMock() + + async def _execute_raw(sql, *args): + recorded.append(sql.strip()) + return 0 + + tx.execute_raw = _execute_raw + tx.query_raw = AsyncMock(return_value=[{"remaining": 0}]) + yield tx + + client.db.tx = _tx + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "maximum_spend_logs_cleanup_batch_timeout": "30s", + } + ) + + # Only 2s of budget left against a 30s batch timeout. + await cleaner._execute_delete_batch(client, "DELETE FROM x", datetime.now(timezone.utc), time.monotonic() + 2) + + timeouts = [sql for sql in recorded if "statement_timeout" in sql] + assert timeouts, f"no statement timeout was issued: {recorded}" + issued_ms = int(timeouts[0].split("=")[1].strip()) + assert issued_ms <= 2000, f"the batch was given {issued_ms}ms with only 2000ms of budget left" + + +@pytest.mark.asyncio +async def test_no_statement_is_issued_once_the_budget_is_spent(): + """ + Every table exits through _finish_table, including the ones a spent run never + started, so an unconditional probe there would put one more statement per + table past the bound. + """ + client = _mock_prisma_for_retention([0, 0]) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + + result = await cleaner._finish_table( + client, + datetime.now(timezone.utc) - timedelta(days=7), + "LiteLLM_SpendLogs", + "startTime", + 123, + "budget_exhausted", + time.monotonic() - 1, + ) + + assert result.rows_deleted == 123 + assert result.stop_reason == "budget_exhausted" + client.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_a_batch_cancelled_by_the_deadline_is_budget_exhaustion_not_a_failure(monkeypatch): + """ + Clamping the timeout means the last batch of a budget-exhausted run is + cancelled by the deadline itself. Counting that as a batch failure would + inflate the failure metric on every such run and walk it toward the abort + threshold, so it has to be classified as the bound working. + """ + failures: list[str] = [] + client = MagicMock() + _wire_tx(client.db) + + # The deadline has to pass DURING the batch, not before it: a deadline + # already spent is caught by the loop's own check and no batch is ever + # issued, which would exercise none of the classification under test. + async def _cancelled_after_the_deadline(sql, *args): + await asyncio.sleep(0.05) + raise Exception("canceling statement due to statement timeout") + + client.db.execute_raw = _cancelled_after_the_deadline + + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + monkeypatch.setattr(SpendLogCleanupMetrics, "record_batch_failure", lambda table: failures.append(table)) + + result = await cleaner._delete_old_logs( + client, datetime.now(timezone.utc) - timedelta(days=7), time.monotonic() + 0.02 + ) + + assert result.stop_reason == "budget_exhausted" + assert failures == [], f"a deadline cancellation was recorded as a batch failure: {failures}" + + +@pytest.mark.asyncio +async def test_partition_maintenance_is_skipped_once_the_run_budget_is_spent(): + """ + Dropping a partition is DDL holding an ACCESS EXCLUSIVE lock, and unlike a + delete batch it cannot be cut short once it has started. A run whose budget is + already gone must therefore not start it at all; the next tick picks it up. + """ + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock(return_value=0) + mock_prisma_client.db = mock_db + + partition_manager = MagicMock() + partition_manager.is_partitioned = AsyncMock(return_value=True) + partition_manager.ensure_partitions = AsyncMock() + partition_manager.drop_partitions_older_than = AsyncMock(return_value=[]) + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "use_spend_logs_partitioning": True, + }, + partition_manager=partition_manager, + ) + cleaner._should_delete_spend_logs() + + # a deadline already in the past is what a run that spent its budget on an + # earlier table looks like + await cleaner._clean_spend_log_tables(mock_prisma_client, time.monotonic() - 1) + + partition_manager.ensure_partitions.assert_not_awaited() + partition_manager.drop_partitions_older_than.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_partition_maintenance_still_runs_while_the_run_has_budget(): + """The skip above must be caused by the spent budget, not by breaking the + partition path outright.""" + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock(return_value=0) + mock_prisma_client.db = mock_db + + partition_manager = MagicMock() + partition_manager.is_partitioned = AsyncMock(return_value=True) + partition_manager.ensure_partitions = AsyncMock() + partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"]) + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "use_spend_logs_partitioning": True, + }, + partition_manager=partition_manager, + ) + cleaner._should_delete_spend_logs() + + await cleaner._clean_spend_log_tables(mock_prisma_client, _far_deadline()) + + partition_manager.ensure_partitions.assert_awaited_once() + partition_manager.drop_partitions_older_than.assert_awaited_once() + + +@pytest.mark.parametrize( + "stop_reasons, expected", + [ + (("exhausted",), "completed"), + (("exhausted", "exhausted"), "completed"), + (("exhausted", "batch_cap_reached"), "batch_cap_reached"), + (("batch_cap_reached", "exhausted"), "batch_cap_reached"), + (("exhausted", "budget_exhausted"), "budget_exhausted"), + (("budget_exhausted", "exhausted"), "budget_exhausted"), + (("batch_cap_reached", "budget_exhausted"), "budget_exhausted"), + (("budget_exhausted", "batch_cap_reached"), "budget_exhausted"), + (("exhausted", "aborted"), "aborted"), + (("aborted", "exhausted"), "aborted"), + (("budget_exhausted", "aborted"), "aborted"), + (("aborted", "budget_exhausted"), "aborted"), + (("aborted", "budget_exhausted", "batch_cap_reached"), "aborted"), + ], +) +def test_the_reported_run_outcome_is_the_most_significant_reason_in_any_order(stop_reasons, expected): + """ + The run outcome answers "why did this run stop", so a table that merely ran + dry must never mask one that hit a bound, and an abort must outrank both. + + Both orders of every pair are covered because this folds several per-table + results into one answer: a first-match-wins implementation would pass on + whichever order happened to be written and fail on its mirror. + """ + results = tuple(TableCleanupResult(rows_deleted=0, stop_reason=reason) for reason in stop_reasons) + assert SpendLogCleanup._run_outcome(results) == expected diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index 87063bdf00b..ed1317e647d 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -17,13 +17,14 @@ import hashlib import json from datetime import datetime, timedelta, timezone from types import SimpleNamespace -from typing import Any +from typing import Any, Final from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException from litellm.proxy._types import LiteLLM_VerificationTokenView +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper from litellm.proxy.utils import PrismaClient @@ -270,6 +271,9 @@ async def test_query_first_with_cached_plan_fallback_reconnects_then_retries_ide assert retry_call.args == first_call.args == (original_query, "abc") reconnect.assert_awaited_once() assert reconnect.await_args.kwargs.get("force", False) is False + # https://github.com/BerriAI/litellm/issues/36418: without this the healthy + # writer probe skips the recreate and the stale plans survive the retry + assert reconnect.await_args.kwargs.get("force_recreate") is True assert [name for name, *_ in manager.mock_calls] == [ "query_first", "attempt_db_reconnect", @@ -564,3 +568,68 @@ async def test_get_data_team_keys_forward_limit_as_take( "where": {"team_id": "team-1"}, "include": {"litellm_budget_table": True}, } + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_reports_pre_query_engine_generation( + prisma_client: PrismaClient, +) -> None: + """The generation is snapshotted before the query, not after it fails: it + names the engine that prepared the stale statement, which is what lets the + reconnect bypass an unrelated cooldown while that engine is still live + (https://github.com/BerriAI/litellm/issues/36418). Reading it after the + failure would miss a recreate that landed in between and force a + needless second one.""" + prisma_client.db.engine_generation = 3 + + async def _fail_then_bump(*args: Any, **kwargs: Any) -> dict[str, str]: + if prisma_client.db.engine_generation == 3: + prisma_client.db.engine_generation = 4 + raise RuntimeError("cached plan must not change result type") + return {"token": "abc"} + + prisma_client.db.query_first = AsyncMock(side_effect=_fail_then_bump) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + kwargs = prisma_client.attempt_db_reconnect.await_args.kwargs + assert kwargs.get("stale_read_engine").generation == 3 + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_reports_the_reader_generation( + prisma_client: PrismaClient, +) -> None: + """With a read replica configured the query runs on the READER, so the + reader's generation is the one that names the engine holding the stale + prepared statement. Snapshotting the writer's instead would let an + unrelated writer reconnect re-arm the cooldown while the reader stayed + poisoned (https://github.com/BerriAI/litellm/issues/36418). The two + generations are deliberately far apart so only the right one matches.""" + writer = MagicMock(name="writer") + writer.engine_generation = 99 + writer.query_first = AsyncMock(return_value={"token": "wrong-engine"}) + reader = MagicMock(name="reader") + reader.engine_generation = 3 + reader.query_first = AsyncMock( + side_effect=[RuntimeError("cached plan must not change result type"), {"token": "abc"}] + ) + prisma_client.db = RoutingPrismaWrapper(writer=writer, reader=reader) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + reported: Final = prisma_client.attempt_db_reconnect.await_args.kwargs.get("stale_read_engine") + pinned = { + "reported_generation": reported.generation, + "reported_the_reader_itself": reported.wrapper is reader, + "reader_served_the_query": reader.query_first.await_count, + "writer_served_the_query": writer.query_first.await_count, + } + assert pinned == { + "reported_generation": 3, + "reported_the_reader_itself": True, + "reader_served_the_query": 2, + "writer_served_the_query": 0, + } diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py index 867554157fd..719d7cc73f5 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py @@ -7,17 +7,34 @@ Symbols pinned here: - ``PrismaClient.start_db_health_watchdog_task`` - ``PrismaClient.stop_db_health_watchdog_task`` - ``PrismaClient._db_health_watchdog_loop`` + +Note on fixtures for the routing tests: the reader and the writer carry +independent generation counters, so a fixture that gives them far-apart values +reads clearly and proves nothing about identity, because comparing the numbers +alone already yields the right answer. Pick values so that ONLY the mechanism +under test can produce the expected result, which for identity means two +engines whose generations deliberately coincide. + +Note on what to assert: pin the requirement, not the mechanism. An assertion +that restates what the implementation currently does can only ever agree with +it, including when it is wrong, so it ends up defending the defect from being +corrected. One here did exactly that, asserting that a declined heavy-path +recreate leaves the dead-engine flag set, which read as a faithful description +and was a reintroduction of #29176. "A later cycle must not kill a healthy +engine" would have failed against it whatever mechanism produced it. """ from __future__ import annotations import asyncio -from typing import Any +import time +from typing import Any, Final from unittest.mock import AsyncMock, MagicMock import pytest -from litellm.proxy.utils import PrismaClient +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper +from litellm.proxy.utils import PrismaClient, _StaleReadEngine @pytest.mark.asyncio @@ -96,6 +113,48 @@ async def test_run_reconnect_cycle_direct_path_recreates_when_probe_fails( } +@pytest.mark.asyncio +async def test_run_reconnect_cycle_force_recreate_skips_probe_and_recreates( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """A healthy writer must not veto the recreate when the caller already + knows the session state is poisoned (stale prepared statements after a + schema change). Regression for + https://github.com/BerriAI/litellm/issues/36418.""" + monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db") + prisma_client._engine_confirmed_dead = False + prisma_client._engine_pid = 0 + prisma_client._start_engine_watcher = AsyncMock() + prisma_client._cleanup_engine_watcher = MagicMock() + + writer = prisma_client.db + writer.recreate_prisma_client = AsyncMock() + writer.query_raw = AsyncMock(return_value=[{"?column?": 1}]) + + await prisma_client._run_reconnect_cycle(timeout_seconds=5, force_recreate=True) + pinned = { + "recreate_called": writer.recreate_prisma_client.await_count, + "writer_query_raw_calls": writer.query_raw.await_count, + } + assert pinned == {"recreate_called": 1, "writer_query_raw_calls": 1} + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_forwards_force_recreate_to_cycle( + prisma_client: PrismaClient, +) -> None: + """Regression for https://github.com/BerriAI/litellm/issues/36418: the flag + has to survive both hops (attempt_db_reconnect -> inside-lock -> cycle), + otherwise the cached-plan caller silently gets a probe-gated reconnect.""" + prisma_client._db_last_reconnect_attempt_ts = 0.0 + prisma_client._run_reconnect_cycle = AsyncMock() + + ok = await prisma_client.attempt_db_reconnect(reason="explicit", force_recreate=True) + + assert ok is True + assert prisma_client._run_reconnect_cycle.await_args.kwargs.get("force_recreate") is True + + @pytest.mark.asyncio async def test_run_reconnect_cycle_passes_writer_generation_to_recreate( prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch @@ -584,3 +643,456 @@ async def test_run_reconnect_cycle_heavy_path_forwards_entry_generation_to_recre kwargs = prisma_client.db.recreate_prisma_client.await_args.kwargs assert kwargs.get("expected_generation") == 4 + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_bypasses_cooldown_for_still_live_stale_engine( + prisma_client: PrismaClient, +) -> None: + """A schema change landing inside the cooldown of an earlier reconnect used + to leave auth failing until the cooldown elapsed. While the engine the + caller's failure came from is still the live one, the cooldown must not + gate the recreate. Regression for + https://github.com/BerriAI/litellm/issues/36418.""" + prisma_client.db.engine_generation = 7 + prisma_client._db_last_reconnect_attempt_ts = time.time() + prisma_client._run_reconnect_cycle = AsyncMock() + + ok = await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7), + ) + + assert ok is True + assert prisma_client._run_reconnect_cycle.await_count == 1 + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_honors_cooldown_once_stale_engine_replaced( + prisma_client: PrismaClient, +) -> None: + """The bypass is scoped to the damaged engine: once a concurrent recreate + has replaced it, the cooldown must still collapse the rest of the burst + onto that recreate instead of killing the fresh engine.""" + prisma_client.db.engine_generation = 8 + prisma_client._db_last_reconnect_attempt_ts = time.time() + prisma_client._run_reconnect_cycle = AsyncMock() + + ok = await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7), + ) + + assert ok is False + prisma_client._run_reconnect_cycle.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_keeps_cooldown_for_callers_without_generation( + prisma_client: PrismaClient, +) -> None: + """Watchdog and transport-error callers name no generation, so they keep + the plain cooldown behaviour.""" + prisma_client._db_last_reconnect_attempt_ts = time.time() + prisma_client._run_reconnect_cycle = AsyncMock() + + ok = await prisma_client.attempt_db_reconnect(reason="watchdog_probe_failed") + + assert ok is False + prisma_client._run_reconnect_cycle.assert_not_awaited() + + +def _routing_client(prisma_client: PrismaClient, reader_generation: int, writer_generation: int) -> tuple[Any, Any]: + """Wire ``prisma_client.db`` to a routing wrapper with distinct engines. + + Returns the (writer, reader) mocks so a test can move either generation + independently, which is the only way to tell the two counters apart. + """ + writer = MagicMock(name="writer") + writer.engine_generation = writer_generation + reader = MagicMock(name="reader") + reader.engine_generation = reader_generation + prisma_client.db = RoutingPrismaWrapper(writer=writer, reader=reader) + return writer, reader + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_reads_generation_from_the_reader_that_served_the_query( + prisma_client: PrismaClient, +) -> None: + """``query_first`` is a top-level read, so with a replica configured the + stale prepared statements are on the READER. A writer reconnect that moved + the writer generation must not re-arm the cooldown while the reader the + query actually failed on is still the live, poisoned one. Regression for + https://github.com/BerriAI/litellm/issues/36418.""" + _routing_client(prisma_client, reader_generation=7, writer_generation=99) + prisma_client._db_last_reconnect_attempt_ts = time.time() + prisma_client._run_reconnect_cycle = AsyncMock() + + ok = await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7), + ) + + assert ok is True + assert prisma_client._run_reconnect_cycle.await_count == 1 + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_honors_cooldown_once_the_reader_itself_was_replaced( + prisma_client: PrismaClient, +) -> None: + """The mirror of the above: once the reader has been replaced, the recreate + the caller needed has already happened, so the cooldown collapses the rest + of the burst even though the writer generation never moved.""" + _routing_client(prisma_client, reader_generation=8, writer_generation=99) + prisma_client._db_last_reconnect_attempt_ts = time.time() + prisma_client._run_reconnect_cycle = AsyncMock() + + ok = await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7), + ) + + assert ok is False + prisma_client._run_reconnect_cycle.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_gates_when_reads_moved_to_an_engine_of_the_same_generation( + prisma_client: PrismaClient, +) -> None: + """The counters are per engine, so the reader and the writer can sit on the + same number at the same time. Once the reader goes unavailable reads move to + the writer, and the caller's poisoned reader is no longer serving anything, + so the cooldown should gate it. Comparing generations alone cannot tell the + two apart and would hand out the waiver here: the generations are equal on + purpose, which is what makes this the case identity has to decide.""" + writer, reader = _routing_client(prisma_client, reader_generation=5, writer_generation=5) + stale: Final = _StaleReadEngine(wrapper=reader, generation=5) + prisma_client.db._reader_unavailable = True + prisma_client._db_last_reconnect_attempt_ts = time.time() + prisma_client._run_reconnect_cycle = AsyncMock() + + ok = await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=stale, + ) + + pinned = { + "reads_now_served_by_the_writer": prisma_client.read_db is writer, + "generations_coincide": reader.engine_generation == writer.engine_generation, + # `_cooldown_applies` gates on the failed-repair record OR on liveness, + # and either alone produces this result. Pin that the record is empty, + # or a stray entry would make this pass while testing the other gate. + "no_failed_repair_recorded": dict(prisma_client._failed_recreate_generations) == {}, + "recovered": ok, + "cycles_run": prisma_client._run_reconnect_cycle.await_count, + } + assert pinned == { + "reads_now_served_by_the_writer": True, + "generations_coincide": True, + "no_failed_repair_recorded": True, + "recovered": False, + "cycles_run": 0, + } + + +@pytest.mark.asyncio +async def test_failed_repair_of_one_engine_is_not_evicted_by_a_failure_on_the_other( + prisma_client: PrismaClient, +) -> None: + """The record is kept per engine. Held in a single slot, a failed writer + repair would evict the reader's record, and the next caller naming the + reader's still-unrepaired generation would get the waiver back and run its + own redundant cycle, which is the burst the record exists to collapse.""" + writer, reader = _routing_client(prisma_client, reader_generation=5, writer_generation=3) + stale_reader: Final = _StaleReadEngine(wrapper=reader, generation=5) + stale_writer: Final = _StaleReadEngine(wrapper=writer, generation=3) + prisma_client._db_last_reconnect_attempt_ts = 0.0 + prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("engine spawn failed")) + + await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", force_recreate=True, stale_read_engine=stale_reader + ) + prisma_client.db._reader_unavailable = True + await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", force_recreate=True, stale_read_engine=stale_writer + ) + prisma_client.db._reader_unavailable = False + cycles_before_the_reader_returns: Final = prisma_client._run_reconnect_cycle.await_count + + await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", force_recreate=True, stale_read_engine=stale_reader + ) + + pinned = { + "cycles_before": cycles_before_the_reader_returns, + "cycles_after": prisma_client._run_reconnect_cycle.await_count, + } + assert pinned == {"cycles_before": 2, "cycles_after": 2} + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_withdraws_the_waiver_after_this_generation_failed_to_repair( + prisma_client: PrismaClient, +) -> None: + """A failed recreate leaves the generation where it was, so without a record + of the failure every queued caller of the same burst would still see its own + generation live and run its own full recreate serially instead of collapsing + onto one attempt. Drives two callers rather than presetting the record, so + the record has to actually be written by the failure.""" + prisma_client.db.engine_generation = 7 + prisma_client._db_last_reconnect_attempt_ts = 0.0 + prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("engine spawn failed")) + + first = await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7), + ) + second = await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7), + ) + + pinned = { + "first": first, + "second": second, + "cycles_run": prisma_client._run_reconnect_cycle.await_count, + } + assert pinned == {"first": False, "second": False, "cycles_run": 1} + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_keeps_the_waiver_after_an_unrelated_reconnect_failure( + prisma_client: PrismaClient, +) -> None: + """The failure record is scoped to the generation it was trying to repair. + A watchdog or transport-error reconnect names no generation, so its failure + says nothing about whether a stale read engine can be repaired and must not + gate it: gating on a global failure count would 503 authentication for the + length of the cooldown.""" + prisma_client.db.engine_generation = 7 + prisma_client._db_last_reconnect_attempt_ts = 0.0 + prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("watchdog reconnect failed")) + + unrelated = await prisma_client.attempt_db_reconnect(reason="watchdog_probe_failed") + # Read before the second call: a global failure gate would be armed here, + # and the recovering reconnect below resets the counter either way. + failures_left_by_the_unrelated_reconnect: Final = prisma_client._consecutive_reconnect_failures + + prisma_client._run_reconnect_cycle = AsyncMock() + cached_plan = await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7), + ) + + pinned = { + "unrelated_failed": unrelated, + "failures_left_by_the_unrelated_reconnect": failures_left_by_the_unrelated_reconnect, + "cached_plan_recovered": cached_plan, + "cycles_run_for_cached_plan": prisma_client._run_reconnect_cycle.await_count, + } + assert pinned == { + "unrelated_failed": False, + "failures_left_by_the_unrelated_reconnect": 1, + "cached_plan_recovered": True, + "cycles_run_for_cached_plan": 1, + } + + +@pytest.mark.asyncio +async def test_forced_recreate_declined_by_the_generation_guard_is_not_reported_as_success( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """``recreate_prisma_client`` declines when the writer generation moved + since cycle entry, and the routing wrapper then leaves the reader untouched + too. A forced caller asked for its engine to be replaced and it was not, so + reporting success would reset the consecutive-failure count and log a repair + that never happened. The declined attempt must equally not count as a + failure, or the caller's own backoff would be gated on its next try.""" + monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db") + prisma_client._engine_confirmed_dead = False + prisma_client._engine_pid = 0 + prisma_client._start_engine_watcher = AsyncMock() + prisma_client._cleanup_engine_watcher = MagicMock() + prisma_client._db_last_reconnect_attempt_ts = 0.0 + prisma_client._consecutive_reconnect_failures = 0 + + writer = prisma_client.db + writer.recreate_prisma_client = AsyncMock(return_value=False) + writer.query_raw = AsyncMock(return_value=[{"?column?": 1}]) + + ok = await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + ) + + pinned = { + "reported_success": ok, + "recreate_attempted": writer.recreate_prisma_client.await_count, + "consecutive_failures": prisma_client._consecutive_reconnect_failures, + } + assert pinned == { + "reported_success": False, + "recreate_attempted": 1, + "consecutive_failures": 0, + } + + +@pytest.mark.asyncio +async def test_unforced_recreate_declined_by_the_generation_guard_still_succeeds( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """The decline is only an error for a caller that forced the recreate. A + transport-blip caller is happy to learn another path already replaced the + engine, so its reconnect still reports success.""" + monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db") + prisma_client._engine_confirmed_dead = False + prisma_client._engine_pid = 0 + prisma_client._start_engine_watcher = AsyncMock() + prisma_client._cleanup_engine_watcher = MagicMock() + prisma_client._db_last_reconnect_attempt_ts = 0.0 + + writer = prisma_client.db + writer.recreate_prisma_client = AsyncMock(return_value=False) + # First call is the liveness probe, which must fail so the recreate is + # reached at all; the second is the post-recreate smoke test. + writer.query_raw = AsyncMock(side_effect=[Exception("probe fails"), [{"?column?": 1}]]) + + ok = await prisma_client.attempt_db_reconnect(reason="transport_blip") + + pinned = {"reported_success": ok, "recreate_attempted": writer.recreate_prisma_client.await_count} + assert pinned == {"reported_success": True, "recreate_attempted": 1} + + +@pytest.mark.asyncio +async def test_heavy_path_forced_recreate_declined_is_not_reported_as_success( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """A forced caller reaches the heavy branch too: the escalation threshold + flips ``_engine_confirmed_dead`` after repeated failures, and every cycle + after that takes the dead-engine path. A decline there has to be treated + exactly as it is on the direct path, or the escalation itself reintroduces + the success that never happened.""" + monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db") + prisma_client._engine_confirmed_dead = True + prisma_client._engine_pid = 1234 + prisma_client._start_engine_watcher = AsyncMock() + prisma_client._cleanup_engine_watcher = MagicMock() + prisma_client._db_last_reconnect_attempt_ts = 0.0 + prisma_client._consecutive_reconnect_failures = 0 + monkeypatch.setattr(PrismaClient, "_reap_all_zombies", staticmethod(lambda: set())) + + prisma_client.db.recreate_prisma_client = AsyncMock(return_value=False) + + ok = await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + ) + + pinned = { + "reported_success": ok, + "recreate_attempted": prisma_client.db.recreate_prisma_client.await_count, + "consecutive_failures": prisma_client._consecutive_reconnect_failures, + # The dead-engine flag must be CLEARED. A raise normally skips the + # clear, which is right for a failure and wrong here: the guard + # declined because another path had already replaced the engine, so it + # is alive. Leaving it set routes the next cycle back down this + # probe-free branch, where the recreate would kill that healthy engine. + "engine_still_confirmed_dead": prisma_client._engine_confirmed_dead, + } + assert pinned == { + "reported_success": False, + "recreate_attempted": 1, + "consecutive_failures": 0, + "engine_still_confirmed_dead": False, + } + + +@pytest.mark.asyncio +async def test_declined_heavy_recreate_disarms_escalation_for_the_next_attempt( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """Clearing the dead-engine flag on a decline is not enough on its own. The + escalation check re-arms that flag whenever the consecutive-failure count is + still at the threshold, so a decline that left the count alone would send + the very next attempt back down the probe-free heavy path and recreate over + the healthy engine another path had just installed. Drives the SECOND + attempt, because the first one alone cannot show this.""" + monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db") + prisma_client._engine_pid = 1234 + prisma_client._start_engine_watcher = AsyncMock() + prisma_client._cleanup_engine_watcher = MagicMock() + prisma_client._db_last_reconnect_attempt_ts = 0.0 + monkeypatch.setattr(PrismaClient, "_reap_all_zombies", staticmethod(lambda: set())) + # Escalation already armed by earlier genuine failures. + prisma_client._consecutive_reconnect_failures = prisma_client._reconnect_escalation_threshold + prisma_client.db.recreate_prisma_client = AsyncMock(return_value=False) + prisma_client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}]) + + armed: Final = prisma_client._engine_confirmed_dead is False and prisma_client._consecutive_reconnect_failures > 0 + + await prisma_client.attempt_db_reconnect(reason="postgres_cached_plan_error", force_recreate=True) + + # Kept as its own assert, not folded into the judgement below. These are two + # claims about two moments, the first being a precondition for the second + # meaning anything, and a single combined comparison would hide which one + # failed from both the traceback and a mutation report. + assert { + "escalation_was_armed_by_the_count": armed, + "failures": prisma_client._consecutive_reconnect_failures, + "engine_confirmed_dead": prisma_client._engine_confirmed_dead, + } == {"escalation_was_armed_by_the_count": True, "failures": 0, "engine_confirmed_dead": False} + + prisma_client._db_last_reconnect_attempt_ts = 0.0 + await prisma_client.attempt_db_reconnect(reason="postgres_cached_plan_error", force_recreate=True) + + # The requirement: a later cycle must not reclassify the healthy replacement + # as dead and restart it through the probe-free path. + assert prisma_client._engine_confirmed_dead is False + + +@pytest.mark.asyncio +async def test_unrelated_reconnect_failure_does_not_erase_the_burst_record( + prisma_client: PrismaClient, +) -> None: + """The failure record names one engine, so a caller that names none must + not overwrite it. Otherwise a watchdog failure landing between two callers + of the same burst clears the record and the second caller runs its own full + recreate against the engine the first one just failed to repair.""" + prisma_client.db.engine_generation = 7 + prisma_client._db_last_reconnect_attempt_ts = 0.0 + stale: Final = _StaleReadEngine(wrapper=prisma_client.read_db, generation=7) + prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("engine spawn failed")) + + await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=stale, + ) + # force=True the way the engine-death callers do, so this one actually + # reaches the failure branch instead of being skipped by the cooldown the + # first caller just stamped. + await prisma_client.attempt_db_reconnect(reason="engine_process_death", force=True) + cycles_before_the_second_burst_caller: Final = prisma_client._run_reconnect_cycle.await_count + + await prisma_client.attempt_db_reconnect( + reason="postgres_cached_plan_error", + force_recreate=True, + stale_read_engine=stale, + ) + + pinned = { + "cycles_before": cycles_before_the_second_burst_caller, + "cycles_after": prisma_client._run_reconnect_cycle.await_count, + } + assert pinned == {"cycles_before": 2, "cycles_after": 2} diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py index 9defb309863..56057dce7e0 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py @@ -76,6 +76,23 @@ def test_convert_mcp_to_llm_format_defaults_model(proxy_logging, make_mcp_reques } +def test_convert_mcp_to_llm_format_exposes_headers_on_metadata(proxy_logging, make_mcp_request_obj): + """Guardrails read the caller's HTTP headers off ``metadata.headers`` on the chat + completions path, so the MCP bridge has to put them in the same place.""" + req = make_mcp_request_obj() + out = proxy_logging._convert_mcp_to_llm_format( + request_obj=req, + kwargs={"headers": {"x-nuid": "nuid-1"}}, + ) + assert out["metadata"]["headers"] == {"x-nuid": "nuid-1"} + + +def test_convert_mcp_to_llm_format_defaults_headers_to_empty(proxy_logging, make_mcp_request_obj): + req = make_mcp_request_obj() + out = proxy_logging._convert_mcp_to_llm_format(request_obj=req, kwargs={}) + assert out["metadata"]["headers"] == {} + + def test_convert_mcp_to_llm_format_missing_request_obj_raises(proxy_logging): with pytest.raises(AttributeError): proxy_logging._convert_mcp_to_llm_format(request_obj=None, kwargs={}) diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index d60fff66c44..4981caa10c3 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -536,6 +536,37 @@ async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch): assert mock_get_tools.await_args.kwargs["request_tags"] == ["team-a"] +@pytest.mark.asyncio +async def test_execute_tool_calls_exposes_sanitized_client_headers_to_logging(monkeypatch): + """The Responses API MCP bridge used to log an empty header dict, hiding the caller's + headers from logging callbacks and hooks.""" + _setup_proxy_logging(monkeypatch) + _setup_mcp_call_environment(monkeypatch) + + captured = {} + + def fake_function_setup(*_args, **kwargs): + captured.update(kwargs) + return None, None + + handler_module = importlib.import_module( + "litellm.responses.mcp.litellm_proxy_mcp_handler" + ) + monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) + + tool_name = "deepwiki-read_wiki_structure" + await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}], + user_api_key_auth=None, + raw_headers={"x-nuid": "nuid-1", "x-litellm-api-key": "sk-proxy", "cookie": "s=1"}, + ) + + expected = {"x-nuid": "nuid-1", "cookie": "***REDACTED***"} + assert captured["metadata"]["headers"] == expected + assert captured["proxy_server_request"]["headers"] == expected + + @pytest.mark.asyncio async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monkeypatch): _setup_proxy_logging(monkeypatch) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 9e160370048..58373df024c 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1,3 +1,5 @@ +import contextlib +import copy import json import os import sys @@ -2461,3 +2463,104 @@ async def test_acompletion_forwards_aws_credentials_through_responses_bridge( finally: litellm.disable_aiohttp_transport = original_disable_aiohttp litellm.in_memory_llm_clients_cache.flush_cache() + + +_GEMINI_RESPONSE_BODY = { + "candidates": [{"content": {"parts": [{"text": "hello"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3}, +} + + +def _gemini_client_returning_a_reply(): + """An injected HTTP client whose post() answers like generativelanguage does.""" + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + request = httpx.Request("POST", "https://generativelanguage.googleapis.com/") + post = MagicMock(return_value=httpx.Response(200, json=_GEMINI_RESPONSE_BODY, request=request)) + return client, post + + +@pytest.fixture +def restore_model_registry(): + """litellm.model_cost and the provider name sets are module-global. + + register_model merges into the existing entry in place, hence the deep copy. + """ + model_cost = copy.deepcopy(litellm.model_cost) + openai_models = set(litellm.open_ai_chat_completion_models) + yield + litellm.model_cost.clear() + litellm.model_cost.update(model_cost) + litellm.open_ai_chat_completion_models.clear() + litellm.open_ai_chat_completion_models.update(openai_models) + + +def test_openai_model_name_does_not_outrank_explicit_provider(): + """`gemini/gpt-4o` goes to Google, not to litellm's OpenAI handler. + + completion() checks `model in litellm.open_ai_chat_completion_models` ahead of + the gemini branch, so the call used to reach the OpenAI handler carrying + VertexGeminiConfig, whose transform_request raises NotImplementedError. + """ + assert "gpt-4o" in litellm.open_ai_chat_completion_models + client, post = _gemini_client_returning_a_reply() + + with patch.object(client, "post", new=post): + response = litellm.completion( + model="gemini/gpt-4o", + messages=[{"role": "user", "content": "hello"}], + api_key="test-api-key", + client=client, + ) + + assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"] + assert "models/gpt-4o" in post.call_args.kwargs["url"] + assert response.choices[0].message.content == "hello" + + +def test_mislabelled_pricing_entry_does_not_reroute_provider(restore_model_registry): + """register_model is the other way into the same failure. + + An entry claiming litellm_provider "openai" adds its name to + open_ai_chat_completion_models, so one mislabelled price reroutes every later + call to that model in the process. + """ + litellm.register_model( + { + "gemini-2.5-pro": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + } + } + ) + assert "gemini-2.5-pro" in litellm.open_ai_chat_completion_models + client, post = _gemini_client_returning_a_reply() + + with patch.object(client, "post", new=post): + response = litellm.completion( + model="gemini/gemini-2.5-pro", + messages=[{"role": "user", "content": "hello"}], + api_key="test-api-key", + client=client, + ) + + assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"] + assert response.choices[0].message.content == "hello" + + +def test_openai_model_without_a_provider_still_routes_to_openai(): + from openai import OpenAI + + client = OpenAI(api_key="fake-key") + raw_response = client.chat.completions.with_raw_response + with patch.object(raw_response, "create") as mock_create, contextlib.suppress(Exception): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hello"}], + client=client, + ) + + mock_create.assert_called() diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index 35d98903226..33baf7474ce 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -224,6 +224,7 @@ def test_nothing_staged_and_no_changes_is_an_explicit_no_op(tmp_path: Path) -> N proc = _run(repo, bin_dir, {}) assert proc.returncode == 0, proc.stdout + proc.stderr assert "nothing to check" in proc.stdout + assert "check: PASS" in proc.stdout assert "linting Python" not in proc.stdout @@ -234,6 +235,7 @@ def test_nothing_staged_without_a_base_ref_fails_with_a_fetch_hint(tmp_path: Pat assert proc.returncode == 1 assert "cannot resolve the merge base" in proc.stdout assert "git fetch origin litellm_internal_staging" in proc.stdout + assert "check: FAIL" in proc.stdout def test_partial_staging_warns_which_checks_were_skipped(tmp_path: Path) -> None: @@ -384,3 +386,43 @@ def test_a_failing_block_fails_the_whole_run(tmp_path: Path, fail: str, message: proc = _run(repo, bin_dir, {"STUB_FAIL": fail}) assert proc.returncode == 1 assert message in proc.stdout + proc.stderr + + +def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "check: summary" in proc.stdout + assert "ran: Python lint (make lint)" in proc.stdout + assert "ran: dashboard lint (prettier + eslint + lint budgets)" in proc.stdout + assert "ran: dashboard API-type sync (npm run gen:api)" in proc.stdout + assert "skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope)" in proc.stdout + assert "check: PASS" in proc.stdout + assert "check: FAIL" not in proc.stdout + + +def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty_log(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + tests_dir = repo / "tests" / "test_litellm" + tests_dir.mkdir(parents=True) + (tests_dir / "test_x.py").write_text("def test_x() -> None: ...\n") + subprocess.run(["git", "add", "tests"], cwd=repo, check=True) + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "no gating lint check matches the files in scope, so nothing ran" in proc.stdout + assert "tests/test_litellm/test_x.py" in proc.stdout + assert "a no-op, not a lint verdict" in proc.stdout + assert "check: PASS" in proc.stdout + assert "linting Python" not in proc.stdout + log = (repo / ".git" / "pre_commit_lint.log").read_text() + assert "check: summary" in log + assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log + + +def test_failing_run_ends_with_a_fail_verdict(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + proc = _run(repo, bin_dir, {"STUB_FAIL": "make-lint"}) + assert proc.returncode == 1 + assert "check: FAIL" in proc.stdout + assert "check: PASS" not in proc.stdout diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 0104733bcf3..894d99c92e0 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22943 + "limit": 22941 }, "LIT002": { - "limit": 27141 + "limit": 27139 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16722 + "limit": 16716 }, "LIT011": { "limit": 5596 diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 5e322598a10..2b903f763d9 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -140,9 +140,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/purity": { "count": 1 }, @@ -228,7 +225,7 @@ "count": 2 }, "no-restricted-imports": { - "count": 2 + "count": 1 } }, "src/app/(dashboard)/cost-tracking/_components/how_it_works.tsx": { @@ -239,9 +236,6 @@ "src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.test.tsx": { @@ -252,9 +246,6 @@ "src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 2 } }, "src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_export_dropdown.tsx": { @@ -275,17 +266,11 @@ "src/app/(dashboard)/cost-tracking/_components/provider_discount_table.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/cost-tracking/_components/use_discount_config.ts": { @@ -313,11 +298,6 @@ "count": 3 } }, - "src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx": { "no-nested-ternary": { "count": 5 @@ -1405,9 +1385,6 @@ "no-nested-ternary": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "prefer-const": { "count": 2 } @@ -1499,9 +1476,6 @@ "src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx": { "local/no-complex-jsx-arrow": { "count": 2 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx": { @@ -1522,9 +1496,6 @@ "no-nested-ternary": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/purity": { "count": 1 }, @@ -1682,21 +1653,11 @@ "count": 1 } }, - "src/app/onboarding/OnboardingErrorView.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/onboarding/OnboardingFormBody.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/app/onboarding/OnboardingLoadingView.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/AIHub/ModelHubTable.test.tsx": { "max-params": { "count": 1 @@ -1709,63 +1670,30 @@ "no-nested-ternary": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "prefer-const": { "count": 4 } }, - "src/components/AIHub/SkillHubDashboard.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/AIHub/UsefulLinksManagement.tsx": { - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/components/AIHub/forms/MakeAgentPublicForm.tsx": { - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/AIHub/forms/MakeMCPPublicForm.test.tsx": { - "react/display-name": { - "count": 1 - } - }, "src/components/AIHub/forms/MakeMCPPublicForm.tsx": { - "no-nested-ternary": { - "count": 2 - }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/components/AIHub/forms/MakeModelPublicForm.tsx": { - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/BetaBadge.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx": { "no-restricted-imports": { "count": 1 @@ -1784,42 +1712,9 @@ "count": 1 } }, - "src/components/DebugWarningBanner.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/DeprecationBanner.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/EntityUsageExport/ExportSummary.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/EntityUsageExport/UsageExportHeader.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/EntityUsageExport/types.ts": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/EntityUsageExport/utils.test.ts": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/EntityUsageExport/utils.ts": { "max-params": { "count": 3 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/GuardrailSettingsView.tsx": { @@ -1830,9 +1725,6 @@ "src/components/GuardrailsMonitor/LogViewer.tsx": { "no-nested-ternary": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/HelpLink.test.tsx": { @@ -1840,54 +1732,16 @@ "count": 1 } }, - "src/components/LicenseExpiryBanner.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/ModelSelect/ModelSelect.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/Navbar/BlogDropdown/BlogDropdown.test.tsx": { "max-nested-callbacks": { "count": 12 } }, - "src/components/Navbar/BlogDropdown/BlogDropdown.tsx": { - "no-restricted-imports": { - "count": 2 - } - }, - "src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/Navbar/NotificationsBell/NotificationsBell.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/Navbar/UserDropdown/UserDropdown.tsx": { - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/Navbar/ViewSwitcher.tsx": { - "no-restricted-imports": { - "count": 2 - } - }, - "src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/SCIM.tsx": { "no-restricted-imports": { "count": 2 @@ -1983,9 +1837,6 @@ } }, "src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx": { - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -1995,11 +1846,6 @@ "count": 1 } }, - "src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx": { "local/no-complex-jsx-arrow": { "count": 1 @@ -2017,9 +1863,6 @@ } }, "src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx": { - "no-restricted-imports": { - "count": 2 - }, "prefer-const": { "count": 2 } @@ -2295,9 +2138,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -2341,9 +2181,6 @@ } }, "src/components/claude_code_plugins/MakeSkillPublicForm.tsx": { - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -2372,20 +2209,7 @@ "count": 2 } }, - "src/components/common_components/AutoRotationView.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/common_components/DefaultProxyAdminTag.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/common_components/DeleteResourceModal.tsx": { - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -2395,16 +2219,6 @@ "count": 1 } }, - "src/components/common_components/IconActionButton/BaseActionButton.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/common_components/KeyLifecycleSettings.tsx": { "local/no-complex-jsx-arrow": { "count": 1 @@ -2413,16 +2227,6 @@ "count": 2 } }, - "src/components/common_components/LabeledField.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/common_components/MemberTable.tsx": { - "no-restricted-imports": { - "count": 2 - } - }, "src/components/common_components/MetadataKeyValueFields.test.tsx": { "no-restricted-imports": { "count": 1 @@ -2434,34 +2238,15 @@ } }, "src/components/common_components/ModelAliasManager.tsx": { - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/components/common_components/ModelSelector.tsx": { - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/common_components/NewBadge.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/common_components/OrganizationDropdown.tsx": { - "local/no-complex-jsx-arrow": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/components/common_components/PassThroughGuardrailsSection.tsx": { "no-restricted-imports": { "count": 2 @@ -2470,29 +2255,11 @@ "count": 1 } }, - "src/components/common_components/PassThroughRoutesSelector.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/common_components/PassThroughSecuritySection.tsx": { "no-restricted-imports": { "count": 2 } }, - "src/components/common_components/PremiumLoggingSettings.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/common_components/ProjectDropdown.tsx": { - "local/no-complex-jsx-arrow": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/components/common_components/RateLimitTypeFormItem.test.tsx": { "no-restricted-imports": { "count": 1 @@ -2503,22 +2270,9 @@ "count": 1 } }, - "src/components/common_components/RouterSettingsAccordion.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/common_components/budget_duration_dropdown.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/common_components/chartUtils.test.tsx": { - "no-restricted-imports": { - "count": 1 } }, "src/components/common_components/chartUtils.tsx": { @@ -2527,9 +2281,6 @@ }, "no-nested-ternary": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/common_components/check_openapi_schema.tsx": { @@ -2554,17 +2305,11 @@ }, "no-nested-ternary": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/common_components/team_dropdown.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/common_components/team_multi_select.tsx": { @@ -2630,11 +2375,6 @@ "count": 1 } }, - "src/components/key_team_helpers/TagRateLimitEditor.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/key_team_helpers/fetch_available_models_team_key.tsx": { "local/filename-pascal-case": { "count": 1 @@ -2696,9 +2436,6 @@ "src/components/logging_settings_view.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/mcp_server_management/MCPServerSelector.tsx": { @@ -2712,9 +2449,6 @@ "src/components/mcp_server_management/MCPToolPermissions.tsx": { "local/no-complex-jsx-arrow": { "count": 1 - }, - "no-restricted-imports": { - "count": 2 } }, "src/components/mcp_tools/ByokCredentialModal.tsx": { @@ -2733,9 +2467,6 @@ "src/components/mcp_tools/McpCrudPermissionPanel.tsx": { "no-nested-ternary": { "count": 3 - }, - "no-restricted-imports": { - "count": 2 } }, "src/components/mcp_tools/types.tsx": { @@ -2764,18 +2495,12 @@ "src/components/model_filters.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/model_group_alias_settings.tsx": { "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -2837,9 +2562,6 @@ "src/components/navbar.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/networking.tsx": { @@ -2865,9 +2587,6 @@ "src/components/object_permissions_view.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/onboarding_link.tsx": { @@ -2914,9 +2633,6 @@ "src/components/organization/organization_view.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/page_utils.test.ts": { @@ -2940,22 +2656,9 @@ "count": 1 } }, - "src/components/permissions/AgentPermissions.tsx": { - "no-restricted-imports": { - "count": 2 - } - }, "src/components/permissions/MCPServerPermissions.tsx": { "no-nested-ternary": { "count": 3 - }, - "no-restricted-imports": { - "count": 2 - } - }, - "src/components/permissions/VectorStorePermissions.tsx": { - "no-restricted-imports": { - "count": 1 } }, "src/components/policies/PolicySelector.tsx": { @@ -2982,9 +2685,6 @@ }, "max-lines": { "count": 1 - }, - "no-restricted-imports": { - "count": 2 } }, "src/components/query_param_input.tsx": { @@ -2997,17 +2697,9 @@ "count": 1 } }, - "src/components/router_settings/LatencyBasedConfiguration.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/router_settings/ReliabilityRetriesSection.tsx": { "no-nested-ternary": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/router_settings/RoutingStrategySelector.tsx": { @@ -3015,18 +2707,10 @@ "count": 1 } }, - "src/components/router_settings/TagFilteringToggle.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/router_settings/index.tsx": { "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "prefer-const": { "count": 2 } @@ -3052,33 +2736,17 @@ "count": 1 } }, - "src/components/settings.test.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/settings.tsx": { "local/filename-pascal-case": { "count": 1 }, - "local/no-complex-jsx-arrow": { - "count": 4 - }, "no-nested-ternary": { "count": 2 }, - "no-restricted-imports": { - "count": 3 - }, "prefer-const": { "count": 4 } }, - "src/components/shared/CreatedKeyDisplay.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/shared/advanced_date_picker.tsx": { "local/filename-pascal-case": { "count": 1 @@ -3143,9 +2811,6 @@ "src/components/shared/numerical_input.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/shared/table_cells/cell_tooltip.tsx": { @@ -3242,9 +2907,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -3259,11 +2921,6 @@ "count": 1 } }, - "src/components/templates/KeyInfoHeader.tsx": { - "no-restricted-imports": { - "count": 2 - } - }, "src/components/templates/key_edit_view.tsx": { "local/filename-pascal-case": { "count": 1 @@ -3293,9 +2950,6 @@ "no-nested-ternary": { "count": 1 }, - "no-restricted-imports": { - "count": 2 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -3475,9 +3129,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -3486,9 +3137,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "prefer-const": { "count": 1 }, @@ -3506,29 +3154,15 @@ "count": 1 } }, - "src/components/view_logs/CostBreakdownViewer.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/view_logs/EvalViewer/EvalViewer.tsx": { - "local/no-complex-jsx-arrow": { - "count": 1 - }, "no-nested-ternary": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/view_logs/GuardrailViewer/CompliancePanel.tsx": { "no-nested-ternary": { "count": 2 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -3541,31 +3175,17 @@ "src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx": { "no-nested-ternary": { "count": 4 - }, - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx": { - "no-restricted-imports": { - "count": 1 } }, "src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx": { "no-nested-ternary": { "count": 3 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx": { "no-nested-ternary": { "count": 2 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 2 } @@ -3575,36 +3195,11 @@ "count": 2 } }, - "src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/view_logs/LogDetailsDrawer/useKeyboardNavigation.ts": { "react-hooks/immutability": { "count": 2 } }, - "src/components/view_logs/ToolsSection/FormattedToolView.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/view_logs/ToolsSection/ToolExpandedContent.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/view_logs/ToolsSection/ToolItem.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/view_logs/VectorStoreViewer.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/view_logs/columns.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.test.tsx index a1484ffb5c5..5e212901bd1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.test.tsx @@ -1,4 +1,5 @@ import { renderWithProviders, screen, within } from "@/../tests/test-utils"; +import { waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { AccessGroupsPage } from "./AccessGroupsPage"; @@ -215,7 +216,9 @@ describe("AccessGroupsPage", () => { await user.click(await openRowMenu(user, "ag-1")); const dialog = screen.getByRole("dialog", { name: "Delete Access Group" }); await user.click(within(dialog).getByRole("button", { name: "Cancel" })); - expect(screen.queryByRole("dialog", { name: "Delete Access Group" })).not.toBeInTheDocument(); + await waitFor(() => { + expect(screen.queryByRole("dialog", { name: "Delete Access Group" })).not.toBeInTheDocument(); + }); expect(mockMutate).not.toHaveBeenCalled(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx index 5167ac16542..e35c0103f7c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx @@ -1,4 +1,4 @@ -import { DateRangePickerValue } from "@tremor/react"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import React, { useEffect, useState } from "react"; import NotificationsManager from "@/components/molecules/notifications_manager"; import UsageDatePicker from "@/components/shared/usage_date_picker"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx index 0dae83ba808..03cef2a66b8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx @@ -1,6 +1,6 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; -import { screen } from "@testing-library/react"; +import { act, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "../../../../../tests/test-utils"; import CostTrackingSettings from "./cost_tracking_settings"; @@ -8,25 +8,29 @@ import CostTrackingSettings from "./cost_tracking_settings"; // Mock sub-hooks so we can control their state without network calls const mockDiscountConfig = vi.fn(() => ({})); const mockMarginConfig = vi.fn(() => ({})); +const mockRemoveDiscount = vi.fn(); +const mockRemoveMargin = vi.fn(); + +const stableDiscountCallbacks = { + fetchDiscountConfig: vi.fn().mockResolvedValue(undefined), + handleAddProvider: vi.fn().mockResolvedValue(true), + handleRemoveProvider: mockRemoveDiscount, + handleDiscountChange: vi.fn().mockResolvedValue(undefined), +}; + +const stableMarginCallbacks = { + fetchMarginConfig: vi.fn().mockResolvedValue(undefined), + handleAddMargin: vi.fn().mockResolvedValue(true), + handleRemoveMargin: mockRemoveMargin, + handleMarginChange: vi.fn().mockResolvedValue(undefined), +}; vi.mock("./use_discount_config", () => ({ - useDiscountConfig: () => ({ - discountConfig: mockDiscountConfig(), - fetchDiscountConfig: vi.fn().mockResolvedValue(undefined), - handleAddProvider: vi.fn().mockResolvedValue(true), - handleRemoveProvider: vi.fn().mockResolvedValue(undefined), - handleDiscountChange: vi.fn().mockResolvedValue(undefined), - }), + useDiscountConfig: () => ({ discountConfig: mockDiscountConfig(), ...stableDiscountCallbacks }), })); vi.mock("./use_margin_config", () => ({ - useMarginConfig: () => ({ - marginConfig: mockMarginConfig(), - fetchMarginConfig: vi.fn().mockResolvedValue(undefined), - handleAddMargin: vi.fn().mockResolvedValue(true), - handleRemoveMargin: vi.fn().mockResolvedValue(undefined), - handleMarginChange: vi.fn().mockResolvedValue(undefined), - }), + useMarginConfig: () => ({ marginConfig: mockMarginConfig(), ...stableMarginCallbacks }), })); vi.mock("./pricing_calculator/index", () => ({ @@ -153,6 +157,79 @@ describe("CostTrackingSettings", () => { }); }); + describe("removing a configured provider", () => { + const expandAndRemove = async (section: string, actionName: string) => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText(section).closest("button")!); + await user.click(await screen.findByRole("button", { name: actionName })); + + return user; + }; + + it("should ask to confirm before removing a discount", async () => { + mockDiscountConfig.mockReturnValue({ openai: 0.05 }); + + await expandAndRemove("Provider Discounts", "Remove discount for openai"); + + expect(await screen.findByRole("button", { name: "Remove" })).toBeInTheDocument(); + expect(screen.getByText(/are you sure you want to remove the discount for openai\?/i)).toBeInTheDocument(); + expect(mockRemoveDiscount).not.toHaveBeenCalled(); + }); + + it("should remove the discount once removal is confirmed", async () => { + mockDiscountConfig.mockReturnValue({ openai: 0.05 }); + + const user = await expandAndRemove("Provider Discounts", "Remove discount for openai"); + await user.click(await screen.findByRole("button", { name: "Remove" })); + + expect(mockRemoveDiscount).toHaveBeenCalledWith("openai"); + }); + + it("should leave the discount in place when the confirmation is cancelled", async () => { + mockDiscountConfig.mockReturnValue({ openai: 0.05 }); + + const user = await expandAndRemove("Provider Discounts", "Remove discount for openai"); + await user.click(await screen.findByRole("button", { name: "Cancel" })); + + expect(mockRemoveDiscount).not.toHaveBeenCalled(); + expect(screen.queryByRole("button", { name: "Remove" })).not.toBeInTheDocument(); + }); + + it("should hold the confirmation open while the removal is still in flight", async () => { + mockDiscountConfig.mockReturnValue({ openai: 0.05 }); + const { promise, resolve: settleRemoval } = Promise.withResolvers(); + mockRemoveDiscount.mockReturnValue(promise); + + const user = await expandAndRemove("Provider Discounts", "Remove discount for openai"); + await user.click(await screen.findByRole("button", { name: "Remove" })); + + const removing = await screen.findByRole("button", { name: "Removing…" }); + expect(removing).toBeDisabled(); + expect(screen.getByRole("button", { name: "Cancel" })).toBeDisabled(); + + await act(async () => { + settleRemoval(); + }); + + await waitFor(() => { + expect(screen.queryByRole("alertdialog")).not.toBeInTheDocument(); + }); + expect(mockRemoveDiscount).toHaveBeenCalledWith("openai"); + }); + + it("should remove the margin once removal is confirmed", async () => { + mockMarginConfig.mockReturnValue({ openai: 0.1 }); + + const user = await expandAndRemove("Fee/Price Margin", "Remove margin for openai"); + expect(screen.getByText(/are you sure you want to remove the margin for openai\?/i)).toBeInTheDocument(); + await user.click(await screen.findByRole("button", { name: "Remove" })); + + expect(mockRemoveMargin).toHaveBeenCalledWith("openai"); + }); + }); + describe("empty state messages", () => { it("should show the empty state message when no discount config is loaded", async () => { mockDiscountConfig.mockReturnValue({}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx index b32e7afd756..7f86bad3028 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx @@ -1,25 +1,24 @@ import React, { useState, useEffect } from "react"; -import { - Title, - Text, - Button, - Accordion, - AccordionHeader, - AccordionBody, - TabGroup, - TabList, - Tab, - TabPanels, - TabPanel, -} from "@tremor/react"; +import { ChevronDown } from "lucide-react"; import { Modal, Form } from "antd"; +import { + AlertDialog, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { Button } from "@/components/ui/button"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { CostTrackingSettingsProps } from "./types"; import ProviderDiscountTable from "./provider_discount_table"; import AddProviderForm from "./add_provider_form"; import ProviderMarginTable from "./provider_margin_table"; import AddMarginForm from "./add_margin_form"; import PricingCalculator from "./pricing_calculator/index"; -import { ExclamationCircleOutlined } from "@ant-design/icons"; import { DocsMenu } from "@/components/HelpLink"; import HowItWorks from "./how_it_works"; import { useDiscountConfig } from "./use_discount_config"; @@ -31,6 +30,29 @@ const DOCS_LINKS = [ { label: "Spend tracking", href: "https://docs.litellm.ai/docs/proxy/cost_tracking" }, ]; +const REMOVAL_COPY = { + discount: { title: "Remove Provider Discount", noun: "discount" }, + margin: { title: "Remove Provider Margin", noun: "margin" }, +} as const; + +interface PendingRemoval { + kind: keyof typeof REMOVAL_COPY; + provider: string; + displayName: string; +} + +const SECTION_HEADER_CLASS = "group/section flex w-full items-center justify-between px-6 py-4 text-left"; + +const SectionHeader: React.FC<{ title: string; description: string }> = ({ title, description }) => ( + +
+ {title} + {description} +
+ + +); + const CostTrackingSettings: React.FC = ({ userID, userRole, accessToken }) => { const [selectedProvider, setSelectedProvider] = useState(undefined); const [newDiscount, setNewDiscount] = useState(""); @@ -42,9 +64,10 @@ const CostTrackingSettings: React.FC = ({ userID, use const [percentageValue, setPercentageValue] = useState(""); const [fixedAmountValue, setFixedAmountValue] = useState(""); const [models, setModels] = useState([]); + const [pendingRemoval, setPendingRemoval] = useState(null); + const [isRemoving, setIsRemoving] = useState(false); const [form] = Form.useForm(); const [marginForm] = Form.useForm(); - const [modal, contextHolder] = Modal.useModal(); const isProxyAdmin = userRole === "proxy_admin" || userRole === "Admin"; @@ -104,16 +127,23 @@ const CostTrackingSettings: React.FC = ({ userID, use handleAddProvider(); }; - const handleRemoveProvider = async (provider: string, providerDisplayName: string) => { - modal.confirm({ - title: "Remove Provider Discount", - icon: , - content: `Are you sure you want to remove the discount for ${providerDisplayName}?`, - okText: "Remove", - okType: "danger", - cancelText: "Cancel", - onOk: () => removeProvider(provider), - }); + const handleRemoveProvider = (provider: string, providerDisplayName: string) => { + setPendingRemoval({ kind: "discount", provider, displayName: providerDisplayName }); + }; + + const handleConfirmRemoval = async () => { + if (!pendingRemoval) return; + setIsRemoving(true); + try { + if (pendingRemoval.kind === "discount") { + await removeProvider(pendingRemoval.provider); + } else { + await removeMargin(pendingRemoval.provider); + } + } finally { + setIsRemoving(false); + setPendingRemoval(null); + } }; const handleAddMargin = async () => { @@ -141,16 +171,8 @@ const CostTrackingSettings: React.FC = ({ userID, use setMarginType("percentage"); }; - const handleRemoveMargin = async (provider: string, providerDisplayName: string) => { - modal.confirm({ - title: "Remove Provider Margin", - icon: , - content: `Are you sure you want to remove the margin for ${providerDisplayName}?`, - okText: "Remove", - okType: "danger", - cancelText: "Cancel", - onOk: () => removeMargin(provider), - }); + const handleRemoveMargin = (provider: string, providerDisplayName: string) => { + setPendingRemoval({ kind: "margin", provider, displayName: providerDisplayName }); }; if (!accessToken) { @@ -159,18 +181,16 @@ const CostTrackingSettings: React.FC = ({ userID, use return (
- {contextHolder} - {/* Header Section - Outside the card */}
- Cost Tracking Settings +

Cost Tracking Settings

- +

Configure cost discounts and margins for different LLM providers. Changes are saved automatically. - +

@@ -178,90 +198,78 @@ const CostTrackingSettings: React.FC = ({ userID, use
{/* Accordion 1: Provider Discounts - Only for proxy admins */} {isProxyAdmin && ( - - -
- Provider Discounts - - Apply percentage-based discounts to reduce costs for specific providers - -
- - - - - Discounts - Test It - - - -
-
- + + + + + + Discounts + Test It + + +
+
+ +
+ {isFetching ? ( +
+

Loading configuration...

- {isFetching ? ( -
- Loading configuration... -
- ) : Object.keys(discountConfig).length > 0 ? ( - - ) : ( -
- - - - No provider discounts configured - - Click "Add Provider Discount" to get started - -
- )} -
- - -
- -
- - - - - + ) : Object.keys(discountConfig).length > 0 ? ( + + ) : ( +
+ + + +

No provider discounts configured

+

Click "Add Provider Discount" to get started

+
+ )} +
+ + +
+ +
+ + + + )} {/* Accordion 2: Fee/Price Margin - Only for proxy admins */} {isProxyAdmin && ( - - -
- Fee/Price Margin - - Add fees or margins to LLM costs for internal billing and cost recovery - -
- - + + +
{isFetching ? (
- Loading configuration... +

Loading configuration...

) : Object.keys(marginConfig).length > 0 ? ( = ({ userID, use d="M12 8c-1.657 0-3 .895-3 2s1.343 2 3 2 3 .895 3 2-1.343 2-3 2m0-8c1.11 0 2.08.402 2.599 1M12 8V7m0 1v8m0 0v1m0-1c-1.11 0-2.08-.402-2.599-1M21 12a9 9 0 11-18 0 9 9 0 0118 0z" /> - No provider margins configured - Click "Add Provider Margin" to get started +

No provider margins configured

+

Click "Add Provider Margin" to get started

)}
- - + + )} {/* Accordion 3: Pricing Calculator - Available to all roles */} - - -
- Pricing Calculator - - Estimate LLM costs based on expected token usage and request volume - -
- - + + +
- - + +
+ {pendingRemoval && ( + !open && !isRemoving && setPendingRemoval(null)}> + + + {REMOVAL_COPY[pendingRemoval.kind].title} + + Are you sure you want to remove the {REMOVAL_COPY[pendingRemoval.kind].noun} for{" "} + {pendingRemoval.displayName}? + + + + Cancel + + + + + )} + @@ -328,10 +352,10 @@ const CostTrackingSettings: React.FC = ({ userID, use }} >
- +

Select a provider and set its discount percentage. Enter a value between 0% and 100% (e.g., 5 for a 5% discount). - +

= ({ userID, use }} >
- +

Select a provider (or "Global" for all providers) and configure the margin. You can use percentage-based or fixed amount. - +

+ within(screen.getByRole("table")) + .getAllByRole("row") + .filter((row) => within(row).queryAllByRole("combobox").length > 0); + +const deleteButtonIn = (row: HTMLElement): HTMLElement => { + const cells = within(row).getAllByRole("cell"); + return within(cells[cells.length - 1]).getByRole("button"); +}; + describe("PricingCalculator", () => { beforeEach(() => { vi.clearAllMocks(); @@ -124,8 +134,31 @@ describe("PricingCalculator", () => { it("should render column headers for Model, Input Tokens, and Output Tokens", () => { renderWithProviders(); - expect(screen.getByText("Model")).toBeInTheDocument(); - expect(screen.getByText("Input Tokens")).toBeInTheDocument(); - expect(screen.getByText("Output Tokens")).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Model" })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Input Tokens" })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Output Tokens" })).toBeInTheDocument(); + }); + + it("should render a numeric field for input tokens, output tokens and requests", () => { + renderWithProviders(); + expect(screen.getAllByRole("spinbutton")).toHaveLength(3); + }); + + it("should offer a model picker per row", () => { + renderWithProviders(); + expect(screen.getAllByRole("combobox")).toHaveLength(1); + }); + + it("should remove a row when its delete button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /add another model/i })); + const withTwoRows = dataRows(); + expect(withTwoRows).toHaveLength(2); + + await user.click(deleteButtonIn(withTwoRows[1])); + + expect(dataRows()).toHaveLength(1); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx index 9b355e55c1c..f3bd74260ad 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx @@ -1,6 +1,10 @@ import React, { useState, useCallback } from "react"; -import { Table, Select, InputNumber, Button, Radio } from "antd"; -import { DeleteOutlined, PlusOutlined } from "@ant-design/icons"; +import { Plus, Trash2 } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; +import { Table, TableBody, TableCell, TableFooter, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { SearchSelect } from "@/components/shared/SearchSelect"; import { PricingCalculatorProps, ModelEntry } from "./types"; import MultiCostResults from "./multi_cost_results"; import { useMultiCostEstimate } from "./use_multi_cost_estimate"; @@ -63,132 +67,115 @@ const PricingCalculator: React.FC = ({ accessToken, mode const multiModelResult = getMultiModelResult(entries); - const columns = [ - { - title: "Model", - dataIndex: "model", - key: "model", - width: "35%", - render: (_: string, record: ModelEntry) => ( - + handleEntryChange(record.id, "input_tokens", e.target.value === "" ? 0 : Number(e.target.value)) + } + /> + + + + handleEntryChange(record.id, "output_tokens", e.target.value === "" ? 0 : Number(e.target.value)) + } + /> + + + + handleEntryChange( + record.id, + requestsField, + e.target.value === "" ? undefined : Number(e.target.value), + ) + } + /> + + + + + + ))} + + + + + + + + +
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.test.tsx index 04ef60469f0..b17dd2cb859 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.test.tsx @@ -85,6 +85,14 @@ function emptyMultiResult(): MultiModelResult { }; } +const expandToggle = (): HTMLElement => screen.getByRole("button", { name: /cost breakdown for / }); + +const shownBreakdown = (): HTMLElement | null => { + const label = screen.queryByText("Total/Request"); + if (label === null) return null; + return label.closest("[style*='display: none']") === null ? label : null; +}; + describe("MultiCostResults", () => { beforeEach(() => { vi.clearAllMocks(); @@ -200,40 +208,78 @@ describe("MultiCostResults", () => { expect(screen.getByRole("button", { name: /export/i })).toBeInTheDocument(); }); + it("should render a column header for each summary column", () => { + renderWithProviders(); + + expect(screen.getByRole("columnheader", { name: "Model" })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Per Request" })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Margin Fee" })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Daily" })).toBeInTheDocument(); + }); + + it("should not show the model breakdown before the row is expanded", () => { + renderWithProviders(); + expect(shownBreakdown()).toBeNull(); + }); + it("should expand the model breakdown row when the expand button is clicked", async () => { const user = userEvent.setup(); renderWithProviders(); - // The expand column renders a button (RightOutlined icon) for rows without errors - const expandButtons = screen.getAllByRole("button"); - // Find the small expand button (not the Export button) - const expandButton = expandButtons.find((btn) => !btn.textContent?.toLowerCase().includes("export")); - expect(expandButton).toBeDefined(); + await user.click(expandToggle()); - await user.click(expandButton!); - - // After expanding, the SingleModelBreakdown should be visible - expect(screen.getByText("Total/Request")).toBeInTheDocument(); + expect(shownBreakdown()).toBeVisible(); + expect(screen.getByText("Daily Total (100 req)")).toBeInTheDocument(); }); - it("should show the collapse icon after expanding a row", async () => { + it("should collapse the model breakdown again on a second click", async () => { const user = userEvent.setup(); renderWithProviders(); - const getExpandButton = () => { - const allButtons = screen.getAllByRole("button"); - return allButtons.find((btn) => !btn.textContent?.toLowerCase().includes("export")); - }; + await user.click(expandToggle()); + expect(shownBreakdown()).toBeVisible(); - // Before expand: button has the "down" aria-label (RightOutlined renders as down in ant icons) - // Just verify clicking works and the breakdown content appears - await user.click(getExpandButton()!); - expect(screen.getByText("Total/Request")).toBeInTheDocument(); + await user.click(expandToggle()); + expect(shownBreakdown()).toBeNull(); + }); - // After a second click, the row collapses — content may be hidden or removed - await user.click(getExpandButton()!); - // The expanded content should no longer be visible - expect(screen.queryByText("Total/Request")).not.toBeVisible(); + it("should name the breakdown toggle and report its expanded state", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const toggle = screen.getByRole("button", { name: "Show cost breakdown for gpt-4" }); + expect(toggle).toHaveAttribute("aria-expanded", "false"); + + await user.click(toggle); + + const collapseToggle = screen.getByRole("button", { name: "Hide cost breakdown for gpt-4" }); + expect(collapseToggle).toHaveAttribute("aria-expanded", "true"); + }); + + it("should not offer an expand toggle for a row that failed", () => { + renderWithProviders( + , + ); + + expect(screen.getAllByRole("button", { name: /cost breakdown for / })).toHaveLength(1); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.tsx index 3ea7ea58127..b8375b930c9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.tsx @@ -1,7 +1,11 @@ import React, { useState } from "react"; -import { Text, Button } from "@tremor/react"; -import { Card, Statistic, Row, Col, Divider, Spin, Table, Tag } from "antd"; -import { LoadingOutlined, DownOutlined, RightOutlined } from "@ant-design/icons"; +import { ChevronDown, ChevronRight } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Card } from "@/components/ui/card"; +import { Separator } from "@/components/ui/separator"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { CostEstimateResponse } from "../types"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { MultiModelResult } from "./types"; @@ -41,55 +45,57 @@ const SingleModelBreakdown: React.FC<{
{loading && (
- } size="small" /> + Updating...
)}
-
- Total/Request - {formatCost(result.cost_per_request)} +
+

Total/Request

+

{formatCost(result.cost_per_request)}

-
- Input Cost - {formatCost(result.input_cost_per_request)} +
+

Input Cost

+

{formatCost(result.input_cost_per_request)}

-
- Output Cost - {formatCost(result.output_cost_per_request)} +
+

Output Cost

+

{formatCost(result.output_cost_per_request)}

-
- Margin Fee - 0 ? "text-amber-600" : ""}`}> +
+

Margin Fee

+

0 ? "text-amber-600" : ""}`}> {formatCost(result.margin_cost_per_request)} - +

{periodCost !== null && (
-
- +
+

{periodLabel} Total ({formatRequests(periodRequests)} req) - - +

+

{formatCost(periodCost)} - +

-
- {periodLabel} Input - {formatCost(periodInputCost)} +
+

{periodLabel} Input

+

{formatCost(periodInputCost)}

-
- {periodLabel} Output - {formatCost(periodOutputCost)} +
+

{periodLabel} Output

+

{formatCost(periodOutputCost)}

-
- {periodLabel} Margin Fee - 0 ? "text-amber-600" : ""}`}> +
+

{periodLabel} Margin Fee

+

0 ? "text-amber-600" : ""}`}> {formatCost(periodMarginCost)} - +

)} @@ -124,7 +130,7 @@ const MultiCostResults: React.FC = ({ multiResult, timePe if (!hasAnyResult && !isAnyLoading && !hasAnyError) { return (
- Select models above to see cost estimates +

Select models above to see cost estimates

); } @@ -133,8 +139,8 @@ const MultiCostResults: React.FC = ({ multiResult, timePe if (!hasAnyResult && isAnyLoading && !hasAnyError) { return (
- } /> - Calculating costs... + +

Calculating costs...

); } @@ -143,10 +149,10 @@ const MultiCostResults: React.FC = ({ multiResult, timePe if (!hasAnyResult && hasAnyError) { return (
- +
- Cost Estimates - {isAnyLoading && } size="small" />} +

Cost Estimates

+ {isAnyLoading && }
{/* Error Messages */} {errorEntries.map((e) => ( @@ -174,102 +180,10 @@ const MultiCostResults: React.FC = ({ multiResult, timePe const hasMargin = multiResult.totals.margin_per_request > 0; const periodLabel = timePeriod === "day" ? "Daily" : "Monthly"; - const periodCostKey = timePeriod === "day" ? "daily_cost" : "monthly_cost"; - - const summaryColumns = [ - { - title: "Model", - dataIndex: "model", - key: "model", - render: ( - text: string, - record: { - id: string; - provider?: string | null; - error?: string | null; - loading?: boolean; - hasZeroCost?: boolean | null; - }, - ) => ( -
-
- {text} - {record.provider && ( - - {record.provider} - - )} - {record.loading && } size="small" />} -
- {record.error &&
⚠️ {record.error}
} - {record.hasZeroCost && !record.error && ( -
- ⚠️ No pricing data found for this model. Set base_model in config. -
- )} -
- ), - }, - { - title: "Per Request", - dataIndex: "cost_per_request", - key: "cost_per_request", - align: "right" as const, - render: (value: number | null, record: { error?: string | null }) => - record.error ? ( - - - ) : ( - {formatCost(value)} - ), - }, - { - title: "Margin Fee", - dataIndex: "margin_cost_per_request", - key: "margin_cost_per_request", - align: "right" as const, - render: (value: number | null, record: { error?: string | null }) => - record.error ? ( - - - ) : ( - 0 ? "text-amber-600" : "text-gray-400"}`}> - {formatCost(value)} - - ), - }, - { - title: periodLabel, - dataIndex: periodCostKey, - key: "period_cost", - align: "right" as const, - render: (value: number | null, record: { error?: string | null }) => - record.error ? ( - - - ) : ( - {formatCost(value)} - ), - }, - { - title: "", - key: "expand", - width: 40, - render: (_: unknown, record: { id: string; error?: string | null }) => - record.error ? null : ( - - ), - }, - ]; // Include both valid results and errors in the table data const allEntriesWithModels = multiResult.entries.filter((e) => e.entry.model); const summaryData = allEntriesWithModels.map((e) => ({ - key: e.entry.id, id: e.entry.id, model: e.result?.model || e.entry.model, provider: e.result?.provider, @@ -284,78 +198,153 @@ const MultiCostResults: React.FC = ({ multiResult, timePe return (
- +
- Cost Estimates +

Cost Estimates

- {isAnyLoading && } size="small" />} + {isAnyLoading && }
{/* Combined Totals - Always show when there are results */} - - - - Total Per Request} - value={formatCost(multiResult.totals.cost_per_request)} - valueStyle={{ color: "#1890ff", fontSize: "18px", fontFamily: "monospace" }} - /> - - - Total {periodLabel}} - value={formatCost(timePeriod === "day" ? multiResult.totals.daily_cost : multiResult.totals.monthly_cost)} - valueStyle={{ - color: timePeriod === "day" ? "#52c41a" : "#722ed1", - fontSize: "18px", - fontFamily: "monospace", - }} - /> - - + +
+
+ Total Per Request +
+ {formatCost(multiResult.totals.cost_per_request)} +
+
+
+ Total {periodLabel} +
+ {formatCost(timePeriod === "day" ? multiResult.totals.daily_cost : multiResult.totals.monthly_cost)} +
+
+
{hasMargin && ( - - +
+
Margin Fee/Request
-
+
{formatCost(multiResult.totals.margin_per_request)}
- - +
+
{periodLabel} Margin Fee
-
+
{formatCost(timePeriod === "day" ? multiResult.totals.daily_margin : multiResult.totals.monthly_margin)}
- - +
+
)} {/* Per-Model Table */} {summaryData.length > 0 && ( - { - const entry = validEntries.find((e) => e.entry.id === record.id); - if (!entry?.result) return null; +
+ + + Model + Per Request + Margin Fee + {periodLabel} + + Cost breakdown + + + + + {summaryData.map((record) => { + const isExpanded = expandedModels.has(record.id); + const periodCost = timePeriod === "day" ? record.daily_cost : record.monthly_cost; + const breakdownEntry = validEntries.find((e) => e.entry.id === record.id); return ( -
- -
+ + + +
+
+ {record.model} + {record.provider && ( + + {record.provider} + + )} + {record.loading && } +
+ {record.error && ( +
⚠️ {record.error}
+ )} + {record.hasZeroCost && !record.error && ( +
+ ⚠️ No pricing data found for this model. Set base_model in config. +
+ )} +
+ + + {record.error ? ( + - + ) : ( + {formatCost(record.cost_per_request)} + )} + + + {record.error ? ( + - + ) : ( + 0 ? "text-amber-600" : "text-gray-400"}`} + > + {formatCost(record.margin_cost_per_request)} + + )} + + + {record.error ? ( + - + ) : ( + {formatCost(periodCost)} + )} + + + {!record.error && ( + + )} + + + {isExpanded && breakdownEntry?.result && ( + + +
+ +
+ + + )} + ); - }, - showExpandColumn: false, - }} - /> + })} + +
)}
); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx index f9a0a40f07d..24280873cf0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx @@ -5,49 +5,21 @@ import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "../../../../../tests/test-utils"; import ProviderDiscountTable from "./provider_discount_table"; -vi.mock("@heroicons/react/outline", () => ({ - TrashIcon: function TrashIcon() { - return null; - }, - PencilAltIcon: function PencilAltIcon() { - return null; - }, - CheckIcon: function CheckIcon() { - return null; - }, - XIcon: function XIcon() { - return null; - }, -})); - -vi.mock("@tremor/react", () => ({ - Table: ({ children }: any) => {children}
, - TableHead: ({ children }: any) => {children}, - TableRow: ({ children }: any) => {children}, - TableHeaderCell: ({ children }: any) => {children}, - TableBody: ({ children }: any) => {children}, - TableCell: ({ children }: any) => {children}, - Text: ({ children }: any) => {children}, - TextInput: ({ value, onValueChange, onKeyDown, placeholder, ...rest }: any) => ( - onValueChange?.(e.target.value)} - onKeyDown={onKeyDown} - placeholder={placeholder} - {...rest} - /> - ), - Icon: ({ icon: IconComponent, onClick }: any) => { - const name = IconComponent?.displayName ?? IconComponent?.name ?? "icon"; - return + + + ) : ( + <> +

{(row.discount * 100).toFixed(1)}%

+ + + )} +
+ ); + }, width: "250px", }, { @@ -125,12 +138,15 @@ const ProviderDiscountTable: React.FC = ({ cell: (row) => { const { displayName } = getProviderLogoAndName(row.provider); return ( - onRemoveProvider(row.provider, displayName)} className="cursor-pointer hover:text-red-600" - /> + > + + ); }, width: "80px", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.test.tsx index 170e61141b6..dd478571568 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.test.tsx @@ -6,43 +6,15 @@ import { renderWithProviders } from "../../../../../tests/test-utils"; import ProviderMarginTable from "./provider_margin_table"; import { Providers, providerLogoMap } from "@/components/provider_info_helpers"; -vi.mock("@heroicons/react/outline", () => ({ - TrashIcon: function TrashIcon() { - return null; - }, - PencilAltIcon: function PencilAltIcon() { - return null; - }, - CheckIcon: function CheckIcon() { - return null; - }, - XIcon: function XIcon() { - return null; - }, -})); +const ROW_ACTION_NAME = { + edit: /^Edit margin for /, + save: /^Save margin for /, + cancel: /^Cancel editing margin for /, + remove: /^Remove margin for /, +} as const; -vi.mock("@tremor/react", () => ({ - Table: ({ children }: any) => {children}
, - TableHead: ({ children }: any) => {children}, - TableRow: ({ children }: any) => {children}, - TableHeaderCell: ({ children }: any) => {children}, - TableBody: ({ children }: any) => {children}, - TableCell: ({ children }: any) => {children}, - Text: ({ children }: any) => {children}, - TextInput: ({ value, onValueChange, placeholder, autoFocus, className }: any) => ( - onValueChange?.(e.target.value)} - placeholder={placeholder} - autoFocus={autoFocus} - className={className} - /> - ), - Icon: ({ icon: IconComponent, onClick }: any) => { - const name = IconComponent?.displayName ?? IconComponent?.name ?? "icon"; - return + + + ) : ( + <> +

{formatMargin(row.margin)}

+ + + )} +
+ ); + }, width: "350px", }, { header: "Actions", cell: (row) => { - const displayName = row.provider === "global" ? "Global" : getProviderLogoAndName(row.provider).displayName; + const displayName = marginRowDisplayName(row.provider); return ( - onRemoveProvider(row.provider, displayName)} className="cursor-pointer hover:text-red-600" - /> + > + + ); }, width: "80px", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx index 14849438135..3b5cc156242 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx @@ -1,4 +1,4 @@ -import type { DateRangePickerValue } from "@tremor/react"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import React, { useCallback, useMemo, useState } from "react"; import { formatDate } from "@/components/networking"; import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts index 485c7cc1f92..02c82611848 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts @@ -17,6 +17,10 @@ export enum ConfigType { */ export enum GeneralSettingsFieldName { MAXIMUM_SPEND_LOGS_RETENTION_PERIOD = "maximum_spend_logs_retention_period", + MAXIMUM_SPEND_LOGS_CLEANUP_BATCH_SIZE = "maximum_spend_logs_cleanup_batch_size", + MAXIMUM_SPEND_LOGS_CLEANUP_MAX_BATCHES = "maximum_spend_logs_cleanup_max_batches", + MAXIMUM_SPEND_LOGS_CLEANUP_RUN_BUDGET = "maximum_spend_logs_cleanup_run_budget", + MAXIMUM_SPEND_LOGS_CLEANUP_BATCH_TIMEOUT = "maximum_spend_logs_cleanup_batch_timeout", // Add more field names here as needed } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts index 88a37b30291..ca57c01679d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts @@ -6,6 +6,10 @@ import { proxyConfigKeys } from "../proxyConfig/useProxyConfig"; export interface StoreRequestInSpendLogsParams { store_prompts_in_spend_logs: boolean; maximum_spend_logs_retention_period?: string; + maximum_spend_logs_cleanup_batch_size?: number; + maximum_spend_logs_cleanup_max_batches?: number; + maximum_spend_logs_cleanup_run_budget?: string; + maximum_spend_logs_cleanup_batch_timeout?: string; } export interface StoreRequestInSpendLogsResponse { @@ -19,6 +23,8 @@ const performStoreRequestInSpendLogs = async ( const proxyBaseUrl = getProxyBaseUrl(); const url = proxyBaseUrl ? `${proxyBaseUrl}/config/update` : `/config/update`; + const { store_prompts_in_spend_logs, ...optionalSettings } = params; + const response = await fetch(url, { method: "POST", headers: { @@ -27,10 +33,8 @@ const performStoreRequestInSpendLogs = async ( }, body: JSON.stringify({ general_settings: { - store_prompts_in_spend_logs: params.store_prompts_in_spend_logs, - ...(params.maximum_spend_logs_retention_period && { - maximum_spend_logs_retention_period: params.maximum_spend_logs_retention_period, - }), + store_prompts_in_spend_logs, + ...optionalSettings, }, }), }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.test.tsx index 9ffcbfc9975..0f3ba6c3471 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.test.tsx @@ -62,6 +62,8 @@ const settingsRow = async (fieldName: string) => { return row as HTMLElement; }; +const numericValueIn = (row: HTMLElement) => Number((within(row).getByRole("spinbutton") as HTMLInputElement).value); + describe("GeneralSettings General tab", () => { beforeEach(() => { vi.mocked(getGeneralSettingsCall).mockResolvedValue([...SETTINGS_FIXTURE.map((s) => ({ ...s }))]); @@ -87,7 +89,7 @@ describe("GeneralSettings General tab", () => { await user.click(screen.getByText("General")); const row = await settingsRow("max_ui_session_budget"); - expect(within(row).getByRole("spinbutton")).toHaveValue("7.50"); + expect(numericValueIn(row)).toBe(7.5); const actionCell = row.querySelectorAll("td")[3]; const resetIcon = actionCell.querySelector("svg"); @@ -95,7 +97,7 @@ describe("GeneralSettings General tab", () => { await user.click(resetIcon as unknown as Element); expect(deleteConfigFieldSetting).toHaveBeenCalledWith("token", "max_ui_session_budget"); - expect(within(row).getByRole("spinbutton")).toHaveValue("1.00"); + expect(numericValueIn(row)).toBe(1); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index ed7b17067d5..2a5b94b2fd7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -1,22 +1,14 @@ import React, { useState, useEffect } from "react"; -import { - Card, - Table, - TableHead, - TableRow, - TableHeaderCell, - TableCell, - TableBody, - Title, - Text, - Button, - Icon, - Switch, -} from "@tremor/react"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardTitle } from "@/components/ui/card"; +import { Input } from "@/components/ui/input"; +import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Switch } from "@/components/ui/switch"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { getGeneralSettingsCall, updateConfigFieldSetting, deleteConfigFieldSetting } from "@/components/networking"; -import { InputNumber, Select as AntdSelect } from "antd"; -import { TrashIcon } from "@heroicons/react/outline"; +import { Trash2 } from "lucide-react"; import { StatusBadge } from "@/components/shared/table_cells"; import RouterSettings from "@/components/router_settings"; @@ -44,16 +36,22 @@ export interface generalSettingsItem { field_default_value?: any; } +const NUMERIC_INPUT_WIDTH = "w-36"; + +const toNumericValue = (raw: string): number | null => (raw === "" ? null : Number(raw)); + const SettingValueEditor: React.FC<{ setting: generalSettingsItem; onChange: (fieldName: string, newValue: any) => void; }> = ({ setting, onChange }) => { if (setting.field_type === "Integer") { return ( - onChange(setting.field_name, newValue)} + className={NUMERIC_INPUT_WIDTH} + value={setting.field_value ?? ""} + onChange={(event) => onChange(setting.field_name, toNumericValue(event.target.value))} /> ); } @@ -61,42 +59,55 @@ const SettingValueEditor: React.FC<{ return ( onChange(setting.field_name, checked)} + onCheckedChange={(checked) => onChange(setting.field_name, checked)} /> ); } if (setting.field_type === "Float") { return ( - onChange(setting.field_name, newValue)} + className={NUMERIC_INPUT_WIDTH} + value={setting.field_value ?? ""} + onChange={(event) => onChange(setting.field_name, toNumericValue(event.target.value))} /> ); } if (setting.field_type === "Dollar") { return ( - onChange(setting.field_name, newValue)} - /> + + $ + onChange(setting.field_name, toNumericValue(event.target.value))} + /> + ); } if (setting.field_type === "Select") { return ( - ({ label: option, value: option }))} - onChange={(newValue) => onChange(setting.field_name, newValue ?? "")} - /> + ); } return null; @@ -131,33 +142,43 @@ export const PromptCachingPanel: React.FC<{ return ( - Prompt Caching + + Prompt Caching -
-
- Automatic Anthropic prompt caching -

{enableSetting.field_description}

-
- persist(ENABLE_ANTHROPIC_PROMPT_CACHING, checked)} /> -
- - {ttlSetting && (
-
- Cache lifetime (TTL) -

{ttlSetting.field_description}

+
+

Automatic Anthropic prompt caching

+

{enableSetting.field_description}

- ({ label: option, value: option }))} - onChange={(newValue) => persist(ANTHROPIC_PROMPT_CACHING_TTL, newValue ?? "")} - /> + persist(ENABLE_ANTHROPIC_PROMPT_CACHING, checked)} />
- )} + + {ttlSetting && ( +
+
+

Cache lifetime (TTL)

+

{ttlSetting.field_description}

+
+ +
+ )} + ); }; @@ -254,55 +275,60 @@ const GeneralSettings: React.FC = ({ accessToken, user - - - - Setting - Value - Status - Action - - - - {generalSettings - .filter((value) => value.field_type !== "TypedDictionary" && value.field_tab !== PROMPT_CACHING_TAB) - .map((value, index) => ( - - - {value.field_name} -

- {value.field_description} -

- - - - - - {value.stored_in_db == true ? ( - - ) : value.stored_in_db == false ? ( - - ) : ( - - )} - - - - handleResetField(value.field_name)}> - Reset - - - - ))} - -
+ + + + + Setting + Value + Status + Action + + + + {generalSettings + .filter((value) => value.field_type !== "TypedDictionary" && value.field_tab !== PROMPT_CACHING_TAB) + .map((value, index) => ( + + +

{value.field_name}

+

+ {value.field_description} +

+ + + + + + {value.stored_in_db == true ? ( + + ) : value.stored_in_db == false ? ( + + ) : ( + + )} + + + + handleResetField(value.field_name)} + className="inline-flex shrink-0 cursor-pointer items-center justify-center px-1.5 py-1.5 text-red-500" + > + + + + + ))} + +
+ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 27c0826d880..d7d971768b1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -14,7 +14,7 @@ import { MoneyCell } from "@/components/shared/table_cells"; import { Card as ShadcnCard, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { hasCapability, type Capability } from "@/utils/capabilities"; import { formatNumberWithCommas } from "@/utils/dataUtils"; -import type { DateRangePickerValue } from "@tremor/react"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import { ChevronDown, ChevronRight, ExternalLink, Info, Loader2 } from "lucide-react"; import type { ColumnDef } from "@tanstack/react-table"; import { Alert, AlertDescription } from "@/components/shared/Alert"; 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 4eeff926ad7..f1bfdc908c7 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 @@ -7,7 +7,7 @@ */ import { ChevronDown, ChevronRight, Download, ExternalLink, Info, Loader2, Sparkles, X } from "lucide-react"; -import type { DateRangePickerValue } from "@tremor/react"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import React, { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { BarChart } from "@/components/shared/charts"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx index 237e3b2c842..29c9b17cb21 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx @@ -244,7 +244,7 @@ describe("UserInfoView", () => { }); // The DeleteResourceModal's OK button has text "Delete" - find it within the modal - const modal = screen.getByText("Remove from Team").closest(".ant-modal") as HTMLElement; + const modal = screen.getByRole("dialog", { name: "Remove from Team" }); const deleteConfirmButton = within(modal).getByRole("button", { name: /delete/i }); await user.click(deleteConfirmButton); diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.test.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.test.tsx index bfbfb8fdd5d..f0bfdfef40d 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.test.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.test.tsx @@ -9,6 +9,11 @@ describe("OnboardingErrorView", () => { expect(screen.getByText("Failed to load invitation")).toBeInTheDocument(); }); + it("should expose the failure as an alert to assistive technology", () => { + render(); + expect(screen.getByRole("alert")).toHaveTextContent("Failed to load invitation"); + }); + it("should show the expiry description", () => { render(); expect(screen.getByText("The invitation link may be invalid or expired.")).toBeInTheDocument(); @@ -16,7 +21,6 @@ describe("OnboardingErrorView", () => { it("should render a Back to Login link pointing to /ui/login/", () => { render(); - // antd Button with href renders as an element const link = screen.getByRole("link", { name: "Back to Login" }); expect(link).toHaveAttribute("href", "/ui/login/"); }); diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.tsx index 3de9a9ffaae..2fa35205a82 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.tsx @@ -1,18 +1,21 @@ import React from "react"; -import { Alert, Button } from "antd"; +import { CircleAlert } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { buttonVariants } from "@/components/ui/button"; import { getLoginUrl } from "@/utils/returnUrlUtils"; export function OnboardingErrorView() { return ( ); diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingLoadingView.test.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingLoadingView.test.tsx index 21c5ccf69d0..755647fa3fc 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingLoadingView.test.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingLoadingView.test.tsx @@ -1,12 +1,12 @@ import React from "react"; -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, it, expect } from "vitest"; import { OnboardingLoadingView } from "./OnboardingLoadingView"; describe("OnboardingLoadingView", () => { - it("should render a spinner container", () => { - const { container } = render(); - expect(container.firstChild).toBeInTheDocument(); + it("should expose the loading state to assistive technology", () => { + render(); + expect(screen.getByRole("status", { name: "Loading invitation" })).toBeInTheDocument(); }); it("should apply centering layout classes", () => { diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingLoadingView.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingLoadingView.tsx index 7efa1d2504f..4d5d2a1371e 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingLoadingView.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingLoadingView.tsx @@ -1,11 +1,10 @@ import React from "react"; -import { Spin } from "antd"; -import { LoadingOutlined } from "@ant-design/icons"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; export function OnboardingLoadingView() { return (
- } size="large" /> +
); } diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 0e4de3f244c..28dac9badcd 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -23,10 +23,12 @@ import { import PublicModelHub from "@/components/public_model_hub"; import { copyToClipboard } from "@/utils/dataUtils"; import { isAdminRole, isProxyAdminRole } from "@/utils/roles"; -import { CopyOutlined } from "@ant-design/icons"; import { SortingState } from "@tanstack/react-table"; -import { Badge, Button, Card, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; -import { Modal } from "antd"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Card } from "@/components/ui/card"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Copy, Inbox } from "lucide-react"; import { useRouter } from "next/navigation"; import React, { useCallback, useEffect, useMemo, useState } from "react"; @@ -393,7 +395,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {/* Header with Title, Description and URL */}
- AI Hub +

AI Hub

{isAdminRole(userRole || "") ? (

Make models, agents, and MCP servers public for developers to know what's available. @@ -403,9 +405,9 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, )}

- Model Hub URL: +

Model Hub URL:

- {`${getProxyBaseUrl()}/ui/model_hub_table`} +

{`${getProxyBaseUrl()}/ui/model_hub_table`}

@@ -568,167 +570,162 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, setSkillHubData(response.plugins); }} /> - - - + +
+
) : ( - - Public Model Hub not enabled. + +

Public Model Hub not enabled.

Ask your proxy admin to enable this on their Admin UI.

 )} {/* Public Page Modal */} - -
-
- Shareable Link: - - {`${getProxyBaseUrl()}/ui/model_hub_table`} - + !open && handleCancel()}> + + + {"Public Model Hub"} + +
+
+

Shareable Link:

+

+ {`${getProxyBaseUrl()}/ui/model_hub_table`} +

+
+
+ +
-
- -
-
- + + {/* Model Details Modal */} - - {selectedModel && ( -
- {/* Model Overview */} -
- Model Overview -
-
- Model Group: - {selectedModel.model_group} + !open && handleCancel()}> + + + {selectedModel?.model_group || "Model Details"} + + {selectedModel && ( +
+ {/* Model Overview */} +
+

Model Overview

+
+
+

Model Group:

+

{selectedModel.model_group}

+
+
+

Mode:

+

{selectedModel.mode || "Not specified"}

+
+
+

Providers:

+
+ {selectedModel.providers.map((provider) => ( + + {provider} + + ))} +
+
-
- Mode: - {selectedModel.mode || "Not specified"} +
+ + {/* Token and Cost Information */} +
+

Token & Cost Information

+
+
+

Max Input Tokens:

+

{selectedModel.max_input_tokens?.toLocaleString() || "Not specified"}

+
+
+

Max Output Tokens:

+

{selectedModel.max_output_tokens?.toLocaleString() || "Not specified"}

+
+
+

Input Cost per 1M Tokens:

+

+ {selectedModel.input_cost_per_token + ? formatCost(selectedModel.input_cost_per_token) + : "Not specified"} +

+
+
+

Output Cost per 1M Tokens:

+

+ {selectedModel.output_cost_per_token + ? formatCost(selectedModel.output_cost_per_token) + : "Not specified"} +

+
+
+ + {/* Capabilities */} +
+

Capabilities

+
+ {(() => { + const capabilities = getModelCapabilities(selectedModel); + const colors = ["green", "blue", "purple", "orange", "red", "yellow"]; + + if (capabilities.length === 0) { + return

No special capabilities listed

; + } + + return capabilities.map((capability, index) => ( + + {formatCapabilityName(capability)} + + )); + })()} +
+
+ + {/* Rate Limits */} + {(selectedModel.tpm || selectedModel.rpm) && (
- Providers: -
- {selectedModel.providers.map((provider) => ( - - {provider} +

Rate Limits

+
+ {selectedModel.tpm && ( +
+

Tokens per Minute:

+

{selectedModel.tpm.toLocaleString()}

+
+ )} + {selectedModel.rpm && ( +
+

Requests per Minute:

+

{selectedModel.rpm.toLocaleString()}

+
+ )} +
+
+ )} + + {/* Supported OpenAI Parameters */} + {selectedModel.supported_openai_params && ( +
+

Supported OpenAI Parameters

+
+ {selectedModel.supported_openai_params.map((param) => ( + + {param} ))}
-
-
+ )} - {/* Token and Cost Information */} -
- Token & Cost Information -
-
- Max Input Tokens: - {selectedModel.max_input_tokens?.toLocaleString() || "Not specified"} -
-
- Max Output Tokens: - {selectedModel.max_output_tokens?.toLocaleString() || "Not specified"} -
-
- Input Cost per 1M Tokens: - - {selectedModel.input_cost_per_token - ? formatCost(selectedModel.input_cost_per_token) - : "Not specified"} - -
-
- Output Cost per 1M Tokens: - - {selectedModel.output_cost_per_token - ? formatCost(selectedModel.output_cost_per_token) - : "Not specified"} - -
-
-
- - {/* Capabilities */} -
- Capabilities -
- {(() => { - const capabilities = getModelCapabilities(selectedModel); - const colors = ["green", "blue", "purple", "orange", "red", "yellow"]; - - if (capabilities.length === 0) { - return No special capabilities listed; - } - - return capabilities.map((capability, index) => ( - - {formatCapabilityName(capability)} - - )); - })()} -
-
- - {/* Rate Limits */} - {(selectedModel.tpm || selectedModel.rpm) && ( + {/* Usage Example */}
- Rate Limits -
- {selectedModel.tpm && ( -
- Tokens per Minute: - {selectedModel.tpm.toLocaleString()} -
- )} - {selectedModel.rpm && ( -
- Requests per Minute: - {selectedModel.rpm.toLocaleString()} -
- )} -
-
- )} - - {/* Supported OpenAI Parameters */} - {selectedModel.supported_openai_params && ( -
- Supported OpenAI Parameters -
- {selectedModel.supported_openai_params.map((param) => ( - - {param} - - ))} -
-
- )} - - {/* Usage Example */} -
- Usage Example - - {`import openai +

Usage Example

+ + {`import openai client = openai.OpenAI( api_key="your_api_key", @@ -746,316 +743,310 @@ response = client.chat.completions.create( ) print(response.choices[0].message.content)`} - + +
-
- )} - + )} + + {/* Agent Details Modal */} - - {selectedAgent && ( -
- {/* Agent Overview */} -
- Agent Overview -
-
- Name: - {selectedAgent.name} -
-
- Version: - v{selectedAgent.version} -
-
- Protocol Version: - {selectedAgent.protocolVersion} -
-
- URL: -
- {selectedAgent.url} - void copyToClipboard(selectedAgent.url)} - className="cursor-pointer text-gray-500 hover:text-blue-500" - /> + !open && handleCancel()}> + + + {selectedAgent?.name || "Agent Details"} + + {selectedAgent && ( +
+ {/* Agent Overview */} +
+

Agent Overview

+
+
+

Name:

+

{selectedAgent.name}

-
-
-
- Description: - {selectedAgent.description} -
-
- - {/* Capabilities */} - {selectedAgent.capabilities && Object.keys(selectedAgent.capabilities).length > 0 && ( -
- Capabilities -
- {Object.entries(selectedAgent.capabilities) - .filter(([_, value]) => value === true) - .map(([key]) => ( - - {key} - - ))} -
-
- )} - - {/* Input/Output Modes */} -
- Input/Output Modes -
-
- Input Modes: -
- {selectedAgent.defaultInputModes?.map((mode) => ( - - {mode} - - )) || Not specified} +
+

Version:

+ v{selectedAgent.version} +
+
+

Protocol Version:

+

{selectedAgent.protocolVersion}

+
+
+

URL:

+
+

{selectedAgent.url}

+ void copyToClipboard(selectedAgent.url)} + className="size-3.5 shrink-0 cursor-pointer text-gray-500 hover:text-blue-500" + /> +
- Output Modes: -
- {selectedAgent.defaultOutputModes?.map((mode) => ( - - {mode} - - )) || Not specified} +

Description:

+

{selectedAgent.description}

+
+
+ + {/* Capabilities */} + {selectedAgent.capabilities && Object.keys(selectedAgent.capabilities).length > 0 && ( +
+

Capabilities

+
+ {Object.entries(selectedAgent.capabilities) + .filter(([_, value]) => value === true) + .map(([key]) => ( + + {key} + + ))} +
+
+ )} + + {/* Input/Output Modes */} +
+

Input/Output Modes

+
+
+

Input Modes:

+
+ {selectedAgent.defaultInputModes?.map((mode) => ( + + {mode} + + )) ||

Not specified

} +
+
+
+

Output Modes:

+
+ {selectedAgent.defaultOutputModes?.map((mode) => ( + + {mode} + + )) ||

Not specified

} +
-
- {/* Skills */} - {selectedAgent.skills && selectedAgent.skills.length > 0 && ( -
- Skills -
- {selectedAgent.skills.map((skill) => ( -
-
-
- {skill.name} - ID: {skill.id} + {/* Skills */} + {selectedAgent.skills && selectedAgent.skills.length > 0 && ( +
+

Skills

+
+ {selectedAgent.skills.map((skill) => ( +
+
+
+

{skill.name}

+

ID: {skill.id}

+
+ {skill.tags && skill.tags.length > 0 && ( +
+ {skill.tags.map((tag) => ( + + {tag} + + ))} +
+ )}
- {skill.tags && skill.tags.length > 0 && ( -
- {skill.tags.map((tag) => ( - - {tag} - - ))} +

{skill.description}

+ {skill.examples && skill.examples.length > 0 && ( +
+

Examples:

+
+ {skill.examples.map((example, idx) => ( + + {example} + + ))} +
)}
- {skill.description} - {skill.examples && skill.examples.length > 0 && ( -
- Examples: -
- {skill.examples.map((example, idx) => ( - - {example} - - ))} -
-
- )} -
- ))} + ))} +
-
- )} + )} - {/* Additional Properties */} - {selectedAgent.supportsAuthenticatedExtendedCard && ( -
- Additional Features - Supports Authenticated Extended Card -
- )} -
- )} - + {/* Additional Properties */} + {selectedAgent.supportsAuthenticatedExtendedCard && ( +
+

Additional Features

+ Supports Authenticated Extended Card +
+ )} +
+ )} + +
{/* MCP Server Details Modal */} - - {selectedMcpServer && ( -
- {/* Server Overview */} -
- Server Overview -
-
- Server Name: - {selectedMcpServer.server_name} -
-
- Server ID: -
- {selectedMcpServer.server_id} - void copyToClipboard(selectedMcpServer.server_id)} - className="cursor-pointer text-gray-500 hover:text-blue-500" - /> + !open && handleCancel()}> + + + {selectedMcpServer?.server_name || "MCP Server Details"} + + {selectedMcpServer && ( +
+ {/* Server Overview */} +
+

Server Overview

+
+
+

Server Name:

+

{selectedMcpServer.server_name}

+
+
+

Server ID:

+
+

{selectedMcpServer.server_id}

+ void copyToClipboard(selectedMcpServer.server_id)} + className="size-3.5 shrink-0 cursor-pointer text-gray-500 hover:text-blue-500" + /> +
+
+ {selectedMcpServer.alias && ( +
+

Alias:

+

{selectedMcpServer.alias}

+
+ )} +
+

Transport:

+ {selectedMcpServer.transport} +
+
+

Auth Type:

+ + {selectedMcpServer.auth_type} + +
+
+

Status:

+ + {selectedMcpServer.status || "unknown"} +
- {selectedMcpServer.alias && ( -
- Alias: - {selectedMcpServer.alias} + {selectedMcpServer.description && ( +
+

Description:

+

{selectedMcpServer.description}

)} -
- Transport: - {selectedMcpServer.transport} -
-
- Auth Type: - - {selectedMcpServer.auth_type} - -
-
- Status: - - {selectedMcpServer.status || "unknown"} - +
+ + {/* Connection Details */} +
+

Connection Details

+
+ {selectedMcpServer.command && ( +
+

Command:

+

{selectedMcpServer.command}

+
+ )}
- {selectedMcpServer.description && ( -
- Description: - {selectedMcpServer.description} + + {/* Tools */} + {selectedMcpServer.allowed_tools && selectedMcpServer.allowed_tools.length > 0 && ( +
+

Allowed Tools

+
+ {selectedMcpServer.allowed_tools.map((tool, idx) => ( + + {tool} + + ))} +
)} -
- {/* Connection Details */} -
- Connection Details -
- {selectedMcpServer.command && ( -
- Command: - - {selectedMcpServer.command} - + {/* Teams */} + {selectedMcpServer.teams && selectedMcpServer.teams.length > 0 && ( +
+

Teams

+
+ {selectedMcpServer.teams.map((team, idx) => ( + + {team} + + ))}
- )} -
-
- - {/* Tools */} - {selectedMcpServer.allowed_tools && selectedMcpServer.allowed_tools.length > 0 && ( -
- Allowed Tools -
- {selectedMcpServer.allowed_tools.map((tool, idx) => ( - - {tool} - - ))} -
-
- )} - - {/* Teams */} - {selectedMcpServer.teams && selectedMcpServer.teams.length > 0 && ( -
- Teams -
- {selectedMcpServer.teams.map((team, idx) => ( - - {team} - - ))} -
-
- )} - - {/* Access Groups */} - {selectedMcpServer.mcp_access_groups && selectedMcpServer.mcp_access_groups.length > 0 && ( -
- Access Groups -
- {selectedMcpServer.mcp_access_groups.map((group, idx) => ( - - {group} - - ))} -
-
- )} - - {/* Metadata */} -
- Metadata -
-
- Created By: - {selectedMcpServer.created_by} -
-
- Updated By: - {selectedMcpServer.updated_by} -
-
- Created At: - {new Date(selectedMcpServer.created_at).toLocaleString()} -
-
- Updated At: - {new Date(selectedMcpServer.updated_at).toLocaleString()} -
- {selectedMcpServer.last_health_check && ( -
- Last Health Check: - {new Date(selectedMcpServer.last_health_check).toLocaleString()} -
- )} -
- {selectedMcpServer.health_check_error && ( -
- Health Check Error: - {selectedMcpServer.health_check_error}
)} -
- {/* Usage Example */} -
- Usage Example - - {`from fastmcp import Client + {/* Access Groups */} + {selectedMcpServer.mcp_access_groups && selectedMcpServer.mcp_access_groups.length > 0 && ( +
+

Access Groups

+
+ {selectedMcpServer.mcp_access_groups.map((group, idx) => ( + + {group} + + ))} +
+
+ )} + + {/* Metadata */} +
+

Metadata

+
+
+

Created By:

+

{selectedMcpServer.created_by}

+
+
+

Updated By:

+

{selectedMcpServer.updated_by}

+
+
+

Created At:

+

{new Date(selectedMcpServer.created_at).toLocaleString()}

+
+
+

Updated At:

+

{new Date(selectedMcpServer.updated_at).toLocaleString()}

+
+ {selectedMcpServer.last_health_check && ( +
+

Last Health Check:

+

{new Date(selectedMcpServer.last_health_check).toLocaleString()}

+
+ )} +
+ {selectedMcpServer.health_check_error && ( +
+

Health Check Error:

+

{selectedMcpServer.health_check_error}

+
+ )} +
+ + {/* Usage Example */} +
+

Usage Example

+ + {`from fastmcp import Client import asyncio # Standard MCP configuration @@ -1088,11 +1079,12 @@ async def main(): if __name__ == "__main__": asyncio.run(main())`} - + +
-
- )} - + )} + +
{/* Make Model Public Form */} = ({ // Derived stats const totalSkills = skills.length; - const domains = useMemo(() => [...new Set(skills.map((s) => s.domain).filter(Boolean))], [skills]); + const domains = useMemo( + () => [...new Set(skills.map((s) => s.domain).filter((domain): domain is string => Boolean(domain)))], + [skills], + ); const namespaces = useMemo(() => [...new Set(skills.map((s) => s.namespace).filter(Boolean))], [skills]); // Filtered table data @@ -73,6 +78,11 @@ const SkillHubDashboard: React.FC = ({ const columns = useMemo(() => getSkillHubTableColumns({ onSkillClick: setSelectedSkill }), []); + const domainItems = useMemo( + () => [{ value: ALL_DOMAINS, label: "All Domains" }, ...domains.map((d) => ({ value: d, label: d }))], + [domains], + ); + const hasActiveFilter = search.trim().length > 0 || domainFilter != null; if (selectedSkill) { @@ -111,21 +121,43 @@ const SkillHubDashboard: React.FC = ({

All {publicPage ? "Public " : ""}Skills

} - placeholder="Search by name, namespace, or tag…" - value={search} - onChange={(e) => setSearch(e.target.value)} - style={{ width: 280 }} - allowClear - /> + items={domainItems} + value={domainFilter ?? ALL_DOMAINS} + onValueChange={(val) => setDomainFilter(val === null || val === ALL_DOMAINS ? undefined : val)} + > + + + + + {domainItems.map((item) => ( + + {item.label} + + ))} + + + + + + + setSearch(e.target.value)} + /> + {search !== "" && ( + + setSearch("")} + > + + + + )} +
= ({ accessTok }; return ( - +
setIsExpanded(!isExpanded)}>
- Link Management +

Link Management

Manage the links that are displayed under 'Useful Links' on the public model hub.

@@ -243,7 +244,7 @@ const UsefulLinksManagement: React.FC = ({ accessTok {isExpanded && (
- Add New Link +

Add New Link

@@ -288,7 +289,7 @@ const UsefulLinksManagement: React.FC = ({ accessTok
- Manage Existing Links +

Manage Existing Links

= ({ accessTok
- + - Display Name - URL - Actions + Display Name + URL + Actions - + {links.map((link, index) => ( diff --git a/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.test.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.test.tsx index 67c6d7d6cc9..a55beaf517f 100644 --- a/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.test.tsx @@ -12,67 +12,8 @@ vi.mock("../../networking", () => ({ import { makeAgentsPublicCall } from "../../networking"; const mockMakeAgentsPublicCall = vi.mocked(makeAgentsPublicCall); -// Mock antd components -vi.mock("antd", () => ({ - Modal: ({ open, title, children, onCancel, footer }: any) => - open ? ( -
-
{title}
- {children} - {footer} -
- ) : null, - Form: Object.assign(({ children, form }: any) => {children}, { - useForm: () => [ - { - resetFields: vi.fn(), - validateFields: vi.fn(), - getFieldsValue: vi.fn(), - setFieldsValue: vi.fn(), - }, - vi.fn(), - ], - Item: ({ children }: any) =>
{children}
, - }), - Steps: Object.assign( - ({ children, current, className }: any) => ( -
- {children} -
- ), - { - Step: ({ title }: any) =>
{title}
, - }, - ), - Button: ({ children, onClick, disabled, loading, ...props }: any) => ( - - ), - Checkbox: ({ checked, indeterminate, onChange, children, disabled }: any) => ( - - ), -})); - -// Mock @tremor/react components -vi.mock("@tremor/react", () => ({ - Text: ({ children, className }: any) => {children}, - Title: ({ children }: any) =>

{children}

, - Badge: ({ children, color, size }: any) => ( - - {children} - - ), -})); +const expectDisabledControl = (element: HTMLElement) => + expect(element.hasAttribute("disabled") || element.getAttribute("aria-disabled") === "true").toBe(true); describe("MakeAgentPublicForm", () => { const mockProps = { @@ -143,7 +84,7 @@ describe("MakeAgentPublicForm", () => { expect(screen.getByText("Select Agents to Make Public")).toBeInTheDocument(); // Select all agents using the select all checkbox - const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + const selectAllCheckbox = screen.getByRole("checkbox", { name: "Select All (2)" }); await act(async () => { fireEvent.click(selectAllCheckbox); }); @@ -169,12 +110,11 @@ describe("MakeAgentPublicForm", () => { render(); // Select all agents - const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + const selectAllCheckbox = screen.getByRole("checkbox", { name: "Select All (2)" }); await act(async () => { fireEvent.click(selectAllCheckbox); }); - // Navigate to confirm step const nextButton = screen.getByRole("button", { name: "Next" }); await act(async () => { fireEvent.click(nextButton); @@ -185,7 +125,6 @@ describe("MakeAgentPublicForm", () => { expect(screen.getByText("Confirm Making Agents Public")).toBeInTheDocument(); }); - // Submit const submitButton = screen.getByRole("button", { name: "Make Public" }); await act(async () => { fireEvent.click(submitButton); @@ -232,6 +171,8 @@ describe("MakeAgentPublicForm", () => { const checkboxes = screen.getAllByRole("checkbox"); await act(async () => { fireEvent.click(checkboxes[0]); // Click select all to select all + }); + await act(async () => { fireEvent.click(checkboxes[0]); // Click select all again to deselect all }); @@ -256,8 +197,8 @@ describe("MakeAgentPublicForm", () => { expect(screen.getByText("No agents available.")).toBeInTheDocument(); // Select All checkbox should be disabled - const selectAllCheckbox = screen.getByLabelText("Select All"); - expect(selectAllCheckbox).toBeDisabled(); + const selectAllCheckbox = screen.getByRole("checkbox", { name: "Select All" }); + expectDisabledControl(selectAllCheckbox); // Next button should be disabled const nextButton = screen.getByRole("button", { name: "Next" }); @@ -332,7 +273,7 @@ describe("MakeAgentPublicForm", () => { // Select all should be indeterminate now const selectAllCheckbox = checkboxes[0]; - expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + expect(selectAllCheckbox).toBePartiallyChecked(); }); it("should display skills overflow text when agent has more than 3 skills", () => { @@ -369,7 +310,6 @@ describe("MakeAgentPublicForm", () => { render(); - // Navigate to confirm step const nextButton = screen.getByRole("button", { name: "Next" }); await act(async () => { fireEvent.click(nextButton); @@ -379,7 +319,6 @@ describe("MakeAgentPublicForm", () => { expect(screen.getByText("Confirm Making Agents Public")).toBeInTheDocument(); }); - // Submit const submitButton = screen.getByRole("button", { name: "Make Public" }); await act(async () => { fireEvent.click(submitButton); @@ -395,7 +334,7 @@ describe("MakeAgentPublicForm", () => { expect(mockProps.onClose).not.toHaveBeenCalled(); }); - it("should show loading state during submit", async () => { + it("should not complete the flow until the submit request resolves", async () => { let resolvePromise: (value: any) => void = () => {}; const pendingPromise = new Promise((resolve) => { resolvePromise = resolve; @@ -404,7 +343,6 @@ describe("MakeAgentPublicForm", () => { render(); - // Navigate to confirm step const nextButton = screen.getByRole("button", { name: "Next" }); await act(async () => { fireEvent.click(nextButton); @@ -414,17 +352,20 @@ describe("MakeAgentPublicForm", () => { expect(screen.getByText("Confirm Making Agents Public")).toBeInTheDocument(); }); - // Submit const submitButton = screen.getByRole("button", { name: "Make Public" }); await act(async () => { fireEvent.click(submitButton); }); - // Check loading state - expect(submitButton).toHaveAttribute("data-loading", "true"); - expect(submitButton).toBeDisabled(); + expectDisabledControl(submitButton); + await act(async () => { + fireEvent.click(submitButton); + }); + expect(mockMakeAgentsPublicCall).toHaveBeenCalledTimes(1); + expect(mockProps.onSuccess).not.toHaveBeenCalled(); + expect(mockProps.onClose).not.toHaveBeenCalled(); + expect(screen.getByText("Confirm Making Agents Public")).toBeInTheDocument(); - // Resolve the promise resolvePromise({}); await waitFor(() => { expect(mockProps.onSuccess).toHaveBeenCalled(); @@ -441,7 +382,7 @@ describe("MakeAgentPublicForm", () => { render(); // Modal should not be rendered - expect(screen.queryByTestId("modal")).not.toBeInTheDocument(); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); expect(screen.queryByText("Make Agents Public")).not.toBeInTheDocument(); }); @@ -500,6 +441,6 @@ describe("MakeAgentPublicForm", () => { // Select all should be indeterminate const selectAllCheckbox = checkboxes[0]; - expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + expect(selectAllCheckbox).toBePartiallyChecked(); }); }); diff --git a/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.tsx index 0ed73872cee..82336206858 100644 --- a/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.tsx @@ -1,11 +1,15 @@ import React, { useState, useEffect } from "react"; -import { Modal, Form, Steps, Button, Checkbox } from "antd"; -import { Text, Title, Badge } from "@tremor/react"; +import { Loader2 } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { cn } from "@/lib/cva.config"; import { makeAgentsPublicCall } from "../../networking"; import NotificationsManager from "../../molecules/notifications_manager"; import { AgentHubData } from "@/components/AIHub/AgentHubTableColumns"; -const { Step } = Steps; +const STEP_TITLES = ["Select Agents", "Confirm"]; interface MakeAgentPublicFormProps { visible: boolean; @@ -25,12 +29,10 @@ const MakeAgentPublicForm: React.FC = ({ const [currentStep, setCurrentStep] = useState(0); const [selectedAgents, setSelectedAgents] = useState>(new Set()); const [loading, setLoading] = useState(false); - const [form] = Form.useForm(); const handleClose = () => { setCurrentStep(0); setSelectedAgents(new Set()); - form.resetFields(); onClose(); }; @@ -113,29 +115,30 @@ const MakeAgentPublicForm: React.FC = ({ return (
- Select Agents to Make Public +

Select Agents to Make Public

- handleSelectAll(e.target.checked)} - disabled={agentHubData.length === 0} - > +
- +

Select the agents you want to be visible on the public model hub. Users will still require a valid Virtual Key to use these agents. - +

{agentHubData.length === 0 ? (
- No agents available. +

No agents available.

) : ( agentHubData.map((agent) => { @@ -144,25 +147,23 @@ const MakeAgentPublicForm: React.FC = ({
handleAgentSelection(agentId, e.target.checked)} + onCheckedChange={(checked) => handleAgentSelection(agentId, checked === true)} /> -
+
- {agent.name} - - v{agent.version} - +

{agent.name}

+ v{agent.version}
- {agent.description} +

{agent.description}

{agent.skills && agent.skills.length > 0 && (
{agent.skills.slice(0, 3).map((skill) => ( - + {skill.name} ))} {agent.skills.length > 3 && ( - +{agent.skills.length - 3} more +

+{agent.skills.length - 3} more

)}
)} @@ -176,9 +177,9 @@ const MakeAgentPublicForm: React.FC = ({ {selectedAgents.size > 0 && (
- +

{selectedAgents.size} agent{selectedAgents.size !== 1 ? "s" : ""} selected - +

)}
@@ -188,33 +189,31 @@ const MakeAgentPublicForm: React.FC = ({ const renderStep2Content = () => { return (
- Confirm Making Agents Public +

Confirm Making Agents Public

- +

Warning: Once you make these agents public, anyone who can go to the{" "} /ui/model_hub_table will be able to know they exist on the proxy. - +

- Agents to be made public: +

Agents to be made public:

{Array.from(selectedAgents).map((agentId) => { const agent = agentHubData.find((a) => (a.agent_id || a.name) === agentId); return (
-
+
- {agent?.name || agentId} - {agent && ( - - v{agent.version} - - )} +

{agent?.name || agentId}

+ {agent && v{agent.version}}
- {agent?.description && {agent.description}} + {agent?.description && ( +

{agent.description}

+ )}
); @@ -224,10 +223,10 @@ const MakeAgentPublicForm: React.FC = ({
- +

Total: {selectedAgents.size} agent{selectedAgents.size !== 1 ? "s" : ""} will be made public - +

); @@ -247,7 +246,7 @@ const MakeAgentPublicForm: React.FC = ({ const renderStepButtons = () => { return (
- @@ -259,7 +258,8 @@ const MakeAgentPublicForm: React.FC = ({ )} {currentStep === 1 && ( - )} @@ -269,24 +269,42 @@ const MakeAgentPublicForm: React.FC = ({ }; return ( - -
- - - - + !open && handleClose()} disablePointerDismissal> + + + Make Agents Public + - {renderStepContent()} - {renderStepButtons()} - - +
+
    + {STEP_TITLES.map((title, index) => ( +
  1. + + {index + 1} + + + {title} + +
    2. + ))} +
+ + {renderStepContent()} + {renderStepButtons()} +
+ + ); }; diff --git a/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.test.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.test.tsx index 994a920b2e4..ff385b3ed7c 100644 --- a/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.test.tsx @@ -12,83 +12,8 @@ vi.mock("../../networking", () => ({ import { makeMCPPublicCall } from "../../networking"; const mockMakeMCPPublicCall = vi.mocked(makeMCPPublicCall); -// Mock antd components -vi.mock("antd", () => ({ - Modal: ({ open, title, children, onCancel, footer }: any) => - open ? ( -
-
{title}
- {children} - {footer} -
- ) : null, - Form: Object.assign(({ children, form }: any) =>
{children}, { - useForm: () => [ - { - resetFields: vi.fn(), - validateFields: vi.fn(), - getFieldsValue: vi.fn(), - setFieldsValue: vi.fn(), - }, - vi.fn(), - ], - Item: ({ children }: any) =>
{children}
, - }), - Steps: Object.assign( - ({ children, current, className }: any) => ( -
- {children} -
- ), - { - Step: ({ title }: any) =>
{title}
, - }, - ), - Button: ({ children, onClick, disabled, loading, ...props }: any) => ( - - ), - Checkbox: ({ checked, indeterminate, onChange, children, disabled }: any) => ( - - ), -})); - -// Additional @tremor/react mocks. -// NOTE: the comment used to say "Button is already mocked globally" — that was -// incorrect. A file-level vi.mock fully replaces the setup-level mock from -// tests/setupTests.ts, so we must re-apply the Button/Tooltip overrides here. -// Without them, the real Tremor Button leaks through and its useTooltip(300) -// schedules a native setTimeout that can fire post-teardown -> "window is not defined". -vi.mock("@tremor/react", async (importOriginal) => { - const actual = await importOriginal(); - const React = await import("react"); - return { - ...actual, - Text: ({ children, className }: any) => {children}, - Title: ({ children }: any) =>

{children}

, - Badge: ({ children, color, size }: any) => ( - - {children} - - ), - Button: React.forwardRef(({ children, ...props }, ref) => ( - - )), - Tooltip: ({ children }: any) => <>{children}, - }; -}); +const expectDisabledControl = (element: HTMLElement) => + expect(element.hasAttribute("disabled") || element.getAttribute("aria-disabled") === "true").toBe(true); describe("MakeMCPPublicForm", () => { const mockProps = { @@ -182,7 +107,7 @@ describe("MakeMCPPublicForm", () => { expect(screen.getByText("Select MCP Servers to Make Public")).toBeInTheDocument(); // Select all servers using the select all checkbox - const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + const selectAllCheckbox = screen.getByRole("checkbox", { name: "Select All (2)" }); await act(async () => { fireEvent.click(selectAllCheckbox); }); @@ -208,12 +133,11 @@ describe("MakeMCPPublicForm", () => { render(); // Select all servers - const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + const selectAllCheckbox = screen.getByRole("checkbox", { name: "Select All (2)" }); await act(async () => { fireEvent.click(selectAllCheckbox); }); - // Navigate to confirm step const nextButton = screen.getByRole("button", { name: "Next" }); await act(async () => { fireEvent.click(nextButton); @@ -224,7 +148,6 @@ describe("MakeMCPPublicForm", () => { expect(screen.getByText("Confirm Making MCP Servers Public")).toBeInTheDocument(); }); - // Submit const submitButton = screen.getByRole("button", { name: "Make Public" }); await act(async () => { fireEvent.click(submitButton); @@ -271,6 +194,8 @@ describe("MakeMCPPublicForm", () => { const checkboxes = screen.getAllByRole("checkbox"); await act(async () => { fireEvent.click(checkboxes[0]); // Click select all to select all + }); + await act(async () => { fireEvent.click(checkboxes[0]); // Click select all again to deselect all }); @@ -295,8 +220,8 @@ describe("MakeMCPPublicForm", () => { expect(screen.getByText("No MCP servers available.")).toBeInTheDocument(); // Select All checkbox should be disabled - const selectAllCheckbox = screen.getByLabelText("Select All"); - expect(selectAllCheckbox).toBeDisabled(); + const selectAllCheckbox = screen.getByRole("checkbox", { name: "Select All" }); + expectDisabledControl(selectAllCheckbox); // Next button should be disabled const nextButton = screen.getByRole("button", { name: "Next" }); @@ -371,7 +296,7 @@ describe("MakeMCPPublicForm", () => { // Select all should be indeterminate now const selectAllCheckbox = checkboxes[0]; - expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + expect(selectAllCheckbox).toBePartiallyChecked(); }); it("should display tools overflow text when server has more than 3 tools", () => { @@ -402,7 +327,6 @@ describe("MakeMCPPublicForm", () => { render(); - // Navigate to confirm step const nextButton = screen.getByRole("button", { name: "Next" }); await act(async () => { fireEvent.click(nextButton); @@ -412,7 +336,6 @@ describe("MakeMCPPublicForm", () => { expect(screen.getByText("Confirm Making MCP Servers Public")).toBeInTheDocument(); }); - // Submit const submitButton = screen.getByRole("button", { name: "Make Public" }); await act(async () => { fireEvent.click(submitButton); @@ -428,7 +351,7 @@ describe("MakeMCPPublicForm", () => { expect(mockProps.onClose).not.toHaveBeenCalled(); }); - it("should show loading state during submit", async () => { + it("should not complete the flow until the submit request resolves", async () => { let resolvePromise: (value: any) => void = () => {}; const pendingPromise = new Promise((resolve) => { resolvePromise = resolve; @@ -437,7 +360,6 @@ describe("MakeMCPPublicForm", () => { render(); - // Navigate to confirm step const nextButton = screen.getByRole("button", { name: "Next" }); await act(async () => { fireEvent.click(nextButton); @@ -447,17 +369,20 @@ describe("MakeMCPPublicForm", () => { expect(screen.getByText("Confirm Making MCP Servers Public")).toBeInTheDocument(); }); - // Submit const submitButton = screen.getByRole("button", { name: "Make Public" }); await act(async () => { fireEvent.click(submitButton); }); - // Check loading state - expect(submitButton).toHaveAttribute("data-loading", "true"); - expect(submitButton).toBeDisabled(); + expectDisabledControl(submitButton); + await act(async () => { + fireEvent.click(submitButton); + }); + expect(mockMakeMCPPublicCall).toHaveBeenCalledTimes(1); + expect(mockProps.onSuccess).not.toHaveBeenCalled(); + expect(mockProps.onClose).not.toHaveBeenCalled(); + expect(screen.getByText("Confirm Making MCP Servers Public")).toBeInTheDocument(); - // Resolve the promise resolvePromise({}); await waitFor(() => { expect(mockProps.onSuccess).toHaveBeenCalled(); @@ -474,7 +399,7 @@ describe("MakeMCPPublicForm", () => { render(); // Modal should not be rendered - expect(screen.queryByTestId("modal")).not.toBeInTheDocument(); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); expect(screen.queryByText("Make MCP Servers Public")).not.toBeInTheDocument(); }); @@ -569,6 +494,6 @@ describe("MakeMCPPublicForm", () => { // Select all should be indeterminate const selectAllCheckbox = checkboxes[0]; - expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + expect(selectAllCheckbox).toBePartiallyChecked(); }); }); diff --git a/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.tsx index b590c3cc1dd..7ef42883400 100644 --- a/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.tsx @@ -1,11 +1,25 @@ import React, { useState, useEffect } from "react"; -import { Modal, Form, Steps, Button, Checkbox } from "antd"; -import { Text, Title, Badge } from "@tremor/react"; +import { Loader2 } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { cn } from "@/lib/cva.config"; import { makeMCPPublicCall } from "../../networking"; import NotificationsManager from "../../molecules/notifications_manager"; import { MCPServerData } from "@/components/AIHub/MCPHubTableColumns"; -const { Step } = Steps; +const STEP_TITLES = ["Select Servers", "Confirm"]; + +const statusVariant = (status?: string) => { + if (status === "active" || status === "healthy") { + return "default" as const; + } + if (status === "inactive" || status === "unhealthy") { + return "destructive" as const; + } + return "outline" as const; +}; interface MakeMCPPublicFormProps { visible: boolean; @@ -25,12 +39,10 @@ const MakeMCPPublicForm: React.FC = ({ const [currentStep, setCurrentStep] = useState(0); const [selectedServers, setSelectedServers] = useState>(new Set()); const [loading, setLoading] = useState(false); - const [form] = Form.useForm(); const handleClose = () => { setCurrentStep(0); setSelectedServers(new Set()); - form.resetFields(); onClose(); }; @@ -114,29 +126,30 @@ const MakeMCPPublicForm: React.FC = ({ return (
- Select MCP Servers to Make Public +

Select MCP Servers to Make Public

- handleSelectAll(e.target.checked)} - disabled={mcpHubData.length === 0} - > +
- +

Select the MCP servers you want to be visible on the public model hub. Users will still require a valid Virtual Key to use these servers. - +

{mcpHubData.length === 0 ? (
- No MCP servers available. +

No MCP servers available.

) : ( mcpHubData.map((server) => { @@ -148,42 +161,25 @@ const MakeMCPPublicForm: React.FC = ({ > handleServerSelection(server.server_id, e.target.checked)} + onCheckedChange={(checked) => handleServerSelection(server.server_id, checked === true)} /> -
-
- {server.server_name} - {isPublic && ( - - Public - - )} - - {server.transport} - - - {server.status || "unknown"} - +
+
+

{server.server_name}

+ {isPublic && Public} + {server.transport} + {server.status || "unknown"}
- {server.description || server.url} +

{server.description || server.url}

{server.allowed_tools && server.allowed_tools.length > 0 && (
{server.allowed_tools.slice(0, 3).map((tool, idx) => ( - + {tool} ))} {server.allowed_tools.length > 3 && ( - +{server.allowed_tools.length - 3} more +

+{server.allowed_tools.length - 3} more

)}
)} @@ -197,9 +193,9 @@ const MakeMCPPublicForm: React.FC = ({ {selectedServers.size > 0 && (
- +

{selectedServers.size} MCP server{selectedServers.size !== 1 ? "s" : ""} selected - +

)}
@@ -209,48 +205,37 @@ const MakeMCPPublicForm: React.FC = ({ const renderStep2Content = () => { return (
- Confirm Making MCP Servers Public +

Confirm Making MCP Servers Public

- +

Warning: Once you make these MCP servers public, anyone who can go to the{" "} /ui/model_hub_table will be able to know they exist on the proxy. - +

- MCP Servers to be made public: +

MCP Servers to be made public:

{Array.from(selectedServers).map((serverId) => { const server = mcpHubData.find((s) => s.server_id === serverId); return (
-
-
- {server?.server_name || serverId} +
+
+

{server?.server_name || serverId}

{server && ( <> - - {server.transport} - - - {server.status || "unknown"} - + {server.transport} + {server.status || "unknown"} )}
- {server?.description && {server.description}} - {server?.url && {server.url}} + {server?.description && ( +

{server.description}

+ )} + {server?.url &&

{server.url}

}
); @@ -260,10 +245,10 @@ const MakeMCPPublicForm: React.FC = ({
- +

Total: {selectedServers.size} MCP server{selectedServers.size !== 1 ? "s" : ""} will be made public - +

); @@ -283,7 +268,7 @@ const MakeMCPPublicForm: React.FC = ({ const renderStepButtons = () => { return (
- @@ -295,7 +280,8 @@ const MakeMCPPublicForm: React.FC = ({ )} {currentStep === 1 && ( - )} @@ -305,24 +291,42 @@ const MakeMCPPublicForm: React.FC = ({ }; return ( - -
- - - - + !open && handleClose()} disablePointerDismissal> + + + Make MCP Servers Public + - {renderStepContent()} - {renderStepButtons()} - - +
+
    + {STEP_TITLES.map((title, index) => ( +
  1. + + {index + 1} + + + {title} + +
    2. + ))} +
+ + {renderStepContent()} + {renderStepButtons()} +
+ + ); }; diff --git a/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.test.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.test.tsx index 2b57535f3ad..ac0df137f6a 100644 --- a/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.test.tsx @@ -29,67 +29,8 @@ vi.mock("../../networking", () => ({ import { makeModelGroupPublic } from "../../networking"; const mockMakeModelGroupPublic = vi.mocked(makeModelGroupPublic); -// Mock antd components -vi.mock("antd", () => ({ - Modal: ({ open, title, children, onCancel, footer }: any) => - open ? ( -
-
{title}
- {children} - {footer} -
- ) : null, - Form: Object.assign(({ children, form }: any) =>
{children}, { - useForm: () => [ - { - resetFields: vi.fn(), - validateFields: vi.fn(), - getFieldsValue: vi.fn(), - setFieldsValue: vi.fn(), - }, - vi.fn(), - ], - Item: ({ children }: any) =>
{children}
, - }), - Steps: Object.assign( - ({ children, current, className }: any) => ( -
- {children} -
- ), - { - Step: ({ title }: any) =>
{title}
, - }, - ), - Button: ({ children, onClick, disabled, loading, ...props }: any) => ( - - ), - Checkbox: ({ checked, indeterminate, onChange, children, disabled }: any) => ( - - ), -})); - -// Mock @tremor/react components -vi.mock("@tremor/react", () => ({ - Text: ({ children, className }: any) => {children}, - Title: ({ children }: any) =>

{children}

, - Badge: ({ children, color, size }: any) => ( - - {children} - - ), -})); +const expectDisabledControl = (element: HTMLElement) => + expect(element.hasAttribute("disabled") || element.getAttribute("aria-disabled") === "true").toBe(true); // Mock ModelFilters component vi.mock("../../model_filters", () => ({ @@ -190,7 +131,7 @@ describe("MakeModelPublicForm", () => { expect(screen.getByText("Select Models to Make Public")).toBeInTheDocument(); // Select all models using the select all checkbox - const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + const selectAllCheckbox = screen.getByRole("checkbox", { name: "Select All (2)" }); await act(async () => { fireEvent.click(selectAllCheckbox); }); @@ -216,12 +157,11 @@ describe("MakeModelPublicForm", () => { render(); // Select all models - const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + const selectAllCheckbox = screen.getByRole("checkbox", { name: "Select All (2)" }); await act(async () => { fireEvent.click(selectAllCheckbox); }); - // Navigate to confirm step const nextButton = screen.getByRole("button", { name: "Next" }); await act(async () => { fireEvent.click(nextButton); @@ -232,7 +172,6 @@ describe("MakeModelPublicForm", () => { expect(screen.getByText("Confirm Making Models Public")).toBeInTheDocument(); }); - // Submit const submitButton = screen.getByRole("button", { name: "Make Public" }); await act(async () => { fireEvent.click(submitButton); @@ -279,6 +218,8 @@ describe("MakeModelPublicForm", () => { const checkboxes = screen.getAllByRole("checkbox"); await act(async () => { fireEvent.click(checkboxes[0]); // Click select all to select all + }); + await act(async () => { fireEvent.click(checkboxes[0]); // Click select all again to deselect all }); @@ -303,8 +244,8 @@ describe("MakeModelPublicForm", () => { expect(screen.getByText("No models match the current filters.")).toBeInTheDocument(); // Select All checkbox should be disabled - const selectAllCheckbox = screen.getByLabelText("Select All"); - expect(selectAllCheckbox).toBeDisabled(); + const selectAllCheckbox = screen.getByRole("checkbox", { name: "Select All" }); + expectDisabledControl(selectAllCheckbox); // Next button should be disabled const nextButton = screen.getByRole("button", { name: "Next" }); @@ -379,7 +320,7 @@ describe("MakeModelPublicForm", () => { // Select all should be indeterminate now const selectAllCheckbox = checkboxes[0]; - expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + expect(selectAllCheckbox).toBePartiallyChecked(); }); it("should display model badges and information", () => { @@ -402,7 +343,6 @@ describe("MakeModelPublicForm", () => { render(); - // Navigate to confirm step const nextButton = screen.getByRole("button", { name: "Next" }); await act(async () => { fireEvent.click(nextButton); @@ -412,7 +352,6 @@ describe("MakeModelPublicForm", () => { expect(screen.getByText("Confirm Making Models Public")).toBeInTheDocument(); }); - // Submit const submitButton = screen.getByRole("button", { name: "Make Public" }); await act(async () => { fireEvent.click(submitButton); @@ -428,7 +367,7 @@ describe("MakeModelPublicForm", () => { expect(mockProps.onClose).not.toHaveBeenCalled(); }); - it("should show loading state during submit", async () => { + it("should not complete the flow until the submit request resolves", async () => { let resolvePromise: (value: any) => void = () => {}; const pendingPromise = new Promise((resolve) => { resolvePromise = resolve; @@ -437,7 +376,6 @@ describe("MakeModelPublicForm", () => { render(); - // Navigate to confirm step const nextButton = screen.getByRole("button", { name: "Next" }); await act(async () => { fireEvent.click(nextButton); @@ -447,17 +385,20 @@ describe("MakeModelPublicForm", () => { expect(screen.getByText("Confirm Making Models Public")).toBeInTheDocument(); }); - // Submit const submitButton = screen.getByRole("button", { name: "Make Public" }); await act(async () => { fireEvent.click(submitButton); }); - // Check loading state - expect(submitButton).toHaveAttribute("data-loading", "true"); - expect(submitButton).toBeDisabled(); + expectDisabledControl(submitButton); + await act(async () => { + fireEvent.click(submitButton); + }); + expect(mockMakeModelGroupPublic).toHaveBeenCalledTimes(1); + expect(mockProps.onSuccess).not.toHaveBeenCalled(); + expect(mockProps.onClose).not.toHaveBeenCalled(); + expect(screen.getByText("Confirm Making Models Public")).toBeInTheDocument(); - // Resolve the promise resolvePromise({}); await waitFor(() => { expect(mockProps.onSuccess).toHaveBeenCalled(); @@ -474,7 +415,7 @@ describe("MakeModelPublicForm", () => { render(); // Modal should not be rendered - expect(screen.queryByTestId("modal")).not.toBeInTheDocument(); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); expect(screen.queryByText("Make Models Public")).not.toBeInTheDocument(); }); @@ -521,21 +462,19 @@ describe("MakeModelPublicForm", () => { // Select all should be indeterminate const selectAllCheckbox = checkboxes[0]; - expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + expect(selectAllCheckbox).toBePartiallyChecked(); }); it("should show selected count", () => { render(); // Should show that 1 model is selected (gpt-3.5-turbo is preselected) - expect(screen.getByText("1")).toBeInTheDocument(); - expect(screen.getByText("model selected")).toBeInTheDocument(); + expect(screen.getByText("model selected")).toHaveTextContent("1 model selected"); }); it("should show confirmation step with selected models", async () => { render(); - // Navigate to confirm step const nextButton = screen.getByRole("button", { name: "Next" }); await act(async () => { fireEvent.click(nextButton); diff --git a/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.tsx index 2d0ae1a0e2b..28a34ee1f1a 100644 --- a/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.tsx @@ -1,11 +1,15 @@ import React, { useState, useCallback, useEffect } from "react"; -import { Modal, Form, Steps, Button, Checkbox } from "antd"; -import { Text, Title, Badge } from "@tremor/react"; +import { Loader2 } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { cn } from "@/lib/cva.config"; import { makeModelGroupPublic } from "../../networking"; import ModelFilters from "../../model_filters"; import NotificationsManager from "../../molecules/notifications_manager"; -const { Step } = Steps; +const STEP_TITLES = ["Select Models", "Confirm"]; interface ModelGroupInfo { model_group: string; @@ -44,13 +48,11 @@ const MakeModelPublicForm: React.FC = ({ const [selectedModels, setSelectedModels] = useState>(new Set()); const [filteredData, setFilteredData] = useState([]); const [loading, setLoading] = useState(false); - const [form] = Form.useForm(); const handleClose = () => { setCurrentStep(0); setSelectedModels(new Set()); setFilteredData([]); - form.resetFields(); onClose(); }; @@ -138,23 +140,24 @@ const MakeModelPublicForm: React.FC = ({ return (
- Select Models to Make Public +

Select Models to Make Public

- handleSelectAll(e.target.checked)} - disabled={filteredData.length === 0} - > +
- +

Select the models you want to be visible on the public model hub. Users will still require a valid Virtual Key to use these models. - +

{/* Filters */} = ({
{filteredData.length === 0 ? (
- No models match the current filters. +

No models match the current filters.

) : ( filteredData.map((model) => ( @@ -178,20 +181,16 @@ const MakeModelPublicForm: React.FC = ({ > handleModelSelection(model.model_group, e.target.checked)} + onCheckedChange={(checked) => handleModelSelection(model.model_group, checked === true)} /> -
-
- {model.model_group} - {model.mode && ( - - {model.mode} - - )} +
+
+

{model.model_group}

+ {model.mode && {model.mode}}
{model.providers.map((provider) => ( - + {provider} ))} @@ -205,9 +204,9 @@ const MakeModelPublicForm: React.FC = ({ {selectedModels.size > 0 && (
- +

{selectedModels.size} model{selectedModels.size !== 1 ? "s" : ""} selected - +

)}
@@ -217,29 +216,29 @@ const MakeModelPublicForm: React.FC = ({ const renderStep2Content = () => { return (
- Confirm Making Models Public +

Confirm Making Models Public

- +

Warning: Once you make these models public, anyone who can go to the{" "} /ui/model_hub_table will be able to know they exist on the proxy. - +

- Models to be made public: +

Models to be made public:

{Array.from(selectedModels).map((modelGroup) => { const model = modelHubData.find((m) => m.model_group === modelGroup); return (
-
- {modelGroup} +
+

{modelGroup}

{model && (
{model.providers.map((provider) => ( - + {provider} ))} @@ -254,10 +253,10 @@ const MakeModelPublicForm: React.FC = ({
- +

Total: {selectedModels.size} model{selectedModels.size !== 1 ? "s" : ""} will be made public - +

); @@ -277,7 +276,7 @@ const MakeModelPublicForm: React.FC = ({ const renderStepButtons = () => { return (
- @@ -289,7 +288,8 @@ const MakeModelPublicForm: React.FC = ({ )} {currentStep === 1 && ( - )} @@ -299,24 +299,42 @@ const MakeModelPublicForm: React.FC = ({ }; return ( - -
- - - - + !open && handleClose()} disablePointerDismissal> + + + Make Models Public + - {renderStepContent()} - {renderStepButtons()} - - +
+
    + {STEP_TITLES.map((title, index) => ( +
  1. + + {index + 1} + + + {title} + +
    2. + ))} +
+ + {renderStepContent()} + {renderStepButtons()} +
+ + ); }; diff --git a/ui/litellm-dashboard/src/components/BetaBadge.tsx b/ui/litellm-dashboard/src/components/BetaBadge.tsx index 7c4ef04417e..4e2195d1c36 100644 --- a/ui/litellm-dashboard/src/components/BetaBadge.tsx +++ b/ui/litellm-dashboard/src/components/BetaBadge.tsx @@ -1,4 +1,4 @@ -import { Badge } from "antd"; +import { Badge } from "@/components/ui/badge"; import { useDisableShowNewBadge } from "@/app/(dashboard)/hooks/useDisableShowNewBadge"; export default function BetaBadge({ children, dot = false }: { children?: React.ReactNode; dot?: boolean }) { @@ -8,11 +8,14 @@ export default function BetaBadge({ children, dot = false }: { children?: React. return children ? <>{children} : null; } + const badge = dot ? : Beta; + return children ? ( - + {children} - + {badge} + ) : ( - + badge ); } diff --git a/ui/litellm-dashboard/src/components/DebugWarningBanner.tsx b/ui/litellm-dashboard/src/components/DebugWarningBanner.tsx index 94474e78b14..9591fe4bb1f 100644 --- a/ui/litellm-dashboard/src/components/DebugWarningBanner.tsx +++ b/ui/litellm-dashboard/src/components/DebugWarningBanner.tsx @@ -1,7 +1,8 @@ "use client"; import React from "react"; -import { Alert } from "antd"; +import { TriangleAlert } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; interface DebugWarningBannerProps { @@ -17,19 +18,14 @@ export const DebugWarningBanner: React.FC = ({ accessTo } return ( - - Detailed debug logging (LITELLM_LOG=DEBUG) is currently enabled. This mode logs extensive - diagnostic information and will significantly degrade performance. It should only be used for troubleshooting - and disabled in production environments. - - } - type="warning" - showIcon - banner - style={{ marginBottom: 0, borderRadius: 0 }} - /> + + + Performance Warning: Detailed Debug Mode Active + + Detailed debug logging (LITELLM_LOG=DEBUG) is currently enabled. This mode logs extensive + diagnostic information and will significantly degrade performance. It should only be used for troubleshooting + and disabled in production environments. + + ); }; diff --git a/ui/litellm-dashboard/src/components/DeprecationBanner.test.tsx b/ui/litellm-dashboard/src/components/DeprecationBanner.test.tsx new file mode 100644 index 00000000000..596ad2a1626 --- /dev/null +++ b/ui/litellm-dashboard/src/components/DeprecationBanner.test.tsx @@ -0,0 +1,44 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it } from "vitest"; +import { DeprecationBanner } from "./DeprecationBanner"; + +describe("DeprecationBanner", () => { + it("names the deprecated feature in the heading and the body", () => { + render(); + + expect(screen.getByText("Memory is on a draft deprecation list")).toBeInTheDocument(); + expect(screen.getByText(/Memory is one of several experimental features/)).toBeInTheDocument(); + }); + + it("states the target removal date and that the list is not final", () => { + render(); + + expect(screen.getByText(/as early as September 1, 2026/)).toBeInTheDocument(); + expect(screen.getByText(/This list is a draft and is not final/)).toBeInTheDocument(); + }); + + it("links to the deprecation discussion in a new tab without leaking the opener", () => { + render(); + + const link = screen.getByRole("link", { name: "deprecation discussion" }); + expect(link).toHaveAttribute("href", "https://github.com/BerriAI/litellm/discussions/32090"); + expect(link).toHaveAttribute("target", "_blank"); + expect(link).toHaveAttribute("rel", "noopener noreferrer"); + }); + + it("exposes a named close control", () => { + render(); + + expect(screen.getByRole("button", { name: /close/i })).toBeInTheDocument(); + }); + + it("hides the banner once the close control is used", async () => { + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("button", { name: /close/i })); + + expect(screen.queryByText("Memory is on a draft deprecation list")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/DeprecationBanner.tsx b/ui/litellm-dashboard/src/components/DeprecationBanner.tsx index 33c75ec22e3..9df34636f82 100644 --- a/ui/litellm-dashboard/src/components/DeprecationBanner.tsx +++ b/ui/litellm-dashboard/src/components/DeprecationBanner.tsx @@ -1,8 +1,8 @@ "use client"; -import React from "react"; +import React, { useState } from "react"; import Link from "next/link"; -import { Alert } from "antd"; +import { Info, X } from "lucide-react"; const DEPRECATION_DISCUSSION_URL = "https://github.com/BerriAI/litellm/discussions/32090"; const DEPRECATION_TARGET_DATE = "September 1, 2026"; @@ -11,21 +11,42 @@ interface DeprecationBannerProps { featureName: string; } -export const DeprecationBanner: React.FC = ({ featureName }) => ( - - {`${featureName} is one of several experimental features we're considering removing, potentially as early as ${DEPRECATION_TARGET_DATE}. This list is a draft and is not final. If you rely on this feature, please share feedback on the `} - - deprecation discussion - - . - - } - type="info" - showIcon - closable - style={{ marginBottom: 16 }} - /> -); +export const DeprecationBanner: React.FC = ({ featureName }) => { + const [isClosed, setIsClosed] = useState(false); + + if (isClosed) { + return null; + } + + return ( +
+ +
+

{`${featureName} is on a draft deprecation list`}

+

+ {`${featureName} is one of several experimental features we're considering removing, potentially as early as ${DEPRECATION_TARGET_DATE}. This list is a draft and is not final. If you rely on this feature, please share feedback on the `} + + deprecation discussion + + . +

+
+ +
+ ); +}; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportSummary.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportSummary.tsx index bec65db9309..e5d96e0b145 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportSummary.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportSummary.tsx @@ -1,5 +1,5 @@ import React from "react"; -import type { DateRangePickerValue } from "@tremor/react"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; interface ExportSummaryProps { dateRange: DateRangePickerValue; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx index 93c3ffc3766..e605733a03a 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx @@ -1,4 +1,4 @@ -import type { DateRangePickerValue } from "@tremor/react"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import { Download } from "lucide-react"; import React, { useState } from "react"; import { Button } from "@/components/ui/button"; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts index 5d7cf652cc9..2fbbb29a22c 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts @@ -1,4 +1,4 @@ -import type { DateRangePickerValue } from "@tremor/react"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import type { Team } from "@/components/key_team_helpers/key_list"; export type ExportFormat = "csv" | "json"; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index 08ca298c1f1..97f14e2d3d0 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -1,4 +1,4 @@ -import type { DateRangePickerValue } from "@tremor/react"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import Papa from "papaparse"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { EntitySpendData, ExportScope } from "./types"; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 9adcb50206d..de637d5d627 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -1,5 +1,5 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; -import type { DateRangePickerValue } from "@tremor/react"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import Papa from "papaparse"; import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types"; diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx index c3671fc9e2c..d41b29f8218 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx @@ -1,8 +1,9 @@ -import { CheckCircleOutlined, CloseOutlined, DownOutlined, WarningOutlined } from "@ant-design/icons"; +import { CircleCheck, ChevronDown, TriangleAlert, X } from "lucide-react"; import { useQuery } from "@tanstack/react-query"; import moment from "moment"; -import { Button, Spin } from "antd"; import React, { useState } from "react"; +import { Button } from "@/components/ui/button"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { uiSpendLogsCall } from "@/components/networking"; import { LogDetailsDrawer } from "@/components/view_logs/LogDetailsDrawer"; import type { LogEntry as ViewLogsLogEntry } from "@/components/view_logs/columns"; @@ -13,21 +14,21 @@ const actionConfig: Record< { icon: React.ElementType; color: string; bg: string; border: string; label: string } > = { blocked: { - icon: CloseOutlined, + icon: X, color: "text-red-600", bg: "bg-red-50", border: "border-red-200", label: "Blocked", }, passed: { - icon: CheckCircleOutlined, + icon: CircleCheck, color: "text-green-600", bg: "bg-green-50", border: "border-green-200", label: "Passed", }, flagged: { - icon: WarningOutlined, + icon: TriangleAlert, color: "text-amber-600", bg: "bg-amber-50", border: "border-amber-200", @@ -125,8 +126,8 @@ export function LogViewer({ {filters.map((f) => (

{log.input_snippet ?? log.input ?? "—"}

- + ); })} diff --git a/ui/litellm-dashboard/src/components/LicenseExpiryBanner.test.tsx b/ui/litellm-dashboard/src/components/LicenseExpiryBanner.test.tsx index d6b419ace7c..627d108e65f 100644 --- a/ui/litellm-dashboard/src/components/LicenseExpiryBanner.test.tsx +++ b/ui/litellm-dashboard/src/components/LicenseExpiryBanner.test.tsx @@ -42,18 +42,22 @@ describe("LicenseExpiryBannerView", () => { expect(container).toBeEmptyDOMElement(); }); - it("shows a dismissible amber warning within 30 days", () => { + it("shows a dismissible warning within 30 days", () => { const { container } = render(); + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(container.querySelector(".lucide-triangle-alert")).toBeInTheDocument(); expect(screen.getByText(/expires in 20 days/)).toBeInTheDocument(); - expect(container.querySelector(".ant-alert-warning")).toBeInTheDocument(); - expect(screen.queryByRole("button")).toBeInTheDocument(); + expect(screen.getByText(/Renew before it lapses to keep enterprise features/)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /close/i })).toBeInTheDocument(); expect(screen.getByRole("link", { name: "sales@berri.ai" })).toHaveAttribute("href", "mailto:sales@berri.ai"); }); - it("shows a non-dismissible red critical alert within 7 days", () => { + it("shows a non-dismissible critical alert within 7 days", () => { const { container } = render(); + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(container.querySelector(".lucide-circle-alert")).toBeInTheDocument(); expect(screen.getByText(/expires in 5 days/)).toBeInTheDocument(); - expect(container.querySelector(".ant-alert-error")).toBeInTheDocument(); + expect(screen.getByText(/Renew now to avoid losing enterprise features/)).toBeInTheDocument(); expect(screen.queryByRole("button")).not.toBeInTheDocument(); }); @@ -62,18 +66,19 @@ describe("LicenseExpiryBannerView", () => { expect(screen.getByText(/expires today/)).toBeInTheDocument(); }); - it("shows a non-dismissible red expired alert stating features are disabled", () => { + it("shows a non-dismissible expired alert stating features are disabled", () => { const { container } = render(); + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(container.querySelector(".lucide-circle-alert")).toBeInTheDocument(); expect(screen.getByText(/expired on/)).toBeInTheDocument(); expect(screen.getByText(/features are now disabled/i)).toBeInTheDocument(); - expect(container.querySelector(".ant-alert-error")).toBeInTheDocument(); expect(screen.queryByRole("button")).not.toBeInTheDocument(); }); it("hides the warning after dismissal and stays hidden within the session", () => { const expiration = daysFromNow(20); const { unmount } = render(); - fireEvent.click(screen.getByRole("button")); + fireEvent.click(screen.getByRole("button", { name: /close/i })); expect(screen.queryByText(/expires in 20 days/)).not.toBeInTheDocument(); unmount(); diff --git a/ui/litellm-dashboard/src/components/LicenseExpiryBanner.tsx b/ui/litellm-dashboard/src/components/LicenseExpiryBanner.tsx index c3b20b5fac0..5867a45bc31 100644 --- a/ui/litellm-dashboard/src/components/LicenseExpiryBanner.tsx +++ b/ui/litellm-dashboard/src/components/LicenseExpiryBanner.tsx @@ -1,7 +1,9 @@ "use client"; import React, { useState } from "react"; -import { Alert } from "antd"; +import { CircleAlert, TriangleAlert, X } from "lucide-react"; +import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { Button } from "@/components/ui/button"; import { LicenseInfo } from "@/components/networking"; import { useLicenseInfo } from "@/app/(dashboard)/hooks/license/useLicenseInfo"; import { formatExpiryDate, getDaysUntilExpiration, getLicenseExpiryTier } from "@/utils/licenseUtils"; @@ -76,16 +78,22 @@ export const LicenseExpiryBannerView: React.FC = ( }; return ( - + + {tier === "warning" ? ( + + ) : ( + + )} + {message} + {description} + {isDismissible && ( + + + + )} + ); }; diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx index e253bc4c0ef..34a21122027 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx @@ -1,6 +1,6 @@ import type { ProxyModel } from "@/app/(dashboard)/hooks/models/useModels"; import type { Organization } from "@/components/networking"; -import { screen, waitFor } from "@testing-library/react"; +import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; @@ -22,64 +22,6 @@ vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({ useCurrentUser: vi.fn(), })); -vi.mock("antd", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - Select: ({ - value, - onChange, - options, - "data-testid": dataTestId, - allowClear, - maxTagCount, - maxTagPlaceholder, - mode, - ...props - }: any) => { - // Simulate maxTagCount responsive behavior - if value length > 5, call maxTagPlaceholder - const shouldShowPlaceholder = maxTagCount === "responsive" && Array.isArray(value) && value.length > 5; - const visibleValues = shouldShowPlaceholder ? value.slice(0, 5) : value; - const omittedValues = shouldShowPlaceholder ? value.slice(5).map((v: string) => ({ value: v, label: v })) : []; - - return ( -
- - {shouldShowPlaceholder && maxTagPlaceholder && ( -
{maxTagPlaceholder(omittedValues)}
- )} -
- ); - }, - Skeleton: { - Input: ({ active, block }: any) =>
, - }, - Tooltip: ({ children }: { children: React.ReactNode }) => <>{children}, - }; -}); - import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; @@ -108,6 +50,14 @@ const createMockOrganization = (models: string[]): Organization => ({ members: null, }); +const openModelList = async (user: ReturnType) => { + await user.click(screen.getAllByRole("combobox")[0]); + await screen.findByRole("listbox"); +}; + +const expectOffered = (label: string) => expect(screen.queryAllByText(label).length).toBeGreaterThan(0); +const expectNotOffered = (label: string) => expect(screen.queryAllByText(label)).toHaveLength(0); + describe("ModelSelect", () => { const mockProxyModels: ProxyModel[] = [ { id: "gpt-4", object: "model", created: 1234567890, owned_by: "openai" }, @@ -138,21 +88,26 @@ describe("ModelSelect", () => { } as any); }); - it("should render with all option groups", async () => { + it("should offer every model and wildcard under its group heading", async () => { + const user = userEvent.setup(); renderWithProviders( , ); - await waitFor(() => { - expect(screen.getByTestId("model-select")).toBeInTheDocument(); - expect(screen.getByText("gpt-4")).toBeInTheDocument(); - expect(screen.getByText("claude-3")).toBeInTheDocument(); - expect(screen.getByText("All Openai models")).toBeInTheDocument(); - expect(screen.getByText("All Anthropic models")).toBeInTheDocument(); - }); + await openModelList(user); + + expectOffered("Wildcard Options"); + expectOffered("gpt-4"); + expectOffered("claude-3"); + expectOffered("All Openai models"); + expectOffered("All Anthropic models"); }); - it("should show skeleton loader when any data is loading", () => { + it("should offer nothing to select while any dependency is loading", () => { + const { unmount: unmountReady } = renderWithProviders(); + expect(screen.getAllByRole("combobox")).toHaveLength(1); + unmountReady(); + const loadingScenarios = [ { hook: mockUseAllProxyModels, context: "user" as const }, { hook: mockUseTeam, context: "team" as const, props: { teamID: "team-1" } }, @@ -168,30 +123,41 @@ describe("ModelSelect", () => { const { unmount } = renderWithProviders(); - expect(screen.getByTestId("skeleton-input")).toBeInTheDocument(); + expect(screen.queryAllByRole("combobox")).toHaveLength(0); unmount(); }); }); - it("should handle model selection and onChange", async () => { + it("should report the picked model to onChange", async () => { const user = userEvent.setup(); renderWithProviders( , ); - await waitFor(() => { - expect(screen.getByTestId("model-select")).toBeInTheDocument(); - }); + await openModelList(user); + await user.click(screen.getAllByText("gpt-4")[0]); - const select = screen.getByRole("listbox"); - await user.selectOptions(select, "gpt-4"); expect(mockOnChange).toHaveBeenCalledWith(["gpt-4"]); - - await user.selectOptions(select, ["gpt-4", "claude-3"]); - expect(mockOnChange).toHaveBeenCalled(); }); - it("should handle special options correctly", async () => { + it("should append a second model to the existing selection", async () => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await openModelList(user); + await user.click(screen.getAllByText("claude-3")[0]); + + expect(mockOnChange).toHaveBeenCalledWith(["gpt-4", "claude-3"]); + }); + + it("should offer both special options when they are enabled", async () => { const user = userEvent.setup(); mockUseOrganization.mockReturnValue({ data: createMockOrganization(["all-proxy-models"]), @@ -207,33 +173,32 @@ describe("ModelSelect", () => { />, ); - await waitFor(() => { - expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); - expect(screen.getByText("No Default Models")).toBeInTheDocument(); - }); + await openModelList(user); - const select = screen.getByRole("listbox"); - await user.selectOptions(select, ["all-proxy-models", "no-default-models"]); - expect(mockOnChange).toHaveBeenCalledWith(["no-default-models"]); + expectOffered("Special Options"); + expectOffered("All Proxy Models"); + expectOffered("No Default Models"); }); - it("should disable models when special option is selected", async () => { + it("should replace an existing selection when a special option is picked", async () => { + const user = userEvent.setup(); + renderWithProviders( , ); - await waitFor(() => { - expect(screen.getByRole("option", { name: "gpt-4" })).toBeDisabled(); - expect(screen.getByRole("option", { name: "All Openai models" })).toBeDisabled(); - }); + await openModelList(user); + await user.click(screen.getAllByText("No Default Models")[0]); + + expect(mockOnChange).toHaveBeenCalledWith(["no-default-models"]); }); - it("should filter models based on context", async () => { + it("should filter the offered models by context", async () => { const testCases = [ { name: "user context with includeUserModels", @@ -340,6 +305,7 @@ describe("ModelSelect", () => { ]; for (const testCase of testCases) { + const user = userEvent.setup(); testCase.setup(); const { unmount } = renderWithProviders( { />, ); - await waitFor(() => { - testCase.expectedVisible.forEach((model) => { - expect(screen.getByText(model)).toBeInTheDocument(); - }); - testCase.expectedHidden.forEach((model) => { - expect(screen.queryByText(model)).not.toBeInTheDocument(); - }); - }); + await openModelList(user); + testCase.expectedVisible.forEach(expectOffered); + testCase.expectedHidden.forEach(expectNotOffered); unmount(); vi.clearAllMocks(); @@ -368,7 +329,7 @@ describe("ModelSelect", () => { } }); - it("should show All Proxy Models option based on conditions", async () => { + it("should offer All Proxy Models only when the context allows it", async () => { const testCases = [ { name: "when showAllProxyModelsOverride is true", @@ -426,6 +387,7 @@ describe("ModelSelect", () => { ]; for (const testCase of testCases) { + const user = userEvent.setup(); testCase.setup(); const { unmount } = renderWithProviders( { />, ); - await waitFor(() => { - if (testCase.shouldShow) { - expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); - } else { - expect(screen.queryByText("All Proxy Models")).not.toBeInTheDocument(); - expect(screen.getByText("No Default Models")).toBeInTheDocument(); - } - }); + await openModelList(user); + if (testCase.shouldShow) { + expectOffered("All Proxy Models"); + } else { + expectNotOffered("All Proxy Models"); + expectOffered("No Default Models"); + } unmount(); vi.clearAllMocks(); @@ -454,27 +415,6 @@ describe("ModelSelect", () => { } }); - it("should deduplicate models with same id", async () => { - const duplicateModels: ProxyModel[] = [ - { id: "gpt-4", object: "model", created: 1234567890, owned_by: "openai" }, - { id: "gpt-4", object: "model", created: 1234567890, owned_by: "openai" }, - ]; - - mockUseAllProxyModels.mockReturnValue({ - data: { data: duplicateModels }, - isLoading: false, - } as any); - - renderWithProviders( - , - ); - - await waitFor(() => { - const gpt4Options = screen.getAllByText("gpt-4"); - expect(gpt4Options.length).toBeGreaterThan(0); - }); - }); - it("should use custom dataTestId when provided", async () => { renderWithProviders( { />, ); - await waitFor(() => { - expect(screen.getByTestId("custom-test-id")).toBeInTheDocument(); - }); + expect(await screen.findByTestId("custom-test-id")).toBeInTheDocument(); }); it("should return all proxy models for team context when organization has empty models array", async () => { + const user = userEvent.setup(); mockUseTeam.mockReturnValue({ data: { team_id: "team-1", team_alias: "Test Team", models: [] }, isLoading: false, @@ -503,52 +442,66 @@ describe("ModelSelect", () => { renderWithProviders(); - await waitFor(() => { - expect(screen.getByText("gpt-4")).toBeInTheDocument(); - expect(screen.getByText("claude-3")).toBeInTheDocument(); - }); + await openModelList(user); + + expectOffered("gpt-4"); + expectOffered("claude-3"); }); - it("should disable No Default Models when all-proxy-models is selected", async () => { - mockUseOrganization.mockReturnValue({ - data: createMockOrganization(["all-proxy-models"]), - isLoading: false, - } as any); + it("should not offer a special options group when includeSpecialOptions is omitted", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await openModelList(user); + + expectNotOffered("Special Options"); + expectNotOffered("All Proxy Models"); + expectNotOffered("No Default Models"); + expectOffered("Models"); + }); + + it("should mark models and wildcards unselectable while a special option is selected", async () => { + const user = userEvent.setup(); renderWithProviders( , ); - await waitFor(() => { - const noDefaultOption = screen.getByRole("option", { name: "No Default Models" }); - expect(noDefaultOption).toBeDisabled(); - }); + await openModelList(user); + + expect(screen.getByRole("option", { name: "gpt-4" })).toHaveAttribute("aria-disabled", "true"); + expect(screen.getByRole("option", { name: "All Openai models" })).toHaveAttribute("aria-disabled", "true"); + expect(screen.getByRole("option", { name: "No Default Models" })).toHaveAttribute("aria-disabled", "true"); + expect(screen.getByRole("option", { name: "All Proxy Models" })).not.toHaveAttribute("aria-disabled", "true"); }); - it("should not render an empty optgroup when includeSpecialOptions is omitted", async () => { - renderWithProviders(); + it("should list a duplicated proxy model only once", async () => { + const user = userEvent.setup(); + mockUseAllProxyModels.mockReturnValue({ + data: { + data: [ + { id: "gpt-4", object: "model", created: 1234567890, owned_by: "openai" }, + { id: "gpt-4", object: "model", created: 1234567890, owned_by: "openai" }, + ], + }, + isLoading: false, + } as any); - await waitFor(() => { - expect(screen.getByTestId("model-select")).toBeInTheDocument(); - }); + renderWithProviders( + , + ); - const optgroups = document.querySelectorAll("optgroup"); - // Wildcard Options + Models — no blank leading group - expect(optgroups.length).toBe(2); - optgroups.forEach((g) => { - expect(g.getAttribute("label")).toBeTruthy(); - }); + await openModelList(user); + + expect(screen.getAllByRole("option", { name: "gpt-4" })).toHaveLength(1); }); - it("should render maxTagPlaceholder when many items are selected", async () => { - // Create many models to trigger maxTagCount responsive behavior - const manyModels: ProxyModel[] = Array.from({ length: 20 }, (_, i) => ({ + it("should collapse selections past the chip limit into a labelled overflow count", async () => { + const manyModels: ProxyModel[] = Array.from({ length: 8 }, (_, i) => ({ id: `model-${i}`, object: "model", created: 1234567890, @@ -560,22 +513,18 @@ describe("ModelSelect", () => { isLoading: false, } as any); - const selectedValues = manyModels.slice(0, 10).map((m) => m.id); - renderWithProviders( m.id)} context="user" options={{ showAllProxyModelsOverride: true }} />, ); - await waitFor(() => { - expect(screen.getByTestId("model-select")).toBeInTheDocument(); - // Verify maxTagPlaceholder is rendered with omitted values - expect(screen.getByTestId("max-tag-placeholder")).toBeInTheDocument(); - expect(screen.getByText(/\+5 more/)).toBeInTheDocument(); - }); + expect(await screen.findByText("+3 more")).toBeInTheDocument(); + expect(screen.getByLabelText("model-0")).toBeInTheDocument(); + expect(screen.getByLabelText("model-4")).toBeInTheDocument(); + expect(screen.queryByLabelText("model-5")).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx index e993fed2408..55aa1f1ec5f 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx @@ -2,7 +2,22 @@ import { ProxyModel, useAllProxyModels } from "@/app/(dashboard)/hooks/models/us import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; -import { Select, Skeleton, Tooltip } from "antd"; +import { + Combobox, + ComboboxChip, + ComboboxChips, + ComboboxChipsInput, + ComboboxCollection, + ComboboxContent, + ComboboxEmpty, + ComboboxGroup, + ComboboxItem, + ComboboxLabel, + ComboboxList, + ComboboxValue, +} from "@/components/ui/combobox"; +import { Skeleton } from "@/components/ui/skeleton"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { Organization, Team } from "../networking"; import { splitWildcardModels } from "./modelUtils"; @@ -21,6 +36,8 @@ export const MODEL_SENTINEL_OPTIONS = [ MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE, ] as const; +const MAX_VISIBLE_MODEL_CHIPS = 5; + export interface ModelSelectProps { teamID?: string; organizationID?: string; @@ -37,6 +54,17 @@ export interface ModelSelectProps { style?: React.CSSProperties; } +type ModelOption = { + label: string; + value: string; + disabled?: boolean; +}; + +type ModelOptionGroup = { + label: string; + items: ModelOption[]; +}; + type FilterContextArgs = { allProxyModels: string[]; selectedTeam?: Team; @@ -109,10 +137,11 @@ export const ModelSelect = (props: ModelSelectProps) => { showAllProxyModelsOverride || (organizationHasAllProxyModels && includeSpecialOptions) || context === "global"; if (isLoading) { - return ; + return ; } - const handleChange = (values: string[]) => { + const handleChange = (selected: ModelOption[]) => { + const values = selected.map((option) => option.value); const specialValues = values.filter(isSpecialOption); let finalValues: string[]; @@ -133,85 +162,122 @@ export const ModelSelect = (props: ModelSelectProps) => { }); const { wildcard, regular } = splitWildcardModels(filteredModels); - return ( - onChange?.(e.target.value)}> - {options?.map((opt: any) => ( - - ))} - - ), -})); - -// Mock icon -vi.mock("@ant-design/icons", () => ({ - CloudServerOutlined: () => , -})); - import WorkerDropdown from "./WorkerDropdown"; +async function openWorkerList(user: ReturnType) { + await user.click(screen.getByRole("combobox")); + await waitFor(() => { + expect(screen.getByRole("combobox")).toHaveAttribute("aria-expanded", "true"); + }); +} + describe("WorkerDropdown", () => { const mockOnWorkerSwitch = vi.fn(); const workers = [ @@ -61,31 +50,7 @@ describe("WorkerDropdown", () => { expect(container).toBeEmptyDOMElement(); }); - it("renders the select when isControlPlane and selectedWorker exist", () => { - mockUseWorker.mockReturnValue({ - isControlPlane: true, - selectedWorker: workers[0], - workers, - }); - - render(); - expect(screen.getByTestId("worker-select")).toBeInTheDocument(); - }); - - it("renders all worker options", () => { - mockUseWorker.mockReturnValue({ - isControlPlane: true, - selectedWorker: workers[0], - workers, - }); - - render(); - expect(screen.getByText("Worker 1")).toBeInTheDocument(); - expect(screen.getByText("Worker 2")).toBeInTheDocument(); - expect(screen.getByText("Worker 3")).toBeInTheDocument(); - }); - - it("sets current worker as selected value", () => { + it("renders a collapsed worker combobox when isControlPlane and selectedWorker exist", () => { mockUseWorker.mockReturnValue({ isControlPlane: true, selectedWorker: workers[1], @@ -93,37 +58,109 @@ describe("WorkerDropdown", () => { }); render(); - const select = screen.getByTestId("worker-select") as HTMLSelectElement; - expect(select.value).toBe("w2"); + expect(screen.getByRole("combobox")).toHaveAttribute("aria-expanded", "false"); }); - it("disables the currently selected worker in options", () => { + it("reveals every worker only once the combobox is opened", async () => { mockUseWorker.mockReturnValue({ isControlPlane: true, - selectedWorker: workers[0], + selectedWorker: workers[1], workers, }); - - render(); - const options = screen.getAllByRole("option"); - const selectedOption = options.find((opt) => (opt as HTMLOptionElement).value === "w1"); - expect(selectedOption).toBeDisabled(); - }); - - it("calls onWorkerSwitch when selection changes", async () => { - mockUseWorker.mockReturnValue({ - isControlPlane: true, - selectedWorker: workers[0], - workers, - }); - - render(); - const select = screen.getByTestId("worker-select"); - - const { default: userEvent } = await import("@testing-library/user-event"); const user = userEvent.setup(); - await user.selectOptions(select, "w2"); - expect(mockOnWorkerSwitch).toHaveBeenCalledWith("w2"); + render(); + expect(screen.queryAllByRole("option")).toHaveLength(0); + expect(screen.queryByText("Worker 1")).not.toBeInTheDocument(); + expect(screen.queryByText("Worker 3")).not.toBeInTheDocument(); + + await openWorkerList(user); + + await waitFor(() => { + expect(screen.getByText("Worker 1")).toBeInTheDocument(); + }); + expect(screen.getAllByText("Worker 2").length).toBeGreaterThan(0); + expect(screen.getByText("Worker 3")).toBeInTheDocument(); + }); + + it("marks exactly one option as selected, the current worker", async () => { + mockUseWorker.mockReturnValue({ + isControlPlane: true, + selectedWorker: workers[1], + workers, + }); + const user = userEvent.setup(); + + render(); + await openWorkerList(user); + + await waitFor(() => { + const selected = screen.getAllByRole("option").filter((o) => o.getAttribute("aria-selected") === "true"); + expect(selected).toHaveLength(1); + expect(selected[0]).toHaveAccessibleName("Worker 2"); + }); + }); + + it("calls onWorkerSwitch with the id of the worker that was picked", async () => { + mockUseWorker.mockReturnValue({ + isControlPlane: true, + selectedWorker: workers[1], + workers, + }); + const user = userEvent.setup(); + + render(); + await openWorkerList(user); + await waitFor(() => { + expect(screen.getByText("Worker 3")).toBeInTheDocument(); + }); + + fireEvent.click(screen.getByText("Worker 3")); + + expect(mockOnWorkerSwitch).toHaveBeenCalledWith("w3"); + }); + + it("does not call onWorkerSwitch when the already-current worker is picked", async () => { + mockUseWorker.mockReturnValue({ + isControlPlane: true, + selectedWorker: workers[1], + workers, + }); + const user = userEvent.setup(); + + render(); + await openWorkerList(user); + await waitFor(() => { + expect(screen.getByText("Worker 3")).toBeInTheDocument(); + }); + + for (const currentWorkerNode of screen.getAllByText("Worker 2")) { + fireEvent.click(currentWorkerNode); + } + + expect(mockOnWorkerSwitch).not.toHaveBeenCalled(); + }); + + it("filters the worker options by the typed search text", async () => { + mockUseWorker.mockReturnValue({ + isControlPlane: true, + selectedWorker: workers[1], + workers, + }); + const user = userEvent.setup(); + + render(); + await openWorkerList(user); + await waitFor(() => { + expect(screen.getByText("Worker 1")).toBeInTheDocument(); + }); + + await user.clear(screen.getByRole("combobox")); + await user.type(screen.getByRole("combobox"), "worker 3"); + + await waitFor(() => { + expect(screen.queryByText("Worker 1")).not.toBeInTheDocument(); + }); + expect(screen.getByText("Worker 3")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx b/ui/litellm-dashboard/src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx index 186cc611117..432bab8c9ef 100644 --- a/ui/litellm-dashboard/src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx @@ -1,35 +1,66 @@ "use client"; import React from "react"; -import { Select } from "antd"; -import { CloudServerOutlined } from "@ant-design/icons"; +import { Server } from "lucide-react"; +import { + Combobox, + ComboboxContent, + ComboboxEmpty, + ComboboxInput, + ComboboxItem, + ComboboxList, +} from "@/components/ui/combobox"; +import { InputGroupAddon } from "@/components/ui/input-group"; import { useWorker } from "@/hooks/useWorker"; interface WorkerDropdownProps { onWorkerSwitch: (workerId: string) => void; } +interface WorkerOption { + label: string; + value: string; + disabled: boolean; +} + const WorkerDropdown: React.FC = ({ onWorkerSwitch }) => { const { isControlPlane, selectedWorker, workers } = useWorker(); if (!isControlPlane || !selectedWorker) return null; + const options: WorkerOption[] = workers.map((w) => ({ + label: w.name, + value: w.worker_id, + disabled: w.worker_id === selectedWorker.worker_id, + })); + return ( - } /> - + {OPTIONAL_FIELDS.map((field) => ( + + {field.kind === "duration" ? ( + } /> + ) : ( + + )} + + ))} 0 && (
- -
diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.tsx index 938e1104301..3efdcfd6b51 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.tsx @@ -4,9 +4,9 @@ * Reuses FallbackGroupConfig with the primary model locked */ -import { Button } from "antd"; +import { Button } from "@/components/ui/button"; import { useQuery } from "@tanstack/react-query"; -import { Pencil } from "lucide-react"; +import { LoaderCircle, Pencil } from "lucide-react"; import React, { useMemo, useState } from "react"; import { fetchAvailableModels } from "@/components/llm_calls/fetch_models"; import NotificationManager from "../../../molecules/notifications_manager"; @@ -88,16 +88,11 @@ export default function EditFallbacks({ disablePrimaryModel />
- -
diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx index 4aa9fb15705..f82780f0c73 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx @@ -1,7 +1,7 @@ import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; -import { ArrowRightIcon, PencilAltIcon, PlayIcon, TrashIcon } from "@heroicons/react/outline"; -import { Icon, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow } from "@tremor/react"; -import { Tooltip, Typography } from "antd"; +import { ArrowRight, Pencil, Play, Trash2 } from "lucide-react"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import openai from "openai"; import React, { useEffect, useState } from "react"; import DeleteResourceModal from "../../../common_components/DeleteResourceModal"; @@ -18,12 +18,14 @@ type Fallbacks = FallbackEntry[]; const modelCardClass = "inline-flex items-center gap-2 px-2.5 py-1 rounded-md border border-gray-200 bg-gray-50 text-sm font-medium text-gray-800 shrink-0"; +const iconWrapperClass = "inline-flex shrink-0 items-center justify-center px-1.5 py-1.5"; + function renderModelNameCell(modelName: string, getProviderFromModel?: (modelName: string) => string): React.ReactNode { const provider = getProviderFromModel?.(modelName) ?? modelName; return ( - {modelName} + {modelName} ); } @@ -41,19 +43,23 @@ function renderFallbacksChain( return ( - {modelName} + {modelName} ); }; return ( - + {list.map((model, i) => ( - {i > 0 && } + {i > 0 && ( + + + + )} ))} @@ -248,7 +254,7 @@ const Fallbacks: React.FC = ({ accessToken, userRole, userID }) const canModify = isProxyAdminRole(userRole ?? ""); return ( - <> + {canModify && ( = ({ accessToken, userRole, userID }) )} {!hasFallbacks ? (
- + No fallbacks configured. Add fallbacks to automatically try another model when the primary fails. - +
) : (
- + - Model Name - Fallbacks - Actions + Model Name + Fallbacks + Actions - + {routerSettings["fallbacks"].map((item: FallbackEntry, index: number) => Object.entries(item).map(([key, value]) => ( - {renderModelNameCell(key, getProviderFromModel)} - + + {renderModelNameCell(key, getProviderFromModel)} + + {renderFallbacksChain(key, Array.isArray(value) ? value : [], getProviderFromModel)} {canModify && ( <> - - testFallbackModelResponse(Object.keys(item)[0], accessToken || "")} - className="cursor-pointer hover:text-blue-600" - /> - - - handleEditClick(item)} - onKeyDown={(e) => e.key === "Enter" && handleEditClick(item)} - className="cursor-pointer inline-flex" + + testFallbackModelResponse(Object.keys(item)[0], accessToken || "")} + className={`${iconWrapperClass} cursor-pointer hover:text-blue-600`} + /> + } > - - + + + Test fallback - - handleDeleteClick(item)} - onKeyDown={(e) => e.key === "Enter" && handleDeleteClick(item)} - className="cursor-pointer inline-flex" + + handleEditClick(item)} + onKeyDown={(e) => e.key === "Enter" && handleEditClick(item)} + className={`${iconWrapperClass} cursor-pointer hover:text-blue-600`} + /> + } > - - + + + Edit fallback + + + handleDeleteClick(item)} + onKeyDown={(e) => e.key === "Enter" && handleDeleteClick(item)} + className={`${iconWrapperClass} cursor-pointer hover:text-red-600`} + /> + } + > + + + Delete fallback )} @@ -350,7 +373,7 @@ const Fallbacks: React.FC = ({ accessToken, userRole, userID }) onOk={handleDeleteConfirm} confirmLoading={isDeleting} /> - + ); }; diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx index 5d7196a84fe..df20924d321 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx @@ -14,6 +14,7 @@ import { PTU_RATE_FIELD, PTU_START_FIELD, ptuCountRules, + ptuNoUsageCostRule, ptuPairRule, ptuRateRules, ptuStartRequiredRule, @@ -261,7 +262,8 @@ const AdvancedSettings: React.FC = ({ @@ -269,7 +271,8 @@ const AdvancedSettings: React.FC = ({ @@ -277,7 +280,8 @@ const AdvancedSettings: React.FC = ({ @@ -286,7 +290,8 @@ const AdvancedSettings: React.FC = ({ @@ -297,7 +302,8 @@ const AdvancedSettings: React.FC = ({ diff --git a/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx b/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx index 7397eaa6b24..fff03ed8e82 100644 --- a/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx +++ b/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx @@ -1,4 +1,5 @@ -import { render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi } from "vitest"; import BulkCreateUsersButton from "./bulk_create_users_button"; @@ -20,9 +21,55 @@ vi.mock("./molecules/notifications_manager", () => ({ }, })); +const csvFile = () => + new File(["user_email,user_role\nnew.hire@example.com,internal_user\n"], "users.csv", { type: "text/csv" }); + +const openUploadStep = async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByText("+ Bulk Invite Users")); + return user; +}; + describe("BulkCreateUsersButton", () => { it("should render", () => { const { getByText } = render(); expect(getByText("+ Bulk Invite Users")).toBeInTheDocument(); }); + + it("parses a CSV chosen through the file input", async () => { + await openUploadStep(); + + const fileInput = document.querySelector('input[type="file"]') as HTMLInputElement; + fireEvent.change(fileInput, { target: { files: [csvFile()] } }); + + expect(await screen.findByText("new.hire@example.com")).toBeInTheDocument(); + }); + + it("parses a CSV dropped onto the drop zone", async () => { + await openUploadStep(); + + const dropZone = screen.getByLabelText(/drag and drop your csv file here/i).closest("label"); + fireEvent.drop(dropZone as HTMLLabelElement, { dataTransfer: { files: [csvFile()], types: ["Files"] } }); + + expect(await screen.findByText("new.hire@example.com")).toBeInTheDocument(); + }); + + it("exposes the drop zone as a label for a keyboard-reachable file input", async () => { + await openUploadStep(); + + const fileInput = screen.getByLabelText(/drag and drop your csv file here/i) as HTMLInputElement; + expect(fileInput).toHaveAttribute("type", "file"); + expect(fileInput).toHaveAttribute("accept", ".csv"); + expect(fileInput).toBeVisible(); + + const dropZone = fileInput.closest("label") as HTMLLabelElement; + expect(fileInput.id).not.toBe(""); + expect(dropZone.htmlFor).toBe(fileInput.id); + + const danglingLabels = [...document.querySelectorAll("label[for]")].filter( + (label) => document.getElementById(label.getAttribute("for") as string) === null, + ); + expect(danglingLabels).toEqual([]); + }); }); diff --git a/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx b/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx index 8faff9ca72f..faaf2cbf7a5 100644 --- a/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx +++ b/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx @@ -1,14 +1,8 @@ import React, { useState, useEffect } from "react"; -import { Text } from "@tremor/react"; -import { Button, Modal, Table, Upload, Typography } from "antd"; -import { - UploadOutlined, - DownloadOutlined, - WarningOutlined, - FileTextOutlined, - DeleteOutlined, - FileExclamationOutlined, -} from "@ant-design/icons"; +import { Button, buttonVariants } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { Download, FileText, FileWarning, Trash2, TriangleAlert, Upload } from "lucide-react"; import { userCreateCall, invitationCreateCall, getProxyUISettings } from "./networking"; import Papa from "papaparse"; import { CheckCircleIcon, XCircleIcon, ExclamationIcon } from "@heroicons/react/outline"; @@ -38,6 +32,8 @@ interface UserData { invitation_link?: string; } +const PREVIEW_PAGE_SIZE = 5; + // Define an interface for the UI settings interface UISettings { PROXY_BASE_URL: string | null; @@ -61,6 +57,9 @@ const BulkCreateUsersButton: React.FC = ({ const [selectedFile, setSelectedFile] = useState(null); const [uiSettings, setUISettings] = useState(null); const [baseUrl, setBaseUrl] = useState("http://localhost:4000"); + const [isDraggingOver, setIsDraggingOver] = useState(false); + const [pageIndex, setPageIndex] = useState(0); + const csvInputId = React.useId(); useEffect(() => { // Get UI settings @@ -93,7 +92,7 @@ const BulkCreateUsersButton: React.FC = ({ if (file.type !== "text/csv" && !file.name.endsWith(".csv")) { setFileError(`Invalid file type: ${file.name}. Please upload a CSV file (.csv extension).`); NotificationsManager.fromBackend("Invalid file type. Please upload a CSV file."); - return false; + return; } // Check file size (limit to 5MB) @@ -101,7 +100,7 @@ const BulkCreateUsersButton: React.FC = ({ setFileError( `File is too large (${(file.size / (1024 * 1024)).toFixed(1)} MB). Please upload a CSV file smaller than 5MB.`, ); - return false; + return; } Papa.parse(file, { @@ -262,7 +261,27 @@ const BulkCreateUsersButton: React.FC = ({ }, header: false, }); - return false; + }; + + const handleFileInputChange = (event: React.ChangeEvent) => { + const file = event.target.files?.[0]; + if (file) { + handleFileUpload(file); + } + }; + + const handleDragOver = (event: React.DragEvent) => { + event.preventDefault(); + setIsDraggingOver(true); + }; + + const handleDrop = (event: React.DragEvent) => { + event.preventDefault(); + setIsDraggingOver(false); + const file = event.dataTransfer.files?.[0]; + if (file) { + handleFileUpload(file); + } }; const removeSelectedFile = () => { @@ -273,6 +292,12 @@ const BulkCreateUsersButton: React.FC = ({ setFileError(null); }; + const resetParsedData = () => { + setParsedData([]); + setParseError(null); + setPageIndex(0); + }; + const handleBulkCreate = async () => { setIsProcessing(true); const updatedData = parsedData.map((user) => ({ ...user, status: "pending" })); @@ -434,340 +459,395 @@ const BulkCreateUsersButton: React.FC = ({ window.URL.revokeObjectURL(url); }; - const columns = [ - { - title: "Row", - dataIndex: "rowNumber", - key: "rowNumber", - width: 80, - }, - { - title: "Email", - dataIndex: "user_email", - key: "user_email", - }, - { - title: "Role", - dataIndex: "user_role", - key: "user_role", - }, - { - title: "Teams", - dataIndex: "teams", - key: "teams", - }, - { - title: "Budget", - dataIndex: "max_budget", - key: "max_budget", - }, - { - title: "Status", - key: "status", - render: (_: any, record: UserData) => { - if (!record.isValid) { - return ( -
-
- - Invalid -
- {record.error && {record.error}} -
- ); - } - if (!record.status || record.status === "pending") { - return Pending; - } - if (record.status === "success") { - return ( -
-
- - Success -
- {record.invitation_link && ( -
-
- {record.invitation_link} - NotificationsManager.success("Invitation link copied!")} - > - - -
-
- )} -
- ); - } - return ( -
-
- - Failed -
- {record.error && {JSON.stringify(record.error)}} + const renderStatusCell = (record: UserData) => { + if (!record.isValid) { + return ( +
+
+ + Invalid
- ); - }, - }, - ]; + {record.error && {record.error}} +
+ ); + } + if (!record.status || record.status === "pending") { + return Pending; + } + if (record.status === "success") { + return ( +
+
+ + Success +
+ {record.invitation_link && ( +
+
+ {record.invitation_link} + NotificationsManager.success("Invitation link copied!")} + > + + +
+
+ )} +
+ ); + } + return ( +
+
+ + Failed +
+ {record.error && {JSON.stringify(record.error)}} +
+ ); + }; + + const pageCount = Math.max(1, Math.ceil(parsedData.length / PREVIEW_PAGE_SIZE)); + const currentPage = Math.min(pageIndex, pageCount - 1); + const visibleRows = parsedData.slice(currentPage * PREVIEW_PAGE_SIZE, (currentPage + 1) * PREVIEW_PAGE_SIZE); return ( <> - - setIsModalVisible(false)} - bodyStyle={{ maxHeight: "70vh", overflow: "auto" }} - footer={null} - > -
- {/* Step indicator */} - {parsedData.length === 0 ? ( -
-
-
- 1 -
-

Download and fill the template

-
- -
-

Add multiple users at once by following these steps:

-
    -
  1. Download our CSV template
    2. -
  2. Add your users' information to the spreadsheet
    3. -
  3. Save the file and upload it here
    4. -
  4. After creation, download the results file containing the Virtual Keys for each user
    5. -
- -
-

Template Column Names

-
-
-
-
-

user_email

-

User's email address (required)

-
-
-
-
-
-

user_role

-

- User's role (one of: "proxy_admin", "proxy_admin_viewer", - "internal_user", "internal_user_viewer") -

-
-
-
-
-
-

teams

-

- Comma-separated team IDs (e.g., "team-1,team-2") -

-
-
-
-
-
-

max_budget

-

Maximum budget as a number (e.g., "100")

-
-
-
-
-
-

budget_duration

-

- Budget reset period (e.g., "30d", "1mo") -

-
-
-
-
-
-

models

-

- Comma-separated allowed models (e.g., "gpt-3.5-turbo,gpt-4") -

-
-
+ !open && setIsModalVisible(false)}> + + + Bulk Invite Users + +
+ {/* Step indicator */} + {parsedData.length === 0 ? ( +
+
+
+ 1
+

Download and fill the template

- -
+
+

Add multiple users at once by following these steps:

+
    +
  1. Download our CSV template
    2. +
  2. Add your users' information to the spreadsheet
    3. +
  3. Save the file and upload it here
    4. +
  4. After creation, download the results file containing the Virtual Keys for each user
    5. +
-
-
- 2 -
-

Upload your completed CSV

-
- -
- {selectedFile ? ( -
-
-
- {fileError ? ( - - ) : ( - - )} +
+

Template Column Names

+
+
+
- - {selectedFile.name} - - - {(selectedFile.size / 1024).toFixed(1)} KB • {new Date().toLocaleDateString()} - +

user_email

+

User's email address (required)

- -
- - {fileError ? ( -
- - {fileError} -
- ) : ( - !csvStructureError && ( -
-
-
-
- Processing... +
+
+
+

user_role

+

+ User's role (one of: "proxy_admin", "proxy_admin_viewer", + "internal_user", "internal_user_viewer") +

+
+
+
+
+
+

teams

+

+ Comma-separated team IDs (e.g., "team-1,team-2") +

+
+
+
+
+
+

max_budget

+

Maximum budget as a number (e.g., "100")

+
+
+
+
+
+

budget_duration

+

+ Budget reset period (e.g., "30d", "1mo") +

+
+
+
+
+
+

models

+

+ Comma-separated allowed models (e.g., "gpt-3.5-turbo,gpt-4") +

- ) - )} -
- ) : ( - -
- -

Drag and drop your CSV file here

-

or

- -

Only CSV files (.csv) are supported

-
- - )} - - {csvStructureError && ( -
-
- -
- - CSV Structure Error - - - {csvStructureError} - - - Please download our template and ensure your CSV follows the required format. -
- )} -
-
- ) : ( -
-
-
- 3 + +
-

- {parsedData.some((user) => user.status === "success" || user.status === "failed") - ? "User Creation Results" - : "Review and create users"} -

-
- {parseError && ( -
-
- -
- {parseError} - {parsedData.some((user) => !user.isValid) && ( -
    -
  • Check the table below for specific errors in each row
    • -
  • - Common issues include invalid email formats, missing required fields, or incorrect role - values -
    • -
  • Fix these issues in your CSV file and upload again
    • -
+
+
+ 2 +
+

Upload your completed CSV

+
+ +
+ {selectedFile ? ( +
+
+
+ {fileError ? ( + + ) : ( + + )} +
+ + {selectedFile.name} + + + {(selectedFile.size / 1024).toFixed(1)} KB • {new Date().toLocaleDateString()} + +
+
+ +
+ + {fileError ? ( +
+ + {fileError} +
+ ) : ( + !csvStructureError && ( +
+
+
+
+ Processing... +
+ ) )}
-
-
- )} + ) : ( + + )} -
-
-
- {parsedData.some((user) => user.status === "success" || user.status === "failed") ? ( -
- Creation Summary - - {parsedData.filter((d) => d.status === "success").length} Successful - - {parsedData.some((d) => d.status === "failed") && ( - - {parsedData.filter((d) => d.status === "failed").length} Failed - + {csvStructureError && ( +
+
+ +
+ CSV Structure Error +

{csvStructureError}

+

+ Please download our template and ensure your CSV follows the required format. +

+
+
+
+ )} +
+
+ ) : ( +
+
+
+ 3 +
+

+ {parsedData.some((user) => user.status === "success" || user.status === "failed") + ? "User Creation Results" + : "Review and create users"} +

+
+ + {parseError && ( +
+
+ +
+

{parseError}

+ {parsedData.some((user) => !user.isValid) && ( +
    +
  • Check the table below for specific errors in each row
    • +
  • + Common issues include invalid email formats, missing required fields, or incorrect role + values +
    • +
  • Fix these issues in your CSV file and upload again
    • +
)}
- ) : ( -
- User Preview - - {parsedData.filter((d) => d.isValid).length} of {parsedData.length} users valid - +
+
+ )} + +
+
+
+ {parsedData.some((user) => user.status === "success" || user.status === "failed") ? ( +
+

Creation Summary

+

+ {parsedData.filter((d) => d.status === "success").length} Successful +

+ {parsedData.some((d) => d.status === "failed") && ( +

+ {parsedData.filter((d) => d.status === "failed").length} Failed +

+ )} +
+ ) : ( +
+

User Preview

+

+ {parsedData.filter((d) => d.isValid).length} of {parsedData.length} users valid +

+
+ )} +
+ + {!parsedData.some((user) => user.status === "success" || user.status === "failed") && ( +
+ +
)}
- {!parsedData.some((user) => user.status === "success" || user.status === "failed") && ( -
+ {parsedData.some((user) => user.status === "success") && ( +
+
+
+ +
+
+

User creation complete

+

+ Next step: Download the credentials file containing + Virtual Keys and invitation links. Users will need these Virtual Keys to make LLM requests + through LiteLLM. +

+
+
+
+ )} + +
+
+ + + Row + Email + Role + Teams + Budget + Status + + + + {visibleRows.map((record) => ( + + {record.rowNumber} + {record.user_email} + {record.user_role} + {record.teams} + {record.max_budget} + {renderStatusCell(record)} + + ))} + +
+
+ + {pageCount > 1 && ( +
+ + Page {currentPage + 1} of {pageCount} + + +
+ )} + + {!parsedData.some((user) => user.status === "success" || user.status === "failed") && ( +
+
)} -
- {parsedData.some((user) => user.status === "success") && ( -
-
-
- -
-
- User creation complete - - Next step: Download the credentials file containing - Virtual Keys and invitation links. Users will need these Virtual Keys to make LLM requests - through LiteLLM. - -
+ {parsedData.some((user) => user.status === "success" || user.status === "failed") && ( +
+ +
-
- )} - - (!record.isValid ? "bg-red-50" : "")} - /> - - {!parsedData.some((user) => user.status === "success" || user.status === "failed") && ( -
- - -
- )} - - {parsedData.some((user) => user.status === "success" || user.status === "failed") && ( -
- - -
- )} + )} + - - )} - - + )} + + + ); }; diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx index 7abe67f3a4f..58dc62ee6b2 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx @@ -1,11 +1,15 @@ import React, { useState, useEffect } from "react"; -import { Modal, Form, Steps, Button, Checkbox } from "antd"; -import { Text, Title, Badge } from "@tremor/react"; +import { Loader2 } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { cn } from "@/lib/cva.config"; import { enableClaudeCodePlugin, disableClaudeCodePlugin } from "../networking"; import NotificationsManager from "../molecules/notifications_manager"; import { Plugin } from "./types"; -const { Step } = Steps; +const STEP_TITLES = ["Select Skills", "Confirm"]; interface MakeSkillPublicFormProps { visible: boolean; @@ -25,12 +29,10 @@ const MakeSkillPublicForm: React.FC = ({ const [currentStep, setCurrentStep] = useState(0); const [selectedSkills, setSelectedSkills] = useState>(new Set()); const [loading, setLoading] = useState(false); - const [form] = Form.useForm(); const handleClose = () => { setCurrentStep(0); setSelectedSkills(new Set()); - form.resetFields(); onClose(); }; @@ -106,52 +108,44 @@ const MakeSkillPublicForm: React.FC = ({ const renderStep1 = () => (
- Select Skills to Publish - handleSelectAll(e.target.checked)} - disabled={skillsList.length === 0} - > +

Select Skills to Publish

+
- +

Selected skills will be visible to all users in the Skill Hub. Deselected skills will be unpublished. - +

{skillsList.length === 0 ? (
- No skills registered yet. +

No skills registered yet.

) : ( skillsList.map((skill) => (
handleSkillSelection(skill.name, e.target.checked)} + onCheckedChange={(checked) => handleSkillSelection(skill.name, checked === true)} /> -
+
- {skill.name} - {skill.enabled && ( - - Public - - )} +

{skill.name}

+ {skill.enabled && Public}
- {skill.description && ( - {skill.description} - )} + {skill.description &&

{skill.description}

}
- {skill.domain && ( - - {skill.domain} - - )} + {skill.domain && {skill.domain}}
)) )} @@ -160,9 +154,9 @@ const MakeSkillPublicForm: React.FC = ({ {selectedSkills.size > 0 && (
- +

{selectedSkills.size} skill{selectedSkills.size !== 1 ? "s" : ""} will be published - +

)}
@@ -170,29 +164,25 @@ const MakeSkillPublicForm: React.FC = ({ const renderStep2 = () => (
- Confirm Publish to Skill Hub +

Confirm Publish to Skill Hub

- +

Note: Published skills will be visible to all users in the Skill Hub tab. Skills not in the list below will be unpublished. - +

- Skills to be published: +

Skills to be published:

{Array.from(selectedSkills).map((name) => { const skill = skillsList.find((s) => s.name === name); return ( -
- {name} - {skill?.domain && ( - - {skill.domain} - - )} +
+

{name}

+ {skill?.domain && {skill.domain}}
); })} @@ -201,49 +191,68 @@ const MakeSkillPublicForm: React.FC = ({
- +

Total: {selectedSkills.size} skill{selectedSkills.size !== 1 ? "s" : ""} will be published - +

); return ( - -
- - - - + !open && handleClose()} disablePointerDismissal> + + + Publish to Skill Hub + - {currentStep === 0 ? renderStep1() : renderStep2()} +
+
    + {STEP_TITLES.map((title, index) => ( +
  1. + + {index + 1} + + + {title} + +
    2. + ))} +
-
- -
- {currentStep === 0 && ( - - )} - {currentStep === 1 && ( - - )} + {currentStep === 0 ? renderStep1() : renderStep2()} + +
+ +
+ {currentStep === 0 && ( + + )} + {currentStep === 1 && ( + + )} +
- - + +
); }; diff --git a/ui/litellm-dashboard/src/components/common_components/AutoRotationView.tsx b/ui/litellm-dashboard/src/components/common_components/AutoRotationView.tsx index d6445d1472d..4e8e0dc8121 100644 --- a/ui/litellm-dashboard/src/components/common_components/AutoRotationView.tsx +++ b/ui/litellm-dashboard/src/components/common_components/AutoRotationView.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { Text, Badge } from "@tremor/react"; +import { StatusBadge } from "@/components/shared/table_cells"; import { RefreshIcon, ClockIcon } from "@heroicons/react/outline"; interface AutoRotationViewProps { @@ -38,63 +38,55 @@ const AutoRotationView: React.FC = ({ const content = (
- {/* Status Section */}
- Auto-Rotation - - {autoRotate ? "Enabled" : "Disabled"} - +

Auto-Rotation

+ {autoRotate && rotationInterval && ( <> - • - Every {rotationInterval} +

•

+

Every {rotationInterval}

)}
- {/* Rotation History - Show if there's any rotation data OR if auto-rotation is enabled */} {(autoRotate || lastRotationAt || keyRotationAt || nextRotationAt) && (
- {/* Last Rotation - Show when available */} {lastRotationAt && ( -
- +
+
- Last Rotation - {formatTimestamp(lastRotationAt)} +

Last Rotation

+

{formatTimestamp(lastRotationAt)}

)} - {/* Next Scheduled Rotation - Show when available */} {(keyRotationAt || nextRotationAt) && ( -
- +
+
- Next Scheduled Rotation - {formatTimestamp(nextRotationAt || keyRotationAt || "")} +

Next Scheduled Rotation

+

{formatTimestamp(nextRotationAt || keyRotationAt || "")}

)} - {/* No rotation data message - Only show if auto-rotation is enabled but no data */} {autoRotate && !lastRotationAt && !keyRotationAt && !nextRotationAt && ( -
- - No rotation history available +
+ +

No rotation history available

)}
)} - {/* Disabled State - Only show if auto-rotation is disabled AND there's no rotation history */} {!autoRotate && !lastRotationAt && !keyRotationAt && !nextRotationAt && ( -
- - Auto-rotation is not enabled for this key +
+ +

Auto-rotation is not enabled for this key

)}
@@ -102,11 +94,11 @@ const AutoRotationView: React.FC = ({ if (variant === "card") { return ( -
-
+
+
- Auto-Rotation - Automatic key rotation settings and status for this key +

Auto-Rotation

+

Automatic key rotation settings and status for this key

{content} @@ -116,7 +108,7 @@ const AutoRotationView: React.FC = ({ return (
- Auto-Rotation +

Auto-Rotation

{content}
); diff --git a/ui/litellm-dashboard/src/components/common_components/DefaultProxyAdminTag.tsx b/ui/litellm-dashboard/src/components/common_components/DefaultProxyAdminTag.tsx index e42e3cee6da..9ec24bb929b 100644 --- a/ui/litellm-dashboard/src/components/common_components/DefaultProxyAdminTag.tsx +++ b/ui/litellm-dashboard/src/components/common_components/DefaultProxyAdminTag.tsx @@ -1,6 +1,4 @@ -import { Tag, Typography } from "antd"; - -const { Text } = Typography; +import { Badge } from "@/components/ui/badge"; const DEFAULT_USER_ID = "default_user_id"; @@ -8,15 +6,10 @@ interface DefaultProxyAdminTagProps { userId: string | null | undefined; } -/** - * Renders "Default Proxy Admin" as a blue Tag when the given userId is - * the well-known `default_user_id`, otherwise renders the raw value as - * plain text. - */ export default function DefaultProxyAdminTag({ userId }: DefaultProxyAdminTagProps) { if (userId === DEFAULT_USER_ID) { - return Default Proxy Admin; + return Default Proxy Admin; } - return {userId}; + return {userId}; } diff --git a/ui/litellm-dashboard/src/components/common_components/DeleteResourceModal.test.tsx b/ui/litellm-dashboard/src/components/common_components/DeleteResourceModal.test.tsx index e27a60cc866..465f7fcfcf0 100644 --- a/ui/litellm-dashboard/src/components/common_components/DeleteResourceModal.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/DeleteResourceModal.test.tsx @@ -159,6 +159,20 @@ describe("DeleteResourceModal", () => { expect(cancelButton).toBeDisabled(); }); + it("should call onCancel when escape is pressed and no deletion is in flight", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.keyboard("{Escape}"); + expect(mockOnCancel).toHaveBeenCalled(); + }); + + it("should ignore escape while confirmLoading is true so the modal cannot close mid-deletion", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.keyboard("{Escape}"); + expect(mockOnCancel).not.toHaveBeenCalled(); + }); + it("should disable delete button when confirmLoading is true even if requiredConfirmation matches", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/common_components/DeleteResourceModal.tsx b/ui/litellm-dashboard/src/components/common_components/DeleteResourceModal.tsx index 5a6160483a0..42baae3d86c 100644 --- a/ui/litellm-dashboard/src/components/common_components/DeleteResourceModal.tsx +++ b/ui/litellm-dashboard/src/components/common_components/DeleteResourceModal.tsx @@ -1,6 +1,10 @@ -import { Alert, Card, Descriptions, Input, Modal, Typography, theme } from "antd"; -import { ExclamationCircleOutlined } from "@ant-design/icons"; +import { CircleAlert } from "lucide-react"; import React, { useState, useEffect } from "react"; +import { Alert, AlertTitle } from "@/components/shared/Alert"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; interface DeleteResourceModalProps { isOpen: boolean; @@ -8,12 +12,11 @@ interface DeleteResourceModalProps { alertMessage?: string; message: string; resourceInformationTitle?: string; - resourceInformation?: Array< - { - label: string; - value: string | number | undefined | null; - } & Omit, "children"> - >; + resourceInformation?: Array<{ + label: string; + value: string | number | undefined | null; + code?: boolean; + }>; onCancel: () => void; onOk: () => void; confirmLoading: boolean; @@ -32,8 +35,6 @@ export default function DeleteResourceModal({ confirmLoading, requiredConfirmation, }: DeleteResourceModalProps) { - const { Text } = Typography; - const { token } = theme.useToken(); const [requiredConfirmationInput, setRequiredConfirmationInput] = useState(""); useEffect(() => { @@ -43,69 +44,69 @@ export default function DeleteResourceModal({ }, [isOpen]); return ( - -
- {alertMessage && } - - - {resourceInformation && - resourceInformation.map(({ label, value, ...textProps }) => ( - {label}}> - {value ?? "-"} - - ))} - - -
- {message} -
- {requiredConfirmation && ( -
- - Type - - {requiredConfirmation} - - to confirm deletion: - - setRequiredConfirmationInput(e.target.value)} - placeholder={requiredConfirmation} - className="rounded-md" - prefix={} - autoFocus - /> + !open && !confirmLoading && onCancel()}> + + + {title} + +
+ {alertMessage && ( + + {alertMessage} + + )} + + {resourceInformationTitle && ( + + {resourceInformationTitle} + + )} + +
+ {resourceInformation?.map(({ label, value, code }) => ( + +
{label}
+
{code ? {value ?? "-"} : value ?? "-"}
+ + ))} +
+ + +
+ {message}
- )} -
- + {requiredConfirmation && ( +
+

+ Type {requiredConfirmation} to confirm deletion: +

+ + + + + setRequiredConfirmationInput(e.target.value)} + placeholder={requiredConfirmation} + autoFocus + /> + +
+ )} +
+ + + + + + ); } diff --git a/ui/litellm-dashboard/src/components/common_components/IconActionButton/BaseActionButton.tsx b/ui/litellm-dashboard/src/components/common_components/IconActionButton/BaseActionButton.tsx index 1d6aa0ec73e..7e5f3fe5342 100644 --- a/ui/litellm-dashboard/src/components/common_components/IconActionButton/BaseActionButton.tsx +++ b/ui/litellm-dashboard/src/components/common_components/IconActionButton/BaseActionButton.tsx @@ -1,5 +1,4 @@ import { cx } from "@/lib/cva.config"; -import { Icon } from "@tremor/react"; import React from "react"; interface BaseActionButtonProps { @@ -10,16 +9,27 @@ interface BaseActionButtonProps { dataTestId?: string; } -export default function BaseActionButton({ icon, onClick, className, disabled, dataTestId }: BaseActionButtonProps) { +export default function BaseActionButton({ + icon: Icon, + onClick, + className, + disabled, + dataTestId, +}: BaseActionButtonProps) { return disabled ? ( - - ) : ( - + > + + + ) : ( + + + ); } diff --git a/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.test.tsx b/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.test.tsx index 718c458b455..9793a0e19ab 100644 --- a/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.test.tsx @@ -1,5 +1,6 @@ -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; -import { describe, expect, it } from "vitest"; +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; import TableIconActionButton, { TableIconActionButtonMap } from "./TableIconActionButton"; describe("TableIconActionButton", () => { @@ -12,27 +13,32 @@ describe("TableIconActionButton", () => { }); }); - it("should have a tooltip", () => { - render( {}} dataTestId="test-button" tooltipText="Edit" />); - const button = screen.getByTestId("test-button"); - const tooltipWrapper = button.closest("span"); - expect(tooltipWrapper).toBeInTheDocument(); + it("should call onClick when clicked", async () => { + const user = userEvent.setup(); + const onClick = vi.fn(); + render(); + + await user.click(screen.getByTestId("test-button")); + + expect(onClick).toHaveBeenCalledTimes(1); }); - it("should show tooltip when tooltipText is provided", async () => { + it("should not show the tooltip before the button is hovered", () => { render( {}} dataTestId="test-button" tooltipText="Edit item" />, ); - const button = screen.getByTestId("test-button"); - const buttonWrapper = button.closest("span"); + expect(screen.queryByText("Edit item")).not.toBeInTheDocument(); + }); - act(() => { - fireEvent.mouseEnter(buttonWrapper!); - }); + it("should show tooltip when tooltipText is provided", async () => { + const user = userEvent.setup(); + render( + {}} dataTestId="test-button" tooltipText="Edit item" />, + ); - await waitFor(() => { - expect(screen.getByText("Edit item")).toBeInTheDocument(); - }); + await user.hover(screen.getByTestId("test-button")); + + expect(await screen.findByText("Edit item")).toBeInTheDocument(); }); it("should render disabled state with disabled styling", () => { @@ -44,7 +50,27 @@ describe("TableIconActionButton", () => { expect(button).toHaveClass("cursor-not-allowed"); }); + it("should not call onClick when disabled", async () => { + const user = userEvent.setup(); + const onClick = vi.fn(); + render( + , + ); + + await user.click(screen.getByTestId("test-button")); + + expect(onClick).not.toHaveBeenCalled(); + }); + it("should show disabledTooltipText when disabled and disabledTooltipText is provided", async () => { + const user = userEvent.setup(); render( { disabledTooltipText="Cannot edit" />, ); - const button = screen.getByTestId("test-button"); - const buttonWrapper = button.closest("span"); - act(() => { - fireEvent.mouseEnter(buttonWrapper!); - }); + await user.hover(screen.getByTestId("test-button")); - await waitFor(() => { - expect(screen.getByText("Cannot edit")).toBeInTheDocument(); - }); + expect(await screen.findByText("Cannot edit")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx b/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx index 2f146aab723..9eeb6922116 100644 --- a/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx +++ b/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx @@ -1,3 +1,4 @@ +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { PencilAltIcon, PlayIcon, @@ -8,7 +9,6 @@ import { ExternalLinkIcon, ClipboardCopyIcon, } from "@heroicons/react/outline"; -import { Tooltip } from "antd"; import BaseActionButton from "../BaseActionButton"; export interface TableIconActionButtonProps { @@ -45,17 +45,21 @@ export default function TableIconActionButton({ variant, }: TableIconActionButtonProps) { const { icon, className } = TableIconActionButtonMap[variant]; + const title = disabled ? disabledTooltipText : tooltipText; + const button = ( + + ); + + if (!title) { + return {button}; + } + return ( - - - - - + + + }>{button} + {title} + + ); } diff --git a/ui/litellm-dashboard/src/components/common_components/LabeledField.test.tsx b/ui/litellm-dashboard/src/components/common_components/LabeledField.test.tsx index db6eb3356a8..a55381b344f 100644 --- a/ui/litellm-dashboard/src/components/common_components/LabeledField.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/LabeledField.test.tsx @@ -32,18 +32,22 @@ describe("LabeledField", () => { }); it("should not be copyable when value is empty", () => { - const { container } = render(); - // antd adds a .ant-typography-copy element when copyable; should not be present - expect(container.querySelector(".ant-typography-copy")).not.toBeInTheDocument(); + render(); + expect(screen.queryByRole("button", { name: "Copy User ID" })).not.toBeInTheDocument(); }); it("should not be copyable when value is default_user_id and defaultUserIdCheck is true", () => { - const { container } = render(); - expect(container.querySelector(".ant-typography-copy")).not.toBeInTheDocument(); + render(); + expect(screen.queryByRole("button", { name: "Copy User ID" })).not.toBeInTheDocument(); + }); + + it("should not be copyable when copyable is false", () => { + render(); + expect(screen.queryByRole("button", { name: "Copy User ID" })).not.toBeInTheDocument(); }); it("should be copyable when copyable is true and value is present", () => { - const { container } = render(); - expect(container.querySelector(".ant-typography-copy")).toBeInTheDocument(); + render(); + expect(screen.getByRole("button", { name: "Copy User ID" })).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/common_components/LabeledField.tsx b/ui/litellm-dashboard/src/components/common_components/LabeledField.tsx index 75a8236a2fa..6107448babe 100644 --- a/ui/litellm-dashboard/src/components/common_components/LabeledField.tsx +++ b/ui/litellm-dashboard/src/components/common_components/LabeledField.tsx @@ -1,9 +1,8 @@ import React from "react"; -import { Typography, Space } from "antd"; +import CopyButton from "@/components/shared/CopyButton"; +import { cx } from "@/lib/cva.config"; import DefaultProxyAdminTag from "./DefaultProxyAdminTag"; -const { Text } = Typography; - interface LabeledFieldProps { label: string; value: string; @@ -29,24 +28,20 @@ export default function LabeledField({ const valueEl = isDefaultUser ? ( ) : ( - - {displayValue} - + + + {displayValue} + + {isCopyable && } + ); return ( -
- - {icon} - - {label} - - -
{valueEl}
+
+
+ {icon} + {label} +
+
{valueEl}
); } diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx index d58094642d2..19c2377bafe 100644 --- a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx @@ -1,11 +1,18 @@ +import { Tooltip } from "@/components/atoms/Tooltip"; import { Member } from "@/components/networking"; -import { CrownOutlined, InfoCircleOutlined, UserAddOutlined, UserOutlined } from "@ant-design/icons"; -import { Button, Space, Table, Tag, Tooltip, Typography } from "antd"; -import type { ColumnsType } from "antd/es/table"; +import { StatusBadge } from "@/components/shared/table_cells"; +import { Button } from "@/components/ui/button"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { Crown, Info, User, UserPlus } from "lucide-react"; import React from "react"; import TableIconActionButton from "./IconActionButton/TableIconActionButtons/TableIconActionButton"; -const { Text } = Typography; +export interface MemberTableColumn { + title: React.ReactNode; + key: React.Key; + dataIndex?: keyof Member; + render?: (value: Member[keyof Member], member: Member, index: number) => React.ReactNode; +} export interface MemberTableProps { members: Member[]; @@ -15,11 +22,18 @@ export interface MemberTableProps { onAddMember?: () => void; roleColumnTitle?: string; roleTooltip?: string; - extraColumns?: ColumnsType; + extraColumns?: MemberTableColumn[]; showDeleteForMember?: (member: Member) => boolean; emptyText?: string; } +const extraColumnCell = (column: MemberTableColumn, member: Member, index: number): React.ReactNode => { + const value = column.dataIndex ? member[column.dataIndex] : undefined; + return column.render ? column.render(value, member, index) : value; +}; + +const STICKY_ACTIONS_CLASS = "sticky right-0 w-[120px] bg-background"; + export default function MemberTable({ members, canEdit, @@ -32,91 +46,96 @@ export default function MemberTable({ showDeleteForMember, emptyText, }: MemberTableProps) { - const baseColumns: ColumnsType = [ - { - title: "User Email", - dataIndex: "user_email", - key: "user_email", - render: (email: string | null) => {email || "-"}, - }, - { - title: "User ID", - dataIndex: "user_id", - key: "user_id", - render: (userId: string | null) => - userId === "default_user_id" ? Default Proxy Admin : {userId || "-"}, - }, - { - title: roleTooltip ? ( - - {roleColumnTitle} - - - - - ) : ( - roleColumnTitle - ), - dataIndex: "role", - key: "role", - render: (role: string) => ( - - {role?.toLowerCase() === "admin" || role?.toLowerCase() === "org_admin" ? ( - - ) : ( - - )} - {role || "-"} - - ), - }, - ...extraColumns, - { - title: "Actions", - key: "actions", - fixed: "right" as const, - width: 120, - render: (_: unknown, record: Member) => - canEdit ? ( - - onEdit(record)} - /> - {(!showDeleteForMember || showDeleteForMember(record)) && ( - onDelete(record)} - /> - )} - - ) : null, - }, - ]; - return ( - +
{members.length} Member{members.length !== 1 ? "s" : ""} -
record.user_id ?? record.user_email ?? JSON.stringify(record)} - pagination={false} - size="small" - scroll={{ x: "max-content" }} - locale={emptyText ? { emptyText } : undefined} - /> +
+ + + User Email + User ID + + {roleTooltip ? ( + + {roleColumnTitle} + + + + + ) : ( + roleColumnTitle + )} + + {extraColumns.map((column) => ( + {column.title} + ))} + Actions + + + + {members.length === 0 ? ( + + + {emptyText ?? "No data"} + + + ) : ( + members.map((member, memberIndex) => ( + + {member.user_email || "-"} + + {member.user_id === "default_user_id" ? ( + + ) : ( + member.user_id || "-" + )} + + + + {member.role?.toLowerCase() === "admin" || member.role?.toLowerCase() === "org_admin" ? ( + + ) : ( + + )} + {member.role || "-"} + + + {extraColumns.map((column) => ( + {extraColumnCell(column, member, memberIndex)} + ))} + + {canEdit ? ( + + onEdit(member)} + /> + {(!showDeleteForMember || showDeleteForMember(member)) && ( + onDelete(member)} + /> + )} + + ) : null} + + + )) + )} + +
{onAddMember && canEdit && ( - )} - +
); } diff --git a/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx b/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx index c3540ce757a..9b89d85507b 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx @@ -1,6 +1,7 @@ import React, { useState, useEffect } from "react"; import { PlusCircleIcon, PencilIcon, TrashIcon } from "@heroicons/react/outline"; -import { Card, Title, Text, Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell } from "@tremor/react"; +import { Card, CardTitle } from "@/components/ui/card"; +import { Table, TableHeader, TableHead, TableBody, TableRow, TableCell } from "@/components/ui/table"; import ModelSelector from "./ModelSelector"; import NotificationsManager from "../molecules/notifications_manager"; @@ -141,7 +142,7 @@ const ModelAliasManager: React.FC = ({ return (
- Add New Alias +

Add New Alias

@@ -186,17 +187,17 @@ const ModelAliasManager: React.FC = ({
- Manage Existing Aliases +

Manage Existing Aliases

- + - Alias Name - Target Model - Actions + Alias Name + Target Model + Actions - + {aliases.map((alias) => ( @@ -284,9 +285,9 @@ const ModelAliasManager: React.FC = ({ {/* Configuration Example */} {showExampleConfig && ( - - Configuration Example - Here's how your current aliases would look in the config: + + Configuration Example +

Here's how your current aliases would look in the config:

model_aliases: diff --git a/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx b/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx index 857e4d3296b..bd7b56f1817 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx @@ -1,28 +1,20 @@ import { act, fireEvent, render, screen } from "@testing-library/react"; -import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import userEvent from "@testing-library/user-event"; +import { afterEach, describe, expect, it, vi } from "vitest"; import ModelSelector from "./ModelSelector"; vi.mock("@/components/llm_calls/fetch_models", () => ({ 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")); - }); +const openCustomModelInput = async () => { + const user = userEvent.setup(); + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("Enter custom model")); return screen.getByPlaceholderText("Enter custom model name"); }; describe("ModelSelector custom model debounce", () => { - beforeEach(() => { - vi.useFakeTimers(); - }); - afterEach(() => { act(() => { vi.runOnlyPendingTimers(); @@ -30,11 +22,12 @@ describe("ModelSelector custom model debounce", () => { vi.useRealTimers(); }); - it("does not call onChange before the debounce wait elapses", () => { + it("does not call onChange before the debounce wait elapses", async () => { const onChange = vi.fn(); render(); - const input = openCustomModelInput(); + const input = await openCustomModelInput(); + vi.useFakeTimers(); act(() => { fireEvent.change(input, { target: { value: "gpt-4o" } }); @@ -49,11 +42,12 @@ describe("ModelSelector custom model debounce", () => { expect(onChange).not.toHaveBeenCalled(); }); - it("calls onChange exactly once with the last typed value after the wait", () => { + it("calls onChange exactly once with the last typed value after the wait", async () => { const onChange = vi.fn(); render(); - const input = openCustomModelInput(); + const input = await openCustomModelInput(); + vi.useFakeTimers(); act(() => { fireEvent.change(input, { target: { value: "g" } }); @@ -71,11 +65,12 @@ describe("ModelSelector custom model debounce", () => { expect(onChange).toHaveBeenCalledWith("gpt-5.2"); }); - it("does not call onChange when unmounted mid-wait", () => { + it("does not call onChange when unmounted mid-wait", async () => { const onChange = vi.fn(); const { unmount } = render(); - const input = openCustomModelInput(); + const input = await openCustomModelInput(); + vi.useFakeTimers(); act(() => { fireEvent.change(input, { target: { value: "gpt-4o" } }); diff --git a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx index f2621cd1acb..a50131256fa 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx @@ -1,8 +1,8 @@ import React, { useState, useEffect } from "react"; -import { TextInput, Text } from "@tremor/react"; -import { Select } from "antd"; -import { RobotOutlined } from "@ant-design/icons"; +import { Bot } from "lucide-react"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; +import { Input } from "@/components/ui/input"; +import { SearchSelect } from "@/components/shared/SearchSelect"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; const MODEL_SELECT_DEBOUNCE_MS = 500; @@ -80,32 +80,30 @@ const ModelSelector: React.FC = ({ return (
{showLabel && ( - - {labelText} - +

+ {labelText} +

)} - { - if (!option) return false; - const org = organizations?.find((o) => o.organization_id === option.key); - if (!org) return false; - - const searchTerm = input.toLowerCase().trim(); - const orgAlias = (org.organization_alias || "").toLowerCase(); - const orgId = (org.organization_id || "").toLowerCase(); - - return orgAlias.includes(searchTerm) || orgId.includes(searchTerm); - }} - > - {organizations?.map((org) => ( - - {org.organization_alias}{" "} - ({org.organization_id}) - - ))} - +
+ ({ + label: org.organization_alias || org.organization_id, + value: org.organization_id, + sublabel: org.organization_id, + }))} + value={value} + onValueChange={(organizationId) => onChange?.(organizationId)} + placeholder={placeholder} + emptyText={loading ? "Loading organizations…" : "No organizations found"} + disabled={disabled} + inputId={id} + /> +
); }; diff --git a/ui/litellm-dashboard/src/components/common_components/PassThroughRoutesSelector.tsx b/ui/litellm-dashboard/src/components/common_components/PassThroughRoutesSelector.tsx index e02125dea56..330f852383d 100644 --- a/ui/litellm-dashboard/src/components/common_components/PassThroughRoutesSelector.tsx +++ b/ui/litellm-dashboard/src/components/common_components/PassThroughRoutesSelector.tsx @@ -1,5 +1,5 @@ import React, { useEffect, useState } from "react"; -import { Select } from "antd"; +import { MultiSelect, type MultiSelectOption } from "@/components/shared/MultiSelect"; import { getPassThroughEndpointsCall } from "../networking"; interface PassThroughRoutesSelectorProps { @@ -17,6 +17,11 @@ interface PassThroughEndpoint { methods?: string[]; } +const routeOption = (endpoint: PassThroughEndpoint): MultiSelectOption => ({ + label: endpoint.methods?.length ? `${endpoint.methods.join(", ")} ${endpoint.path}` : endpoint.path, + value: endpoint.path, +}); + const PassThroughRoutesSelector: React.FC = ({ onChange, value, @@ -26,7 +31,7 @@ const PassThroughRoutesSelector: React.FC = ({ disabled = false, teamId, }) => { - const [passThroughRoutes, setPassThroughRoutes] = useState>([]); + const [passThroughRoutes, setPassThroughRoutes] = useState([]); const [loading, setLoading] = useState(false); useEffect(() => { @@ -37,27 +42,7 @@ const PassThroughRoutesSelector: React.FC = ({ try { const response = await getPassThroughEndpointsCall(accessToken, teamId); if (response.endpoints) { - const routes = response.endpoints.flatMap((endpoint: PassThroughEndpoint) => { - const path = endpoint.path; - const methods = endpoint.methods; - - // If methods are specified, create one entry per method - if (methods && methods.length > 0) { - return methods.map((method) => ({ - label: `${method} ${path}`, - value: path, // Keep value as path for backward compatibility - })); - } - - // If no methods specified, show just the path (all methods supported) - return [ - { - label: path, - value: path, - }, - ]; - }); - setPassThroughRoutes(routes); + setPassThroughRoutes(response.endpoints.map(routeOption)); } } catch (error) { console.error("Error fetching pass through routes:", error); @@ -70,19 +55,16 @@ const PassThroughRoutesSelector: React.FC = ({ }, [accessToken, teamId]); return ( - ({ + label: project.project_alias || project.project_id, + value: project.project_id, + sublabel: project.project_id, + })) + } value={value} - onChange={onChange} + onValueChange={(projectId) => onChange?.(projectId)} + placeholder="Search or select a project" + emptyText={loading ? "Loading projects…" : "No projects found"} disabled={disabled} - loading={loading} - allowClear - notFoundContent={loading ? } size="small" /> : undefined} - filterOption={(input, option) => { - if (!option) return false; - const project = filtered?.find((p) => p.project_id === option.key); - if (!project) return false; - - const searchTerm = input.toLowerCase().trim(); - const alias = (project.project_alias || "").toLowerCase(); - const id = (project.project_id || "").toLowerCase(); - - return alias.includes(searchTerm) || id.includes(searchTerm); - }} - optionFilterProp="children" - > - {!loading && - filtered?.map((project) => ( - - {project.project_alias || project.project_id}{" "} - ({project.project_id}) - - ))} - + inputId={id} + /> ); }; diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx index 5ac3b8b2b64..6f819607e3f 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx @@ -21,14 +21,6 @@ vi.mock("../Settings/RouterSettings/Fallbacks/FallbackSelectionForm", () => ({ ), })); -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, diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx index 56227abe9ea..7570806182d 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx @@ -1,5 +1,5 @@ import React, { useEffect, useState, useImperativeHandle, forwardRef, useRef } from "react"; -import { TabPanel, TabPanels, TabGroup, TabList, Tab } from "@tremor/react"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { useQuery } from "@tanstack/react-query"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { getRouterSettingsCall } from "../networking"; @@ -344,13 +344,13 @@ const RouterSettingsAccordion = forwardRef - - - Loadbalancing - Fallbacks - - - + + + Loadbalancing + Fallbacks + +
+ - - + + - - - + +
+
); }, diff --git a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx index 4db36f27553..e283f0550ec 100644 --- a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx @@ -1,10 +1,16 @@ import React from "react"; -import { Select } from "antd"; - -const { Option } = Select; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; export const NEVER_RESETS_BUDGET_DURATION = "none"; +const DURATION_LABELS: Record = { + [NEVER_RESETS_BUDGET_DURATION]: "Never resets", + "1h": "hourly", + "24h": "daily", + "7d": "weekly", + "30d": "monthly", +}; + interface BudgetDurationDropdownProps { value?: string | null; onChange?: (value: string | undefined) => void; @@ -24,18 +30,21 @@ const BudgetDurationDropdown: React.FC = ({ }) => { return ( ); }; diff --git a/ui/litellm-dashboard/src/components/common_components/chartUtils.test.tsx b/ui/litellm-dashboard/src/components/common_components/chartUtils.test.tsx index b924021863a..01a4d8373c4 100644 --- a/ui/litellm-dashboard/src/components/common_components/chartUtils.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/chartUtils.test.tsx @@ -1,9 +1,11 @@ import { render, screen } from "@testing-library/react"; import { describe, expect, it } from "vitest"; import { CustomLegend, CustomTooltip } from "./chartUtils"; -import type { CustomTooltipProps } from "@tremor/react"; +import type { ChartTooltipProps } from "@/components/shared/charts/chart_tooltip"; import { SpendMetrics } from "../UsagePage/types"; +type TooltipPayload = NonNullable; + describe("CustomTooltip", () => { const mockPayload = [ { @@ -28,9 +30,9 @@ describe("CustomTooltip", () => { ]; it("should render", () => { - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: mockPayload, + payload: mockPayload as unknown as TooltipPayload, label: "2024-01-15", }; render(); @@ -38,9 +40,9 @@ describe("CustomTooltip", () => { }); it("should return null when not active", () => { - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: false, - payload: mockPayload, + payload: mockPayload as unknown as TooltipPayload, label: "2024-01-15", }; const { container } = render(); @@ -48,9 +50,9 @@ describe("CustomTooltip", () => { }); it("should return null when payload is empty", () => { - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: [], + payload: [] as unknown as TooltipPayload, label: "2024-01-15", }; const { container } = render(); @@ -58,9 +60,9 @@ describe("CustomTooltip", () => { }); it("should display formatted category names", () => { - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: mockPayload, + payload: mockPayload as unknown as TooltipPayload, label: "2024-01-15", }; render(); @@ -89,9 +91,9 @@ describe("CustomTooltip", () => { }, }, ]; - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: payloadWithUnderscores, + payload: payloadWithUnderscores as unknown as TooltipPayload, label: "2024-01-15", }; render(); @@ -120,9 +122,9 @@ describe("CustomTooltip", () => { }, }, ]; - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: spendPayload, + payload: spendPayload as unknown as TooltipPayload, label: "2024-01-15", }; render(); @@ -130,9 +132,9 @@ describe("CustomTooltip", () => { }); it("should format non-spend numeric values with locale string", () => { - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: mockPayload, + payload: mockPayload as unknown as TooltipPayload, label: "2024-01-15", }; render(); @@ -161,9 +163,9 @@ describe("CustomTooltip", () => { }, }, ]; - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: payloadWithUndefined, + payload: payloadWithUndefined as unknown as TooltipPayload, label: "2024-01-15", }; render(); @@ -211,9 +213,9 @@ describe("CustomTooltip", () => { }, }, ]; - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: multiplePayload, + payload: multiplePayload as unknown as TooltipPayload, label: "2024-01-15", }; render(); @@ -222,9 +224,9 @@ describe("CustomTooltip", () => { }); it("should convert color names to hex values", () => { - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: mockPayload, + payload: mockPayload as unknown as TooltipPayload, label: "2024-01-15", }; const { container } = render(); @@ -254,9 +256,9 @@ describe("CustomTooltip", () => { }, }, ]; - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: payloadWithHexColor, + payload: payloadWithHexColor as unknown as TooltipPayload, label: "2024-01-15", }; const { container } = render(); @@ -286,9 +288,9 @@ describe("CustomTooltip", () => { }, }, ]; - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: payloadWithoutDataKey as any, + payload: payloadWithoutDataKey as unknown as TooltipPayload, label: "2024-01-15", }; render(); @@ -304,9 +306,9 @@ describe("CustomTooltip", () => { payload: undefined, }, ]; - const props: CustomTooltipProps = { + const props: ChartTooltipProps = { active: true, - payload: payloadWithoutPayload as any, + payload: payloadWithoutPayload as unknown as TooltipPayload, label: "2024-01-15", }; render(); diff --git a/ui/litellm-dashboard/src/components/common_components/chartUtils.tsx b/ui/litellm-dashboard/src/components/common_components/chartUtils.tsx index c0930f290f6..0abf004803b 100644 --- a/ui/litellm-dashboard/src/components/common_components/chartUtils.tsx +++ b/ui/litellm-dashboard/src/components/common_components/chartUtils.tsx @@ -1,4 +1,4 @@ -import type { CustomTooltipProps } from "@tremor/react"; +import type { ChartTooltipProps } from "@/components/shared/charts/chart_tooltip"; import { SpendMetrics } from "../UsagePage/types"; interface ChartDataPoint { @@ -16,7 +16,7 @@ const colorNameToHex: { [key: string]: string } = { emerald: "#37bc7d", }; -export const CustomTooltip = ({ active, payload, label }: CustomTooltipProps) => { +export const CustomTooltip = ({ active, payload, label }: ChartTooltipProps) => { if (active && payload && payload.length) { const formatCategoryName = (name: string): string => { return name diff --git a/ui/litellm-dashboard/src/components/common_components/routerSettingsWiring.test.tsx b/ui/litellm-dashboard/src/components/common_components/routerSettingsWiring.test.tsx index d6ed0ae07db..15c52e71e8f 100644 --- a/ui/litellm-dashboard/src/components/common_components/routerSettingsWiring.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/routerSettingsWiring.test.tsx @@ -1,6 +1,6 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { render, screen, waitFor } from "@testing-library/react"; -import type { ReactElement, ReactNode } from "react"; +import type { ReactElement } from "react"; import { describe, expect, it, vi } from "vitest"; import type { FallbackGroup } from "../Settings/RouterSettings/Fallbacks/FallbackGroupConfig"; import type { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; @@ -16,14 +16,6 @@ vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModelsForTeam: vi.fn().mockResolvedValue([]), })); -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 }: { value: RouterSettingsFormValue }) => (
{JSON.stringify(value.routerSettings)}
diff --git a/ui/litellm-dashboard/src/components/common_components/simple_table.tsx b/ui/litellm-dashboard/src/components/common_components/simple_table.tsx index 4a858a346d4..17e8d46d21d 100644 --- a/ui/litellm-dashboard/src/components/common_components/simple_table.tsx +++ b/ui/litellm-dashboard/src/components/common_components/simple_table.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { Table, TableHead, TableRow, TableHeaderCell, TableBody, TableCell, Text } from "@tremor/react"; +import { Table, TableHeader, TableRow, TableHead, TableBody, TableCell } from "@/components/ui/table"; export interface SimpleTableColumn { header: string; @@ -31,20 +31,20 @@ export function SimpleTable({ }: SimpleTableProps) { return (
- + {columns.map((column, index) => ( - + {column.header} - + ))} - + {isLoading ? ( - {loadingMessage} + {loadingMessage} ) : data.length > 0 ? ( @@ -60,7 +60,7 @@ export function SimpleTable({ ) : ( - {emptyMessage} + {emptyMessage} )} 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 7d27886c7f5..35121f41598 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx @@ -1,13 +1,8 @@ -import React, { useMemo, useState, type UIEvent } from "react"; -import { Select, Typography } from "antd"; -import { LoadingOutlined } from "@ant-design/icons"; -import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; +import React, { useMemo, useState } from "react"; +import { PaginatedSearchSelect } from "@/components/shared/PaginatedSearchSelect"; 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; - interface TeamDropdownProps { value?: string; onChange?: (value: string) => void; @@ -17,10 +12,9 @@ interface TeamDropdownProps { /** Filter teams by organization. */ organizationId?: string | null; pageSize?: number; + id?: string; } -const SCROLL_THRESHOLD = 0.8; - const TeamDropdown: React.FC = ({ value, onChange, @@ -28,15 +22,13 @@ const TeamDropdown: React.FC = ({ disabled, organizationId, pageSize = 20, + id, }) => { - const [searchInput, setSearchInput] = useState(""); - const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { - wait: DEBOUNCE_WAIT_MS, - }); + const [search, setSearch] = useState(""); const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteTeams( pageSize, - debouncedSearch || undefined, + search || undefined, organizationId, ); @@ -54,59 +46,35 @@ const TeamDropdown: React.FC = ({ return result; }, [data]); - const handlePopupScroll = (e: UIEvent) => { - const target = e.currentTarget; - const scrollRatio = (target.scrollTop + target.clientHeight) / target.scrollHeight; - if (scrollRatio >= SCROLL_THRESHOLD && hasNextPage && !isFetchingNextPage) { - fetchNextPage(); - } - }; - - const handleSearch = (val: string) => { - setSearchInput(val); - setDebouncedSearch(val); - }; - - const handleChange = (teamId: string | undefined) => { - onChange?.(teamId ?? ""); + const handleChange = (teamId: string) => { + onChange?.(teamId); if (onTeamSelect) { - const team = teamId ? teams.find((t) => t.team_id === teamId) ?? null : null; - onTeamSelect(team); + onTeamSelect(teamId ? teams.find((t) => t.team_id === teamId) ?? null : null); } }; return ( - +
+ ({ + label: team.team_alias || team.team_id, + value: team.team_id, + sublabel: team.team_id, + }))} + value={value || undefined} + onValueChange={handleChange} + onSearchChange={setSearch} + onLoadMore={fetchNextPage} + hasNextPage={hasNextPage} + isLoading={isLoading} + isFetchingNextPage={isFetchingNextPage} + placeholder="Search or select a team" + emptyText="No teams found" + loadingText="Loading teams…" + disabled={disabled} + inputId={id} + /> +
); }; diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/TagRateLimitEditor.test.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/TagRateLimitEditor.test.tsx new file mode 100644 index 00000000000..85f02f2045e --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_team_helpers/TagRateLimitEditor.test.tsx @@ -0,0 +1,129 @@ +import React, { useState } from "react"; +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, it, expect } from "vitest"; +import { TagRateLimitEditor, TagRateLimitEntry, tagLimitsToRows, tagRowsToLimits } from "./TagRateLimitEditor"; + +function Harness({ initial = [] as TagRateLimitEntry[], onValue }: { initial?: TagRateLimitEntry[]; onValue?: any }) { + const [rows, setRows] = useState(initial); + return ( + { + setRows(next); + onValue?.(next); + }} + /> + ); +} + +const rowsWith = (tag: string, rpm: number | null): TagRateLimitEntry[] => [{ id: "r1", tag, rpm_limit: rpm }]; + +describe("TagRateLimitEditor", () => { + it("should render one tag and one RPM field per row", () => { + render(); + expect(screen.getByRole("textbox", { name: "Tag" })).toHaveValue("cell-1"); + expect(screen.getByRole("spinbutton", { name: "RPM limit" })).toHaveValue(100); + }); + + it("should add a row when Add Tag Limit is clicked", async () => { + const user = userEvent.setup(); + render(); + expect(screen.queryAllByRole("textbox", { name: "Tag" })).toHaveLength(0); + + await user.click(screen.getByRole("button", { name: /add tag limit/i })); + + expect(screen.getAllByRole("textbox", { name: "Tag" })).toHaveLength(1); + }); + + it("should let the user type a tag name", async () => { + const user = userEvent.setup(); + render(); + + await user.type(screen.getByRole("textbox", { name: "Tag" }), "cell-2"); + + expect(screen.getByRole("textbox", { name: "Tag" })).toHaveValue("cell-2"); + }); + + it("should record the typed RPM limit as a number, not a string", async () => { + const user = userEvent.setup(); + const seen: TagRateLimitEntry[][] = []; + render( seen.push(v)} />); + + await user.type(screen.getByRole("spinbutton", { name: "RPM limit" }), "60"); + + const latest = seen[seen.length - 1][0]; + expect(latest.rpm_limit).toBe(60); + expect(typeof latest.rpm_limit).toBe("number"); + }); + + it("should reset the RPM limit to null when the field is cleared", async () => { + const user = userEvent.setup(); + const seen: TagRateLimitEntry[][] = []; + render( seen.push(v)} />); + + await user.clear(screen.getByRole("spinbutton", { name: "RPM limit" })); + + expect(seen[seen.length - 1][0].rpm_limit).toBeNull(); + }); + + it("should remove only the clicked row", async () => { + const user = userEvent.setup(); + const initial: TagRateLimitEntry[] = [ + { id: "r1", tag: "keep-me", rpm_limit: 10 }, + { id: "r2", tag: "delete-me", rpm_limit: 20 }, + ]; + render(); + + await user.click(screen.getAllByRole("button", { name: "Remove tag limit" })[1]); + + const tags = screen.getAllByRole("textbox", { name: "Tag" }); + expect(tags).toHaveLength(1); + expect(tags[0]).toHaveValue("keep-me"); + }); + + it("should not submit the surrounding form when a row is removed", async () => { + const user = userEvent.setup(); + let submitted = false; + render( +
{ + submitted = true; + }} + > + + , + ); + + await user.click(screen.getByRole("button", { name: "Remove tag limit" })); + + expect(submitted).toBe(false); + expect(screen.queryAllByRole("textbox", { name: "Tag" })).toHaveLength(0); + }); +}); + +describe("tagRowsToLimits", () => { + it("should map named rows with numeric limits into the rpm map", () => { + expect(tagRowsToLimits([{ id: "a", tag: "cell-1", rpm_limit: 60 }])).toEqual({ tag_rpm_limit: { "cell-1": 60 } }); + }); + + it("should drop rows with a blank tag or a null limit", () => { + const rows: TagRateLimitEntry[] = [ + { id: "a", tag: " ", rpm_limit: 60 }, + { id: "b", tag: "cell-2", rpm_limit: null }, + ]; + expect(tagRowsToLimits(rows)).toEqual({ tag_rpm_limit: {} }); + }); +}); + +describe("tagLimitsToRows", () => { + it("should rebuild rows from a stored rpm map", () => { + const rows = tagLimitsToRows({ "cell-1": 60 }); + expect(rows).toHaveLength(1); + expect(rows[0]).toMatchObject({ tag: "cell-1", rpm_limit: 60 }); + }); + + it("should ignore non-numeric entries", () => { + expect(tagLimitsToRows({ "cell-1": "sixty" })).toEqual([]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/TagRateLimitEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/TagRateLimitEditor.tsx index ee022ee9a75..151e1593765 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/TagRateLimitEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/TagRateLimitEditor.tsx @@ -1,5 +1,6 @@ -import { Button, Input, InputNumber } from "antd"; import React from "react"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; export interface TagRateLimitEntry { // Stable identity for React list keys so deleting a middle row doesn't shift @@ -72,25 +73,29 @@ export function TagRateLimitEditor({ value, onChange }: TagRateLimitEditorProps) {value.map((row, idx) => (
updateRow(idx, "tag", e.target.value)} placeholder="Tag (e.g. cell-1)" style={{ width: 180 }} /> - updateRow(idx, "rpm_limit", v ?? null)} + value={row.rpm_limit ?? ""} + onChange={(e) => updateRow(idx, "rpm_limit", e.target.value === "" ? null : Number(e.target.value))} placeholder="RPM" style={{ width: 120 }} /> -
))}
- + - Alias Name - Target Model Group - Actions + Alias Name + Target Model Group + Actions - + {aliases.map((alias) => ( @@ -275,8 +276,12 @@ const ModelGroupAliasSettings: React.FC = ({ ) : ( <> - {alias.aliasName} - {alias.targetModelGroup} + + {alias.aliasName} + + + {alias.targetModelGroup} +
{/* Configuration Example */} - - Configuration Example - - Here's how your current aliases would look in the config.yaml: - + + Configuration Example +

Here's how your current aliases would look in the config.yaml:

router_settings: diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index a3ce40494eb..5eb17fee2c4 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -617,9 +617,14 @@ describe("ModelInfoView", () => { describe("PTU cost attribution gate", () => { const ptuModelData = { ...defaultModelData, + // Zero per-token pricing is what the backend stores for a PTU deployment, since the flat + // cost of its reserved capacity already covers the traffic that capacity serves. + litellm_params: { ...defaultModelData.litellm_params, input_cost_per_token: 0, output_cost_per_token: 0 }, model_info: { ...defaultModelData.model_info, team_id: "team-1", + input_cost_per_token: 0, + output_cost_per_token: 0, ptu_count: 15, cost_per_ptu_per_hour: 2, ptu_effective_from: "2026-07-01T00:00:00+00:00", @@ -683,6 +688,70 @@ describe("ModelInfoView", () => { expect(modelInfo).not.toHaveProperty("ptu_effective_to"); }); + it("shows a zeroed PTU price as 0.0000 rather than Not Set", async () => { + mockUsePtuCostAttributionEnabled.mockReturnValue(true); + renderWithPtuModel(); + + await waitFor(() => { + expect(screen.getByText("Input Cost (per 1M tokens)")).toBeInTheDocument(); + }); + for (const label of ["Input Cost (per 1M tokens)", "Output Cost (per 1M tokens)"]) { + expect(screen.getByText(label).parentElement).toHaveTextContent("0.0000"); + } + }); + + it("blocks the save once the operator types a non-zero per-token cost alongside PTU config", async () => { + mockUsePtuCostAttributionEnabled.mockReturnValue(true); + const user = userEvent.setup(); + renderWithPtuModel(); + + 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.getByPlaceholderText("Enter input cost")).toBeInTheDocument(); + }); + await user.clear(screen.getByPlaceholderText("Enter input cost")); + await user.type(screen.getByPlaceholderText("Enter input cost"), "2.5"); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(screen.getByText(/bills by reserved capacity/i)).toBeInTheDocument(); + }); + expect(mockModelPatchUpdateCall).not.toHaveBeenCalled(); + }); + + it("lets the operator put a cost-map-priced deployment on PTU without clearing the seeded rate", async () => { + // A rate the form seeded from /model/info is the server's own, so refusing it blocked + // every attempt to enable PTU from the dashboard. + mockUsePtuCostAttributionEnabled.mockReturnValue(true); + const seededModel = { + ...defaultModelData, + model_info: { ...defaultModelData.model_info, team_id: "team-1", input_cost_per_token: 0.0000003 }, + }; + mockUseModelsInfo.mockReturnValue({ data: { data: [seededModel] }, isLoading: false, error: null }); + mockModelInfoV1Call.mockResolvedValue({ data: [seededModel] }); + const user = userEvent.setup(); + render(, { wrapper }); + + 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.getByPlaceholderText("e.g. 15")).toBeInTheDocument(); + }); + await user.type(screen.getByPlaceholderText("e.g. 15"), "15"); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(screen.queryByText(/bills by reserved capacity/i)).not.toBeInTheDocument(); + }); + }); + it("sends the PTU fields on save when enabled", async () => { mockUsePtuCostAttributionEnabled.mockReturnValue(true); const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index e1a4d4311ab..a0c545e17aa 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -25,6 +25,7 @@ import { PTU_COUNT_FIELD, PTU_RATE_FIELD, ptuCountRules, + ptuNoUsageCostRule, ptuPairRule, ptuRateRules, ptuStartRequiredRule, @@ -128,6 +129,12 @@ const PTU_EDIT_FIELDS: PtuEditField[] = [ }, ]; +/** Per-1M-token rate for the first rate that is set, so a deliberate 0 seeds the form as 0. */ +const perMillionTokens = (...rates: (number | null | undefined)[]): number | null => { + const rate = rates.find((candidate) => candidate != null); + return rate == null ? null : rate * 1_000_000; +}; + const ptuFieldDependencies = ({ isStart, pairedWith, windowPeer }: PtuEditField): string[] | undefined => { const deps = [ ...(isStart ? [PTU_COUNT_FIELD] : []), @@ -227,6 +234,8 @@ export default function ModelInfoView({ const { data: modelHubData } = useModelHub(); const { data: teams } = useTeams(); const ptuCostAttributionEnabled = usePtuCostAttributionEnabled(); + const ptuCostRule = (field: string) => + ptuCostAttributionEnabled ? [ptuNoUsageCostRule(PTU_COUNT_FIELD, field)] : []; // Transform the model data const getProviderFromModel = (model: string) => { @@ -835,12 +844,14 @@ export default function ModelInfoView({ max_retries: localModelData.litellm_params.max_retries, timeout: localModelData.litellm_params.timeout, stream_timeout: localModelData.litellm_params.stream_timeout, - input_cost: localModelData.litellm_params.input_cost_per_token - ? localModelData.litellm_params.input_cost_per_token * 1_000_000 - : localModelData.model_info?.input_cost_per_token * 1_000_000 || null, - output_cost: localModelData.litellm_params?.output_cost_per_token - ? localModelData.litellm_params.output_cost_per_token * 1_000_000 - : localModelData.model_info?.output_cost_per_token * 1_000_000 || null, + input_cost: perMillionTokens( + localModelData.litellm_params.input_cost_per_token, + localModelData.model_info?.input_cost_per_token, + ), + output_cost: perMillionTokens( + localModelData.litellm_params?.output_cost_per_token, + localModelData.model_info?.output_cost_per_token, + ), ptu_count: localModelData.model_info?.ptu_count ?? null, cost_per_ptu_per_hour: localModelData.model_info?.cost_per_ptu_per_hour ?? null, ptu_effective_from: utcIsoToPickerValue(localModelData.model_info?.ptu_effective_from), @@ -917,14 +928,19 @@ export default function ModelInfoView({
Input Cost (per 1M tokens) {isEditing ? ( - + ) : (
- {localModelData?.litellm_params?.input_cost_per_token - ? (localModelData.litellm_params?.input_cost_per_token * 1_000_000).toFixed(4) - : localModelData?.model_info?.input_cost_per_token + {localModelData?.litellm_params?.input_cost_per_token != null + ? (localModelData.litellm_params.input_cost_per_token * 1_000_000).toFixed(4) + : localModelData?.model_info?.input_cost_per_token != null ? (localModelData.model_info.input_cost_per_token * 1_000_000).toFixed(4) : "Not Set"}
@@ -934,14 +950,19 @@ export default function ModelInfoView({
Output Cost (per 1M tokens) {isEditing ? ( - + ) : (
- {localModelData?.litellm_params?.output_cost_per_token + {localModelData?.litellm_params?.output_cost_per_token != null ? (localModelData.litellm_params.output_cost_per_token * 1_000_000).toFixed(4) - : localModelData?.model_info?.output_cost_per_token + : localModelData?.model_info?.output_cost_per_token != null ? (localModelData.model_info.output_cost_per_token * 1_000_000).toFixed(4) : "Not Set"}
@@ -995,6 +1016,8 @@ export default function ModelInfoView({ @@ -1018,6 +1041,8 @@ export default function ModelInfoView({ diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index c999ee8035e..6ce92fb9449 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -8,8 +8,8 @@ import { useTheme } from "@/contexts/ThemeContext"; import { clearTokenCookies } from "@/utils/cookieUtils"; import { clearStoredReturnUrl, getLoginUrl } from "@/utils/returnUrlUtils"; import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; -import { DownOutlined, MenuFoldOutlined, MenuUnfoldOutlined } from "@ant-design/icons"; -import { Tag } from "antd"; +import { Badge } from "@/components/ui/badge"; +import { ChevronDown, PanelLeftClose, PanelLeftOpen } from "lucide-react"; import Link from "next/link"; import React from "react"; import { BlogDropdown } from "./Navbar/BlogDropdown/BlogDropdown"; @@ -71,7 +71,13 @@ const Navbar: React.FC = ({ className="mr-2 flex h-9 w-9 items-center justify-center rounded-md text-gray-600 transition-colors hover:bg-gray-100 hover:text-gray-900" title={sidebarCollapsed ? "Expand sidebar" : "Collapse sidebar"} > - {sidebarCollapsed ? : } + + {sidebarCollapsed ? ( + + ) : ( + + )} + )} @@ -98,7 +104,7 @@ const Navbar: React.FC = ({ 🌑 )} - + = ({ > v{version} - +
)}
@@ -138,7 +144,7 @@ const Navbar: React.FC = ({ > Docs {/* Layout parity with Blog chevron — intentional single-level link */} - + diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx index 327e127e8fe..c7baa3d52c2 100644 --- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx +++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx @@ -1,5 +1,4 @@ import React from "react"; -import { Text } from "@tremor/react"; import VectorStorePermissions from "./permissions/VectorStorePermissions"; import MCPServerPermissions from "./permissions/MCPServerPermissions"; import AgentPermissions from "./permissions/AgentPermissions"; @@ -38,14 +37,14 @@ export function ObjectPermissionsView({ accessToken={accessToken} /> -
- Search tools +
+

Search tools

{searchTools.length === 0 ? ( - +

No restriction — all configured search tools are allowed for this team. - +

) : ( - {searchTools.join(", ")} +

{searchTools.join(", ")}

)}
@@ -56,8 +55,8 @@ export function ObjectPermissionsView({
- Object Permissions - Access control for Vector Stores and MCP Servers +

Object Permissions

+

Access control for Vector Stores and MCP Servers

{content} @@ -67,7 +66,7 @@ export function ObjectPermissionsView({ return (
- Object Permissions +

Object Permissions

{content}
); diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.tsx index a9f79810e1d..8c1607088c9 100644 --- a/ui/litellm-dashboard/src/components/organization/organization_view.tsx +++ b/ui/litellm-dashboard/src/components/organization/organization_view.tsx @@ -11,10 +11,9 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import { teamDetailHref } from "@/utils/entityLinks"; import { createTeamAliasMap } from "@/utils/teamUtils"; import { BadgeLink } from "@/components/shared/BadgeLink"; -import type { ColumnsType } from "antd/es/table"; import { ArrowLeft } from "lucide-react"; import React, { useMemo, useState } from "react"; -import MemberTable from "../common_components/MemberTable"; +import MemberTable, { type MemberTableColumn } from "../common_components/MemberTable"; import UserSearchModal from "../common_components/user_search_modal"; import NotificationsManager from "../molecules/notifications_manager"; import { @@ -122,7 +121,7 @@ const OrganizationInfoView: React.FC = ({ return
Organization not found
; } - const orgExtraColumns: ColumnsType = [ + const orgExtraColumns: MemberTableColumn[] = [ { title: "Spend (USD)", key: "spend", diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx index 11951b2decb..ee6fbbd89f5 100644 --- a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx @@ -1,7 +1,7 @@ import React, { useState, useEffect } from "react"; -import { Text, Badge } from "@tremor/react"; import { UserGroupIcon } from "@heroicons/react/outline"; -import { Tooltip } from "antd"; +import { Badge } from "@/components/ui/badge"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { getAgentsList } from "../networking"; interface Agent { @@ -58,10 +58,8 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }
- Agents - - {totalCount} - +

Agents

+ {totalCount}
{totalCount > 0 ? ( @@ -71,14 +69,17 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }
{item.type === "agent" ? ( - -
- - - {getAgentDisplayName(item.value)} - -
- + + + }> + + + {getAgentDisplayName(item.value)} + + + {`Full ID: ${item.value}`} + + ) : (
@@ -96,7 +97,7 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken } ) : (
- No agents or access groups configured +

No agents or access groups configured

)}
diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx index b16b094e0e1..02980cd4c3e 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx @@ -1,7 +1,7 @@ import React, { useState, useEffect } from "react"; -import { Text, Badge } from "@tremor/react"; import { ServerIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; -import { Tooltip } from "antd"; +import { Badge } from "@/components/ui/badge"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { fetchMCPServers, fetchMCPToolsets } from "../networking"; import { MCPServer, MCPToolset } from "../mcp_tools/types"; import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; @@ -111,8 +111,8 @@ export function MCPServerPermissions({
- MCP Servers - +

MCP Servers

+ {blocksAllMcpServers ? "Blocked" : grantsAllProxyMcpServers ? "All" : totalCount}
@@ -120,14 +120,14 @@ export function MCPServerPermissions({ {blocksAllMcpServers ? (
- +

No MCP servers — this key is blocked from all MCP servers, including its team's servers - +

) : grantsAllProxyMcpServers ? (
- All Proxy MCP Servers +

All Proxy MCP Servers

) : totalCount > 0 ? (
@@ -146,13 +146,14 @@ export function MCPServerPermissions({ >
{item.type === "server" ? ( - -
+ + }> {getMCPServerDisplayName(item.value)} -
+ + {`Full ID: ${item.value}`} ) : (
@@ -256,7 +257,7 @@ export function MCPServerPermissions({ ) : (
- No MCP servers, access groups, or toolsets configured +

No MCP servers, access groups, or toolsets configured

)}
diff --git a/ui/litellm-dashboard/src/components/permissions/VectorStorePermissions.tsx b/ui/litellm-dashboard/src/components/permissions/VectorStorePermissions.tsx index 8541d65e11f..6bf79d8a632 100644 --- a/ui/litellm-dashboard/src/components/permissions/VectorStorePermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/VectorStorePermissions.tsx @@ -1,6 +1,6 @@ import React, { useState, useEffect } from "react"; -import { Text, Badge } from "@tremor/react"; import { DatabaseIcon } from "@heroicons/react/outline"; +import { Badge } from "@/components/ui/badge"; import { vectorStoreListCall } from "../networking"; interface VectorStoreDetails { @@ -52,10 +52,8 @@ export function VectorStorePermissions({ vectorStores, accessToken }: VectorStor
- Vector Stores - - {vectorStores.length} - +

Vector Stores

+ {vectorStores.length}
{vectorStores.length > 0 ? ( @@ -63,7 +61,7 @@ export function VectorStorePermissions({ vectorStores, accessToken }: VectorStor {vectorStores.map((store, index) => (
{getVectorStoreDisplayName(store)}
@@ -72,7 +70,7 @@ export function VectorStorePermissions({ vectorStores, accessToken }: VectorStor ) : (
- No vector stores configured +

No vector stores configured

)}
diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 777cdc62987..b163a6b341c 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -94,6 +94,14 @@ describe("provider_info_helpers", () => { expect(result.displayName).toBe(Providers.ZAI); }); + it("should resolve the nvidia_riva provider value to the Nvidia Riva display name and logo", () => { + const result = getProviderLogoAndName("nvidia_riva"); + expect(result.displayName).toBe(Providers.NVIDIA_RIVA); + expect(provider_map.NVIDIA_RIVA).toBe("nvidia_riva"); + expect(result.logo).toBe(providerLogoMap[Providers.NVIDIA_RIVA]); + expect(result.logo).toBeTruthy(); + }); + it("should return provider value as display name when no mapping exists", () => { const unknownProvider = "unknown_provider"; const result = getProviderLogoAndName(unknownProvider); @@ -225,6 +233,15 @@ describe("provider_info_helpers", () => { expect(getPlaceholder(Providers.ZAI)).toBe("zai/glm-4.5"); }); + it("should return the riva asr placeholder for NVIDIA_RIVA provider", () => { + expect(getPlaceholder(Providers.NVIDIA_RIVA)).toBe("nvidia_riva/nvidia/parakeet-ctc-1_1b-asr"); + }); + + it("should resolve enum keys from the provider dropdown, not just enum values", () => { + expect(getPlaceholder("NVIDIA_RIVA")).toBe("nvidia_riva/nvidia/parakeet-ctc-1_1b-asr"); + expect(getPlaceholder("WATSONX")).toBe("watsonx/ibm/granite-3-3-8b-instruct"); + }); + it("should return default gpt-3.5-turbo placeholder for unknown provider", () => { expect(getPlaceholder("UnknownProvider" as any)).toBe("gpt-3.5-turbo"); }); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index fa6b3c79230..ae311070c07 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -130,6 +130,7 @@ export enum Providers { NOVITA = "Novita", NSCALE = "Nscale", NVIDIA_NIM = "Nvidia Nim", + NVIDIA_RIVA = "Nvidia Riva", Ollama = "Ollama", OLLAMA_CHAT = "Ollama Chat", OOBABOOGA = "Oobabooga", @@ -238,6 +239,7 @@ export const provider_map: Record = { NOVITA: "novita", NSCALE: "nscale", NVIDIA_NIM: "nvidia_nim", + NVIDIA_RIVA: "nvidia_riva", Ollama: "ollama", OLLAMA_CHAT: "ollama_chat", OOBABOOGA: "oobabooga", @@ -334,6 +336,7 @@ export const providerLogoMap: Partial> = { [Providers.NEBIUS]: nebiusLogo.src, [Providers.NOVITA]: novitaLogo.src, [Providers.NVIDIA_NIM]: nvidiaNimLogo.src, + [Providers.NVIDIA_RIVA]: nvidiaNimLogo.src, [Providers.Ollama]: ollamaLogo.src, [Providers.OLLAMA_CHAT]: ollamaLogo.src, [Providers.OOBABOOGA]: openaiSmallLogo.src, @@ -400,48 +403,32 @@ export const getProviderLogoAndName = (providerValue: string): { logo: string; d return { logo, displayName }; }; +const providerPlaceholderMap: Partial> = { + [Providers.AIML]: "aiml/flux-pro/v1.1", + [Providers.Anthropic]: "claude-3-opus", + [Providers.Azure]: "my-deployment", + [Providers.Azure_AI_Studio]: "azure_ai/command-r-plus", + [Providers.Bedrock]: "claude-3-opus", + [Providers.Cursor]: "cursor/claude-4-sonnet", + [Providers.DeepInfra]: "deepinfra/", + [Providers.FalAI]: "fal_ai/fal-ai/flux-pro/v1.1-ultra", + [Providers.Google_AI_Studio]: "gemini-pro", + [Providers.JinaAI]: "jina_ai/", + [Providers.NVIDIA_RIVA]: "nvidia_riva/nvidia/parakeet-ctc-1_1b-asr", + [Providers.Oracle]: "oci/xai.grok-4", + [Providers.RunwayML]: "runwayml/gen4_turbo", + [Providers.SageMaker]: "sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b", + [Providers.Snowflake]: "snowflake/mistral-7b", + [Providers.Vertex_AI]: "gemini-pro", + [Providers.VolcEngine]: "volcengine/", + [Providers.Voyage]: "voyage/", + [Providers.WATSONX]: "watsonx/ibm/granite-3-3-8b-instruct", + [Providers.ZAI]: "zai/glm-4.5", +}; + export const getPlaceholder = (selectedProvider: string): string => { - if (selectedProvider === Providers.AIML) { - return "aiml/flux-pro/v1.1"; - } else if (selectedProvider === Providers.Vertex_AI) { - return "gemini-pro"; - } else if (selectedProvider == Providers.Anthropic) { - return "claude-3-opus"; - } else if (selectedProvider == Providers.Bedrock) { - return "claude-3-opus"; - } else if (selectedProvider == Providers.SageMaker) { - return "sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b"; - } else if (selectedProvider == Providers.Google_AI_Studio) { - return "gemini-pro"; - } else if (selectedProvider == Providers.Azure_AI_Studio) { - return "azure_ai/command-r-plus"; - } else if (selectedProvider == Providers.Azure) { - return "my-deployment"; - } else if (selectedProvider == Providers.Oracle) { - return "oci/xai.grok-4"; - } else if (selectedProvider == Providers.Snowflake) { - return "snowflake/mistral-7b"; - } else if (selectedProvider == Providers.Voyage) { - return "voyage/"; - } else if (selectedProvider == Providers.JinaAI) { - return "jina_ai/"; - } else if (selectedProvider == Providers.VolcEngine) { - return "volcengine/"; - } else if (selectedProvider == Providers.DeepInfra) { - return "deepinfra/"; - } else if (selectedProvider == Providers.FalAI) { - return "fal_ai/fal-ai/flux-pro/v1.1-ultra"; - } else if (selectedProvider == Providers.RunwayML) { - return "runwayml/gen4_turbo"; - } else if (selectedProvider === Providers.WATSONX) { - return "watsonx/ibm/granite-3-3-8b-instruct"; - } else if (selectedProvider === Providers.Cursor) { - return "cursor/claude-4-sonnet"; - } else if (selectedProvider === Providers.ZAI) { - return "zai/glm-4.5"; - } else { - return "gpt-3.5-turbo"; - } + const resolvedProvider = Providers[selectedProvider as keyof typeof Providers] ?? (selectedProvider as Providers); + return providerPlaceholderMap[resolvedProvider] ?? "gpt-3.5-turbo"; }; export const getProviderModels = (provider: Providers, modelMap: any): Array => { diff --git a/ui/litellm-dashboard/src/components/public_model_hub.test.tsx b/ui/litellm-dashboard/src/components/public_model_hub.test.tsx index 43788c1dfd8..875f89b5adc 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.test.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.test.tsx @@ -226,4 +226,19 @@ describe("public hub MCP details modal", () => { await screen.findByText("Server Overview"); expect(screen.queryByText(PUBLIC_SERVER_URL)).not.toBeInTheDocument(); }); + + it("closes the server details modal from its close control", async () => { + const networkingModule = await import("./networking"); + vi.mocked(networkingModule.mcpHubPublicServersCall).mockResolvedValue([mockMcpServer]); + + render(); + + fireEvent.click(await screen.findByRole("tab", { name: /MCP Hub/i })); + fireEvent.click(await screen.findByRole("button", { name: "exa_test" })); + await screen.findByText("Server Overview"); + + fireEvent.click(screen.getByRole("button", { name: /close/i })); + + await waitFor(() => expect(screen.queryByText("Server Overview")).not.toBeInTheDocument()); + }); }); diff --git a/ui/litellm-dashboard/src/components/public_model_hub.tsx b/ui/litellm-dashboard/src/components/public_model_hub.tsx index 2bdb3055835..ab1a10918e7 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.tsx @@ -1,10 +1,25 @@ import { ThemeProvider } from "@/contexts/ThemeContext"; import { ExternalLinkIcon, SearchIcon } from "@heroicons/react/outline"; import { SortingState } from "@tanstack/react-table"; -import { Card, Text, Title } from "@tremor/react"; -import { Modal, Select, Tabs, Tag, Tooltip } from "antd"; import { Copy, Inbox, Info } from "lucide-react"; import React, { useCallback, useEffect, useMemo, useState } from "react"; +import { Badge } from "@/components/ui/badge"; +import { Card } from "@/components/ui/card"; +import { + Combobox, + ComboboxChip, + ComboboxChips, + ComboboxChipsInput, + ComboboxContent, + ComboboxEmpty, + ComboboxItem, + ComboboxList, + ComboboxValue, +} from "@/components/ui/combobox"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import { MultiSelect } from "./shared/MultiSelect"; import { DataTable } from "./shared/DataTable"; import NotificationsManager from "./molecules/notifications_manager"; import Navbar from "./navbar"; @@ -32,8 +47,6 @@ import { getEndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; import { MessageType } from "@/components/chat_ui/types"; import { getProviderLogoAndName } from "./provider_info_helpers"; -const { TabPane } = Tabs; - interface PublicModelHubProps { accessToken?: string | null; isEmbedded?: boolean; // When true, hides navbar and adjusts layout for embedding in dashboard @@ -397,11 +410,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded setIsModalVisible(true); }, []); - const handleModalOk = () => { - setIsModalVisible(false); - setSelectedModel(null); - }; - const handleModalCancel = () => { setIsModalVisible(false); setSelectedModel(null); @@ -412,11 +420,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded setIsAgentModalVisible(true); }, []); - const handleAgentModalOk = () => { - setIsAgentModalVisible(false); - setSelectedAgent(null); - }; - const handleAgentModalCancel = () => { setIsAgentModalVisible(false); setSelectedAgent(null); @@ -427,11 +430,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded setIsMcpModalVisible(true); }, []); - const handleMcpModalOk = () => { - setIsMcpModalVisible(false); - setSelectedMcpServer(null); - }; - const handleMcpModalCancel = () => { setIsMcpModalVisible(false); setSelectedMcpServer(null); @@ -468,758 +466,784 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded const agentColumns = useMemo(() => getPublicAgentHubColumns({ onAgentClick: showAgentModal }), [showAgentModal]); const mcpColumns = useMemo(() => getPublicMCPHubColumns({ onServerClick: showMcpModal }), [showMcpModal]); + const hasAgents = Array.isArray(agentHubData) && agentHubData.length > 0; + const hasMcpServers = Array.isArray(mcpHubData) && mcpHubData.length > 0; + + const providerOptions = useMemo( + () => (Array.isArray(modelHubData) ? getUniqueProviders(modelHubData) : []), + [modelHubData], + ); + const modeOptions = useMemo( + () => + Array.isArray(modelHubData) ? getUniqueModes(modelHubData).map((mode) => ({ label: mode, value: mode })) : [], + [modelHubData], + ); + const featureOptions = useMemo( + () => + Array.isArray(modelHubData) + ? getUniqueFeatures(modelHubData).map((feature) => ({ label: feature, value: feature })) + : [], + [modelHubData], + ); + const agentSkillOptions = useMemo( + () => + Array.isArray(agentHubData) + ? getUniqueAgentSkills(agentHubData).map((skill) => ({ label: skill, value: skill })) + : [], + [agentHubData], + ); + const mcpTransportOptions = useMemo( + () => + Array.isArray(mcpHubData) + ? getUniqueMcpTransports(mcpHubData).map((transport) => ({ label: transport, value: transport })) + : [], + [mcpHubData], + ); + return ( -
- {/* Navigation - only show when not embedded */} - {!isEmbedded && } + +
+ {/* Navigation - only show when not embedded */} + {!isEmbedded && } -
- {/* Embedded Explainer - only shown when embedded in dashboard */} - {isEmbedded && ( -
-

- These are models, agents, and MCP servers your proxy admin has indicated are available in your company. -

-
- )} - - {/* About Section - only shown when not embedded */} - {!isEmbedded && ( - - About -

- {customDocsDescription ? customDocsDescription : "Proxy Server to call 100+ LLMs in the OpenAI format."} -

-
- - 🔧 - Built with litellm: v{litellmVersion} - +
+ {/* Embedded Explainer - only shown when embedded in dashboard */} + {isEmbedded && ( +
+

+ These are models, agents, and MCP servers your proxy admin has indicated are available in your + company. +

- - )} + )} - {/* Useful Links - only shown when not embedded */} - {usefulLinks && Object.keys(usefulLinks).length > 0 && ( - - Useful Links -
- {Object.entries(usefulLinks || {}) - .map(([title, value]) => { - // Handle both old format (string) and new format ({url, index}) - const url = typeof value === "string" ? value : value.url; - const index = typeof value === "string" ? 0 : value.index ?? 0; - return { title, url, index }; - }) - .sort((a, b) => a.index - b.index) - .map(({ title, url }) => ( - - ))} -
- - )} - - {/* Health and Endpoint Status - only shown when not embedded */} - {!isEmbedded && ( - - Health and Endpoint Status -
- Service status: {serviceStatus} -
- - )} - - {/* Tabs for Models and Agents */} - - - {/* Models Tab */} - -
- Available Models + {/* About Section - only shown when not embedded */} + {!isEmbedded && ( + +

About

+

+ {customDocsDescription + ? customDocsDescription + : "Proxy Server to call 100+ LLMs in the OpenAI format."} +

+
+ + 🔧 + Built with litellm: v{litellmVersion} +
+ + )} - {/* Filters */} -
-
-
- Search Models: - 0 && ( + +

Useful Links

+
+ {Object.entries(usefulLinks || {}) + .map(([title, value]) => { + // Handle both old format (string) and new format ({url, index}) + const url = typeof value === "string" ? value : value.url; + const index = typeof value === "string" ? 0 : value.index ?? 0; + return { title, url, index }; + }) + .sort((a, b) => a.index - b.index) + .map(({ title, url }) => ( +
-
- - setSearchTerm(e.target.value)} - className="border border-gray-300 rounded-lg pl-10 pr-4 py-2 w-full text-sm focus:outline-hidden focus:ring-2 focus:ring-blue-500 focus:border-transparent bg-white" - /> -
-
-
- Provider: - -
-
- Mode: - -
-
- Features: - -
-
- - model.model_group || String(index)} - sortingMode="client" - sorting={modelSorting} - onSortingChange={setModelSorting} - isLoading={loading} - loadingMessage="Loading models…" - noDataMessage={ - - } - size="compact" - /> - -
- - Showing {filteredData.length} of {modelHubData?.length || 0} models - -
- - - {/* Agents Tab */} - {agentHubData && Array.isArray(agentHubData) && agentHubData.length > 0 && ( - -
- Available Agents -
- - {/* Filters */} -
-
-
- Search Agents: - - - -
-
- - setAgentSearchTerm(e.target.value)} - className="border border-gray-300 rounded-lg pl-10 pr-4 py-2 w-full text-sm focus:outline-hidden focus:ring-2 focus:ring-blue-500 focus:border-transparent bg-white" - /> -
-
-
- Skills: - -
-
- - agent.name || String(index)} - sortingMode="client" - sorting={agentSorting} - onSortingChange={setAgentSorting} - isLoading={agentLoading} - loadingMessage="Loading agents…" - noDataMessage={ - - } - size="compact" - /> - -
- - Showing {filteredAgentData.length} of {agentHubData?.length || 0} agents - -
- - )} - - {/* MCP Servers Tab */} - {mcpHubData && Array.isArray(mcpHubData) && mcpHubData.length > 0 && ( - -
- Available MCP Servers -
- - {/* Filters */} -
-
-
- Search MCP Servers: - - - -
-
- - setMcpSearchTerm(e.target.value)} - className="border border-gray-300 rounded-lg pl-10 pr-4 py-2 w-full text-sm focus:outline-hidden focus:ring-2 focus:ring-blue-500 focus:border-transparent bg-white" - /> -
-
-
- Transport: - -
-
- - server.server_id || String(index)} - sortingMode="client" - sorting={mcpSorting} - onSortingChange={setMcpSorting} - isLoading={mcpLoading} - loadingMessage="Loading MCP servers…" - noDataMessage={ - - } - size="compact" - /> - -
- - Showing {filteredMcpData.length} of {mcpHubData?.length || 0} MCP servers - -
- - )} - - {/* Skill Hub Tab */} - - - - - -
- - {/* Model Details Modal */} - - {selectedModel?.model_group || "Model Details"} - {selectedModel && ( - - copyToClipboard(selectedModel.model_group)} - className="cursor-pointer text-gray-500 hover:text-blue-500 w-4 h-4" - /> - - )} -
- } - width={1000} - open={isModalVisible} - footer={null} - onOk={handleModalOk} - onCancel={handleModalCancel} - > - {selectedModel && ( -
- {/* Model Overview */} -
- Model Overview -
-
- Model Name: - {selectedModel.model_group} -
-
- Mode: - {selectedModel.mode || "Not specified"} -
-
- Providers: -
- {(selectedModel.providers ?? []).map((provider) => { - const { logo } = getProviderLogoAndName(provider); - return ( - -
- {logo && ( - {provider} { - (e.target as HTMLImageElement).style.display = "none"; - }} - /> - )} - {provider} -
- - ); - })} -
-
-
- - {/* Wildcard Routing Note */} - {selectedModel.model_group.includes("*") && ( -
-
- -
- Wildcard Routing - - This model uses wildcard routing. You can pass any value where you see the{" "} - * symbol. - - - For example, with{" "} - - {selectedModel.model_group} - - , you can use any string ( - - {selectedModel.model_group.replaceAll("*", "my-custom-value")} - - ) that matches this pattern. - -
-
-
- )} -
- - {/* Token and Cost Information */} -
- Token & Cost Information -
-
- Max Input Tokens: - {selectedModel.max_input_tokens?.toLocaleString() || "Not specified"} -
-
- Max Output Tokens: - {selectedModel.max_output_tokens?.toLocaleString() || "Not specified"} -
-
- Input Cost per 1M Tokens: - - {selectedModel.input_cost_per_token - ? formatCost(selectedModel.input_cost_per_token) - : "Not specified"} - -
-
- Output Cost per 1M Tokens: - - {selectedModel.output_cost_per_token - ? formatCost(selectedModel.output_cost_per_token) - : "Not specified"} - -
-
-
- - {/* Capabilities */} -
- Capabilities -
- {(() => { - const capabilities = getModelCapabilities(selectedModel); - const colors = ["green", "blue", "purple", "orange", "red", "yellow"]; - - if (capabilities.length === 0) { - return No special capabilities listed; - } - - return capabilities.map((capability, index) => ( - - {formatCapabilityName(capability)} - - )); - })()} -
-
- - {/* Rate Limits */} - {(selectedModel.tpm || selectedModel.rpm) && ( -
- Rate Limits -
- {selectedModel.tpm && ( -
- Tokens per Minute: - {selectedModel.tpm.toLocaleString()} -
- )} - {selectedModel.rpm && ( -
- Requests per Minute: - {selectedModel.rpm.toLocaleString()} -
- )} -
-
- )} - - {/* Supported OpenAI Parameters */} - {selectedModel.supported_openai_params && selectedModel.supported_openai_params.length > 0 && ( -
- Supported OpenAI Parameters -
- {selectedModel.supported_openai_params.map((param) => ( - - {param} - + +

{title}

+ ))} -
- )} + + )} - {/* Usage Example */} -
- Usage Example -
- 
-                    {(() => {
-                      const codeSnippet = generateCodeSnippet({
-                        apiKeySource: "custom",
-                        accessToken: null,
-                        apiKey: "your_api_key",
-                        inputMessage: "Hello, how are you?",
-                        chatHistory: [{ role: "user", content: "Hello, how are you?", isImage: false } as MessageType],
-                        selectedTags: [],
-                        selectedVectorStores: [],
-                        selectedGuardrails: [],
-                        selectedPolicies: [],
-                        selectedMCPServers: [],
-                        endpointType: getEndpointType(selectedModel.mode || "chat"),
-                        selectedModel: selectedModel.model_group,
-                        selectedSdk: "openai",
-                      });
-                      return codeSnippet;
-                    })()}
-
+ {/* Health and Endpoint Status - only shown when not embedded */} + {!isEmbedded && ( + +

Health and Endpoint Status

+
+

Service status: {serviceStatus}

-
- -
-
-
- )} - + + )} - {/* Agent Details Modal */} - - {selectedAgent?.name || "Agent Details"} - {selectedAgent && ( - - copyToClipboard(selectedAgent.name)} - className="cursor-pointer text-gray-500 hover:text-blue-500 w-4 h-4" - /> - - )} -
- } - width={1000} - open={isAgentModalVisible} - footer={null} - onOk={handleAgentModalOk} - onCancel={handleAgentModalCancel} - > - {selectedAgent && ( -
- {/* Agent Overview */} -
- Agent Overview -
-
- Name: - {selectedAgent.name} + {/* Tabs for Models and Agents */} + + + + Model Hub + {hasAgents && Agent Hub} + {hasMcpServers && MCP Hub} + Skill Hub + + + {/* Models Tab */} + +
+

Available Models

-
- Version: - {selectedAgent.version} -
-
- Description: - {selectedAgent.description} -
- {selectedAgent.url && ( + + {/* Filters */} +
- URL: - - {selectedAgent.url} - +
+

Search Models:

+ + } /> + + Smart search with relevance ranking - finds models containing your search terms, ranked by + relevance. Try searching 'xai grok-4', 'claude-4', 'gpt-4', or + 'sonnet' + + +
+
+ + setSearchTerm(e.target.value)} + className="border border-gray-300 rounded-lg pl-10 pr-4 py-2 w-full text-sm focus:outline-hidden focus:ring-2 focus:ring-blue-500 focus:border-transparent bg-white" + /> +
+
+
+

Provider:

+ setSelectedProviders(values)} + > + + + {(values: string[]) => + values.map((provider) => ( + + {provider} + + )) + } + + + + + No providers found + + {(provider: string) => { + const { logo } = getProviderLogoAndName(provider); + return ( + + + {logo && ( + {provider} { + (e.target as HTMLImageElement).style.display = "none"; + }} + /> + )} + {provider} + + + ); + }} + + + +
+
+

Mode:

+ +
+
+

Features:

+
- )} -
-
- - {/* Capabilities */} - {selectedAgent.capabilities && ( -
- Capabilities -
- {Object.entries(selectedAgent.capabilities) - .filter(([_, value]) => value === true) - .map(([key]) => ( - - {key} - - ))}
-
- )} - {/* Skills */} - {selectedAgent.skills && selectedAgent.skills.length > 0 && ( -
- Skills -
- {selectedAgent.skills.map((skill, index) => ( -
-
+ model.model_group || String(index)} + sortingMode="client" + sorting={modelSorting} + onSortingChange={setModelSorting} + isLoading={loading} + loadingMessage="Loading models…" + noDataMessage={ + + } + size="compact" + /> + +
+

+ Showing {filteredData.length} of {modelHubData?.length || 0} models +

+
+ + + {/* Agents Tab */} + {hasAgents && ( + +
+

Available Agents

+
+ + {/* Filters */} +
+
+
+

Search Agents:

+ + } /> + Search agents by name or description + +
+
+ + setAgentSearchTerm(e.target.value)} + className="border border-gray-300 rounded-lg pl-10 pr-4 py-2 w-full text-sm focus:outline-hidden focus:ring-2 focus:ring-blue-500 focus:border-transparent bg-white" + /> +
+
+
+

Skills:

+ +
+
+ + agent.name || String(index)} + sortingMode="client" + sorting={agentSorting} + onSortingChange={setAgentSorting} + isLoading={agentLoading} + loadingMessage="Loading agents…" + noDataMessage={ + + } + size="compact" + /> + +
+

+ Showing {filteredAgentData.length} of {agentHubData?.length || 0} agents +

+
+ + )} + + {/* MCP Servers Tab */} + {hasMcpServers && ( + +
+

Available MCP Servers

+
+ + {/* Filters */} +
+
+
+

Search MCP Servers:

+ + } /> + Search MCP servers by name or description + +
+
+ + setMcpSearchTerm(e.target.value)} + className="border border-gray-300 rounded-lg pl-10 pr-4 py-2 w-full text-sm focus:outline-hidden focus:ring-2 focus:ring-blue-500 focus:border-transparent bg-white" + /> +
+
+
+

Transport:

+ +
+
+ + server.server_id || String(index)} + sortingMode="client" + sorting={mcpSorting} + onSortingChange={setMcpSorting} + isLoading={mcpLoading} + loadingMessage="Loading MCP servers…" + noDataMessage={ + + } + size="compact" + /> + +
+

+ Showing {filteredMcpData.length} of {mcpHubData?.length || 0} MCP servers +

+
+ + )} + + {/* Skill Hub Tab */} + + + + + +
+ + {/* Model Details Modal */} + !open && handleModalCancel()}> + + + + {selectedModel?.model_group || "Model Details"} + {selectedModel && ( + + copyToClipboard(selectedModel.model_group)} + className="cursor-pointer text-gray-500 hover:text-blue-500 w-4 h-4 shrink-0" + /> + } + /> + Copy model name + + )} + + + {selectedModel && ( +
+ {/* Model Overview */} +
+

Model Overview

+
+
+

Model Name:

+

{selectedModel.model_group}

+
+
+

Mode:

+

{selectedModel.mode || "Not specified"}

+
+
+

Providers:

+
+ {(selectedModel.providers ?? []).map((provider) => { + const { logo } = getProviderLogoAndName(provider); + return ( + +
+ {logo && ( + {provider} { + (e.target as HTMLImageElement).style.display = "none"; + }} + /> + )} + {provider} +
+ + ); + })} +
+
+
+ + {/* Wildcard Routing Note */} + {selectedModel.model_group.includes("*") && ( +
+
+
- {skill.name} - {skill.description} +

Wildcard Routing

+

+ This model uses wildcard routing. You can pass any value where you see the{" "} + * symbol. +

+

+ For example, with{" "} + + {selectedModel.model_group} + + , you can use any string ( + + {selectedModel.model_group.replaceAll("*", "my-custom-value")} + + ) that matches this pattern. +

- {skill.tags && skill.tags.length > 0 && ( -
- {skill.tags.map((tag) => ( - - {tag} - - ))} +
+ )} +
+ + {/* Token and Cost Information */} +
+

Token & Cost Information

+
+
+

Max Input Tokens:

+

{selectedModel.max_input_tokens?.toLocaleString() || "Not specified"}

+
+
+

Max Output Tokens:

+

{selectedModel.max_output_tokens?.toLocaleString() || "Not specified"}

+
+
+

Input Cost per 1M Tokens:

+

+ {selectedModel.input_cost_per_token + ? formatCost(selectedModel.input_cost_per_token) + : "Not specified"} +

+
+
+

Output Cost per 1M Tokens:

+

+ {selectedModel.output_cost_per_token + ? formatCost(selectedModel.output_cost_per_token) + : "Not specified"} +

+
+
+
+ + {/* Capabilities */} +
+

Capabilities

+
+ {(() => { + const capabilities = getModelCapabilities(selectedModel); + + if (capabilities.length === 0) { + return

No special capabilities listed

; + } + + return capabilities.map((capability) => ( + + {formatCapabilityName(capability)} + + )); + })()} +
+
+ + {/* Rate Limits */} + {(selectedModel.tpm || selectedModel.rpm) && ( +
+

Rate Limits

+
+ {selectedModel.tpm && ( +
+

Tokens per Minute:

+

{selectedModel.tpm.toLocaleString()}

+
+ )} + {selectedModel.rpm && ( +
+

Requests per Minute:

+

{selectedModel.rpm.toLocaleString()}

)}
- ))} -
-
- )} - - {/* Input/Output Modes */} -
- Input/Output Modes -
-
- Input Modes: -
- {(selectedAgent.defaultInputModes ?? []).map((mode) => ( - - {mode} - - ))}
-
+ )} + + {/* Supported OpenAI Parameters */} + {selectedModel.supported_openai_params && selectedModel.supported_openai_params.length > 0 && ( +
+

Supported OpenAI Parameters

+
+ {selectedModel.supported_openai_params.map((param) => ( + + {param} + + ))} +
+
+ )} + + {/* Usage Example */}
- Output Modes: -
- {(selectedAgent.defaultOutputModes ?? []).map((mode) => ( - - {mode} - - ))} +

Usage Example

+
+ 
+                        {(() => {
+                          const codeSnippet = generateCodeSnippet({
+                            apiKeySource: "custom",
+                            accessToken: null,
+                            apiKey: "your_api_key",
+                            inputMessage: "Hello, how are you?",
+                            chatHistory: [
+                              { role: "user", content: "Hello, how are you?", isImage: false } as MessageType,
+                            ],
+                            selectedTags: [],
+                            selectedVectorStores: [],
+                            selectedGuardrails: [],
+                            selectedPolicies: [],
+                            selectedMCPServers: [],
+                            endpointType: getEndpointType(selectedModel.mode || "chat"),
+                            selectedModel: selectedModel.model_group,
+                            selectedSdk: "openai",
+                          });
+                          return codeSnippet;
+                        })()}
+
+
+
+
-
- - {/* Documentation */} - {selectedAgent.documentationUrl && ( -
- Documentation - - - View Documentation - -
)} + +
- {/* A2A Usage Example */} -
- Usage Example (A2A Protocol) + {/* Agent Details Modal */} + !open && handleAgentModalCancel()}> + + + + {selectedAgent?.name || "Agent Details"} + {selectedAgent && ( + + copyToClipboard(selectedAgent.name)} + className="cursor-pointer text-gray-500 hover:text-blue-500 w-4 h-4 shrink-0" + /> + } + /> + Copy agent name + + )} + + + {selectedAgent && ( +
+ {/* Agent Overview */} +
+

Agent Overview

+
+
+

Name:

+

{selectedAgent.name}

+
+
+

Version:

+

{selectedAgent.version}

+
+
+

Description:

+

{selectedAgent.description}

+
+ {selectedAgent.url && ( + + )} +
+
- {/* Step 1: Retrieve Agent Card */} -
- Step 1: Retrieve Agent Card -
- 
-                      {`base_url = '${selectedAgent.url}'
+                  {/* Capabilities */}
+                  {selectedAgent.capabilities && (
+                    

+                      
Capabilities

+                      

+                        {Object.entries(selectedAgent.capabilities)
+                          .filter(([_, value]) => value === true)
+                          .map(([key]) => (
+                            
+                              {key}
+                            
+                          ))}
+                      

+                    

+                  )}
+
+                  {/* Skills */}
+                  {selectedAgent.skills && selectedAgent.skills.length > 0 && (
+                    

+                      
Skills

+                      

+                        {selectedAgent.skills.map((skill, index) => (
+                          

+                            

+                              

+                                
{skill.name}

+                                
{skill.description}

+                              

+                            

+                            {skill.tags && skill.tags.length > 0 && (
+                              

+                                {skill.tags.map((tag) => (
+                                  
+                                    {tag}
+                                  
+                                ))}
+                              

+                            )}
+                          

+                        ))}
+                      

+                    

+                  )}
+
+                  {/* Input/Output Modes */}
+                  

+                    
Input/Output Modes

+                    

+                      

+                        
Input Modes:

+                        

+                          {(selectedAgent.defaultInputModes ?? []).map((mode) => (
+                            
+                              {mode}
+                            
+                          ))}
+                        

+                      

+                      

+                        
Output Modes:

+                        

+                          {(selectedAgent.defaultOutputModes ?? []).map((mode) => (
+                            
+                              {mode}
+                            
+                          ))}
+                        

+                      

+                    

+                  

+
+                  {/* Documentation */}
+                  {selectedAgent.documentationUrl && (
+                    

+                      
Documentation

+                      
+                        
+                        View Documentation
+                      
+                    

+                  )}
+
+                  {/* A2A Usage Example */}
+                  

+                    
Usage Example (A2A Protocol)

+
+                    {/* Step 1: Retrieve Agent Card */}
+                    

+                      
Step 1: Retrieve Agent Card

+                      

+                        
+                          {`base_url = '${selectedAgent.url}'
 
 resolver = A2ACardResolver(
     httpx_client=httpx_client,
@@ -1251,12 +1275,12 @@ if _public_card.supports_authenticated_extended_card:
             f'Failed to fetch extended agent card: {e_extended}. Will proceed with public card.',
             exc_info=True,
         )`}
-                    

-                  

-                  

-
+
+
+ -
-
+ copyToClipboard(codeSnippet); + }} + className="text-sm text-blue-600 hover:text-blue-800 cursor-pointer" + > + Copy to clipboard + +
+
- {/* Step 2: Call the Agent */} -
- Step 2: Call the Agent -
- 
-                      {`client = A2AClient(
+                    {/* Step 2: Call the Agent */}
+                    

+                      
Step 2: Call the Agent

+                      

+                        
+                          {`client = A2AClient(
     httpx_client=httpx_client, agent_card=final_agent_card_to_use
 )
 
@@ -1333,12 +1357,12 @@ request = SendMessageRequest(
 
 response = await client.send_message(request)
 print(response.model_dump(mode='json', exclude_none=True))`}
-                    

-                  

-                  

-
+
+
+ + copyToClipboard(codeSnippet); + }} + className="text-sm text-blue-600 hover:text-blue-800 cursor-pointer" + > + Copy to clipboard + +
+
-
-
- )} - - - {/* MCP Server Details Modal */} - - {selectedMcpServer?.server_name || "MCP Server Details"} - {selectedMcpServer && ( - - copyToClipboard(selectedMcpServer.server_name)} - className="cursor-pointer text-gray-500 hover:text-blue-500 w-4 h-4" - /> - )} -
- } - width={1000} - open={isMcpModalVisible} - footer={null} - onOk={handleMcpModalOk} - onCancel={handleMcpModalCancel} - > - {selectedMcpServer && ( -
- {/* Server Overview */} -
- Server Overview -
+ + + + {/* MCP Server Details Modal */} + !open && handleMcpModalCancel()}> + + + + {selectedMcpServer?.server_name || "MCP Server Details"} + {selectedMcpServer && ( + + copyToClipboard(selectedMcpServer.server_name)} + className="cursor-pointer text-gray-500 hover:text-blue-500 w-4 h-4 shrink-0" + /> + } + /> + Copy server name + + )} + + + {selectedMcpServer && ( +
+ {/* Server Overview */}
- Server Name: - {selectedMcpServer.server_name} +

Server Overview

+
+
+

Server Name:

+

{selectedMcpServer.server_name}

+
+
+

Transport:

+ {selectedMcpServer.transport} +
+ {selectedMcpServer.alias && ( +
+

Alias:

+

{selectedMcpServer.alias}

+
+ )} +
+

Auth Type:

+ + {selectedMcpServer.auth_type} + +
+
+

Description:

+

{selectedMcpServer.mcp_info?.description || "-"}

+
+
-
- Transport: - {selectedMcpServer.transport} -
- {selectedMcpServer.alias && ( + + {/* Additional Info */} + {selectedMcpServer.mcp_info && Object.keys(selectedMcpServer.mcp_info).length > 0 && (
- Alias: - {selectedMcpServer.alias} +

Additional Information

+
+ 
+                          {JSON.stringify(selectedMcpServer.mcp_info, null, 2)}
+
+
)} + + {/* Usage Example */}
- Auth Type: - - {selectedMcpServer.auth_type} - -
-
- Description: - {selectedMcpServer.mcp_info?.description || "-"} -
-
-
- - {/* Additional Info */} - {selectedMcpServer.mcp_info && Object.keys(selectedMcpServer.mcp_info).length > 0 && ( -
- Additional Information -
- 
{JSON.stringify(selectedMcpServer.mcp_info, null, 2)}
-
-
- )} - - {/* Usage Example */} -
- Usage Example -
- 
-                    {`# Using MCP Server with Python FastMCP
+                    
Usage Example

+                    

+                      
+                        {`# Using MCP Server with Python FastMCP
 
 from fastmcp import Client
 import asyncio
@@ -1474,12 +1501,12 @@ async def main():
 
 if __name__ == "__main__":
     asyncio.run(main())`}
-                  

-                

-                

-
+
+
+ + copyToClipboard(codeSnippet); + }} + className="text-sm text-blue-600 hover:text-blue-800 cursor-pointer" + > + Copy to clipboard + +
+
-
-
- )} - -
+ )} + + +
+ ); }; diff --git a/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx b/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx index 1d1f0c92ad1..f140e4cc371 100644 --- a/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { Input } from "antd"; +import { Input } from "@/components/ui/input"; interface routingStrategyArgs { ttl?: number; diff --git a/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.tsx b/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.tsx index da089552b11..dc6b397827a 100644 --- a/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { Input } from "antd"; +import { Input } from "@/components/ui/input"; interface ReliabilityRetriesSectionProps { routerSettings: { [key: string]: any }; diff --git a/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx b/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx index 76965afd603..2c3c7808b0f 100644 --- a/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx @@ -22,6 +22,11 @@ describe("TagFilteringToggle", () => { expect(screen.getByText("Enable Tag Filtering")).toBeInTheDocument(); }); + it("should name the switch with the metadata label so it is reachable by accessible name", () => { + render(); + expect(screen.getByRole("switch", { name: "Tag Filtering" })).toBeInTheDocument(); + }); + it("should display the label from metadata when provided", () => { render(); expect(screen.getByText("Tag Filtering")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.tsx b/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.tsx index 44bac326bb0..3c3ba4e3160 100644 --- a/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.tsx @@ -1,5 +1,5 @@ -import React from "react"; -import { Switch } from "antd"; +import React, { useId } from "react"; +import { Switch } from "@/components/ui/switch"; interface TagFilteringToggleProps { enabled: boolean; @@ -8,11 +8,13 @@ interface TagFilteringToggleProps { } const TagFilteringToggle: React.FC = ({ enabled, routerFieldsMetadata, onToggle }) => { + const toggleId = useId(); + return (
-
- +
); diff --git a/ui/litellm-dashboard/src/components/router_settings/index.test.tsx b/ui/litellm-dashboard/src/components/router_settings/index.test.tsx index 94cbb94d164..72657069cde 100644 --- a/ui/litellm-dashboard/src/components/router_settings/index.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/index.test.tsx @@ -134,6 +134,30 @@ describe("RouterSettings", () => { ); }); + it("should send the edited input value, not the loaded one, on Save Changes", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByTestId("strategy-select")).toBeInTheDocument(); + }); + + const numRetries = await screen.findByRole("textbox", { name: /num_retries/i }); + await user.clear(numRetries); + await user.type(numRetries, "42"); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => + expect(setCallbacksCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + router_settings: expect.objectContaining({ num_retries: 42 }), + }), + ), + ); + }); + it("should show a success notification after saving", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/router_settings/index.tsx b/ui/litellm-dashboard/src/components/router_settings/index.tsx index aea08425388..b2668d98d69 100644 --- a/ui/litellm-dashboard/src/components/router_settings/index.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/index.tsx @@ -1,5 +1,5 @@ -import { Button } from "antd"; import React, { useEffect, useState } from "react"; +import { Button } from "@/components/ui/button"; import NotificationsManager from "../molecules/notifications_manager"; import { getCallbacksCall, getRouterSettingsCall, setCallbacksCall } from "../networking"; import RouterSettingsForm, { RouterSettingsFormValue } from "./RouterSettingsForm"; @@ -190,10 +190,10 @@ const RouterSettings: React.FC = ({ accessToken, userRole, {/* Actions - Sticky at bottom */}
- - +
); diff --git a/ui/litellm-dashboard/src/components/settings.test.tsx b/ui/litellm-dashboard/src/components/settings.test.tsx index c9bcc1eb5b9..62efa1dc372 100644 --- a/ui/litellm-dashboard/src/components/settings.test.tsx +++ b/ui/litellm-dashboard/src/components/settings.test.tsx @@ -1,8 +1,8 @@ -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { act, render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { Form } from "antd"; +import { FormProvider, useForm } from "react-hook-form"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; -import { alertingSettingsCall, getCallbackConfigsCall, getCallbacksCall } from "./networking"; +import { alertingSettingsCall, getCallbackConfigsCall, getCallbacksCall, setCallbacksCall } from "./networking"; import Settings, { backendCallbackLogoSrc, CallbackSelector } from "./settings"; vi.mock("./networking", () => ({ @@ -114,42 +114,20 @@ describe("Settings", () => { }); }); - it("should display edit modal with fields when edit is clicked", async () => { - const mockCallback = { - name: "langfuse", - variables: { - LANGFUSE_PUBLIC_KEY: "test-public-key", - LANGFUSE_SECRET_KEY: "test-secret-key", - LANGFUSE_HOST: "https://test.langfuse.com", - SLACK_WEBHOOK_URL: null, - OPENMETER_API_KEY: null, - }, - }; - - const mockCallbackConfig = { - id: "langfuse", - displayName: "Langfuse", - dynamic_params: { - LANGFUSE_PUBLIC_KEY: { - type: "text", - ui_name: "Public Key", - required: true, - }, - LANGFUSE_SECRET_KEY: { - type: "password", - ui_name: "Secret Key", - required: true, - }, - LANGFUSE_HOST: { - type: "text", - ui_name: "Host", - required: false, - }, - }, - }; - + const openLangfuseEditModal = async () => { mockGetCallbacksCall.mockResolvedValue({ - callbacks: [mockCallback], + callbacks: [ + { + name: "langfuse", + variables: { + LANGFUSE_PUBLIC_KEY: "test-public-key", + LANGFUSE_SECRET_KEY: "test-secret-key", + LANGFUSE_HOST: "https://test.langfuse.com", + SLACK_WEBHOOK_URL: null, + OPENMETER_API_KEY: null, + }, + }, + ], available_callbacks: { langfuse: { litellm_callback_name: "langfuse", @@ -160,30 +138,118 @@ describe("Settings", () => { alerts: [], }); - mockGetCallbackConfigsCall.mockResolvedValue([mockCallbackConfig]); + mockGetCallbackConfigsCall.mockResolvedValue([ + { + id: "langfuse", + displayName: "Langfuse", + dynamic_params: { + LANGFUSE_PUBLIC_KEY: { type: "text", ui_name: "Public Key", required: true }, + LANGFUSE_SECRET_KEY: { type: "password", ui_name: "Secret Key", required: true }, + LANGFUSE_HOST: { type: "text", ui_name: "Host", required: false }, + }, + }, + ]); const user = userEvent.setup(); - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("Active Logging Callbacks")).toBeInTheDocument(); + expect(screen.getByText("Active Logging Callbacks")).toBeInTheDocument(); }); await waitFor(() => { - expect(getByText("Langfuse")).toBeInTheDocument(); + expect(screen.getByText("Langfuse")).toBeInTheDocument(); }); await user.click(screen.getByTestId("callback-actions-langfuse-success")); await user.click(await screen.findByTestId("callback-action-edit")); await waitFor(() => { - expect(getByText("Edit Callback Settings")).toBeInTheDocument(); + expect(screen.getByText("Edit Callback Settings")).toBeInTheDocument(); + }); + + return user; + }; + + it("should display edit modal with fields when edit is clicked", async () => { + await openLangfuseEditModal(); + + await waitFor(() => { + expect(screen.getByText("Public Key")).toBeInTheDocument(); + expect(screen.getByText("Secret Key")).toBeInTheDocument(); + expect(screen.getByText("Host")).toBeInTheDocument(); }); await waitFor(() => { - expect(getByText("Public Key")).toBeInTheDocument(); - expect(getByText("Secret Key")).toBeInTheDocument(); - expect(getByText("Host")).toBeInTheDocument(); + expect(screen.getByLabelText("Public Key")).toHaveValue("test-public-key"); + }); + expect(screen.getByLabelText("Secret Key")).toHaveValue("test-secret-key"); + expect(screen.getByLabelText("Host")).toHaveValue("https://test.langfuse.com"); + + const danglingLabels = [...document.querySelectorAll("label[for]")].filter( + (label) => document.getElementById(label.getAttribute("for") as string) === null, + ); + expect(danglingLabels).toEqual([]); + }); + + it("should post the edited callback variables when the edit modal is saved", async () => { + const user = await openLangfuseEditModal(); + + await waitFor(() => { + expect(screen.getByLabelText("Host")).toHaveValue("https://test.langfuse.com"); + }); + + await user.clear(screen.getByLabelText("Host")); + await user.type(screen.getByLabelText("Host"), "https://edited.langfuse.com"); + await user.click(within(screen.getByRole("dialog")).getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(vi.mocked(setCallbacksCall)).toHaveBeenCalledWith("token", { + environment_variables: { + callback: "langfuse", + LANGFUSE_PUBLIC_KEY: "test-public-key", + LANGFUSE_SECRET_KEY: "test-secret-key", + LANGFUSE_HOST: "https://edited.langfuse.com", + }, + litellm_settings: { success_callback: ["langfuse"] }, + }); + }); + }); + + it("should block the edit submit when a required field is emptied", async () => { + const user = await openLangfuseEditModal(); + + await waitFor(() => { + expect(screen.getByLabelText("Public Key")).toHaveValue("test-public-key"); + }); + + await user.clear(screen.getByLabelText("Public Key")); + await user.click(within(screen.getByRole("dialog")).getByRole("button", { name: "Save Changes" })); + + expect(await screen.findByText("Please enter the public key")).toBeInTheDocument(); + expect(vi.mocked(setCallbacksCall)).not.toHaveBeenCalled(); + }); + + it("should send the typed webhook url for an alert type when the alerting tab is saved", async () => { + const user = userEvent.setup(); + render(); + + await user.click(await screen.findByRole("tab", { name: "Alerting Types" })); + + const webhookInput = document.querySelector('input[name="llm_exceptions"]') as HTMLInputElement; + expect(webhookInput).not.toBeNull(); + await user.type(webhookInput, "https://hooks.example.com/llm-exceptions"); + + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(vi.mocked(setCallbacksCall)).toHaveBeenCalledWith("token", { + general_settings: expect.objectContaining({ + alert_to_webhook_url: expect.objectContaining({ + llm_exceptions: "https://hooks.example.com/llm-exceptions", + }), + }), + }); }); }); @@ -252,6 +318,19 @@ describe("backendCallbackLogoSrc", () => { }); }); +const CallbackSelectorHarness = ({ + callbackConfigs, +}: { + callbackConfigs: { id: string; displayName: string; logo?: string }[]; +}) => { + const form = useForm>(); + return ( + + + + ); +}; + describe("CallbackSelector logos", () => { it("resolves backend logos per entry: bare filename, external url, and missing logo", async () => { const callbackConfigs = [ @@ -260,13 +339,9 @@ describe("CallbackSelector logos", () => { { id: "nologo", displayName: "NoLogo" }, ]; - render( -
- - , - ); + render(); - fireEvent.mouseDown(screen.getByRole("combobox")); + await userEvent.click(screen.getByRole("combobox")); expect(await screen.findByAltText("Langfuse logo")).toHaveAttribute("src", "/ui/assets/logos/langfuse.png"); expect(screen.getByAltText("Hosted logo")).toHaveAttribute("src", "https://logos.example.com/hosted.png"); diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 904fd4d611e..34fd9af06db 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -1,31 +1,26 @@ -import { - Button, - Card, - Grid, - SelectItem, - Switch, - Tab, - TabGroup, - Table, - TableBody, - TableCell, - TableHead, - TableHeaderCell, - TableRow, - TabList, - TabPanel, - TabPanels, - Text, - TextInput, -} from "@tremor/react"; import React, { useEffect, useState } from "react"; +import { Controller, FormProvider, useForm, useFormContext } from "react-hook-form"; -import { Button as Button2, Form, Input, Modal, Select } from "antd"; +import { Field, FieldError, FieldLabel } from "@/components/shared/form/field"; +import { Button } from "@/components/ui/button"; +import { Card } from "@/components/ui/card"; +import { + Combobox, + ComboboxContent, + ComboboxEmpty, + ComboboxInput, + ComboboxItem, + ComboboxList, +} from "@/components/ui/combobox"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Switch } from "@/components/ui/switch"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import EmailSettings from "./email_settings"; import { Logo } from "@/components/molecules/logo/Logo"; import NotificationsManager from "./molecules/notifications_manager"; -import FormItem from "antd/es/form/FormItem"; import AlertingSettings from "./alerting/alerting_settings"; import CloudZeroCostTracking from "./CloudZeroCostTracking/CloudZeroCostTracking"; import DeleteResourceModal from "./common_components/DeleteResourceModal"; @@ -46,6 +41,8 @@ interface SettingsPageProps { premiumUser: boolean; } +type CallbackFormValues = Record; + const assetsLogoFolder = "/ui/assets/logos/"; export const backendCallbackLogoSrc = (logo: string | null | undefined): string | undefined => { @@ -61,6 +58,9 @@ interface DynamicParamsFieldsProps { } const DynamicParamsFields: React.FC = ({ params, callbackConfigs, selectedCallback }) => { + const { register, formState } = useFormContext(); + const fieldIdPrefix = React.useId(); + if (!params || params.length === 0) { return null; } @@ -73,54 +73,51 @@ const DynamicParamsFields: React.FC = ({ params, callb const paramType = paramConfig.type || "text"; const fieldLabel = paramConfig.ui_name || param.replace(/_/g, " ").replace(/\b\w/g, (l) => l.toUpperCase()); const isRequired = paramConfig.required || false; + const fieldId = `${fieldIdPrefix}-${param}`; + const registration = register( + param, + isRequired ? { required: `Please enter the ${fieldLabel.toLowerCase()}` } : undefined, + ); return ( - {fieldLabel} } - name={param} - key={param} - className="mb-4" - rules={ - isRequired - ? [ - { - required: true, - message: `Please enter the ${fieldLabel.toLowerCase()}`, - }, - ] - : undefined - } - > + + + {fieldLabel} + {paramType === "password" ? ( - ) : paramType === "number" ? ( ) : ( - + )} - + + ); })}
); }; +interface CallbackConfigOption { + id: string; + displayName: string; + logo?: string | null; +} + // Shared component for rendering callback selector interface CallbackSelectorProps { callbackConfigs: any[]; @@ -135,42 +132,64 @@ export const CallbackSelector: React.FC = ({ onCallbackChange, disabled = false, }) => { + const { control } = useFormContext(); + const inputId = React.useId(); + const selectedConfig = callbackConfigs.find((config) => config.id === selectedCallback) ?? null; + return ( - - - + rules={disabled ? undefined : { required: "Please select a callback" }} + render={({ field, fieldState }) => ( + + Callback + { + field.onChange(config?.id ?? ""); + onCallbackChange(config?.id ?? ""); + }} + isItemEqualToValue={(a: CallbackConfigOption, b: CallbackConfigOption) => a.id === b.id} + itemToStringLabel={(config: CallbackConfigOption) => config.displayName} + filter={(config: CallbackConfigOption, query: string) => + config.id.toLowerCase().includes(query.trim().toLowerCase()) + } + disabled={disabled} + > + + + No results + + {(callbackConfig: CallbackConfigOption) => ( + +
+
+ +
+ {callbackConfig.displayName} +
+ + )} + + + + + + )} + /> ); }; @@ -206,8 +225,8 @@ const Settings: React.FC = ({ accessToken, userRole, userID, const [callbacks, setCallbacks] = useState([]); const [isLoadingCallbacks, setIsLoadingCallbacks] = useState(true); const [alerts, setAlerts] = useState([]); - const [addForm] = Form.useForm(); - const [editForm] = Form.useForm(); + const addForm = useForm({ shouldUnregister: true }); + const editForm = useForm({ shouldUnregister: true }); const [selectedCallback, setSelectedCallback] = useState(null); const [catchAllWebhookURL, setCatchAllWebhookURL] = useState(""); const [alertToWebhooks, setAlertToWebhooks] = useState>({}); @@ -254,7 +273,7 @@ const Settings: React.FC = ({ accessToken, userRole, userID, const normalized = Object.fromEntries( Object.entries(selectedEditCallback.variables || {}).map(([k, v]) => [k, v ?? ""]), ); - editForm.setFieldsValue({ + editForm.reset({ ...normalized, callback: selectedEditCallback.name, }); @@ -337,11 +356,11 @@ const Settings: React.FC = ({ accessToken, userRole, userID, if (isEdit) { setShowEditCallback(false); - editForm.resetFields(); + editForm.reset(); setSelectedEditCallback(null); } else { setShowAddCallbacksModal(false); - addForm.resetFields(); + addForm.reset(); setSelectedCallback(null); setSelectedCallbackParams([]); } @@ -383,6 +402,23 @@ const Settings: React.FC = ({ accessToken, userRole, userID, setSelectedCallbackParams(params); }; + const closeAddCallbackModal = () => { + setShowAddCallbacksModal(false); + setSelectedCallback(null); + setSelectedCallbackParams([]); + }; + + const cancelAddCallback = () => { + closeAddCallbackModal(); + addForm.reset(); + }; + + const closeEditCallbackModal = () => { + setShowEditCallback(false); + setSelectedEditCallback(null); + editForm.reset(); + }; + const handleSaveAlerts = async () => { if (!accessToken) { return; @@ -447,257 +483,216 @@ const Settings: React.FC = ({ accessToken, userRole, userID, return (
- - - - Logging Callbacks - CloudZero Cost Tracking - Alerting Types - Alerting Settings - Email Alerts - - - - setShowAddCallbacksModal(true)} - onEdit={(cb) => { - setSelectedEditCallback(cb); - setShowEditCallback(true); - }} - onDelete={(cb) => handleDeleteCallback(cb)} - onTest={async (cb) => { - try { - await serviceHealthCheck(accessToken, cb.name); - NotificationsManager.success("Health check triggered"); - } catch (error) { - NotificationsManager.fromBackend(parseErrorMessage(error)); - } - }} - /> - - -
- -
- - - - - Alerts are only supported for Slack Webhook URLs. Get your webhook urls from{" "} - - here - - -
- - - - - Slack Webhook URL - - +
+ + + Logging Callbacks + CloudZero Cost Tracking + Alerting Types + Alerting Settings + Email Alerts + + + setShowAddCallbacksModal(true)} + onEdit={(cb) => { + setSelectedEditCallback(cb); + setShowEditCallback(true); + }} + onDelete={(cb) => handleDeleteCallback(cb)} + onTest={async (cb) => { + try { + await serviceHealthCheck(accessToken, cb.name); + NotificationsManager.success("Health check triggered"); + } catch (error) { + NotificationsManager.fromBackend(parseErrorMessage(error)); + } + }} + /> + + +
+ +
+ + + +

+ Alerts are only supported for Slack Webhook URLs. Get your webhook urls from{" "} + + here + +

+
+ + + + + Slack Webhook URL + + - - {Object.entries(alerts_to_UI_NAME).map(([key, value], index) => ( - - - {key == "region_outage_alerts" ? ( - premiumUser ? ( - handleSwitchChange(key)} - /> - ) : ( - - ) - ) : ( + + {Object.entries(alerts_to_UI_NAME).map(([key, value], index) => ( + + + {key == "region_outage_alerts" ? ( + premiumUser ? ( handleSwitchChange(key)} + onCheckedChange={() => handleSwitchChange(key)} /> - )} - - - {value} - - - - - - ))} - -
- + ) : ( + + ) + ) : ( + handleSwitchChange(key)} + /> + )} + + +

{value}

+ + + + + + ))} + + + - - - - - - - - - - - - + + + + + + + + + + +
- { - setShowAddCallbacksModal(false); - setSelectedCallback(null); - setSelectedCallbackParams([]); - }} - footer={null} - > - - {" "} - LiteLLM Docs: Logging - + !open && closeAddCallbackModal()}> + + + Add Logging Callback + + + {" "} + LiteLLM Docs: Logging + -
- - - - -
- { - setShowAddCallbacksModal(false); - setSelectedCallback(null); - setSelectedCallbackParams([]); - addForm.resetFields(); - }} - disabled={isAddingCallback} - > - Cancel - - - {isAddingCallback ? "Adding..." : "Add Callback"} - -
- -
- - { - setShowEditCallback(false); - setSelectedEditCallback(null); - editForm.resetFields(); - }} - footer={null} - > -
- {selectedEditCallback && ( - <> + + {}} - disabled={true} + selectedCallback={selectedCallback} + onCallbackChange={handleSelectedCallbackChange} /> - - )} -
- { - setShowEditCallback(false); - setSelectedEditCallback(null); - editForm.resetFields(); - }} - disabled={isUpdatingCallback} - > - Cancel - - { - editForm.submit(); - }} - loading={isUpdatingCallback} - disabled={isUpdatingCallback} - > - {isUpdatingCallback ? "Saving..." : "Save Changes"} - -
- -
+
+ + +
+ + + + + + !open && closeEditCallbackModal()}> + + + Edit Callback Settings + + +
+ {selectedEditCallback && ( + <> + {}} + disabled={true} + /> + + + + )} + +
+ + +
+ +
+ + = ({ apiKey }) => {
- +
); diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx index acf50d282b4..5e8d63eda08 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx @@ -21,6 +21,13 @@ describe("SearchSelect", () => { expect(screen.getByRole("combobox")).toHaveValue("Growth"); }); + it("shows a value the options do not carry yet instead of blanking the field", () => { + const { rerender } = render(); + expect(screen.getByRole("combobox")).toHaveValue("team-2"); + rerender(); + 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(); diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx index f67a1cffa1d..c6ae11f5729 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx @@ -24,6 +24,7 @@ interface SearchSelectProps { emptyText?: string; disabled?: boolean; className?: string; + inputId?: string; } const matchesQuery = (option: SearchSelectOption, query: string): boolean => { @@ -40,12 +41,18 @@ export function SearchSelect({ emptyText = "No results", disabled = false, className, + inputId, }: SearchSelectProps) { - const selected = options.find((option) => option.value === value) ?? null; + const selected = + value === undefined || value === "" + ? null + : options.find((option) => option.value === value) ?? { label: value, value }; + const items = + selected !== null && !options.some((option) => option.value === selected.value) ? [selected, ...options] : options; return ( onValueChange(item?.value ?? "")} isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value} @@ -54,6 +61,7 @@ export function SearchSelect({ disabled={disabled} > = ({ step = 0.01, @@ -32,7 +32,8 @@ const NumericalInput: React.FC = ({ ...rest }) => { return ( - event.currentTarget.blur()} step={step} style={style} diff --git a/ui/litellm-dashboard/src/components/shared/usage_date_picker.tsx b/ui/litellm-dashboard/src/components/shared/usage_date_picker.tsx index 821dee65a3e..390da4c8f88 100644 --- a/ui/litellm-dashboard/src/components/shared/usage_date_picker.tsx +++ b/ui/litellm-dashboard/src/components/shared/usage_date_picker.tsx @@ -1,5 +1,6 @@ import React, { useCallback, useState, useRef } from "react"; -import { DateRangePicker, DateRangePickerValue, Text } from "@tremor/react"; +import { DateRangePicker, Text } from "@tremor/react"; +import type { DateRangePickerValue } from "./date_picker_types"; interface UsageDatePickerProps { value: DateRangePickerValue; diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index df7523dfe4e..cfb7ad1a5ad 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -989,9 +989,8 @@ describe("TeamInfoView", () => { const user = userEvent.setup({ delay: null }); const resetBudgetItem = await openSettingsEditorForTeam(user, { budget_duration: "30d" }); - const clearIcon = resetBudgetItem.querySelector(".ant-select-clear"); - expect(clearIcon).not.toBeNull(); - fireEvent.mouseDown(clearIcon as Element); + await user.click(within(resetBudgetItem).getByRole("combobox")); + await user.click(await screen.findByText("Never resets")); await waitFor(() => { expect(within(resetBudgetItem).getByText("Never resets")).toBeInTheDocument(); @@ -1554,13 +1553,14 @@ describe("TeamInfoView", () => { await user.click(within(routesFormItem).getByRole("combobox")); - const option = await screen.findByTitle("POST /bedrock-passthrough"); + const option = await screen.findByText("POST /bedrock-passthrough"); await user.click(option); await waitFor(() => { expect(within(routesFormItem).getByText(/\/bedrock-passthrough/)).toBeInTheDocument(); }); + await user.keyboard("{Escape}"); await user.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => { diff --git a/ui/litellm-dashboard/src/components/team/member_permissions.test.tsx b/ui/litellm-dashboard/src/components/team/member_permissions.test.tsx index 10c78331c68..652d4f8e685 100644 --- a/ui/litellm-dashboard/src/components/team/member_permissions.test.tsx +++ b/ui/litellm-dashboard/src/components/team/member_permissions.test.tsx @@ -1,5 +1,5 @@ import * as networking from "@/components/networking"; -import { act, fireEvent, screen, waitFor } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor, within } from "@testing-library/react"; import { renderWithProviders } from "../../../tests/test-utils"; import { afterEach, describe, expect, it, vi } from "vitest"; import MemberPermissions from "./member_permissions"; @@ -9,6 +9,9 @@ vi.mock("@/components/networking", () => ({ teamPermissionsUpdateCall: vi.fn(), })); +const checkboxFor = (endpoint: string) => + within(screen.getByText(endpoint).closest("tr") as HTMLElement).getByRole("checkbox"); + describe("MemberPermissions", () => { afterEach(() => { vi.clearAllMocks(); @@ -69,32 +72,27 @@ describe("MemberPermissions", () => { expect(screen.getByText("Member Permissions")).toBeInTheDocument(); }); - const checkboxes = screen.getAllByRole("checkbox"); - const unselectedCheckbox = checkboxes.find((cb) => !(cb as HTMLInputElement).checked); + expect(checkboxFor("/key/generate")).toBeChecked(); + expect(checkboxFor("/key/list")).not.toBeChecked(); - if (unselectedCheckbox) { - await act(async () => { - fireEvent.click(unselectedCheckbox); - }); + await act(async () => { + fireEvent.click(checkboxFor("/key/list")); + }); - await waitFor(() => { - const saveButton = screen.getByRole("button", { name: /save changes/i }); - expect(saveButton).toBeInTheDocument(); - }); + expect(checkboxFor("/key/list")).toBeChecked(); - const saveButton = screen.getByRole("button", { name: /save changes/i }); - await act(async () => { - fireEvent.click(saveButton); - }); + const saveButton = await screen.findByRole("button", { name: /save changes/i }); + await act(async () => { + fireEvent.click(saveButton); + }); - await waitFor(() => { - expect(networking.teamPermissionsUpdateCall).toHaveBeenCalledWith( - "token-123", - "team-123", - expect.arrayContaining(["/key/generate", "/key/list"]), - ); - }); - } + await waitFor(() => { + expect(networking.teamPermissionsUpdateCall).toHaveBeenCalledWith( + "token-123", + "team-123", + expect.arrayContaining(["/key/generate", "/key/list"]), + ); + }); }); it("should render team daily activity permission with correct method and description", async () => { @@ -123,11 +121,13 @@ describe("MemberPermissions", () => { expect(screen.getByText("Member Permissions")).toBeInTheDocument(); }); - const checkboxes = screen.getAllByRole("checkbox"); - checkboxes.forEach((checkbox) => { - expect(checkbox).toBeDisabled(); + expect(checkboxFor("/key/list")).not.toBeChecked(); + + await act(async () => { + fireEvent.click(checkboxFor("/key/list")); }); + expect(checkboxFor("/key/list")).not.toBeChecked(); expect(screen.queryByRole("button", { name: /save changes/i })).not.toBeInTheDocument(); }); @@ -143,32 +143,27 @@ describe("MemberPermissions", () => { expect(screen.getByText("Member Permissions")).toBeInTheDocument(); }); - const checkboxes = screen.getAllByRole("checkbox"); - const unselectedCheckbox = checkboxes.find((cb) => !(cb as HTMLInputElement).checked); + await act(async () => { + fireEvent.click(checkboxFor("/key/list")); + }); - if (unselectedCheckbox) { - await act(async () => { - fireEvent.click(unselectedCheckbox); - }); + expect(checkboxFor("/key/list")).toBeChecked(); - await waitFor(() => { - const resetButton = screen.getByRole("button", { name: /reset/i }); - expect(resetButton).toBeInTheDocument(); - }); + vi.mocked(networking.getTeamPermissionsCall).mockResolvedValueOnce({ + all_available_permissions: ["/key/generate", "/key/list"], + team_member_permissions: ["/key/generate"], + }); - vi.mocked(networking.getTeamPermissionsCall).mockResolvedValueOnce({ - all_available_permissions: ["/key/generate", "/key/list"], - team_member_permissions: ["/key/generate"], - }); + const resetButton = await screen.findByRole("button", { name: /reset/i }); + await act(async () => { + fireEvent.click(resetButton); + }); - const resetButton = screen.getByRole("button", { name: /reset/i }); - await act(async () => { - fireEvent.click(resetButton); - }); + await waitFor(() => { + expect(networking.getTeamPermissionsCall).toHaveBeenCalledTimes(2); + }); - await waitFor(() => { - expect(networking.getTeamPermissionsCall).toHaveBeenCalledTimes(2); - }); - } + expect(checkboxFor("/key/list")).not.toBeChecked(); + expect(screen.queryByRole("button", { name: /save changes/i })).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/team/member_permissions.tsx b/ui/litellm-dashboard/src/components/team/member_permissions.tsx index 5bbd82f4a5d..62c7d1f96da 100644 --- a/ui/litellm-dashboard/src/components/team/member_permissions.tsx +++ b/ui/litellm-dashboard/src/components/team/member_permissions.tsx @@ -1,7 +1,9 @@ import { getTeamPermissionsCall, teamPermissionsUpdateCall } from "@/components/networking"; -import { ReloadOutlined, SaveOutlined } from "@ant-design/icons"; -import { Card, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Text, Title } from "@tremor/react"; -import { Button, Checkbox, Empty } from "antd"; +import { Button } from "@/components/ui/button"; +import { Card } from "@/components/ui/card"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { RotateCw, Save } from "lucide-react"; import React, { useEffect, useState } from "react"; import NotificationsManager from "../molecules/notifications_manager"; import { getPermissionInfo } from "./permission_definitions"; @@ -75,36 +77,38 @@ const MemberPermissions: React.FC = ({ teamId, accessTok const hasPermissions = permissions.length > 0; return ( - +
- Member Permissions +

Member Permissions

{canEditTeam && hasChanges && (
- -
)}
- Control what team members can do when they are not team admins. +

Control what team members can do when they are not team admins.

{hasPermissions ? (
- - +
+ - Method - Endpoint - Description - + Method + Endpoint + Description + Allow Access - + - + {permissions.map((permission) => { const permInfo = getPermissionInfo(permission); @@ -125,8 +129,9 @@ const MemberPermissions: React.FC = ({ teamId, accessTok {permInfo.description} handlePermissionChange(permission, e.target.checked)} + onCheckedChange={(checked) => handlePermissionChange(permission, checked)} disabled={!canEditTeam} /> @@ -138,7 +143,7 @@ const MemberPermissions: React.FC = ({ teamId, accessTok ) : (
- +

No permissions available

)} diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx index d0dd782a697..f31a265da87 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx @@ -1,27 +1,35 @@ import React from "react"; -import { Button, Typography, Tooltip, Space, Divider, Flex, Popover, Dropdown, Tag } from "antd"; -import type { MenuProps } from "antd"; import { - ArrowLeftOutlined, - SyncOutlined, - DeleteOutlined, - PlusOutlined, - UserOutlined, - CalendarOutlined, - ClockCircleOutlined, - ThunderboltOutlined, - SafetyCertificateOutlined, - TransactionOutlined, - FieldTimeOutlined, - MoreOutlined, - StopOutlined, - CheckCircleOutlined, -} from "@ant-design/icons"; + ArrowLeft, + ArrowLeftRight, + Ban, + Calendar, + CircleCheck, + Clock, + MoreVertical, + Plus, + RefreshCw, + ShieldCheck, + Timer, + Trash2, + User, + Zap, +} from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card"; +import { Separator } from "@/components/ui/separator"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import CopyButton from "@/components/shared/CopyButton"; import LabeledField from "../common_components/LabeledField"; import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; -const { Title, Text } = Typography; - export interface KeyInfoData { keyName: string; keyId: string; @@ -52,14 +60,12 @@ interface KeyInfoHeaderProps { function UserField({ userAlias, userEmail, userId }: { userAlias?: string | null; userEmail: string; userId: string }) { const labelEl = ( - - - - - - User - - +
+ + + + User +
); const isEmpty = !userAlias && !userEmail && !userId; @@ -68,7 +74,7 @@ function UserField({ userAlias, userEmail, userId }: { userAlias?: string | null
{labelEl}
- - + -
); @@ -87,14 +93,12 @@ function UserField({ userAlias, userEmail, userId }: { userAlias?: string | null
{label} {value ? ( - - {value} - +
+ + {value} + + +
) : ( - )} @@ -108,11 +112,18 @@ function UserField({ userAlias, userEmail, userId }: { userAlias?: string | null
{labelEl}
- - - - - + + + + + } + /> + + {popoverContent} + +
); @@ -122,11 +133,14 @@ function UserField({ userAlias, userEmail, userId }: { userAlias?: string | null
{labelEl}
- - - {displayValue} - - + + {displayValue}} + /> + + {popoverContent} + +
); @@ -146,104 +160,124 @@ export function KeyInfoHeader({ regenerateDisabled = false, regenerateTooltip, }: KeyInfoHeaderProps) { - const destructiveActionItems: MenuProps["items"] = [ - ...(onToggleBlocked - ? [ - isBlocked - ? { key: "unblock", label: "Unblock Key", icon: } - : { key: "block", label: "Block Key", icon: , danger: true }, - ] - : []), - ...(onResetSpend - ? [{ key: "reset-spend", label: "Reset Spend", icon: , danger: true }] - : []), - { key: "delete", label: "Delete Key", icon: , danger: true }, - ]; - - const handleDestructiveActionClick: MenuProps["onClick"] = ({ key }) => { - if (key === "block" || key === "unblock") onToggleBlocked?.(); - if (key === "reset-spend") onResetSpend?.(); - if (key === "delete") onDelete?.(); - }; + const regenerateButton = ( + + + + ); return (
{onCreateNew && (
-
)}
-
- -
- - + <div className="flex items-start justify-between" style={{ marginBottom: 20 }}> + <div className="min-w-0"> + <div className="flex items-center gap-2"> + <h3 className="m-0 flex items-center gap-1 text-2xl font-semibold"> {data.keyName} - + + {isBlocked && ( - }> + + Blocked - + )} - - - Key ID: {data.keyId} - +
+
+ Key ID: {data.keyId} + +
{canModifyKey && ( - - - - - - - -
- - +
+
- } /> - + } /> +
- + - - } /> +
+ } /> } + icon={} truncate copyable defaultUserIdCheck /> - +
- + - - } /> - } /> - - +
+ } /> + } /> +
+
); } diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 374e36029a0..42d1884e563 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -182,19 +182,6 @@ vi.mock("@heroicons/react/outline", async () => { return { ArrowLeftIcon, TrashIcon, RefreshIcon }; }); -vi.mock("lucide-react", async () => { - const React = await import("react"); - function CopyIcon() { - return React.createElement("span"); - } - (CopyIcon as any).displayName = "CopyIcon"; - function CheckIcon() { - return React.createElement("span"); - } - (CheckIcon as any).displayName = "CheckIcon"; - return { CopyIcon, CheckIcon }; -}); - // Heavy children -> async factories & local React vi.mock("../organisms/RegenerateKeyModal", () => { function RegenerateKeyModal() { diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index caad6fc30fc..cdbd3197f7d 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -961,9 +961,8 @@ describe("KeyEditView", () => { ); const resetBudgetItem = (await screen.findByText("Reset Budget")).closest(".ant-form-item") as HTMLElement; - const clearIcon = resetBudgetItem.querySelector(".ant-select-clear"); - expect(clearIcon).not.toBeNull(); - fireEvent.mouseDown(clearIcon as Element); + await userEvent.click(within(resetBudgetItem).getByRole("combobox")); + await userEvent.click(await screen.findByText("Never resets")); await waitFor(() => { expect(within(resetBudgetItem).getByText("Never resets")).toBeInTheDocument(); @@ -995,7 +994,8 @@ describe("KeyEditView", () => { ); const resetBudgetItem = (await screen.findByText("Reset Budget")).closest(".ant-form-item") as HTMLElement; - fireEvent.mouseDown(resetBudgetItem.querySelector(".ant-select-clear") as Element); + await userEvent.click(within(resetBudgetItem).getByRole("combobox")); + await userEvent.click(await screen.findByText("Never resets")); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); @@ -1251,9 +1251,10 @@ describe("KeyEditView", () => { expect(screen.getByText("Organization")).toBeInTheDocument(); }); - const orgFormItem = screen.getByText("Organization").closest(".ant-form-item"); - const disabledSelect = orgFormItem?.querySelector(".ant-select-disabled"); - expect(disabledSelect).toBeTruthy(); + const orgFormItem = screen.getByText("Organization").closest(".ant-form-item") as HTMLElement; + await userEvent.click(within(orgFormItem).getByRole("combobox")); + + expect(screen.queryByText("Engineering")).not.toBeInTheDocument(); }); it("should not disable the organization dropdown for admin users", async () => { @@ -1273,9 +1274,10 @@ describe("KeyEditView", () => { expect(screen.getByText("Organization")).toBeInTheDocument(); }); - const orgFormItem = screen.getByText("Organization").closest(".ant-form-item"); - const disabledSelect = orgFormItem?.querySelector(".ant-select-disabled"); - expect(disabledSelect).toBeFalsy(); + const orgFormItem = screen.getByText("Organization").closest(".ant-form-item") as HTMLElement; + await userEvent.click(within(orgFormItem).getByRole("combobox")); + + expect(await screen.findByText("Engineering")).toBeInTheDocument(); }); it("should initialize organization from keyData", async () => { @@ -1296,8 +1298,9 @@ describe("KeyEditView", () => { />, ); + const orgFormItem = (await screen.findByText("Organization")).closest(".ant-form-item") as HTMLElement; await waitFor(() => { - expect(screen.getByText("Engineering")).toBeInTheDocument(); + expect(within(orgFormItem).getByRole("combobox")).toHaveValue("Engineering"); }); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 15d14d5abf5..a2d926dff8a 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -4,9 +4,12 @@ import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings" import useTeams from "@/app/(dashboard)/hooks/useTeams"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils"; -import { ArrowLeftIcon } from "@heroicons/react/outline"; -import { Badge, Button, Card, Grid, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; -import { Modal, Tag } from "antd"; +import { ArrowLeft } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Card } from "@/components/ui/card"; +import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { KeyInfoHeader } from "./KeyInfoHeader"; import { useEffect, useState } from "react"; import { isProxyAdminRole, isUserTeamAdminForSingleTeam, rolesWithWriteAccess } from "../../utils/roles"; @@ -150,10 +153,11 @@ export default function KeyInfoView({ if (!currentKeyData) { return (
- - Key not found +

Key not found

); } @@ -534,93 +538,111 @@ export default function KeyInfoView({ /> {/* Reset Spend Confirmation Modal */} - setIsResetSpendModalOpen(false)} - okText="Reset" - okButtonProps={{ danger: true }} - confirmLoading={resetSpendLoading} - > -

- Reset spend for {currentKeyData?.key_alias || currentKeyData?.token_id || "this key"} to{" "} - $0? -

-

- Current spend: ${formatNumberWithCommas(currentKeyData.spend, 4)}. Spend history is preserved - in logs. This resets the current period spend counter, the same as an automatic budget reset. -

- + setIsResetSpendModalOpen(open)}> + + + Reset Key Spend + +

+ Reset spend for {currentKeyData?.key_alias || currentKeyData?.token_id || "this key"} to{" "} + $0? +

+

+ Current spend: ${formatNumberWithCommas(currentKeyData.spend, 4)}. Spend history is + preserved in logs. This resets the current period spend counter, the same as an automatic budget reset. +

+ + + + + + - setIsBlockModalOpen(false)} - okText={isBlocked ? "Unblock" : "Block"} - okButtonProps={isBlocked ? undefined : { danger: true }} - confirmLoading={blockLoading} - > -

- {isBlocked ? "Unblock" : "Block"}{" "} - {currentKeyData?.key_alias || currentKeyData?.token_id || "this key"}? -

-

- {isBlocked - ? "Requests using this key will be accepted again." - : "Requests using this key will be rejected with a 401 error until it is unblocked. The key is not deleted and can be unblocked at any time."} -

- + setIsBlockModalOpen(open)}> + + + {isBlocked ? "Unblock Key" : "Block Key"} + +

+ {isBlocked ? "Unblock" : "Block"}{" "} + {currentKeyData?.key_alias || currentKeyData?.token_id || "this key"}? +

+

+ {isBlocked + ? "Requests using this key will be accepted again." + : "Requests using this key will be rejected with a 401 error until it is unblocked. The key is not deleted and can be unblocked at any time."} +

+ + + + + + - - - Overview - Settings - + + + Overview + Settings + - +
{/* Overview Panel */} - - - - Spend + +
+ +

Spend

- ${formatNumberWithCommas(currentKeyData.spend, 4)} - of {budgetDisplay} +

${formatNumberWithCommas(currentKeyData.spend, 4)}

+

of {budgetDisplay}

{currentKeyData.budget_reset_at && ( - Resets {formatTimestamp(currentKeyData.budget_reset_at)} +

Resets {formatTimestamp(currentKeyData.budget_reset_at)}

)}
- - Rate Limits + +

Rate Limits

- TPM: {currentKeyData.tpm_limit !== null ? currentKeyData.tpm_limit : "Unlimited"} - RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"} +

+ TPM: {currentKeyData.tpm_limit !== null ? currentKeyData.tpm_limit : "Unlimited"} +

+

+ RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"} +

{Boolean(currentKeyData.metadata?.throttle_on_budget_exceeded) && ( - Throttle on budget exceeded: Yes +

Throttle on budget exceeded: Yes

)}
- - Models + +

Models

{currentKeyData.models && currentKeyData.models.length > 0 ? ( currentKeyData.models.map((model, index) => ( - + {model} )) ) : ( - No models specified +

No models specified

)}
- + - - Guardrails + +

Guardrails

{Array.isArray(currentKeyData.metadata?.guardrails) && currentKeyData.metadata.guardrails.length > 0 ? (
{currentKeyData.metadata.guardrails.map((guardrail: string, index: number) => ( - + {guardrail} ))}
) : ( - No guardrails configured +

No guardrails configured

)} {typeof currentKeyData.metadata?.disable_global_guardrails === "boolean" && currentKeyData.metadata.disable_global_guardrails === true && (
- Global Guardrails Disabled + Global Guardrails Disabled
)} - - Policies + +

Policies

{Array.isArray(currentKeyData.metadata?.policies) && currentKeyData.metadata.policies.length > 0 ? (
{currentKeyData.metadata.policies.map((policy: string, index: number) => (
- {policy} - {loadingPolicies && Loading guardrails...} + + {policy} + + {loadingPolicies &&

Loading guardrails...

}
{!loadingPolicies && policyGuardrails[policy] && policyGuardrails[policy].length > 0 && (
- Resolved Guardrails: +

Resolved Guardrails:

{policyGuardrails[policy].map((guardrail: string, gIndex: number) => ( - + {guardrail} ))} @@ -675,7 +699,7 @@ export default function KeyInfoView({ ))}
) : ( - No policies configured +

No policies configured

)} @@ -697,15 +721,19 @@ export default function KeyInfoView({ nextRotationAt={currentKeyData.next_rotation_at} variant="card" /> - - +
+ {/* Settings Panel */} - - + +
- Key Settings - {!isEditing && canModifyKey && } +

Key Settings

+ {!isEditing && canModifyKey && ( + + )}
{isEditing ? ( @@ -722,29 +750,29 @@ export default function KeyInfoView({ ) : (
- Key ID - {currentKeyData.token_id || currentKeyData.token} +

Key ID

+

{currentKeyData.token_id || currentKeyData.token}

- Key Alias - {currentKeyData.key_alias || "Not Set"} +

Key Alias

+

{currentKeyData.key_alias || "Not Set"}

- Secret Key - {currentKeyData.key_name} +

Secret Key

+

{currentKeyData.key_name}

- Team ID - {currentKeyData.team_id || "Not Set"} +

Team ID

+

{currentKeyData.team_id || "Not Set"}

{enableProjectsUI && (
- Project - +

Project

+

{currentKeyData.project_id ? (() => { const project = projects?.find((p) => p.project_id === currentKeyData.project_id); @@ -753,41 +781,43 @@ export default function KeyInfoView({ : currentKeyData.project_id; })() : "Not Set"} - +

)}
- Organization - {(currentKeyData.organization_id ?? currentKeyData.org_id) || "Not Set"} +

Organization

+

{(currentKeyData.organization_id ?? currentKeyData.org_id) || "Not Set"}

- Created - {formatTimestamp(currentKeyData.created_at)} +

Created

+

{formatTimestamp(currentKeyData.created_at)}

{lastRegeneratedAt && (
- Last Regenerated +

Last Regenerated

- {formatTimestamp(lastRegeneratedAt)} - - Recent - +

{formatTimestamp(lastRegeneratedAt)}

+ Recent
)}
- Expires - {currentKeyData.expires ? formatTimestamp(currentKeyData.expires) : "Never"} +

Expires

+

+ {currentKeyData.expires ? formatTimestamp(currentKeyData.expires) : "Never"} +

{Boolean(currentKeyData.metadata?.enable_prompt_caching) && (
- Prompt Caching - Enabled (auto-injects cache_control markers on Anthropic and Bedrock Claude requests) +

Prompt Caching

+

+ Enabled (auto-injects cache_control markers on Anthropic and Bedrock Claude requests) +

)} @@ -802,31 +832,31 @@ export default function KeyInfoView({ />
- Spend - ${formatNumberWithCommas(currentKeyData.spend, 4)} USD +

Spend

+

${formatNumberWithCommas(currentKeyData.spend, 4)} USD

- Budget - +

Budget

+

{currentKeyData.max_budget !== null ? `$${formatNumberWithCommas(currentKeyData.max_budget, 2)}` : "Unlimited"} - +

- Budget Reset - +

Budget Reset

+

{currentKeyData.budget_reset_at ? `${currentKeyData.budget_duration ? `Every ${currentKeyData.budget_duration}, next ` : ""}${formatTimestamp(currentKeyData.budget_reset_at)}` : "Never"} - +

{currentKeyData.budget_fallbacks && Object.keys(currentKeyData.budget_fallbacks).length > 0 && (
- Budget Fallbacks +

Budget Fallbacks

{Object.entries(currentKeyData.budget_fallbacks).map(([model, fallbacks]) => (
@@ -841,7 +871,7 @@ export default function KeyInfoView({ {hasRouterSettings(currentKeyData.router_settings) && (
- Router Settings +

Router Settings

@@ -849,7 +879,7 @@ export default function KeyInfoView({ )}
- Tags +

Tags

{Array.isArray(currentKeyData.metadata?.tags) && currentKeyData.metadata.tags.length > 0 ? currentKeyData.metadata.tags.map((tag, index) => ( @@ -862,8 +892,8 @@ export default function KeyInfoView({
- Prompts - +

Prompts

+

{Array.isArray(currentKeyData.metadata?.prompts) && currentKeyData.metadata.prompts.length > 0 ? currentKeyData.metadata.prompts.map((prompt, index) => ( @@ -871,11 +901,11 @@ export default function KeyInfoView({ )) : "No prompts specified"} - +

- Allowed Routes +

Allowed Routes

{Array.isArray(currentKeyData.allowed_routes) && currentKeyData.allowed_routes.length > 0 ? ( currentKeyData.allowed_routes.map((route, index) => ( @@ -884,14 +914,14 @@ export default function KeyInfoView({ )) ) : ( - All routes allowed + All routes allowed )}
- Allowed Pass Through Routes - +

Allowed Pass Through Routes

+

{Array.isArray(currentKeyData.metadata?.allowed_passthrough_routes) && currentKeyData.metadata.allowed_passthrough_routes.length > 0 ? currentKeyData.metadata.allowed_passthrough_routes.map((route, index) => ( @@ -900,22 +930,22 @@ export default function KeyInfoView({ )) : "No pass through routes specified"} - +

- Disable Global Guardrails - +

Disable Global Guardrails

+

{currentKeyData.metadata?.disable_global_guardrails === true ? ( - Enabled - Global guardrails bypassed + Enabled - Global guardrails bypassed ) : ( - Disabled - Global guardrails active + Disabled - Global guardrails active )} - +

- Models +

Models

{currentKeyData.models && currentKeyData.models.length > 0 ? ( currentKeyData.models.map((model, index) => ( @@ -924,56 +954,60 @@ export default function KeyInfoView({ )) ) : ( - No models specified +

No models specified

)}
- Rate Limits - TPM: {currentKeyData.tpm_limit !== null ? currentKeyData.tpm_limit : "Unlimited"} - RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"} - +

Rate Limits

+

+ TPM: {currentKeyData.tpm_limit !== null ? currentKeyData.tpm_limit : "Unlimited"} +

+

+ RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"} +

+

Max Parallel Requests:{" "} {currentKeyData.max_parallel_requests !== null ? currentKeyData.max_parallel_requests : "Unlimited"} - - +

+

Model TPM Limits:{" "} {currentKeyData.metadata?.model_tpm_limit ? JSON.stringify(currentKeyData.metadata.model_tpm_limit) : "Unlimited"} - - +

+

Model RPM Limits:{" "} {currentKeyData.metadata?.model_rpm_limit ? JSON.stringify(currentKeyData.metadata.model_rpm_limit) : "Unlimited"} - - +

+

Tag RPM Limits:{" "} {currentKeyData.metadata?.tag_rpm_limit && Object.keys(currentKeyData.metadata.tag_rpm_limit).length > 0 ? JSON.stringify(currentKeyData.metadata.tag_rpm_limit) : "Unlimited"} - - +

+

Estimated Output Tokens:{" "} {currentKeyData.metadata?.default_estimated_output_tokens != null ? String(currentKeyData.metadata.default_estimated_output_tokens) : "Default"} - - +

+

Estimated Output Tokens Per Model:{" "} {currentKeyData.metadata?.default_estimated_output_tokens_per_model ? JSON.stringify(currentKeyData.metadata.default_estimated_output_tokens_per_model) : "Default"} - +

- Metadata +

Metadata

                       {formatMetadataForDisplay(stripTagsFromMetadata(currentKeyData.metadata))}
@@ -999,9 +1033,9 @@ export default function KeyInfoView({
)} - - - + +
+
); } diff --git a/ui/litellm-dashboard/src/components/user_agent_activity.tsx b/ui/litellm-dashboard/src/components/user_agent_activity.tsx index a05db3313ab..ca6c2953dfe 100644 --- a/ui/litellm-dashboard/src/components/user_agent_activity.tsx +++ b/ui/litellm-dashboard/src/components/user_agent_activity.tsx @@ -17,7 +17,7 @@ import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip import { BarChart } from "@/components/shared/charts"; import { userAgentSummaryCall, tagDauCall, tagWauCall, tagMauCall, tagDistinctCall } from "./networking"; import PerUserUsage from "./per_user_usage"; -import type { DateRangePickerValue } from "@tremor/react"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import { ChartLoader } from "./shared/chart_loader"; // New interfaces for the updated API response diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 1de232fadb8..ce5337aa7a7 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -1,6 +1,5 @@ "use client"; import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; -import { Col, Grid } from "@tremor/react"; import { jwtDecode } from "jwt-decode"; import React, { useEffect, useState } from "react"; import { fetchTeams } from "./common_components/fetch_teams"; @@ -218,8 +217,8 @@ const UserDashboard: React.FC = ({ return (
- -
+
+
= ({ ) : undefined } /> - - +
+
); }; diff --git a/ui/litellm-dashboard/src/components/view_logs/CostBreakdownViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/CostBreakdownViewer.tsx index 96ac965cf8b..cc2696a85b3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/CostBreakdownViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/CostBreakdownViewer.tsx @@ -1,5 +1,6 @@ -import React from "react"; -import { Collapse } from "antd"; +import React, { useState } from "react"; +import { ChevronDown, ChevronRight } from "lucide-react"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { formatNumberWithCommas } from "@/utils/dataUtils"; export interface CostBreakdown { @@ -49,6 +50,7 @@ export const CostBreakdownViewer: React.FC = ({ cacheReadTokens, cacheCreationTokens, }) => { + const [open, setOpen] = useState(false); const isCached = cacheHit?.toLowerCase() === "true"; const hasTokenCounts = promptTokens !== undefined || completionTokens !== undefined; @@ -90,197 +92,195 @@ export const CostBreakdownViewer: React.FC = ({ return (
- -

Cost Breakdown

-
- Total: - - {formatCost(totalSpend)} - {isCached && " (Cached)"} - -
-
- ), - children: ( -
- {/* Step 1: Base Token Costs */} -
- {(() => { - const hasCacheBreakdown = - costBreakdown?.cache_read_cost !== undefined || costBreakdown?.cache_creation_cost !== undefined; - if (hasCacheBreakdown) { - // Separate line items: Input / Cache Read / Cache Write - const rawCost = isCached - ? 0 - : (inputCost ?? 0) - - (costBreakdown?.cache_read_cost ?? 0) - - (costBreakdown?.cache_creation_cost ?? 0); - return ( - <> -
- Input Cost: - - {formatCost(rawCost)} - {rawInputTokens !== undefined && rawInputTokens !== null && ( - - ({rawInputTokens.toLocaleString()} tokens) - - )} - -
- {(costBreakdown?.cache_read_cost ?? 0) > 0 && ( -
- Prompt Cache Read Cost: - - {formatCost(isCached ? 0 : costBreakdown?.cache_read_cost)} - {(cacheReadTokens ?? 0) > 0 && ( - - ({(cacheReadTokens ?? 0).toLocaleString()} tokens) - - )} - -
- )} - {(costBreakdown?.cache_creation_cost ?? 0) > 0 && ( -
- Prompt Cache Write Cost: - - {formatCost(isCached ? 0 : costBreakdown?.cache_creation_cost)} - {(cacheCreationTokens ?? 0) > 0 && ( - - ({(cacheCreationTokens ?? 0).toLocaleString()} tokens) - - )} - -
- )} - - ); - } - return ( + + + {open ? ( + + ) : ( + + )} +
+

Cost Breakdown

+
+ Total: + + {formatCost(totalSpend)} + {isCached && " (Cached)"} + +
+
+ + +
+ {/* Step 1: Base Token Costs */} +
+ {(() => { + const hasCacheBreakdown = + costBreakdown?.cache_read_cost !== undefined || costBreakdown?.cache_creation_cost !== undefined; + if (hasCacheBreakdown) { + // Separate line items: Input / Cache Read / Cache Write + const rawCost = isCached + ? 0 + : (inputCost ?? 0) - + (costBreakdown?.cache_read_cost ?? 0) - + (costBreakdown?.cache_creation_cost ?? 0); + return ( + <>
Input Cost: - {formatCost(inputCost)} - {promptTokens !== undefined && ( + {formatCost(rawCost)} + {rawInputTokens !== undefined && rawInputTokens !== null && ( - ({promptTokens.toLocaleString()} prompt tokens) + ({rawInputTokens.toLocaleString()} tokens) )}
- ); - })()} + {(costBreakdown?.cache_read_cost ?? 0) > 0 && ( +
+ Prompt Cache Read Cost: + + {formatCost(isCached ? 0 : costBreakdown?.cache_read_cost)} + {(cacheReadTokens ?? 0) > 0 && ( + + ({(cacheReadTokens ?? 0).toLocaleString()} tokens) + + )} + +
+ )} + {(costBreakdown?.cache_creation_cost ?? 0) > 0 && ( +
+ Prompt Cache Write Cost: + + {formatCost(isCached ? 0 : costBreakdown?.cache_creation_cost)} + {(cacheCreationTokens ?? 0) > 0 && ( + + ({(cacheCreationTokens ?? 0).toLocaleString()} tokens) + + )} + +
+ )} + + ); + } + return (
- Output Cost: + Input Cost: - {formatCost(outputCost)} - {completionTokens !== undefined && ( + {formatCost(inputCost)} + {promptTokens !== undefined && ( - ({completionTokens.toLocaleString()} completion tokens) + ({promptTokens.toLocaleString()} prompt tokens) )}
- {costBreakdown?.tool_usage_cost !== undefined && costBreakdown.tool_usage_cost > 0 && ( -
- Tool Usage Cost: - {formatCost(costBreakdown.tool_usage_cost)} -
- )} - {costBreakdown?.additional_costs && - Object.entries(costBreakdown.additional_costs) - .filter(([, value]) => value != null && value !== 0) - .map(([key, value]) => ( -
- {key}: - {formatCost(value)} -
- ))} -
- - {/* Subtotal / Original Cost - hide when cached since it would be $0 */} - {!isCached && ( -
-
- Original LLM Cost: - {formatCost(originalCost)} -
-
- )} - - {/* Step 2: Adjustments (Discount & Margin) */} - {(hasDiscount || hasMargin) && ( -
- {/* Discounts */} - {hasDiscount && ( -
- {costBreakdown.discount_percent !== undefined && costBreakdown.discount_percent !== 0 && ( -
- - Discount ({formatPercent(costBreakdown.discount_percent)}): - - -{formatCost(costBreakdown.discount_amount)} -
- )} - {costBreakdown.discount_amount !== undefined && - costBreakdown.discount_percent === undefined && ( -
- Discount Amount: - -{formatCost(costBreakdown.discount_amount)} -
- )} -
- )} - - {/* Margins */} - {hasMargin && ( -
- {costBreakdown.margin_percent !== undefined && costBreakdown.margin_percent !== 0 && ( -
- - Margin ({formatPercent(costBreakdown.margin_percent)}): - - - + - {formatCost( - (costBreakdown.margin_total_amount || 0) - (costBreakdown.margin_fixed_amount || 0), - )} - -
- )} - {costBreakdown.margin_fixed_amount !== undefined && costBreakdown.margin_fixed_amount !== 0 && ( -
- Margin: - +{formatCost(costBreakdown.margin_fixed_amount)} -
- )} -
- )} -
- )} - - {/* Final Summary */} -
-
- Final Calculated Cost: - - {formatCost(totalCost)} - {isCached && " (Cached)"} + ); + })()} +
+ Output Cost: + + {formatCost(outputCost)} + {completionTokens !== undefined && ( + + ({completionTokens.toLocaleString()} completion tokens) -
+ )} + +
+ {costBreakdown?.tool_usage_cost !== undefined && costBreakdown.tool_usage_cost > 0 && ( +
+ Tool Usage Cost: + {formatCost(costBreakdown.tool_usage_cost)} +
+ )} + {costBreakdown?.additional_costs && + Object.entries(costBreakdown.additional_costs) + .filter(([, value]) => value != null && value !== 0) + .map(([key, value]) => ( +
+ {key}: + {formatCost(value)} +
+ ))} +
+ + {/* Subtotal / Original Cost - hide when cached since it would be $0 */} + {!isCached && ( +
+
+ Original LLM Cost: + {formatCost(originalCost)}
- ), - }, - ]} - /> + )} + + {/* Step 2: Adjustments (Discount & Margin) */} + {(hasDiscount || hasMargin) && ( +
+ {/* Discounts */} + {hasDiscount && ( +
+ {costBreakdown.discount_percent !== undefined && costBreakdown.discount_percent !== 0 && ( +
+ + Discount ({formatPercent(costBreakdown.discount_percent)}): + + -{formatCost(costBreakdown.discount_amount)} +
+ )} + {costBreakdown.discount_amount !== undefined && costBreakdown.discount_percent === undefined && ( +
+ Discount Amount: + -{formatCost(costBreakdown.discount_amount)} +
+ )} +
+ )} + + {/* Margins */} + {hasMargin && ( +
+ {costBreakdown.margin_percent !== undefined && costBreakdown.margin_percent !== 0 && ( +
+ + Margin ({formatPercent(costBreakdown.margin_percent)}): + + + + + {formatCost( + (costBreakdown.margin_total_amount || 0) - (costBreakdown.margin_fixed_amount || 0), + )} + +
+ )} + {costBreakdown.margin_fixed_amount !== undefined && costBreakdown.margin_fixed_amount !== 0 && ( +
+ Margin: + +{formatCost(costBreakdown.margin_fixed_amount)} +
+ )} +
+ )} +
+ )} + + {/* Final Summary */} +
+
+ Final Calculated Cost: + + {formatCost(totalCost)} + {isCached && " (Cached)"} + +
+
+
+ +
); }; diff --git a/ui/litellm-dashboard/src/components/view_logs/EvalViewer/EvalViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/EvalViewer/EvalViewer.tsx index 60b27eb0e3a..217efc7ac27 100644 --- a/ui/litellm-dashboard/src/components/view_logs/EvalViewer/EvalViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/EvalViewer/EvalViewer.tsx @@ -1,8 +1,9 @@ import React from "react"; -import { Card, Tag, Table, Typography, Space, Tooltip } from "antd"; -import { CheckCircleOutlined, CloseCircleOutlined, ExperimentOutlined } from "@ant-design/icons"; - -const { Text } = Typography; +import { CircleCheck, CircleX, FlaskConical } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Card, CardAction, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Table, TableBody, TableCell, TableFooter, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; interface EvalVerdict { criterion_name: string; @@ -36,10 +37,10 @@ export default function EvalViewer({ data }: EvalViewerProps) { return (
- - + + LLM Judge Results - +
{entries.map((entry, idx) => ( @@ -56,151 +57,159 @@ function EvalEntryCard({ entry }: { entry: EvalInformation }) { // Filter out synthetic "Overall" row the judge sometimes appends — it's already in the header const verdicts = (entry.verdicts || []).filter((v) => (v.criterion_name || "").toLowerCase() !== "overall"); - const columns = [ - { - title: "Criterion", - dataIndex: "criterion_name", - key: "criterion_name", - width: 160, - render: (v: string) => ( - - {v} - - ), - }, - { - title: "Weight", - dataIndex: "weight", - key: "weight", - width: 65, - render: (v: number) => - v != null ? ( - - {v}% - - ) : null, - }, - { - title: "Score", - dataIndex: "score", - key: "score", - width: 65, - render: (v: number) => ( - = 70 ? "#52c41a" : v >= 50 ? "#faad14" : "#ff4d4f", fontWeight: 600 }}>{v} - ), - }, - { - title: ( - - Weighted - - ), - key: "weighted", - width: 75, - render: (_: unknown, row: EvalVerdict) => { - if (row.weight == null) return null; - const contrib = (row.score * row.weight) / 100; - return ( - - {contrib % 1 === 0 ? contrib : contrib.toFixed(1)} - - ); - }, - }, - { - title: "Comment", - dataIndex: "reasoning", - key: "reasoning", - ellipsis: { showTitle: false }, - render: (v: string) => ( - - {v} - - ), - }, - ]; + const hasWeights = verdicts.some((v) => v.weight != null); + const weightedTotal = verdicts.reduce((sum, v) => sum + (v.weight != null ? (v.score * v.weight) / 100 : 0), 0); return ( - - {passed ? ( - - ) : ( - - )} - {entry.eval_name} - {passed ? "PASSED" : "FAILED"} - - - {entry.overall_score?.toFixed(0)} / 100 - {entry.threshold != null && ` (threshold: ${entry.threshold})`} - - - - } - extra={ - - {entry.judge_model && ( - - Judge: {entry.judge_model} - - )} - {entry.iteration != null && ( - - Iter: {entry.iteration + 1} - - )} - - } - > - {entry.eval_error && ( - - Judge error: {entry.eval_error} - - )} + + + +
+ {passed ? ( + + ) : ( + + )} + {entry.eval_name} + {passed ? "PASSED" : "FAILED"} + + + + } + > + {entry.overall_score?.toFixed(0)} / 100 + {entry.threshold != null && ` (threshold: ${entry.threshold})`} + + + Weighted average of all criterion scores. Each criterion has a weight (%) set when the eval was + created — higher-weight criteria count more toward the final score. + + + +
+ + +
+ {entry.judge_model && ( + + Judge: {entry.judge_model} + + )} + {entry.iteration != null && ( + + Iter: {entry.iteration + 1} + + )} +
+ + - {verdicts.length > 0 ? ( -
{ - const hasWeights = verdicts.some((v) => v.weight != null); - if (!hasWeights) return null; - const total = verdicts.reduce((sum, v) => sum + (v.weight != null ? (v.score * v.weight) / 100 : 0), 0); - return ( - - - - Total - - - - - - - {total % 1 === 0 ? total : total.toFixed(1)} - - - - - ); - }} - /> - ) : ( - - Score: {entry.overall_score?.toFixed(1)} — no per-criterion breakdown available. - - )} + + {entry.eval_error && ( + + Judge error: {entry.eval_error} + + )} + + {verdicts.length > 0 ? ( +
+ + + Criterion + Weight + Score + + + + }> + Weighted + + + Score × Weight — how much each criterion contributes to the final score + + + + + Comment + + + + {verdicts.map((row) => { + const contrib = row.weight != null ? (row.score * row.weight) / 100 : null; + return ( + + + + {row.criterion_name} + + + + {row.weight != null ? ( + + {row.weight}% + + ) : null} + + + = 70 ? "#52c41a" : row.score >= 50 ? "#faad14" : "#ff4d4f", + fontWeight: 600, + }} + > + {row.score} + + + + {contrib != null ? ( + + {contrib % 1 === 0 ? contrib : contrib.toFixed(1)} + + ) : null} + + + + + }>{row.reasoning} + {row.reasoning} + + + + + ); + })} + + {hasWeights && ( + + + + + Total + + + + + + + {weightedTotal % 1 === 0 ? weightedTotal : weightedTotal.toFixed(1)} + + + + + + )} +
+ ) : ( + + Score: {entry.overall_score?.toFixed(1)} — no per-criterion breakdown available. + + )} + ); } diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/CompliancePanel.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/CompliancePanel.tsx index ebe89f12a7b..1d74927a264 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/CompliancePanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/CompliancePanel.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Tooltip } from "antd"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { checkEuAiActCompliance, checkGdprCompliance, @@ -66,9 +66,12 @@ const ComplianceCard = ({ {loading ? ( ) : error ? ( - - -- - + + + }>-- + {error} + + ) : data?.compliant ? ( ) : ( diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx index 4625f500c6d..32db72a0a25 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx @@ -1,5 +1,5 @@ import React, { useState, useMemo } from "react"; -import { Tooltip } from "antd"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import PresidioDetectedEntities from "./PresidioDetectedEntities"; import BedrockGuardrailDetails, { BedrockGuardrailResponse, @@ -517,13 +517,20 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { )} {riskScore != null && success && ( - - - Risk {riskScore}/10 - - + + + + } + > + Risk {riskScore}/10 + + {`Risk score: ${riskScore}/10`} + + )}
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx index ee619291f66..e80b447f402 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx @@ -1,6 +1,9 @@ -import { Button, Space, Tag, Tooltip, Typography } from "antd"; -import { CloseOutlined, UpOutlined, DownOutlined } from "@ant-design/icons"; +import { useState } from "react"; +import { Check, ChevronDown, ChevronUp, Copy, X } from "lucide-react"; import moment from "moment"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { LogEntry } from "../columns"; import { AutoRouterTag } from "@/components/shared/table_cells"; import { ClassifyTag } from "./ClassifyTag"; @@ -10,15 +13,11 @@ import { COLOR_BORDER, COLOR_BACKGROUND, SPACING_MEDIUM, - SPACING_LARGE, FONT_SIZE_HEADER, FONT_SIZE_MEDIUM, FONT_FAMILY_MONO, - SPACING_SMALL, } from "./constants"; -const { Text } = Typography; - interface DrawerHeaderProps { log: LogEntry; onClose: () => void; @@ -96,7 +95,7 @@ function ModelProviderSection({ providerName?: string; }) { return ( - +
{providerLogo && ( )} - - +
+ {model} - + {providerName && ( - + {providerName} - + )} - - +
+
); } @@ -128,24 +127,50 @@ function ModelProviderSection({ * Request ID display with copy functionality */ function RequestIdSection({ requestId }: { requestId: string }) { + const [copied, setCopied] = useState(false); + + const handleCopy = async () => { + try { + await navigator.clipboard.writeText(requestId); + setCopied(true); + setTimeout(() => setCopied(false), 1200); + } catch { + /* clipboard unavailable in non-secure contexts */ + } + }; + return (
- - - {requestId} - - + + + + } + > + {requestId} + + + {requestId} + +
); } @@ -172,21 +197,29 @@ function NavigationSection({ marginLeft: 4, background: "#fafafa", }; + const splitStyle = { width: 1, height: 20, background: COLOR_BORDER }; return ( - }> - - - - + ); +} + function ErrorDescription({ errorInfo }: { errorInfo: any }) { return (
{errorInfo.error_code && (
- Error Code: {errorInfo.error_code} + Error Code: {errorInfo.error_code}
)} {errorInfo.error_message && (
- Message: {errorInfo.error_message} + Message: {errorInfo.error_message}
)}
@@ -241,16 +299,16 @@ function ErrorDescription({ errorInfo }: { errorInfo: any }) { function TagsSection({ tags }: { tags: Record }) { return (
- + Tags - - + +
{Object.entries(tags).map(([key, value]) => ( - + {key}: {String(value)} - + ))} - +
); } @@ -262,12 +320,12 @@ function GuardrailLabel({ label, maskedCount }: { label: string; maskedCount: nu }; return ( - + {label} - {maskedCount > 0 && {maskedCount} masked} - + {maskedCount > 0 && {maskedCount} masked} + ); } @@ -291,26 +349,24 @@ const PROMPT_CACHE_DOCS_URL = "https://docs.litellm.ai/docs/completion/prompt_ca function MetricLabel({ label, tooltip, docsUrl }: { label: string; tooltip: string; docsUrl: string }) { return ( - + {label} - + + + } + > + + + {tooltip}{" "} - + Docs - - } - > - - - + + + + ); } @@ -333,102 +389,111 @@ function MetricsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: return (
- - - {showAnthropicMessagesInputOutput ? ( - <> - {formatNumberWithCommas(uncachedInputTokens)} - - {formatNumberWithCommas(logEntry.completion_tokens)} - - - ) : ( - - - - )} - ${formatNumberWithCommas(logEntry.spend || 0, 8)} - - {logEntry.request_duration_ms != null ? (logEntry.request_duration_ms / 1000).toFixed(3) : "-"} s - - {ttftMs != null && ttftMs > 0 && ( - {(ttftMs / 1000).toFixed(3)} s - )} - - {showResponseCache && ( - - } - > - {isResponseCacheHit ? "Hit" : "Miss"} - - )} - {promptCacheReadTokens > 0 && ( - - } - > - {formatNumberWithCommas(promptCacheReadTokens)} - - )} - {promptCacheCreationTokens > 0 && ( - - } - > - {formatNumberWithCommas(promptCacheCreationTokens)} - - )} - - {metadata?.litellm_overhead_time_ms !== undefined && metadata.litellm_overhead_time_ms !== null && ( - - {metadata.litellm_overhead_time_ms.toFixed(2)} ms - - )} - - - {metadata?.attempted_retries !== undefined && metadata?.attempted_retries !== null ? ( - metadata.attempted_retries > 0 ? ( - <> - {metadata.attempted_retries} - {metadata.max_retries !== undefined && metadata.max_retries !== null - ? ` / ${metadata.max_retries}` - : ""} - - ) : ( - None - ) + + + Metrics + + + + {showAnthropicMessagesInputOutput ? ( + <> + {formatNumberWithCommas(uncachedInputTokens)} + + {formatNumberWithCommas(logEntry.completion_tokens)} + + ) : ( - "-" + + + + )} + ${formatNumberWithCommas(logEntry.spend || 0, 8)} + + {logEntry.request_duration_ms != null ? (logEntry.request_duration_ms / 1000).toFixed(3) : "-"} s + + {ttftMs != null && ttftMs > 0 && ( + {(ttftMs / 1000).toFixed(3)} s )} - - - {moment(logEntry.startTime).format("YYYY-MM-DDTHH:mm:ss.SSS[Z]")} - - - {moment(logEntry.endTime).format("YYYY-MM-DDTHH:mm:ss.SSS[Z]")} - - + {showResponseCache && ( + + } + > + + {isResponseCacheHit ? "Hit" : "Miss"} + + + )} + {promptCacheReadTokens > 0 && ( + + } + > + {formatNumberWithCommas(promptCacheReadTokens)} + + )} + {promptCacheCreationTokens > 0 && ( + + } + > + {formatNumberWithCommas(promptCacheCreationTokens)} + + )} + + {metadata?.litellm_overhead_time_ms !== undefined && metadata.litellm_overhead_time_ms !== null && ( + + {metadata.litellm_overhead_time_ms.toFixed(2)} ms + + )} + + + {metadata?.attempted_retries !== undefined && metadata?.attempted_retries !== null ? ( + metadata.attempted_retries > 0 ? ( + <> + {metadata.attempted_retries} + {metadata.max_retries !== undefined && metadata.max_retries !== null + ? ` / ${metadata.max_retries}` + : ""} + + ) : ( + + None + + ) + ) : ( + "-" + )} + + + + {moment(logEntry.startTime).format("YYYY-MM-DDTHH:mm:ss.SSS[Z]")} + + + {moment(logEntry.endTime).format("YYYY-MM-DDTHH:mm:ss.SSS[Z]")} + + +
); @@ -449,6 +514,7 @@ function RequestResponseSection({ getFormattedResponse, logEntry, }: RequestResponseSectionProps) { + const [open, setOpen] = useState(true); const [activeTab, setActiveTab] = useState(TAB_REQUEST); const [viewMode, setViewMode] = useState<"pretty" | "json">("pretty"); @@ -476,90 +542,76 @@ function RequestResponseSection({ return (
- { - const target = e.target as HTMLElement; - if (target.closest(".ant-radio-group")) { - e.stopPropagation(); - } - }} - > -

- Request & Response -

- setViewMode(e.target.value)}> - Pretty - JSON - -
- ), - children: ( -
- {viewMode === "pretty" ? ( - - ) : ( - setActiveTab(key as typeof TAB_REQUEST | typeof TAB_RESPONSE)} - tabBarExtraContent={ - - } - items={[ - { - key: TAB_REQUEST, - label: "Request", - children: ( -
- -
- ), - }, - { - key: TAB_RESPONSE, - label: "Response", - children: ( -
- {hasResponse || hasError ? ( - - ) : ( -
- Response data not available -
- )} -
- ), - }, - ]} - /> - )} -
- ), - }, - ]} - /> + + setViewMode(value as "pretty" | "json")}> +
+ + {open ? ( + + ) : ( + + )} +

+ Request & Response +

+ + + Pretty + JSON + +
+ +
+ + + + + setActiveTab(key as typeof TAB_REQUEST | typeof TAB_RESPONSE)} + > +
+ + Request + Response + + +
+ +
+ +
+ + +
+ {hasResponse || hasError ? ( + + ) : ( +
+ Response data not available +
+ )} +
+ + + +
+ + +
); } @@ -602,43 +654,40 @@ export function GuardrailJumpLink({ guardrailEntries }: { guardrailEntries: any[ } function MetadataSection({ metadata }: { metadata: Record }) { + const [open, setOpen] = useState(true); + return (
- Metadata, - children: ( -
-
- -
- 
-                  {JSON.stringify(metadata, null, 2)}
-
-
- ), - }, - ]} - /> + + + {open ? ( + + ) : ( + + )} +

Metadata

+ + +
+
+ JSON.stringify(metadata, null, 2)} label="Copy Metadata" /> +
+ 
+              {JSON.stringify(metadata, null, 2)}
+
+
+ +
); } diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx index 5b27cfc1d0f..83049a12fd2 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx @@ -1,7 +1,8 @@ import { useEffect, useMemo, useState } from "react"; -import { Button, Drawer, Segmented } from "antd"; -import { CheckOutlined, CopyOutlined, LeftOutlined, RightOutlined } from "@ant-design/icons"; -import { Bot, Sparkles, Wrench } from "lucide-react"; +import { Bot, Check, ChevronLeft, ChevronRight, Copy, Sparkles, Wrench } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Sheet, SheetContent, SheetTitle } from "@/components/ui/sheet"; +import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { LogEntry } from "../columns"; import { AutoRouterIcon, useIsAutoRoutedModelGroup } from "@/components/shared/table_cells"; import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "../constants"; @@ -297,181 +298,186 @@ export function LogDetailsDrawer({ if (!currentLog || !enrichedLog) return null; return ( - { + if (!nextOpen) onClose(); }} > -
- {!isSidebarCollapsed ? ( - + + + {logEntry?.request_id ? `Request ${logEntry.request_id} details` : "Request details"} + +
+ {!isSidebarCollapsed ? ( + + ) : ( + + )} + {!isSidebarCollapsed && ( +
+
+
+
+
+ {isSessionMode ? "Session" : "Trace"} +
+
+ {leftPanelDisplayId} + +
-
-
- {logsForList.length} req - {[ - isSessionMode - ? llmCount - : logsForList.filter( - (row) => !MCP_CALL_TYPES.includes(row.call_type) && !AGENT_CALL_TYPES.includes(row.call_type), - ).length, - isSessionMode - ? agentCount - : logsForList.filter((row) => AGENT_CALL_TYPES.includes(row.call_type)).length, - isSessionMode ? mcpCount : logsForList.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length, - ].map((count, i) => { - const label = [" LLM", " Agent", " MCP"][i]; - return count > 0 ? ( - +
+ {logsForList.length} req + {[ + isSessionMode + ? llmCount + : logsForList.filter( + (row) => !MCP_CALL_TYPES.includes(row.call_type) && !AGENT_CALL_TYPES.includes(row.call_type), + ).length, + isSessionMode + ? agentCount + : logsForList.filter((row) => AGENT_CALL_TYPES.includes(row.call_type)).length, + isSessionMode + ? mcpCount + : logsForList.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length, + ].map((count, i) => { + const label = [" LLM", " Agent", " MCP"][i]; + return count > 0 ? ( + + · + {count} + {label} + + ) : null; + })} + · + {isSessionMode ? getSpendString(totalSessionCost) : getSpendString(currentLog.spend || 0)} + {isSessionMode && ( + <> · - {count} - {label} - - ) : null; - })} - · - {isSessionMode ? getSpendString(totalSessionCost) : getSpendString(currentLog.spend || 0)} + {sessionDurationSeconds}s + + )} +
+ {isSessionMode && sessionTruncated && ( +
+ Showing most recent {logsForList.length} of {sessionTotalCount} +
+ )} {isSessionMode && ( - <> - · - {sessionDurationSeconds}s - + setSessionSortMode(value as SessionLogSortMode)} + > + + + Duration + + + Start time + + + )}
- {isSessionMode && sessionTruncated && ( -
- Showing most recent {logsForList.length} of {sessionTotalCount} -
- )} - {isSessionMode && ( - setSessionSortMode(value as SessionLogSortMode)} - /> - )} -
-
- {normalizeGuardrailEntries(metadata?.guardrail_information).length > 0 && ( -
- -
- )} - {isSessionMode ? ( -
- {/* Child events — vertical tree line with horizontal connectors */} -
-
- {logsForList.map((row, idx) => { - const isLast = idx === logsForList.length - 1; - return ( -
-
- {isLast &&
} - { - setSelectedSessionRequestId(row.request_id); - onSelectLog?.(row); - }} - /> -
- ); - })} +
+ {normalizeGuardrailEntries(metadata?.guardrail_information).length > 0 && ( +
+
-
- ) : ( -
- {logsForList.map((row) => ( - onSelectLog?.(row)} - /> - ))} -
- )} + )} + {isSessionMode ? ( +
+ {/* Child events — vertical tree line with horizontal connectors */} +
+
+ {logsForList.map((row, idx) => { + const isLast = idx === logsForList.length - 1; + return ( +
+
+ {isLast &&
} + { + setSelectedSessionRequestId(row.request_id); + onSelectLog?.(row); + }} + /> +
+ ); + })} +
+
+ ) : ( +
+ {logsForList.map((row) => ( + onSelectLog?.(row)} + /> + ))} +
+ )} +
-
- )} + )} -
- -
- + +
+ +
-
- + + ); } diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.test.tsx index 4f8f662aa16..201c9e74ce8 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.test.tsx @@ -4,16 +4,6 @@ import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { RealtimePrettyView, isRealtimeResponse } from "./RealtimePrettyView"; -vi.mock("antd", async () => { - const actual = await vi.importActual("antd"); - return { - ...actual, - message: { - success: vi.fn(), - }, - }; -}); - const sampleRealtimeResponse = { usage: { total_tokens: 587, diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.tsx index 5552c6decbe..5441bcbc4cf 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.tsx @@ -5,19 +5,11 @@ */ import { useState } from "react"; -import { Typography, Tag, Tooltip } from "antd"; -import { - SoundOutlined, - MessageOutlined, - SettingOutlined, - AudioOutlined, - DownOutlined, - UpOutlined, -} from "@ant-design/icons"; +import { ChevronDown, ChevronUp, MessageSquare, Mic, Settings, Volume2 } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { SectionHeader } from "./SectionHeader"; -const { Text } = Typography; - interface RealtimeEvent { type: string; event_id?: string; @@ -163,34 +155,34 @@ function SessionCard({ session, turnCount }: { session: RealtimeSession; turnCou
{isCollapsed ? ( - + ) : ( - + )}
- - Session + + Session
- + {session.model} - + {turnCount > 0 && ( - + {turnCount} {turnCount === 1 ? "turn" : "turns"} - + )} {session.voice && ( - - {session.voice} - + + {session.voice} + )} {session.modalities && (
{session.modalities.map((m) => ( - - {m === "audio" ? : } {m} - + + {m === "audio" ? : } {m} + ))}
)} @@ -228,8 +220,8 @@ function SessionCard({ session, turnCount }: { session: RealtimeSession; turnCou {session.instructions && (
- Instructions - +
- + {response.status || "unknown"} - + {usage && ( - + {usage.input_tokens ?? 0} in / {usage.output_tokens ?? 0} out tokens - + )} {response.conversation_id && ( - - - conv: {response.conversation_id.slice(0, 12)}... - - + + + } + > + conv: {response.conversation_id.slice(0, 12)}... + + {response.conversation_id} + + )}
@@ -381,8 +378,8 @@ function OutputMessage({ output }: { output: RealtimeOutputItem }) { return (
- {output.role?.toUpperCase() || "ASSISTANT"} - + {contents.map((c, cIdx) => { const text = c.transcript || c.text; if (!text) return null; @@ -407,20 +404,18 @@ function OutputMessage({ output }: { output: RealtimeOutputItem }) { }} > {c.type === "audio" && ( - )} {c.type === "text" && ( - - + {label} Token Breakdown - +
{ if (typeof value === "number") { return ( - + {formatTokenLabel(key)}: {value.toLocaleString()} - + ); } return null; @@ -483,9 +481,9 @@ function ConfigRow({ label, value }: { label: string; value: any }) { if (value === undefined || value === null) return null; return (
- + {label} - +
{String(value)}
); diff --git a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/FormattedToolView.tsx b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/FormattedToolView.tsx index c58bde72995..b548aa2f368 100644 --- a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/FormattedToolView.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/FormattedToolView.tsx @@ -2,11 +2,9 @@ * Formatted view of tool definition with parameters table and call data */ -import { Typography, Table } from "antd"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; import { ParsedTool, ParameterRow } from "./types"; -const { Text } = Typography; - interface FormattedToolViewProps { tool: ParsedTool; } @@ -23,57 +21,27 @@ export function FormattedToolView({ tool }: FormattedToolViewProps) { }), ); - const columns = [ - { - title: "Parameter", - dataIndex: "name", - key: "name", - render: (name: string, record: ParameterRow) => ( - - {name} - {record.required && *} - - ), - }, - { - title: "Type", - dataIndex: "type", - key: "type", - render: (type: string) => ( - - {type} - - ), - }, - { - title: "Description", - dataIndex: "description", - key: "description", - render: (desc: string) => {desc}, - }, - ]; - return (
{/* Description */} {tool.description && (
- {tool.description} - +
)} {/* Parameters Table */} {parameterRows.length > 0 && (
- Parameters - - + +
+ + + Parameter + Type + Description + + + + {parameterRows.map((row) => ( + + + + {row.name} + {row.required && *} + + + + {row.type} + + + {row.description} + + + ))} + +
)} {/* If tool was called, show the arguments used */} {tool.called && tool.callData && (
- Called With - +
- - Description - - setViewMode(e.target.value)}> - Formatted - JSON - + Description + setViewMode(value as ViewMode)}> + + Formatted + JSON + +
{viewMode === "formatted" ? : } diff --git a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/ToolItem.tsx b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/ToolItem.tsx index 78525e1a661..112364f5ff3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/ToolItem.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/ToolItem.tsx @@ -3,13 +3,11 @@ */ import { useState } from "react"; -import { Typography, Tag } from "antd"; -import { ToolOutlined, RightOutlined, DownOutlined } from "@ant-design/icons"; +import { ChevronDown, ChevronRight, Wrench } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; import { ParsedTool } from "./types"; import { ToolExpandedContent } from "./ToolExpandedContent"; -const { Text } = Typography; - interface ToolItemProps { tool: ParsedTool; } @@ -39,18 +37,18 @@ export function ToolItem({ tool }: ToolItemProps) { }} >
- - + + {tool.index}. {tool.name} - +
- {tool.called ? "called" : "not called"} + {tool.called ? "called" : "not called"} {expanded ? ( - + ) : ( - + )}
diff --git a/ui/litellm-dashboard/src/components/view_logs/VectorStoreViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/VectorStoreViewer.tsx index bdafc817020..fd46ba34d67 100644 --- a/ui/litellm-dashboard/src/components/view_logs/VectorStoreViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/VectorStoreViewer.tsx @@ -1,5 +1,6 @@ import React, { useState } from "react"; -import { Collapse } from "antd"; +import { ChevronDown, ChevronRight } from "lucide-react"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { getProviderLogoAndName } from "../provider_info_helpers"; interface VectorStoreContent { @@ -31,6 +32,7 @@ interface VectorStoreViewerProps { } export function VectorStoreViewer({ data }: VectorStoreViewerProps) { + const [open, setOpen] = useState(true); const [expandedResults, setExpandedResults] = useState>({}); if (!data || data.length === 0) { @@ -57,110 +59,110 @@ export function VectorStoreViewer({ data }: VectorStoreViewerProps) { return (
- Vector Store Requests, - children: ( -
- {data.map((request, index) => ( -
-
-
-
-
- Query: - {request.query} -
-
- Vector Store ID: - {request.vector_store_id} -
-
- Provider: - - {(() => { - const { logo, displayName } = getProviderLogoAndName(request.custom_llm_provider); - return ( - <> - {logo && {`${displayName}} - {displayName} - - ); - })()} + + + {open ? ( + + ) : ( + + )} +

Vector Store Requests

+ + +
+ {data.map((request, index) => ( +
+
+
+
+
+ Query: + {request.query} +
+
+ Vector Store ID: + {request.vector_store_id} +
+
+ Provider: + + {(() => { + const { logo, displayName } = getProviderLogoAndName(request.custom_llm_provider); + return ( + <> + {logo && {`${displayName}} + {displayName} + + ); + })()} + +
+
+
+
+ Start Time: + {formatTime(request.start_time)} +
+
+ End Time: + {formatTime(request.end_time)} +
+
+ Duration: + {calculateDuration(request.start_time, request.end_time)} +
+
+
+
+ +

Search Results

+
+ {request.vector_store_search_response.data.map((result, resultIndex) => { + const isExpanded = expandedResults[`${index}-${resultIndex}`] || false; + + return ( +
+
toggleResult(index, resultIndex)} + > + + + +
+ Result {resultIndex + 1} + + Score: {result.score.toFixed(4)}
-
-
- Start Time: - {formatTime(request.start_time)} + + {isExpanded && ( +
+ {result.content.map((content, contentIndex) => ( +
+
{content.type}
+ 
+                                  {content.text}
+
+
+ ))}
-
- End Time: - {formatTime(request.end_time)} -
-
- Duration: - {calculateDuration(request.start_time, request.end_time)} -
-
+ )}
-
- -

Search Results

-
- {request.vector_store_search_response.data.map((result, resultIndex) => { - const isExpanded = expandedResults[`${index}-${resultIndex}`] || false; - - return ( -
-
toggleResult(index, resultIndex)} - > - - - -
- Result {resultIndex + 1} - - Score: {result.score.toFixed(4)} - -
-
- - {isExpanded && ( -
- {result.content.map((content, contentIndex) => ( -
-
{content.type}
- 
-                                      {content.text}
-
-
- ))} -
- )} -
- ); - })} -
-
- ))} + ); + })} +
- ), - }, - ]} - /> + ))} +
+ +
); } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b6dbe0dfc6d..1c5fe408203 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23675,6 +23675,26 @@ export interface components { * @description Maximum retention period for auto-router benchmark session rollup rows (e.g., '365d'). Rows whose last turn is older than this are deleted by the spend log cleanup job, on that job's schedule. Unset means rollup rows are never deleted. */ maximum_autorouter_session_retention_period?: string | null; + /** + * Maximum Spend Logs Cleanup Batch Size + * @description Rows deleted per DELETE statement by the spend log cleanup job. Defaults to 1000. + */ + maximum_spend_logs_cleanup_batch_size?: number | null; + /** + * Maximum Spend Logs Cleanup Batch Timeout + * @description Postgres statement_timeout and lock_timeout applied to each spend log cleanup delete batch (e.g. '30s'), so cleanup cannot hold row locks or a connection indefinitely. Defaults to '30s'. + */ + maximum_spend_logs_cleanup_batch_timeout?: string | null; + /** + * Maximum Spend Logs Cleanup Max Batches + * @description Maximum DELETE statements the spend log cleanup job issues per table per run. Defaults to 500. + */ + maximum_spend_logs_cleanup_max_batches?: number | null; + /** + * Maximum Spend Logs Cleanup Run Budget + * @description Wall-clock budget for one spend log cleanup run (e.g. '5m'), shared across every table it prunes. A run that hits the budget stops and the next run resumes from where it left off. Defaults to '5m'. + */ + maximum_spend_logs_cleanup_run_budget?: string | null; /** * Maximum Spend Logs Retention Period * @description Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted. diff --git a/ui/litellm-dashboard/src/utils/ptuValidation.test.ts b/ui/litellm-dashboard/src/utils/ptuValidation.test.ts index 1338c2a39ae..ed30e6a7b20 100644 --- a/ui/litellm-dashboard/src/utils/ptuValidation.test.ts +++ b/ui/litellm-dashboard/src/utils/ptuValidation.test.ts @@ -3,6 +3,7 @@ import { PTU_COUNT_FIELD, PTU_RATE_FIELD, ptuCountRules, + ptuNoUsageCostRule, ptuPairRule, ptuRateRules, ptuStartRequiredRule, @@ -40,6 +41,39 @@ describe("ptuCountRules", () => { }); }); +describe("ptuNoUsageCostRule", () => { + const check = (value: unknown, count: unknown) => + ptuNoUsageCostRule(PTU_COUNT_FIELD)({ getFieldValue: () => count }).validator(null, value); + + it("leaves a deployment without PTU config free to carry any price", async () => { + await expect(check("2.5", "")).resolves.toBeUndefined(); + await expect(check(0.5, null)).resolves.toBeUndefined(); + }); + + it("accepts a blank or zero price alongside PTU config, which is what the backend stores", async () => { + await expect(check("", 15)).resolves.toBeUndefined(); + await expect(check(null, 15)).resolves.toBeUndefined(); + await expect(check(0, 15)).resolves.toBeUndefined(); + await expect(check("0", 15)).resolves.toBeUndefined(); + }); + + it("rejects a non-zero price alongside PTU config, which the backend answers with a 400", async () => { + await expect(check("2.5", 15)).rejects.toThrow("bills by reserved capacity"); + await expect(check(0.000001, 15)).rejects.toThrow("bills by reserved capacity"); + }); + + it("reads the count by the field name it was given", () => { + const seen: string[] = []; + ptuNoUsageCostRule(PTU_COUNT_FIELD)({ + getFieldValue: (name: string) => { + seen.push(name); + return 15; + }, + }).validator(null, 0); + expect(seen).toEqual([PTU_COUNT_FIELD]); + }); +}); + describe("ptuPairRule", () => { const rule = (sibling: unknown) => ptuPairRule(PTU_RATE_FIELD)({ getFieldValue: () => sibling }); const check = (value: unknown, sibling: unknown) => rule(sibling).validator(null, value); diff --git a/ui/litellm-dashboard/src/utils/ptuValidation.ts b/ui/litellm-dashboard/src/utils/ptuValidation.ts index 9fabf39ae23..da367f233b6 100644 --- a/ui/litellm-dashboard/src/utils/ptuValidation.ts +++ b/ui/litellm-dashboard/src/utils/ptuValidation.ts @@ -4,6 +4,7 @@ interface ValidatorRule { interface FormInstance { getFieldValue: (name: string) => unknown; + isFieldTouched?: (name: string) => boolean; } export const PTU_COUNT_FIELD = "ptu_count"; @@ -71,6 +72,25 @@ export const ptuPairRule = : Promise.reject(new Error("PTU Count and Cost per PTU / Hour must be set together")), }); +/** + * A PTU deployment is billed by the flat cost of its reserved capacity, so the backend refuses + * a non-zero per-token price alongside PTU config and stores 0 when none is given. Pair this + * with `dependencies` on the count so the error clears once the price or the PTU config goes. + */ +export const ptuNoUsageCostRule = + (countField: string, thisField?: string) => + ({ getFieldValue, isFieldTouched }: FormInstance): ValidatorRule => ({ + validator: (_, value) => { + // A cost the operator never typed was seeded from the rate /model/info resolved, which + // for an unpriced deployment is the public cost map. Refusing it would block every + // attempt to put an existing deployment on PTU, and the save omits it anyway. + const echoed = thisField !== undefined && isFieldTouched !== undefined && !isFieldTouched(thisField); + return echoed || !isFilled(getFieldValue(countField)) || !isFilled(value) || Number(value) === 0 + ? Promise.resolve() + : Promise.reject(new Error("A PTU deployment bills by reserved capacity, so this cost must be 0 or blank")); + }, + }); + /** * The backend requires an effective start whenever PTU is configured, since flat cost * accrues from that instant and an inferred start would bill days a deployment did not