diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 9363db12385..f2bd18cd046 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -20,6 +20,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import ( OTELSemconvCategory, parse_semconv_opt_in, ) +from litellm.integrations.otel.model.db_endpoint import db_span_attributes from litellm.integrations.otel.model.semconv import Metric from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.secret_redaction import redact_string @@ -718,6 +719,28 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): self._handle_failure(kwargs, response_obj, start_time, end_time) + def _start_service_span(self, payload: ServiceLoggerPayload, parent_otel_span: Span, start_time_ns: int) -> Span: + """Open a service span, named and classified by what the service is. + + A datastore call is an outbound CLIENT span carrying ``db.*`` semconv. + Without those a Postgres span says only ``service=postgres``, so the + backend falls back to the transport peer, which for Prisma is the local + query engine on loopback. Everything else stays an INTERNAL span. + """ + from opentelemetry import trace + from opentelemetry.trace import SpanKind + + attributes: Final = db_span_attributes(payload.service.value, payload.call_type) + span: Final = self.tracer.start_span( + name=payload.service, + context=trace.set_span_in_context(parent_otel_span), + start_time=start_time_ns, + kind=SpanKind.CLIENT if attributes else SpanKind.INTERNAL, + ) + for key, value in attributes.items(): + self.safe_set_attribute(span=span, key=key, value=value) + return span + async def async_service_success_hook( self, payload: ServiceLoggerPayload, @@ -726,7 +749,6 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): end_time: datetime | float | None = None, event_metadata: dict | None = None, ): - from opentelemetry import trace from opentelemetry.trace import Status, StatusCode _start_time_ns = 0 @@ -743,12 +765,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): _end_time_ns = self._to_ns(end_time) if parent_otel_span is not None: - _span_name: Final = payload.service - service_logging_span: Final = self.tracer.start_span( - name=_span_name, - context=trace.set_span_in_context(parent_otel_span), - start_time=_start_time_ns, - ) + service_logging_span: Final = self._start_service_span(payload, parent_otel_span, _start_time_ns) self.safe_set_attribute( span=service_logging_span, key="call_type", @@ -786,7 +803,6 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): end_time: float | datetime | None = None, event_metadata: dict | None = None, ): - from opentelemetry import trace from opentelemetry.trace import Status, StatusCode _start_time_ns = 0 @@ -803,12 +819,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): _end_time_ns = self._to_ns(end_time) if parent_otel_span is not None: - _span_name: Final = payload.service - service_logging_span: Final = self.tracer.start_span( - name=_span_name, - context=trace.set_span_in_context(parent_otel_span), - start_time=_start_time_ns, - ) + service_logging_span: Final = self._start_service_span(payload, parent_otel_span, _start_time_ns) self.safe_set_attribute( span=service_logging_span, key="call_type", diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 032441535e0..79487e69ac4 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -18,6 +18,7 @@ from litellm.integrations.otel.mappers.utils import ( serialize_messages, tool_definition_attrs, ) +from litellm.integrations.otel.model.db_endpoint import db_span_attributes from litellm.integrations.otel.model.payloads import ( GuardrailSpanData, LLMCallSpanData, @@ -27,7 +28,6 @@ from litellm.integrations.otel.model.payloads import ( ToolDefinition, ) from litellm.integrations.otel.model.semconv import ( - DB, MCP, Error, GenAI, @@ -36,7 +36,6 @@ from litellm.integrations.otel.model.semconv import ( RpcSystem, Server, ) -from litellm.integrations.otel.model.spans import db_system class GenAIMapper: @@ -182,12 +181,8 @@ class GenAIMapper: def _service(cls, data: ServiceSpanData) -> AttributeMap: attrs: Final = collect(cls._SERVICE_ATTRS, data) # An outbound datastore call (DB_CALL / CLIENT span) also carries db.* - # semconv. Internal services (router, budget jobs, …) have no db.system, - # so they get only the litellm.service.* keys above. - system: Final = db_system(data.service_name) - if system is not None: - attrs[DB.SYSTEM_NAME] = system - if data.call_type: - attrs[DB.OPERATION_NAME] = data.call_type + # semconv naming the server it reached. Internal services (router, budget + # jobs, …) have no db.system, so they get only the litellm.service.* keys. + attrs.update(db_span_attributes(data.service_name, data.call_type)) attrs.update({f"{LiteLLM.METADATA_PREFIX}{key}": value for key, value in data.event_metadata.items()}) return attrs diff --git a/litellm/integrations/otel/model/db_endpoint.py b/litellm/integrations/otel/model/db_endpoint.py new file mode 100644 index 00000000000..562162a8f31 --- /dev/null +++ b/litellm/integrations/otel/model/db_endpoint.py @@ -0,0 +1,164 @@ +"""OTel ``db.*`` / ``server.*`` attributes naming the database litellm talks to. + +Prisma reaches PostgreSQL through a query engine listening on loopback, so +transport-level instrumentation attributes the work to ``localhost`` and an +operator cannot tell it is a PostgreSQL call or correlate it with the database's +own metrics. These attributes name the real server on litellm's DB spans. + +Only the host, port, database and schema of the DSN are read, so no credential +can reach an exporter. +""" + +from __future__ import annotations + +import os +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final +from urllib.parse import ParseResult, parse_qs, unquote, urlparse + +from litellm.integrations.otel.model.semconv import DB, Server +from litellm.integrations.otel.model.spans import POSTGRESQL, db_system + +_DATABASE_URL_ENV: Final = "DATABASE_URL" +_READ_REPLICA_ENV: Final = "DATABASE_URL_READ_REPLICA" +_DEFAULT_POSTGRES_PORT: Final = 5432 +_DEFAULT_POSTGRES_SCHEMA: Final = "public" +_POSTGRES_SCHEMES: Final = frozenset({"postgres", "postgresql"}) +_EMPTY_ATTRIBUTES: Final[Mapping[str, str | int]] = MappingProxyType({}) + + +@dataclass(frozen=True, slots=True) +class DatabaseEndpoint: + """The non-sensitive identity of a PostgreSQL server, parsed from a DSN.""" + + address: str | None + port: int | None + namespace: str | None + + +def parse_database_endpoint(url: str | None) -> DatabaseEndpoint | None: + """Parse a PostgreSQL DSN into its exportable endpoint identity. + + Returns ``None`` for an absent, malformed or non-PostgreSQL URL rather than + raising: an unparseable DSN must degrade to a span without endpoint + attributes, never break the request that emitted it. + """ + if not url: + return None + try: + parsed: Final = urlparse(url) + if parsed.scheme not in _POSTGRES_SCHEMES: + return None + query: Final = parse_qs(parsed.query) + raw_database: Final = (parsed.path or "").lstrip("/") + if _is_misparsed_authority(parsed, url, raw_database): + return None + # ``host=`` beats the netloc: it is how libpq names a Unix socket + # directory and how the Cloud SQL connector sits behind a localhost + # netloc, where the netloc is the very answer this module replaces. + address: Final = _first(query.get("host")) or parsed.hostname + # ``port=`` accompanies ``host=`` in a libpq URI, so honour it the same way. + port: Final = _port(_first(query.get("port")), parsed.port) if address else None + namespace: Final = _namespace(unquote(raw_database), _first(query.get("schema"))) + except ValueError: + return None + if address is None and namespace is None: + return None + return DatabaseEndpoint(address=address, port=port, namespace=namespace) + + +def _is_misparsed_authority(parsed: ParseResult, url: str, raw_database: str) -> bool: + """Whether the URL authority may have been truncated by an unencoded character. + + ``/``, ``#`` or ``?`` in a password ends the netloc early, so urlparse hands + back the username as the host, the leading digits of the password as the + port, and the rest of the credential as the path, query or fragment. The + stranded userinfo ``@`` is the only surviving evidence. + + A database name cannot hold an unencoded slash either, so a second path + segment is the same evidence. + + A DSN that carries the at-sign in a query parameter instead, such as + ``?application_name=svc@prod``, is indistinguishable from a mis-split by any + property of the parse: both leave no userinfo, a host, a port and a path. + Since guessing wrong publishes a credential fragment to a tracing backend, + that ambiguity resolves to refusing the endpoint. Such a DSN loses + ``server.address`` and ``db.namespace`` and keeps the rest of the span, + which is the cheaper error of the two. Percent-encode the at-sign to keep + them. + """ + if "/" in raw_database: + return True + return "@" in url and "@" not in parsed.netloc + + +def _first(values: Sequence[str] | None) -> str: + return values[0] if values else "" + + +def _port(from_query: str, from_netloc: int | None) -> int: + return int(from_query) if from_query.isdigit() else (from_netloc or _DEFAULT_POSTGRES_PORT) + + +def _namespace(database: str, schema: str) -> str | None: + """``{database}|{schema}`` per the PostgreSQL semconv, dropping absent halves. + + Only Prisma's literal default schema stays implicit. The match is + case-sensitive because Prisma quotes the name, so ``?schema=PUBLIC`` builds + a second schema alongside ``public`` and the two must not collapse to one + namespace. + """ + qualifier: Final = "" if schema == _DEFAULT_POSTGRES_SCHEMA else schema + return "|".join(part for part in (database, qualifier) if part) or None + + +def postgres_endpoint() -> DatabaseEndpoint | None: + """The PostgreSQL endpoint the process is currently connected to. + + Read from ``os.environ`` on every span, deliberately, on both counts. + + The environment is what Prisma itself connects with, so the span cannot + disagree with the connection; ``get_secret_str`` would consult a configured + secret manager first and could name a different server than the one serving + the query. And the value is not static: the RDS IAM refresh rebuilds the URL + from ``DATABASE_HOST``/``PORT``/``NAME``/``SCHEMA`` every rotation, the + reconnect path re-reads ``DATABASE_URL``, and the DB-backed + ``environment_variables`` config overlay can rewrite any of them after + startup, so a value cached for the process lifetime goes stale against a + connection that has genuinely moved. Nothing is memoized either: a cache + keyed on the URL would hold a rotated credential past its rotation, and the + parse is a single ``urlparse`` on a short string. + + A configured read replica yields ``None``: ``RoutingPrismaWrapper`` picks + reader or writer per Prisma call, underneath the span, so naming the writer + would attribute replica reads to the primary. + """ + if os.environ.get(_READ_REPLICA_ENV): + return None + return parse_database_endpoint(os.environ.get(_DATABASE_URL_ENV, "")) + + +def db_span_attributes(service_name: str, call_type: str | None = None) -> Mapping[str, str | int]: + """The ``db.*``/``server.*`` attributes for a datastore service call. + + Empty for services that are not outbound datastore calls. Endpoint + attributes are PostgreSQL-only: ``DATABASE_URL`` says nothing about where + the redis-backed services point. ``db.system`` rides alongside the current + ``db.system.name`` because Datadog's OTLP intake still types a database span + from the older key. + """ + system: Final = db_system(service_name) + if system is None: + return _EMPTY_ATTRIBUTES + endpoint: Final = postgres_endpoint() if system == POSTGRESQL else None + pairs: Final[tuple[tuple[str, str | int | None], ...]] = ( + (DB.SYSTEM_NAME, system), + (DB.SYSTEM_LEGACY, system), + (DB.OPERATION_NAME, call_type), + (Server.ADDRESS, endpoint.address if endpoint is not None else None), + (Server.PORT, endpoint.port if endpoint is not None else None), + (DB.NAMESPACE, endpoint.namespace if endpoint is not None else None), + ) + return MappingProxyType({key: value for key, value in pairs if value}) diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index 3d585c36b67..ada2822ba66 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -238,7 +238,11 @@ class DB: """ SYSTEM_NAME: Final = "db.system.name" + # Superseded by SYSTEM_NAME, dual-emitted because Datadog's OTLP intake + # still infers a span's database type from this key. + SYSTEM_LEGACY: Final = "db.system" OPERATION_NAME: Final = "db.operation.name" + NAMESPACE: Final = "db.namespace" class HTTP: diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 0f67f0e7a7c..08318f78b7c 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -115,10 +115,12 @@ SPAN_REGISTRY: Final[dict[SpanRole, SpanSpec]] = { # redis-backed spend queues. Any service not mapped here is litellm-internal work # and stays an INTERNAL ``SERVICE`` span. This table is the single source of # datastore knowledge — both the role classifier and the mapper read it. +POSTGRESQL: Final = "postgresql" + _DB_SYSTEM_BY_SERVICE: Final[dict[str, str]] = { "redis": "redis", - "postgres": "postgresql", - "batch_write_to_db": "postgresql", + "postgres": POSTGRESQL, + "batch_write_to_db": POSTGRESQL, } diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 07d5e6314dd..e4c1c9fc5cf 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1816,16 +1816,19 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]: This helper uses ``json.JSONDecoder.raw_decode()`` to walk the string and extract each JSON object individually. + The walk degrades gracefully: if the string is malformed or truncated + (e.g. a stream that ended mid-tool-call), whatever complete objects were + parsed before the bad tail are returned and the remainder is discarded + with a warning, rather than raising. The sole caller + (``_convert_to_bedrock_tool_call_invoke``) treats an empty result as + ``input={}`` so the conversation can continue instead of hard-failing. + Returns ------- list[dict] A list of parsed dicts – one per JSON object found. If *raw* is - empty or whitespace-only, an empty list is returned. - - Raises - ------ - json.JSONDecodeError - If the string contains text that cannot be parsed as JSON at all. + empty, whitespace-only, or wholly unparseable, an empty list is + returned. """ import json @@ -1845,7 +1848,17 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]: if idx >= length: break - obj, end_idx = decoder.raw_decode(raw, idx) + try: + obj, end_idx = decoder.raw_decode(raw, idx) + except json.JSONDecodeError as e: + verbose_logger.warning( + "split_concatenated_json_objects: discarding unparseable tool-call " + "arguments tail after %d complete object(s); decode_start=%d error=%s", + len(results), + idx, + e, + ) + break if isinstance(obj, dict): results.append(obj) else: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 2ffe015c727..0ed15c43ccf 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3712,7 +3712,13 @@ def _convert_to_bedrock_tool_call_invoke( _parts_list.append(cache_point_block) return _parts_list except Exception as e: - raise Exception(f"Unable to convert openai tool calls={tool_calls} to bedrock tool calls. Received error={e}") + tool_call_ids: Final = tuple(tool.get("id") for tool in tool_calls if isinstance(tool, dict)) + raise litellm.BadRequestError( + message=f"Unable to convert openai tool calls with ids={tool_call_ids} to bedrock tool calls. " + f"Received error={e}", + model=model or "", + llm_provider="bedrock", + ) from e def _append_bedrock_tool_result_media_block( diff --git a/litellm/llms/tinyfish/search/transformation.py b/litellm/llms/tinyfish/search/transformation.py index ba9ca2e1bde..b688dc2cd01 100644 --- a/litellm/llms/tinyfish/search/transformation.py +++ b/litellm/llms/tinyfish/search/transformation.py @@ -14,6 +14,7 @@ import httpx from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( @@ -22,7 +23,7 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str -_UrlEncodableParams: Final = TypeAdapter(dict[str, str | int | bool]) +_UrlEncodableParams: Final = TypeAdapter(dict[str, str | int | float | bool]) _StrList: Final = TypeAdapter(list[str]) _StrFrozenSet: Final = TypeAdapter(frozenset[str]) @@ -94,16 +95,16 @@ class TinyfishSearchConfig(BaseSearchConfig): TinyFish equivalents: - ``query`` (str or list[str]) → ``query`` (list joined by spaces) - ``country`` → ``location`` - - ``search_domain_filter`` (list[str]) → folded into the query as - ``() (site:a OR site:b ...)`` (TinyFish has no first-class - field today; see ML-2084 for the planned ``include_domains``) + - ``search_domain_filter`` (list[str]) → folded into the query using + search operators - ``max_results`` → not sent on the wire; stashed on ``self._caller_max_results`` for client-side response truncation (TinyFish doesn't honor it server-side) - ``max_tokens_per_page`` → silently dropped (no TinyFish equivalent) Any other ``optional_params`` keys are forwarded to TinyFish as-is. - dict/list values are JSON-encoded so they survive ``urlencode``. + dict and list values are JSON-encoded so structured payloads survive + ``urlencode``. Returns: ``{_TINYFISH_PARAMS_KEY: }``. @@ -144,14 +145,12 @@ class TinyfishSearchConfig(BaseSearchConfig): supported_perplexity: Final = _StrFrozenSet.validate_python(raw_supported) for param, value in optional_params.items(): if param not in supported_perplexity and param not in request_data: - # `fetch` expects a JSON-encoded object on the wire; accept the - # natural Python dict form and serialize here so callers don't - # have to pre-stringify. - if isinstance(value, dict): + # Serialize dicts/lists as JSON so structured params survive urlencode. + if isinstance(value, (dict, list)): value = json.dumps(value, separators=(",", ":")) # `urlencode` would render Python bool as "True"/"False" - # (capitalized). ux-labs validators require lowercase - # "true"/"false" (e.g. `include_thumbnail`); normalize here. + # (capitalized). TinyFish Search's bool params require lowercase + # "true"/"false" strings on the wire; normalize here. elif isinstance(value, bool): value = "true" if value else "false" request_data[param] = value @@ -167,17 +166,35 @@ class TinyfishSearchConfig(BaseSearchConfig): """ Transform a TinyFish response to LiteLLM's unified ``SearchResponse``. - Mappings (per-result): - - ``title`` → ``SearchResult.title`` (defaults to ``""`` if missing/null) - - ``url`` → ``SearchResult.url`` (defaults to ``""``) - - ``snippet`` → ``SearchResult.snippet`` (defaults to ``""``) - - all other per-result fields (``position``, ``site_name``, - ``thumbnail_url``, ``fetch``, ``fetch_error``, ...) ride through as - extras on ``SearchResult`` via its ``extra="allow"`` config. + Per-result field handling: + - ``title``, ``url``, ``snippet`` are declared on ``SearchResult`` and + populated by ``SearchResponse.model_validate`` when present. Missing + or ``None`` values are defaulted to ``""`` beforehand by + ``_default_missing_result_fields`` so a degraded result flows through + instead of failing the whole call. + - All undeclared per-result fields (``position``, ``site_name``, and + any others TinyFish returns) ride through as extras via + ``SearchResult``'s ``extra="allow"`` config — accessible as + attributes on the result object or enumerable via + ``result.model_extra``. - Top-level ``parameter_warnings`` (see ML-2085) is read when present and - each entry is re-fired via ``verbose_logger.warning``. Absent or - malformed entries are silently skipped — never throws. + Top-level ``parameter_warnings`` is read when present and each entry + is re-fired via ``verbose_logger.warning``. Absent or malformed + entries are silently skipped — never throws. + + Top-level extras (``query``, ``total_results``, ``page``, and any + future TinyFish additions) ride through via + ``SearchResponse.extra="allow"``. The validated response is returned + in place after truncating ``results`` to the caller's ``max_results``, + so every field pydantic populated survives regardless of which + storage bucket (declared attribute or ``__pydantic_extra__``) holds it. + + TinyFish response headers (e.g. ``x-request-id``, ``retry-after``, + ``x-ratelimit-limit`` — httpx normalizes header names to lowercase) + are stashed on ``response._hidden_params["headers"]`` (raw) and + ``response._hidden_params["additional_headers"]`` (sanitized via + ``process_response_headers``) so callers can correlate a search with + server-side logs. Error paths routed through ``self._wrap_error`` for uniform ``"TinyFish Search: . See for details."`` wrapping: @@ -223,7 +240,12 @@ class TinyfishSearchConfig(BaseSearchConfig): _emit_parameter_warnings(parsed) max_results: Final = self._caller_max_results or _TINYFISH_RESULT_CAP - return SearchResponse(results=list(parsed.results[:max_results])) + parsed.results = parsed.results[:max_results] + raw_headers: Final = dict(raw_response.headers) + hidden: Final = parsed._hidden_params # pyright: ignore[reportPrivateUsage] # sole hidden-params channel + hidden["headers"] = raw_headers + hidden["additional_headers"] = process_response_headers(raw_headers) + return parsed def _wrap_error( self, @@ -243,9 +265,9 @@ class TinyfishSearchConfig(BaseSearchConfig): carry the ``TinyFish Search:`` prefix — the bare error already names the host in the URL, so attribution is implicit there. """ - # ux-labs frontend wraps every error body as {"error": {"code", "message", "details"?}}. + # TinyFish Search wraps every error body as {"error": {"code", "message", "details"?}}. # Best-effort unwrap to surface the inner message; fall back to the raw body - # for non-ux-labs responses (CDN HTML pages, other JSON envelopes, plain text). + # for other envelope shapes (CDN HTML pages, other JSON envelopes, plain text). inner_message = error_message try: body: Final[object] = json.loads(error_message) # any-ok: json.loads -> Any @@ -290,7 +312,7 @@ def _default_missing_result_fields(raw_json: object) -> None: def _emit_parameter_warnings(parsed: SearchResponse) -> None: - """Re-fire TinyFish-side ``parameter_warnings`` (see ML-2085) as warnings. + """Re-fire TinyFish-side ``parameter_warnings`` as warnings. Defensive: skip silently on any shape we don't recognize so a malformed entry (or an early/partial rollout of the field) never throws. diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 8edb3b42dce..781fe264eb8 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -594,13 +594,16 @@ def _build_aggregated_where_clause( sql_params.append(adjusted_end) p += 1 - # Optional entity filter + # Optional entity filter; an empty list must match nothing, not everything if entity_id is not None: if isinstance(entity_id, list): - placeholders = ", ".join(f"${p + i}" for i in range(len(entity_id))) - sql_conditions.append(f'"{entity_id_field}" IN ({placeholders})') - sql_params.extend(entity_id) - p += len(entity_id) + if entity_id: + placeholders = ", ".join(f"${p + i}" for i in range(len(entity_id))) + sql_conditions.append(f'"{entity_id_field}" IN ({placeholders})') + sql_params.extend(entity_id) + p += len(entity_id) + else: + sql_conditions.append("FALSE") else: sql_conditions.append(f'"{entity_id_field}" = ${p}') sql_params.append(entity_id) 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 1c8bce28454..5f6489a69ca 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 @@ -229,6 +229,25 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): verbose_proxy_logger.warning("Error calculating image editing cost: %s", e) return 0.0 + @staticmethod + def _calculate_embeddings_cost( + litellm_model_response: EmbeddingResponse, + model: str, + custom_llm_provider: str, + ) -> float: + try: + return litellm.completion_cost( + completion_response=litellm_model_response, + model=model, + custom_llm_provider=custom_llm_provider, + call_type="aembedding", + ) + except Exception as e: # noqa: BLE001 # completion_cost raises bare Exception for unmapped models; cost failure must never drop the spend log + verbose_proxy_logger.warning( + "Error calculating embeddings cost for model %s, logging spend with cost 0: %s", model, e + ) + return 0.0 + @staticmethod def _build_responses_api_response_and_cost( model: str, @@ -351,11 +370,10 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): model_response_object=EmbeddingResponse(), response_type="embedding", ) - response_cost = litellm.completion_cost( - completion_response=litellm_model_response, + response_cost = OpenAIPassthroughLoggingHandler._calculate_embeddings_cost( + litellm_model_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: @@ -471,6 +489,12 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): except Exception as e: verbose_proxy_logger.error("Error in OpenAI passthrough cost tracking: %s", e) + if not is_chat_completions: + unbilled_result: Final[PassThroughEndpointLoggingTypedDict] = { + "result": None, + "kwargs": kwargs, + } + return unbilled_result # Fall back to base handler without cost tracking base_handler = OpenAIPassthroughLoggingHandler() return base_handler.passthrough_chat_handler( diff --git a/litellm/router.py b/litellm/router.py index c5881960c80..efd3b5a527e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -11463,8 +11463,10 @@ class Router: deployment the strategy was registered from via its (model_name, tags) pair. - With tag filtering enabled, strategies that all carry real tags matching - none of the request's do not capture it when the name also has plain + With tag filtering enabled, router-wide or by the request's + enable_tag_filtering (which the proxy sets from key/team + router_settings), strategies that all carry real tags matching none of + the request's do not capture it when the name also has plain deployments: returning None hands the request to ordinary tag-aware deployment selection. """ @@ -11487,8 +11489,9 @@ class Router: for tagged in candidates: if "default" in tagged.tags: return tagged + request_scoped_filtering: Final = request_kwargs.get("enable_tag_filtering") is True if ( - self.enable_tag_filtering + (self.enable_tag_filtering or request_scoped_filtering) and all(tagged.tags for tagged in candidates) and self._model_name_has_plain_deployments(model) ): diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 63bc5203417..3c9a4097321 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -253,6 +253,7 @@ def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[li PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file") +PROVIDER_SCOPED_CREATION_FUNCTION_NAMES: Final = frozenset({"_acreate_file"}) def _get_fallback_target_model_group(fallback_entry: str | Mapping[str, object]) -> str | None: @@ -274,6 +275,18 @@ def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: return any(kwargs.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS) +def creates_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: + """ + True when the request creates a resource that will live under one provider's credentials. + + A file uploaded for batches or fine-tuning is stored in the account of the deployment + that handled it, and its id is only usable against the model group the caller named. + Letting the upload fall back to a different model group silently stores the file with + the wrong provider, and every later use of the returned id fails. + """ + return getattr(kwargs.get("original_function"), "__name__", None) in PROVIDER_SCOPED_CREATION_FUNCTION_NAMES + + async def run_async_fallback( *args: tuple[Any], litellm_router: LitellmRouter, @@ -322,7 +335,9 @@ async def run_async_fallback( metadata_variable_name: Final = _get_router_metadata_variable_name( function_name=getattr(kwargs.get("original_function"), "__name__", None) ) - same_model_group_only: Final = references_provider_scoped_resource(kwargs) + same_model_group_only: Final = references_provider_scoped_resource(kwargs) or creates_provider_scoped_resource( + kwargs + ) # Read out of kwargs and narrowed here rather than declared as a parameter: every caller # reaches this function by spreading a loosely-typed kwargs dict, so a declared parameter # would carry an annotation that no call site can actually be checked against. diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d44d4cca6c4..07005d7f9ad 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1631,8 +1631,20 @@ class PromptTokensDetailsWrapper( def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) + extra_fields: Final = self.model_extra + nested_cache_creation_input_tokens: Final = ( + extra_fields.get("cache_creation_input_tokens") if extra_fields is not None else None + ) self.cache_write_tokens = ( - self.cache_write_tokens if self.cache_write_tokens is not None else self.cache_creation_tokens + self.cache_write_tokens + if self.cache_write_tokens is not None + else ( + self.cache_creation_tokens + if self.cache_creation_tokens is not None + else ( + nested_cache_creation_input_tokens if isinstance(nested_cache_creation_input_tokens, int) else None + ) + ) ) if self.character_count is None: del self.character_count diff --git a/tests/batches_tests/test_batch_rate_limits.py b/tests/batches_tests/test_batch_rate_limits.py index 2c804d21ace..ae02c1be12c 100644 --- a/tests/batches_tests/test_batch_rate_limits.py +++ b/tests/batches_tests/test_batch_rate_limits.py @@ -27,6 +27,16 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( from litellm.proxy.utils import InternalUsageCache +def _build_batch_limiter() -> _PROXY_BatchRateLimiter: + internal_usage_cache = InternalUsageCache(dual_cache=DualCache()) + return _PROXY_BatchRateLimiter( + internal_usage_cache=internal_usage_cache, + parallel_request_limiter=_PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ), + ) + + def get_expected_batch_file_usage(file_path: str) -> tuple[int, int]: """ Helper function to calculate expected request count and token count from a batch JSONL file. @@ -69,10 +79,7 @@ async def test_batch_rate_limits(): """ litellm._turn_on_debug() CUSTOM_LLM_PROVIDER = "openai" - BATCH_LIMITER = _PROXY_BatchRateLimiter( - internal_usage_cache=None, - parallel_request_limiter=None, - ) + BATCH_LIMITER = _build_batch_limiter() file_name = "openai_batch_completions.jsonl" _current_dir = os.path.dirname(os.path.abspath(__file__)) @@ -580,10 +587,7 @@ async def test_batch_rate_limiter_without_user_context(tmp_path): CUSTOM_LLM_PROVIDER = "openai" # Setup - BATCH_LIMITER = _PROXY_BatchRateLimiter( - internal_usage_cache=None, - parallel_request_limiter=None, - ) + BATCH_LIMITER = _build_batch_limiter() # Create a simple batch file batch_content = """{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}}""" diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index ebbfd3415a5..b50551ec105 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -18,6 +18,16 @@ - {id: reliability.routing.usage_based.picks_under_tpm, module: reliability, tier: P0, behavior: routing, variant: usage_based, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "router_strategy/lowest_tpm_rpm_v2.py", rationale: "Routes to lowest-TPM deployment; prevents over-allocation"} - {id: reliability.routing.least_busy.picks_lowest_traffic, module: reliability, tier: P1, behavior: routing, variant: least_busy, assertions: [picks_lowest_traffic], exercised_on: [chat_completions, messages], source: "router_strategy/least_busy.py", rationale: "Fewest in-flight requests"} - {id: reliability.routing.complexity_llm_classifier.routes_by_llm_tier, module: reliability, tier: P1, behavior: routing, variant: complexity_llm_classifier, assertions: [routes_by_llm_tier], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py", fail_before_fix: proven, rationale: "v2 auto-router LLM complexity classifier runs over the proxy and routes by semantic tier instead of silently crashing on absent litellm_metadata and falling back to heuristic scoring"} +- {id: reliability.routing.tagged_marker.request_tag_selects_marker, module: reliability, tier: P0, behavior: routing, variant: tagged_marker, assertions: [request_tag_selects_marker], exercised_on: [chat_completions], source: "litellm/router.py:11445", rationale: "Tagged request selects the tagged strategy marker under a shared model_name instead of the plain deployment registered first (GitHub issue #36619)"} +- {id: reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment, module: reliability, tier: P0, behavior: routing, variant: tagged_marker, assertions: [untagged_request_served_by_plain_deployment], exercised_on: [chat_completions, messages, responses], source: "litellm/router.py:11445", rationale: "Untagged requests to a shared model_name are served by the plain deployment on every call, never captured or errored by the tagged marker (GitHub issue #36620)"} +- {id: reliability.routing.tagged_marker.header_tag_selects_marker, module: reliability, tier: P0, behavior: routing, variant: tagged_marker, assertions: [header_tag_selects_marker], exercised_on: [messages], source: "litellm/router.py:11445", rationale: "A request tagged only via the x-litellm-tags header selects the tagged marker on Anthropic-native /v1/messages (GitHub issue #36621)"} +- {id: reliability.routing.tagged_marker.untagged_tier_deployments_still_served, module: reliability, tier: P1, behavior: routing, variant: tagged_marker, assertions: [untagged_tier_deployments_still_served], exercised_on: [chat_completions, messages], source: "litellm/router_strategy/tag_based_routing.py:433", rationale: "Routing tags the marker consumed no longer constrain deployment selection inside the routed tier group, so untagged tier deployments serve the rewrite (GitHub issue #36621)"} +- {id: reliability.routing.tagged_marker.tag_semantics_stay_strict, module: reliability, tier: P1, behavior: routing, variant: tagged_marker, assertions: [tag_semantics_stay_strict], exercised_on: [chat_completions], source: "litellm/router_strategy/tag_based_routing.py:299", rationale: "Tag consumption must not loosen strict semantics: a tagged call aimed straight at an untagged deployment still gets the 401 tags-configuration denial"} +- {id: reliability.routing.tagged_marker.responses_input_routes_through_marker, module: reliability, tier: P0, behavior: routing, variant: tagged_marker, assertions: [responses_input_routes_through_marker], exercised_on: [responses], source: "litellm/router.py:11489", rationale: "Tagged /v1/responses (header or litellm_metadata.tags, string or list input) routes through the marker to its tier, extending the GitHub issues #36620/#36621 tag split to the Responses surface"} +- {id: reliability.routing.tagged_marker.alias_connection_params_stay_with_tier, module: reliability, tier: P0, behavior: routing, variant: tagged_marker, assertions: [alias_connection_params_stay_with_tier], exercised_on: [chat_completions], source: "litellm/router.py:11567", rationale: "An api_key or api_base on the marker alias is never forwarded onto the routed request; the tier deployment calls its provider with its own credential (GitHub PR #36626)"} +- {id: reliability.routing.semantic_auto_router.responses_input_routed, module: reliability, tier: P0, behavior: routing, variant: semantic_auto_router, assertions: [responses_input_routed], exercised_on: [responses], source: "litellm/router_strategy/auto_router/auto_router.py:131", fail_before_fix: proven, rationale: "/v1/responses input is resolved into messages for the semantic auto-router pre-routing hook instead of failing 400 Unmapped LLM provider auto_router (GitHub PR #37333)"} +- {id: reliability.routing.strategy_alias.custom_pricing_ignored, module: reliability, tier: P1, behavior: routing, variant: strategy_alias, assertions: [custom_pricing_ignored], exercised_on: [chat_completions], source: "litellm/router.py:11489", rationale: "Custom pricing on a strategy-router alias never prices the routed request; spend logs at the routed tier deployment's own rate (GitHub PR #36691)"} +- {id: reliability.routing.complexity_heuristic.scores_current_ask_only, module: reliability, tier: P1, behavior: routing, variant: complexity_heuristic, assertions: [scores_current_ask_only], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py:942", rationale: "The heuristic complexity classifier scores the caller's current ask only, so a keyword-heavy agent system prompt cannot inflate the tier (GitHub PR #36721)"} - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 734e63a94e6..df5cb841fad 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -73,6 +73,7 @@ class KeyGenerateBody(BaseModel): allowed_passthrough_routes: list[str] | None = None metadata: KeyMetadata | None = None object_permission: ObjectPermission | None = None + router_settings: "RouterSettingsOverride | None" = None class KeyGenerateResponse(BaseModel): @@ -234,16 +235,18 @@ class ChatBody(BaseModel): class RouterSettingsOverride(BaseModel): - """Per-request `router_settings_override` in a /chat/completions body: the - reliability knobs (fallbacks by trigger, retry count) the reliability suite - drives per call instead of via static router config. Serialized exclude_none, so - an override sets only the strategies a test exercises. Each fallbacks map is - model_name -> the ordered fallback model_names to try.""" + """Router settings a test scopes below the global config: sent per request as + `router_settings_override` in a /chat/completions body (the reliability suite's + fallback and retry knobs) or stored on a key as `router_settings` at + /key/generate (the auto-router suite's tag filtering switch). Serialized + exclude_none, so an override sets only the knobs a test exercises. Each + fallbacks map is model_name -> the ordered fallback model_names to try.""" fallbacks: list[dict[str, list[str]]] | None = None context_window_fallbacks: list[dict[str, list[str]]] | None = None content_policy_fallbacks: list[dict[str, list[str]]] | None = None num_retries: int | None = None + enable_tag_filtering: bool | None = None class ReliabilityChatBody(ChatBody): @@ -744,6 +747,10 @@ class LiteLLMParamsBody(BaseModel): extra_headers: dict[str, str] | None = None use_in_pass_through: bool | None = None complexity_router_config: dict[str, object] | None = None + auto_router_config: str | None = None + auto_router_default_model: str | None = None + auto_router_embedding_model: str | None = None + tags: list[str] | None = None mock_response: str | None = None timeout: float | None = None tpm: int | None = None diff --git a/tests/e2e/router/test_auto_router_regressions_e2e.py b/tests/e2e/router/test_auto_router_regressions_e2e.py new file mode 100644 index 00000000000..c6ef9cda05d --- /dev/null +++ b/tests/e2e/router/test_auto_router_regressions_e2e.py @@ -0,0 +1,616 @@ +"""Live e2e regression pins for strategy-router (auto-router) routing. + +A strategy marker (an ``auto_router/complexity_router`` deployment) and a plain +deployment can share one ``model_name``, split by tags once +``enable_tag_filtering`` is on: tagged requests route through the marker to its +tier models, untagged requests go to the plain deployment. That split, and the +strategy-router alias behaviors around it, regressed repeatedly; each test here +pins one fixed behavior: + +- GitHub issue #36619: a tagged request selects the tagged marker under a + shared name even when a plain deployment was registered first. +- GitHub issue #36620: untagged requests keep being served by the plain + deployment on every call, never captured or 400'd by the tagged marker. +- GitHub issue #36621: a request tagged via the ``x-litellm-tags`` header + routes through the marker even when the tier deployments carry no tags + (the marker consumes the routing tags before deployment selection), while a + tagged call aimed straight at an untagged deployment stays denied. +- GitHub issues #36620/#36621 on /v1/responses: the same tag split holds for + string and list input, whether the tag arrives in litellm_metadata or the + x-litellm-tags header. +- GitHub PR #37333: /v1/responses input is resolved into messages for a + semantic ``auto_router`` deployment's pre-routing hook; such requests used + to fail with 400 "Unmapped LLM provider auto_router" because only chat + messages fed the route matcher. +- GitHub PR #36691: custom pricing on the marker alias never prices the routed + request; spend logs at the routed tier deployment's own rate. +- GitHub PR #36721: the heuristic complexity classifier scores the caller's + current ask only, so a large agent system prompt cannot inflate the tier. +- GitHub PR #36626: connection params on the marker alias (``api_key``, + ``api_base``) stay with the alias; the routed tier calls its provider with + its own credentials. + +Every deployment is registered via /model/new (stage has no static config for +these) and ``enable_tag_filtering`` is enabled through key-level +``router_settings`` on the keys the tag tests mint, so the switch rides only +this module's own requests and the rest of the suite is never filtered. +The served deployment is always read back from the spend log's ``model``, +which stores either the registered alias or the provider-prefixed form. +""" + +import json +import os +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from pydantic import BaseModel, ConfigDict, Field + +from e2e_config import unique_marker +from e2e_http import AnthropicHeaders, AuthHeaders, UnauthorizedError, unwrap +from lifecycle import ResourceManager +from models import ( + AnthropicMessagesBody, + AnthropicMessagesResponse, + ChatBody, + ChatMessage, + ChatMetadata, + KeyGenerateBody, + LiteLLMParamsBody, + RouterSettingsOverride, + SpendLogRow, +) +from proxy_client import ProxyClient + +pytestmark = pytest.mark.e2e + +PLAIN_MODEL = "anthropic/claude-sonnet-5" +CHEAP_MODEL = "anthropic/claude-haiku-4-5" +STRONG_MODEL = "openai/gpt-5.6" +MAX_TOKENS = 16 +PLAIN_SERVED = frozenset({PLAIN_MODEL, "claude-sonnet-5"}) +CHEAP_SERVED = frozenset({CHEAP_MODEL, "claude-haiku-4-5"}) +EMBEDDING_MODEL = "openai/text-embedding-3-small" +SEMANTIC_ROUTE_UTTERANCE = "summarize this quarterly revenue report into three bullet points" + +KEYWORD_HEAVY_SYSTEM_PROMPT = ( + "You are the principal architecture assistant for a distributed systems platform. " + "Analyze every request step by step: design the algorithm, prove its correctness, " + "evaluate time and space complexity, and reason about concurrency, consistency, and " + "fault tolerance tradeoffs. When asked, refactor and debug multi-threaded code, " + "optimize database query plans, derive mathematical proofs, and explain the theorem " + "or lemma behind each optimization. Think through edge cases rigorously before answering. " +) * 4 + + +class TaggedAuthHeaders(AuthHeaders): + x_litellm_tags: str | None = Field(default=None, serialization_alias="x-litellm-tags") + + +class TaggedAnthropicHeaders(AnthropicHeaders): + x_litellm_tags: str | None = Field(default=None, serialization_alias="x-litellm-tags") + + +class ResponsesTagMetadata(BaseModel): + tags: list[str] + + +class ResponsesInputItem(BaseModel): + role: str + content: str + + +class ResponsesBody(BaseModel): + model: str + input: str | list[ResponsesInputItem] + max_output_tokens: int | None = None + litellm_metadata: ResponsesTagMetadata | None = None + + +class ResponsesApiResponse(BaseModel): + """Minimal /v1/responses answer shape; routing is proven from spend logs, + so only the fields the assertions read are modeled.""" + + model_config = ConfigDict(extra="allow") + id: str | None = None + status: str | None = None + model: str | None = None + + +@dataclass(frozen=True, slots=True) +class TagSplitDeployments: + """Scenario A mirrors the customer-shaped config from GitHub issue #36619: + plain deployment registered first, tier deployment and marker both tagged. + Scenario B flips both axes for GitHub issue #36621: marker registered first + and its tier deployment left untagged, so routing depends neither on + registration order nor on tier deployments carrying tags.""" + + tag_a: str + shared_a: str + tier_a: str + tag_b: str + shared_b: str + tier_b: str + + +@dataclass(frozen=True, slots=True) +class ZeroPricedAlias: + alias: str + tier: str + + +@dataclass(frozen=True, slots=True) +class HeuristicSplit: + alias: str + cheap: str + strong: str + + +@dataclass(frozen=True, slots=True) +class SemanticAutoRouter: + marker: str + target: str + fallback: str + embedding: str + + +@dataclass(frozen=True, slots=True) +class CredentialedAlias: + alias: str + tier: str + + +def _provider_key(env_var: str) -> str: + return os.environ.get(env_var) or f"os.environ/{env_var}" + + +def _uniform_tier_config(tier_model: str) -> dict[str, object]: + return { + "classifier_type": "heuristic", + "tiers": {"SIMPLE": tier_model, "MEDIUM": tier_model, "COMPLEX": tier_model, "REASONING": tier_model}, + } + + +def _key_for( + proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False +) -> str: + key: Final = proxy.generate_key( + KeyGenerateBody( + models=models, + user_id="e2e-auto-router-regressions", + router_settings=RouterSettingsOverride(enable_tag_filtering=True) if tag_filtering else None, + ) + ) + resources.defer(lambda: proxy.delete_key(key)) + return key + + +def _hello_chat_body(model: str, tags: list[str] | None = None) -> ChatBody: + return ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"say hello {unique_marker()}")], + max_tokens=MAX_TOKENS, + metadata=ChatMetadata(tags=tags) if tags is not None else None, + ) + + +def _hello_messages_body(model: str) -> AnthropicMessagesBody: + return AnthropicMessagesBody( + model=model, + messages=[ChatMessage(role="user", content=f"say hello {unique_marker()}")], + max_tokens=MAX_TOKENS, + ) + + +def _assert_served_only_by(rows: list[SpendLogRow], allowed: frozenset[str], context: str) -> None: + served: Final = tuple(row.model for row in rows) + assert served and all(model in allowed for model in served), ( + f"{context}: expected every request to be served by one of {sorted(allowed)}, spend logs show {served}" + ) + + +@pytest.fixture(scope="module") +def split(proxy: ProxyClient) -> Iterator[TagSplitDeployments]: + marker: Final = unique_marker() + deployments: Final = TagSplitDeployments( + tag_a=f"e2e-split-a-{marker}", + shared_a=f"e2e-autoroute-a-{marker}", + tier_a=f"e2e-tier-a-{marker}", + tag_b=f"e2e-split-b-{marker}", + shared_b=f"e2e-autoroute-b-{marker}", + tier_b=f"e2e-tier-b-{marker}", + ) + anthropic_key: Final = _provider_key("ANTHROPIC_API_KEY") + marker_params_a: Final = LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config=_uniform_tier_config(deployments.tier_a), + tags=[deployments.tag_a], + ) + marker_params_b: Final = LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config=_uniform_tier_config(deployments.tier_b), + tags=[deployments.tag_b], + ) + registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = ( + (deployments.shared_a, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), + (deployments.tier_a, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=[deployments.tag_a])), + (deployments.shared_a, marker_params_a), + (deployments.shared_b, marker_params_b), + (deployments.tier_b, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key)), + (deployments.shared_b, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), + ) + created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) + try: + yield deployments + finally: + for model_id in created: + proxy.delete_model(model_id) + + +@pytest.fixture(scope="module") +def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: + marker: Final = unique_marker() + named: Final = ZeroPricedAlias(alias=f"e2e-priced-alias-{marker}", tier=f"e2e-priced-tier-{marker}") + alias_params: Final = LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config=_uniform_tier_config(named.tier), + input_cost_per_token=0.0, + output_cost_per_token=0.0, + ) + registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = ( + (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), + (named.alias, alias_params), + ) + created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) + try: + yield named + finally: + for model_id in created: + proxy.delete_model(model_id) + + +@pytest.fixture(scope="module") +def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: + marker: Final = unique_marker() + named: Final = HeuristicSplit( + alias=f"e2e-heuristic-router-{marker}", + cheap=f"e2e-heuristic-cheap-{marker}", + strong=f"e2e-heuristic-strong-{marker}", + ) + config: Final[dict[str, object]] = { + "classifier_type": "heuristic", + "token_thresholds": {"simple": 15, "complex": 400}, + "tiers": {"SIMPLE": named.cheap, "MEDIUM": named.strong, "COMPLEX": named.strong, "REASONING": named.strong}, + } + registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = ( + (named.cheap, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), + (named.strong, LiteLLMParamsBody(model=STRONG_MODEL, api_key=_provider_key("OPENAI_API_KEY"))), + (named.alias, LiteLLMParamsBody(model="auto_router/complexity_router", complexity_router_config=config)), + ) + created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) + try: + yield named + finally: + for model_id in created: + proxy.delete_model(model_id) + + +@pytest.fixture(scope="module") +def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: + marker: Final = unique_marker() + named: Final = SemanticAutoRouter( + marker=f"e2e-semantic-router-{marker}", + target=f"e2e-semantic-target-{marker}", + fallback=f"e2e-semantic-fallback-{marker}", + embedding=f"e2e-semantic-embedding-{marker}", + ) + router_config: Final = json.dumps( + {"routes": [{"name": named.target, "utterances": [SEMANTIC_ROUTE_UTTERANCE], "score_threshold": 0.3}]} + ) + marker_params: Final = LiteLLMParamsBody( + model=f"auto_router/{named.marker}", + auto_router_config=router_config, + auto_router_default_model=named.fallback, + auto_router_embedding_model=named.embedding, + ) + registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = ( + (named.embedding, LiteLLMParamsBody(model=EMBEDDING_MODEL, api_key=_provider_key("OPENAI_API_KEY"))), + (named.target, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), + (named.fallback, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), + (named.marker, marker_params), + ) + created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) + try: + yield named + finally: + for model_id in created: + proxy.delete_model(model_id) + + +@pytest.fixture(scope="module") +def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: + marker: Final = unique_marker() + named: Final = CredentialedAlias(alias=f"e2e-cred-alias-{marker}", tier=f"e2e-cred-tier-{marker}") + alias_params: Final = LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config=_uniform_tier_config(named.tier), + api_key=f"sk-alias-never-used-{marker}", + ) + registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = ( + (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), + (named.alias, alias_params), + ) + created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) + try: + yield named + finally: + for model_id in created: + proxy.delete_model(model_id) + + +class TestTagSplitRouting: + @pytest.mark.covers("reliability.routing.tagged_marker.request_tag_selects_marker") + def test_body_tagged_chat_routes_through_the_marker_to_its_tier( + self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + ) -> None: + """Pins GitHub issue #36619: with tag filtering on, a chat request whose + body metadata tags match the tagged marker under a shared model name is + answered by the marker's tier deployment, not by the plain deployment + that was registered under the name first.""" + key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a, tags=[split.tag_a]))) + assert chat.choices, "tagged chat through the shared name returned no choices" + rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged chat on the shared name") + + @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") + def test_untagged_chat_is_always_served_by_the_plain_deployment( + self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + ) -> None: + """Pins GitHub issue #36620: untagged chat requests to the shared name + succeed on every call and are all served by the plain deployment; the + tagged marker never captures them, so no intermittent auto-router + errors and no tier hijacking.""" + key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + for _ in range(5): + chat = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a))) + assert chat.choices, "untagged chat through the shared name returned no choices" + rows: Final = proxy.poll_logs_for_key(key, min_rows=5) + _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged chat on the shared name") + + @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") + def test_untagged_messages_is_served_by_the_plain_deployment( + self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + ) -> None: + """Pins GitHub issue #36620 on the /v1/messages surface: an untagged + Anthropic-native request to the shared name is served by the plain + deployment, not captured by the tagged marker.""" + key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + answer: Final = unwrap(proxy.messages(key, _hello_messages_body(split.shared_a))) + assert answer.content or answer.choices, "untagged /v1/messages returned neither content nor choices" + rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/messages on the shared name") + + +class TestUntaggedTierDeployments: + @pytest.mark.covers("reliability.routing.tagged_marker.header_tag_selects_marker") + def test_header_tagged_messages_routes_through_the_marker_to_an_untagged_tier( + self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + ) -> None: + """Pins GitHub issue #36621: a /v1/messages request tagged only via the + x-litellm-tags header selects the tagged marker, and the rewrite still + lands on the tier deployment even though that deployment carries no + tags, because the marker consumed the routing tags.""" + key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) + headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_b) + answer: Final = unwrap( + proxy.transport.post( + "/v1/messages", + headers=headers, + json=_hello_messages_body(split.shared_b), + response_type=AnthropicMessagesResponse, + ) + ) + assert answer.content or answer.choices, "header-tagged /v1/messages returned neither content nor choices" + rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "header-tagged /v1/messages on the shared name") + + @pytest.mark.covers("reliability.routing.tagged_marker.untagged_tier_deployments_still_served") + def test_body_tagged_chat_reaches_the_untagged_tier_after_marker_rewrite( + self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + ) -> None: + """Pins the tag-consumption half of GitHub issue #36621: after the + tagged marker rewrites the request to its tier model, the consumed + routing tags no longer constrain deployment selection, so the untagged + tier deployment serves the request instead of a strict-tag denial.""" + key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) + chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_b, tags=[split.tag_b]))) + assert chat.choices, "body-tagged chat through the marker-first shared name returned no choices" + rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "body-tagged chat with untagged tier") + + @pytest.mark.covers("reliability.routing.tagged_marker.tag_semantics_stay_strict") + def test_tagged_call_straight_at_an_untagged_deployment_stays_denied( + self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + ) -> None: + """The tag-consumption fix must not loosen strict tag semantics: a + tagged request aimed directly at an untagged deployment (no marker + involved) is still rejected with the 401 tags-configuration error.""" + key: Final = _key_for(proxy, resources, [split.tier_b], tag_filtering=True) + result: Final = proxy.chat(key, _hello_chat_body(split.tier_b, tags=[split.tag_b])) + assert isinstance(result, UnauthorizedError), ( + f"expected the tagged direct call to an untagged deployment to be denied with 401, got {result}" + ) + + +class TestResponsesApiTagRouting: + @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") + def test_header_tagged_responses_with_string_input_routes_to_the_tier( + self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + ) -> None: + """Pins the /v1/responses surface of the tag split (GitHub issues + #36620/#36621): a /v1/responses request with string input, tagged via + the x-litellm-tags header, succeeds and routes through the tagged + marker to its tier.""" + key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_a) + body: Final = ResponsesBody( + model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + ) + answer: Final = unwrap( + proxy.transport.post("/v1/responses", headers=headers, json=body, response_type=ResponsesApiResponse) + ) + assert answer.id, "header-tagged /v1/responses returned no response id" + rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "header-tagged /v1/responses string input") + + @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") + def test_body_tagged_responses_with_list_input_routes_to_the_tier( + self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + ) -> None: + """Pins the body-tag and list-input combination of the same split: + /v1/responses with litellm_metadata.tags and structured input items + routes through the tagged marker to its tier.""" + key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + body: Final = ResponsesBody( + model=split.shared_a, + input=[ResponsesInputItem(role="user", content=f"say hello {unique_marker()}")], + max_output_tokens=64, + litellm_metadata=ResponsesTagMetadata(tags=[split.tag_a]), + ) + answer: Final = unwrap( + proxy.transport.post( + "/v1/responses", + headers=proxy.transport.bearer(key), + json=body, + response_type=ResponsesApiResponse, + ) + ) + assert answer.id, "body-tagged /v1/responses returned no response id" + rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged /v1/responses list input") + + @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") + def test_untagged_responses_is_served_by_the_plain_deployment( + self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + ) -> None: + """Pins the untagged half of the /v1/responses tag split: an untagged + request to the shared name is served by the plain deployment, matching + the chat and messages surfaces.""" + key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + body: Final = ResponsesBody( + model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + ) + answer: Final = unwrap( + proxy.transport.post( + "/v1/responses", + headers=proxy.transport.bearer(key), + json=body, + response_type=ResponsesApiResponse, + ) + ) + assert answer.id, "untagged /v1/responses returned no response id" + rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/responses on the shared name") + + +class TestStrategyAliasPricing: + @pytest.mark.covers("reliability.routing.strategy_alias.custom_pricing_ignored") + def test_zero_priced_alias_still_logs_spend_at_the_tier_rate( + self, proxy: ProxyClient, resources: ResourceManager, zero_priced_alias: ZeroPricedAlias + ) -> None: + """Pins GitHub PR #36691: custom pricing registered on a strategy-router + alias never prices the routed request. The alias here carries explicit + zero pricing, so any zero-spend row would prove the alias pricing was + applied; the routed tier deployment's real rate must produce spend > 0.""" + key: Final = _key_for(proxy, resources, [zero_priced_alias.alias, zero_priced_alias.tier]) + chat: Final = unwrap(proxy.chat(key, _hello_chat_body(zero_priced_alias.alias))) + assert chat.choices, "chat through the zero-priced alias returned no choices" + rows: Final = proxy.poll_logs_for_key( + key, min_rows=1, predicate=lambda logged: all((row.spend or 0.0) > 0.0 for row in logged) + ) + _assert_served_only_by(rows, CHEAP_SERVED | {zero_priced_alias.tier}, "chat through the zero-priced alias") + priced: Final = tuple((row.model, row.spend) for row in rows) + assert all((row.spend or 0.0) > 0.0 for row in rows), ( + f"expected spend at the tier deployment's own rate, got zero-spend rows: {priced}" + ) + + +class TestComplexityHeuristicScope: + @pytest.mark.covers("reliability.routing.complexity_heuristic.scores_current_ask_only") + def test_trivial_ask_behind_keyword_heavy_system_prompt_stays_on_the_cheap_tier( + self, proxy: ProxyClient, resources: ResourceManager, heuristic_split: HeuristicSplit + ) -> None: + """Pins GitHub PR #36721: the heuristic complexity classifier scores the + caller's current ask alone. The trivial ask scores SIMPLE on its own, + while the accompanying ~2KB agent system prompt is packed with enough + reasoning and complexity keywords that scoring the combined text lands + in REASONING; only ask-only scoring keeps this on the cheap tier.""" + key: Final = _key_for( + proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong] + ) + body: Final = ChatBody( + model=heuristic_split.alias, + messages=[ + ChatMessage(role="system", content=KEYWORD_HEAVY_SYSTEM_PROMPT), + ChatMessage(role="user", content=f"hi {unique_marker()}"), + ], + max_tokens=MAX_TOKENS, + ) + chat: Final = unwrap(proxy.chat(key, body)) + assert chat.choices, "chat through the heuristic router returned no choices" + rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + _assert_served_only_by( + rows, CHEAP_SERVED | {heuristic_split.cheap}, "trivial ask behind a keyword-heavy system prompt" + ) + + +class TestSemanticAutoRouterResponses: + @pytest.mark.covers("reliability.routing.semantic_auto_router.responses_input_routed") + def test_responses_input_reaches_the_semantic_auto_router( + self, proxy: ProxyClient, resources: ResourceManager, semantic_auto_router: SemanticAutoRouter + ) -> None: + """Pins GitHub PR #37333: /v1/responses input is resolved into messages + for the semantic auto-router's pre-routing hook, so the marker embeds + the input, matches its route, and the target deployment serves the + request; before the fix the hook saw no messages and the request + failed with 400 "Unmapped LLM provider auto_router".""" + key: Final = _key_for( + proxy, + resources, + [semantic_auto_router.marker, semantic_auto_router.target, semantic_auto_router.fallback], + ) + body: Final = ResponsesBody( + model=semantic_auto_router.marker, input=SEMANTIC_ROUTE_UTTERANCE, max_output_tokens=64 + ) + answer: Final = unwrap( + proxy.transport.post( + "/v1/responses", + headers=proxy.transport.bearer(key), + json=body, + response_type=ResponsesApiResponse, + ) + ) + assert answer.id, "/v1/responses through the semantic auto-router returned no response id" + rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + _assert_served_only_by( + rows, CHEAP_SERVED | {semantic_auto_router.target}, "semantic auto-router /v1/responses string input" + ) + + +class TestAliasParamForwarding: + @pytest.mark.covers("reliability.routing.tagged_marker.alias_connection_params_stay_with_tier") + def test_alias_api_key_never_overrides_the_tier_credential( + self, proxy: ProxyClient, resources: ResourceManager, credentialed_alias: CredentialedAlias + ) -> None: + """Pins GitHub PR #36626: an api_key set on the marker alias entry is + never forwarded onto the routed request, so the tier deployment calls + its provider with its own credential. Before the fix the alias's key + was copied into the request, overriding the tier's credential, and + every routed call failed provider auth.""" + key: Final = _key_for(proxy, resources, [credentialed_alias.alias, credentialed_alias.tier]) + chat: Final = unwrap(proxy.chat(key, _hello_chat_body(credentialed_alias.alias))) + assert chat.choices, "chat through the credentialed alias returned no choices" + rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + _assert_served_only_by(rows, CHEAP_SERVED | {credentialed_alias.tier}, "chat through the credentialed alias") diff --git a/tests/e2e/ui/tests/mcp/mcpTools.spec.ts b/tests/e2e/ui/tests/mcp/mcpTools.spec.ts index edaeab196aa..225ca8b9449 100644 --- a/tests/e2e/ui/tests/mcp/mcpTools.spec.ts +++ b/tests/e2e/ui/tests/mcp/mcpTools.spec.ts @@ -63,7 +63,7 @@ test.describe("MCP Tools", () => { // The form is generated from the tool's inputSchema, so `repoName` proves the schema // round-tripped through the proxy instead of the panel falling back to a generic field. - const repoInput = page.locator('input[id="repoName"]'); + const repoInput = page.getByLabel(/repoName/); await expect(repoInput).toBeVisible(); await repoInput.fill(TOOL_ARG_REPO); diff --git a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts index 1b11ea69f97..bd5373e1569 100644 --- a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -338,8 +338,8 @@ test.describe("Add Model", () => { await page.getByRole("button", { name: "Add Model" }).last().click(); - // Scope to antd's notification container so a stale toast can't satisfy this. - await expect(page.locator(".ant-notification").getByText("created successfully").last()).toBeVisible({ + // Scope to the toast container so a stale toast can't satisfy this. + await expect(page.locator("[data-sonner-toast]").getByText("created successfully").last()).toBeVisible({ timeout: 15_000, }); diff --git a/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts index e67dcb96f36..c532641b238 100644 --- a/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts @@ -85,9 +85,9 @@ test.describe("Clear custom pricing on a deployment", () => { const inputCost = page.getByPlaceholder("Enter input cost"); const outputCost = page.getByPlaceholder("Enter output cost"); // Both cache fields share the same placeholder ("Defaults to Input Cost if blank"), - // so disambiguate via the Form.Item id (AntD assigns the `name` prop as input id). - const cacheReadCost = page.locator("#cache_read_cost"); - const cacheWriteCost = page.locator("#cache_write_cost"); + // so disambiguate via their labels. + const cacheReadCost = page.getByLabel(/Cache Read Cost/); + const cacheWriteCost = page.getByLabel(/Cache Write Cost/); await inputCost.waitFor({ timeout: 15_000 }); for (const field of [inputCost, outputCost, cacheReadCost, cacheWriteCost]) { await field.click({ clickCount: 3 }); diff --git a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts index d9b0f959c9f..004fedb3263 100644 --- a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts @@ -36,9 +36,8 @@ test.describe("Proxy Admin - Keys", () => { // Wait for the key creation modal await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); - // Fill key name (has data-testid="base-input" in the built UI) const keyName = `e2e-admin-key-${Date.now()}`; - await page.getByTestId("base-input").fill(keyName); + await page.getByLabel(/Key Name/).fill(keyName); // Select team const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); @@ -192,7 +191,7 @@ test.describe("Proxy Admin - Keys", () => { await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); const keyName = `e2e-admin-allproxy-${Date.now()}`; - await page.getByTestId("base-input").fill(keyName); + await page.getByLabel(/Key Name/).fill(keyName); // No team selection — leave team dropdown empty so the key is owned by the admin user @@ -220,7 +219,7 @@ test.describe("Proxy Admin - Keys", () => { await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); const keyName = `e2e-admin-specific-${Date.now()}`; - await page.getByTestId("base-input").fill(keyName); + await page.getByLabel(/Key Name/).fill(keyName); // Open the model multi-select and pick a single specific model. Use // getByRole("option", ...) to avoid the strict-mode collision between diff --git a/tests/e2e/ui/tests/proxy-admin/teams.spec.ts b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts index d7c8eb6237e..172fde173df 100644 --- a/tests/e2e/ui/tests/proxy-admin/teams.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts @@ -79,7 +79,7 @@ test.describe("Proxy Admin - Teams", () => { await expect(modal).toBeVisible({ timeout: 5_000 }); // The email field is a Select — type to search, then select from dropdown - await modal.locator(".ant-select").first().click(); + await modal.getByRole("combobox").first().click(); await page.keyboard.type("invitable@test.local"); // Wait for the option to appear, then select via keyboard (avoids viewport issues) diff --git a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts index d71d5e6c0fe..0ef74e71529 100644 --- a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts +++ b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts @@ -66,7 +66,7 @@ test.describe("Team Admin", () => { // Use a dedicated invitee user so this doesn't race with the proxy-admin // "Invite a user" test that adds invitable@test.local to the same team. - await modal.locator(".ant-select").first().click(); + await modal.getByRole("combobox").first().click(); await page.keyboard.type("invitable-team@test.local"); const emailOption = page.getByRole("option", { name: "invitable-team@test.local" }).first(); @@ -136,7 +136,7 @@ test.describe("Team Admin", () => { await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); const keyName = `e2e-team-admin-key-${Date.now()}`; - await page.getByTestId("base-input").fill(keyName); + await page.getByLabel(/Key Name/).fill(keyName); // Team selector — same locator pattern as the proxy-admin keys test. const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index f0a73325afa..33dcdbb57a5 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -269,7 +269,7 @@ class TestAimlImageGeneration(BaseImageGenTest): class TestGoogleImageGen(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: - return {"model": "gemini/imagen-4.0-generate-001"} + return {"model": "gemini/gemini-3.1-flash-image"} @pytest.mark.skip(reason="Runwayml image generation API only tested locally") diff --git a/tests/llm_translation/test_groq.py b/tests/llm_translation/test_groq.py index cf4be9e801e..c720f818eaf 100644 --- a/tests/llm_translation/test_groq.py +++ b/tests/llm_translation/test_groq.py @@ -20,7 +20,7 @@ from litellm.llms.groq.chat.transformation import ( class TestGroq(BaseLLMChatTest): def get_base_completion_call_args(self) -> dict: return { - "model": "groq/llama-3.3-70b-versatile", + "model": "groq/openai/gpt-oss-120b", } def test_tool_call_no_arguments(self, tool_call_no_arguments): diff --git a/tests/local_testing/test_model_alias_map.py b/tests/local_testing/test_model_alias_map.py index 14c1de2f6a7..9ef0448e7c6 100644 --- a/tests/local_testing/test_model_alias_map.py +++ b/tests/local_testing/test_model_alias_map.py @@ -15,7 +15,7 @@ from litellm import completion, embedding litellm.set_verbose = True -model_alias_map = {"good-model": "groq/llama-3.1-8b-instant"} +model_alias_map = {"good-model": "groq/openai/gpt-oss-120b"} def test_model_alias_map(caplog): @@ -34,7 +34,7 @@ def test_model_alias_map(caplog): if rec.levelname == "ERROR" and rec.name.startswith("LiteLLM"): pytest.fail(f"Unexpected litellm ERROR log: {rec.getMessage()}") - assert "llama-3.1-8b-instant" in response.model + assert "gpt-oss-120b" in response.model except litellm.ServiceUnavailableError: pass except Exception as e: diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index af1a052e86d..7bc29517f8c 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -120,7 +120,7 @@ async def test_router_provider_wildcard_routing(): print("response 2 = ", response2) response3 = await router.acompletion( - model="groq/llama-3.1-8b-instant", + model="groq/openai/gpt-oss-120b", messages=[{"role": "user", "content": "hello"}], ) diff --git a/tests/local_testing/test_router_batch_completion.py b/tests/local_testing/test_router_batch_completion.py index f7a1b41ca29..bb9e1851c61 100644 --- a/tests/local_testing/test_router_batch_completion.py +++ b/tests/local_testing/test_router_batch_completion.py @@ -44,7 +44,7 @@ async def test_batch_completion_multiple_models(mode): { "model_name": "groq-llama", "litellm_params": { - "model": "groq/llama-3.1-8b-instant", + "model": "groq/openai/gpt-oss-120b", }, }, ] @@ -143,7 +143,7 @@ async def test_batch_completion_fastest_response_streaming(): { "model_name": "groq-llama", "litellm_params": { - "model": "groq/llama-3.1-8b-instant", + "model": "groq/openai/gpt-oss-120b", }, }, ] @@ -179,7 +179,7 @@ async def test_batch_completion_multiple_models_multiple_messages(): { "model_name": "groq-llama", "litellm_params": { - "model": "groq/llama-3.1-8b-instant", + "model": "groq/openai/gpt-oss-120b", }, }, ] diff --git a/tests/local_testing/test_stream_chunk_builder.py b/tests/local_testing/test_stream_chunk_builder.py index 38e04b93f18..664fd936205 100644 --- a/tests/local_testing/test_stream_chunk_builder.py +++ b/tests/local_testing/test_stream_chunk_builder.py @@ -871,7 +871,7 @@ def load_env(): } LLAMA3_3 = { "messages": messages, - "model": "groq/llama-3.3-70b-versatile", + "model": "groq/openai/gpt-oss-120b", "api_base": "https://api.groq.com/openai/v1", "temperature": 0.0, "tools": tools, diff --git a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py index d56d5e51b04..3b42595b959 100644 --- a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py +++ b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py @@ -317,110 +317,69 @@ async def test_redaction_responses_api_stream(): @pytest.mark.asyncio async def test_redaction_responses_api_with_reasoning_summary(): """Test that reasoning summary in ResponsesAPIResponse output is properly redacted""" + import litellm from litellm.litellm_core_utils.redact_messages import perform_redaction - # Create a simple mock object with output items that have reasoning summaries - class MockResponsesAPIResponse: - def __init__(self): - self.output = [ - # Reasoning item with summary - type( - "obj", - (object,), + response = litellm.ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "reasoning", + "id": "rs_123", + "summary": [ { - "type": "reasoning", - "id": "rs_123", - "summary": [ - type( - "obj", - (object,), - { - "text": "This is a detailed reasoning summary that should be redacted", - "type": "summary_text", - }, - )() - ], - }, - )(), - # Message item with content - type( - "obj", - (object,), + "type": "summary_text", + "text": "This is a detailed reasoning summary that should be redacted", + } + ], + }, + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ { - "type": "message", - "id": "msg_123", - "content": [ - type( - "obj", - (object,), - { - "text": "This is the actual message content", - "type": "output_text", - }, - )() - ], - }, - )(), - ] - self.reasoning = {"effort": "low", "summary": "auto"} + "type": "output_text", + "text": "This is the actual message content", + "annotations": [], + } + ], + }, + ], + reasoning={"effort": "low", "summary": "auto"}, + ) - # Mock as ResponsesAPIResponse so perform_redaction recognizes it - mock_response = MockResponsesAPIResponse() - mock_response.__class__.__name__ = "ResponsesAPIResponse" + model_call_details = { + "messages": [{"role": "user", "content": "test"}], + "prompt": "test prompt", + "input": "test input", + } - # Patch isinstance to recognize our mock as ResponsesAPIResponse - import litellm + redacted_result = perform_redaction(model_call_details, response) - original_isinstance = isinstance + assert isinstance( + redacted_result, litellm.ResponsesAPIResponse + ), "Redaction should preserve the ResponsesAPIResponse type" - def patched_isinstance(obj, cls): - if ( - cls == litellm.ResponsesAPIResponse - and obj.__class__.__name__ == "ResponsesAPIResponse" - ): - return True - return original_isinstance(obj, cls) + reasoning_item = redacted_result.output[0] + assert ( + reasoning_item.summary[0].text == "redacted-by-litellm" + ), "Reasoning summary text should be redacted" - import builtins + message_item = redacted_result.output[1] + assert ( + message_item.content[0].text == "redacted-by-litellm" + ), "Message content text should be redacted" - builtins.isinstance = patched_isinstance + assert ( + redacted_result.reasoning is None + ), "Top-level reasoning field should be None" - try: - model_call_details = { - "messages": [{"role": "user", "content": "test"}], - "prompt": "test prompt", - "input": "test input", - } - - # Perform redaction - redacted_result = perform_redaction(model_call_details, mock_response) - - # Verify reasoning summary text is redacted - reasoning_item = redacted_result.output[0] - assert ( - reasoning_item.summary[0].text == "redacted-by-litellm" - ), "Reasoning summary text should be redacted" - - # Verify message content is also redacted - message_item = redacted_result.output[1] - assert ( - message_item.content[0].text == "redacted-by-litellm" - ), "Message content text should be redacted" - - # Verify top-level reasoning field is removed - assert ( - redacted_result.reasoning is None - ), "Top-level reasoning field should be None" - - # Verify input messages are redacted - assert ( - model_call_details["messages"][0]["content"] == "redacted-by-litellm" - ), "Input messages should be redacted" - - print("✓ Reasoning summary redaction test passed") - finally: - # Restore original isinstance - builtins.isinstance = original_isinstance + assert ( + model_call_details["messages"][0]["content"] == "redacted-by-litellm" + ), "Input messages should be redacted" @pytest.mark.asyncio diff --git a/tests/ocr_tests/test_ocr_azure_document_intelligence.py b/tests/ocr_tests/test_ocr_azure_document_intelligence.py index 09c21842ad7..5736bd797e3 100644 --- a/tests/ocr_tests/test_ocr_azure_document_intelligence.py +++ b/tests/ocr_tests/test_ocr_azure_document_intelligence.py @@ -62,7 +62,7 @@ class TestAzureDocumentIntelligencePagesParam: return AzureDocumentIntelligenceOCRConfig() def test_get_supported_ocr_params_includes_pages_and_features(self, cfg): - assert cfg.get_supported_ocr_params("prebuilt-layout") == ["pages", "features"] + assert cfg.get_supported_ocr_params("prebuilt-layout") == ["pages", "features", "req_format"] def test_map_ocr_params_mistral_zero_based_int_list(self, cfg): mapped = cfg.map_ocr_params({"pages": [0, 1, 2]}, {}, "prebuilt-layout") diff --git a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py index c263b8ce381..77fb924c085 100644 --- a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py +++ b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py @@ -413,35 +413,54 @@ async def test_pass_through_request_logging_failure_with_stream( assert response.body == b'{"mock": "response"}' +PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES = { + "/comprehendmedical": {"POST"}, + "/comprehendmedical/{operation}": {"POST"}, +} + + def test_pass_through_routes_support_all_methods(): """ - Test that all pass-through routes support GET, POST, PUT, DELETE, PATCH methods + A pass-through route fronts a whole provider API, so narrowing its method + set turns a request the upstream would have accepted into a 405. The + exceptions are providers whose wire protocol admits only one method: Amazon + Comprehend Medical speaks AWS JSON 1.1, which is POST-only, so there is no + other method to forward. """ - # Import the routers from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( router as llm_router, ) - # Expected HTTP methods expected_methods = {"GET", "POST", "PUT", "DELETE", "PATCH"} - # Function to check routes in a router def check_router_methods(router): for route in router.routes: if isinstance(route, APIRoute): - # Get path and methods for this route path = route.path methods = set(route.methods) - print("supported methods for route", path, "are", methods) - # Assert all expected methods are supported + allowed = PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES.get(path, expected_methods) assert ( - methods == expected_methods - ), f"Route {path} does not support all methods. Supported: {methods}, Expected: {expected_methods}" + methods == allowed + ), f"Route {path} does not support all methods. Supported: {methods}, Expected: {allowed}" - # Check both routers check_router_methods(llm_router) +def test_protocol_constrained_pass_through_exemptions_are_not_stale(): + """ + The exemption list above weakens the method contract, so it must not + outlive the routes it covers: a renamed or deleted route has to fail here + rather than sit in the list silently exempting nothing. + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + router as llm_router, + ) + + registered_paths = {route.path for route in llm_router.routes if isinstance(route, APIRoute)} + unmatched = set(PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES) - registered_paths + assert not unmatched, f"Exempted pass-through routes no longer exist: {sorted(unmatched)}" + + def test_is_bedrock_agent_runtime_route(): """ Test that _is_bedrock_agent_runtime_route correctly identifies bedrock agent runtime endpoints diff --git a/tests/search_tests/test_tinyfish_search.py b/tests/search_tests/test_tinyfish_search.py index aca28544513..becb8287a29 100644 --- a/tests/search_tests/test_tinyfish_search.py +++ b/tests/search_tests/test_tinyfish_search.py @@ -35,11 +35,16 @@ MOCK_TINYFISH_RESPONSE = { def _make_mock_response( - json_data: dict, status_code: int = 200, request_url: str | None = None + json_data: dict, + status_code: int = 200, + request_url: str | None = None, + headers: dict | None = None, ) -> MagicMock: mock = MagicMock() mock.status_code = status_code mock.json.return_value = json_data + # httpx.Headers normalizes keys to lowercase — mirror production behavior. + mock.headers = httpx.Headers(headers or {}) if request_url: mock.request = MagicMock() mock.request.url = httpx.URL(request_url) @@ -163,7 +168,7 @@ class TestTinyfishSearch: @pytest.mark.asyncio async def test_fetch_param_round_trip(self): - # End-to-end check: caller passes `fetch=...` (JSON-encoded tf-fetch + # End-to-end check: caller passes `fetch=...` (JSON-encoded fetch # config); param reaches TinyFish on the request side and the nested # `fetch` object on each result surfaces back to the SearchResult on the # response side. No LiteLLM-side support code is required. @@ -235,6 +240,58 @@ class TestTinyfishSearch: assert result.results[0].title == "Result 0" assert result.results[2].title == "Result 2" + @pytest.mark.asyncio + async def test_top_level_extras_surface_end_to_end(self): + # Envelope extras (`query`, `total_results`, `page`) must survive the + # full asearch dispatch — proves LiteLLM's entry-point plumbing outside + # our transformer doesn't accidentally strip them. + os.environ["TINYFISH_API_KEY"] = "sk-tinyfish-test" + + mock_response = _make_mock_response(MOCK_TINYFISH_RESPONSE) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", + new_callable=AsyncMock, + ) as mock_get: + mock_get.return_value = mock_response + + response = await litellm.asearch( + query="web automation tools", + search_provider="tinyfish", + ) + + assert getattr(response, "query", None) == "web automation tools" + assert getattr(response, "total_results", None) == 2 + assert getattr(response, "page", None) == 0 + + @pytest.mark.asyncio + async def test_response_headers_surface_end_to_end(self): + # Response headers must land on `_hidden_params` after the full + # asearch dispatch (both raw and sanitized channels). + os.environ["TINYFISH_API_KEY"] = "sk-tinyfish-test" + + mock_response = _make_mock_response( + MOCK_TINYFISH_RESPONSE, + headers={"X-Request-ID": "req-e2e-1"}, + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", + new_callable=AsyncMock, + ) as mock_get: + mock_get.return_value = mock_response + + response = await litellm.asearch( + query="test", + search_provider="tinyfish", + ) + + raw = response._hidden_params["headers"] + add = response._hidden_params["additional_headers"] + # httpx lowercases; both channels agree on the value. + assert raw["x-request-id"] == "req-e2e-1" + assert add["llm_provider-x-request-id"] == "req-e2e-1" + @pytest.mark.asyncio async def test_empty_results(self): os.environ["TINYFISH_API_KEY"] = "sk-tinyfish-test" diff --git a/tests/test_litellm/integrations/otel/test_db_endpoint.py b/tests/test_litellm/integrations/otel/test_db_endpoint.py new file mode 100644 index 00000000000..5ab0a927b52 --- /dev/null +++ b/tests/test_litellm/integrations/otel/test_db_endpoint.py @@ -0,0 +1,312 @@ +"""Tests for litellm/integrations/otel/model/db_endpoint.py + +Prisma talks to PostgreSQL through a loopback query engine, so a DB span with no +``server.address`` gets attributed to ``localhost`` by the backend. These cover +the endpoint derivation that names the real server, for the local engine and for +remote and read-replica deployments, and pin the rule that no credential is ever +exported. +""" + +import os +from unittest.mock import patch + +import pytest + +from litellm.integrations.otel.model.db_endpoint import ( + DatabaseEndpoint, + db_span_attributes, + parse_database_endpoint, + postgres_endpoint, +) + +LOCAL_DSN = "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm" +REMOTE_DSN = "postgresql://llmproxy:s3cr3t@litellm-prod.abc123.us-east-1.rds.amazonaws.com:6432/litellm?schema=reporting&sslmode=require" +REPLICA_DSN = "postgresql://reader:r3ad0nly@litellm-prod-ro.abc123.us-east-1.rds.amazonaws.com/litellm_replica" + + +def _resolve(service, call_type=None, database_url=None, read_replica_url=None): + """Resolve attributes with the two DB env vars set, as the proxy sets them.""" + env = {k: v for k, v in (("DATABASE_URL", database_url), ("DATABASE_URL_READ_REPLICA", read_replica_url)) if v} + with patch.dict(os.environ, env, clear=False): + for absent in {"DATABASE_URL", "DATABASE_URL_READ_REPLICA"} - set(env): + os.environ.pop(absent, None) + return dict(db_span_attributes(service, call_type)) + + +def test_local_prisma_engine_endpoint_is_the_postgres_server_not_the_engine(): + assert parse_database_endpoint(LOCAL_DSN) == DatabaseEndpoint( + address="localhost", port=5432, namespace="litellm" + ) + + +def test_remote_endpoint_keeps_host_port_and_schema_qualified_namespace(): + assert parse_database_endpoint(REMOTE_DSN) == DatabaseEndpoint( + address="litellm-prod.abc123.us-east-1.rds.amazonaws.com", + port=6432, + namespace="litellm|reporting", + ) + + +def test_read_replica_dsn_parses_to_the_replica_host_and_database(): + assert parse_database_endpoint(REPLICA_DSN) == DatabaseEndpoint( + address="litellm-prod-ro.abc123.us-east-1.rds.amazonaws.com", + port=5432, + namespace="litellm_replica", + ) + + +def test_default_schema_is_not_spelled_out_in_the_namespace(): + """``?schema=public`` and no schema at all are the same deployment, so they + must not split a group-by on db.namespace.""" + assert parse_database_endpoint("postgresql://u:p@db.internal/litellm?schema=public") == parse_database_endpoint( + "postgresql://u:p@db.internal/litellm" + ) + + +def test_unix_socket_host_parameter_wins_over_the_netloc(): + """libpq and the Cloud SQL connector both put the real target in ``host=`` + behind a localhost netloc, which is the attribution this module removes.""" + assert parse_database_endpoint( + "postgresql://u:p@localhost:5432/litellm?host=/cloudsql/proj:us-east1:inst" + ) == DatabaseEndpoint(address="/cloudsql/proj:us-east1:inst", port=5432, namespace="litellm") + + +def test_socket_only_dsn_without_a_netloc_host_still_resolves(): + assert parse_database_endpoint("postgresql:///litellm?host=/var/run/postgresql") == DatabaseEndpoint( + address="/var/run/postgresql", port=5432, namespace="litellm" + ) + + +def test_percent_encoded_database_name_is_decoded(): + endpoint = parse_database_endpoint("postgresql://u:p@db.internal/litellm%20prod") + assert endpoint is not None and endpoint.namespace == "litellm prod" + + +MISPARSED_AUTHORITY_DSNS = ( + ("postgresql://litellm:/kJ8xQz+9wT@db.internal:5432/litellm", "kJ8xQz+9wT"), + ("postgresql://litellm:12345/aBcD@db.internal:5432/litellm", "aBcD"), + # '#' sends the tail to the fragment and '?' to the query, so the path is + # empty and only the stranded userinfo '@' reveals the mis-split. + ("postgresql://litellm:12345#aBcD@db.internal/litellm", "aBcD"), + ("postgresql://litellm:12345?aBcD@db.internal/litellm", "aBcD"), + # A '?'-stranded tail that happens to parse as parameters, including one + # that hijacks the host= parameter into server.address. + ("postgresql://litellm:12345?a=aBcD@db.internal/litellm", "aBcD"), + ("postgresql://litellm:12345?host=aBcD@db.internal/litellm", "aBcD"), + # Both '/' and '?key=value' together: the slash leaves a clean path holding + # the password remainder and the query still parses, so only the stranded + # at-sign gives it away. + ("postgresql://litellm:12345/aBcD?x=1@db.internal/litellm", "aBcD"), +) + + +@pytest.mark.parametrize(("dsn", "secret"), MISPARSED_AUTHORITY_DSNS) +def test_unencoded_slash_in_password_never_yields_an_endpoint(dsn, secret): + """An unencoded '/' truncates the authority, so urlparse reports the username + as the host and the password tail as the database. Postgres drivers reject + such a DSN outright, so the only safe reading is no endpoint at all.""" + assert parse_database_endpoint(dsn) is None + + +@pytest.mark.parametrize(("dsn", "secret"), MISPARSED_AUTHORITY_DSNS) +def test_unencoded_slash_in_password_never_reaches_a_span(dsn, secret): + attrs = _resolve("postgres", "get_data", database_url=dsn) + exported = " ".join(str(value) for value in attrs.values()) + assert secret not in exported + assert "db.namespace" not in attrs + assert "server.address" not in attrs + + +def test_extra_path_segment_yields_no_endpoint(): + """A database name cannot hold an unencoded '/', so a second path segment + means the authority was mis-split even when no '@' survived into the path.""" + assert parse_database_endpoint("postgresql://db.internal:5432/litellm/extra") is None + + +@pytest.mark.parametrize("dsn", [d for d, _ in MISPARSED_AUTHORITY_DSNS]) +def test_a_mis_split_authority_never_exports_the_database_username(dsn): + """The username lands in ``parsed.hostname`` when the authority truncates, so + a span would name the DB user as the server.""" + attrs = _resolve("postgres", "get_data", database_url=dsn) + assert "server.address" not in attrs + assert "litellm" not in " ".join(str(v) for v in attrs.values()) + + +@pytest.mark.parametrize( + "dsn", + [ + "postgresql://db.internal:5432/litellm?application_name=svc@prod", + "postgresql://db.internal:5432/litellm?user=admin@company.com", + ], +) +def test_an_unencoded_at_sign_in_a_query_forfeits_the_endpoint(dsn): + """This shape is byte-for-byte indistinguishable from a mis-split password, + so it resolves to no endpoint rather than risking a credential fragment. + Percent-encoding the at-sign restores the attributes.""" + assert parse_database_endpoint(dsn) is None + assert parse_database_endpoint(dsn.replace("@", "%40")) is not None + + +def test_host_and_port_query_parameters_are_honoured_together(): + assert parse_database_endpoint("postgresql://ignored/litellm?host=real.internal&port=6543") == DatabaseEndpoint( + address="real.internal", port=6543, namespace="litellm" + ) + + +def test_percent_encoded_password_still_resolves_the_endpoint(): + """The encoded spelling is the one a driver accepts, so it must keep working.""" + assert parse_database_endpoint("postgresql://litellm:pa%2Fssw0rd@db.internal:5432/litellm") == DatabaseEndpoint( + address="db.internal", port=5432, namespace="litellm" + ) + + +def test_hostless_socket_dsn_still_names_the_database(): + """``postgresql:///litellm`` is a valid local-socket DSN that Prisma accepts, + so the database is knowable even though no server address is.""" + assert parse_database_endpoint("postgresql:///litellm") == DatabaseEndpoint( + address=None, port=None, namespace="litellm" + ) + + +def test_hostless_socket_dsn_emits_namespace_without_a_server(): + attrs = _resolve("postgres", "get_data", database_url="postgresql:///litellm") + assert attrs["db.namespace"] == "litellm" + assert "server.address" not in attrs + assert "server.port" not in attrs + + +def test_dsn_with_neither_host_nor_database_yields_no_endpoint(): + assert parse_database_endpoint("postgresql://") is None + + +def test_prisma_default_schema_is_left_implicit(): + endpoint = parse_database_endpoint("postgresql://u:p@db.internal/litellm?schema=public") + assert endpoint is not None and endpoint.namespace == "litellm" + + +@pytest.mark.parametrize("spelling", ["PUBLIC", "Public", "reporting"]) +def test_a_non_default_schema_stays_in_the_namespace(spelling): + """Prisma quotes the schema name, so ``?schema=PUBLIC`` provisions a second + schema alongside ``public`` with its own tables. Case-folding them into one + namespace would report two different schemas as the same database.""" + endpoint = parse_database_endpoint(f"postgresql://u:p@db.internal/litellm?schema={spelling}") + assert endpoint is not None and endpoint.namespace == f"litellm|{spelling}" + + +def test_postgres_scheme_alias_is_accepted(): + assert parse_database_endpoint("postgres://u:p@db.internal/litellm") == DatabaseEndpoint( + address="db.internal", port=5432, namespace="litellm" + ) + + +@pytest.mark.parametrize( + "dsn", + [ + None, + "", + "mysql://u:p@db.internal:3306/litellm", + "postgresql://u:p@db.internal:not-a-port/litellm", + "not a url at all", + ], +) +def test_unusable_dsn_degrades_to_no_endpoint(dsn): + assert parse_database_endpoint(dsn) is None + + +def test_database_without_name_or_schema_has_no_namespace(): + assert parse_database_endpoint("postgresql://u:p@db.internal:5432/") == DatabaseEndpoint( + address="db.internal", port=5432, namespace=None + ) + + +def test_postgres_service_span_carries_system_operation_and_endpoint(): + assert _resolve("postgres", "get_data", database_url=REMOTE_DSN) == { + "db.system.name": "postgresql", + "db.system": "postgresql", + "db.operation.name": "get_data", + "server.address": "litellm-prod.abc123.us-east-1.rds.amazonaws.com", + "server.port": 6432, + "db.namespace": "litellm|reporting", + } + + +def test_legacy_db_system_is_dual_emitted_for_datadog(): + """Datadog's OTLP intake infers the database span type from ``db.system``, + not from the semconv-current ``db.system.name``.""" + assert _resolve("postgres", "get_data", database_url=LOCAL_DSN)["db.system"] == "postgresql" + assert _resolve("redis", "set")["db.system"] == "redis" + + +def test_batch_write_service_is_also_attributed_to_postgres(): + attrs = _resolve("batch_write_to_db", "_PROXY_track_cost_callback", database_url=REMOTE_DSN) + assert attrs["db.system.name"] == "postgresql" + assert attrs["server.address"] == "litellm-prod.abc123.us-east-1.rds.amazonaws.com" + + +def test_redis_service_never_borrows_the_postgres_endpoint(): + assert _resolve("redis", "set", database_url=REMOTE_DSN) == { + "db.system.name": "redis", + "db.system": "redis", + "db.operation.name": "set", + } + + +def test_non_datastore_service_gets_no_db_attributes(): + assert _resolve("reset_budget_job", "reset_budget", database_url=REMOTE_DSN) == {} + + +def test_configured_read_replica_suppresses_the_endpoint_rather_than_naming_the_primary(): + """Reads are routed to the replica per Prisma call, underneath the span, so + naming the writer would pin replica latency onto the primary.""" + attrs = _resolve("postgres", "get_data", database_url=REMOTE_DSN, read_replica_url=REPLICA_DSN) + assert attrs == { + "db.system.name": "postgresql", + "db.system": "postgresql", + "db.operation.name": "get_data", + } + + +def test_endpoint_attributes_are_omitted_when_database_url_is_unset(): + assert _resolve("postgres", "get_data") == { + "db.system.name": "postgresql", + "db.system": "postgresql", + "db.operation.name": "get_data", + } + + +def test_blank_call_type_does_not_emit_an_empty_operation_attribute(): + assert "db.operation.name" not in _resolve("postgres", "") + assert "db.operation.name" not in _resolve("postgres", None) + + +@pytest.mark.parametrize( + ("dsn", "secrets"), + [ + (LOCAL_DSN, ("dbpassword9090", "llmproxy")), + (REMOTE_DSN, ("s3cr3t", "llmproxy", "sslmode")), + (REPLICA_DSN, ("r3ad0nly", "reader")), + ], +) +def test_no_credential_reaches_any_exported_attribute(dsn, secrets): + attrs = _resolve("postgres", "get_data", database_url=dsn) + assert attrs["server.address"] + exported = " ".join(str(value) for value in attrs.values()) + for secret in secrets: + assert secret not in exported + + +def test_a_runtime_endpoint_change_is_reflected_on_the_next_span(): + """The RDS IAM refresh, the reconnect path and the DB-backed + environment_variables overlay can all rewrite DATABASE_URL after startup, so + a value cached for the process lifetime would report a server the process no + longer talks to.""" + first = _resolve("postgres", "get_data", database_url=LOCAL_DSN) + assert first["server.address"] == "localhost" + moved = _resolve("postgres", "get_data", database_url=REMOTE_DSN) + assert moved["server.address"] == "litellm-prod.abc123.us-east-1.rds.amazonaws.com" + + +def test_a_replica_configured_after_the_first_span_suppresses_the_endpoint(): + assert _resolve("postgres", "get_data", database_url=REMOTE_DSN)["server.address"] + later = _resolve("postgres", "get_data", database_url=REMOTE_DSN, read_replica_url=REPLICA_DSN) + assert "server.address" not in later diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 82b074220fa..bb2d970e9c7 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -8,6 +8,8 @@ hooks, proxy SERVER span lifecycle (start + setters), parent-context resolution import asyncio import contextlib +import os +from unittest.mock import patch from datetime import datetime, timedelta, timezone import pytest @@ -1521,6 +1523,36 @@ def test_async_service_success_hook_emits_service_span(): assert span.status.status_code is StatusCode.UNSET +def test_postgres_db_span_names_the_database_server_not_the_prisma_engine(): + """Prisma reaches Postgres over loopback, so without server.address the + backend attributes the wait to localhost.""" + dsn = "postgresql://llmproxy:dbpassword9090@litellm-prod.abc123.us-east-1.rds.amazonaws.com:6432/litellm?schema=reporting" + logger, exporter = _logger() + parent = _service_parent(logger) + try: + with patch.dict(os.environ, {"DATABASE_URL": dsn}, clear=False): + os.environ.pop("DATABASE_URL_READ_REPLICA", None) + asyncio.run( + logger.async_service_success_hook( + payload=_ServicePayload("postgres", "get_data"), + parent_otel_span=parent, + ) + ) + finally: + parent.end() + span = {s.name: s for s in exporter.get_finished_spans()}["postgres get_data"] + assert span.kind is SpanKind.CLIENT + assert span.attributes["db.system.name"] == "postgresql" + assert span.attributes["db.operation.name"] == "get_data" + assert span.attributes["server.address"] == "litellm-prod.abc123.us-east-1.rds.amazonaws.com" + assert span.attributes["server.port"] == 6432 + assert span.attributes["db.namespace"] == "litellm|reporting" + assert span.attributes["db.system"] == "postgresql" + exported = " ".join(str(value) for value in span.attributes.values()) + assert "dbpassword9090" not in exported + assert "llmproxy" not in exported + + def test_async_service_failure_hook_marks_error_status(): logger, exporter = _logger() parent = _service_parent(logger) diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index fa1c9fa8a79..28d2cb22e7c 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -34,6 +34,7 @@ from litellm.integrations.opentelemetry import ( _normalize_team_metadata_keys, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.types.services import ServiceLoggerPayload, ServiceTypes class TestOpenTelemetryGuardrails(unittest.TestCase): @@ -6212,3 +6213,110 @@ class TestDynamicTracerProviderCache(unittest.TestCase): self.assertTrue(entry.owns_exporter) self.assertIsNotNone(entry.provider._atexit_handler) +class TestOpenTelemetryDatabaseSemconvAttributes(unittest.TestCase): + """A Postgres service span must name the PostgreSQL server it reached. + + Without ``db.system`` and ``server.address``, the only host in the trace is + the loopback address of Prisma's local query engine, so the backend + attributes the wait to ``localhost`` and it cannot be correlated with the + database's own metrics. + """ + + DSN = "postgresql://llmproxy:dbpassword9090@litellm-prod.abc123.us-east-1.rds.amazonaws.com:6432/litellm?schema=reporting" + REPLICA_DSN = "postgresql://reader:r3ad0nly@litellm-prod-ro.abc123.us-east-1.rds.amazonaws.com/litellm" + + def _service_span(self, service, call_type, dsn, error=None, replica_dsn=None): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + otel = OpenTelemetry() + otel.tracer = provider.get_tracer(__name__) + parent = otel.tracer.start_span("Received Proxy Server Request") + payload = ServiceLoggerPayload( + is_error=error is not None, + error=error, + service=service, + duration=0.25, + call_type=call_type, + event_metadata=None, + ) + hook = otel.async_service_failure_hook if error else otel.async_service_success_hook + kwargs = {"error": error} if error else {} + env = {k: v for k, v in (("DATABASE_URL", dsn), ("DATABASE_URL_READ_REPLICA", replica_dsn)) if v} + with patch.dict(os.environ, env, clear=False): + for absent in {"DATABASE_URL", "DATABASE_URL_READ_REPLICA"} - set(env): + os.environ.pop(absent, None) + asyncio.run( + hook( + payload=payload, + parent_otel_span=parent, + start_time=datetime.now(), + end_time=datetime.now(), + **kwargs, + ) + ) + parent.end() + return next(s for s in exporter.get_finished_spans() if s.name == service.value) + + def test_postgres_span_names_the_database_server(self): + span = self._service_span(ServiceTypes.DB, "get_data", self.DSN) + self.assertEqual(span.attributes["db.system.name"], "postgresql") + self.assertEqual(span.attributes["db.operation.name"], "get_data") + self.assertEqual( + span.attributes["server.address"], + "litellm-prod.abc123.us-east-1.rds.amazonaws.com", + ) + self.assertEqual(span.attributes["server.port"], 6432) + self.assertEqual(span.attributes["db.namespace"], "litellm|reporting") + + def test_datastore_span_is_a_client_span_carrying_the_legacy_db_system(self): + """Datadog types a span as a database call from CLIENT kind plus + ``db.system``; an INTERNAL span is classified as custom work.""" + span = self._service_span(ServiceTypes.DB, "get_data", self.DSN) + self.assertEqual(span.kind, trace.SpanKind.CLIENT) + self.assertEqual(span.attributes["db.system"], "postgresql") + + def test_internal_service_span_stays_internal(self): + span = self._service_span(ServiceTypes.RESET_BUDGET_JOB, "reset_budget", self.DSN) + self.assertEqual(span.kind, trace.SpanKind.INTERNAL) + self.assertNotIn("db.system.name", span.attributes) + self.assertNotIn("server.address", span.attributes) + + def test_existing_service_and_call_type_attributes_are_unchanged(self): + span = self._service_span(ServiceTypes.DB, "get_data", self.DSN) + self.assertEqual(span.attributes["service"], "postgres") + self.assertEqual(span.attributes["call_type"], "get_data") + + def test_failed_postgres_span_also_names_the_database_server(self): + span = self._service_span(ServiceTypes.DB, "get_data", self.DSN, error="connection refused") + self.assertEqual(span.attributes["db.system.name"], "postgresql") + self.assertEqual(span.kind, trace.SpanKind.CLIENT) + self.assertEqual( + span.attributes["server.address"], + "litellm-prod.abc123.us-east-1.rds.amazonaws.com", + ) + self.assertEqual(span.attributes["error"], "connection refused") + + def test_no_credential_from_the_dsn_lands_on_the_span(self): + span = self._service_span(ServiceTypes.DB, "get_data", self.DSN) + exported = " ".join(str(value) for value in span.attributes.values()) + self.assertIn("litellm-prod.abc123.us-east-1.rds.amazonaws.com", exported) + self.assertNotIn("dbpassword9090", exported) + self.assertNotIn("llmproxy", exported) + + def test_redis_span_does_not_borrow_the_postgres_endpoint(self): + span = self._service_span(ServiceTypes.REDIS, "async_set_cache", self.DSN) + self.assertEqual(span.attributes["db.system.name"], "redis") + self.assertEqual(span.kind, trace.SpanKind.CLIENT) + self.assertNotIn("server.address", span.attributes) + + def test_configured_read_replica_suppresses_the_endpoint(self): + span = self._service_span(ServiceTypes.DB, "get_data", self.DSN, replica_dsn=self.REPLICA_DSN) + self.assertEqual(span.attributes["db.system.name"], "postgresql") + self.assertNotIn("server.address", span.attributes) + self.assertNotIn("db.namespace", span.attributes) + + def test_unset_database_url_leaves_the_span_without_endpoint_attributes(self): + span = self._service_span(ServiceTypes.DB, "get_data", None) + self.assertEqual(span.attributes["db.system.name"], "postgresql") + self.assertNotIn("server.address", span.attributes) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index af40245ebfa..f9311497729 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -241,10 +241,32 @@ def test_split_concatenated_json_non_dict_value(): assert result == [{}] -def test_split_concatenated_json_invalid_raises(): - """Completely invalid JSON raises JSONDecodeError.""" - with pytest.raises(json.JSONDecodeError): - split_concatenated_json_objects("not json at all") +def test_split_concatenated_json_wholly_invalid_returns_empty(): + """ + Wholly unparseable JSON degrades to an empty list instead of raising. + + Regression for https://github.com/BerriAI/litellm/issues/18667: a raise + here propagated out of `_convert_to_bedrock_tool_call_invoke` and turned + every replayed conversation into a 500. + """ + assert split_concatenated_json_objects("not json at all") == [] + + +def test_split_concatenated_json_malformed_object_returns_empty(): + """ + A single malformed object (missing comma between keys) degrades to an + empty list rather than raising `Expecting ',' delimiter`. + """ + assert split_concatenated_json_objects('{"location": "Boston" "unit": "celsius"}') == [] + + +def test_split_concatenated_json_salvages_prefix_before_truncated_tail(): + """ + Complete objects parsed before an unparseable/truncated tail are kept; + only the bad tail is discarded. + """ + result = split_concatenated_json_objects('{"a": 1}{"b": 2}{"c":') + assert result == [{"a": 1}, {"b": 2}] # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index de5d0a180c6..fffbc884782 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -2287,6 +2287,116 @@ def test_bedrock_tool_call_invoke_non_dict_arguments(): assert result[0]["toolUse"]["input"] == {} +def test_bedrock_tool_call_invoke_malformed_json_does_not_raise(): + """ + Regression for https://github.com/BerriAI/litellm/issues/18667. + + When the model emits malformed JSON in tool-call arguments (here a + missing comma between keys), replaying that history must NOT raise + `Unable to convert openai tool calls ... Expecting ',' delimiter`. + It degrades to an empty-object input so the conversation can continue. + """ + tool_calls = [ + { + "id": "toolu_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "Boston" "unit": "celsius"}', + }, + } + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + assert len(result) == 1 + assert result[0]["toolUse"]["toolUseId"] == "toolu_abc123" + assert result[0]["toolUse"]["name"] == "get_weather" + assert result[0]["toolUse"]["input"] == {} + + +def test_bedrock_tool_call_invoke_salvages_valid_prefix_before_truncated_tail(): + """ + A valid leading object followed by a truncated tail keeps the valid + object rather than dropping everything or raising. + """ + tool_calls = [ + { + "id": "call_partial", + "type": "function", + "function": {"name": "shell", "arguments": '{"cmd": "ls"}{"cmd":'}, + } + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + assert len(result) == 1 + assert result[0]["toolUse"]["input"] == {"cmd": "ls"} + + +def test_bedrock_tool_call_invoke_mixed_turn_survives_one_malformed_call(): + """ + Regression for LIT-4574: an assistant turn with several tool calls where only one + has malformed/truncated arguments must keep the valid calls intact and degrade just + the bad one to empty input, instead of killing the entire turn. + """ + tool_calls = [ + { + "id": "t_good", + "type": "function", + "function": { + "name": "good_tool", + "arguments": '{"item_type": "email", "item_id": "AAMkAD=="}', + }, + }, + { + "id": "t_bad", + "type": "function", + "function": {"name": "bad_tool", "arguments": '{"item_type": "email"'}, + }, + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + tool_uses = [block["toolUse"] for block in result if "toolUse" in block] + assert len(tool_uses) == 2 + by_name = {tool_use["name"]: tool_use for tool_use in tool_uses} + assert by_name["good_tool"]["input"] == {"item_type": "email", "item_id": "AAMkAD=="} + assert by_name["bad_tool"]["input"] == {} + + +def test_bedrock_tool_call_invoke_truncated_json_arguments(): + """ + Truncated tool call arguments (issue #35303) must not raise. A client replaying a + partially streamed tool call would otherwise trigger a pre-network exception that the + router maps to a retryable APIConnectionError and retries through the fallback graph. + """ + tool_calls = [ + { + "id": "tooluse_MAh2QLVjBRkvi5QJkLQ08V", + "type": "function", + "function": { + "name": "replace_note_content", + "arguments": '{"note_id": "999af35c-4061-4ece-8581-7d43fc988ba4", "title": "WG"', + }, + } + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + assert len(result) == 1 + assert result[0]["toolUse"]["toolUseId"] == "tooluse_MAh2QLVjBRkvi5QJkLQ08V" + assert result[0]["toolUse"]["input"] == {} + + +def test_bedrock_tool_call_invoke_unconvertible_raises_non_retryable_bad_request(): + """ + Conversion failures are client input errors, so they must surface as a non-retryable + BadRequestError instead of a bare Exception that maps to APIConnectionError, and the + message must not embed the tool call payload (issue #35303). + """ + tool_calls = [{"id": "call_bad", "type": "function", "function": None}] + + with pytest.raises(litellm.BadRequestError) as exc_info: + _convert_to_bedrock_tool_call_invoke(tool_calls) + + assert exc_info.value.status_code == 400 + assert "call_bad" in str(exc_info.value) + assert "function" not in str(exc_info.value).split("Received error=")[0] + + def test_make_valid_bedrock_tool_name_preserves_hyphens(): assert make_valid_bedrock_tool_name("my-tool") == "my-tool" assert ( diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py index 6f5aaabae06..510776ddfdf 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -271,6 +271,47 @@ class TestDashscopeCostCalculator: assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + def test_dashscope_nested_cache_creation_input_tokens_bill_at_cache_write_rate(self): + """ + Regression (LIT-5757): DashScope nests cache_creation_input_tokens inside + prompt_tokens_details; those tokens must bill at the tier's cache-creation + rate instead of being folded into text tokens at the input rate. + """ + self._register_tiered_model( + "dashscope/qwen-nested-cache-write-test", + [ + { + "range": [0, 128000], + "input_cost_per_token": 4e-07, + "cache_read_input_token_cost": 1.6e-07, + "cache_creation_input_token_cost": 5e-07, + "output_cost_per_token": 1.6e-06, + } + ], + ) + + usage = Usage( + prompt_tokens=2059, + completion_tokens=201, + total_tokens=2260, + prompt_tokens_details={ + "cached_tokens": 0, + "text_tokens": 2059, + "cache_type": "ephemeral", + "cache_creation_input_tokens": 2048, + "cache_creation": {"ephemeral_5m_input_tokens": 2048}, + }, + completion_tokens_details={"reasoning_tokens": 170}, + ) + + prompt_cost, _ = dashscope_cost_per_token( + model="qwen-nested-cache-write-test", usage=usage + ) + + assert math.isclose( + prompt_cost, (2048 * 5e-07) + (11 * 4e-07), rel_tol=1e-10 + ) + def test_dashscope_tiered_cache_creation_falls_back_to_tier_input_rate(self): """ Tiers without a cache_creation_input_token_cost bill cache-creation tokens at diff --git a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py index 58363e3baea..2dcccb8ea7e 100644 --- a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py +++ b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py @@ -47,7 +47,9 @@ def _make_mock_response( mock = MagicMock() mock.status_code = status_code - mock.headers = headers or {} + # httpx.Headers normalizes keys to lowercase — mirror production so tests + # assert what callers actually see. + mock.headers = httpx.Headers(headers or {}) if json_data is not None: mock.json.return_value = json_data mock.text = text if text is not None else _json.dumps(json_data) @@ -222,7 +224,7 @@ class TestTransformSearchRequest: assert param not in result["_tinyfish_params"] def test_arbitrary_param_passed_through(self): - # `fetch` is a TinyFish-specific param (JSON-encoded tf-fetch config). + # `fetch` is a TinyFish-specific param (JSON-encoded fetch config). # The passthrough loop should forward it verbatim without LiteLLM needing # to know about it. config = TinyfishSearchConfig() @@ -237,26 +239,49 @@ class TestTransformSearchRequest: config = TinyfishSearchConfig() result = config.transform_search_request( query="test", - optional_params={"fetch": {"format": "html", "fetch_path": "fast"}}, - ) - assert ( - result["_tinyfish_params"]["fetch"] - == '{"format":"html","fetch_path":"fast"}' + optional_params={"fetch": {"format": "html"}}, ) + assert result["_tinyfish_params"]["fetch"] == '{"format":"html"}' def test_bool_param_serialized_as_lowercase(self): - # urlencode renders Python bool as capitalized "True"/"False"; ux-labs - # rejects those (e.g. include_thumbnail must be literal "true"/"false"). - # Normalize before passing through. + # urlencode renders Python bool as capitalized "True"/"False"; TinyFish + # Search's bool params require lowercase "true"/"false" strings on the + # wire. Normalize before passing through. config = TinyfishSearchConfig() true_result = config.transform_search_request( - query="test", optional_params={"include_thumbnail": True} + query="test", optional_params={"some_bool_param": True} ) false_result = config.transform_search_request( - query="test", optional_params={"include_thumbnail": False} + query="test", optional_params={"some_bool_param": False} ) - assert true_result["_tinyfish_params"]["include_thumbnail"] == "true" - assert false_result["_tinyfish_params"]["include_thumbnail"] == "false" + assert true_result["_tinyfish_params"]["some_bool_param"] == "true" + assert false_result["_tinyfish_params"]["some_bool_param"] == "false" + + def test_float_param_passes_through(self): + # Float values pass the urlencode adapter and land on the wire as + # their decimal string form. If TinyFish's server rejects a float + # for a param it expects as int, the server's 400 response is + # attributed via _wrap_error (`TinyFish Search: ...`) — better than + # a client-side pydantic ValidationError with no context. + config = TinyfishSearchConfig() + result = config.transform_search_request( + query="test", + optional_params={"some_float_param": 0.5}, + ) + assert result["_tinyfish_params"]["some_float_param"] == 0.5 + + def test_list_param_auto_json_encoded(self): + # TinyFish Search's JSON-array params arrive on the wire as JSON- + # encoded strings. Accept the natural Python list form and serialize + # so the caller doesn't have to pre-stringify. Params whose wire + # format is a plain comma-separated string are the caller's + # responsibility to pass as a Python str. + config = TinyfishSearchConfig() + result = config.transform_search_request( + query="test", + optional_params={"some_list_param": ["a.example", "b.example"]}, + ) + assert result["_tinyfish_params"]["some_list_param"] == '["a.example","b.example"]' def test_pre_stringified_param_passed_unchanged(self): # If the caller already JSON-encoded, don't re-encode. @@ -422,10 +447,104 @@ class TestTransformSearchResponse: assert getattr(first, "position", None) == 1 assert getattr(first, "site_name", None) == "tinyfish.ai" + def test_top_level_extras_flow_through(self): + # TinyFish returns `query`, `total_results`, `page` at the envelope + # level. These must ride through to the caller via SearchResponse's + # extra="allow" so pagination logic, echo checks, etc. work. + config = TinyfishSearchConfig() + mock_response = _make_mock_response(MOCK_TINYFISH_RESPONSE) + result = config.transform_search_response( + raw_response=mock_response, logging_obj=None + ) + assert getattr(result, "query", None) == "web automation tools" + assert getattr(result, "total_results", None) == 2 + assert getattr(result, "page", None) == 0 + + def test_top_level_future_extras_flow_through(self): + # Any future TinyFish top-level field must ride through unchanged + # (design contract: no LiteLLM code change needed for new fields). + config = TinyfishSearchConfig() + body = { + "results": [ + {"title": "x", "url": "https://x", "snippet": "x"}, + ], + "query": "test", + "example_int_extra": 123, # hypothetical future field + "example_str_extra": "value", # hypothetical future field + "example_id_extra": "abc-def", # hypothetical future field + } + result = config.transform_search_response( + raw_response=_make_mock_response(body), logging_obj=None + ) + assert getattr(result, "example_int_extra", None) == 123 + assert getattr(result, "example_str_extra", None) == "value" + assert getattr(result, "example_id_extra", None) == "abc-def" + + def test_response_headers_stashed_on_hidden_params(self): + # TinyFish Search sets X-Request-ID on every success response. Confirm it + # lands on both `_hidden_params["headers"]` (raw) and + # `_hidden_params["additional_headers"]` (sanitized/prefixed). + # httpx.Headers lowercases every key, so assertions use lowercase. + config = TinyfishSearchConfig() + mock_response = _make_mock_response( + MOCK_TINYFISH_RESPONSE, + headers={"X-Request-ID": "req-abc-123", "Content-Type": "application/json"}, + ) + result = config.transform_search_response( + raw_response=mock_response, logging_obj=None + ) + # Raw copy — httpx has normalized keys to lowercase. + assert result._hidden_params["headers"]["x-request-id"] == "req-abc-123" + # process_response_headers prefixes non-OpenAI-standard keys with "llm_provider-". + assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-abc-123" + + def test_response_headers_future_headers_flow_through(self): + # "Accept extra": any header TinyFish Search adds later must ride + # through without a LiteLLM code change. + config = TinyfishSearchConfig() + mock_response = _make_mock_response( + MOCK_TINYFISH_RESPONSE, + headers={ + "X-Request-ID": "req-1", + "X-Example-Header-A": "value-a", # hypothetical future header + "X-Example-Header-B": "value-b", # hypothetical future header + }, + ) + result = config.transform_search_response( + raw_response=mock_response, logging_obj=None + ) + raw = result._hidden_params["headers"] + # httpx lowercases header names on read. + assert raw["x-example-header-a"] == "value-a" + assert raw["x-example-header-b"] == "value-b" + + def test_response_headers_strips_x_litellm_spoof(self): + # A provider setting `x-litellm-*` in its response must not be able to + # spoof LiteLLM-internal markers via _hidden_params["additional_headers"]. + # The raw copy preserves the header (opt-in debug view); the sanitized + # copy prefixes it with `llm_provider-` so bare `x-litellm-*` markers + # can't be spoofed (values still survive under the prefixed key for + # observability). + config = TinyfishSearchConfig() + mock_response = _make_mock_response( + MOCK_TINYFISH_RESPONSE, + headers={"x-litellm-attempted-fallbacks": "spoofed", "X-Request-ID": "r1"}, + ) + result = config.transform_search_response( + raw_response=mock_response, logging_obj=None + ) + # Raw view still has the spoof. + assert result._hidden_params["headers"]["x-litellm-attempted-fallbacks"] == "spoofed" + # Sanitized view: the spoof survives only under the llm_provider- prefix + # (never under the bare x-litellm-* key that LiteLLM downstream trusts). + additional = result._hidden_params["additional_headers"] + assert "x-litellm-attempted-fallbacks" not in additional + assert additional.get("llm_provider-x-litellm-attempted-fallbacks") == "spoofed" + def test_fetch_field_rides_through_to_search_result(self): - # Mirrors browser-search's per-result `fetch` nested object (see - # api/src/parser.rs SearchResult.fetch). Confirms `fetch=...` requests - # surface their content to LiteLLM callers without provider changes. + # Mirrors TinyFish Search's per-result `fetch` nested object. + # Confirms `fetch=...` requests surface their content to LiteLLM + # callers without provider changes. config = TinyfishSearchConfig() fetched = { "results": [ @@ -568,7 +687,7 @@ class TestTransformSearchResponse: class TestErrorHandling: def test_4xx_response_raises_with_attribution_and_unwrapped_message(self): - # Reproduces ux-labs' error envelope shape for an INVALID_INPUT response. + # Reproduces TinyFish Search's error envelope shape for an INVALID_INPUT response. config = TinyfishSearchConfig() body = { "error": { @@ -590,7 +709,7 @@ class TestErrorHandling: def test_429_preserves_status_code_and_headers(self): config = TinyfishSearchConfig() - body = {"error": {"code": "RATE_LIMIT_EXCEEDED", "message": "60 rpm"}} + body = {"error": {"code": "RATE_LIMIT_EXCEEDED", "message": "rate limit exceeded"}} mock_response = _make_mock_response( body, status_code=429, headers={"Retry-After": "60"} ) @@ -600,10 +719,12 @@ class TestErrorHandling: ) assert getattr(exc_info.value, "status_code", None) == 429 headers = getattr(exc_info.value, "headers", {}) or {} - assert headers.get("Retry-After") == "60" + # httpx lowercases; the exception carries the same dict shape. + assert headers.get("retry-after") == "60" - def test_5xx_with_non_ux_labs_body_falls_back_to_raw_text(self): - # Cloudflare-style JSON or any other envelope: unwrap fails, fall back to raw. + def test_5xx_with_non_tinyfish_envelope_shape_falls_back_to_raw_text(self): + # A JSON body that doesn't match TinyFish Search's error envelope shape: + # unwrap fails, fall back to the raw body text. config = TinyfishSearchConfig() body = {"errors": [{"code": "10000", "message": "Internal"}]} mock_response = _make_mock_response(body, status_code=502) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 412e6f0f83c..62c05841197 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -14,6 +14,7 @@ sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory from litellm.proxy.management_endpoints.common_daily_activity import ( _adjust_dates_for_timezone, _build_aggregated_sql_query, + _build_entity_rollup_sql_query, _is_user_agent_tag, _record_to_spend_metrics, get_api_key_metadata, @@ -982,6 +983,58 @@ class TestBuildAggregatedSqlQuery: assert "COALESCE(model_group, model)" not in normalized +class TestAggregatedEmptyEntityFilter: + _BUILDERS: Final = (_build_aggregated_sql_query, _build_entity_rollup_sql_query) + + @pytest.mark.parametrize("build", _BUILDERS) + def test_empty_entity_list_emits_no_degenerate_in_clause(self, build): + sql, params = build( + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id=[], + start_date="2026-08-01", + end_date="2026-08-19", + model=None, + api_key=None, + ) + + normalized = " ".join(sql.split()) + assert "IN ()" not in normalized + assert '"team_id" IN' not in normalized + assert params == ["2026-08-01", "2026-08-19"] + + @pytest.mark.parametrize("build", _BUILDERS) + def test_empty_entity_list_matches_nothing_rather_than_everything(self, build): + sql, _ = build( + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id=[], + start_date="2026-08-01", + end_date="2026-08-19", + model=None, + api_key=None, + ) + + assert "FALSE" in " ".join(sql.split()) + + @pytest.mark.parametrize("build", _BUILDERS) + def test_populated_entity_list_still_filters_on_its_ids(self, build): + sql, params = build( + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id=["team-alpha", "team-beta"], + start_date="2026-08-01", + end_date="2026-08-19", + model=None, + api_key=None, + ) + + normalized = " ".join(sql.split()) + assert '"team_id" IN ($3, $4)' in normalized + assert "FALSE" not in normalized + assert params == ["2026-08-01", "2026-08-19", "team-alpha", "team-beta"] + + @pytest.mark.asyncio async def test_get_daily_activity_aggregated_empty_result_set(): """Regression test for the empty-range 500. 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 664015003e4..69819318800 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 @@ -1286,6 +1286,75 @@ class TestOpenAIPassthroughIntegration: mock_chat_handler.assert_called_once() assert result == {"result": None, "kwargs": {}} + def test_openai_passthrough_handler_embeddings_unmapped_model_logs_zero_cost(self): + response_body = { + "object": "list", + "model": "lit5787-unmapped-embeddings-deployment", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "usage": {"prompt_tokens": 9, "total_tokens": 9}, + } + mock_logging_obj = self._create_mock_logging_obj() + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=self._create_mock_httpx_response(response_body), + response_body=response_body, + logging_obj=mock_logging_obj, + url_route="https://my-resource.openai.azure.com/openai/v1/embeddings", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={ + "model": "lit5787-unmapped-embeddings-deployment", + "input": "spend probe", + }, + passthrough_logging_payload=PassthroughStandardLoggingPayload( + url="https://my-resource.openai.azure.com/openai/v1/embeddings", + request_body={ + "model": "lit5787-unmapped-embeddings-deployment", + "input": "spend probe", + }, + request_method="POST", + ), + litellm_params={}, + ) + + assert result["result"] is not None + assert result["result"].usage.prompt_tokens == 9 + assert result["kwargs"]["response_cost"] == 0.0 + assert result["kwargs"]["model"] == "lit5787-unmapped-embeddings-deployment" + assert result["result"]._hidden_params["response_cost"] == 0.0 + assert mock_logging_obj.model_call_details["response_cost"] == 0.0 + + def test_openai_passthrough_handler_embeddings_error_skips_chat_fallback(self): + response_body = { + "object": "list", + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 9, "total_tokens": 9}, + } + kwargs_in = { + "passthrough_logging_payload": PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/embeddings", + request_body={"model": "text-embedding-3-small", "input": "spend probe"}, + request_method="POST", + ), + "litellm_params": {}, + } + 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={"model": "text-embedding-3-small", "input": "spend probe"}, + **kwargs_in, + ) + + assert result["result"] is None + assert result["kwargs"]["passthrough_logging_payload"] == kwargs_in["passthrough_logging_payload"] + @patch( "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler" ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 33d835652cd..e5add059260 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -124,6 +124,30 @@ def test_get_logging_payload_maps_openai_cache_write_tokens_to_cache_creation_in assert additional_usage_values["prompt_tokens_details"]["cache_write_tokens"] == 800 +def test_get_logging_payload_maps_nested_cache_creation_input_tokens(): + """ + Regression (LIT-5757): DashScope nests cache_creation_input_tokens inside + prompt_tokens_details; SpendLogs must record it as cache_creation_input_tokens. + """ + additional_usage_values: Final = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=2059, + completion_tokens=31, + total_tokens=2090, + prompt_tokens_details={ + "cached_tokens": 0, + "text_tokens": 2059, + "cache_type": "ephemeral", + "cache_creation_input_tokens": 2048, + "cache_creation": {"ephemeral_5m_input_tokens": 2048}, + }, + ) + ) + + assert additional_usage_values["cache_creation_input_tokens"] == 2048 + assert additional_usage_values["prompt_tokens_details"]["cache_write_tokens"] == 2048 + + def test_get_logging_payload_preserves_anthropic_cache_creation_input_tokens(): additional_usage_values = _get_additional_usage_values_for_usage( litellm.Usage( diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 68395737469..24477248a8a 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -1,4 +1,5 @@ import json +from typing import NoReturn from unittest.mock import MagicMock, patch import httpx @@ -167,6 +168,10 @@ async def _acreate_batch(*args, **kwargs): raise AssertionError("only used for its __name__") +async def _acreate_file(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + @pytest.mark.asyncio async def test_run_async_fallback_keeps_uploaded_file_requests_in_their_model_group(): """An input_file_id only exists under the credentials of the group it was uploaded @@ -229,6 +234,46 @@ async def test_run_async_fallback_allows_same_model_group_retry_for_uploaded_fil assert router.attempted_model_groups == ["openai-group"] +@pytest.mark.asyncio +async def test_run_async_fallback_keeps_file_creation_in_its_model_group(): + """A file created for batches lands in the account of the deployment that stored it, + and its id is only usable against the model group the caller named. A cross-group + fallback silently stores the file with the wrong provider.""" + router = AttemptRecordingRouter() + + with pytest.raises(RuntimeError, match="azure connection error"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["openai-group"], + original_model_group="azure-group", + original_exception=RuntimeError("azure connection error"), + max_fallbacks=3, + fallback_depth=0, + model="azure-group", + original_function=_acreate_file, + ) + + assert router.attempted_model_groups == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_allows_same_model_group_retry_for_file_creation(): + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[{"model": "azure-group", "_target_order": 2}], + original_model_group="azure-group", + original_exception=RuntimeError("first deployment failed"), + max_fallbacks=3, + fallback_depth=0, + model="azure-group", + original_function=_acreate_file, + ) + + assert router.attempted_model_groups == ["azure-group"] + + @pytest.mark.asyncio async def test_run_async_fallback_still_crosses_model_groups_without_an_uploaded_file(): router = AttemptRecordingRouter() diff --git a/tests/test_litellm/test_responses_api_bridge_non_stream.py b/tests/test_litellm/test_responses_api_bridge_non_stream.py index 25a3bc2dbba..c272b151865 100644 --- a/tests/test_litellm/test_responses_api_bridge_non_stream.py +++ b/tests/test_litellm/test_responses_api_bridge_non_stream.py @@ -1,6 +1,6 @@ import os import sys -from typing import Optional +from typing import Final, Optional from unittest.mock import Mock import pytest @@ -192,6 +192,32 @@ def test_transform_usage_with_cached_tokens_only(): print("✓ Transformation works with cached_tokens only") +def test_transform_usage_maps_nested_cache_creation_input_tokens(): + """ + Regression (LIT-5757): DashScope nests cache_creation_input_tokens inside + prompt_tokens_details; the bridge must surface it as cache_write_tokens. + """ + usage: Final = Usage( + prompt_tokens=2059, + completion_tokens=31, + total_tokens=2090, + prompt_tokens_details={ + "cached_tokens": 0, + "text_tokens": 2059, + "cache_type": "ephemeral", + "cache_creation_input_tokens": 2048, + "cache_creation": {"ephemeral_5m_input_tokens": 2048}, + }, + ) + + responses_usage: Final = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( + usage + ) + + assert responses_usage.input_tokens_details is not None + assert responses_usage.input_tokens_details.cache_write_tokens == 2048 + + def test_transform_usage_with_reasoning_tokens_only(): """ Test transformation when only reasoning_tokens is provided (no cached_tokens). diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 49ed236356c..d2d5d9268be 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -492,6 +492,54 @@ async def test_async_router_acreate_file_with_jsonl(): assert first_call_content == non_jsonl_content +@pytest.mark.asyncio +async def test_async_router_acreate_file_does_not_fall_back_across_model_groups(): + """A file created for batches only exists under the credentials of the model group + the caller named. A cross-group fallback silently stores it with the wrong provider + and the later batch create against the named group permanently fails.""" + from unittest.mock import MagicMock, patch + + router = litellm.Router( + model_list=[ + { + "model_name": "azure-gpt", + "litellm_params": { + "model": "azure/my-azure-deployment", + "api_base": "http://127.0.0.1:9", + "api_key": "dummy-key", + "api_version": "2024-06-01", + }, + }, + { + "model_name": "openai-gpt", + "litellm_params": {"model": "gpt-4o-mini"}, + }, + ], + fallbacks=[{"azure-gpt": ["openai-gpt"]}], + ) + + def fail_azure(*args: object, **kwargs: object) -> MagicMock: + if kwargs.get("model") == "azure/my-azure-deployment": + raise litellm.APIConnectionError( + message="Connection error.", + llm_provider="azure", + model="azure/my-azure-deployment", + ) + return MagicMock() + + with patch("litellm.acreate_file", side_effect=fail_azure) as mock_acreate_file: + with pytest.raises(litellm.APIConnectionError): + await router.acreate_file( + model="azure-gpt", + purpose="batch", + file=MagicMock(), + ) + + called_models = [call.kwargs.get("model") for call in mock_acreate_file.call_args_list] + assert "azure/my-azure-deployment" in called_models + assert "gpt-4o-mini" not in called_models + + @pytest.mark.asyncio async def test_async_router_acreate_file_uses_deployment_custom_llm_provider(): """ @@ -7794,6 +7842,21 @@ class TestTaggedAutoRouterOnSharedModelName: assert response is not None assert response.model == "gemini-flash" + @pytest.mark.asyncio + async def test_request_level_tag_filtering_from_key_settings_bypasses_the_marker(self): + router = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=False) + + assert await self._hook_response(router, {"enable_tag_filtering": True}) is None + + @pytest.mark.asyncio + async def test_globally_disabled_filtering_still_lets_the_marker_capture_untagged_requests(self): + router = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=False) + + response = await self._hook_response(router, {}) + + assert response is not None + assert response.model == "gemini-flash" + @pytest.mark.asyncio async def test_marker_only_alias_still_captures_untagged_requests(self): router = self._router(marker_tags=["route"], include_plain_sibling=False, enable_tag_filtering=True) diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index cd5e8dda012..672aa84cc73 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -1,5 +1,6 @@ import os import sys +from typing import Final import pytest @@ -75,6 +76,29 @@ def test_usage_dump(): assert new_usage.prompt_tokens_details.web_search_requests == 1 +def test_prompt_tokens_details_maps_nested_cache_creation_input_tokens(): + """Regression (LIT-5757): DashScope nests the Anthropic-spelled + cache_creation_input_tokens inside prompt_tokens_details. It must populate + the canonical cache_write_tokens/cache_creation_tokens pair, without + overriding an explicitly provided canonical value.""" + from litellm.types.utils import PromptTokensDetailsWrapper + + nested: Final = PromptTokensDetailsWrapper( + cached_tokens=0, text_tokens=2059, cache_creation_input_tokens=2048 + ) + assert nested.cache_write_tokens == 2048 + assert nested.cache_creation_tokens == 2048 + + explicit: Final = PromptTokensDetailsWrapper( + cache_write_tokens=100, cache_creation_input_tokens=2048 + ) + assert explicit.cache_write_tokens == 100 + assert explicit.cache_creation_tokens == 100 + + non_int: Final = PromptTokensDetailsWrapper(cache_creation_input_tokens=None) + assert not hasattr(non_int, "cache_write_tokens") + + def test_usage_server_tool_use_dict_is_coerced_and_round_trips(): from litellm.types.utils import ServerToolUse, Usage diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index dfaa417cae3..3f32048cac1 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2328,14 +2328,6 @@ "count": 1 } }, - "src/components/organisms/create_key_button.test.tsx": { - "@typescript-eslint/no-require-imports": { - "count": 1 - }, - "react/display-name": { - "count": 8 - } - }, "src/components/organisms/create_key_button.tsx": { "local/filename-pascal-case": { "count": 1 @@ -2636,15 +2628,6 @@ "src/components/templates/key_edit_view.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "local/no-complex-jsx-arrow": { - "count": 2 - }, - "no-nested-ternary": { - "count": 2 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/templates/key_info_view.tsx": { @@ -3031,4 +3014,4 @@ "count": 1 } } -} \ No newline at end of file +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/TagsInput.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/TagsInput.tsx index 0d799474399..6fb1508497e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/TagsInput.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/TagsInput.tsx @@ -28,6 +28,7 @@ interface TagsInputProps { emptyText?: string; tokenSeparators?: string[]; loading?: boolean; + disabled?: boolean; id?: string; } @@ -48,6 +49,7 @@ export const TagsInput = ({ emptyText = "No matching options", tokenSeparators = [], loading = false, + disabled = false, id, }: TagsInputProps) => { const anchor = useComboboxAnchor(); @@ -103,6 +105,7 @@ export const TagsInput = ({ itemToStringLabel={(option: TagsInputOption) => option.label} filter={matchesQuery} openOnInputClick + disabled={disabled || loading} > } className="min-h-8 py-1 text-sm"> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index 320e0b5baef..b6ee130d50a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -1,5 +1,6 @@ import * as networking from "@/components/networking"; import { fireEvent, render, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { afterEach, describe, expect, it, vi } from "vitest"; import GuardrailInfoView from "./guardrail_info"; @@ -136,7 +137,7 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { getByText, container } = render( + const { getByText, findByText, container } = render( {}} accessToken="123" isAdmin={true} />, ); @@ -152,18 +153,9 @@ describe("Guardrail Info", () => { expect(getByText("Guardrail Settings")).toBeInTheDocument(); }); - // Find the info icon and hover over it - const infoIcon = within(container).getByRole("img", { name: "info-circle" }); - expect(infoIcon).toBeInTheDocument(); + await userEvent.hover(within(container).getByRole("img", { name: "info-circle" })); - if (infoIcon) { - fireEvent.mouseEnter(infoIcon); - - // Wait for the tooltip to appear - await waitFor(() => { - expect(getByText("Guardrail is defined in the config file and cannot be edited.")).toBeInTheDocument(); - }); - } + expect(await findByText("Guardrail is defined in the config file and cannot be edited.")).toBeInTheDocument(); }); it("should render the guardrail info", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index caac09a1c83..e80ddac932f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -6,7 +6,8 @@ import { } from "@/components/networking"; import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils"; import { CodeOutlined, EyeInvisibleOutlined, InfoCircleOutlined, StopOutlined } from "@ant-design/icons"; -import { Button as AntdButton, Tooltip } from "antd"; + +import { Button as AntdButton } from "antd"; import { ArrowLeft, CheckIcon, CopyIcon } from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Card } from "@/components/ui/card"; @@ -21,7 +22,7 @@ import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Separator } from "@/components/ui/separator"; import { Textarea } from "@/components/ui/textarea"; -import { TooltipProvider } from "@/components/ui/tooltip"; +import { SimpleTooltip, TooltipProvider } from "@/components/ui/tooltip"; import { asText, GuardrailField, @@ -663,9 +664,9 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Guardrail Settings

{isConfigGuardrail && ( - + - + )} {!isEditing && !isConfigGuardrail && diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx index 6db19c6d090..5c561f38197 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx @@ -1,7 +1,8 @@ "use client"; import React, { useState, useEffect } from "react"; -import { Tooltip, Button as AntdButton } from "antd"; +import { SimpleTooltip } from "@/components/ui/tooltip"; +import { Button as AntdButton } from "antd"; import { z } from "zod/v4"; import { fetchUserModels } from "@/components/organisms/create_key_button"; import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; @@ -280,7 +281,9 @@ const TagInfoView: React.FC = ({ tagId, onClose, accessToken, ) : ( tagDetails.models.map((modelId) => ( - {tagDetails.model_info?.[modelId] || modelId} + + {tagDetails.model_info?.[modelId] || modelId} + )) )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx index a862e59d793..5ea15df13f0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx @@ -17,7 +17,8 @@ import { teamMemberDeleteCall, Member, } from "@/components/networking"; -import { Button as AntdButton, Modal, Tooltip } from "antd"; +import { SimpleTooltip } from "@/components/ui/tooltip"; +import { Button as AntdButton, Modal } from "antd"; import { Field, FieldGroup, FieldLabel } from "@/components/shared/form/field"; import { Combobox, @@ -747,10 +748,10 @@ export default function UserInfoView({ {MEMBER_ROLE_OPTIONS.map((option) => ( - + {option.value} - {option.hint} - + ))} diff --git a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx index e58139c80a0..fff732fe270 100644 --- a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx +++ b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx @@ -403,28 +403,27 @@ const ModelInfoEditForm: React.FC = ({
); - const pricingField = (name: TouchedPricingField, label: string, placeholder: string, description?: string) => ( -
- {label} - {isEditing ? ( - - {({ value, onChange, ...control }) => ( - ) => { - markTouched(name); - onChange(event); - }} - /> - )} - - ) : ( + const pricingField = (name: TouchedPricingField, label: string, placeholder: string, description?: string) => + isEditing ? ( + + {({ value, onChange, ...control }) => ( + ) => { + markTouched(name); + onChange(event); + }} + /> + )} + + ) : ( +
+ {label} {displayCost(localModelData, name)} - )} -
- ); +
+ ); const tagsField = ( name: "model_access_group" | "guardrails" | "tags", diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx index 381818a53f5..e3a36f4fbf2 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx @@ -3,7 +3,8 @@ * Handles primary model selection and fallback chain configuration */ -import { Select, Tooltip } from "antd"; +import { SimpleTooltip } from "@/components/ui/tooltip"; +import { Select } from "antd"; import { AlertCircle, ArrowDown, X } from "lucide-react"; import React from "react"; @@ -142,12 +143,9 @@ export function FallbackGroupConfig({ }} maxTagCount="responsive" maxTagPlaceholder={(omittedValues) => ( - value).join(", ")} - > + value).join(", ")}> +{omittedValues.length} more - + )} showSearch filterOption={(input, option) => (option?.label ?? "").toLowerCase().includes(input.toLowerCase())} diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx index 2993f549c65..b3c1be62e88 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx @@ -3,7 +3,8 @@ import { BarChart } from "@/components/shared/charts"; import { DataTable } from "@/components/shared/DataTable"; import { IdCell, MoneyCell } from "@/components/shared/table_cells"; import { ChevronDownIcon, ChevronUpIcon } from "@heroicons/react/outline"; -import { Segmented, Tooltip } from "antd"; +import { SimpleTooltip } from "@/components/ui/tooltip"; +import { Segmented } from "antd"; import React, { useState } from "react"; import { formatNumberWithCommas } from "../../../../utils/dataUtils"; import { transformKeyInfo } from "../../../key_team_helpers/transform_key_info"; @@ -113,9 +114,9 @@ const TopKeyView: React.FC = ({ topKeys, teams, showTags = fals
{displayTags.map((tag, index) => ( -
Tag Name: {tag.tag} @@ -128,7 +129,7 @@ const TopKeyView: React.FC = ({ topKeys, teams, showTags = fals } > {tag.tag.slice(0, 7)}... - + ))} {hasMoreTags && (
Optional: when a user message contains one of these phrases, the request is bumped one tier higher than it would diff --git a/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx b/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx index 72d6e9f4fef..49fea3d30e7 100644 --- a/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx +++ b/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx @@ -1,5 +1,6 @@ import { DeleteOutlined, InfoCircleOutlined, PlusOutlined } from "@ant-design/icons"; -import { Button, Card, Empty, Select as AntdSelect, Tooltip, Typography } from "antd"; +import { SimpleTooltip } from "@/components/ui/tooltip"; +import { Button, Card, Empty, Select as AntdSelect, Typography } from "antd"; import React from "react"; import { emptyKeywordTierRuleIndexes } from "./complexity_router_keywords"; @@ -70,9 +71,9 @@ const KeywordTierRules: React.FC = ({ rules, onChange, ti Keyword Tier Overrides - + - +
) ) : ( - + - + )}
diff --git a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts new file mode 100644 index 00000000000..67729d63fc3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts @@ -0,0 +1,556 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { buildKeyCreatePayload, type KeyCreateInput, type KeyPayloadResult } from "./createKeyPayload"; + +const baseInput: KeyCreateInput = { + formValues: {}, + existingKeys: null, + keyOwner: "you", + userID: "test-user", + selectedAgentId: null, + loggingSettings: [], + disabledCallbacks: [], + autoRotationEnabled: false, + rotationInterval: "30d", + modelAliases: {}, + routerSettings: null, + budgetLimits: [], + tagRateLimits: [], + budgetFallbacks: {}, +}; + +const build = (formValues: Record, overrides: Partial = {}): KeyPayloadResult => + buildKeyCreatePayload({ ...baseInput, ...overrides, formValues }); + +const payloadOf = (result: KeyPayloadResult): Record => { + expect(result.kind).toBe("ok"); + if (result.kind !== "ok") throw new Error("unreachable"); + return result.payload; +}; + +const wireKeys = (payload: Record): string[] => + Object.keys(JSON.parse(JSON.stringify(payload)) as Record); + +const DROPPED_AT_SERIALISATION = [ + "access_group_ids", + "allowed_passthrough_routes", + "allowed_vector_store_ids", + "budget_duration", + "enable_prompt_caching", + "guardrails", + "max_budget", + "organization_id", + "policies", + "prompts", + "rpm_limit", + "tags", + "throttle_on_budget_exceeded", + "tpm_limit", +]; + +const CLOSED_SECTIONS_VALUES = { + organization_id: undefined, + team_id: null, + key_alias: "my-key", + models: [], + key_type: "llm_api", +}; + +const OPTIONAL_SETTINGS_VALUES = { + ...CLOSED_SECTIONS_VALUES, + max_budget: undefined, + budget_duration: undefined, + tpm_limit: undefined, + tpm_limit_type: "key", + rpm_limit: undefined, + rpm_limit_type: "key", + throttle_on_budget_exceeded: undefined, + enable_prompt_caching: undefined, + guardrails: undefined, + disable_global_guardrails: undefined, + policies: undefined, + prompts: undefined, + access_group_ids: undefined, + allowed_passthrough_routes: undefined, + allowed_vector_store_ids: undefined, + tags: undefined, +}; + +const aliasOnly = (overrides: Record = {}): Record => ({ + key_alias: "my-key", + user_id: "test-user", + duration: null, + metadata: "{}", + ...overrides, +}); + +afterEach(() => { + vi.restoreAllMocks(); +}); + +describe("always-present keys", () => { + it("emits the eight keys the closed form sends, and nothing else", () => { + const closedFormValues = { ...CLOSED_SECTIONS_VALUES }; + const closedFormPayload = { + ...closedFormValues, + user_id: "test-user", + duration: null, + metadata: "{}", + }; + expect(payloadOf(build(closedFormValues))).toStrictEqual(closedFormPayload); + }); + + it("injects user_id, duration and metadata even when the form reported none of them", () => { + expect(payloadOf(build({ key_alias: "my-key" }))).toStrictEqual(aliasOnly()); + }); + + it("keeps a mounted-but-untouched field as an undefined-valued key rather than dropping it", () => { + expect(payloadOf(build({ key_alias: "my-key", guardrails: undefined, tags: undefined }))).toStrictEqual( + aliasOnly({ guardrails: undefined, tags: undefined }), + ); + }); +}); + +describe("duration", () => { + it("forwards a typed duration by value", () => { + expect(payloadOf(build({ key_alias: "my-key", duration: "45d" }))).toStrictEqual(aliasOnly({ duration: "45d" })); + }); + + it.each([ + ["an empty string", ""], + ["a whitespace-only string", " "], + ["undefined", undefined], + ])("coalesces %s to null", (_label, duration) => { + expect(payloadOf(build({ key_alias: "my-key", duration }))).toStrictEqual(aliasOnly({ duration: null })); + }); + + it("does not coerce a non-string duration", () => { + expect(() => build({ key_alias: "my-key", duration: 30 })).toThrow(TypeError); + }); +}); + +describe("key ownership", () => { + it("overwrites user_id with the signed-in user when the key is owned by you", () => { + expect(payloadOf(build({ key_alias: "my-key", user_id: "someone-else" }))).toStrictEqual( + aliasOnly({ user_id: "test-user" }), + ); + }); + + it("leaves the form's user_id alone for another_user", () => { + const expected = { key_alias: "my-key", user_id: "someone-else", duration: null, metadata: "{}" }; + expect( + payloadOf(build({ key_alias: "my-key", user_id: "someone-else" }, { keyOwner: "another_user" })), + ).toStrictEqual(expected); + }); + + it("adds the selected agent id for an agent-owned key", () => { + const expected = { key_alias: "my-key", agent_id: "agent-1", duration: null, metadata: "{}" }; + expect(payloadOf(build({ key_alias: "my-key" }, { keyOwner: "agent", selectedAgentId: "agent-1" }))).toStrictEqual( + expected, + ); + }); + + it("reports agent_not_selected instead of building a payload when no agent is selected", () => { + expect(build({ key_alias: "my-key" }, { keyOwner: "agent", selectedAgentId: null })).toStrictEqual({ + kind: "agent_not_selected", + }); + }); + + it("stamps the alias into metadata as the service account id and sends no user_id", () => { + expect(payloadOf(build({ key_alias: "svc-key" }, { keyOwner: "service_account" }))).toStrictEqual({ + key_alias: "svc-key", + duration: null, + metadata: '{"service_account_id":"svc-key"}', + }); + }); +}); + +describe("metadata", () => { + it("re-serialises the parsed form value", () => { + expect(payloadOf(build({ key_alias: "my-key", metadata: '{"team":"core"}' }))).toStrictEqual( + aliasOnly({ metadata: '{"team":"core"}' }), + ); + }); + + it("falls back to an empty object and logs when the JSON is malformed", () => { + const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}); + expect(payloadOf(build({ key_alias: "my-key", metadata: "{not json" }))).toStrictEqual(aliasOnly()); + expect(consoleError).toHaveBeenCalledWith("Error parsing metadata:", expect.any(SyntaxError)); + }); + + it("merges logging configs, dropping rows with no callback selected", () => { + expect( + payloadOf( + build( + { key_alias: "my-key", metadata: '{"team":"core"}' }, + { loggingSettings: [{ callback_name: "langfuse" }, { callback_name: "" }] }, + ), + ), + ).toStrictEqual(aliasOnly({ metadata: '{"team":"core","logging":[{"callback_name":"langfuse"}]}' })); + }); + + it("maps disabled callbacks from display names to internal names", () => { + expect(payloadOf(build({ key_alias: "my-key" }, { disabledCallbacks: ["Langfuse"] }))).toStrictEqual( + aliasOnly({ metadata: '{"litellm_disabled_callbacks":["langfuse"]}' }), + ); + }); + + it("keeps an array metadata as an array when stamping the service account id", () => { + expect( + payloadOf(build({ key_alias: "svc-key", metadata: '["a"]' }, { keyOwner: "service_account" })), + ).toStrictEqual({ key_alias: "svc-key", duration: null, metadata: '["a"]' }); + }); + + it("rejects a service account whose metadata parses to a primitive", () => { + expect(() => build({ key_alias: "svc-key", metadata: "5" }, { keyOwner: "service_account" })).toThrow(TypeError); + }); + + it("keeps every metadata contributor in one object", () => { + expect( + payloadOf( + build( + { key_alias: "svc-key", metadata: '{"team":"core"}' }, + { + keyOwner: "service_account", + loggingSettings: [{ callback_name: "otel" }], + disabledCallbacks: ["Datadog"], + }, + ), + ), + ).toStrictEqual({ + key_alias: "svc-key", + duration: null, + metadata: + '{"team":"core","service_account_id":"svc-key","logging":[{"callback_name":"otel"}],"litellm_disabled_callbacks":["datadog"]}', + }); + }); +}); + +describe("object_permission", () => { + it("is absent when nothing contributes to it", () => { + expect(payloadOf(build({ key_alias: "my-key", allowed_vector_store_ids: [] }))).toStrictEqual( + aliasOnly({ allowed_vector_store_ids: [] }), + ); + }); + + it("moves selected vector stores off the top level", () => { + expect(payloadOf(build({ key_alias: "my-key", allowed_vector_store_ids: ["vs-1"] }))).toStrictEqual( + aliasOnly({ object_permission: { vector_stores: ["vs-1"] } }), + ); + }); + + it("splits an MCP selection into servers, access groups and toolsets", () => { + expect( + payloadOf( + build({ + key_alias: "my-key", + allowed_mcp_servers_and_groups: { servers: ["s-1"], accessGroups: ["g-1"], toolsets: ["t-1"] }, + }), + ), + ).toStrictEqual( + aliasOnly({ + object_permission: { mcp_servers: ["s-1"], mcp_access_groups: ["g-1"], mcp_toolsets: ["t-1"] }, + }), + ); + }); + + it("omits the empty halves of an MCP selection", () => { + expect( + payloadOf( + build({ + key_alias: "my-key", + allowed_mcp_servers_and_groups: { servers: ["s-1"], accessGroups: [], toolsets: [] }, + }), + ), + ).toStrictEqual(aliasOnly({ object_permission: { mcp_servers: ["s-1"] } })); + }); + + it("leaves an all-empty MCP selection on the top level", () => { + expect( + payloadOf( + build({ key_alias: "my-key", allowed_mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] } }), + ), + ).toStrictEqual(aliasOnly({ allowed_mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] } })); + }); + + it("nests configured MCP tool permissions", () => { + expect(payloadOf(build({ key_alias: "my-key", mcp_tool_permissions: { "s-1": ["read"] } }))).toStrictEqual( + aliasOnly({ object_permission: { mcp_tool_permissions: { "s-1": ["read"] } } }), + ); + }); + + it("always strips mcp_tool_permissions from the top level, even when empty", () => { + expect(payloadOf(build({ key_alias: "my-key", mcp_tool_permissions: {} }))).toStrictEqual(aliasOnly()); + }); + + it("lets a standalone access group list win over the one from the MCP selection", () => { + expect( + payloadOf( + build({ + key_alias: "my-key", + allowed_mcp_servers_and_groups: { servers: ["s-1"], accessGroups: ["from-selector"] }, + allowed_mcp_access_groups: ["standalone"], + }), + ), + ).toStrictEqual(aliasOnly({ object_permission: { mcp_servers: ["s-1"], mcp_access_groups: ["standalone"] } })); + }); + + it("splits an agent selection into agents and agent access groups", () => { + expect( + payloadOf(build({ key_alias: "my-key", allowed_agents_and_groups: { agents: ["a-1"], accessGroups: ["ag-1"] } })), + ).toStrictEqual(aliasOnly({ object_permission: { agents: ["a-1"], agent_access_groups: ["ag-1"] } })); + }); + + it("merges every source into a single object_permission", () => { + const everySource = { + key_alias: "my-key", + allowed_vector_store_ids: ["vs-1"], + allowed_mcp_servers_and_groups: { servers: ["s-1"], accessGroups: ["g-1"], toolsets: ["t-1"] }, + mcp_tool_permissions: { "s-1": ["read"] }, + allowed_agents_and_groups: { agents: ["a-1"], accessGroups: ["ag-1"] }, + }; + expect(payloadOf(build(everySource))).toStrictEqual( + aliasOnly({ + object_permission: { + vector_stores: ["vs-1"], + mcp_servers: ["s-1"], + mcp_access_groups: ["g-1"], + mcp_toolsets: ["t-1"], + mcp_tool_permissions: { "s-1": ["read"] }, + agents: ["a-1"], + agent_access_groups: ["ag-1"], + }, + }), + ); + }); +}); + +describe("premium and rotation flags", () => { + it("drops disable_global_guardrails when it is off", () => { + expect(payloadOf(build({ key_alias: "my-key", disable_global_guardrails: false }))).toStrictEqual(aliasOnly()); + }); + + it("keeps disable_global_guardrails when it is on", () => { + expect(payloadOf(build({ key_alias: "my-key", disable_global_guardrails: true }))).toStrictEqual( + aliasOnly({ disable_global_guardrails: true }), + ); + }); + + it("adds the rotation fields only when auto rotation is enabled", () => { + expect( + payloadOf(build({ key_alias: "my-key" }, { autoRotationEnabled: true, rotationInterval: "7d" })), + ).toStrictEqual(aliasOnly({ auto_rotate: true, rotation_interval: "7d" })); + }); + + it("sends no rotation fields when auto rotation is off", () => { + expect(payloadOf(build({ key_alias: "my-key" }, { rotationInterval: "7d" }))).toStrictEqual(aliasOnly()); + }); +}); + +describe("keys sourced from component state", () => { + it("serialises model aliases", () => { + expect(payloadOf(build({ key_alias: "my-key" }, { modelAliases: { fast: "gpt-4o-mini" } }))).toStrictEqual( + aliasOnly({ aliases: '{"fast":"gpt-4o-mini"}' }), + ); + }); + + it("sends router settings that hold at least one value", () => { + expect( + payloadOf(build({ key_alias: "my-key" }, { routerSettings: { router_settings: { num_retries: 3 } } })), + ).toStrictEqual(aliasOnly({ router_settings: { num_retries: 3 } })); + }); + + it("skips router settings whose every field is blank", () => { + expect( + payloadOf( + build( + { key_alias: "my-key" }, + { routerSettings: { router_settings: { num_retries: null, timeout: undefined, routing_strategy: "" } } }, + ), + ), + ).toStrictEqual(aliasOnly()); + }); + + it("keeps only budget windows that carry both a duration and a limit", () => { + expect( + payloadOf( + build( + { key_alias: "my-key" }, + { + budgetLimits: [ + { budget_duration: "1h", max_budget: 5 }, + { budget_duration: "", max_budget: 3 }, + { budget_duration: "7d", max_budget: null }, + ], + }, + ), + ), + ).toStrictEqual(aliasOnly({ budget_limits: [{ budget_duration: "1h", max_budget: 5 }] })); + }); + + it("keeps a zero budget window rather than treating it as unset", () => { + expect( + payloadOf(build({ key_alias: "my-key" }, { budgetLimits: [{ budget_duration: "1h", max_budget: 0 }] })), + ).toStrictEqual(aliasOnly({ budget_limits: [{ budget_duration: "1h", max_budget: 0 }] })); + }); + + it("omits budget_limits when no window is complete", () => { + expect( + payloadOf(build({ key_alias: "my-key" }, { budgetLimits: [{ budget_duration: "7d", max_budget: null }] })), + ).toStrictEqual(aliasOnly()); + }); + + it("reduces tag rows into a tag_rpm_limit map", () => { + expect( + payloadOf( + build( + { key_alias: "my-key" }, + { + tagRateLimits: [ + { id: "r-1", tag: "prod", rpm_limit: 10 }, + { id: "r-2", tag: " ", rpm_limit: 5 }, + { id: "r-3", tag: "dev", rpm_limit: null }, + ], + }, + ), + ), + ).toStrictEqual(aliasOnly({ tag_rpm_limit: { prod: 10 } })); + }); + + it("omits tag_rpm_limit when no row is complete", () => { + expect( + payloadOf(build({ key_alias: "my-key" }, { tagRateLimits: [{ id: "r-1", tag: "", rpm_limit: 10 }] })), + ).toStrictEqual(aliasOnly()); + }); + + it("sends configured budget fallbacks", () => { + expect(payloadOf(build({ key_alias: "my-key" }, { budgetFallbacks: { "gpt-4": ["gpt-4o"] } }))).toStrictEqual( + aliasOnly({ budget_fallbacks: { "gpt-4": ["gpt-4o"] } }), + ); + }); +}); + +describe("budget duration", () => { + it("turns the never-resets sentinel into null", () => { + expect(payloadOf(build({ key_alias: "my-key", budget_duration: "none" }))).toStrictEqual( + aliasOnly({ budget_duration: null }), + ); + }); + + it("forwards a real budget duration untouched", () => { + expect(payloadOf(build({ key_alias: "my-key", budget_duration: "30d" }))).toStrictEqual( + aliasOnly({ budget_duration: "30d" }), + ); + }); +}); + +describe("purity", () => { + it("leaves the submitted form values untouched", () => { + const values = { + key_alias: "my-key", + mcp_tool_permissions: { "s-1": ["read"] }, + allowed_vector_store_ids: ["vs-1"], + disable_global_guardrails: false, + duration: "", + }; + const before = structuredClone(values); + build(values); + expect(values).toStrictEqual(before); + }); +}); + +describe("serialised wire shape", () => { + it("keeps an untouched closed form at eight object keys and seven wire keys", () => { + const payload = payloadOf(build(CLOSED_SECTIONS_VALUES)); + expect(Object.keys(payload)).toHaveLength(8); + expect(wireKeys(payload)).toStrictEqual([ + "team_id", + "key_alias", + "models", + "key_type", + "user_id", + "duration", + "metadata", + ]); + expect(payload.duration).toBeNull(); + }); + + it("drops the undefined picker and keeps the null one, which is what the two Form.Items differ on", () => { + const payload = payloadOf(build(CLOSED_SECTIONS_VALUES)); + expect(payload.organization_id).toBeUndefined(); + expect(payload.team_id).toBeNull(); + expect(wireKeys(payload)).not.toContain("organization_id"); + expect(wireKeys(payload)).toContain("team_id"); + }); + + it("forwards a selected team by value", () => { + expect(payloadOf(build({ ...CLOSED_SECTIONS_VALUES, team_id: "team-1" })).team_id).toBe("team-1"); + }); + + it("adds fifteen keys to the object and only the two limit types to the wire when Optional Settings opens", () => { + const payload = payloadOf(build(OPTIONAL_SETTINGS_VALUES)); + expect(Object.keys(payload)).toHaveLength(23); + expect(wireKeys(payload)).toStrictEqual([ + "team_id", + "key_alias", + "models", + "key_type", + "tpm_limit_type", + "rpm_limit_type", + "user_id", + "duration", + "metadata", + ]); + }); + + it("never turns an undefined-valued key into null or an empty string", () => { + const payload = payloadOf(build(OPTIONAL_SETTINGS_VALUES)); + DROPPED_AT_SERIALISATION.forEach((key) => { + expect(payload[key]).toBeUndefined(); + }); + expect(wireKeys(payload)).toEqual(expect.not.arrayContaining(DROPPED_AT_SERIALISATION)); + }); +}); + +describe("duplicate alias", () => { + it("reports the clash instead of building a payload", () => { + expect( + build({ key_alias: "taken", team_id: "team-1" }, { existingKeys: [{ team_id: "team-1", key_alias: "taken" }] }), + ).toStrictEqual({ kind: "duplicate_alias", alias: "taken", teamId: "team-1" }); + }); + + it("scopes the clash to the same team", () => { + expect( + payloadOf( + build({ key_alias: "taken", team_id: "team-2" }, { existingKeys: [{ team_id: "team-1", key_alias: "taken" }] }), + ).key_alias, + ).toBe("taken"); + }); + + it("treats a keyless form and a teamless key as the same bucket", () => { + expect(build({}, { existingKeys: [{ team_id: null, key_alias: "" }] })).toStrictEqual({ + kind: "duplicate_alias", + alias: "", + teamId: null, + }); + }); + + it("checks the alias before the agent selection", () => { + expect( + build( + { key_alias: "taken" }, + { existingKeys: [{ team_id: null, key_alias: "taken" }], keyOwner: "agent", selectedAgentId: null }, + ).kind, + ).toBe("duplicate_alias"); + }); +}); + +describe("endpoint", () => { + it.each([ + ["you", "standard"], + ["another_user", "standard"], + ["service_account", "service_account"], + ])("routes a %s key to the %s endpoint", (keyOwner, endpoint) => { + const result = build({ key_alias: "my-key" }, { keyOwner }); + expect(result.kind === "ok" && result.endpoint).toBe(endpoint); + }); +}); diff --git a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts new file mode 100644 index 00000000000..2b14e8372f7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts @@ -0,0 +1,212 @@ +import { mapDisplayToInternalNames } from "../callback_info_helpers"; +import { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown"; +import type { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion"; +import type { BudgetWindowEntry } from "../key_team_helpers/BudgetWindowsEditor"; +import { tagRowsToLimits, type TagRateLimitEntry } from "../key_team_helpers/TagRateLimitEditor"; + +export interface KeyLoggingSetting { + callback_name?: string; +} + +export interface ExistingKey { + readonly team_id?: string | null; + readonly key_alias?: string | null; +} + +export interface KeyCreateInput { + readonly formValues: Record; + readonly existingKeys: readonly ExistingKey[] | null; + readonly keyOwner: string; + readonly userID: string | null; + readonly selectedAgentId: string | null; + readonly loggingSettings: KeyLoggingSetting[]; + readonly disabledCallbacks: string[]; + readonly autoRotationEnabled: boolean; + readonly rotationInterval: string; + readonly modelAliases: Record; + readonly routerSettings: RouterSettingsAccordionValue | null; + readonly budgetLimits: BudgetWindowEntry[]; + readonly tagRateLimits: TagRateLimitEntry[]; + readonly budgetFallbacks: Record; +} + +export type KeyPayloadResult = + | { + readonly kind: "ok"; + readonly payload: Record; + readonly endpoint: "standard" | "service_account"; + } + | { readonly kind: "duplicate_alias"; readonly alias: string; readonly teamId: string | null } + | { readonly kind: "agent_not_selected" }; + +interface McpSelection { + readonly servers?: unknown[]; + readonly accessGroups?: unknown[]; + readonly toolsets?: unknown[]; +} + +interface AgentSelection { + readonly agents?: unknown[]; + readonly accessGroups?: unknown[]; +} + +const nonEmptyList = (raw: unknown): unknown[] | undefined => { + const list = raw as unknown[] | undefined; + return list && list.length > 0 ? list : undefined; +}; + +const readMcpSelection = (raw: unknown): McpSelection | undefined => { + const selection = raw as McpSelection | undefined; + if (!selection) return undefined; + const servers = nonEmptyList(selection.servers); + const accessGroups = nonEmptyList(selection.accessGroups); + const toolsets = nonEmptyList(selection.toolsets); + if (!servers && !accessGroups && !toolsets) return undefined; + return { servers, accessGroups, toolsets }; +}; + +const readAgentSelection = (raw: unknown): AgentSelection | undefined => { + const selection = raw as AgentSelection | undefined; + if (!selection) return undefined; + const agents = nonEmptyList(selection.agents); + const accessGroups = nonEmptyList(selection.accessGroups); + if (!agents && !accessGroups) return undefined; + return { agents, accessGroups }; +}; + +const readToolPermissions = (raw: unknown): unknown | undefined => { + const permissions = raw || {}; + return Object.keys(permissions as object).length > 0 ? permissions : undefined; +}; + +const parseMetadata = (raw: unknown): unknown => { + try { + return JSON.parse((raw as string) || "{}"); + } catch (error) { + console.error("Error parsing metadata:", error); + return {}; + } +}; + +const buildMetadataJson = (values: Record, input: KeyCreateInput): string => { + const parsed = parseMetadata(values.metadata); + if (input.keyOwner === "service_account") { + (parsed as Record).service_account_id = values.key_alias; + } + const logged = + input.loggingSettings.length > 0 + ? { ...(parsed as object), logging: input.loggingSettings.filter((config) => config.callback_name) } + : parsed; + const disabled = + input.disabledCallbacks.length > 0 + ? { ...(logged as object), litellm_disabled_callbacks: mapDisplayToInternalNames(input.disabledCallbacks) } + : logged; + return JSON.stringify(disabled); +}; + +interface PermissionSources { + readonly vectorStores: unknown[] | undefined; + readonly mcp: McpSelection | undefined; + readonly toolPermissions: unknown | undefined; + readonly extraMcpAccessGroups: unknown[] | undefined; + readonly agents: AgentSelection | undefined; +} + +const readPermissionSources = (values: Record): PermissionSources => ({ + vectorStores: nonEmptyList(values.allowed_vector_store_ids), + mcp: readMcpSelection(values.allowed_mcp_servers_and_groups), + toolPermissions: readToolPermissions(values.mcp_tool_permissions), + extraMcpAccessGroups: nonEmptyList(values.allowed_mcp_access_groups), + agents: readAgentSelection(values.allowed_agents_and_groups), +}); + +const buildObjectPermission = ({ + vectorStores, + mcp, + toolPermissions, + extraMcpAccessGroups, + agents, +}: PermissionSources): Record | undefined => { + const permission: Record = { + ...(vectorStores && { vector_stores: vectorStores }), + ...(mcp?.servers && { mcp_servers: mcp.servers }), + ...(mcp?.accessGroups && { mcp_access_groups: mcp.accessGroups }), + ...(mcp?.toolsets && { mcp_toolsets: mcp.toolsets }), + ...(toolPermissions !== undefined && { mcp_tool_permissions: toolPermissions }), + ...(extraMcpAccessGroups && { mcp_access_groups: extraMcpAccessGroups }), + ...(agents?.agents && { agents: agents.agents }), + ...(agents?.accessGroups && { agent_access_groups: agents.accessGroups }), + }; + return Object.keys(permission).length > 0 ? permission : undefined; +}; + +const consumedSourceKeys = ( + values: Record, + { vectorStores, mcp, extraMcpAccessGroups, agents }: PermissionSources, +): ReadonlySet => + new Set([ + "mcp_tool_permissions", + ...(values.disable_global_guardrails ? [] : ["disable_global_guardrails"]), + ...(vectorStores ? ["allowed_vector_store_ids"] : []), + ...(mcp ? ["allowed_mcp_servers_and_groups"] : []), + ...(extraMcpAccessGroups ? ["allowed_mcp_access_groups"] : []), + ...(agents ? ["allowed_agents_and_groups"] : []), + ]); + +const withoutKeys = (values: Record, dropped: ReadonlySet): Record => + Object.fromEntries(Object.entries(values).filter(([key]) => !dropped.has(key))); + +const duplicateAlias = (input: KeyCreateInput): { alias: string; teamId: string | null } | undefined => { + const alias = (input.formValues?.key_alias as string | undefined) ?? ""; + const teamId = (input.formValues?.team_id as string | undefined) ?? null; + const taken = (input.existingKeys ?? []).filter((key) => key.team_id === teamId).map((key) => key.key_alias); + return taken.includes(alias) ? { alias, teamId } : undefined; +}; + +export const buildKeyCreatePayload = (input: KeyCreateInput): KeyPayloadResult => { + const duplicate = duplicateAlias(input); + if (duplicate) { + return { kind: "duplicate_alias", ...duplicate }; + } + if (input.keyOwner === "agent" && !input.selectedAgentId) { + return { kind: "agent_not_selected" }; + } + + const values = input.formValues; + + const sources = readPermissionSources(values); + const objectPermission = buildObjectPermission(sources); + const dropped = consumedSourceKeys(values, sources); + + const duration = values.duration; + const validWindows = input.budgetLimits.filter( + (window) => window.budget_duration && window.max_budget !== null && window.max_budget !== undefined, + ); + const { tag_rpm_limit } = tagRowsToLimits(input.tagRateLimits); + const routerSettings = input.routerSettings?.router_settings; + const configuredRouterSettings = + routerSettings && + Object.values(routerSettings).some((value) => value !== null && value !== undefined && value !== "") + ? routerSettings + : undefined; + + return { + kind: "ok", + endpoint: input.keyOwner === "service_account" ? "service_account" : "standard", + payload: { + ...withoutKeys(values, dropped), + ...(input.keyOwner === "you" && { user_id: input.userID }), + ...(input.keyOwner === "agent" && { agent_id: input.selectedAgentId }), + ...(input.autoRotationEnabled && { auto_rotate: true, rotation_interval: input.rotationInterval }), + duration: !duration || (duration as string).trim() === "" ? null : duration, + metadata: buildMetadataJson(values, input), + ...(objectPermission && { object_permission: objectPermission }), + ...(Object.keys(input.modelAliases).length > 0 && { aliases: JSON.stringify(input.modelAliases) }), + ...(configuredRouterSettings && { router_settings: configuredRouterSettings }), + ...(validWindows.length > 0 && { budget_limits: validWindows }), + ...(Object.keys(tag_rpm_limit).length > 0 && { tag_rpm_limit }), + ...(Object.keys(input.budgetFallbacks).length > 0 && { budget_fallbacks: input.budgetFallbacks }), + ...(values.budget_duration === NEVER_RESETS_BUDGET_DURATION && { budget_duration: null }), + }, + }; +}; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index feea216e4ff..9788d365fd0 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -1,70 +1,114 @@ import userEvent from "@testing-library/user-event"; -import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; +import type { Team } from "../key_team_helpers/key_list"; +import { keyCreateCall, keyCreateServiceAccountCall, modelAvailableCall, userFilterUICall } from "../networking"; import CreateKey from "./create_key_button"; -const { mockKeyCreateCall } = vi.hoisted(() => ({ - mockKeyCreateCall: vi.fn().mockResolvedValue({ key: "sk-created", soft_budget: null }), +const state = vi.hoisted(() => ({ + authorized: { + accessToken: "test-token", + userId: "test-user-id", + userRole: "Admin", + premiumUser: false, + }, + can: {} as Record, + uiSettings: {} as Record, + tags: {} as Record, + teams: [] as { team_id: string; team_alias: string; models: string[] }[], + organizations: [] as { organization_id: string; organization_alias: string }[], + accessGroups: [] as { access_group_id: string; access_group_name: string }[], + projects: [] as { project_id: string; project_alias: string; team_id?: string; models?: string[] }[], })); -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => ({ accessToken: "test-token", userId: "test-user-id", userRole: "Admin", premiumUser: true }), +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => state.authorized })); +vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ + default: (capability: string) => state.can[capability] ?? true, })); - -vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ default: () => false })); - -vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ keyKeys: { lists: () => ["keys"] } })); - vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ - useOrganizations: () => ({ data: [], isLoading: false }), + useOrganizations: () => ({ data: state.organizations, isLoading: false }), })); - vi.mock("@/app/(dashboard)/hooks/projects/useProjects", () => ({ - useProjects: () => ({ data: [], isLoading: false }), + useProjects: () => ({ data: state.projects, isLoading: false }), })); - vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ - useUISettings: () => ({ data: { values: {} } }), + useUISettings: () => ({ data: { values: state.uiSettings } }), +})); +vi.mock("@/app/(dashboard)/hooks/tags/useTags", () => ({ + useTags: () => ({ data: state.tags, isLoading: false }), })); - -vi.mock("@/app/(dashboard)/hooks/tags/useTags", () => ({ useTags: () => ({ data: {} }) })); - vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ - useInfiniteTeams: () => ({ data: { pages: [{ teams: [] }] }, fetchNextPage: vi.fn(), hasNextPage: false }), + useInfiniteTeams: () => ({ + data: { pages: [{ teams: state.teams, total: state.teams.length, page: 1, page_size: 50, total_pages: 1 }] }, + fetchNextPage: vi.fn(), + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + }), })); - vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ - useAccessGroups: () => ({ data: [], isLoading: false, isError: false }), + useAccessGroups: () => ({ data: state.accessGroups, isLoading: false, isError: false }), +})); +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: () => ({ data: [], isLoading: false }), +})); +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups", () => ({ + useMCPAccessGroups: () => ({ data: [], isLoading: false }), +})); +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ + useMCPToolsets: () => ({ data: [], isLoading: false }), })); -vi.mock("../networking", () => ({ - keyCreateCall: mockKeyCreateCall, - keyCreateServiceAccountCall: vi.fn().mockResolvedValue({ key: "sk-sa", soft_budget: null }), - modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }] }), - getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), - getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), - getPromptsList: vi.fn().mockResolvedValue({ prompts: [] }), - getPossibleUserRoles: vi.fn().mockResolvedValue({}), - userFilterUICall: vi.fn().mockResolvedValue([]), - fetchMCPAccessGroups: vi.fn().mockResolvedValue([]), - getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), - getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: [] }), - proxyBaseUrl: "http://localhost:4000", -})); +vi.mock("../networking", async (importOriginal) => { + const actual = await importOriginal(); + const emptyMcpTools = { tools: [], error: null, message: null, stack_trace: null }; + return { + ...actual, + keyCreateCall: vi.fn().mockResolvedValue({ key: "sk-created", soft_budget: null }), + keyCreateServiceAccountCall: vi.fn().mockResolvedValue({ key: "sk-service-account", soft_budget: null }), + modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }] }), + getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), + getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), + getPromptsList: vi.fn().mockResolvedValue({ prompts: [] }), + getPossibleUserRoles: vi.fn().mockResolvedValue({}), + userFilterUICall: vi.fn().mockResolvedValue([]), + getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), + getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: [] }), + vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), + listMCPTools: vi.fn().mockResolvedValue(emptyMcpTools), + getRouterSettingsCall: vi.fn().mockResolvedValue({ router_settings: {} }), + }; +}); -vi.mock("../agent_management/AgentSelector", () => ({ default: () => null })); -vi.mock("../common_components/check_openapi_schema", () => ({ default: () => null })); -vi.mock("../common_components/PremiumLoggingSettings", () => ({ default: () => null })); -vi.mock("../common_components/RouterSettingsAccordion", () => ({ default: () => null })); -vi.mock("../mcp_server_management/MCPServerSelector", () => ({ default: () => null })); -vi.mock("../mcp_server_management/MCPToolPermissions", () => ({ default: () => null })); -vi.mock("../vector_store_management/VectorStoreSelector", () => ({ default: () => null })); -vi.mock("../CreateUserButton", () => ({ CreateUserButton: () => null })); +const OPENAPI_SCHEMA = { + components: { + schemas: { + GenerateKeyRequest: { + properties: { + key: { type: "string", title: "Key" }, + soft_budget: { type: "number", title: "Soft Budget" }, + blocked: { type: "boolean", title: "Blocked" }, + max_budget: { type: "number", title: "Max Budget" }, + }, + }, + }, + }, +}; -const MINIMAL_CREATE_PAYLOAD = { +const SECTIONS = { + mcp: /MCP Settings/i, + agent: /Agent Settings/i, + logging: /Logging Settings/i, + router: /Router Settings/i, + aliases: /Model Aliases/i, + lifecycle: /Key Lifecycle/i, + advanced: /Advanced Settings/i, +} as const; + +const ALL_CLOSED_PAYLOAD = { organization_id: undefined, team_id: null, - key_alias: "probe-key", + key_alias: "contract-key", models: [], key_type: "llm_api", user_id: "test-user-id", @@ -72,100 +116,763 @@ const MINIMAL_CREATE_PAYLOAD = { metadata: "{}", }; -describe("CreateKey submit payload contract", () => { +const OPTIONAL_OPEN_PAYLOAD = { + ...ALL_CLOSED_PAYLOAD, + max_budget: undefined, + budget_duration: undefined, + tpm_limit: undefined, + tpm_limit_type: null, + rpm_limit: undefined, + rpm_limit_type: null, + throttle_on_budget_exceeded: undefined, + enable_prompt_caching: undefined, + guardrails: undefined, + policies: undefined, + prompts: undefined, + access_group_ids: undefined, + allowed_passthrough_routes: undefined, + allowed_vector_store_ids: undefined, + tags: undefined, +}; + +const ROUTER_SETTINGS_DEFAULT = { + routing_strategy: null, + allowed_fails: null, + cooldown_time: null, + num_retries: null, + timeout: null, + retry_after: null, + fallbacks: null, + context_window_fallbacks: null, + retry_policy: null, + model_group_alias: null, + enable_tag_filtering: false, + routing_strategy_args: null, +}; + +const SECTION_PAYLOAD_ADDITIONS: Record> = { + mcp: { allowed_mcp_servers_and_groups: { servers: [], accessGroups: [] } }, + agent: { allowed_agents_and_groups: undefined }, + logging: {}, + router: { router_settings: ROUTER_SETTINGS_DEFAULT }, + aliases: {}, + lifecycle: {}, + advanced: { key: undefined, soft_budget: undefined, blocked: undefined }, +}; + +const ALL_OPEN_PAYLOAD = { + ...OPTIONAL_OPEN_PAYLOAD, + ...SECTION_PAYLOAD_ADDITIONS.mcp, + ...SECTION_PAYLOAD_ADDITIONS.agent, + ...SECTION_PAYLOAD_ADDITIONS.router, + ...SECTION_PAYLOAD_ADDITIONS.advanced, +}; + +const renderCreateKey = (props: Partial> = {}) => + renderWithProviders(); + +const openModal = async (props: Partial> = {}) => { + const view = renderCreateKey(props); + await userEvent.click(screen.getByTestId("create-key-button")); + await screen.findByRole("button", { name: /^create key$/i }); + return view; +}; + +const antdSearchInput = (placeholder: HTMLElement): HTMLInputElement => { + const input = placeholder.parentElement?.querySelector("input"); + if (input == null) { + throw new Error("no antd search input next to the placeholder"); + } + return input; +}; + +const openAntdSelect = async (placeholder: HTMLElement) => { + await userEvent.click(antdSearchInput(placeholder)); +}; + +const openSection = async (name: RegExp) => { + await userEvent.click(await screen.findByRole("button", { name })); +}; + +const nameTheKey = async (alias = "contract-key") => { + await userEvent.type(await screen.findByLabelText(/Key Name/), alias); +}; + +const submit = async () => { + await userEvent.click(screen.getByRole("button", { name: /^create key$/i })); +}; + +const createdPayload = async () => { + await waitFor(() => { + expect(vi.mocked(keyCreateCall)).toHaveBeenCalled(); + }); + return vi.mocked(keyCreateCall).mock.calls[0][2] as Record; +}; + +describe("CreateKey", () => { beforeEach(() => { - mockKeyCreateCall.mockClear(); + testQueryClient.clear(); + state.authorized = { accessToken: "test-token", userId: "test-user-id", userRole: "Admin", premiumUser: false }; + state.can = {}; + state.uiSettings = {}; + state.tags = {}; + state.teams = []; + state.organizations = []; + state.projects = []; + state.accessGroups = []; + vi.mocked(keyCreateCall).mockClear().mockResolvedValue({ key: "sk-created", soft_budget: null }); + vi.mocked(keyCreateServiceAccountCall) + .mockClear() + .mockResolvedValue({ key: "sk-service-account", soft_budget: null }); + vi.mocked(userFilterUICall).mockClear().mockResolvedValue([]); + vi.mocked(modelAvailableCall) + .mockClear() + .mockResolvedValue({ data: [{ id: "gpt-4" }] }); + vi.stubGlobal( + "fetch", + vi.fn(async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes("openapi.json")) { + return { ok: true, status: 200, json: async () => OPENAPI_SCHEMA } as unknown as Response; + } + return { ok: true, status: 200, json: async () => ({}) } as unknown as Response; + }), + ); }); - const openModal = async () => { - renderWithProviders( {}} />); - await userEvent.click(screen.getAllByTestId("create-key-button")[0]); - await screen.findByRole("button", { name: /create key/i }); - }; - - it("sends exactly the bound form fields for a minimal create", async () => { - await openModal(); - - await userEvent.type(screen.getByLabelText(/Key Name/), "probe-key"); - await userEvent.click(screen.getByRole("button", { name: /create key/i })); - - await waitFor(() => { - expect(mockKeyCreateCall).toHaveBeenCalled(); - }); - expect(mockKeyCreateCall.mock.calls[0][2]).toStrictEqual(MINIMAL_CREATE_PAYLOAD); + afterEach(() => { + vi.unstubAllGlobals(); }); - it("keeps a collapsed Optional Settings section out of the payload entirely", async () => { - await openModal(); + describe("submit payload contract", () => { + it("sends only the always-mounted fields when every collapsible section is closed", async () => { + await openModal(); + await nameTheKey(); + await submit(); - await userEvent.type(screen.getByLabelText(/Key Name/), "probe-key"); - await userEvent.click(screen.getByRole("button", { name: /create key/i })); - - await waitFor(() => { - expect(mockKeyCreateCall).toHaveBeenCalled(); + expect(await createdPayload()).toStrictEqual(ALL_CLOSED_PAYLOAD); }); - const payload = mockKeyCreateCall.mock.calls[0][2]; - expect(payload).not.toHaveProperty("tpm_limit_type"); - expect(payload).not.toHaveProperty("rpm_limit_type"); - expect(payload).not.toHaveProperty("max_budget"); - }); - it("carries the shared rate-limit-type control into the payload once its section is open", async () => { - await openModal(); + it("registers the Optional Settings fields as undefined-valued keys once that section is open", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await submit(); - await userEvent.type(screen.getByLabelText(/Key Name/), "probe-key"); - await userEvent.click(screen.getByText("Optional Settings")); - - await userEvent.click(await screen.findByLabelText(/TPM Rate Limit Type/)); - await userEvent.click(await screen.findByText("Guaranteed throughput")); - - await userEvent.click(screen.getByRole("button", { name: /create key/i })); - - await waitFor(() => { - expect(mockKeyCreateCall).toHaveBeenCalled(); + expect(await createdPayload()).toStrictEqual(OPTIONAL_OPEN_PAYLOAD); }); - expect(mockKeyCreateCall.mock.calls[0][2]).toMatchObject({ - tpm_limit_type: "guaranteed_throughput", - rpm_limit_type: null, + + it("sends the full mounted field set when every section is open", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + for (const trigger of Object.values(SECTIONS)) { + await openSection(trigger); + } + await screen.findByLabelText("Soft Budget"); + await submit(); + + expect(await createdPayload()).toStrictEqual(ALL_OPEN_PAYLOAD); + }); + + it.each(Object.keys(SECTIONS) as (keyof typeof SECTIONS)[])( + "adds exactly the %s section's own keys when it is the only nested section open", + async (section) => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await openSection(SECTIONS[section]); + if (section === "advanced") { + await screen.findByLabelText("Soft Budget"); + } + await submit(); + + expect(await createdPayload()).toStrictEqual({ + ...OPTIONAL_OPEN_PAYLOAD, + ...SECTION_PAYLOAD_ADDITIONS[section], + }); + }, + ); + + it.each([ + [ + "every section closed", + false, + ["team_id", "key_alias", "models", "key_type", "user_id", "duration", "metadata"], + ], + [ + "Optional Settings open", + true, + [ + "team_id", + "key_alias", + "models", + "key_type", + "tpm_limit_type", + "rpm_limit_type", + "user_id", + "duration", + "metadata", + ], + ], + ])("serialises to exactly the wire keys with %s", async (_label, openOptional, wireKeys) => { + await openModal(); + await nameTheKey(); + if (openOptional) { + await openSection(/Optional Settings/i); + } + await submit(); + + const serialised = JSON.parse(JSON.stringify(await createdPayload())) as Record; + expect(Object.keys(serialised).sort()).toStrictEqual([...wireKeys].sort()); + }); + + it("omits a budget typed into a section the user closed again, rather than sending it as null", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.type(await screen.findByLabelText(/Max Budget \(USD\)/), "150.75"); + await openSection(/Optional Settings/i); + await submit(); + + const payload = await createdPayload(); + expect(payload).not.toHaveProperty("max_budget"); + expect(payload).toStrictEqual(ALL_CLOSED_PAYLOAD); + }); + + it("restores the typed budget when the closed section is expanded again before submit", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.type(await screen.findByLabelText(/Max Budget \(USD\)/), "150.75"); + await openSection(/Optional Settings/i); + await openSection(/Optional Settings/i); + + expect(await screen.findByLabelText(/Max Budget \(USD\)/)).toHaveValue(150.75); + + await submit(); + + expect(await createdPayload()).toStrictEqual({ ...OPTIONAL_OPEN_PAYLOAD, max_budget: "150.75" }); + }); + + it("sends a typed max budget as a string, not a number", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.type(await screen.findByLabelText(/Max Budget \(USD\)/), "150.75"); + await submit(); + + const payload = await createdPayload(); + expect(payload.max_budget).toBe("150.75"); + }); + + it.each([ + ["Tokens per minute Limit (TPM)", "tpm_limit"], + ["Requests per minute Limit (RPM)", "rpm_limit"], + ])("routes a typed %s into the %s payload key", async (label, key) => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.type(await screen.findByLabelText(label), "42"); + await submit(); + + expect((await createdPayload())[key]).toBe("42"); + }); + + it("routes the shared rate-limit-type control into its own payload key", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.click(await screen.findByLabelText(/TPM Rate Limit Type/)); + await userEvent.click(await screen.findByRole("option", { name: /Guaranteed throughput/ })); + await submit(); + + const payload = await createdPayload(); + expect(payload.tpm_limit_type).toBe("guaranteed_throughput"); + expect(payload.rpm_limit_type).toBeNull(); + }); + + it("routes a typed expiry into duration, which is otherwise coalesced to null", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await openSection(SECTIONS.lifecycle); + await userEvent.type(await screen.findByLabelText("Expire Key"), "45d"); + await submit(); + + expect((await createdPayload()).duration).toBe("45d"); + }); + + it("drops a typed expiry back to null when its section is collapsed before submit", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await openSection(SECTIONS.lifecycle); + await userEvent.type(await screen.findByLabelText("Expire Key"), "45d"); + await openSection(SECTIONS.lifecycle); + await submit(); + + expect((await createdPayload()).duration).toBeNull(); + }); + + it("routes a selected budget reset window into budget_duration", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.click(await screen.findByLabelText("Reset Budget")); + await userEvent.click(await screen.findByRole("option", { name: "daily" })); + await submit(); + + expect((await createdPayload()).budget_duration).toBe("24h"); + }); + + it("sends an explicit null budget_duration when the never-resets window is chosen", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.click(await screen.findByLabelText("Reset Budget")); + await userEvent.click(await screen.findByRole("option", { name: /never resets/i })); + await submit(); + + const payload = await createdPayload(); + expect("budget_duration" in payload).toBe(true); + expect(payload.budget_duration).toBeNull(); + }); + + it("routes typed tags into the tags key", async () => { + state.tags = { production: { name: "production" } }; + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.type(await screen.findByLabelText("Tags"), "production{Enter}"); + await submit(); + + expect((await createdPayload()).tags).toStrictEqual(["production"]); + }); + + it("routes a chosen access group into access_group_ids", async () => { + state.accessGroups = [{ access_group_id: "ag-1", access_group_name: "Group One" }]; + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + + await openAntdSelect(await screen.findByText("Select access groups (optional)")); + await userEvent.click(await screen.findByText("Group One")); + await submit(); + + expect((await createdPayload()).access_group_ids).toStrictEqual(["ag-1"]); + }); + + it("moves the schema-driven Advanced Settings fields onto the payload under their own keys", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await openSection(SECTIONS.advanced); + await userEvent.type(await screen.findByLabelText("Soft Budget"), "12"); + await submit(); + + const payload = await createdPayload(); + expect(payload.soft_budget).toBe(12); + expect(payload).toHaveProperty("key"); + }); + + it("drops the schema-driven custom key field when the proxy disables custom API keys", async () => { + state.uiSettings = { disable_custom_api_keys: true }; + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await openSection(SECTIONS.advanced); + await screen.findByLabelText("Soft Budget"); + await submit(); + + const payload = await createdPayload(); + expect(payload).not.toHaveProperty("key"); + expect(payload).toHaveProperty("soft_budget"); + }); + + it("drops the policy and prompt keys entirely for a role that cannot see those fields", async () => { + state.can = { viewPolicies: false, viewPrompts: false }; + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await submit(); + + const payload = await createdPayload(); + expect(payload).not.toHaveProperty("policies"); + expect(payload).not.toHaveProperty("prompts"); + expect(payload).toHaveProperty("guardrails"); + }); + + it("adds project_id only when the projects UI is enabled", async () => { + state.uiSettings = { enable_projects_ui: true }; + await openModal(); + await nameTheKey(); + await submit(); + + expect(await createdPayload()).toStrictEqual({ ...ALL_CLOSED_PAYLOAD, project_id: undefined }); + }); + + it("keeps disable_global_guardrails out of the payload while the switch is off", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await submit(); + + expect(await createdPayload()).not.toHaveProperty("disable_global_guardrails"); + }); + + it("sends disable_global_guardrails once the switch is on", async () => { + state.authorized = { ...state.authorized, premiumUser: true }; + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.click(await screen.findByLabelText("Disable Global Guardrails")); + await submit(); + + expect((await createdPayload()).disable_global_guardrails).toBe(true); + }); + + it("folds a metadata JSON string back through JSON.stringify", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.type(await screen.findByLabelText("Metadata"), '{{"team":"research"}'); + await submit(); + + expect((await createdPayload()).metadata).toBe('{"team":"research"}'); }); }); - it("carries the shared lifecycle expiry into the payload once its section is open", async () => { - await openModal(); + describe("key ownership", () => { + it("stamps the signed-in user onto user_id when the key is owned by you", async () => { + await openModal(); + await nameTheKey(); + await submit(); - await userEvent.type(screen.getByLabelText(/Key Name/), "probe-key"); - await userEvent.click(screen.getByText("Optional Settings")); - await userEvent.click(await screen.findByText("Key Lifecycle")); - - await userEvent.type(await screen.findByLabelText("Expire Key"), "45d"); - - await userEvent.click(screen.getByRole("button", { name: /create key/i })); - - await waitFor(() => { - expect(mockKeyCreateCall).toHaveBeenCalled(); + expect((await createdPayload()).user_id).toBe("test-user-id"); + }); + + it("mounts the user search control only once Another User is chosen", async () => { + await openModal(); + expect(screen.queryByText("Type email to search for users")).not.toBeInTheDocument(); + + await userEvent.click(screen.getByRole("radio", { name: "Another User" })); + + expect(await screen.findByText("Type email to search for users")).toBeInTheDocument(); + }); + + it("hides the Another User option from a non-admin", async () => { + state.authorized = { ...state.authorized, userRole: "Internal User" }; + await openModal(); + + expect(screen.queryByRole("radio", { name: "Another User" })).not.toBeInTheDocument(); + }); + + it("routes a service account through the service account endpoint and stamps the alias into metadata", async () => { + state.teams = [{ team_id: "team-1", team_alias: "Team One", models: [] }]; + await openModal({ teams: state.teams as unknown as Team[] }); + await userEvent.click(screen.getByRole("radio", { name: "Service Account" })); + await userEvent.type(await screen.findByLabelText(/Service Account ID/), "svc-account-1"); + + await userEvent.click(await screen.findByLabelText("Team")); + await userEvent.click(await screen.findByRole("option", { name: /Team One/ })); + + await submit(); + + await waitFor(() => { + expect(vi.mocked(keyCreateServiceAccountCall)).toHaveBeenCalled(); + }); + expect(vi.mocked(keyCreateCall)).not.toHaveBeenCalled(); + const payload = vi.mocked(keyCreateServiceAccountCall).mock.calls[0][1] as Record; + expect(JSON.parse(String(payload.metadata))).toStrictEqual({ service_account_id: "svc-account-1" }); + expect(payload).not.toHaveProperty("user_id"); }); - expect(mockKeyCreateCall.mock.calls[0][2]).toMatchObject({ duration: "45d" }); }); - it("preserves a value typed in a section that is collapsed and reopened before submit", async () => { - await openModal(); + describe("required field validation", () => { + it("blocks the submit and marks the alias invalid when it is blank", async () => { + await openModal(); + await submit(); - await userEvent.type(screen.getByLabelText(/Key Name/), "probe-key"); - await userEvent.click(screen.getByText("Optional Settings")); - - const maxBudget = await screen.findByLabelText(/Max Budget/); - await userEvent.type(maxBudget, "150.75"); - - await userEvent.click(screen.getByText("Optional Settings")); - await userEvent.click(screen.getByText("Optional Settings")); - - expect(await screen.findByLabelText(/Max Budget/)).toHaveValue(150.75); - - await userEvent.click(screen.getByRole("button", { name: /create key/i })); - - await waitFor(() => { - expect(mockKeyCreateCall).toHaveBeenCalled(); + await waitFor(() => { + expect(screen.getByLabelText(/Key Name/)).toHaveAttribute("aria-invalid", "true"); + }); + expect(vi.mocked(keyCreateCall)).not.toHaveBeenCalled(); + }); + + it("suppresses the required message behind the always-visible help text", async () => { + await openModal(); + await submit(); + + await waitFor(() => { + expect(screen.getByLabelText(/Key Name/)).toHaveAttribute("aria-invalid", "true"); + }); + expect(screen.queryByText("Please input a key name")).not.toBeInTheDocument(); + expect(screen.getByText("required")).toBeInTheDocument(); + }); + + it("blocks a service account submit until a team is chosen, then lets it through", async () => { + state.teams = [{ team_id: "team-1", team_alias: "Team One", models: [] }]; + await openModal({ teams: state.teams as unknown as Team[] }); + await userEvent.click(screen.getByRole("radio", { name: "Service Account" })); + await userEvent.type(await screen.findByLabelText(/Service Account ID/), "svc-account-1"); + await submit(); + + await waitFor(() => { + expect(vi.mocked(keyCreateServiceAccountCall)).not.toHaveBeenCalled(); + }); + + await userEvent.click(await screen.findByLabelText("Team")); + await userEvent.click(await screen.findByRole("option", { name: /Team One/ })); + await submit(); + + await waitFor(() => { + expect(vi.mocked(keyCreateServiceAccountCall)).toHaveBeenCalledTimes(1); + }); + }); + + it("blocks the submit when the budget exceeds the team ceiling", async () => { + await openModal({ team: { team_id: "team-1", max_budget: 10 } as unknown as Team }); + await nameTheKey(); + await openSection(/Optional Settings/i); + await userEvent.type(await screen.findByLabelText(/Max Budget \(USD\)/), "50"); + await submit(); + + await waitFor(() => { + expect(screen.getByLabelText(/Max Budget \(USD\)/)).toHaveAttribute("aria-invalid", "true"); + }); + expect(vi.mocked(keyCreateCall)).not.toHaveBeenCalled(); + }); + }); + + describe("deep link prefill", () => { + it("prefills the key alias", async () => { + renderCreateKey({ autoOpenCreate: true, prefillData: { key_alias: "prefilled-key" } }); + + expect(await screen.findByLabelText(/Key Name/)).toHaveValue("prefilled-key"); + }); + + it("prefills models once the available model list arrives", async () => { + renderCreateKey({ autoOpenCreate: true, prefillData: { models: ["gpt-4"] } }); + + expect(await screen.findByTitle("gpt-4")).toBeInTheDocument(); + }); + + it("ignores a team the user has no access to", async () => { + renderCreateKey({ + teams: [{ team_id: "team-1", models: [] } as unknown as Team], + autoOpenCreate: true, + prefillData: { team_id: "team-404", key_alias: "example-key" }, + }); + + await userEvent.type(await screen.findByLabelText(/Key Name/), "-suffix"); + await submit(); + + const payload = await createdPayload(); + expect(payload.team_id).toBeNull(); + expect(payload.key_alias).toBe("example-key-suffix"); + }); + + it("falls back to you when another_user is requested by a non-admin", async () => { + state.authorized = { ...state.authorized, userRole: "Internal User" }; + renderCreateKey({ autoOpenCreate: true, prefillData: { owned_by: "another_user", key_alias: "example-key" } }); + + expect(await screen.findByRole("radio", { name: "You" })).toBeChecked(); + }); + + it("applies another_user for an admin", async () => { + renderCreateKey({ autoOpenCreate: true, prefillData: { owned_by: "another_user" } }); + + expect(await screen.findByRole("radio", { name: "Another User" })).toBeChecked(); + }); + + it("prefills the key type", async () => { + renderCreateKey({ autoOpenCreate: true, prefillData: { key_type: "management" } }); + + await screen.findByLabelText(/Key Name/); + await userEvent.type(await screen.findByLabelText(/Key Name/), "prefilled-type"); + await submit(); + + expect((await createdPayload()).key_type).toBe("management"); + }); + }); + + describe("models dropdown team gating", () => { + it("offers all-proxy-models but not all-team-models when no team is selected", async () => { + await openModal(); + await userEvent.click(await screen.findByLabelText("Models")); + + expect(await screen.findByTitle("All Proxy Models")).toBeInTheDocument(); + expect(screen.queryByTitle("All Team Models")).not.toBeInTheDocument(); + }); + + it("offers all-team-models but hides all-proxy-models once a team is selected", async () => { + state.teams = [{ team_id: "team-1", team_alias: "Team One", models: ["team-model-1"] }]; + await openModal({ teams: state.teams as unknown as Team[] }); + + await userEvent.click(await screen.findByLabelText("Team")); + await userEvent.click(await screen.findByRole("option", { name: /Team One/ })); + + await userEvent.click(await screen.findByLabelText("Models")); + + expect(await screen.findByTitle("All Team Models")).toBeInTheDocument(); + expect(screen.queryByTitle("All Proxy Models")).not.toBeInTheDocument(); + }); + }); + + describe("organization dropdown", () => { + it("is editable for an admin", async () => { + state.organizations = [{ organization_id: "org-1", organization_alias: "Engineering" }]; + await openModal(); + + expect(await screen.findByLabelText("Organization")).not.toHaveAttribute("data-disabled"); + }); + + it("is read-only for a non-admin", async () => { + state.authorized = { ...state.authorized, userRole: "Internal User" }; + state.organizations = [{ organization_id: "org-1", organization_alias: "Engineering" }]; + await openModal(); + + expect(await screen.findByLabelText("Organization")).toHaveAttribute("data-disabled", ""); + }); + + it("routes a chosen organization into organization_id", async () => { + state.organizations = [{ organization_id: "org-1", organization_alias: "Engineering" }]; + await openModal(); + await nameTheKey(); + + await userEvent.click(await screen.findByLabelText("Organization")); + await userEvent.click(await screen.findByRole("option", { name: /Engineering/ })); + await submit(); + + expect((await createdPayload()).organization_id).toBe("org-1"); + }); + }); + + describe("policy and prompt fields", () => { + it("loads and offers both selectors for a role that can see them", async () => { + await openModal(); + await openSection(/Optional Settings/i); + + expect(await screen.findByLabelText("Policies")).toBeInTheDocument(); + expect(await screen.findByLabelText("Prompts")).toBeInTheDocument(); + }); + + it("omits both selectors for a role that cannot see them", async () => { + state.can = { viewPolicies: false, viewPrompts: false }; + await openModal(); + await openSection(/Optional Settings/i); + + await screen.findByLabelText("Guardrails"); + expect(screen.queryByLabelText("Policies")).not.toBeInTheDocument(); + expect(screen.queryByLabelText("Prompts")).not.toBeInTheDocument(); + }); + }); + + describe("tags dropdown", () => { + it("offers the tags returned by the tags hook", async () => { + state.tags = { production: { name: "production" }, staging: { name: "staging" } }; + await openModal(); + await openSection(/Optional Settings/i); + + await userEvent.click(await screen.findByLabelText("Tags")); + + expect(await screen.findByTitle("production")).toBeInTheDocument(); + expect(await screen.findByTitle("staging")).toBeInTheDocument(); + }); + }); + + describe("user search debounce", () => { + it("fires exactly one search carrying the last typed value", async () => { + const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime }); + vi.useFakeTimers({ shouldAdvanceTime: true }); + try { + renderCreateKey({ autoOpenCreate: true, prefillData: { owned_by: "another_user" } }); + const search = antdSearchInput(await screen.findByText("Type email to search for users")); + + await user.type(search, "ali"); + expect(vi.mocked(userFilterUICall)).not.toHaveBeenCalled(); + + await user.type(search, "ce"); + await vi.advanceTimersByTimeAsync(400); + + expect(vi.mocked(userFilterUICall)).toHaveBeenCalledTimes(1); + const params = vi.mocked(userFilterUICall).mock.calls[0][1] as URLSearchParams; + expect(params.get("user_email")).toBe("alice"); + } finally { + vi.useRealTimers(); + } + }); + }); + + describe("created key display", () => { + it("surfaces the generated key after a successful create", async () => { + await openModal(); + await nameTheKey(); + await submit(); + await createdPayload(); + + expect(await screen.findByText("Save your Key")).toBeInTheDocument(); + }); + + it("rejects a duplicate alias within the same team without calling the API", async () => { + await openModal({ data: [{ team_id: null, key_alias: "contract-key" }] }); + await nameTheKey(); + await submit(); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /^create key$/i })).toBeInTheDocument(); + }); + expect(vi.mocked(keyCreateCall)).not.toHaveBeenCalled(); + }); + }); + + describe("key type gating", () => { + it("labels the llm_api option AI APIs", async () => { + await openModal(); + await userEvent.click(await screen.findByLabelText("Key Type")); + + expect( + await screen.findByText("Can call only AI API routes (chat/completions, embeddings, etc.)"), + ).toBeInTheDocument(); + expect(screen.queryByText("LLM API")).not.toBeInTheDocument(); + }); + + it("clears and disables models when a management key type is chosen", async () => { + await openModal(); + await nameTheKey(); + + await userEvent.click(await screen.findByLabelText("Key Type")); + await userEvent.click(await screen.findByText("Management")); + await submit(); + + const payload = await createdPayload(); + expect(payload.key_type).toBe("management"); + expect(payload.models).toStrictEqual([]); + }); + }); + + describe("mount-gate liveness", () => { + it("keeps the MCP tool permissions key out of the payload even with its section open", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await openSection(SECTIONS.mcp); + await submit(); + + const payload = await createdPayload(); + expect(payload).not.toHaveProperty("mcp_tool_permissions"); + expect(payload).toHaveProperty("allowed_mcp_servers_and_groups"); + }); + + it("registers no Optional Settings field while a team choice is still required", async () => { + vi.mocked(modelAvailableCall).mockResolvedValue({ data: [{ id: "no-default-models" }] }); + renderCreateKey(); + await userEvent.click(screen.getByTestId("create-key-button")); + + expect(await screen.findByText(/Please select a team to continue/)).toBeInTheDocument(); + expect(screen.queryByLabelText(/Key Name/)).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /Optional Settings/i })).not.toBeInTheDocument(); }); - expect(mockKeyCreateCall.mock.calls[0][2]).toMatchObject({ max_budget: "150.75" }); }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx index 7c9eed62c89..e6bd6cddd38 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx @@ -1,891 +1,76 @@ -import { act, fireEvent, within } from "@testing-library/react"; -import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; -import { Team } from "../key_team_helpers/key_list"; -import { getPoliciesList, getPromptsList, userFilterUICall } from "../networking"; -import CreateKey from "./create_key_button"; - -const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall, teamDropdownTeamsRef } = - vi.hoisted(() => { - const formStateRef = { current: {} as Record }; - const teamDropdownTeamsRef = { current: [] as Array<{ team_id: string; team_alias: string; models: string[] }> }; - const mockKeyCreateCall = vi.fn().mockResolvedValue({ - key: "test-api-key", - soft_budget: null, - }); - const formMock = { - setFieldsValue: vi.fn((values: Record) => { - Object.assign(formStateRef.current, values); - }), - setFieldValue: vi.fn((name: string, value: any) => { - formStateRef.current[name] = value; - }), - getFieldValue: vi.fn((name: string) => formStateRef.current[name]), - resetFields: vi.fn(() => { - formStateRef.current = {}; - }), - }; - const radioGroupValueRef = { current: null as string | null }; - return { - formMock, - setFieldsValueMock: formMock.setFieldsValue, - radioGroupValueRef, - formStateRef, - mockKeyCreateCall, - teamDropdownTeamsRef, - }; - }); - -const defaultAuthorizedState = { - accessToken: "test-token", - userId: "test-user-id", - userRole: "Admin", - premiumUser: false, -}; - -let authorizedState = { ...defaultAuthorizedState }; - -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => authorizedState, -})); - -vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ - keyKeys: { - lists: () => ["keys"], - }, -})); - -vi.mock("@ant-design/icons", () => ({ - InfoCircleOutlined: () => null, -})); - -vi.mock("react-copy-to-clipboard", () => ({ - CopyToClipboard: ({ children }: { children: any }) => children, -})); - -vi.mock("antd", () => { - const React = require("react"); - - const getValueFromEvent = (event: any) => { - if (event?.target) { - if (event.target.type === "checkbox") { - return event.target.checked; - } - return event.target.value; - } - return event; - }; - - const Form = ({ - children, - onFinish, - ...props - }: { - children?: any; - onFinish?: (values: Record) => void; - }) => - React.createElement( - "form", - { - ...props, - onSubmit: (event: Event) => { - event.preventDefault(); - onFinish?.({ ...formStateRef.current }); - }, - }, - children, - ); - - Form.Item = ({ children, name }: { children?: any; name?: string }) => { - if (!name || !React.isValidElement(children)) { - return React.createElement(React.Fragment, null, children); - } - - return React.cloneElement(children, { - value: formStateRef.current[name], - onChange: (event: any) => { - formStateRef.current[name] = getValueFromEvent(event); - }, - }); - }; - - Form.useForm = () => [formMock]; - - Form.useWatch = (name: string) => formStateRef.current[name]; - - const Select = ({ - children, - onChange, - onSearch, - options, - ...props - }: { - children?: any; - onChange?: (value: string) => void; - onSearch?: (value: string) => void; - options?: Array<{ value: string; label: string }>; - }) => { - const select = React.createElement( - "select", - { - ...props, - onChange: (event: any) => onChange?.(event.target.value), - }, - children, - options?.map((opt: any) => React.createElement("option", { key: opt.value, value: opt.value }, opt.label)), - ); - - if (!onSearch) { - return select; - } - - return React.createElement( - React.Fragment, - null, - React.createElement("input", { - "data-testid": "select-search-input", - onChange: (event: React.ChangeEvent) => onSearch(event.target.value), - }), - select, - ); - }; - - Select.Option = ({ children, ...props }: { children?: any }) => React.createElement("option", props, children); - - const Input = (props: any) => React.createElement("input", props); - Input.Password = (props: any) => React.createElement("input", { ...props, type: "password" }); - Input.TextArea = (props: any) => React.createElement("textarea", props); - - const Modal = ({ children, open }: { children?: any; open?: boolean }) => - open ? React.createElement("div", null, children) : null; - - const Radio = ({ children, ...props }: { children?: any }) => React.createElement("div", props, children); - - Radio.Group = ({ children, value }: { children?: any; value?: string }) => { - radioGroupValueRef.current = value ?? null; - return React.createElement("div", null, children); - }; - - const Switch = (props: any) => React.createElement("input", { ...props, type: "checkbox" }); - const Tag = ({ children }: { children?: any }) => React.createElement("span", null, children); - const Tooltip = ({ children }: { children?: any }) => React.createElement(React.Fragment, null, children); - - const Button = ({ children, htmlType, ...props }: { children?: any; htmlType?: string }) => - React.createElement("button", { ...props, type: htmlType ?? props.type }, children); - - const Typography = ({ children, ...props }: { children?: any }) => React.createElement("div", props, children); - Typography.Text = ({ children, ...props }: { children?: any }) => React.createElement("span", props, children); - Typography.Paragraph = ({ children, ...props }: { children?: any }) => React.createElement("p", props, children); - Typography.Title = ({ children, ...props }: { children?: any }) => React.createElement("h1", props, children); - - return { - Button, - Form, - Input, - message: { - success: vi.fn(), - error: vi.fn(), - warning: vi.fn(), - info: vi.fn(), - }, - Modal, - Radio, - Select, - Switch, - Tag, - Tooltip, - Typography, - }; -}); +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { modelAvailableCall } from "../networking"; +import { fetchTeamModels, fetchUserModels } from "./create_key_button"; vi.mock("../networking", () => ({ - keyCreateCall: mockKeyCreateCall, - modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }] }), - getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), - getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), - getPromptsList: vi.fn().mockResolvedValue({ prompts: [] }), - proxyBaseUrl: "http://localhost:4000", - getPossibleUserRoles: vi.fn().mockResolvedValue({ - Admin: { ui_label: "Admin" }, - User: { ui_label: "User" }, - }), - userFilterUICall: vi.fn().mockResolvedValue([]), - keyCreateServiceAccountCall: vi.fn().mockResolvedValue({ - key: "test-service-account-key", - soft_budget: null, - }), - fetchMCPAccessGroups: vi.fn().mockResolvedValue([]), - getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), + modelAvailableCall: vi.fn(), })); -vi.mock("../agent_management/AgentSelector", () => ({ default: () => null })); -vi.mock("../common_components/budget_duration_dropdown", () => ({ - NEVER_RESETS_BUDGET_DURATION: "none", - default: ({ - showNeverResets, - placeholder, - onChange, - }: { - showNeverResets?: boolean; - placeholder?: string; - onChange?: (value: string) => void; - }) => ( - - ), -})); -vi.mock("../common_components/check_openapi_schema", () => ({ default: () => null })); -vi.mock("../common_components/KeyLifecycleSettings", () => ({ default: () => null })); -vi.mock("../common_components/ModelAliasManager", () => ({ default: () => null })); -vi.mock("../common_components/PassThroughRoutesSelector", () => ({ default: () => null })); -vi.mock("../common_components/PremiumLoggingSettings", () => ({ default: () => null })); -vi.mock("../common_components/RateLimitTypeFormItem", () => ({ default: () => null })); -vi.mock("../common_components/RouterSettingsAccordion", () => ({ default: () => null })); -vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ - useInfiniteTeams: () => ({ - data: { - pages: [ - { - teams: [ - { team_id: "team-1", team_alias: "Team One" }, - { team_id: "team-2", team_alias: "Team Two" }, - ], - total: 2, - page: 1, - page_size: 50, - total_pages: 1, - }, - ], - }, - fetchNextPage: vi.fn(), - hasNextPage: false, - isFetchingNextPage: false, - isLoading: false, - }), -})); -vi.mock("../common_components/team_dropdown", () => ({ - default: ({ - onTeamSelect, - disabled, - }: { - onTeamSelect?: (team: { team_id: string; team_alias: string; models: string[] } | null) => void; - disabled?: boolean; - }) => ( - - ), -})); -vi.mock("../CreateUserButton", () => ({ CreateUserButton: () => null })); -vi.mock("../mcp_server_management/MCPServerSelector", () => ({ default: () => null })); -vi.mock("../mcp_server_management/MCPToolPermissions", () => ({ default: () => null })); -vi.mock("../shared/numerical_input", () => ({ default: () => null })); -vi.mock("../vector_store_management/VectorStoreSelector", () => ({ default: () => null })); -vi.mock("../key_team_helpers/fetch_available_models_team_key", async () => { - const actual = await vi.importActual("../key_team_helpers/fetch_available_models_team_key"); - return { - ...actual, - getModelDisplayName: (model: string) => model, - }; -}); - -vi.mock("@/app/(dashboard)/hooks/tags/useTags", () => ({ - useTags: vi.fn().mockReturnValue({ - data: [ - { name: "production", description: "Prod tag", models: [], created_at: "2026-01-01", updated_at: "2026-01-01" }, - { name: "staging", description: "Staging tag", models: [], created_at: "2026-01-01", updated_at: "2026-01-01" }, - ], - isLoading: false, - }), -})); - -vi.mock("@/app/(dashboard)/hooks/projects/useProjects", () => ({ - useProjects: vi.fn().mockReturnValue({ data: [], isLoading: false }), -})); - -vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ - useOrganizations: vi.fn().mockReturnValue({ - data: [ - { organization_id: "org-1", organization_alias: "Engineering" }, - { organization_id: "org-2", organization_alias: "Sales" }, - ], - isLoading: false, - }), -})); - -vi.mock("../common_components/OrganizationDropdown", () => ({ - default: ({ value, onChange, disabled }: { value?: string; onChange?: (v: string) => void; disabled?: boolean }) => ( - - ), -})); - -vi.mock("../common_components/ProjectDropdown", () => ({ - default: ({ value, onChange }: { value?: string; onChange?: (v: string) => void }) => ( - onChange?.(e.target.value)} /> - ), -})); - -vi.mock("../common_components/AccessGroupSelector", () => ({ - default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( - onChange?.(event.target.value ? event.target.value.split(",").map((v) => v.trim()) : [])} - /> - ), -})); - -describe("CreateKey", () => { - const defaultProps = { - team: null, - teams: [], - data: [], - addKey: vi.fn(), - }; - - const openOptionalSettings = async () => { - const trigger = await screen.findByRole("button", { name: /optional settings/i }); - act(() => { - fireEvent.click(trigger); - }); - }; - +describe("fetchTeamModels", () => { beforeEach(() => { - vi.clearAllMocks(); - if (typeof window !== "undefined" && window.localStorage && typeof window.localStorage.clear === "function") { - window.localStorage.clear(); - } - authorizedState = { ...defaultAuthorizedState }; - radioGroupValueRef.current = null; - formStateRef.current = {}; - teamDropdownTeamsRef.current = [ - { team_id: "team-1", team_alias: "Team One", models: [] }, - { team_id: "team-2", team_alias: "Team Two", models: [] }, - ]; - mockKeyCreateCall.mockResolvedValue({ - key: "test-api-key", - soft_budget: null, - }); + vi.mocked(modelAvailableCall).mockReset(); }); - it("should render the CreateKey component", () => { - renderWithProviders(); - expect(screen.getByRole("button", { name: /create new key/i })).toBeInTheDocument(); + it("asks the proxy for the team-scoped model list and returns the ids", async () => { + vi.mocked(modelAvailableCall).mockResolvedValue({ data: [{ id: "gpt-4" }, { id: "claude-opus-4" }] }); + + await expect(fetchTeamModels("user-1", "Admin", "token-1", "team-1")).resolves.toStrictEqual([ + "gpt-4", + "claude-opus-4", + ]); + expect(modelAvailableCall).toHaveBeenCalledWith("token-1", "user-1", "Admin", true, "team-1", true); }); - it("should display 'AI APIs' label for the llm_api key type option", async () => { - renderWithProviders(); + it("passes a null team through rather than dropping the argument", async () => { + vi.mocked(modelAvailableCall).mockResolvedValue({ data: [] }); - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); + await fetchTeamModels("user-1", "Admin", "token-1", null); - await waitFor(() => { - expect(screen.getByText("AI APIs")).toBeInTheDocument(); - expect(screen.queryByText("LLM API")).not.toBeInTheDocument(); - }); + expect(modelAvailableCall).toHaveBeenCalledWith("token-1", "user-1", "Admin", true, null, true); }); - it("should include access_group_ids in keyCreateCall payload when access groups are selected", async () => { - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - await openOptionalSettings(); - - await waitFor(() => { - expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); - }); - - act(() => { - fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } }); - formMock.setFieldValue("key_alias", "Test Key"); - }); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create key/i })); - }); - - await waitFor(() => { - expect(mockKeyCreateCall).toHaveBeenCalled(); - const formValues = mockKeyCreateCall.mock.calls[0][2]; - expect(formValues).toHaveProperty("access_group_ids"); - expect(formValues.access_group_ids).toEqual(["ag-1", "ag-2"]); - }); + it("returns an empty list and makes no call when the user id is null", async () => { + await expect(fetchTeamModels(null as unknown as string, "Admin", "token-1", "team-1")).resolves.toStrictEqual([]); + expect(modelAvailableCall).not.toHaveBeenCalled(); }); - it("should include mcp_toolsets in keyCreateCall payload when only toolsets are selected", async () => { - renderWithProviders(); + it("swallows a failed lookup and returns an empty list", async () => { + vi.mocked(modelAvailableCall).mockRejectedValue(new Error("proxy down")); - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /create key/i })).toBeInTheDocument(); - }); - - act(() => { - formMock.setFieldValue("key_alias", "Test Key"); - formMock.setFieldValue("allowed_mcp_servers_and_groups", { - servers: [], - accessGroups: [], - toolsets: ["ts-1"], - }); - }); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create key/i })); - }); - - await waitFor(() => { - expect(mockKeyCreateCall).toHaveBeenCalled(); - }); - expect(mockKeyCreateCall.mock.calls[0][2].object_permission?.mcp_toolsets).toEqual(["ts-1"]); - }); - - it("should prefill models when provided without team_id", async () => { - renderWithProviders( - , - ); - - await waitFor(() => { - expect(setFieldsValueMock).toHaveBeenCalledWith({ models: ["gpt-4"] }); - }); - }); - - it("should prefill team_id when it exists in teams", async () => { - renderWithProviders( - , - ); - - await waitFor(() => { - expect(setFieldsValueMock).toHaveBeenCalledWith({ team_id: "team-1" }); - }); - }); - - it("should ignore team_id when it does not exist in teams", async () => { - renderWithProviders( - , - ); - - await waitFor(() => { - expect(setFieldsValueMock).toHaveBeenCalledWith({ key_alias: "example-key" }); - }); - - expect(setFieldsValueMock).not.toHaveBeenCalledWith({ team_id: "team-404" }); - }); - - it('should fall back to "you" when owned_by is another_user for non-admin', async () => { - authorizedState = { ...defaultAuthorizedState, userRole: "Internal User" }; - - renderWithProviders( - , - ); - - await waitFor(() => { - expect(setFieldsValueMock).toHaveBeenCalledWith({ key_alias: "example-key" }); - }); - - expect(radioGroupValueRef.current).toBe("you"); - }); - - it("should apply owned_by another_user for admin", async () => { - renderWithProviders( - , - ); - - await waitFor(() => { - expect(radioGroupValueRef.current).toBe("another_user"); - }); - }); - - it("should prefill key_type when provided", async () => { - renderWithProviders(); - - await waitFor(() => { - expect(setFieldsValueMock).toHaveBeenCalledWith({ key_type: "management" }); - }); - }); - - describe("organization dropdown", () => { - it("should render the organization dropdown when modal is open", async () => { - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - await waitFor(() => { - expect(screen.getByTestId("org-dropdown")).toBeInTheDocument(); - }); - }); - - it("should disable the organization dropdown for non-admin users", async () => { - authorizedState = { ...defaultAuthorizedState, userRole: "Internal User" }; - - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - await waitFor(() => { - expect(screen.getByTestId("org-dropdown")).toBeDisabled(); - }); - }); - - it("should enable the organization dropdown for admin users", async () => { - authorizedState = { ...defaultAuthorizedState, userRole: "Admin" }; - - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - await waitFor(() => { - expect(screen.getByTestId("org-dropdown")).toBeEnabled(); - }); - }); - - it("should render team dropdown alongside organization dropdown", async () => { - const teamsWithOrg = [{ team_id: "team-1", team_alias: "Team Alpha", organization_id: "org-1", models: [] }]; - - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - await waitFor(() => { - expect(screen.getByTestId("org-dropdown")).toBeInTheDocument(); - expect(screen.getByTestId("team-dropdown")).toBeInTheDocument(); - }); - }); - - it("should set organization_id in form state when org is selected", async () => { - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - await waitFor(() => { - expect(screen.getByTestId("org-dropdown")).toBeInTheDocument(); - }); - - act(() => { - fireEvent.change(screen.getByTestId("org-dropdown"), { target: { value: "org-1" } }); - }); - - expect(formStateRef.current["organization_id"]).toBe("org-1"); - }); - }); - - describe("models dropdown team gating", () => { - const getModelsSelect = async (): Promise => { - return waitFor(() => { - const element = document.querySelector('select[placeholder="Select models"]'); - expect(element).toBeTruthy(); - return element as HTMLElement; - }); - }; - - it("should offer all-proxy-models but not all-team-models when no team is selected", async () => { - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - const modelsSelect = await getModelsSelect(); - - await waitFor(() => { - expect(within(modelsSelect).getByText("gpt-4")).toBeInTheDocument(); - }); - - expect(within(modelsSelect).getByText("All Proxy Models")).toBeInTheDocument(); - expect(within(modelsSelect).queryByText("All Team Models")).not.toBeInTheDocument(); - }); - - it("should offer all-team-models but hide all-proxy-models when a team is selected", async () => { - teamDropdownTeamsRef.current = [ - { team_id: "team-1", team_alias: "Team One", models: ["all-proxy-models", "team-model-1"] }, - ]; - - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - await waitFor(() => { - expect(screen.getByTestId("team-dropdown")).toBeInTheDocument(); - }); - - act(() => { - fireEvent.change(screen.getByTestId("team-dropdown"), { target: { value: "team-1" } }); - }); - - const modelsSelect = await getModelsSelect(); - - await waitFor(() => { - expect(within(modelsSelect).getByText("team-model-1")).toBeInTheDocument(); - }); - - expect(within(modelsSelect).getByText("All Team Models")).toBeInTheDocument(); - expect(within(modelsSelect).queryByText("All Proxy Models")).not.toBeInTheDocument(); - expect(within(modelsSelect).queryByText("all-proxy-models")).not.toBeInTheDocument(); - }); - }); - - describe("user search debounce", () => { - const mockUserFilterUICall = vi.mocked(userFilterUICall); - - beforeEach(() => { - vi.useFakeTimers(); - }); - - afterEach(() => { - vi.runOnlyPendingTimers(); - vi.useRealTimers(); - }); - - const renderUserSearch = () => { - const view = renderWithProviders( - , - ); - return { input: screen.getByTestId("select-search-input"), unmount: view.unmount }; - }; - - it("should not fire the search before the wait elapses", () => { - const { input } = renderUserSearch(); - - act(() => { - fireEvent.change(input, { target: { value: "alice" } }); - }); - - expect(mockUserFilterUICall).not.toHaveBeenCalled(); - - act(() => { - vi.advanceTimersByTime(299); - }); - - expect(mockUserFilterUICall).not.toHaveBeenCalled(); - }); - - it("should fire exactly one search carrying the last value after the wait", async () => { - const { input } = renderUserSearch(); - - act(() => { - fireEvent.change(input, { target: { value: "a" } }); - vi.advanceTimersByTime(100); - fireEvent.change(input, { target: { value: "al" } }); - vi.advanceTimersByTime(100); - fireEvent.change(input, { target: { value: "alice" } }); - }); - - expect(mockUserFilterUICall).not.toHaveBeenCalled(); - - await act(async () => { - vi.advanceTimersByTime(300); - }); - - expect(mockUserFilterUICall).toHaveBeenCalledTimes(1); - const params = mockUserFilterUICall.mock.calls[0][1] as URLSearchParams; - expect(params.get("user_email")).toBe("alice"); - }); - - it("should fire nothing when unmounted mid-wait", () => { - const { input, unmount } = renderUserSearch(); - - act(() => { - fireEvent.change(input, { target: { value: "alice" } }); - }); - - unmount(); - - act(() => { - vi.advanceTimersByTime(1000); - }); - - expect(mockUserFilterUICall).not.toHaveBeenCalled(); - }); - }); - - describe("tags dropdown", () => { - it("should populate tags dropdown with options from useTags hook", async () => { - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - await openOptionalSettings(); - - await waitFor(() => { - expect(screen.getByText("production")).toBeInTheDocument(); - expect(screen.getByText("staging")).toBeInTheDocument(); - }); - }); - }); - - describe("policy and prompt fields", () => { - const POLICIES_PLACEHOLDER = "Premium feature - Upgrade to set policies by key"; - const PROMPTS_PLACEHOLDER = "Premium feature - Upgrade to set prompts by key"; - - const openModal = async () => { - renderWithProviders(); - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - await openOptionalSettings(); - }; - - beforeEach(() => { - vi.mocked(getPoliciesList).mockResolvedValue({ policies: [{ policy_name: "policy-a" }] }); - vi.mocked(getPromptsList).mockResolvedValue({ prompts: [{ prompt_id: "prompt-a" }] } as any); - }); - - it("should load and offer both selectors for an admin", async () => { - await openModal(); - - await waitFor(() => { - expect(screen.getByRole("option", { name: "policy-a" })).toBeInTheDocument(); - expect(screen.getByRole("option", { name: "prompt-a" })).toBeInTheDocument(); - }); - expect(getPoliciesList).toHaveBeenCalledWith("test-token"); - expect(getPromptsList).toHaveBeenCalledWith("test-token"); - expect(screen.getByPlaceholderText(POLICIES_PLACEHOLDER)).toBeInTheDocument(); - expect(screen.getByPlaceholderText(PROMPTS_PLACEHOLDER)).toBeInTheDocument(); - }); - - it("should omit both selectors and fire neither admin-only request for an internal user", async () => { - authorizedState = { ...defaultAuthorizedState, userRole: "Internal User" }; - - await openModal(); - - expect(await screen.findByTestId("org-dropdown")).toBeInTheDocument(); - - expect(getPoliciesList).not.toHaveBeenCalled(); - expect(getPromptsList).not.toHaveBeenCalled(); - expect(screen.queryByPlaceholderText(POLICIES_PLACEHOLDER)).not.toBeInTheDocument(); - expect(screen.queryByPlaceholderText(PROMPTS_PLACEHOLDER)).not.toBeInTheDocument(); - }); - }); - - describe("budget reset", () => { - const openModal = async () => { - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create new key/i })); - }); - - await openOptionalSettings(); - - await waitFor(() => { - expect(screen.getByTestId("budget-duration-dropdown")).toBeInTheDocument(); - }); - }; - - const submit = async () => { - act(() => { - fireEvent.click(screen.getByRole("button", { name: /create key/i })); - }); - - await waitFor(() => { - expect(mockKeyCreateCall).toHaveBeenCalled(); - }); - - return mockKeyCreateCall.mock.calls[0][2]; - }; - - it("should send an explicit null budget_duration when 'Never resets' is selected", async () => { - await openModal(); - - expect(screen.getByRole("option", { name: "Never resets" })).toBeInTheDocument(); - - act(() => { - fireEvent.change(screen.getByTestId("budget-duration-dropdown"), { target: { value: "none" } }); - formMock.setFieldValue("key_alias", "Never Resets Key"); - }); - - const formValues = await submit(); - - expect("budget_duration" in formValues).toBe(true); - expect(formValues.budget_duration).toBeNull(); - }); - - it("should label the omit option distinctly from 'Never resets'", async () => { - await openModal(); - - const optionLabels = Array.from( - screen.getByTestId("budget-duration-dropdown").querySelectorAll("option"), - (option) => option.textContent, - ); - - expect(optionLabels).toContain("Never resets"); - expect(new Set(optionLabels).size).toBe(optionLabels.length); - expect(screen.getByRole("option", { name: "Never resets" })).toHaveValue("none"); - }); - - it("should omit budget_duration entirely when the reset dropdown is untouched", async () => { - await openModal(); - - act(() => { - formMock.setFieldValue("key_alias", "Inherits Default Key"); - }); - - const formValues = await submit(); - - expect("budget_duration" in formValues).toBe(false); - }); + await expect(fetchTeamModels("user-1", "Admin", "token-1", "team-1")).resolves.toStrictEqual([]); + }); +}); + +describe("fetchUserModels", () => { + beforeEach(() => { + vi.mocked(modelAvailableCall).mockReset(); + }); + + it("hands the returned ids to the setter without the team-scoped arguments", async () => { + vi.mocked(modelAvailableCall).mockResolvedValue({ data: [{ id: "gpt-4" }] }); + const setUserModels = vi.fn(); + + await fetchUserModels("user-1", "Admin", "token-1", setUserModels); + + expect(modelAvailableCall).toHaveBeenCalledWith("token-1", "user-1", "Admin"); + expect(setUserModels).toHaveBeenCalledWith(["gpt-4"]); + }); + + it("leaves the setter untouched when the lookup fails", async () => { + vi.mocked(modelAvailableCall).mockRejectedValue(new Error("proxy down")); + const setUserModels = vi.fn(); + + await fetchUserModels("user-1", "Admin", "token-1", setUserModels); + + expect(setUserModels).not.toHaveBeenCalled(); + }); + + it("makes no call when the user role is null", async () => { + const setUserModels = vi.fn(); + + await fetchUserModels("user-1", null as unknown as string, "token-1", setUserModels); + + expect(modelAvailableCall).not.toHaveBeenCalled(); + expect(setUserModels).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 9bb5ed51bfa..77ef31e0ede 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -30,9 +30,8 @@ import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import React, { useEffect, useState } from "react"; import { rolesWithWriteAccess } from "../../utils/roles"; import AgentSelector from "../agent_management/AgentSelector"; -import { mapDisplayToInternalNames } from "../callback_info_helpers"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; -import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown"; +import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; import SchemaFormFields from "../common_components/check_openapi_schema"; import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; import ModelAliasManager from "../common_components/ModelAliasManager"; @@ -46,7 +45,7 @@ import ProjectDropdown from "../common_components/ProjectDropdown"; import { CreateUserButton } from "../CreateUserButton"; import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor"; import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor"; -import { TagRateLimitEditor, TagRateLimitEntry, tagRowsToLimits } from "../key_team_helpers/TagRateLimitEditor"; +import { TagRateLimitEditor, TagRateLimitEntry } from "../key_team_helpers/TagRateLimitEditor"; import { excludeProxyWideSentinel, getModelDisplayName, @@ -72,6 +71,7 @@ import { import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; import NumericalInput from "../shared/numerical_input"; import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; +import { buildKeyCreatePayload, type KeyCreateInput } from "./createKeyPayload"; import { simplifyKeyGenerateError } from "./utils"; const { Option } = Select; @@ -378,198 +378,42 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp const handleCreate = async (formValues: Record) => { try { - const newKeyAlias = formValues?.key_alias ?? ""; - const newKeyTeamId = formValues?.team_id ?? null; - - const existingKeyAliases = data?.filter((k) => k.team_id === newKeyTeamId).map((k) => k.key_alias) ?? []; - - if (existingKeyAliases.includes(newKeyAlias)) { + const input: KeyCreateInput = { + formValues, + existingKeys: data, + keyOwner, + userID, + selectedAgentId, + loggingSettings, + disabledCallbacks, + autoRotationEnabled, + rotationInterval, + modelAliases, + routerSettings, + budgetLimits, + tagRateLimits, + budgetFallbacks, + }; + const built = buildKeyCreatePayload(input); + if (built.kind === "duplicate_alias") { throw new Error( - `Key alias ${newKeyAlias} already exists for team with ID ${newKeyTeamId}, please provide another key alias`, + `Key alias ${built.alias} already exists for team with ID ${built.teamId}, please provide another key alias`, ); } toast.info("Making API Call"); setIsModalVisible(true); - if (keyOwner === "you") { - formValues.user_id = userID; - } else if (keyOwner === "agent") { - if (!selectedAgentId) { - toast.fromError("Please select an agent"); - return; - } - formValues.agent_id = selectedAgentId; + if (built.kind === "agent_not_selected") { + toast.fromError("Please select an agent"); + return; } + const { payload, endpoint } = built; - // Handle metadata for all key types - let metadata: Record = {}; - try { - metadata = JSON.parse(formValues.metadata || "{}"); - } catch (error) { - console.error("Error parsing metadata:", error); - } - - // If it's a service account, add the service_account_id to the metadata - if (keyOwner === "service_account") { - metadata["service_account_id"] = formValues.key_alias; - } - - // Add logging settings to the metadata - if (loggingSettings.length > 0) { - metadata = { - ...metadata, - logging: loggingSettings.filter((config) => config.callback_name), - }; - } - - // Add disabled callbacks to the metadata - if (disabledCallbacks.length > 0) { - // Map display names to internal callback values - const mappedDisabledCallbacks = mapDisplayToInternalNames(disabledCallbacks); - metadata = { - ...metadata, - litellm_disabled_callbacks: mappedDisabledCallbacks, - }; - } - - // Add auto-rotation settings as top-level fields - if (autoRotationEnabled) { - formValues.auto_rotate = true; - formValues.rotation_interval = rotationInterval; - } - - // Handle duration field for key expiry - convert empty string to null - if (!formValues.duration || formValues.duration.trim() === "") { - formValues.duration = null; - } - - // Update the formValues with the final metadata - formValues.metadata = JSON.stringify(metadata); - - // disable_global_guardrails is premium-gated server-side; only send it when enabled - // so non-premium key creation isn't blocked by that gate. - if (!formValues.disable_global_guardrails) { - delete formValues.disable_global_guardrails; - } - - // Transform allowed_vector_store_ids and allowed_mcp_servers_and_groups into object_permission format - if (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) { - formValues.object_permission = { - vector_stores: formValues.allowed_vector_store_ids, - }; - // Remove the original field as it's now part of object_permission - delete formValues.allowed_vector_store_ids; - } - - // Transform allowed_mcp_servers_and_groups into object_permission format - if ( - formValues.allowed_mcp_servers_and_groups && - (formValues.allowed_mcp_servers_and_groups.servers?.length > 0 || - formValues.allowed_mcp_servers_and_groups.accessGroups?.length > 0 || - formValues.allowed_mcp_servers_and_groups.toolsets?.length > 0) - ) { - if (!formValues.object_permission) { - formValues.object_permission = {}; - } - const { servers, accessGroups, toolsets } = formValues.allowed_mcp_servers_and_groups; - if (servers && servers.length > 0) { - formValues.object_permission.mcp_servers = servers; - } - if (accessGroups && accessGroups.length > 0) { - formValues.object_permission.mcp_access_groups = accessGroups; - } - if (toolsets && toolsets.length > 0) { - formValues.object_permission.mcp_toolsets = toolsets; - } - // Remove the original field as it's now part of object_permission - delete formValues.allowed_mcp_servers_and_groups; - } - - // Add MCP tool permissions to object_permission - const mcpToolPermissions = formValues.mcp_tool_permissions || {}; - if (Object.keys(mcpToolPermissions).length > 0) { - if (!formValues.object_permission) { - formValues.object_permission = {}; - } - formValues.object_permission.mcp_tool_permissions = mcpToolPermissions; - } - delete formValues.mcp_tool_permissions; - - // Transform allowed_mcp_access_groups into object_permission format - if (formValues.allowed_mcp_access_groups && formValues.allowed_mcp_access_groups.length > 0) { - if (!formValues.object_permission) { - formValues.object_permission = {}; - } - formValues.object_permission.mcp_access_groups = formValues.allowed_mcp_access_groups; - // Remove the original field as it's now part of object_permission - delete formValues.allowed_mcp_access_groups; - } - - // Transform allowed_agents_and_groups into object_permission format - if ( - formValues.allowed_agents_and_groups && - (formValues.allowed_agents_and_groups.agents?.length > 0 || - formValues.allowed_agents_and_groups.accessGroups?.length > 0) - ) { - if (!formValues.object_permission) { - formValues.object_permission = {}; - } - const { agents, accessGroups } = formValues.allowed_agents_and_groups; - if (agents && agents.length > 0) { - formValues.object_permission.agents = agents; - } - if (accessGroups && accessGroups.length > 0) { - formValues.object_permission.agent_access_groups = accessGroups; - } - // Remove the original field as it's now part of object_permission - delete formValues.allowed_agents_and_groups; - } - - // Add model_aliases if any are defined - if (Object.keys(modelAliases).length > 0) { - formValues.aliases = JSON.stringify(modelAliases); - } - - // Add router_settings if any are defined - if (routerSettings?.router_settings) { - // Only include router_settings if it has at least one non-null value - const hasValues = Object.values(routerSettings.router_settings).some( - (value) => value !== null && value !== undefined && value !== "", - ); - if (hasValues) { - formValues.router_settings = routerSettings.router_settings; - } - } - - // Add multi-window budget limits (filter out incomplete entries) - const validWindows = budgetLimits.filter( - (w) => w.budget_duration && w.max_budget !== null && w.max_budget !== undefined, - ); - if (validWindows.length > 0) { - formValues.budget_limits = validWindows; - } - - // Add per-tag rate limits (only when at least one row is configured) - const { tag_rpm_limit } = tagRowsToLimits(tagRateLimits); - if (Object.keys(tag_rpm_limit).length > 0) { - formValues.tag_rpm_limit = tag_rpm_limit; - } - - if (Object.keys(budgetFallbacks).length > 0) { - formValues.budget_fallbacks = budgetFallbacks; - } - - if (formValues.budget_duration === NEVER_RESETS_BUDGET_DURATION) { - formValues.budget_duration = null; - } - - let response; - if (keyOwner === "service_account") { - response = await keyCreateServiceAccountCall(accessToken, formValues); - } else { - response = await keyCreateCall(accessToken, userID, formValues); - } + const response = + endpoint === "service_account" + ? await keyCreateServiceAccountCall(accessToken, payload) + : await keyCreateCall(accessToken, userID, payload); // Add the data to the state in the parent component // Also directly update the keys list in VirtualKeysTable without an API call diff --git a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx index 6be1dad564f..5fc80758cfd 100644 --- a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx +++ b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx @@ -1,6 +1,6 @@ "use client"; -import { Tooltip } from "@/components/atoms/Tooltip"; +import { SimpleTooltip } from "@/components/ui/tooltip"; import type { Team } from "@/components/key_team_helpers/key_list"; import type { Organization } from "@/components/networking"; import { formatNumberWithCommas } from "@/utils/dataUtils"; @@ -53,7 +53,7 @@ interface InheritedBudgetHintProps { export function InheritedBudgetHint({ gates }: InheritedBudgetHintProps) { if (gates.length === 0) return null; return ( - This key has no budget of its own, but its spend still counts toward: diff --git a/ui/litellm-dashboard/src/components/shared/MultiSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/MultiSelect.test.tsx index 908c22b5d67..6fe2f0fe51c 100644 --- a/ui/litellm-dashboard/src/components/shared/MultiSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/MultiSelect.test.tsx @@ -92,4 +92,18 @@ describe("MultiSelect", () => { expect(screen.queryByText('Create "vs-typed"')).not.toBeInTheDocument(); expect(await screen.findByText("No options found")).toBeInTheDocument(); }); + + it("marks a disabled option as disabled and refuses to select it", async () => { + const { onValueChange, input } = renderMultiSelect({ + options: [OPTIONS[0], { ...OPTIONS[1], disabled: true }], + }); + + await openPopup(input); + const disabledOption = screen.getByRole("option", { name: /beta-kb/ }); + expect(disabledOption).toHaveAttribute("aria-disabled", "true"); + + await userEvent.click(disabledOption); + + expect(onValueChange).not.toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/components/shared/MultiSelect.tsx b/ui/litellm-dashboard/src/components/shared/MultiSelect.tsx index 49695b56d44..c60a019db06 100644 --- a/ui/litellm-dashboard/src/components/shared/MultiSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/MultiSelect.tsx @@ -18,6 +18,7 @@ export interface MultiSelectOption { label: string; value: string; description?: string; + disabled?: boolean; } interface MultiSelectProps { @@ -116,7 +117,7 @@ export function MultiSelect({ {emptyText} {(option: MultiSelectOption) => ( - + {option.label} {option.description && ( diff --git a/ui/litellm-dashboard/src/components/team/LoggingSettings.test.tsx b/ui/litellm-dashboard/src/components/team/LoggingSettings.test.tsx index 17f20ad0dcc..ae02a26ddb0 100644 --- a/ui/litellm-dashboard/src/components/team/LoggingSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/team/LoggingSettings.test.tsx @@ -1,9 +1,19 @@ +import { readFileSync } from "fs"; +import { resolve } from "path"; import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import userEvent from "@testing-library/user-event"; import { renderWithProviders, screen, fireEvent } from "../../../tests/test-utils"; import LoggingSettings from "./LoggingSettings"; +const SOURCE_PATH = resolve(process.cwd(), "src/components/team/LoggingSettings.tsx"); + +const HARDCODED_PALETTE = + /\b(?:text|bg|border|hover:bg|hover:text|hover:border|dark:bg|dark:text|dark:border|ring|divide|fill|stroke)-(?:gray|slate|zinc|neutral|stone|red|blue|green|yellow|amber|orange|indigo|purple|pink|rose|teal|cyan|sky|violet|fuchsia|lime|emerald)-\d+(?:\/\d+)?\b/g; + +const SEMANTIC_TOKEN = + /\b(?:text|bg|border|hover:bg|hover:text|ring|divide|fill|stroke)-(?:foreground|muted-foreground|muted|background|card|popover|primary|secondary|destructive|border|input|accent|ring)(?:-foreground)?(?:\/\d+)?\b/g; + describe("LoggingSettings", () => { beforeEach(() => { vi.clearAllMocks(); @@ -163,6 +173,33 @@ describe("LoggingSettings", () => { expect(screen.getByText("C")).toBeInTheDocument(); }); + it("styles itself from semantic tokens instead of hardcoded palette classes", () => { + const source = readFileSync(SOURCE_PATH, "utf8"); + + expect(source).toContain("const LoggingSettings"); + expect(source.match(SEMANTIC_TOKEN) ?? []).not.toHaveLength(0); + expect(source.match(HARDCODED_PALETTE) ?? []).toHaveLength(0); + }); + + it("reports the chosen event type when a different option is picked", async () => { + const user = userEvent.setup({ delay: null }); + const mockOnChange = vi.fn(); + const initialValue = [ + { + callback_name: "langsmith", + callback_type: "success", + callback_vars: {}, + }, + ]; + + renderWithProviders(); + + await user.click(screen.getByTitle("Success Only")); + await user.click(await screen.findByTitle("Failure Only")); + + expect(mockOnChange).toHaveBeenCalledWith([expect.objectContaining({ callback_type: "failure" })]); + }); + it("correctly handles numerical input with decimal values", () => { const mockOnChange = vi.fn(); diff --git a/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx b/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx index eee6e76569d..159f8fa0b74 100644 --- a/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx +++ b/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx @@ -1,8 +1,10 @@ /* eslint-disable react/no-unescaped-entities */ import React from "react"; -import { Select, Tooltip, Divider } from "antd"; +import { SimpleTooltip } from "@/components/ui/tooltip"; +import { Select, Divider } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button } from "@/components/ui/button"; +import { Badge } from "@/components/ui/badge"; import { Card } from "@/components/ui/card"; import { Input } from "@/components/ui/input"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; @@ -142,30 +144,24 @@ const LoggingSettings: React.FC = ({ if (Object.keys(dynamicParams).length === 0) return null; return ( -
+
-
-
+
+
- Integration Parameters + Integration Parameters
{Object.entries(dynamicParams).map(([paramName, paramType]) => (
-