diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index c09369b9b2c..ba131ecac4c 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 22344 + "limit": 22343 }, "reportArgumentType": { "limit": 2578 diff --git a/helm/litellm-helm/templates/migrations-job.yaml b/helm/litellm-helm/templates/migrations-job.yaml index f8a660e23f8..5a873cbb965 100644 --- a/helm/litellm-helm/templates/migrations-job.yaml +++ b/helm/litellm-helm/templates/migrations-job.yaml @@ -119,4 +119,7 @@ spec: {{- end }} ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }} backoffLimit: {{ .Values.migrationJob.backoffLimit }} + {{- with .Values.migrationJob.activeDeadlineSeconds }} + activeDeadlineSeconds: {{ . }} + {{- end }} {{- end }} diff --git a/helm/litellm-helm/tests/migrations-job_tests.yaml b/helm/litellm-helm/tests/migrations-job_tests.yaml index cb962118a25..e327a3ec201 100644 --- a/helm/litellm-helm/tests/migrations-job_tests.yaml +++ b/helm/litellm-helm/tests/migrations-job_tests.yaml @@ -314,3 +314,31 @@ tests: operator: Equal value: litellm-e2e effect: NoSchedule + + - it: bounds the Job with a deadline by default, so a blocked migration cannot stall the release forever + set: + migrationJob: + enabled: true + asserts: + - equal: + path: spec.activeDeadlineSeconds + value: 1800 + + - it: honours an operator-supplied deadline + set: + migrationJob: + enabled: true + activeDeadlineSeconds: 600 + asserts: + - equal: + path: spec.activeDeadlineSeconds + value: 600 + + - it: omits the deadline entirely when it is nulled out, restoring the unbounded behaviour + set: + migrationJob: + enabled: true + activeDeadlineSeconds: null + asserts: + - notExists: + path: spec.activeDeadlineSeconds diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index df2b55723fe..628ca038339 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -427,6 +427,13 @@ migrationJob: enabled: true # Enable or disable the schema migration Job retries: 3 # Number of retries for the Job in case of failure backoffLimit: 4 # Backoff limit for Job restarts + # Wall-clock budget for the whole Job, shared across every `backoffLimit` + # retry rather than granted per attempt. Without it a migration that blocks + # on the database never fails, and when the Helm hook is enabled the release + # waits on it forever: `helm upgrade` and any GitOps controller driving it + # stop reconciling the whole chart until someone deletes the Job by hand. + # Set to null to opt out and restore the unbounded behaviour. + activeDeadlineSeconds: 1800 disableSchemaUpdate: false # Skip schema migrations for specific environments. When True, the job will exit with code 0. # Optional service account for the migration job. # Only used when migrationJob.hooks.helm.enabled=true and serviceAccount.create=true. diff --git a/helm/litellm/templates/migrations-job.yaml b/helm/litellm/templates/migrations-job.yaml index 2debe8a1e10..9cd8397f794 100644 --- a/helm/litellm/templates/migrations-job.yaml +++ b/helm/litellm/templates/migrations-job.yaml @@ -21,6 +21,9 @@ metadata: spec: backoffLimit: {{ .Values.migrationJob.backoffLimit }} ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }} + {{- with .Values.migrationJob.activeDeadlineSeconds }} + activeDeadlineSeconds: {{ . }} + {{- end }} template: metadata: {{- /* The Job's selector is generated by the controller rather than diff --git a/helm/litellm/tests/migration_job_tests.yaml b/helm/litellm/tests/migration_job_tests.yaml index 12e525c5a8c..c3f3083ece5 100644 --- a/helm/litellm/tests/migration_job_tests.yaml +++ b/helm/litellm/tests/migration_job_tests.yaml @@ -167,3 +167,24 @@ tests: - equal: path: spec.template.metadata.labels['app.kubernetes.io/component'] value: batch-migrations + + - it: bounds the Job with a deadline by default, so a blocked migration cannot stall the release forever + asserts: + - equal: + path: spec.activeDeadlineSeconds + value: 1800 + + - it: honours an operator-supplied deadline + set: + migrationJob.activeDeadlineSeconds: 600 + asserts: + - equal: + path: spec.activeDeadlineSeconds + value: 600 + + - it: omits the deadline entirely when it is nulled out, restoring the unbounded behaviour + set: + migrationJob.activeDeadlineSeconds: null + asserts: + - notExists: + path: spec.activeDeadlineSeconds diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 7820a898ef1..3f8aacfce17 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -56,6 +56,15 @@ migrationJob: enabled: true backoffLimit: 4 ttlSecondsAfterFinished: 120 + # Wall-clock budget for the whole Job, shared across every `backoffLimit` + # retry rather than granted per attempt. Without it a migration that blocks + # on the database never fails, and because this is a pre-upgrade hook the + # release waits on it forever: `helm upgrade` and any GitOps controller + # driving it stop reconciling the whole chart until someone deletes the Job + # by hand. A migration that has exhausted its retries is not going to + # succeed on the next one, so failing is strictly better than hanging. + # Set to null to opt out and restore the unbounded behaviour. + activeDeadlineSeconds: 1800 resources: {} # ServiceAccount for the Job pod only. # diff --git a/litellm/constants.py b/litellm/constants.py index 8f236eba327..75aacd2e6f1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1593,6 +1593,13 @@ DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL: Final = 10 # in a single ``/{name1,name2,...}/mcp`` URL. Bounds the per-request DB / cache # fan-out an authenticated caller can trigger by stuffing the path with tokens. DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: Final = 16 +# Ceilings on the cached auth registries; larger tables fall back to per-row lookups +# instead of holding an unbounded id set in every worker. +TAG_REGISTRY_MAX_SIZE: Final = 5000 +END_USER_RESTRICTED_REGISTRY_MAX_SIZE: Final = 5000 +# How long a failed registry load is remembered as "unusable", so a degraded Postgres +# is not re-scanned on every request on top of the per-id lookups it falls back to. +REGISTRY_ERROR_NEGATIVE_CACHE_TTL: Final = 30 # Sentry Scrubbing Configuration SENTRY_DENYLIST: Final = [ diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index f691a1a4454..e45414b4a73 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -734,6 +734,7 @@ class LiteLLMAnthropicMessagesAdapter: "input_schema", "description", "cache_control", + "strict", "type", ] @@ -763,6 +764,8 @@ class LiteLLMAnthropicMessagesAdapter: function_chunk["parameters"] = tool["input_schema"] if "description" in tool: function_chunk["description"] = tool["description"] + if "strict" in tool: + function_chunk["strict"] = bool(tool["strict"]) for k, v in tool.items(): if k not in mapped_tool_params: # pass additional computer kwargs diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index f12dd979338..e2ad9c9c6d3 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -3,7 +3,7 @@ import json import traceback from collections import deque -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from typing import Any, Final from litellm import verbose_logger @@ -68,6 +68,19 @@ class AnthropicResponsesStreamWrapper: self._current_block_index += 1 return self._current_block_index + def _open_block(self, item_id: str | None, content_block: Mapping[str, Any]) -> int: + block_idx = self._next_block_index() + if item_id: + self._item_id_to_block_index[item_id] = block_idx + self._chunk_queue.append( + { + "type": "content_block_start", + "index": block_idx, + "content_block": content_block, + } + ) + return block_idx + def _process_event(self, event: Any) -> None: """Convert one Responses API event into zero or more Anthropic chunks queued for emission.""" event_type = getattr(event, "type", None) @@ -93,47 +106,22 @@ class AnthropicResponsesStreamWrapper: item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item_type == "message": - block_idx = self._next_block_index() - if item_id: - self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append( - { - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "text", "text": ""}, - } - ) + self._open_block(item_id, {"type": "text", "text": ""}) elif item_type == "function_call": call_id: Final = ( getattr(item, "call_id", None) or (item.get("call_id") if isinstance(item, dict) else None) or "" ) name = getattr(item, "name", None) or (item.get("name") if isinstance(item, dict) else None) or "" - block_idx = self._next_block_index() if item_id: - self._item_id_to_block_index[item_id] = block_idx self._pending_tool_ids[item_id] = call_id - self._chunk_queue.append( + self._open_block( + item_id, { - "type": "content_block_start", - "index": block_idx, - "content_block": { - "type": "tool_use", - "id": call_id, - "name": name, - "input": {}, - }, - } - ) - elif item_type == "reasoning": - block_idx = self._next_block_index() - if item_id: - self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append( - { - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "thinking", "thinking": ""}, - } + "type": "tool_use", + "id": call_id, + "name": name, + "input": {}, + }, ) return @@ -146,16 +134,7 @@ class AnthropicResponsesStreamWrapper: # Some providers (e.g. LMStudio) skip response.output_item.added, # so no text block is open yet; synthesize content_block_start # instead of emitting a delta with index -1 - block_idx = self._next_block_index() - if item_id: - self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append( - { - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "text", "text": ""}, - } - ) + block_idx = self._open_block(item_id, {"type": "text", "text": ""}) self._chunk_queue.append( { "type": "content_block_delta", @@ -169,11 +148,11 @@ class AnthropicResponsesStreamWrapper: if event_type == "response.reasoning_summary_text.delta": item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") - block_idx = ( - self._item_id_to_block_index.get(item_id, self._current_block_index) - if item_id - else self._current_block_index - ) + block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index + if block_idx < 0: + if not delta: + return + block_idx = self._open_block(item_id, {"type": "thinking", "thinking": ""}) self._chunk_queue.append( { "type": "content_block_delta", @@ -207,11 +186,9 @@ class AnthropicResponsesStreamWrapper: item_id = ( getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None ) - block_idx = ( - self._item_id_to_block_index.get(item_id, self._current_block_index) - if item_id - else self._current_block_index - ) + block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index + if block_idx < 0: + return self._chunk_queue.append( { "type": "content_block_stop", diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index be4cef4dfe0..21a8cb9501e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -266,7 +266,13 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if (isinstance(tool_type, str) and tool_type.startswith("web_search")) or tool_name == "web_search": result.append({"type": "web_search_preview"}) continue - func_tool: dict[str, Any] = {"type": "function", "name": tool_name} + # Responses turns strict mode on when `strict` is omitted, silently rewriting + # `required` to every property. Anthropic tools are non-strict unless asked. + func_tool: dict[str, Any] = { + "type": "function", + "name": tool_name, + "strict": bool(tool_dict.get("strict")), + } if "description" in tool_dict: func_tool["description"] = tool_dict["description"] if "input_schema" in tool_dict: diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 86e97b55a8e..3139b679e5c 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2416,7 +2416,7 @@ def _build_oauth_authorization_server_response( _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth authorization server") return { - "issuer": request_base_url, # point to your proxy + "issuer": f"{request_base_url}/{mcp_server_name}" if mcp_server_name else request_base_url, "authorization_endpoint": authorization_endpoint, "token_endpoint": token_endpoint, "response_types_supported": ["code"], @@ -2464,7 +2464,14 @@ async def oauth_authorization_server_mcp(request: Request, mcp_server_name: str # Alias for standard OpenID discovery @router.get("/.well-known/openid-configuration") async def openid_configuration(request: Request): - response = await oauth_authorization_server_mcp(request) + response: Final = await oauth_authorization_server_mcp(request) + if not isinstance(response, dict): + return response + + request_base_url: Final = get_request_base_url(request) + # OIDC verifiers derive this URL from their configured issuer (the proxy origin), + # so keep the origin issuer here even when root resolution scoped the metadata. + unscoped_response: Final = {**response, "issuer": request_base_url} # If MCPJWTSigner is active, augment the discovery doc with JWKS fields so # MCP servers and gateways (e.g. AWS Bedrock AgentCore Gateway) can resolve @@ -2476,17 +2483,15 @@ async def openid_configuration(request: Request): signer: Final = get_mcp_jwt_signer() if signer is not None: - request_base_url: Final = get_request_base_url(request) - if isinstance(response, dict): - response = { - **response, - "jwks_uri": f"{request_base_url}/.well-known/jwks.json", - "id_token_signing_alg_values_supported": ["RS256"], - } + return { + **unscoped_response, + "jwks_uri": f"{request_base_url}/.well-known/jwks.json", + "id_token_signing_alg_values_supported": ["RS256"], + } except ImportError: pass - return response + return unscoped_response @router.get("/.well-known/jwks.json") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c782f0dfa09..97b03b5e60c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,7 +13,7 @@ import json import os import re import time -from collections.abc import AsyncIterator, Callable, Sequence +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextlib import asynccontextmanager from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast from urllib.parse import ParseResult, urlparse @@ -46,6 +46,9 @@ from litellm.constants import ( ) from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth +from litellm.integrations.custom_guardrail import ( + _sync_guardrail_info_to_logging_obj, # pyright: ignore[reportPrivateUsage] - the same bridge @log_guardrail_information uses; reimplementing it here would fork the metadata-key logic +) from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( @@ -162,6 +165,7 @@ if TYPE_CHECKING: from mcp.types import CreateMessageRequestParams from litellm.caching.caching import InMemoryCache + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.mcp_server.mcp_toolset import MCPToolset try: @@ -1233,6 +1237,35 @@ def _create_elicitation_callback(): return _elicitation_callback +def _record_mcp_guardrail_evaluations( + synthetic_llm_data: dict[str, Any], # mutable-ok: `_sync_guardrail_info_to_logging_obj` takes a concrete dict + litellm_logging_obj: "LiteLLMLoggingObj | None", +) -> None: + """Bridge guardrail decision records off an MCP synthetic request onto the request's logger. + + MCP guardrails run against a throwaway LLM-shaped dict from + ``ProxyLogging._convert_mcp_to_llm_format``, so ``@log_guardrail_information`` + files ``standard_logging_guardrail_information`` in that dict's metadata bucket, + which ``get_standard_logging_object_payload`` never reads. Native (non-unified) + guardrails receive no ``logging_obj`` kwarg, so the decorator cannot bridge on + their behalf; this calls the same helper it would have. + + Only the decision records move. The synthetic request's messages and tool + arguments stay behind: they can carry end-user data, and the monitor needs none + of it. + """ + if litellm_logging_obj is None: + return + + try: + _sync_guardrail_info_to_logging_obj(synthetic_llm_data, litellm_logging_obj) + except Exception as e: # noqa: BLE001 # callers run this from a `finally` on the block path + # The breadth is the point. Narrowing to the knowable AttributeError/TypeError + # would let an unexpected type escape that ``finally`` and replace the guardrail's + # block with a bookkeeping error. + verbose_logger.warning("Failed to record MCP guardrail evaluation for logging: %s", e) + + class MCPServerManager: _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") @@ -4573,6 +4606,7 @@ class MCPServerManager: proxy_logging_obj: ProxyLogging | None, server: MCPServer, raw_headers: dict[str, str] | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> dict[str, Any]: """ Run pre-call checks and guardrail hooks for an MCP tool call. @@ -4582,6 +4616,10 @@ class MCPServerManager: present. An absent logger must never be able to turn an authorization decision into a no-op. + ``litellm_logging_obj`` is the request's logger, and it is what lands a + ``pre_mcp_call`` evaluation (or a block) on the spend-log row the Guardrails + Monitor counts. It stays optional so callers that do no logging are unchanged. + Returns a dict that may contain: - "arguments": hook-modified tool arguments (only if changed) - "extra_headers": headers injected by pre_mcp_call guardrail hooks @@ -4640,8 +4678,13 @@ class MCPServerManager: # Create MCP request object for processing mcp_request_obj: Final = proxy_logging_obj._create_mcp_request_object_from_kwargs(pre_hook_kwargs) - # Convert to LLM format for existing guardrail compatibility + # Convert to LLM format for existing guardrail compatibility. + # Unified guardrails read the seeded logger off the request dict and pass it + # into ``apply_guardrail``, so ``@log_guardrail_information`` bridges their + # evaluations itself; the ``finally`` below covers native guardrails, which + # never receive it. Same seeding the pass-through routes do. synthetic_llm_data: Final = proxy_logging_obj._convert_mcp_to_llm_format(mcp_request_obj, pre_hook_kwargs) + synthetic_llm_data["litellm_logging_obj"] = litellm_logging_obj try: # Use standard pre_call_hook @@ -4666,6 +4709,12 @@ class MCPServerManager: # Re-raise guardrail exceptions to properly fail the MCP call verbose_logger.error("Guardrail blocked MCP tool call pre call: %s", e) raise e + finally: + # ``finally`` rather than after the ``try``: a block raises straight out of + # here, and the failure spend-log row that "Total Blocked" counts is built + # from this logger further up the stack, so the record has to be attached + # before the exception leaves this frame. + _record_mcp_guardrail_evaluations(synthetic_llm_data, litellm_logging_obj) return hook_result @@ -4677,8 +4726,14 @@ class MCPServerManager: user_api_key_auth: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging, start_time: datetime.datetime, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ): - """Create and return a during hook task for MCP tool calls.""" + """Create and return a during hook task for MCP tool calls. + + ``litellm_logging_obj`` is the request's logger; see ``pre_call_tool_check``. + The task is awaited before the tool call's success logging runs, so a + ``during_mcp_call`` evaluation recorded on it is serialized with that call. + """ from litellm.types.llms.base import HiddenParams from litellm.types.mcp import MCPDuringCallRequestObject @@ -4697,15 +4752,23 @@ class MCPServerManager: "user_api_key_auth": user_api_key_auth, } + # Seeded for the same reason as in ``pre_call_tool_check``. synthetic_llm_data: Final = proxy_logging_obj._convert_mcp_to_llm_format(request_obj, during_hook_kwargs) + synthetic_llm_data["litellm_logging_obj"] = litellm_logging_obj - return asyncio.create_task( - proxy_logging_obj.during_call_hook( - user_api_key_dict=user_api_key_auth, - data=synthetic_llm_data, - call_type=CallTypes.call_mcp_tool.value, - ) - ) + # Wrapped so the bridge runs inside the task: the caller only holds the task and + # gathers it later, so there is no other point that still sees a block here. + async def _run_during_call_hook() -> Mapping[str, Any] | None: + try: + return await proxy_logging_obj.during_call_hook( + user_api_key_dict=user_api_key_auth, + data=synthetic_llm_data, + call_type=CallTypes.call_mcp_tool.value, + ) + finally: + _record_mcp_guardrail_evaluations(synthetic_llm_data, litellm_logging_obj) + + return asyncio.create_task(_run_during_call_hook()) def _get_call_semaphore(self, mcp_server: MCPServer) -> asyncio.Semaphore | None: limit: Final = mcp_server.max_concurrent_requests @@ -5234,6 +5297,7 @@ class MCPServerManager: oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, host_progress_callback: Callable | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> CallToolResult: """ Call a tool with the given name and arguments @@ -5246,6 +5310,9 @@ class MCPServerManager: mcp_auth_header: MCP auth header (deprecated) mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} proxy_logging_obj: Optional ProxyLogging object for hook integration + litellm_logging_obj: Optional request logger the guardrail hooks record + their evaluations onto, so MCP guardrail activity reaches the + Guardrails Monitor. See ``pre_call_tool_check`` Returns: @@ -5276,6 +5343,7 @@ class MCPServerManager: proxy_logging_obj=proxy_logging_obj, server=mcp_server, raw_headers=raw_headers, + litellm_logging_obj=litellm_logging_obj, ) if "arguments" in hook_result: arguments = hook_result["arguments"] @@ -5290,6 +5358,7 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, proxy_logging_obj=proxy_logging_obj, start_time=start_time, + litellm_logging_obj=litellm_logging_obj, ) tasks.append(during_hook_task) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f237529b319..17457c3362f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2824,6 +2824,7 @@ if MCP_AVAILABLE: proxy_logging_obj=proxy_logging_obj, server=mcp_server, raw_headers=raw_headers, + litellm_logging_obj=litellm_logging_obj, ) # `pre_call_tool_check` may return guardrail-modified # arguments; honor them on the local path too. @@ -2962,6 +2963,7 @@ if MCP_AVAILABLE: proxy_logging_obj=proxy_logging_obj, server=prefix_server, raw_headers=raw_headers, + litellm_logging_obj=litellm_logging_obj, ) if "arguments" in hook_result: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args @@ -3149,6 +3151,20 @@ if MCP_AVAILABLE: traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) from litellm.proxy.proxy_server import proxy_logging_obj + # Ordering is load-bearing. ``_ProxyDBLogger.async_post_call_failure_hook``, + # reached below, writes the failure spend-log row from this logger's + # ``standard_logging_object``, which only exists once the failure handlers + # have run. Flush them first or the row lands with + # ``guardrail_information=None`` and a guardrail block is never counted. + # + # Not double-logged: both handlers gate on ``should_run_logging`` and then + # mark it, so the ``@client`` wrapper's own post-raise logging no-ops on this + # logger, same as ``_fire_mcp_tool_call_logging`` does for ``isError=True``. + if litellm_logging_obj is not None: + end_time: Final = datetime.now() # noqa: DTZ005 # naive to match `start_time`, which it is subtracted from + litellm_logging_obj.failure_handler(e, traceback_str, start_time, end_time) + await litellm_logging_obj.async_failure_handler(e, traceback_str, start_time, end_time) + if proxy_logging_obj and user_api_key_auth: await proxy_logging_obj.post_call_failure_hook( request_data=kwargs, @@ -3326,6 +3342,7 @@ if MCP_AVAILABLE: raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, + litellm_logging_obj=litellm_logging_obj, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) return call_tool_result diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a566d491597..22cc961a7d2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -680,6 +680,7 @@ class LiteLLMRoutes(enum.Enum): # permitted teams exactly like /spend/logs/ui — it belongs to the same # access tier, not to customer management. "/management/v1/spend_logs/end_users", + "/management/v1/spend_logs/users", "/cost/estimate", ] @@ -872,12 +873,13 @@ class LiteLLMRoutes(enum.Enum): # PROXY_ADMIN_VIEW_ONLY — the route gate must match). "/customer/list", "/customer/info", - # UI Logs page detail drawer (single + session) and the end-user filter - # facet. The list endpoint `/spend/logs/ui` is covered via + # UI Logs page detail drawer (single + session) and the filter facets. + # The list endpoint `/spend/logs/ui` is covered via # spend_tracking_routes below. "/spend/logs/ui/{logId}", "/spend/logs/session/ui", "/management/v1/spend_logs/end_users", + "/management/v1/spend_logs/users", # Settings / observability read endpoints exposed in admin-only # sidebar groups (Logging & Alerts, Admin Settings, Budgets, # Invitations). diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index f75899b91dc..a48ef0f08bb 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -193,6 +193,9 @@ async def anthropic_response( ) verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e) + if isinstance(e, ProxyException): + raise + # Extract model_id from request metadata (same as success path) litellm_metadata: Final = data.get("litellm_metadata", {}) or {} model_info: Final = litellm_metadata.get("model_info", {}) or {} diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3d8fed18423..8708f96339f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,7 +13,7 @@ import asyncio import math import re import time -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast @@ -30,6 +30,9 @@ from litellm.constants import ( DEFAULT_IN_MEMORY_TTL, DEFAULT_MAX_RECURSE_DEPTH, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, + END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + REGISTRY_ERROR_NEGATIVE_CACHE_TTL, + TAG_REGISTRY_MAX_SIZE, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -74,9 +77,15 @@ from litellm.proxy.common_utils.http_parsing_utils import ( ) from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import ( + END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, + TAG_REGISTRY_OVERFLOW_SENTINEL, UserApiKeyCache, + end_user_cache_key, + end_user_restricted_registry_cache_key, get_management_object_ttl, object_permission_cache_key, + tag_cache_key, + tag_registry_cache_key, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( @@ -163,7 +172,7 @@ class _PrismaAuthTable(Protocol[RowT_co]): async def find_many( self, *, - where: Mapping[str, object], + where: Mapping[str, object] | None = None, include: Mapping[str, object] | None = None, take: int | None = None, ) -> Sequence[RowT_co]: ... @@ -220,6 +229,16 @@ def _tag_table(repo: _PrismaTableHolder[_PrismaTagRow]) -> _PrismaAuthTable[_Pri return repo.table +class _PrismaEndUserRow(Protocol): + user_id: str + + def dict(self) -> Mapping[str, object]: ... + + +def _end_user_table(repo: _PrismaTableHolder[_PrismaEndUserRow]) -> _PrismaAuthTable[_PrismaEndUserRow]: + return repo.table + + class _RawCacheRead(Protocol): async def async_get_cache(self, *, key: str) -> object: ... @@ -1284,6 +1303,191 @@ async def _check_end_user_budget( ) +#: Columns whose non-null value makes an end-user row restrict something auth enforces. ``blocked`` +#: is separate: it restricts when true rather than when merely set. +_RESTRICTED_COLUMNS: Final = ("budget_id", "allowed_model_region", "default_model", "object_permission_id") + + +def _column_is_set(column: str) -> Mapping[str, object]: + """``column IS NOT NULL`` as a plain dict, which is the only shape prisma's builder accepts.""" + return {column: {"not": None}} # mutable-ok: prisma's query builder isinstance-checks for dict + + +def _restricted_end_user_where() -> Mapping[str, object]: + """Prisma filter selecting every end-user row that carries a restriction auth enforces.""" + return {"OR": [{"blocked": True}, *map(_column_is_set, _RESTRICTED_COLUMNS)]} # mutable-ok: prisma needs dict/list + + +class _RegistryNotCached: + """No cached registry answer, as distinct from the cached answer ``None`` (registry unusable).""" + + +_REGISTRY_NOT_CACHED: Final = _RegistryNotCached() + +#: One lock per registry; module-level because the stampede to collapse is worker-wide. +_TAG_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() +_END_USER_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() + + +async def _cached_registry( + cache_key: str, + overflow_sentinel: str, + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None | _RegistryNotCached: + """The cached registry answer, or ``_REGISTRY_NOT_CACHED`` when the caller has to query.""" + cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key) + if cached == overflow_sentinel: + return None + # Memory hands back the tuple that was written; Redis round-trips it through JSON as a list. + if isinstance(cached, (list, tuple)): + return frozenset(entry for entry in cached if isinstance(entry, str)) + return _REGISTRY_NOT_CACHED + + +async def _cache_registry_answer( + cache_key: str, + value: tuple[str, ...] | str, + ttl: float, + user_api_key_cache: UserApiKeyCache, +) -> None: + """Best-effort: a cache backend failure must not turn a registry load into a failed request.""" + try: + await user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl) + except Exception as e: # noqa: BLE001 # best-effort cache write: auth must survive a cache backend error + verbose_proxy_logger.warning("Failed to cache registry %s: %s", cache_key, e) + + +async def _fetch_and_cache_registry( + cache_key: str, + overflow_sentinel: str, + max_size: int, + fetch_ids: Callable[[], Awaitable[tuple[str, ...]]], + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """The registry as the database has it, cached whole, or ``None`` when it is unusable.""" + try: + registry_ids: Final = await fetch_ids() + except Exception as e: # noqa: BLE001 # fail-safe: any registry load error must degrade to per-id lookups, never break auth + verbose_proxy_logger.warning( + "Registry %s could not be loaded from the database, so per-id lookups will run and the " + "registry query is suppressed for %ss: %s", + cache_key, + REGISTRY_ERROR_NEGATIVE_CACHE_TTL, + e, + ) + await _cache_registry_answer( + cache_key=cache_key, + value=overflow_sentinel, + ttl=REGISTRY_ERROR_NEGATIVE_CACHE_TTL, + user_api_key_cache=user_api_key_cache, + ) + return None + + if len(registry_ids) > max_size: + await _cache_registry_answer( + cache_key=cache_key, + value=overflow_sentinel, + ttl=get_management_object_ttl(user_api_key_cache), + user_api_key_cache=user_api_key_cache, + ) + return None + + await _cache_registry_answer( + cache_key=cache_key, + value=registry_ids, + ttl=get_management_object_ttl(user_api_key_cache), + user_api_key_cache=user_api_key_cache, + ) + return frozenset(registry_ids) + + +async def _load_bounded_registry( + cache_key: str, + overflow_sentinel: str, + max_size: int, + load_lock: asyncio.Lock, + fetch_ids: Callable[[], Awaitable[tuple[str, ...]]], + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """ + A bounded id set under one cache key, so an id outside it costs no DB read. + + ``None`` = unusable (overflow or recent DB error): fall back to per-id lookups. An empty + frozenset is a real, cacheable answer. Loads are single-flighted to stop TTL-expiry stampedes. + """ + cached: Final = await _cached_registry(cache_key, overflow_sentinel, user_api_key_cache) + if not isinstance(cached, _RegistryNotCached): + return cached + + async with load_lock: + # The request that held the lock has since cached an answer for everyone waiting on it. + cached_after_wait: Final = await _cached_registry(cache_key, overflow_sentinel, user_api_key_cache) + if not isinstance(cached_after_wait, _RegistryNotCached): + return cached_after_wait + + return await _fetch_and_cache_registry( + cache_key=cache_key, + overflow_sentinel=overflow_sentinel, + max_size=max_size, + fetch_ids=fetch_ids, + user_api_key_cache=user_api_key_cache, + ) + + +async def _load_end_user_restricted_registry( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """The set of end-user ids whose ``LiteLLM_EndUserTable`` row carries a restriction.""" + + async def fetch_ids() -> tuple[str, ...]: + restricted_rows: Final = await _end_user_table(EndUserRepository(prisma_client)).find_many( + where=_restricted_end_user_where(), + take=END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1, + ) + return tuple(row.user_id for row in restricted_rows) + + return await _load_bounded_registry( + cache_key=end_user_restricted_registry_cache_key(), + overflow_sentinel=END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, + max_size=END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + load_lock=_END_USER_REGISTRY_LOAD_LOCK, + fetch_ids=fetch_ids, + user_api_key_cache=user_api_key_cache, + ) + + +async def _end_user_is_known_unrestricted( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + token_end_user_max_budget: float | None, +) -> bool: + """ + True when the cached registry proves the id restricts nothing, so its row need not be read. + + Every field ``get_end_user_object`` callers consume (budget, spend under that budget, region, + default model, object permission, blocked) is part of the registry predicate, so an id outside + it is indistinguishable from one with no row at all. The skip is off whenever mere existence of + the row is meaningful: ``max_end_user_budget_id`` grafts a default budget onto any row that + exists, ``validate_end_user_id_in_db`` rejects ids that resolve to no row, and a token-supplied + ``end_user_max_budget`` (a ``user_custom_auth`` callable can set one against an otherwise + unrestricted row) is enforced against the row's recorded spend. + """ + if ( + litellm.max_end_user_budget_id is not None + or litellm.validate_end_user_id_in_db + or token_end_user_max_budget is not None + ): + return False + + registry: Final = await _load_end_user_restricted_registry( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return registry is not None and end_user_id not in registry + + @log_db_metrics async def get_end_user_object( end_user_id: str | None, @@ -1292,6 +1496,7 @@ async def get_end_user_object( route: str | None = "", parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, + token_end_user_max_budget: float | None = None, ) -> LiteLLM_EndUserTable | None: """ Returns end user object from database or cache. @@ -1306,6 +1511,9 @@ async def get_end_user_object( route: The request route parent_otel_span: Optional OpenTelemetry span for tracing proxy_logging_obj: Optional proxy logging object + token_end_user_max_budget: ``valid_token.end_user_max_budget``, when the caller holds a + token. Budget enforcement reads the row's spend, so a row that restricts nothing on + its own must still be loaded when the token carries a budget for it. Returns: LiteLLM_EndUserTable if found, None otherwise @@ -1316,7 +1524,7 @@ async def get_end_user_object( if end_user_id is None: return None - _key: Final = f"end_user_id:{end_user_id}" + _key: Final = end_user_cache_key(end_user_id) # Check cache first cached_user_obj: Final = await user_api_key_cache.async_get_cache( @@ -1335,6 +1543,14 @@ async def get_end_user_object( return return_obj + if await _end_user_is_known_unrestricted( + end_user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + token_end_user_max_budget=token_end_user_max_budget, + ): + return None + # Fetch from database try: response: Final = await _dictable_table(EndUserRepository(prisma_client)).find_unique( @@ -1358,9 +1574,10 @@ async def get_end_user_object( # Save to cache await user_api_key_cache.async_set_cache( - key=f"end_user_id:{end_user_id}", + key=_key, value=_response, model_type=LiteLLM_EndUserTable, + ttl=get_management_object_ttl(user_api_key_cache), ) return _response @@ -1480,6 +1697,67 @@ async def _end_user_id_exists_in_db( return False +async def _load_tag_registry( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """The set of tag names that have a row in ``LiteLLM_TagTable``.""" + + async def fetch_ids() -> tuple[str, ...]: + registry_rows: Final = await _tag_table(TagRepository(prisma_client)).find_many( + take=TAG_REGISTRY_MAX_SIZE + 1, + ) + return tuple(row.tag_name for row in registry_rows) + + return await _load_bounded_registry( + cache_key=tag_registry_cache_key(), + overflow_sentinel=TAG_REGISTRY_OVERFLOW_SENTINEL, + max_size=TAG_REGISTRY_MAX_SIZE, + load_lock=_TAG_REGISTRY_LOAD_LOCK, + fetch_ids=fetch_ids, + user_api_key_cache=user_api_key_cache, + ) + + +async def _fetch_uncached_tags( + uncached_tags: Sequence[str], + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> tuple[tuple[str, LiteLLM_TagTable], ...]: + """Rows for the tags a cache probe missed; names absent from the registry never reach the DB.""" + if not uncached_tags: + return () + + registry: Final = await _load_tag_registry( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + tags_to_fetch: Final = ( + tuple(uncached_tags) if registry is None else tuple(tag for tag in uncached_tags if tag in registry) + ) + if not tags_to_fetch: + return () + + try: + db_tags: Final = await _tag_table(TagRepository(prisma_client)).find_many( + where={"tag_name": {"in": list(tags_to_fetch)}}, + include={"litellm_budget_table": True}, + ) + fetched: Final = tuple((db_tag.tag_name, LiteLLM_TagTable.model_validate(db_tag.dict())) for db_tag in db_tags) + for fetched_name, fetched_obj in fetched: + await user_api_key_cache.async_set_cache( + key=tag_cache_key(fetched_name), + value=fetched_obj, + model_type=LiteLLM_TagTable, + ttl=get_management_object_ttl(user_api_key_cache), + ) + except Exception as e: # noqa: BLE001 # fail-safe: a tag fetch error must yield "no budget objects", never break auth + verbose_proxy_logger.debug("Error batch fetching tags from database: %s", e) + return () + else: + return fetched + + @log_db_metrics async def get_tag_objects_batch( tag_names: list[str], @@ -1492,8 +1770,9 @@ async def get_tag_objects_batch( Batch fetch multiple tag objects from cache and db. Optimizes for latency by: - 1. Fetching all cached tags in parallel - 2. Batch fetching uncached tags in one DB query + 1. Serving already-cached tags without touching the DB + 2. Skipping tags that no ``LiteLLM_TagTable`` row exists for, via the cached name registry + 3. Batch fetching the remaining uncached tags in one DB query Args: tag_names: List of tag names to fetch @@ -1505,50 +1784,22 @@ async def get_tag_objects_batch( Returns: Dictionary mapping tag_name to LiteLLM_TagTable object """ - if prisma_client is None: + if prisma_client is None or not tag_names: return {} - if not tag_names: - return {} - - tag_objects: Final = dict[str, LiteLLM_TagTable]() - uncached_tags: Final = list[str]() - - # Try to get all tags from cache first - for tag_name in tag_names: - cache_key = f"tag:{tag_name}" - cached_tag = await user_api_key_cache.async_get_cache( - key=cache_key, - model_type=LiteLLM_TagTable, + probed: Final = [ + ( + tag_name, + await user_api_key_cache.async_get_cache(key=tag_cache_key(tag_name), model_type=LiteLLM_TagTable), ) - if cached_tag is not None: - tag_objects[tag_name] = cached_tag - else: - uncached_tags.append(tag_name) - - # Batch fetch uncached tags from DB in one query - if uncached_tags: - try: - db_tags: Final = await _tag_table(TagRepository(prisma_client)).find_many( - where={"tag_name": {"in": uncached_tags}}, - include={"litellm_budget_table": True}, - ) - - # Cache and add to tag_objects - for db_tag in db_tags: - tag_name = db_tag.tag_name - cache_key = f"tag:{tag_name}" - _tag_obj = LiteLLM_TagTable.model_validate(db_tag.dict()) - await user_api_key_cache.async_set_cache( - key=cache_key, - value=_tag_obj, - model_type=LiteLLM_TagTable, - ) - tag_objects[tag_name] = _tag_obj - except Exception as e: - verbose_proxy_logger.debug("Error batch fetching tags from database: %s", e) - - return tag_objects + for tag_name in tag_names + ] + fetched: Final = await _fetch_uncached_tags( + uncached_tags=tuple(tag_name for tag_name, tag_obj in probed if tag_obj is None), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return {tag_name: tag_obj for tag_name, tag_obj in (*probed, *fetched) if tag_obj is not None} @log_db_metrics @@ -4573,25 +4824,15 @@ async def delete_cached_project_object( user_api_key_cache: UserApiKeyCache, ) -> None: """ - Every endpoint that mutates litellm_projecttable must call this: get_project_object - serves auth cache-first with no freshness check, so without invalidation a stale - project (e.g. a pre-update empty model allowlist) keeps being enforced until the - TTL expires (LIT-3803). Best-effort on both steps: the DB write has already - committed, so a cache backend error must not fail the endpoint; the stale entry - then expires via TTL. + Every endpoint that mutates litellm_projecttable must call this, or a stale project (e.g. a + pre-update empty model allowlist) keeps being enforced until the TTL expires (LIT-3803). """ - from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast - cache_key: Final = _project_cache_key(project_id) - try: - await user_api_key_cache.async_delete_cache(key=cache_key) - except Exception as e: # noqa: BLE001 # best-effort eviction: any cache backend error must not fail the mutation - verbose_proxy_logger.warning( - "Failed to evict cached project entry %s; a stale project may be served until its TTL expires: %s", - cache_key, - e, - ) - await publish_auth_cache_invalidation(cache_key=cache_key) + await evict_and_broadcast( + cache_keys=(_project_cache_key(project_id),), + user_api_key_cache=user_api_key_cache, + ) async def _organization_max_budget_check( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f7a04ba79e7..99592d44f9b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2307,6 +2307,7 @@ async def _run_centralized_common_checks( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, route=route, + token_end_user_max_budget=user_api_key_auth_obj.end_user_max_budget, ), ) ) @@ -2841,6 +2842,7 @@ async def _lookup_end_user_and_apply_budget( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, route=route, + token_end_user_max_budget=valid_token.end_user_max_budget, ) if end_user_object is not None: end_user_params = { diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 6952c0c6f89..623900322be 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -40,6 +40,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( update_batch_in_database, validate_managed_id_requirement, ) +from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.proxy.utils import handle_exception_on_proxy, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository from litellm.types.llms.openai import LiteLLMBatchCreateRequest @@ -140,6 +141,8 @@ async def create_batch( ) data["metadata"] = sanitize_openai_provider_metadata(data.get("metadata")) + raise_if_required_body_param_missing(route_type="acreate_batch", data=data) + ## check if model is a loadbalanced model router_model: str | None = None is_router_model = False diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 7fc8da42a3d..acdc9728390 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,5 +1,6 @@ import asyncio import json +from collections.abc import Sequence from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Final @@ -72,6 +73,27 @@ async def publish_auth_cache_invalidation(cache_key: str) -> None: verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) +async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "UserApiKeyCache") -> None: + """ + Drop cached management objects here and on every other worker. + + Every endpoint that mutates a cached object must call this: auth serves those objects + cache-first with no freshness check, so a mutation that leaves the entry in place keeps the + stale object enforced until its TTL expires (LIT-3803). Best-effort on both steps: the DB write + has already committed, so a cache backend error must not fail the endpoint. + """ + for cache_key in cache_keys: + try: + await user_api_key_cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # best-effort eviction: any cache backend error must not fail the mutation + verbose_proxy_logger.warning( + "Failed to evict cached entry %s; a stale object may be served until its TTL expires: %s", + cache_key, + e, + ) + await publish_auth_cache_invalidation(cache_key=cache_key) + + class AuthCacheInvalidationSubscriber: __slots__ = ("_redis_cache", "_task", "_user_api_key_cache") diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bf760a92d88..7b7cba5fc42 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -28,6 +28,7 @@ from litellm.proxy.common_utils.timezone_utils import ( compute_budget_reset_at, get_budget_reset_settings, ) +from litellm.proxy.common_utils.user_api_key_cache import tag_cache_key from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import ReadOnlyTable, SpendLinkedTable @@ -112,7 +113,7 @@ def _tag_counter_key(row: _TagRow) -> str: def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]: - return (f"tag:{row.tag_name}",) + return (tag_cache_key(row.tag_name),) def _budget_link_where( diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 22c3741d1a2..93d51bdd461 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -170,6 +170,36 @@ def object_permission_cache_key(object_permission_id: str) -> str: return f"object_permission_id:{object_permission_id}" +#: Cached under ``tag_registry_cache_key`` when the table exceeds ``TAG_REGISTRY_MAX_SIZE``: +#: registry unusable, fall back to the per-tag lookup. +TAG_REGISTRY_OVERFLOW_SENTINEL: Final = "__tag_registry_overflow__" + + +def tag_cache_key(tag_name: str) -> str: + """Cache key one tag row is stored under; shared so its five reader/writer modules cannot drift.""" + return f"tag:{tag_name}" + + +def tag_registry_cache_key() -> str: + """Cache key for the set of tag names that exist in ``LiteLLM_TagTable``.""" + return "tag_registry" + + +#: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds +#: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch. +END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__" + + +def end_user_cache_key(end_user_id: str) -> str: + """Cache key one end-user row is stored under; shared so auth and spend tracking cannot drift.""" + return f"end_user_id:{end_user_id}" + + +def end_user_restricted_registry_cache_key() -> str: + """Cache key for the set of end-user ids whose row carries a restriction auth enforces.""" + return "end_user_restricted_registry" + + def get_management_object_ttl(cache: DualCache) -> float: """ In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index c4a350fb285..6172f3a9158 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -38,6 +38,8 @@ from litellm.proxy._types import ( CommonProxyErrors, LitellmDataForBackendLLMCall, LitellmUserRoles, + ProxyErrorTypes, + ProxyException, SpecialHeaders, TeamCallbackMetadata, UserAPIKeyAuth, @@ -348,6 +350,36 @@ def reject_url_valued_destination(field: str, value: str) -> None: ) +_METADATA_JSON_TYPE_NAMES: Final[Mapping[type, str]] = MappingProxyType( + {bool: "a boolean", int: "an integer", float: "a number", str: "a string", list: "an array"} +) + + +def _invalid_metadata_type_error(field: str, value: object) -> ProxyException: + received_type: Final = _METADATA_JSON_TYPE_NAMES.get(type(value), f"a {type(value).__name__}") + return ProxyException( + message=f"Invalid type for '{field}': expected an object, but got {received_type} instead.", + type=ProxyErrorTypes.bad_request_error, + param=field, + code=400, + ) + + +def _normalized_metadata_object(field: str, value: object) -> Mapping[str, Any]: + """Return ``value`` as a metadata object or raise a 400 like OpenAI does. + + A JSON string that parses to an object is accepted because multipart/form-data + and ``extra_body`` callers can only send metadata as a string. The caller pops + the raw value from the request body before validating so the failure-logging + hooks that inspect the body afterwards don't crash on it and mask the 400 as a 500. + """ + if isinstance(value, dict): + return value + if isinstance(value, str) and isinstance((parsed := safe_json_loads(value)), dict): + return parsed + raise _invalid_metadata_type_error(field=field, value=value) + + def _strip_untrusted_request_header_controls( headers: Any, *, @@ -1572,6 +1604,13 @@ async def add_litellm_data_to_request( continue data.pop(_internal_key, None) _reject_url_valued_destinations(data) + _raw_metadata_by_field: Final = { + _metadata_field: data.pop(_metadata_field) + for _metadata_field in ("metadata", "litellm_metadata") + if data.get(_metadata_field) is not None + } + for _metadata_field, _raw_metadata in _raw_metadata_by_field.items(): + data[_metadata_field] = _normalized_metadata_object(_metadata_field, _raw_metadata) # Strip spoofable auth metadata from user-supplied metadata dict _user_metadata = data.get("metadata") if isinstance(_user_metadata, dict): @@ -1711,29 +1750,10 @@ async def add_litellm_data_to_request( verbose_proxy_logger.debug("receiving data: %s", data) - # Parse metadata if it's a string (e.g., from multipart/form-data) - if "metadata" in data and data["metadata"] is not None: - if isinstance(data["metadata"], str): - data["metadata"] = safe_json_loads(data["metadata"]) - if not isinstance(data["metadata"], dict): - verbose_proxy_logger.warning( - "Failed to parse 'metadata' as JSON dict. Received value: %s", data["metadata"] - ) - # requester_metadata is snapshotted AFTER the strip below so - # downstream consumers (e.g. PANW guardrail reading user_ip / - # profile_id) don't see attacker-injected admin slots preserved in - # the deepcopy. - - # Parse litellm_metadata if it's a string (e.g., from multipart/form-data or extra_body) - if "litellm_metadata" in data and data["litellm_metadata"] is not None: - if isinstance(data["litellm_metadata"], str): - parsed_litellm_metadata: Final = safe_json_loads(data["litellm_metadata"]) - if not isinstance(parsed_litellm_metadata, dict): - verbose_proxy_logger.warning( - "Failed to parse 'litellm_metadata' as JSON dict. Received value: %s", data["litellm_metadata"] - ) - else: - data["litellm_metadata"] = parsed_litellm_metadata + # requester_metadata is snapshotted AFTER the strip below so + # downstream consumers (e.g. PANW guardrail reading user_ip / + # profile_id) don't see attacker-injected admin slots preserved in + # the deepcopy. # Strip internal pipeline state and admin-injection slots from user input. # Runs AFTER the string-to-dict parse above so JSON-string metadata (sent diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 6c25f096532..9ef3d2defef 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -29,6 +29,10 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import ( + end_user_cache_key, + end_user_restricted_registry_cache_key, +) from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.proxy.management_endpoints.common_utils import validate_budget_duration from litellm.proxy.management_helpers.object_permission_utils import ( @@ -99,6 +103,25 @@ def _typed_table(repo: EndUserRepository | BudgetRepository) -> object: router: Final = APIRouter() +async def _evict_end_user_cache_keys(cache_keys: Sequence[str]) -> None: + """ + Every endpoint that mutates an end-user row must call this, or a newly blocked or budgeted + customer keeps being served unrestricted until the TTL expires: auth reads end users + cache-first, and the cached restricted-id registry decides whether the row is read at all. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + ) + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache) + + +def _end_user_cache_keys(user_ids: Sequence[str]) -> tuple[str, ...]: + """The per-id entries plus the registry, which any restriction change can move ids in or out of.""" + return (*(end_user_cache_key(user_id) for user_id in user_ids), end_user_restricted_registry_cache_key()) + + def _to_customer_response(record: BaseModel) -> CustomerResponse: """Validate a raw end-user DB row into the typed customer response. @@ -152,6 +175,7 @@ async def block_user(data: BlockUsers): }, ) records.append(record) + await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) else: raise HTTPException( status_code=500, @@ -448,6 +472,8 @@ async def new_end_user( include={"litellm_budget_table": True, "object_permission": True}, ) + await _evict_end_user_cache_keys(_end_user_cache_keys((data.user_id,))) + return _to_customer_response(end_user_record) except Exception as e: verbose_proxy_logger.exception( @@ -691,6 +717,8 @@ async def update_end_user( raise ValueError(f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}") verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response) + await _evict_end_user_cache_keys(_end_user_cache_keys((data.user_id,))) + return _to_customer_response(response) else: raise ValueError(f"user_id is required, passed user_id = {data.user_id}") @@ -764,6 +792,9 @@ async def delete_end_user( where={"user_id": {"in": data.user_ids}} ) verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response) + + await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) + return DeleteCustomersResponse( deleted_customers=response, message="Successfully deleted customers with ids: " + str(data.user_ids), diff --git a/litellm/proxy/management_endpoints/management_v1/spend_logs.py b/litellm/proxy/management_endpoints/management_v1/spend_logs.py index 96e60fcfdfc..5fee8eaede3 100644 --- a/litellm/proxy/management_endpoints/management_v1/spend_logs.py +++ b/litellm/proxy/management_endpoints/management_v1/spend_logs.py @@ -1,7 +1,7 @@ """`/management/v1/spend_logs` facets.""" from datetime import datetime, timezone -from typing import Annotated, Any, Final +from typing import Annotated, Any, Final, Literal from fastapi import APIRouter, Depends, Query, Request @@ -35,7 +35,7 @@ def _as_utc(value: datetime) -> datetime: return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc) -async def _end_user_scope_clause( +async def _spend_log_scope_clause( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, next_param_index: int, @@ -43,8 +43,8 @@ async def _end_user_scope_clause( """SQL predicate restricting the facet to spend logs this caller may read. Returns ``(None, ())`` for a proxy admin. Mirrors the scoping ``/spend/logs/ui`` - applies, so the dropdown can never offer an end user whose rows the caller - could not open. + applies, so a dropdown can never offer a value from a row the caller could + not open. """ from litellm.proxy.spend_tracking.spend_management_endpoints import ( _get_permitted_team_ids_for_spend_logs, @@ -77,6 +77,98 @@ async def _end_user_scope_clause( return f"({' OR '.join(clauses)})", params +async def _list_spend_log_facet( + request: Request, + user_api_key_dict: UserAPIKeyAuth, + start_time: datetime, + end_time: datetime, + q: str | None, + page: int, + page_size: int, + column: Literal["end_user", "user"], +) -> FacetListResponse: + try: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}database-not-connected", + title="Database not connected", + status=503, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + ) + + column_sql: Final = "end_user" if column == "end_user" else '"user"' + window_params: Final[tuple[Any, ...]] = (_as_utc(start_time), _as_utc(end_time)) + search_params: Final[tuple[Any, ...]] = (f"%{escape_like(q)}%",) if q else () + search_clause: Final = (f"{column_sql} ILIKE ${len(window_params) + 1} ESCAPE '\\'",) if q else () + + scope_clause, scope_params = await _spend_log_scope_clause( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + next_param_index=len(window_params) + len(search_params) + 1, + ) + + where_parts: Final = ( + ( + "\"startTime\" >= ($1::timestamptz AT TIME ZONE 'UTC')", + "\"startTime\" <= ($2::timestamptz AT TIME ZONE 'UTC')", + f"{column_sql} IS NOT NULL", + f"{column_sql} != ''", + ) + + search_clause + + ((scope_clause,) if scope_clause is not None else ()) + ) + + # The inner LIMIT walks the startTime index newest first and bounds the + # rows DISTINCT can inspect. request_id makes the cut-off deterministic, + # and page_size + 1 reveals has_more without a COUNT(*). + params: Final = ( + window_params + + search_params + + scope_params + + (SPEND_LOGS_FACET_SCAN_CAP, page_size + 1, (page - 1) * page_size) + ) + scan_idx: Final = len(params) - 2 + facet_sql: Final = ( + f"SELECT DISTINCT {column_sql} FROM (" + f" SELECT {column_sql}" + f' FROM "LiteLLM_SpendLogs"' + f" WHERE {' AND '.join(where_parts)}" + f' ORDER BY "startTime" DESC, request_id DESC' + f" LIMIT ${scan_idx}" + f") recent" + f" ORDER BY {column_sql} ASC" + f" LIMIT ${scan_idx + 1} OFFSET ${scan_idx + 2}" + ) + rows: Final = await prisma_client.db.query_raw(facet_sql, *params) + values: Final[list[str]] = [row[column] for row in rows if row.get(column)] + has_more: Final = len(values) > page_size + + return FacetListResponse( + data=values[:page_size], + meta=PageMeta(page=page, page_size=page_size, has_more=has_more), + links=build_page_links(request=request, page=page, has_more=has_more), + ) + except ManagementProblem: + raise + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.management_endpoints.management_v1.spend_logs._list_spend_log_facet(): Exception occured - %s", + e, + ) + raise ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}internal-server-error", + title="Internal server error", + status=500, + detail=f"Failed to list spend log {column.replace('_', ' ')}s.", + ) + ) + + @router.get( "/spend_logs/end_users", tags=["Budget & Spend Tracking"], @@ -116,85 +208,47 @@ async def list_spend_log_end_users( --header 'Authorization: Bearer sk-1234' ``` """ - try: - from litellm.proxy.proxy_server import prisma_client + return await _list_spend_log_facet( + request=request, + user_api_key_dict=user_api_key_dict, + start_time=start_time, + end_time=end_time, + q=q, + page=page, + page_size=page_size, + column="end_user", + ) - if prisma_client is None: - raise ManagementProblem( - ProblemDetail( - type=f"{PROBLEM_TYPE_BASE}database-not-connected", - title="Database not connected", - status=503, - detail=CommonProxyErrors.db_not_connected_error.value, - ) - ) - window_params: Final[tuple[Any, ...]] = (_as_utc(start_time), _as_utc(end_time)) - search_params: Final[tuple[Any, ...]] = (f"%{escape_like(q)}%",) if q else () - search_clause: Final = (f"end_user ILIKE ${len(window_params) + 1} ESCAPE '\\'",) if q else () - - scope_clause, scope_params = await _end_user_scope_clause( - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - next_param_index=len(window_params) + len(search_params) + 1, - ) - - where_parts: Final = ( - ( - "\"startTime\" >= ($1::timestamptz AT TIME ZONE 'UTC')", - "\"startTime\" <= ($2::timestamptz AT TIME ZONE 'UTC')", - "end_user IS NOT NULL", - "end_user != ''", - ) - + search_clause - + ((scope_clause,) if scope_clause is not None else ()) - ) - - # The inner LIMIT is the safety bound: it walks the startTime index newest - # first and stops, so DISTINCT never runs over an unbounded row set. - # request_id breaks startTime ties so the cut-off row is deterministic and - # successive OFFSET pages agree on the set they are paging through. - # page_size + 1: one row beyond the page reveals has_more without a COUNT(*). - params: Final = ( - window_params - + search_params - + scope_params - + (SPEND_LOGS_FACET_SCAN_CAP, page_size + 1, (page - 1) * page_size) - ) - scan_idx: Final = len(params) - 2 - facet_sql: Final = ( - f"SELECT DISTINCT end_user FROM (" - f" SELECT end_user" - f' FROM "LiteLLM_SpendLogs"' - f" WHERE {' AND '.join(where_parts)}" - f' ORDER BY "startTime" DESC, request_id DESC' - f" LIMIT ${scan_idx}" - f") recent" - f" ORDER BY end_user ASC" - f" LIMIT ${scan_idx + 1} OFFSET ${scan_idx + 2}" - ) - rows: Final = await prisma_client.db.query_raw(facet_sql, *params) - end_users: Final[list[str]] = [row["end_user"] for row in rows if row.get("end_user")] - has_more: Final = len(end_users) > page_size - - return FacetListResponse( - data=end_users[:page_size], - meta=PageMeta(page=page, page_size=page_size, has_more=has_more), - links=build_page_links(request=request, page=page, has_more=has_more), - ) - - except ManagementProblem: - raise - except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.management_endpoints.management_v1.spend_logs.list_spend_log_end_users(): Exception occured - %s", - e, - ) - raise ManagementProblem( - ProblemDetail( - type=f"{PROBLEM_TYPE_BASE}internal-server-error", - title="Internal server error", - status=500, - detail="Failed to list spend log end users.", - ) - ) +@router.get( + "/spend_logs/users", + tags=["Budget & Spend Tracking"], + dependencies=[Depends(user_api_key_auth), Depends(reject_unknown_query_params)], + response_model=FacetListResponse, +) +async def list_spend_log_users( + request: Request, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + start_time: Annotated[ + datetime, + Query(alias="filter[startTime][gte]", description="Window start (UTC when no offset is given)"), + ], + end_time: Annotated[ + datetime, + Query(alias="filter[startTime][lte]", description="Window end (UTC when no offset is given)"), + ], + q: Annotated[str | None, Query(description="Case-insensitive partial match on the internal user id")] = None, + page: Annotated[int, Query(ge=1, description="Page number")] = 1, + page_size: Annotated[int, Query(ge=1, le=100, description="Page size")] = 50, +) -> FacetListResponse: + """The distinct internal users appearing in spend logs the caller can read.""" + return await _list_spend_log_facet( + request=request, + user_api_key_dict=user_api_key_dict, + start_time=start_time, + end_time=end_time, + q=q, + page=page, + page_size=page_size, + column="user", + ) diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 894ba116f25..7aeb5039687 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -21,6 +21,10 @@ from fastapi import APIRouter, Depends, HTTPException, Query from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import ( + tag_cache_key, + tag_registry_cache_key, +) from litellm.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, get_daily_activity, @@ -133,6 +137,20 @@ def _table( return prisma_table +async def _evict_tag_cache_keys(cache_keys: Sequence[str]) -> None: + """ + Every endpoint that mutates a tag row must call this, or a deleted tag keeps its budget + enforced and a newly created one stays invisible to the cached name registry until the TTL + expires: auth reads tags cache-first, with no freshness check. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + ) + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache) + + async def _get_internal_user_api_keys( prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, @@ -294,6 +312,8 @@ async def new_tag( } ) + await _evict_tag_cache_keys((tag_cache_key(tag.name), tag_registry_cache_key())) + # Update models with new tag if tag.models: tasks: Final = [] @@ -440,6 +460,8 @@ async def update_tag( data=update_data, ) + await _evict_tag_cache_keys((tag_cache_key(tag.name),)) + # Build response tag_config: Final = TagConfig( name=updated_tag_record.tag_name, @@ -689,6 +711,8 @@ async def delete_tag( # Delete tag from database await _table(TagRepository(prisma_client)).delete(where={"tag_name": data.name}) + await _evict_tag_cache_keys((tag_cache_key(data.name), tag_registry_cache_key())) + return {"message": f"Tag {data.name} deleted successfully"} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index ab159e84b6a..6a0b3c6bfb2 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -916,6 +916,7 @@ class ProxyInitializationHelpers: "path that can cause schema thrashing during rolling deploys where two " "LiteLLM versions contend for the same DB. Default is the v1 resolver." ), + envvar="USE_V2_MIGRATION_RESOLVER", ) @click.option( "--reload", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e07870c6867..9692d4449d4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -358,7 +358,9 @@ from litellm.proxy.common_utils.timezone_utils import ( ) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + end_user_cache_key, get_management_object_ttl, + tag_cache_key, ) from litellm.proxy.config_resolvers import resolve_fields from litellm.proxy.config_resolvers.alerting import ( @@ -2780,7 +2782,7 @@ async def _increment_end_user_and_tag_spend_counters( if end_user_id is not None: await _init_and_increment_unreserved_spend_counter( counter_key=f"spend:end_user:{end_user_id}", - source_cache_key=f"end_user_id:{end_user_id}", + source_cache_key=end_user_cache_key(end_user_id), increment=response_cost, reserved_counter_keys=reserved_counter_keys, ) @@ -2795,7 +2797,7 @@ async def _increment_end_user_and_tag_spend_counters( seen_tags.add(tag_name) await _init_and_increment_unreserved_spend_counter( counter_key=f"spend:tag:{tag_name}", - source_cache_key=f"tag:{tag_name}", + source_cache_key=tag_cache_key(tag_name), increment=response_cost, reserved_counter_keys=reserved_counter_keys, ) @@ -3134,7 +3136,7 @@ async def update_cache( if end_user_id is None or response_cost is None: return - _id: Final = f"end_user_id:{end_user_id}" + _id: Final = end_user_cache_key(end_user_id) try: # Fetch the existing cost for the given user cached_end_user: Final = await user_api_key_cache.async_get_cache(key=_id) @@ -3226,7 +3228,7 @@ async def update_cache( if not tag_name or not isinstance(tag_name, str): continue - cache_key = f"tag:{tag_name}" + cache_key = tag_cache_key(tag_name) # Fetch the existing tag object from cache cached_tag = await user_api_key_cache.async_get_cache(key=cache_key) if cached_tag is None: @@ -10371,6 +10373,8 @@ async def moderations( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception("litellm.proxy.proxy_server.moderations(): Exception occured - %s", e) + if isinstance(e, ProxyException): + raise if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index b347360a939..c85325b6fa9 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -6,7 +6,7 @@ import httpx from fastapi import HTTPException, status import litellm -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router_utils.common_utils import _is_proxy_admin_request # Client-supplied params that make the router or the call path fabricate a @@ -141,6 +141,7 @@ ROUTE_ENDPOINT_MAPPING: Final = { "aget_run": "/evals/{eval_id}/runs/{run_id}", "acancel_run": "/evals/{eval_id}/runs/{run_id}/cancel", "adelete_run": "/evals/{eval_id}/runs/{run_id}", + "acreate_batch": "/batches", } @@ -152,27 +153,33 @@ class ProxyModelNotFoundError(HTTPException): super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail) -REQUIRED_BODY_PARAM_BY_ROUTE: Final[Mapping[str, str]] = { - "acompletion": "messages", - "aembedding": "input", +REQUIRED_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, ...]]] = { + "acompletion": ("messages",), + "aembedding": ("input",), + "acreate_batch": ("input_file_id", "endpoint", "completion_window"), } -class ProxyMissingRequiredParamError(HTTPException): +class ProxyMissingRequiredParamError(ProxyException): def __init__(self, route: str, param: str): - detail: Final = {"error": f"{route}: Missing required parameter: '{param}'."} - super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail) - self.type = "invalid_request_error" - self.param = param + super().__init__( + message=f"{route}: Missing required parameter: '{param}'.", + type="invalid_request_error", + param=param, + code=status.HTTP_400_BAD_REQUEST, + ) def raise_if_required_body_param_missing(route_type: str, data: Mapping[str, object]) -> None: - required_param: Final = REQUIRED_BODY_PARAM_BY_ROUTE.get(route_type) - if required_param is None or data.get(required_param) is not None: + missing_param: Final = next( + (param for param in REQUIRED_BODY_PARAMS_BY_ROUTE.get(route_type, ()) if data.get(param) is None), + None, + ) + if missing_param is None: return raise ProxyMissingRequiredParamError( route=ROUTE_ENDPOINT_MAPPING.get(route_type, route_type), - param=required_param, + param=missing_param, ) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 58a85171cc7..4c0dbdc0f45 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -24,6 +24,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key, tag_cache_key from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -448,7 +449,7 @@ async def _get_end_user_budget_counter( if end_user_id is None: return None - source_cache_key: Final = f"end_user_id:{end_user_id}" + source_cache_key: Final = end_user_cache_key(end_user_id) max_budget = _to_float(valid_token.end_user_max_budget) fallback_spend = 0.0 if end_user_object is not None: @@ -502,7 +503,7 @@ async def _get_tag_budget_counters( counters.append( _BudgetCounter( counter_key=f"spend:tag:{tag_name}", - source_cache_key=f"tag:{tag_name}", + source_cache_key=tag_cache_key(tag_name), max_budget=max_budget, fallback_spend=_to_float(_get_value(tag_object, "spend")) or 0.0, entity_type="Tag", diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 8fb5570965b..ed2ecd8325a 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2426,7 +2426,23 @@ async def ui_view_spend_logs( user_api_key_dict=user_api_key_dict, request_id=request_id, ) - permitted_team_ids: list[str] | None = None + user_scope_applies: Final = ( + not is_request_id_lookup + and not is_admin_view + and team_id is None + and _can_user_view_spend_log(user_api_key_dict=user_api_key_dict) + ) + permitted_team_ids: Final = ( + await _get_permitted_team_ids_for_spend_logs_or_empty( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + ) + if user_scope_applies + else () + ) + explicit_user_requires_caller_scope: Final = ( + user_scope_applies and not permitted_team_ids and user_id is not None + ) if not is_request_id_lookup and not is_admin_view: if team_id is not None: can_view_team: Final = await _can_team_member_view_log( @@ -2440,25 +2456,22 @@ async def ui_view_spend_logs( detail={"error": f"Not authorized to view team spend for team_id={team_id}"}, ) where_conditions["team_id"] = team_id - where_conditions.pop("user", None) - else: - if _can_user_view_spend_log(user_api_key_dict=user_api_key_dict): - try: - permitted_team_ids = await _get_permitted_team_ids_for_spend_logs( - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - ) - except Exception: - permitted_team_ids = [] - if permitted_team_ids: + elif user_scope_applies: + if permitted_team_ids: + if user_id is None: where_conditions.pop("user", None) - where_conditions["OR"] = [ - {"user": user_api_key_dict.user_id}, - {"team_id": {"in": permitted_team_ids}}, - ] - else: + where_conditions["OR"] = [ + {"user": user_api_key_dict.user_id}, + {"team_id": {"in": permitted_team_ids}}, + ] + else: + if user_id is None: where_conditions["user"] = user_api_key_dict.user_id - where_conditions.pop("team_id", None) + else: + where_conditions["AND"] = where_conditions.get("AND", []) + [ + {"user": user_api_key_dict.user_id} + ] + where_conditions.pop("team_id", None) # Calculate skip value for pagination skip: Final = (page - 1) * page_size @@ -2502,12 +2515,16 @@ async def ui_view_spend_logs( p += 1 # Multi-team OR filter: (user = $X OR team_id = ANY($Y)) - if permitted_team_ids is not None and len(permitted_team_ids) > 0: + if permitted_team_ids: or_clause: Final = f'("user" = ${p} OR team_id = ANY(${p + 1}::text[]))' sql_params.append(user_api_key_dict.user_id) sql_params.append(permitted_team_ids) p += 2 sql_conditions.append(or_clause) + elif explicit_user_requires_caller_scope: + sql_conditions.append(f'"user" = ${p}') + sql_params.append(user_api_key_dict.user_id) + p += 1 if session_id is not None and isinstance(session_id, str): like_escaped_session_id: Final = session_id.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") @@ -4272,3 +4289,19 @@ async def _get_permitted_team_ids_for_spend_logs( ): permitted.append(team_obj.team_id) return permitted + + +async def _get_permitted_team_ids_for_spend_logs_or_empty( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[str, ...]: + """Resolve permitted teams once, falling back to the caller's own-user scope.""" + try: + return tuple( + await _get_permitted_team_ids_for_spend_logs( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + ) + ) + except Exception: + return () diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index b3725d86ec3..197d0c02ba8 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -806,6 +806,7 @@ class LiteLLM_Proxy_MCP_Handler: oauth2_headers=oauth2_headers, raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, + litellm_logging_obj=litellm_logging_obj, ) if proxy_logging_obj: diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 69d291eebd0..17ba78b0190 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -3,7 +3,7 @@ from enum import Enum from typing import Any, Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict -from typing_extensions import NotRequired, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from .openai import ( ChatCompletionCachedContent, @@ -48,6 +48,7 @@ class AnthropicMessagesTool(TypedDict, total=False): name: Required[str] description: str input_schema: AnthropicInputSchema | None + strict: ReadOnly[bool] type: Literal["custom"] cache_control: dict | ChatCompletionCachedContent | None defer_loading: bool diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 22d3119d457..15396a95632 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -201,7 +201,7 @@ "limit": 58 }, "SIM102": { - "limit": 319 + "limit": 317 }, "SIM103": { "limit": 119 diff --git a/tests/e2e/claude_code/_probe_unit_tests/__init__.py b/tests/e2e/claude_code/_probe_unit_tests/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/claude_code/_probe_unit_tests/test_http_probe.py b/tests/e2e/claude_code/_probe_unit_tests/test_http_probe.py new file mode 100644 index 00000000000..6989868ba57 --- /dev/null +++ b/tests/e2e/claude_code/_probe_unit_tests/test_http_probe.py @@ -0,0 +1,106 @@ +"""Unit tests for the tool-search replay assertion in `http_probe`. + +Markerless harness tests: they exercise probe plumbing over hand-built +`Result` values, not a product feature, so they run without a proxy and carry +no `e2e` marker. + +The red paths are what these are for. A live cell only ever executes the green +one, so a broken diagnostic in the failure branch would sit undetected until +the day the provider actually rejects the history, which is the day the +diagnostic has to be right. +""" + +from __future__ import annotations + +from e2e_http import Result, Success, UnknownApiError +from models import ( + AnthropicContentBlock, + AnthropicMessagesResponse, + AnthropicToolResultTurn, + ChatMessage, +) + +from claude_code.http_probe import ( + ToolSearchReplay, + _replay_history, + assert_tool_search_replay_shape, +) + +_REJECTED: Result[AnthropicMessagesResponse] = UnknownApiError( + status_code=400, + body="server_tool_use blocks are not supported", +) +_ACCEPTED: Result[AnthropicMessagesResponse] = Success( + status_code=200, + data=AnthropicMessagesResponse(content=[AnthropicContentBlock(type="text", text="done")]), +) + + +def _replay(block_types: tuple[str, ...], second_turn: Result[AnthropicMessagesResponse]) -> ToolSearchReplay: + answer = AnthropicMessagesResponse( + content=[AnthropicContentBlock(type=block_type, id="srvtoolu_01") for block_type in block_types] + ) + return ToolSearchReplay( + first_turn=Success(status_code=200, data=answer), + history=_replay_history(answer), + second_turn=second_turn, + ) + + +def test_accepts_a_replayed_server_tool_pair() -> None: + replay = _replay(("text", "server_tool_use", "tool_search_tool_result"), _ACCEPTED) + assert assert_tool_search_replay_shape(replay) is None + + +def test_reports_the_status_when_the_replayed_history_is_rejected() -> None: + replay = _replay(("server_tool_use", "tool_search_tool_result"), _REJECTED) + error = assert_tool_search_replay_shape(replay) + assert error is not None + assert "status 400" in error + assert "server_tool_use" in error + + +def test_a_turn_truncated_before_the_result_block_is_not_a_pass() -> None: + replay = _replay(("server_tool_use",), _ACCEPTED) + error = assert_tool_search_replay_shape(replay) + assert error is not None + assert "tool_search_tool_result" in error + + +def test_a_history_with_no_server_tool_block_is_not_a_pass() -> None: + replay = _replay(("text",), _ACCEPTED) + error = assert_tool_search_replay_shape(replay) + assert error is not None + assert "server_tool_use" in error + + +def test_a_failed_first_turn_is_reported_as_the_first_turn() -> None: + replay = ToolSearchReplay(first_turn=_REJECTED, history=(), second_turn=None) + error = assert_tool_search_replay_shape(replay) + assert error is not None + assert error.startswith("first turn: ") + + +def test_a_pending_tool_use_is_answered_with_the_id_the_model_returned() -> None: + answer = AnthropicMessagesResponse( + content=[ + AnthropicContentBlock(type="server_tool_use", id="srvtoolu_01"), + AnthropicContentBlock(type="tool_search_tool_result", id=None), + AnthropicContentBlock(type="tool_use", id="toolu_99"), + ] + ) + last_turn = _replay_history(answer)[-1] + assert isinstance(last_turn, AnthropicToolResultTurn) + assert [block.tool_use_id for block in last_turn.content] == ["toolu_99"] + + +def test_a_turn_with_no_pending_tool_use_gets_a_plain_follow_up() -> None: + answer = AnthropicMessagesResponse( + content=[ + AnthropicContentBlock(type="server_tool_use", id="srvtoolu_01"), + AnthropicContentBlock(type="tool_search_tool_result"), + ] + ) + last_turn = _replay_history(answer)[-1] + assert isinstance(last_turn, ChatMessage) + assert last_turn.role == "user" diff --git a/tests/e2e/claude_code/http_probe.py b/tests/e2e/claude_code/http_probe.py index c77020acd6e..8aba54576c4 100644 --- a/tests/e2e/claude_code/http_probe.py +++ b/tests/e2e/claude_code/http_probe.py @@ -28,6 +28,7 @@ the upstream, or LiteLLM 500 on a transformation bug). from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING from pydantic import BaseModel @@ -42,10 +43,14 @@ from e2e_http import ( ValidationError, ) from models import ( + AnthropicAssistantTurn, AnthropicCustomTool, + AnthropicMessage, AnthropicMessagesBody, AnthropicMessagesResponse, AnthropicTool, + AnthropicToolResultBlock, + AnthropicToolResultTurn, AnthropicToolSearchTool, ChatMessage, CountTokensBody, @@ -132,6 +137,7 @@ def probe_tool_search( client: ProxyClient, api_key: str, model: str, + max_tokens: int = 64, rate_limiter: RateLimiter | None = None, ) -> Result[AnthropicMessagesResponse]: """POST to `/v1/messages` with a `tool_search_tool_regex_20251119` tool @@ -155,13 +161,115 @@ def probe_tool_search( api_key, AnthropicMessagesBody( model=model, - max_tokens=64, + max_tokens=max_tokens, messages=[ChatMessage(role="user", content=_TOOL_SEARCH_PROMPT)], tools=list(_TOOL_SEARCH_TOOLS), ), ) +_TOOL_SEARCH_FOLLOW_UP = "Thanks. Now reply with the word 'done'." +_TOOL_RESULT_STUB = "3" +# A `server_tool_use` block and the `tool_search_tool_result` answering it are +# one indivisible pair: replaying the request without its result is malformed +# Anthropic and 400s on any provider. 64 output tokens is not enough room for +# both, so the turn we replay is generated with a budget that fits the whole +# discovery round trip. +_REPLAY_SOURCE_MAX_TOKENS = 1024 +_REPLAYED_SERVER_BLOCKS = frozenset({"server_tool_use", "tool_search_tool_result"}) + + +@dataclass(frozen=True, slots=True) +class ToolSearchReplay: + """Both turns of the multi-turn probe plus the history the second turn + carried, so a failing cell can report which turn broke and what was on the + wire when it did.""" + + first_turn: Result[AnthropicMessagesResponse] + history: tuple[AnthropicMessage, ...] + second_turn: Result[AnthropicMessagesResponse] | None + + +def _replayed_server_block_types(history: tuple[AnthropicMessage, ...]) -> frozenset[str]: + return frozenset( + block.type + for turn in history + if isinstance(turn, AnthropicAssistantTurn) + for block in turn.content + if block.type in _REPLAYED_SERVER_BLOCKS + ) + + +def _replay_history(answer: AnthropicMessagesResponse) -> tuple[AnthropicMessage, ...]: + """Turn a real first-turn answer into a well-formed two-turn history. + + Every client-side `tool_use` the model emitted gets a `tool_result` keyed on + the id the model actually returned; a turn with none gets a plain follow-up + instead. An unanswered `tool_use`, or a `tool_result` pointing at an invented + id, is malformed Anthropic and 400s on any provider, which would make this + probe measure our own request rather than the provider's handling of the + replayed server-tool blocks.""" + blocks = tuple(answer.content or ()) + pending = tuple(block.id for block in blocks if block.type == "tool_use" and block.id is not None) + reply: AnthropicMessage = ( + AnthropicToolResultTurn( + content=[ + AnthropicToolResultBlock(tool_use_id=tool_use_id, content=_TOOL_RESULT_STUB) + for tool_use_id in pending + ] + ) + if pending + else ChatMessage(role="user", content=_TOOL_SEARCH_FOLLOW_UP) + ) + return ( + ChatMessage(role="user", content=_TOOL_SEARCH_PROMPT), + AnthropicAssistantTurn(content=list(blocks)), + reply, + ) + + +def probe_tool_search_multiturn( + *, + client: ProxyClient, + api_key: str, + model: str, + rate_limiter: RateLimiter | None = None, +) -> ToolSearchReplay: + """Run `probe_tool_search`, then send the real assistant turn back as + history with the same tools still declared. + + The first turn only proves the proxy attaches the tool-search beta header on + the way out. Nothing proves the provider accepts the `server_tool_use` and + `tool_search_tool_result` blocks it produced when they come back in + `messages`, which is every turn of a real Claude Code session after the + first.""" + first_turn = probe_tool_search( + client=client, + api_key=api_key, + model=model, + max_tokens=_REPLAY_SOURCE_MAX_TOKENS, + rate_limiter=rate_limiter, + ) + if not isinstance(first_turn, Success): + return ToolSearchReplay(first_turn=first_turn, history=(), second_turn=None) + + history = _replay_history(first_turn.data) + _acquire(model, rate_limiter) + return ToolSearchReplay( + first_turn=first_turn, + history=history, + second_turn=client.messages( + api_key, + AnthropicMessagesBody( + model=model, + max_tokens=64, + messages=list(history), + tools=list(_TOOL_SEARCH_TOOLS), + ), + ), + ) + + def _failure_diagnostic[R: BaseModel](result: Result[R], route: str) -> str: """Map a non-success `Result` to a one-line diagnostic. The `status 429` wording is load-bearing: the compat conftest classifies a rate-limited cell @@ -207,6 +315,40 @@ def assert_tool_search_shape(result: Result[AnthropicMessagesResponse]) -> str | return _failure_diagnostic(result, "/v1/messages") +def assert_tool_search_replay_shape(replay: ToolSearchReplay) -> str | None: + """Return None on success, else describe the first violation. + + Acceptance criteria: + + 1. The first turn succeeded, on the same terms as `assert_tool_search_shape`. + 2. That turn produced a complete `server_tool_use` / `tool_search_tool_result` + pair to replay. Without both the second turn carries either an ordinary + text history or a half-finished tool call, and the cell would report on + our own request rather than on the provider's handling of server-tool + blocks in history. + 3. The provider accepted the history containing those blocks. + """ + first_error = assert_tool_search_shape(replay.first_turn) + if first_error is not None: + return f"first turn: {first_error}" + + replayed = _replayed_server_block_types(replay.history) + missing = _REPLAYED_SERVER_BLOCKS - replayed + if missing: + return ( + f"first turn returned no {' or '.join(sorted(missing))} block to replay, so the history " + "proves nothing about server-tool handling; a turn truncated at max_tokens looks like this" + ) + + if replay.second_turn is None: + return "second turn was never sent" + + second_error = assert_tool_search_shape(replay.second_turn) + if second_error is not None: + return f"history replaying {sorted(replayed)} rejected: {second_error}" + return None + + def assert_count_tokens_shape(result: Result[CountTokensResponse]) -> str | None: """Return None on success, or an error string describing the first violation. diff --git a/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py b/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py index c4735c78f0c..5b4c50e9dc5 100644 --- a/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py @@ -1,12 +1,17 @@ """tool_search x Bedrock (Invoke). -HTTP-probe row. Sends a single `/v1/messages` request whose `tools` -array includes a `tool_search_tool_regex_20251119` discovery tool, and +HTTP-probe row. Sends a `/v1/messages` request whose `tools` array +includes a `tool_search_tool_regex_20251119` discovery tool, and asserts the proxy round-trips it to the upstream without a 400. This verifies LiteLLM's tool-search beta-header translation (`advanced-tool-use-2025-11-20` for Anthropic-shape providers, `tool-search-tool-2025-10-19` for Vertex/Bedrock) survives end-to-end. +A second probe then replays that turn's answer as history, which is +what every turn of a real session after the first looks like: the +first turn only exercises the outbound header, and the blocks the +model sends back have to be accepted on the way in too. + The (feature, provider) for this cell is inferred from the file path by `tests/e2e/claude_code/conftest.py`: @@ -47,8 +52,10 @@ import pytest from claude_code._env import require_proxy_client from claude_code.http_probe import ( + assert_tool_search_replay_shape, assert_tool_search_shape, probe_tool_search, + probe_tool_search_multiturn, ) @@ -80,3 +87,31 @@ def test_tool_search_bedrock_invoke(compat_result): if failures: pytest.fail("; ".join(failures), pytrace=False) + + +@pytest.mark.covers("llm.messages.bedrock_invoke.tool_search_history.nonstream.works") +def test_tool_search_history_bedrock_invoke(compat_result): + """Send the tool-search request, take the real assistant turn back, and + replay it as history with the tools still declared. + + Every turn of a real Claude Code session after the first carries the + `server_tool_use` and `tool_search_tool_result` blocks the previous turn + produced. The single-turn probe above never sends them, so it cannot see a + provider or a transformation that accepts tool_search on the way out and + rejects the blocks it gets back.""" + client, api_key = require_proxy_client(compat_result) + + failures = [] + for model in BEDROCK_INVOKE_MODELS: + replay = probe_tool_search_multiturn(client=client, api_key=api_key, model=model) + shape_error = assert_tool_search_replay_shape(replay) + if shape_error is not None: + error = f"[{model}] tool_search history replay failed: {shape_error}" + compat_result.add({"status": "fail", "error": error}) + failures.append(error) + continue + + compat_result.add({"status": "pass"}) + + if failures: + pytest.fail("; ".join(failures), pytrace=False) diff --git a/tests/e2e/coverage_registry/llm_claude_code_compat.yaml b/tests/e2e/coverage_registry/llm_claude_code_compat.yaml index c2c17a6e764..d78f07564aa 100644 --- a/tests/e2e/coverage_registry/llm_claude_code_compat.yaml +++ b/tests/e2e/coverage_registry/llm_claude_code_compat.yaml @@ -8,7 +8,8 @@ # route : anthropic | azure_foundry | bedrock_converse | bedrock_invoke | vertex # capability : basic | tool_use | vision | thinking | prompt_cache_5m | prompt_cache_1h # | structured_output | pdf_input | long_context_1m -# | thinking_with_tool_use | tool_search | count_tokens | web_search +# | thinking_with_tool_use | tool_search | tool_search_history | count_tokens +# | web_search # streaming : stream | nonstream # ---- basic / non-streaming ---- @@ -94,6 +95,7 @@ - {id: llm.messages.bedrock_converse.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search discovery tool over Bedrock Converse"} - {id: llm.messages.bedrock_invoke.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search discovery tool over Bedrock Invoke"} - {id: llm.messages.vertex.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search discovery tool over Vertex AI"} +- {id: llm.messages.bedrock_invoke.tool_search_history.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: tool_search_history, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "A real server_tool_use / tool_search_tool_result pair replayed as history over Bedrock Invoke"} # ---- count_tokens ---- - {id: llm.messages.anthropic.count_tokens.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: count_tokens, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "/v1/messages/count_tokens over Anthropic direct"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 76844c039f1..a5c723f8965 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -76,6 +76,7 @@ LlmCapability = Literal[ "thinking", "thinking_with_tool_use", "tool_search", + "tool_search_history", "tool_use", "vision", "web_search", diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9ba191d7f0e..734e63a94e6 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -384,9 +384,45 @@ class AnthropicCustomTool(BaseModel): type AnthropicTool = AnthropicToolSearchTool | AnthropicWebSearchTool | AnthropicCustomTool +class AnthropicContentBlock(BaseModel): + """One block of a `content` array. Only the fields a test reads are + declared; `extra="allow"` keeps the rest (a `server_tool_use` block's + `input`, a `tool_search_tool_result` block's nested `content`) so an + assistant turn read off the wire can be replayed into history verbatim + instead of being silently flattened to its text.""" + + model_config = ConfigDict(extra="allow") + type: str | None = None + text: str | None = None + id: str | None = None + + +class AnthropicToolResultBlock(BaseModel): + """The user-turn answer to a client-side `tool_use`. `tool_use_id` must be + the id the model actually emitted; an invented one is rejected by + Anthropic's own schema validator, which Bedrock inherits.""" + + type: Literal["tool_result"] = "tool_result" + tool_use_id: str + content: str + + +class AnthropicAssistantTurn(BaseModel): + role: Literal["assistant"] = "assistant" + content: list[AnthropicContentBlock] + + +class AnthropicToolResultTurn(BaseModel): + role: Literal["user"] = "user" + content: list[AnthropicToolResultBlock] + + +type AnthropicMessage = ChatMessage | AnthropicAssistantTurn | AnthropicToolResultTurn + + class AnthropicMessagesBody(BaseModel): model: str - messages: list[ChatMessage] + messages: list[AnthropicMessage] max_tokens: int stream: bool | None = None tools: list[AnthropicTool] | None = None @@ -401,11 +437,6 @@ class CountTokensBody(BaseModel): messages: list[ChatMessage] -class AnthropicContentBlock(BaseModel): - type: str | None = None - text: str | None = None - - class AnthropicMessagesResponse(BaseModel): """A /v1/messages answer. `content` is the Anthropic-native passthrough shape; `choices` is the OpenAI-normalized shape LiteLLM emits for some diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 9145829ecb2..1185893d428 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -3518,6 +3518,53 @@ def test_translate_anthropic_tools_to_openai_preserves_parameters_type(): assert new_tools[0]["type"] == "function" +def test_translate_anthropic_tools_to_openai_maps_strict_onto_function_not_parameters(): + """A tool-level `strict` lands on the OpenAI function, leaving the caller's `input_schema` untouched.""" + adapter = LiteLLMAnthropicMessagesAdapter() + input_schema = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, + } + tools = [{"type": "custom", "name": "get_weather", "strict": True, "input_schema": input_schema}] + + new_tools, _ = adapter.translate_anthropic_tools_to_openai(tools=tools) + + function = new_tools[0]["function"] + assert function["strict"] is True + assert "strict" not in function["parameters"] + assert input_schema == { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, + } + + +def test_translate_anthropic_tools_to_openai_omits_unset_strict(): + """Chat Completions already defaults to non-strict, so an unset `strict` stays unset.""" + adapter = LiteLLMAnthropicMessagesAdapter() + tools = [ + { + "type": "custom", + "name": "search", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}, "cursor": {"type": "string"}}, + "required": ["query"], + }, + } + ] + + new_tools, _ = adapter.translate_anthropic_tools_to_openai(tools=tools) + + function = new_tools[0]["function"] + assert "strict" not in function + assert "strict" not in function["parameters"] + assert function["parameters"]["required"] == ["query"] + + TOOL_RESULT_IMAGE_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" TOOL_RESULT_IMAGE_URL = "https://example.com/screenshot.png" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index 73b58e71009..b1ae865fde1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -76,6 +76,78 @@ class TestProcessEventResponseCreatedGuard: assert len(message_starts) == 1 +class TestReasoningItemWithoutSummaryText: + """Regression: a reasoning item whose summary never produces text must not + surface as a thinking content block. + + OpenAI emits ``response.output_item.added`` with ``type: "reasoning"`` on + every reasoning turn, but only emits + ``response.reasoning_summary_text.delta`` when a summary was requested and + the model actually produced one. Eagerly opening the block on + ``output_item.added`` left ``{"type": "thinking", "thinking": ""}`` in the + assistant turn, which clients persist in their session transcript. Replaying + that transcript against an Anthropic model (what ``claude --resume`` does + once the resumed session falls back to the default Anthropic model) fails + with:: + + 400 invalid_request_error - messages.2.content.0.thinking: + each thinking block must contain thinking + + So the thinking block is opened on the first non-empty summary delta. + """ + + @staticmethod + def _gpt_turn(reasoning_summary_deltas: list) -> list: + return [ + {"type": "response.created"}, + {"type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}}, + *( + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": delta} + for delta in reasoning_summary_deltas + ), + {"type": "response.output_item.done", "item": {"type": "reasoning", "id": "rs_1"}}, + {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}, + {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hello"}, + {"type": "response.output_item.done", "item": {"type": "message", "id": "msg_1"}}, + ] + + def test_reasoning_without_summary_emits_no_thinking_block(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=[])) + + assert not [ + c for c in chunks if c["type"] == "content_block_start" and c["content_block"]["type"] == "thinking" + ] + assert [(c["type"], c.get("index")) for c in chunks[1:]] == [ + ("content_block_start", 0), + ("content_block_delta", 0), + ("content_block_stop", 0), + ] + assert chunks[1]["content_block"] == {"type": "text", "text": ""} + + def test_reasoning_with_only_empty_summary_deltas_emits_no_thinking_block(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=["", ""])) + + assert not [c for c in chunks if c["type"] == "content_block_delta" and c["delta"]["type"] == "thinking_delta"] + assert not [ + c for c in chunks if c["type"] == "content_block_start" and c["content_block"]["type"] == "thinking" + ] + + def test_reasoning_with_summary_text_still_emits_a_thinking_block(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=["Weigh", "ing options"])) + + assert [(c["type"], c.get("index")) for c in chunks[1:]] == [ + ("content_block_start", 0), + ("content_block_delta", 0), + ("content_block_delta", 0), + ("content_block_stop", 0), + ("content_block_start", 1), + ("content_block_delta", 1), + ("content_block_stop", 1), + ] + assert chunks[1]["content_block"] == {"type": "thinking", "thinking": ""} + assert "".join(c["delta"]["thinking"] for c in chunks[2:4]) == "Weighing options" + + class TestProcessEventTextDeltaWithoutOutputItemAdded: """Streams that skip response.output_item.added (e.g. LMStudio) must still open a text block before any delta and never emit index -1.""" @@ -110,12 +182,13 @@ class TestProcessEventTextDeltaWithoutOutputItemAdded: "type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}, }, + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": "hm"}, {"type": "response.output_text.delta", "item_id": "m1", "delta": "Hi"}, ] ) - assert chunks[1]["type"] == "content_block_start" - assert chunks[1]["content_block"] == {"type": "text", "text": ""} - assert [c["index"] for c in chunks[1:]] == [1, 1] + assert chunks[2]["type"] == "content_block_start" + assert chunks[2]["content_block"] == {"type": "text", "text": ""} + assert [c["index"] for c in chunks[2:]] == [1, 1] def test_process_event_registered_item_id_does_not_synthesize_start(self): chunks = _process_all( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 73d636fbc4b..876213eda3f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -22,7 +22,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) -from litellm.types.llms.anthropic import AnthropicMessagesRequest +from litellm.types.llms.anthropic import ( + AllAnthropicToolsValues, + AnthropicMessagesRequest, +) from litellm.types.llms.openai import ResponseAPIUsage @@ -606,6 +609,7 @@ class TestTranslateToolsToResponsesAPI: { "type": "function", "name": "get_weather", + "strict": False, "description": "Get current weather for a city.", "parameters": { "type": "object", @@ -615,6 +619,60 @@ class TestTranslateToolsToResponsesAPI: } ] + def test_tool_with_optional_properties_stays_non_strict(self): + """Regression: an unset Anthropic `strict` must not become the Responses strict default, + which would rewrite `required` to include every optional property.""" + tools: List[AllAnthropicToolsValues] = [ + { + "name": "search", + "input_schema": { + "type": "object", + "properties": { + "query": {"type": "string"}, + "cursor": {"type": "string"}, + }, + "required": ["query"], + "additionalProperties": False, + }, + } + ] + + result = _ADAPTER.translate_tools_to_responses_api(tools) + + assert result[0]["strict"] is False + assert result[0]["parameters"]["required"] == ["query"] + + def test_tool_forwards_explicit_strict_true(self): + """An explicit Anthropic `strict: True` still reaches Responses as True.""" + tools: List[AllAnthropicToolsValues] = [ + { + "name": "search", + "strict": True, + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + "additionalProperties": False, + }, + } + ] + + result = _ADAPTER.translate_tools_to_responses_api(tools) + + assert result == [ + { + "type": "function", + "name": "search", + "strict": True, + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + "additionalProperties": False, + }, + } + ] + def test_tool_without_description(self): """Tool without a description omits the description key.""" tools = [{"name": "ping", "input_schema": {"type": "object", "properties": {}}}] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 9bc84b43fc5..6bfd83b82b8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -8163,8 +8163,40 @@ async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): ) # per-server, not aggregate: the single server's name is in the endpoints assert "/test_oauth/authorize" in authorization_response["authorization_endpoint"] - assert authorization_response["issuer"] == "https://llm.example.com" - assert resource_response["authorization_servers"] == ["https://llm.example.com/test_oauth"] + expected_issuer = "https://llm.example.com/test_oauth" + assert authorization_response["issuer"] == expected_issuer + assert resource_response["authorization_servers"] == [expected_issuer] + finally: + global_mcp_server_manager.registry.clear() + + +@pytest.mark.asyncio +async def test_openid_configuration_alias_keeps_origin_issuer_on_root_resolution(): + """OIDC verifiers derive /.well-known/openid-configuration from their configured issuer + (the proxy origin), so the alias must keep the origin issuer even when single-server + root resolution scopes the underlying authorization-server metadata.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + openid_configuration, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + response = await openid_configuration(mock_request) + assert isinstance(response, dict) + assert response["issuer"] == "https://llm.example.com" + assert "/test_oauth/authorize" in response["authorization_endpoint"] finally: global_mcp_server_manager.registry.clear() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py new file mode 100644 index 00000000000..64d926bc5e3 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py @@ -0,0 +1,126 @@ +"""Tests for guardrail-block recording in +``litellm.proxy._experimental.mcp_server.server.call_mcp_tool``. + +A pre-call MCP guardrail block *raises* into ``call_mcp_tool``'s +``except Exception``. The failure spend-log row that the Guardrails Monitor's +"Total Blocked" counts is written by ``_ProxyDBLogger.async_post_call_failure_hook`` +(reached via ``proxy_logging_obj.post_call_failure_hook``), which reads +``standard_logging_object`` off the request's logging obj -- and that only exists +once ``failure_handler`` / ``async_failure_handler`` have run. So the failure +handlers must run *before* ``post_call_failure_hook``, otherwise the row persists +with ``guardrail_information=None`` and the block is never counted. These tests +pin that ordering. + +``call_mcp_tool`` is wrapped by ``@client`` (``litellm.utils.client``), which uses +``functools.wraps`` and therefore exposes the raw undecorated coroutine as +``__wrapped__``. The tests drive ``__wrapped__`` directly so the except-block +ordering is observed in isolation, without the wrapper's own post-raise logging +firing. Note that this means they do not exercise the wrapper's dedup path; that +dedup rests on ``should_run_logging("sync_failure")`` / ``("async_failure")``, +which has its own coverage in the logging tests. + +``proxy_logging_obj`` is imported lazily inside the except block via +``from litellm.proxy.proxy_server import proxy_logging_obj``; the real +``proxy_server`` module is heavy, so a fake module is injected into ``sys.modules`` +to satisfy that lazy import without loading it. +""" + +import contextlib +import sys +import types +from unittest import mock + +import pytest +from fastapi import HTTPException + +from litellm.proxy._experimental.mcp_server import server + + +class _RecordingLoggingObj: + """Stands in for ``LiteLLMLoggingObj``, recording the failure flush the fix + makes so the test can assert it happens before ``post_call_failure_hook``.""" + + def __init__(self, order: list) -> None: + self._order = order + self.failure_calls = 0 + self.async_failure_calls = 0 + + def failure_handler(self, *_args, **_kwargs) -> None: + self.failure_calls += 1 + self._order.append("failure_handler") + + async def async_failure_handler(self, *_args, **_kwargs) -> None: + self.async_failure_calls += 1 + self._order.append("async_failure_handler") + + +async def _call_block(logging_obj, order: list, *, user_api_key_auth=mock.sentinel.auth): + """Drive ``call_mcp_tool`` into its except path via ``arguments=None``, which + raises ``HTTPException(400)`` before any server-manager call, and return once it + re-raises.""" + + async def _record_post_call_failure_hook(**_kwargs) -> None: + order.append("post_call_failure_hook") + + proxy_logging_obj = mock.MagicMock() + proxy_logging_obj.post_call_failure_hook.side_effect = _record_post_call_failure_hook + + fake_proxy_server = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy_server.proxy_logging_obj = proxy_logging_obj # pyright: ignore[reportAttributeAccessIssue] + + with mock.patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}): + with contextlib.suppress(HTTPException): + await server.call_mcp_tool.__wrapped__( + name="t", + arguments=None, + user_api_key_auth=user_api_key_auth, + litellm_logging_obj=logging_obj, + ) + + +@pytest.mark.asyncio +async def test_block_flushes_failure_before_post_call_failure_hook(): + order: list = [] + await _call_block(_RecordingLoggingObj(order), order) + + assert order == ["failure_handler", "async_failure_handler", "post_call_failure_hook"], order + + +@pytest.mark.asyncio +async def test_block_flushes_each_handler_exactly_once(): + """Each handler runs once, so the block yields exactly one counted row rather + than double-counting on the shared logging obj.""" + order: list = [] + obj = _RecordingLoggingObj(order) + await _call_block(obj, order) + + assert (obj.failure_calls, obj.async_failure_calls) == (1, 1) + + +@pytest.mark.asyncio +async def test_block_flushes_failure_for_anonymous_calls(): + """With no ``user_api_key_auth`` the failure handlers still run, so OTel and the + other failure sinks see the block. + + ``post_call_failure_hook`` stays gated on auth, matching the pre-existing + contract: SpendLogs rows are attributable billing/audit records and the + downstream DB logger dereferences authenticated key, budget, and route data. + Counting anonymous MCP blocks needs a counter that does not live in SpendLogs, + which is a separate design change, not part of this fix. + """ + order: list = [] + obj = _RecordingLoggingObj(order) + await _call_block(obj, order, user_api_key_auth=None) + + assert order == ["failure_handler", "async_failure_handler"], order + + +@pytest.mark.asyncio +async def test_absent_logging_obj_still_calls_hook_and_skips_flush(): + """Without a logging obj the flush is skipped (no crash) but + ``post_call_failure_hook`` still fires. Byte-equivalent to stock behavior for + that branch; its value is as a mutation-killer for the ``is not None`` guard.""" + order: list = [] + await _call_block(None, order) + + assert order == ["post_call_failure_hook"], order diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py new file mode 100644 index 00000000000..24e6d2de10d --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py @@ -0,0 +1,301 @@ +"""Tests for MCP guardrail evaluations reaching the Guardrails Monitor. + +MCP tool calls run their guardrails against a throwaway LLM-shaped dict built by +``ProxyLogging._convert_mcp_to_llm_format``, not against the dict the tool call is +logged from. ``@log_guardrail_information`` therefore appends +``standard_logging_guardrail_information`` to that throwaway dict's metadata +bucket, where ``get_standard_logging_object_payload`` never sees it, so the +Guardrails Monitor reported zero evaluations and zero blocks for MCP traffic. + +``pre_call_tool_check`` and ``_create_during_hook_task`` now take the request's +``litellm_logging_obj`` and bridge those records onto it. These tests pin both the +seeding (which unified guardrails consume off ``data["litellm_logging_obj"]``) and +the bridge (which native guardrails depend on), including on the block path. +""" + +import asyncio +import datetime +from typing import Any +from unittest import mock + +import pytest + +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy._experimental.mcp_server import mcp_server_manager as MOD + + +class _FakeLoggingObj: + """Minimal stand-in for ``LiteLLMLoggingObj``. + + ``_sync_guardrail_info_to_logging_obj`` reads exactly these two attributes, + and the spend-log payload is built from ``litellm_params["metadata"]``, so a + real ``Logging`` instance would add setup cost without adding coverage. + """ + + def __init__(self) -> None: + self.litellm_params: dict[str, Any] = {"metadata": {}} + self.model_call_details: dict[str, Any] = {"litellm_params": self.litellm_params} + + @property + def recorded_guardrails(self) -> list: + return self.litellm_params["metadata"].get("standard_logging_guardrail_information", []) + + +def _bare_manager() -> MOD.MCPServerManager: + """An ``MCPServerManager`` without running ``__init__``. + + The authorization/validation helpers on the path are stubbed out so the test + reaches the guardrail hooks; they have their own coverage elsewhere. + """ + mgr = MOD.MCPServerManager.__new__(MOD.MCPServerManager) + mgr.check_allowed_or_banned_tools = lambda name, server: True + mgr.validate_allowed_params = lambda tool_name, arguments, server: None + + async def _ok(*_args, **_kwargs) -> None: + return None + + mgr.check_tool_permission_for_key_team = _ok + return mgr + + +def _fake_proxy_logging(capture: dict, *, guardrail_effect=None): + """A ``proxy_logging_obj`` double whose hooks capture the data they receive. + + ``guardrail_effect`` stands in for a guardrail: it is handed the synthetic + request dict so it can append a guardrail record (and optionally raise, the + way a blocking guardrail does). + """ + plo = mock.MagicMock() + plo._create_mcp_request_object_from_kwargs.return_value = mock.MagicMock() + # Mirror the real conversion's metadata bucket so a test can prove it survives. + plo._convert_mcp_to_llm_format.side_effect = lambda *_a, **_k: { + "metadata": {"headers": {"x-forwarded-for": "1.2.3.4"}} + } + + async def _hook(*, user_api_key_dict, data, call_type) -> None: + del user_api_key_dict # captured shape is what matters, not the auth double + capture["data"] = data + capture["call_type"] = call_type + if guardrail_effect is not None: + guardrail_effect(data) + + plo.pre_call_hook.side_effect = _hook + plo.during_call_hook.side_effect = _hook + return plo + + +def _record_guardrail(status: str = "success"): + """Write a guardrail record the way ``@log_guardrail_information`` does.""" + + def _effect(data: dict) -> None: + data.setdefault("metadata", {}).setdefault("standard_logging_guardrail_information", []).append( + {"guardrail_name": "test-guardrail", "guardrail_status": status} + ) + + return _effect + + +def _blocking_guardrail(): + record = _record_guardrail(status="guardrail_intervened") + + def _effect(data: dict) -> None: + record(data) + raise GuardrailRaisedException(guardrail_name="test-guardrail", message="blocked") + + return _effect + + +async def _run_pre_call(mgr, plo, logging_obj) -> dict: + return await mgr.pre_call_tool_check( + name="t", + arguments={}, + server_name="s", + user_api_key_auth=None, + proxy_logging_obj=plo, + server=mock.MagicMock(), + raw_headers={}, + litellm_logging_obj=logging_obj, + ) + + +@pytest.mark.asyncio +async def test_pre_call_seeds_request_logging_obj_for_unified_guardrails(): + """Unified guardrails read ``data["litellm_logging_obj"]`` and pass it into + ``apply_guardrail``, whose ``@log_guardrail_information`` wrapper bridges the + evaluation onto that logger itself. Drop the seed and that path records + nothing.""" + capture: dict = {} + logging_obj = _FakeLoggingObj() + await _run_pre_call(_bare_manager(), _fake_proxy_logging(capture), logging_obj) + + assert capture["data"]["litellm_logging_obj"] is logging_obj + + +@pytest.mark.asyncio +async def test_pre_call_keeps_synthetic_request_headers_metadata(): + """The seed must not clobber the metadata bucket ``_convert_mcp_to_llm_format`` + builds: guardrails such as ``MCPJWTSigner`` read ``metadata["headers"]`` off + it.""" + capture: dict = {} + await _run_pre_call(_bare_manager(), _fake_proxy_logging(capture), _FakeLoggingObj()) + + assert capture["data"]["metadata"]["headers"] == {"x-forwarded-for": "1.2.3.4"} + + +@pytest.mark.asyncio +async def test_pre_call_bridges_allowed_evaluation_onto_request_logger(): + """An allowed ``pre_mcp_call`` evaluation must land on the request logger, which + is what the monitor's "Total Evaluations" counts.""" + capture: dict = {} + logging_obj = _FakeLoggingObj() + plo = _fake_proxy_logging(capture, guardrail_effect=_record_guardrail()) + + await _run_pre_call(_bare_manager(), plo, logging_obj) + + assert logging_obj.recorded_guardrails == [{"guardrail_name": "test-guardrail", "guardrail_status": "success"}] + + +@pytest.mark.asyncio +async def test_pre_call_bridges_blocked_evaluation_before_reraising(): + """A block raises straight out of ``pre_call_tool_check``, and the failure + spend-log row that "Total Blocked" counts is built from this logger further up + the stack. So the record has to be attached before the exception leaves the + frame -- hence the bridge lives in a ``finally``.""" + capture: dict = {} + logging_obj = _FakeLoggingObj() + plo = _fake_proxy_logging(capture, guardrail_effect=_blocking_guardrail()) + + with pytest.raises(GuardrailRaisedException): + await _run_pre_call(_bare_manager(), plo, logging_obj) + + assert logging_obj.recorded_guardrails == [ + {"guardrail_name": "test-guardrail", "guardrail_status": "guardrail_intervened"} + ] + + +@pytest.mark.asyncio +async def test_pre_call_without_logging_obj_is_unchanged(): + """Callers that thread no logger are unaffected: the seed is an explicit + ``None`` (which every consumer reads via ``.get``) and nothing is bridged. + Guards against the bridge assuming a logger exists.""" + capture: dict = {} + plo = _fake_proxy_logging(capture, guardrail_effect=_record_guardrail()) + mgr = _bare_manager() + + result = await mgr.pre_call_tool_check( + name="t", + arguments={}, + server_name="s", + user_api_key_auth=None, + proxy_logging_obj=plo, + server=mock.MagicMock(), + raw_headers={}, + ) + + assert result == {} + assert capture["data"]["litellm_logging_obj"] is None + + +@pytest.mark.asyncio +async def test_during_hook_seeds_and_bridges_onto_request_logger(): + """``during_mcp_call`` evaluations need the same treatment. The task is awaited + before the tool call's success logging runs, so the record is serialized with + that call.""" + capture: dict = {} + logging_obj = _FakeLoggingObj() + plo = _fake_proxy_logging(capture, guardrail_effect=_record_guardrail()) + + await _bare_manager()._create_during_hook_task( + name="t", + arguments={}, + server_name_from_prefix="s", + user_api_key_auth=None, + proxy_logging_obj=plo, + start_time=datetime.datetime(2026, 7, 14), + litellm_logging_obj=logging_obj, + ) + + assert capture["data"]["litellm_logging_obj"] is logging_obj + assert logging_obj.recorded_guardrails == [{"guardrail_name": "test-guardrail", "guardrail_status": "success"}] + + +@pytest.mark.asyncio +async def test_during_hook_bridges_even_when_hook_raises(): + """A during-call guardrail block must still be recorded before the task's + exception propagates to the ``asyncio.gather`` in ``call_tool``.""" + capture: dict = {} + logging_obj = _FakeLoggingObj() + plo = _fake_proxy_logging(capture, guardrail_effect=_blocking_guardrail()) + + task = _bare_manager()._create_during_hook_task( + name="t", + arguments={}, + server_name_from_prefix="s", + user_api_key_auth=None, + proxy_logging_obj=plo, + start_time=datetime.datetime(2026, 7, 14), + litellm_logging_obj=logging_obj, + ) + with pytest.raises(GuardrailRaisedException): + await task + + assert logging_obj.recorded_guardrails == [ + {"guardrail_name": "test-guardrail", "guardrail_status": "guardrail_intervened"} + ] + + +@pytest.mark.asyncio +async def test_bridge_failure_does_not_mask_a_guardrail_block(): + """Recording is best-effort bookkeeping. If the bridge itself raises, the guardrail's + block must still be what the caller sees, not a bookkeeping error. + + The bridge is forced to fail by making the logger's ``model_call_details`` raise, and + the swallow is asserted (not just the surviving exception type) so the test cannot go + vacuous if a refactor stops the bridge from touching that attribute. + """ + capture: dict = {} + plo = _fake_proxy_logging(capture, guardrail_effect=_blocking_guardrail()) + + broken_logging_obj = mock.MagicMock() + type(broken_logging_obj).model_call_details = mock.PropertyMock(side_effect=RuntimeError("boom")) + + with mock.patch.object(MOD.verbose_logger, "warning") as warn: + with pytest.raises(GuardrailRaisedException): + await _run_pre_call(_bare_manager(), plo, broken_logging_obj) + + assert warn.call_count == 1, "the bridge did not actually fail, so this test proves nothing" + assert "boom" in str(warn.call_args) + + +@pytest.mark.asyncio +async def test_call_tool_threads_logging_obj_into_both_hooks(): + """``call_tool`` is the single entry point every MCP dispatch route funnels + through, so it must hand the logger to both guardrail hook sites.""" + mgr = _bare_manager() + logging_obj = _FakeLoggingObj() + seen: dict = {} + + async def _fake_pre_call_tool_check(**kwargs): + seen["pre_call"] = kwargs.get("litellm_logging_obj") + return {} + + def _fake_during_hook_task(**kwargs): + seen["during_call"] = kwargs.get("litellm_logging_obj") + return asyncio.get_running_loop().create_future() + + mgr.pre_call_tool_check = _fake_pre_call_tool_check + mgr._create_during_hook_task = _fake_during_hook_task + mgr._resolve_mcp_server_for_tool_call = lambda server_name, name: mock.MagicMock(spec_path=None) + mgr._resolve_oauth2_headers_for_tool_call = mock.AsyncMock(return_value=None) + mgr._call_regular_mcp_tool = mock.AsyncMock(return_value=mock.MagicMock()) + + with mock.patch.object(MOD, "_resolve_byok_mcp_auth_header", mock.AsyncMock(return_value=None)): + await mgr.call_tool( + server_name="s", + name="t", + arguments={}, + proxy_logging_obj=mock.MagicMock(), + litellm_logging_obj=logging_obj, + ) + + assert seen == {"pre_call": logging_obj, "during_call": logging_obj} diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index 69d90a8b59b..0a427df0cb7 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -125,6 +125,45 @@ class TestBlockedResponseUsage: mock_logging.post_call_failure_hook.assert_awaited_once() +class TestProxyExceptionPassthrough: + @pytest.mark.asyncio + async def test_anthropic_response_reraises_proxy_exception_unwrapped(self): + """A 400 ProxyException from request validation must surface as-is, + not be re-wrapped into a code-500 ProxyException.""" + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import ProxyErrorTypes, ProxyException + + exc = ProxyException( + message="Invalid type for 'metadata': expected an object, but got a string instead.", + type=ProxyErrorTypes.bad_request_error, + param="metadata", + code=400, + ) + + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})), + patch.object( + ep.ProxyBaseLLMRequestProcessing, + "base_process_llm_request", + new=AsyncMock(side_effect=exc), + ), + patch.object(proxy_server, "proxy_logging_obj") as mock_logging, + ): + mock_logging.post_call_failure_hook = AsyncMock() + with pytest.raises(ProxyException) as exc_info: + await ep.anthropic_response( + fastapi_response=MagicMock(), + request=MagicMock(), + user_api_key_dict=MagicMock(), + ) + + assert exc_info.value is exc + assert exc_info.value.code == "400" + assert exc_info.value.param == "metadata" + mock_logging.post_call_failure_hook.assert_awaited_once() + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 28eda6633e8..270f3eca0f9 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2,6 +2,7 @@ import asyncio import json import os import sys +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch sys.path.insert( @@ -51,9 +52,22 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache -from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL +from litellm.constants import ( + DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + REGISTRY_ERROR_NEGATIVE_CACHE_TTL, + TAG_REGISTRY_MAX_SIZE, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import ( + END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, + TAG_REGISTRY_OVERFLOW_SENTINEL, + UserApiKeyCache, + end_user_cache_key, + end_user_restricted_registry_cache_key, + tag_cache_key, + tag_registry_cache_key, +) from litellm.utils import get_utc_datetime @@ -2075,22 +2089,342 @@ async def test_get_tag_objects_batch(): assert tag_objects["uncached-2"].spend == 40.0 assert tag_objects["uncached-3"].spend == 50.0 - # Verify DB was called ONCE with all 3 uncached tags - mock_prisma.db.litellm_tagtable.find_many.assert_called_once() - call_args = mock_prisma.db.litellm_tagtable.find_many.call_args - assert call_args.kwargs["where"]["tag_name"]["in"] == [ + # Verify the DB saw exactly the registry query plus ONE batch query for all 3 uncached tags + assert mock_prisma.db.litellm_tagtable.find_many.call_count == 2 + registry_call, batch_call = mock_prisma.db.litellm_tagtable.find_many.call_args_list + assert "where" not in registry_call.kwargs + assert batch_call.kwargs["where"]["tag_name"]["in"] == [ "uncached-1", "uncached-2", "uncached-3", ] - # Verify uncached tags were cached after fetching - assert mock_cache.async_set_cache.call_count == 3 + # Verify uncached tags were cached after fetching, alongside the tag-name registry cache_calls = mock_cache.async_set_cache.call_args_list cached_keys = [call.kwargs["key"] for call in cache_calls] - assert "tag:uncached-1" in cached_keys - assert "tag:uncached-2" in cached_keys - assert "tag:uncached-3" in cached_keys + assert sorted(cached_keys) == [ + "tag:uncached-1", + "tag:uncached-2", + "tag:uncached-3", + "tag_registry", + ] + # Every write is TTL-bounded; an unbounded tag entry would outlive budget updates. + assert all("ttl" in call.kwargs for call in cache_calls) + + +class _TtlRecordingCache(UserApiKeyCache): + """A real cache that also records the ttl each write carried, so tests can catch unbounded entries.""" + + def __init__(self): + super().__init__() + self.writes = [] + + async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): + self.writes.append((key, kwargs.get("ttl"))) + return await super().async_set_cache(key, value, local_only=local_only, **kwargs) + + +def _tag_registry_row(tag_name: str): + """A row as the names-only registry query sees it: only ``tag_name`` is read off it.""" + return SimpleNamespace(tag_name=tag_name) + + +def _tag_db_row(tag_name: str, max_budget=None): + row = MagicMock() + row.tag_name = tag_name + budget = None if max_budget is None else {"max_budget": max_budget} + row.dict = MagicMock( + return_value={ + "tag_name": tag_name, + "spend": 0.0, + "models": [], + "litellm_budget_table": budget, + } + ) + return row + + +def _registry_calls(find_many): + return [call for call in find_many.call_args_list if "where" not in call.kwargs] + + +def _batch_calls(find_many): + return [call for call in find_many.call_args_list if "where" in call.kwargs] + + +@pytest.mark.asyncio +async def test_get_tag_objects_batch_never_queries_db_for_unregistered_tags(): + """ + Regression: a request tag with no LiteLLM_TagTable row must not cost a DB read per request. + + Cost-attribution tags are free-form, so most carry no tag row. Before the cached name + registry, every request carrying one ran its own Postgres find_many, forever, which is what + saturated a customer's Prisma pool. + """ + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock( + return_value=[_tag_registry_row("some-other-tag")] + ) + cache = UserApiKeyCache() + + first = await get_tag_objects_batch( + tag_names=["unregistered-tag"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert first == {} + + # The only query is the names-only registry fetch; the tag itself is never looked up. + mock_prisma.db.litellm_tagtable.find_many.assert_called_once_with( + take=TAG_REGISTRY_MAX_SIZE + 1 + ) + + second = await get_tag_objects_batch( + tag_names=["unregistered-tag"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert second == {} + assert mock_prisma.db.litellm_tagtable.find_many.call_count == 1 + + +@pytest.mark.asyncio +async def test_get_tag_objects_batch_fetches_only_registered_uncached_tags(): + """Cached tags skip the DB, registered ones are batch-fetched, unregistered ones are dropped.""" + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + cache = UserApiKeyCache() + await cache.async_set_cache( + key=tag_cache_key("cached-tag"), + value=LiteLLM_TagTable(tag_name="cached-tag", spend=7.0, models=[]), + model_type=LiteLLM_TagTable, + ) + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + return [_tag_registry_row("cached-tag"), _tag_registry_row("registered-tag")] + requested = kwargs["where"]["tag_name"]["in"] + return [_tag_db_row(name) for name in requested if name == "registered-tag"] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + + tag_objects = await get_tag_objects_batch( + tag_names=["cached-tag", "registered-tag", "unregistered-tag"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + + assert sorted(tag_objects) == ["cached-tag", "registered-tag"] + assert tag_objects["cached-tag"].spend == 7.0 + + batch_calls = _batch_calls(mock_prisma.db.litellm_tagtable.find_many) + assert len(batch_calls) == 1 + assert batch_calls[0].kwargs["where"]["tag_name"]["in"] == ["registered-tag"] + + +@pytest.mark.asyncio +async def test_get_tag_objects_batch_caches_empty_registry(): + """An empty tag table is a valid registry answer and must be cached, not re-queried.""" + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(return_value=[]) + cache = UserApiKeyCache() + + assert ( + await get_tag_objects_batch( + tag_names=["tag-a", "tag-b"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + == {} + ) + # "No tags registered" is a cached answer, not a cache miss (which would be None). + cached_registry = await cache.async_get_cache(key=tag_registry_cache_key()) + assert cached_registry is not None + assert tuple(cached_registry) == () + + assert ( + await get_tag_objects_batch( + tag_names=["tag-a", "tag-b"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + == {} + ) + assert mock_prisma.db.litellm_tagtable.find_many.call_count == 1 + + +@pytest.mark.asyncio +async def test_get_tag_objects_batch_registry_db_error_negative_caches_and_keeps_per_tag_fetch(): + """ + A degraded database must not be re-asked for the registry on every request. + + Without the negative cache the failing scan re-runs per request on top of the per-tag fallback + it triggers, doubling load exactly when Postgres is least able to take it. Tag budgets keep + being enforced through the per-tag path throughout, and the registry is retried once the + negative entry expires. + """ + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + raise Exception("registry query failed") + requested = kwargs["where"]["tag_name"]["in"] + return [_tag_db_row(name) for name in requested] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + cache = _TtlRecordingCache() + + first = await get_tag_objects_batch( + tag_names=["tag-a"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert list(first) == ["tag-a"] + assert await cache.async_get_cache(key=tag_registry_cache_key()) == TAG_REGISTRY_OVERFLOW_SENTINEL + assert (tag_registry_cache_key(), REGISTRY_ERROR_NEGATIVE_CACHE_TTL) in cache.writes + + second = await get_tag_objects_batch( + tag_names=["tag-b"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert list(second) == ["tag-b"] + assert len(_registry_calls(mock_prisma.db.litellm_tagtable.find_many)) == 1 + + # The window closing (here: the entry expiring) puts the registry back in play. + await cache.async_delete_cache(key=tag_registry_cache_key()) + third = await get_tag_objects_batch( + tag_names=["tag-c"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert list(third) == ["tag-c"] + assert len(_registry_calls(mock_prisma.db.litellm_tagtable.find_many)) == 2 + + +@pytest.mark.asyncio +async def test_tag_registry_load_is_single_flighted_across_concurrent_requests(): + """ + A cold registry under load must run one scan, not one per in-flight request. + + The registry query is an unindexed table scan; a TTL expiry on a busy worker would otherwise + fan out into as many identical scans as there are concurrent requests. + """ + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + await asyncio.sleep(0) + return [_tag_registry_row("registered-tag")] + return [_tag_db_row(name) for name in kwargs["where"]["tag_name"]["in"]] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + cache = UserApiKeyCache() + + results = await asyncio.gather( + *( + get_tag_objects_batch( + tag_names=["registered-tag"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + for _ in range(8) + ) + ) + + assert all(list(result) == ["registered-tag"] for result in results) + assert len(_registry_calls(mock_prisma.db.litellm_tagtable.find_many)) == 1 + + +@pytest.mark.asyncio +async def test_get_tag_objects_batch_oversized_registry_falls_back_and_stops_refetching(): + """Past the cap the registry is unusable: keep the old per-tag path, but stop rebuilding it.""" + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + oversized = [ + _tag_registry_row(f"tag-{index}") for index in range(TAG_REGISTRY_MAX_SIZE + 1) + ] + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + return oversized + return [_tag_db_row(name) for name in kwargs["where"]["tag_name"]["in"]] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + cache = UserApiKeyCache() + + first = await get_tag_objects_batch( + tag_names=["tag-a"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert list(first) == ["tag-a"] + assert ( + await cache.async_get_cache(key=tag_registry_cache_key()) + == TAG_REGISTRY_OVERFLOW_SENTINEL + ) + + second = await get_tag_objects_batch( + tag_names=["tag-b"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert list(second) == ["tag-b"] + + find_many = mock_prisma.db.litellm_tagtable.find_many + assert len(_registry_calls(find_many)) == 1 + assert [call.kwargs["where"]["tag_name"]["in"] for call in _batch_calls(find_many)] == [ + ["tag-a"], + ["tag-b"], + ] + + +@pytest.mark.asyncio +async def test_tag_max_budget_check_still_enforces_registered_tag_over_budget(): + """The registry filter must not swallow a real tag: an over-budget tag still raises.""" + from litellm.proxy.utils import ProxyLogging + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + return [_tag_registry_row("paid-tag")] + return [ + _tag_db_row(name, max_budget=1.0) + for name in kwargs["where"]["tag_name"]["in"] + ] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): + if counter_key == "spend:tag:paid-tag": + return 1.5 + return fallback_spend + + with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _tag_max_budget_check( + request_body={"metadata": {"tags": ["paid-tag", "unregistered-tag"]}}, + prisma_client=mock_prisma, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(token="test-token"), + ) + assert exc_info.value.current_cost == 1.5 + assert exc_info.value.entity_id == "paid-tag" + + # The unregistered tag alongside it never reached the DB. + batch_calls = _batch_calls(mock_prisma.db.litellm_tagtable.find_many) + assert [call.kwargs["where"]["tag_name"]["in"] for call in batch_calls] == [["paid-tag"]] @pytest.mark.asyncio @@ -5390,6 +5724,400 @@ async def test_get_end_user_object_db_fetch_returns_validated_end_user(): assert result.spend == 3.0 +def _end_user_registry_row(user_id: str): + """A row as the restricted-id registry query sees it: only ``user_id`` is read off it.""" + return SimpleNamespace(user_id=user_id) + + +def _end_user_db_row(user_id: str, **fields): + row = MagicMock() + row.user_id = user_id + row.dict = lambda: {"user_id": user_id, "blocked": False, "spend": 0.0, **fields} + return row + + +_RESTRICTED_END_USER_WHERE = { + "OR": [ + {"blocked": True}, + {"budget_id": {"not": None}}, + {"allowed_model_region": {"not": None}}, + {"default_model": {"not": None}}, + {"object_permission_id": {"not": None}}, + ] +} + + +@pytest.fixture +def end_user_registry_skip_enabled(monkeypatch): + """Both bypass gates off: the default deployment, and the only state the registry skip runs in.""" + monkeypatch.setattr(litellm, "max_end_user_budget_id", None) + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", False) + + +@pytest.mark.asyncio +async def test_get_end_user_object_never_queries_db_for_unrestricted_end_users( + end_user_registry_skip_enabled, +): + """ + Regression: an end user carrying no restriction must not cost a DB read per request. + + Spend tracking auto-creates a row for every distinct caller-supplied ``user`` id with every + restriction field null, so a high-cardinality deployment misses the per-pod cache on virtually + every request. Before the cached registry each miss ran its own Postgres find_unique, twice per + request, and under Prisma pool contention those queued for minutes inside user_api_key_auth. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[_end_user_registry_row("eu-blocked")]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1")) + cache = UserApiKeyCache() + + assert ( + await get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + is None + ) + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + mock_prisma.db.litellm_endusertable.find_many.assert_awaited_once() + registry_call = mock_prisma.db.litellm_endusertable.find_many.call_args + assert registry_call.kwargs["take"] == END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1 + # Every field the callers of get_end_user_object consume has to be in this predicate, or an id + # the registry calls unrestricted would silently lose a restriction that is actually enforced. + assert registry_call.kwargs["where"] == _RESTRICTED_END_USER_WHERE + + mock_prisma.db.litellm_endusertable.find_many.reset_mock() + assert ( + await get_end_user_object( + end_user_id="eu-anon-2", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + is None + ) + # A second, different unknown id inside the TTL costs nothing: no rebuild, no row fetch. + mock_prisma.db.litellm_endusertable.find_many.assert_not_awaited() + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_end_user_object_still_fetches_restricted_end_user(end_user_registry_skip_enabled): + """An id in the registry keeps today's path: fetched, TTL-bounded in cache, then served cached.""" + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[_end_user_registry_row("eu-blocked")]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock( + return_value=_end_user_db_row("eu-blocked", blocked=True) + ) + cache = _TtlRecordingCache() + + blocked = await get_end_user_object( + end_user_id="eu-blocked", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert isinstance(blocked, LiteLLM_EndUserTable) + assert blocked.blocked is True + mock_prisma.db.litellm_endusertable.find_unique.assert_awaited_once() + # Without a ttl the Redis entry never expires, so a later unblock would never be picked up. + assert (end_user_cache_key("eu-blocked"), DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL) in cache.writes + + mock_prisma.db.litellm_endusertable.find_unique.reset_mock() + again = await get_end_user_object( + end_user_id="eu-blocked", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert again is not None and again.blocked is True + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_end_user_object_caches_empty_restricted_registry(end_user_registry_skip_enabled): + """No restricted end users at all is a valid answer and must be cached, not re-queried.""" + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1")) + cache = UserApiKeyCache() + + assert ( + await get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + is None + ) + # "Nobody is restricted" is a cached answer, not a cache miss (which would read back as None). + cached_registry = await cache.async_get_cache(key=end_user_restricted_registry_cache_key()) + assert cached_registry is not None + assert tuple(cached_registry) == () + + assert ( + await get_end_user_object( + end_user_id="eu-anon-2", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + is None + ) + mock_prisma.db.litellm_endusertable.find_many.assert_awaited_once() + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_end_user_object_registry_db_error_negative_caches_and_keeps_per_id_fetch( + end_user_registry_skip_enabled, +): + """ + A degraded database must not be re-asked for the registry on every request. + + Restrictions keep being enforced through the per-id fetch, exactly as before the registry + existed, but the failing scan is suppressed for the negative-cache window instead of running + again on every request on top of that fetch. It is retried once the window closes. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(side_effect=Exception("registry query failed")) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock( + side_effect=lambda **kwargs: _end_user_db_row(kwargs["where"]["user_id"], blocked=True) + ) + cache = _TtlRecordingCache() + + first = await get_end_user_object( + end_user_id="eu-blocked-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert first is not None and first.blocked is True + assert ( + await cache.async_get_cache(key=end_user_restricted_registry_cache_key()) + == END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL + ) + assert (end_user_restricted_registry_cache_key(), REGISTRY_ERROR_NEGATIVE_CACHE_TTL) in cache.writes + + second = await get_end_user_object( + end_user_id="eu-blocked-2", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert second is not None and second.blocked is True + mock_prisma.db.litellm_endusertable.find_many.assert_awaited_once() + + # The window closing (here: the entry expiring) puts the registry back in play. + await cache.async_delete_cache(key=end_user_restricted_registry_cache_key()) + third = await get_end_user_object( + end_user_id="eu-blocked-3", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert third is not None and third.blocked is True + assert mock_prisma.db.litellm_endusertable.find_many.await_count == 2 + + +@pytest.mark.asyncio +async def test_registry_db_error_is_logged_at_warning(end_user_registry_skip_enabled): + """ + A registry that stops loading is a silent enforcement degradation, so seeing it must not + require debug logging: per-id lookups still enforce restrictions, but an operator has no other + signal that the database is failing the scan and that every request is paying for it. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(side_effect=Exception("registry query failed")) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-1", blocked=True)) + + with patch("litellm.proxy.auth.auth_checks.verbose_proxy_logger") as mock_logger: + await get_end_user_object( + end_user_id="eu-1", + prisma_client=mock_prisma, + user_api_key_cache=UserApiKeyCache(), + ) + + warnings = [_rendered_log_message(call) for call in mock_logger.warning.call_args_list] + assert any( + end_user_restricted_registry_cache_key() in message and "registry query failed" in message + for message in warnings + ) + + +@pytest.mark.asyncio +async def test_end_user_registry_load_is_single_flighted_across_concurrent_requests( + end_user_registry_skip_enabled, +): + """ + A cold registry under load must run one scan, not one per in-flight request. + + The registry query is an unindexed scan over the end-user table, which for the deployments this + exists for holds hundreds of thousands of rows; a TTL expiry on a busy worker would otherwise + fan it out across every concurrent request. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + async def fake_find_many(**kwargs): + await asyncio.sleep(0) + return [_end_user_registry_row("eu-blocked")] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(side_effect=fake_find_many) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1")) + cache = UserApiKeyCache() + + results = await asyncio.gather( + *( + get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + for _ in range(8) + ) + ) + + assert all(result is None for result in results) + mock_prisma.db.litellm_endusertable.find_many.assert_awaited_once() + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_end_user_object_oversized_registry_falls_back_and_stops_refetching( + end_user_registry_skip_enabled, +): + """Past the cap the registry is unusable: keep the per-id path, but stop rebuilding the set.""" + from litellm.proxy.auth.auth_checks import get_end_user_object + + oversized = [_end_user_registry_row(f"eu-{index}") for index in range(END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1)] + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=oversized) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock( + side_effect=lambda **kwargs: _end_user_db_row(kwargs["where"]["user_id"], blocked=True) + ) + cache = UserApiKeyCache() + + first = await get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert first is not None and first.blocked is True + assert ( + await cache.async_get_cache(key=end_user_restricted_registry_cache_key()) + == END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL + ) + + second = await get_end_user_object( + end_user_id="eu-anon-2", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert second is not None and second.blocked is True + mock_prisma.db.litellm_endusertable.find_many.assert_awaited_once() + assert mock_prisma.db.litellm_endusertable.find_unique.await_count == 2 + + +@pytest.mark.asyncio +async def test_get_end_user_object_default_budget_gate_keeps_fetching_unrestricted_end_users(monkeypatch): + """ + With ``max_end_user_budget_id`` set, an existing unrestricted row is not equivalent to a missing + one: the default budget is grafted onto whatever row exists and is then enforced, so the skip + has to stay off entirely. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + monkeypatch.setattr(litellm, "max_end_user_budget_id", "default-eu-budget") + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", False) + + budget_row = MagicMock() + budget_row.dict = lambda: {"budget_id": "default-eu-budget", "max_budget": 25.0} + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1")) + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(return_value=budget_row) + + result = await get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=UserApiKeyCache(), + ) + + assert result is not None + assert result.litellm_budget_table is not None + assert result.litellm_budget_table.max_budget == 25.0 + mock_prisma.db.litellm_endusertable.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_end_user_object_token_budget_gate_keeps_fetching_unrestricted_end_users( + end_user_registry_skip_enabled, +): + """ + A token-supplied end-user budget is enforced against the row's recorded spend, so the row has + to be loaded even though nothing on it is restricted. + + A ``user_custom_auth`` callable can set ``end_user_max_budget`` on the returned token for an + end user whose row carries no budget of its own, which keeps it out of the registry. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock( + return_value=_end_user_db_row("eu-anon-1", spend=100.0) + ) + cache = UserApiKeyCache() + + result = await get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + token_end_user_max_budget=50.0, + ) + + assert result is not None + assert result.spend == 100.0 + mock_prisma.db.litellm_endusertable.find_unique.assert_awaited_once() + mock_prisma.db.litellm_endusertable.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_end_user_id_validation_gate_still_resolves_unrestricted_end_users(monkeypatch): + """ + With ``validate_end_user_id_in_db`` on, existence itself is the answer, so the skip stays off. + + Skipping here would turn every unrestricted customer into an unknown id and drop it from the + request, which for a deployment with no default budget means the id silently stops being tracked. + """ + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + monkeypatch.setattr(litellm, "max_end_user_budget_id", None) + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-known-1")) + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + resolved = await resolve_and_validate_end_user_id( + raw_end_user_id="eu-known-1", + prisma_client=mock_prisma, + user_api_key_cache=UserApiKeyCache(), + ) + + assert resolved == "eu-known-1" + mock_prisma.db.litellm_endusertable.find_many.assert_not_awaited() + + @pytest.mark.asyncio async def test_get_team_membership_db_fetch_returns_validated_membership(): from litellm.proxy._types import LiteLLM_TeamMembership diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index 3203878a1e0..cf1f665ad21 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -173,6 +173,55 @@ async def test_custom_auth_defers_end_user_budget_to_common_checks_when_enabled( mock_check.assert_not_awaited() +@pytest.mark.asyncio +async def test_custom_auth_token_budget_still_loads_and_caches_unrestricted_end_user(): + """ + A token-supplied end_user_max_budget must leave the end-user row in cache. + + Custom auth can set that budget for a customer whose own row carries no budget, block, region + or permission, which keeps the row out of the cached restricted-id registry that lets auth skip + the read. The end-user spend counter seeds from this cache entry, so skipping the read would + cold-start the counter at 0 and under-count a customer who has already spent 100. + """ + from unittest.mock import MagicMock + + from litellm.proxy.auth.user_api_key_auth import _lookup_end_user_and_apply_budget + from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + end_user_cache_key, + ) + + end_user_row = MagicMock() + end_user_row.user_id = "customer-1" + end_user_row.dict = lambda: { + "user_id": "customer-1", + "blocked": False, + "spend": 100.0, + } + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=end_user_row) + cache = UserApiKeyCache() + + _, end_user_object = await _lookup_end_user_and_apply_budget( + valid_token=UserAPIKeyAuth( + token="test_token", + end_user_id="customer-1", + end_user_max_budget=50.0, + ), + route="/v1/chat/completions", + parent_otel_span=None, + prisma_client=mock_prisma, + user_api_key_cache=cache, + proxy_logging_obj=MagicMock(), + ) + + assert end_user_object is not None + assert end_user_object.spend == 100.0 + assert await cache.async_get_cache(key=end_user_cache_key("customer-1")) is not None + + def test_update_valid_token_does_not_override_custom_auth_values_with_none(): """ Greptile feedback: if custom auth sets end_user_model_max_budget on the token, diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 129813d806c..ab7e3d9701c 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -2,6 +2,7 @@ import asyncio import json import os import sys +from contextlib import contextmanager from datetime import datetime, timedelta from types import SimpleNamespace from unittest.mock import ANY, AsyncMock, MagicMock, patch @@ -3739,6 +3740,140 @@ async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): setattr(_proxy_server_mod, k, v) +def _unrestricted_end_user_prisma(spend: float): + """Prisma stand-in where "customer-1" exists but restricts nothing: no row matches the + restricted-registry query, and the row itself carries only spend.""" + end_user_row = MagicMock() + end_user_row.user_id = "customer-1" + end_user_row.dict = lambda: {"user_id": "customer-1", "blocked": False, "spend": spend} + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=end_user_row) + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + return mock_prisma + + +@contextmanager +def _custom_auth_end_user_world(mock_prisma): + """The proxy globals a custom-auth deployment running the centralized gate reads, with cold + spend counters. Real caches, so the end user's spend reaches the counter the way it does in + production: through the cache entry get_end_user_object writes.""" + import litellm.proxy.proxy_server as _proxy_server_mod + + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + key_cache = UserApiKeyCache() + attrs = { + **_proxy_attrs_for_centralized_checks(user_custom_auth=AsyncMock(), flag=True), + "prisma_client": mock_prisma, + "user_api_key_cache": key_cache, + "spend_counter_cache": DualCache(), + "proxy_logging_obj": ProxyLogging(user_api_key_cache=key_cache), + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + yield + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +def _chat_request(): + from fastapi import Request + from starlette.datastructures import URL + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + return request + + +@pytest.mark.asyncio +async def test_centralized_checks_enforce_token_end_user_budget_against_row_spend(): + """ + Regression: a token-supplied end-user budget must still be checked against the end user's + recorded spend. + + A user_custom_auth callable can set end_user_max_budget on the token for an end user whose own + row carries no budget, which keeps that row out of the restricted-id registry. Auth must still + load it, because the reservation counter cold-starts from the spend on the loaded row; skipping + the load admits a customer who is already double their budget. + """ + mock_prisma = _unrestricted_end_user_prisma(spend=100.0) + token = UserAPIKeyAuth( + api_key="sk-test", + token="hashed-token", + user_id="u1", + end_user_id="customer-1", + end_user_max_budget=50.0, + ) + + with _custom_auth_end_user_world(mock_prisma): + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ), + pytest.raises(litellm.BudgetExceededError) as exc_info, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=_chat_request(), + request_data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + }, + route="/chat/completions", + ) + + assert exc_info.value.max_budget == 50.0 + assert exc_info.value.current_cost == pytest.approx(100.6) + + +@pytest.mark.asyncio +async def test_centralized_checks_skip_end_user_lookup_without_a_token_budget(): + """The companion case: with no token budget an unrestricted end user costs zero row reads.""" + mock_prisma = _unrestricted_end_user_prisma(spend=100.0) + token = UserAPIKeyAuth( + api_key="sk-test", + token="hashed-token", + user_id="u1", + end_user_id="customer-1", + ) + + with _custom_auth_end_user_world(mock_prisma): + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=_chat_request(), + request_data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + }, + route="/chat/completions", + ) + + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + @pytest.mark.asyncio async def test_centralized_common_checks_runs_for_custom_auth_with_flag(): """Custom-auth deployments that opt in via custom_auth_run_common_checks diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 824654dcf66..b720ea16968 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -758,6 +758,30 @@ async def test_create__model_encoded_beats_loadbalancing(harness): harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body, missing_param", + [ + ({"endpoint": "/v1/chat/completions", "completion_window": "24h"}, "input_file_id"), + ({"input_file_id": "file-abc", "completion_window": "24h"}, "endpoint"), + ({"input_file_id": "file-abc", "endpoint": "/v1/chat/completions"}, "completion_window"), + ({}, "input_file_id"), + ], +) +async def test_create__missing_required_param_is_400(harness, body, missing_param): + set_body(harness, body) + + with pytest.raises(ProxyException) as exc_info: + await call_create(harness) + + assert exc_info.value.code == "400" + assert exc_info.value.type == "invalid_request_error" + assert exc_info.value.param == missing_param + assert exc_info.value.message == f"/batches: Missing required parameter: '{missing_param}'." + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + # =========================================================================== # # Team-level batch expiry enforcement (independent of routing). # =========================================================================== # diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_spend_logs.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_spend_logs.py index 79f13a6f703..35fcd3b6cd7 100644 --- a/tests/test_litellm/proxy/management_endpoints/management_v1/test_spend_logs.py +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_spend_logs.py @@ -1,5 +1,4 @@ from datetime import datetime, timezone -from typing import List from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -45,6 +44,7 @@ app.include_router(router) client = TestClient(app) END_USERS_PATH = f"{MANAGEMENT_V1_PREFIX}/spend_logs/end_users" +USERS_PATH = f"{MANAGEMENT_V1_PREFIX}/spend_logs/users" WINDOW = "filter[startTime][gte]=2026-07-23T00:00:00Z&filter[startTime][lte]=2026-07-24T00:00:00Z" @@ -65,7 +65,7 @@ def as_proxy_admin(): app.dependency_overrides.clear() -def _mock_rows(mock_prisma_client, end_users: List[str]) -> AsyncMock: +def _mock_rows(mock_prisma_client, end_users: list[str]) -> AsyncMock: query_raw = AsyncMock(return_value=[{"end_user": eu} for eu in end_users]) mock_prisma_client.db.query_raw = query_raw return query_raw @@ -82,6 +82,11 @@ def _get(query: str = WINDOW): return client.get(f"{END_USERS_PATH}{suffix}", headers={"Authorization": "Bearer k"}) +def _get_users(query: str = WINDOW): + suffix = f"?{query}" if query else "" + return client.get(f"{USERS_PATH}{suffix}", headers={"Authorization": "Bearer k"}) + + def test_returns_the_control_plane_envelope(mock_prisma_client, as_proxy_admin): """`{data, meta, links}` is the contract; a bare list or a legacy `aliases` key is not.""" _mock_rows(mock_prisma_client, ["a", "b"]) @@ -213,7 +218,7 @@ def test_requires_a_time_window(mock_prisma_client, as_proxy_admin, query): def test_rejects_a_malformed_window_as_a_problem_document(mock_prisma_client, as_proxy_admin): _mock_rows(mock_prisma_client, []) - response = _get(f"filter[startTime][gte]=yesterday&filter[startTime][lte]=2026-07-24T00:00:00Z") + response = _get("filter[startTime][gte]=yesterday&filter[startTime][lte]=2026-07-24T00:00:00Z") assert response.status_code == 400 assert response.headers["content-type"].startswith("application/problem+json") @@ -400,6 +405,49 @@ def test_q_placeholder_precedes_the_scan_limit_and_offset(mock_prisma_client, as assert query_raw.call_args.args[5:] == (11, 0) +def test_user_facet_reads_internal_users_from_spend_logs(mock_prisma_client, as_proxy_admin): + query_raw = AsyncMock(return_value=[{"user": "alice@example.com"}, {"user": "user-42"}]) + mock_prisma_client.db.query_raw = query_raw + + response = _get_users() + + assert response.status_code == 200 + assert response.json()["data"] == ["alice@example.com", "user-42"] + sql = query_raw.call_args.args[0] + assert 'SELECT DISTINCT "user"' in sql + assert '"user" IS NOT NULL' in sql + assert "end_user IS NOT NULL" not in sql + + +def test_user_facet_uses_the_same_team_scope_as_request_logs(mock_prisma_client): + query_raw = AsyncMock(return_value=[{"user": "member@example.com"}]) + mock_prisma_client.db.query_raw = query_raw + original = _as_role(LitellmUserRoles.INTERNAL_USER, user_id="team-admin-1") + try: + with patch( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + new=AsyncMock(return_value=["team-a"]), + ): + response = _get_users() + finally: + app.dependency_overrides = original + + assert response.status_code == 200 + assert '("user" = $3 OR team_id = ANY($4::text[]))' in query_raw.call_args.args[0] + assert query_raw.call_args.args[3] == "team-admin-1" + assert query_raw.call_args.args[4] == ["team-a"] + + +def test_user_facet_searches_the_internal_user_value(mock_prisma_client, as_proxy_admin): + query_raw = AsyncMock(return_value=[]) + mock_prisma_client.db.query_raw = query_raw + + _get_users(f"{WINDOW}&q=alice%40example.com") + + assert '"user" ILIKE $3 ESCAPE' in query_raw.call_args.args[0] + assert query_raw.call_args.args[3] == "%alice@example.com%" + + @pytest.mark.parametrize( "role", [ @@ -416,18 +464,19 @@ def test_is_reachable_by_every_role_that_can_open_the_logs_page(role): """ from litellm.proxy.auth.route_checks import RouteChecks - for allowed in ( - LiteLLMRoutes.internal_user_routes.value, - LiteLLMRoutes.internal_user_view_only_routes.value, - ): - assert ("/spend/logs/ui" in allowed) == (END_USERS_PATH in allowed) + for facet_path in (END_USERS_PATH, USERS_PATH): + for allowed in ( + LiteLLMRoutes.internal_user_routes.value, + LiteLLMRoutes.internal_user_view_only_routes.value, + ): + assert ("/spend/logs/ui" in allowed) == (facet_path in allowed) - if role in (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY): - allowed_routes = ( - LiteLLMRoutes.internal_user_routes.value - if role == LitellmUserRoles.INTERNAL_USER - else LiteLLMRoutes.internal_user_view_only_routes.value - ) - assert RouteChecks.check_route_access(route=END_USERS_PATH, allowed_routes=allowed_routes) - else: - assert END_USERS_PATH in LiteLLMRoutes.admin_viewer_routes.value + if role in (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY): + allowed_routes = ( + LiteLLMRoutes.internal_user_routes.value + if role == LitellmUserRoles.INTERNAL_USER + else LiteLLMRoutes.internal_user_view_only_routes.value + ) + assert RouteChecks.check_route_access(route=facet_path, allowed_routes=allowed_routes) + else: + assert facet_path in LiteLLMRoutes.admin_viewer_routes.value diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 5efed8de325..5c163c44cb3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -1,3 +1,4 @@ +from contextlib import contextmanager from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -815,3 +816,129 @@ def test_char_delete_body(mock_prisma_client, mock_user_api_key_auth): "deleted_customers": 2, "message": "Successfully deleted customers with ids: ['c1', 'c2']", } + + +class _RecordingAuthCache: + """Captures the keys an endpoint evicts, so tests assert on cache keys not mock plumbing.""" + + def __init__(self): + self.deleted: list[str] = [] + + async def async_delete_cache(self, key: str) -> None: + self.deleted.append(key) + + +@contextmanager +def _end_user_cache_doubles(): + """Swaps in the auth cache and the cross-worker publisher a customer mutation is expected to hit.""" + recording_cache = _RecordingAuthCache() + mock_publish = AsyncMock() + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", recording_cache), + patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + mock_publish, + ), + ): + yield recording_cache, mock_publish + + +def _published_keys(mock_publish) -> list[str]: + return [call.kwargs["cache_key"] for call in mock_publish.call_args_list] + + +def test_customer_new_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): + """ + A customer created on one worker must be visible to every worker's auth path immediately. + + Auth serves end users cache-first, and the cached restricted-id registry is what decides whether + the row is read at all, so a create that leaves both entries stale means the new customer's + budget or block goes unenforced until the TTL expires. + """ + mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW)) + + with _end_user_cache_doubles() as (recording_cache, mock_publish): + response = client.post( + "/customer/new", + json={"user_id": "c1", "blocked": True}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200, response.text + assert recording_cache.deleted == ["end_user_id:c1", "end_user_restricted_registry"] + assert _published_keys(mock_publish) == ["end_user_id:c1", "end_user_restricted_registry"] + + +def test_customer_update_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): + """An update can add or drop a budget, block, region or permission, moving the id in the registry.""" + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( + return_value=_row({"user_id": "c1", "blocked": False}) + ) + mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=_row(_FULL_DB_ROW)) + + with _end_user_cache_doubles() as (recording_cache, mock_publish): + response = client.post( + "/customer/update", + json={"user_id": "c1", "budget_id": "b1"}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200, response.text + assert recording_cache.deleted == ["end_user_id:c1", "end_user_restricted_registry"] + assert _published_keys(mock_publish) == ["end_user_id:c1", "end_user_restricted_registry"] + + +def test_customer_block_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): + """Blocking is the one mutation that must take effect instantly; a stale registry keeps serving it.""" + mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock( + return_value=LiteLLM_EndUserTable(user_id="c1", blocked=True) + ) + + with _end_user_cache_doubles() as (recording_cache, mock_publish): + response = client.post( + "/customer/block", + json={"user_ids": ["c1", "c2"]}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200, response.text + assert recording_cache.deleted == [ + "end_user_id:c1", + "end_user_id:c2", + "end_user_restricted_registry", + ] + assert _published_keys(mock_publish) == [ + "end_user_id:c1", + "end_user_id:c2", + "end_user_restricted_registry", + ] + + +def test_customer_delete_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): + """Without this a deleted customer keeps its cached budget and block enforced until the TTL expires.""" + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock( + return_value=[ + LiteLLM_EndUserTable(user_id="c1", blocked=False), + LiteLLM_EndUserTable(user_id="c2", blocked=False), + ] + ) + mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=2) + + with _end_user_cache_doubles() as (recording_cache, mock_publish): + response = client.post( + "/customer/delete", + json={"user_ids": ["c1", "c2"]}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200, response.text + assert recording_cache.deleted == [ + "end_user_id:c1", + "end_user_id:c2", + "end_user_restricted_registry", + ] + assert _published_keys(mock_publish) == [ + "end_user_id:c1", + "end_user_id:c2", + "end_user_restricted_registry", + ] diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index 4fe1b54694f..018979aa19b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -14,7 +14,8 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from unittest.mock import Mock, patch +from contextlib import contextmanager +from unittest.mock import AsyncMock, Mock, patch import litellm from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -275,6 +276,190 @@ async def test_delete_tag(): app.dependency_overrides.clear() +class _RecordingAuthCache: + """Captures the keys an endpoint evicts, so tests assert on cache keys not mock plumbing.""" + + def __init__(self): + self.deleted: list[str] = [] + + async def async_delete_cache(self, key: str) -> None: + self.deleted.append(key) + + +@contextmanager +def _tag_cache_doubles(): + """Swaps in the auth cache and the cross-worker publisher a tag mutation is expected to hit.""" + recording_cache = _RecordingAuthCache() + mock_publish = AsyncMock() + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", recording_cache), + patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + mock_publish, + ), + ): + yield recording_cache, mock_publish + + +def _published_keys(mock_publish) -> list[str]: + return [call.kwargs["cache_key"] for call in mock_publish.call_args_list] + + +@pytest.mark.asyncio +async def test_new_tag_invalidates_tag_and_registry_caches(): + """ + A tag created on one worker must be visible to every worker's auth path immediately. + + Auth serves tags cache-first, and the cached tag-name registry is what decides whether a + request tag is looked up at all, so a create that leaves both entries stale means the new + tag's budget goes unenforced until the TTL expires. + """ + from datetime import datetime + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + try: + with ( + _tag_cache_doubles() as (recording_cache, mock_publish), + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.llm_router"), + patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id" + ), + patch( + "litellm.proxy.management_endpoints.tag_management_endpoints.get_deployments_by_model" + ) as mock_get_deployments, + ): + mock_db = Mock() + mock_prisma.db = mock_db + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=None) + mock_db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_get_deployments.return_value = [] + + created_tag = Mock() + created_tag.tag_name = "cache-tag" + created_tag.description = None + created_tag.models = [] + created_tag.model_info = {} + created_tag.spend = 0.0 + created_tag.budget_id = None + created_tag.created_at = datetime.now() + created_tag.updated_at = datetime.now() + created_tag.created_by = "test-user-123" + mock_db.litellm_tagtable.create = AsyncMock(return_value=created_tag) + + response = client.post( + "/tag/new", + json={"name": "cache-tag"}, + headers={"Authorization": "Bearer sk-1234"}, + ) + assert response.status_code == 200 + + assert recording_cache.deleted == ["tag:cache-tag", "tag_registry"] + assert _published_keys(mock_publish) == ["tag:cache-tag", "tag_registry"] + finally: + app.dependency_overrides.clear() + + +@pytest.mark.asyncio +async def test_update_tag_invalidates_only_the_tag_cache(): + """An update can change the tag's budget but never the set of names, so the registry stands.""" + from datetime import datetime + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + try: + with ( + _tag_cache_doubles() as (recording_cache, mock_publish), + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id" + ), + ): + mock_db = Mock() + mock_prisma.db = mock_db + + existing_tag = Mock() + existing_tag.tag_name = "cache-tag" + existing_tag.budget_id = None + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag) + mock_db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + updated_tag = Mock() + updated_tag.tag_name = "cache-tag" + updated_tag.description = "updated" + updated_tag.models = [] + updated_tag.model_info = {} + updated_tag.spend = 0.0 + updated_tag.budget_id = None + updated_tag.created_at = datetime.now() + updated_tag.updated_at = datetime.now() + updated_tag.created_by = "test-user-123" + mock_db.litellm_tagtable.update = AsyncMock(return_value=updated_tag) + + response = client.post( + "/tag/update", + json={"name": "cache-tag", "description": "updated"}, + headers={"Authorization": "Bearer sk-1234"}, + ) + assert response.status_code == 200 + + assert recording_cache.deleted == ["tag:cache-tag"] + assert _published_keys(mock_publish) == ["tag:cache-tag"] + finally: + app.dependency_overrides.clear() + + +@pytest.mark.asyncio +async def test_delete_tag_invalidates_tag_and_registry_caches(): + """Without this a deleted tag keeps its cached budget enforced until the TTL expires.""" + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + try: + with ( + _tag_cache_doubles() as (recording_cache, mock_publish), + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + ): + mock_db = Mock() + mock_prisma.db = mock_db + + existing_tag = Mock() + existing_tag.tag_name = "cache-tag" + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag) + mock_db.litellm_tagtable.delete = AsyncMock(return_value=existing_tag) + + response = client.post( + "/tag/delete", + json={"name": "cache-tag"}, + headers={"Authorization": "Bearer sk-1234"}, + ) + assert response.status_code == 200 + + assert recording_cache.deleted == ["tag:cache-tag", "tag_registry"] + assert _published_keys(mock_publish) == ["tag:cache-tag", "tag_registry"] + finally: + app.dependency_overrides.clear() + + @pytest.mark.asyncio async def test_list_tags_with_dynamic_tags(): """ diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 057193a69db..7052e050806 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -150,7 +150,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params): return where -def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None): +def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None): """ Create a MockPrismaClient for /spend/logs/ui endpoint tests. @@ -177,6 +177,8 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No return [{col: value, "_count": {col: n}} for value, n in tallied.items()] async def query_raw(self, sql_query, *params): + if query_observer is not None: + query_observer(sql_query, params) if "mcp_tool_call_count" in sql_query: return [] filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params)) @@ -1321,7 +1323,128 @@ async def test_ui_view_spend_logs_internal_user_scoped_without_user_id( @pytest.mark.asyncio -async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeypatch): +async def test_ui_view_spend_logs_explicit_user_filter_cannot_escape_own_scope(client, monkeypatch): + caller_log = { + "id": "log1", + "request_id": "req1", + "api_key": "sk-test-key", + "user": "caller@example.com", + "team_id": None, + "spend": 0.05, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + } + observed_queries = [] + + def observe_query(sql_query, params): + if 'FROM "LiteLLM_SpendLogs"' in sql_query: + observed_queries.append((sql_query, params)) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([caller_log], lambda _where: [], query_observer=observe_query), + ) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + AsyncMock(return_value=[]), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="caller@example.com" + ) + + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "user_id": "someone-else@example.com", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert response.json()["data"] == [] + page_sql, page_params = next((sql, params) for sql, params in observed_queries if "SELECT\n" in sql) + assert page_sql.count('"user" = $') == 2 + assert page_params[2:4] == ("someone-else@example.com", "caller@example.com") + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_without_user_filter_includes_permitted_team_scope(client, monkeypatch): + caller_log = { + "id": "log1", + "request_id": "req1", + "api_key": "sk-test-key", + "user": "team-admin@example.com", + "team_id": None, + "spend": 0.05, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + } + member_log = {**caller_log, "id": "log2", "request_id": "req2", "user": "member@example.com", "team_id": "team-9"} + outside_log = { + **caller_log, + "id": "log3", + "request_id": "req3", + "user": "outside@example.com", + "team_id": "outside-team", + } + + def filter_by_scope(where): + if {"multi_team": True} in where.get("OR", []) and "user" not in where: + return [caller_log, member_log] + return [caller_log, member_log, outside_log] + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([caller_log, member_log, outside_log], filter_by_scope), + ) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + AsyncMock(return_value=["team-9"]), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin@example.com" + ) + + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={"start_date": start_date, "end_date": end_date}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert [row["request_id"] for row in response.json()["data"]] == ["req1", "req2"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_permitted_team_scope_falls_back_to_own_user_when_lookup_fails(monkeypatch): + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + AsyncMock(side_effect=RuntimeError("database unavailable")), + ) + + permitted_team_ids = await spend_management_endpoints._get_permitted_team_ids_for_spend_logs_or_empty( + prisma_client=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="caller@example.com", + ), + ) + + assert permitted_team_ids == () + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_team_admin_can_filter_team_spend_by_user(client, monkeypatch): """ Team admins should be able to view team-wide spend when team_id is provided. """ @@ -1346,11 +1469,23 @@ async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeyp "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-4", }, + { + "id": "log3", + "request_id": "req3", + "api_key": "sk-test-key", + "user": "member3", + "team_id": "team_admin_team", + "spend": 0.15, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + }, ] def filter_by_team(where): - if "team_id" in where and where["team_id"] == "team_admin_team": + if where.get("team_id") == "team_admin_team" and where.get("user") == "member1": return [mock_spend_logs[0]] + if where.get("team_id") == "team_admin_team": + return [mock_spend_logs[0], mock_spend_logs[2]] return mock_spend_logs class TeamTable: @@ -1383,6 +1518,7 @@ async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeyp "/spend/logs/ui", params={ "team_id": "team_admin_team", + "user_id": "member1", "start_date": start_date, "end_date": end_date, }, @@ -1398,6 +1534,66 @@ async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeyp app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_user_filter_intersects_permitted_team_scope(client, monkeypatch): + member_log = { + "id": "log1", + "request_id": "req1", + "api_key": "sk-test-key", + "user": "member@example.com", + "team_id": "team-9", + "spend": 0.05, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + } + other_team_log = { + **member_log, + "id": "log2", + "request_id": "req2", + "team_id": "team-outside-scope", + } + seen_where = [] + + def filter_by_user_and_scope(where): + seen_where.append(where) + if where.get("user") == "member@example.com" and {"multi_team": True} in where.get("OR", []): + return [member_log] + return [member_log, other_team_log] + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([member_log, other_team_log], filter_by_user_and_scope), + ) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + AsyncMock(return_value=["team-9"]), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin" + ) + + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "user_id": "member@example.com", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert [row["request_id"] for row in response.json()["data"]] == ["req1"] + assert any( + where.get("user") == "member@example.com" and {"multi_team": True} in where.get("OR", []) + for where in seen_where + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_pagination(client, monkeypatch): mock_spend_logs = [ diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index e31058f402e..d3b1e089489 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -11,7 +11,7 @@ from pydantic import ValidationError as PydanticValidationError from starlette.datastructures import Headers import litellm -from litellm.proxy._types import AddTeamCallback, TeamCallbackMetadata, UserAPIKeyAuth +from litellm.proxy._types import AddTeamCallback, ProxyException, TeamCallbackMetadata, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import ( KeyAndTeamLoggingSettings, LiteLLMProxyRequestSetup, @@ -417,6 +417,91 @@ async def test_add_litellm_data_to_request_string_metadata_does_not_crash(): assert updated["metadata"].get("generation_name") == "test" +def _batches_request_mock() -> MagicMock: + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/batches" + request_mock.url.path = "/v1/batches" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + return request_mock + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "field,value,received_type", + [ + ("metadata", "abc", "a string"), + ("litellm_metadata", "abc", "a string"), + ("metadata", 42, "an integer"), + ("litellm_metadata", [1, 2], "an array"), + ("metadata", True, "a boolean"), + ], +) +async def test_add_litellm_data_to_request_rejects_non_object_metadata(field, value, received_type): + """Regression for https://github.com/BerriAI/litellm/issues/37147: a + non-object metadata was silently dropped with a 200, and a non-object + litellm_metadata crashed later with a 500 ('str' object has no attribute + 'update'). Both must be a 400 naming the field, like OpenAI returns.""" + data = {"input_file_id": "file-abc", "endpoint": "/v1/chat/completions", field: value} + + with pytest.raises(ProxyException) as exc_info: + await add_litellm_data_to_request( + data=data, + request=_batches_request_mock(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert exc_info.value.code == "400" + assert exc_info.value.param == field + assert exc_info.value.message == f"Invalid type for '{field}': expected an object, but got {received_type} instead." + assert field not in data + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_removes_every_invalid_metadata_field_before_raising(): + """When both fields are invalid, the raise for the first must not leave the + second invalid value in data, or failure-logging hooks that inspect the body + can crash on it and mask the 400 as a 500.""" + data = {"input_file_id": "file-abc", "metadata": "abc", "litellm_metadata": "xyz"} + + with pytest.raises(ProxyException) as exc_info: + await add_litellm_data_to_request( + data=data, + request=_batches_request_mock(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert exc_info.value.param == "metadata" + assert "metadata" not in data + assert "litellm_metadata" not in data + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_parses_json_object_string_litellm_metadata(): + data = {"input_file_id": "file-abc", "litellm_metadata": json.dumps({"cost_centre": "research"})} + + updated = await add_litellm_data_to_request( + data=data, + request=_batches_request_mock(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["litellm_metadata"]["cost_centre"] == "research" + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_proxy_server_request_body_is_post_strip(): """Regression: proxy_server_request['body'] used to be snapshotted before diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 2ccaa0df440..20d17b5a510 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1925,6 +1925,64 @@ class TestRunServerDbSetup: assert exc_info.value.code == 1 mock_setup_database.assert_not_called() + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") + def test_v2_migration_resolver_opts_in_via_env_var( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_atexit_register, + mock_subprocess_run, + ): + """USE_V2_MIGRATION_RESOLVER must select the v2 resolver. + + The Helm migrations Job runs `python litellm/proxy/prisma_migration.py`, + which calls run_server with a fixed argv, so a deployment has no way to + pass --use_v2_migration_resolver and an env var is the only route in. + """ + from litellm.proxy.proxy_cli import run_server + + mock_subprocess_run.return_value = MagicMock(returncode=0) + mock_should_update_schema.return_value = True + mock_setup_database.return_value = True + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" + clean_env["USE_V2_MIGRATION_RESOLVER"] = "true" + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + ): + run_server.main( + ["--local", "--skip_server_startup"], standalone_mode=False + ) + + mock_setup_database.assert_called_once_with( + use_migrate=True, use_v2_resolver=True + ) + # --- Module-level helpers for worker startup hook tests --- diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index ba207242e29..5545ee92e84 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11064,3 +11064,37 @@ async def test_ptu_rollup_job_not_registered_without_opt_in(monkeypatch): assert scheduler.get_job(PTU_ROLLUP_JOB_ID) is None assert len(scheduler.get_jobs()) > 0 + + +@pytest.mark.asyncio +async def test_moderations_reraises_proxy_exception_unwrapped(): + """A 400 ProxyException from request validation must surface as-is, + not be re-wrapped into a code-500 ProxyException.""" + from litellm.proxy._types import ProxyErrorTypes, ProxyException + + exc = ProxyException( + message="Invalid type for 'metadata': expected an object, but got a string instead.", + type=ProxyErrorTypes.bad_request_error, + param="metadata", + code=400, + ) + + request = MagicMock() + request.body = AsyncMock(return_value=b'{"input": "hi", "metadata": "abc"}') + + with ( + patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), + patch.object(proxy_server_module, "proxy_logging_obj") as mock_logging, + ): + mock_logging.post_call_failure_hook = AsyncMock() + with pytest.raises(ProxyException) as exc_info: + await proxy_server_module.moderations( + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=MagicMock(), + ) + + assert exc_info.value is exc + assert exc_info.value.code == "400" + assert exc_info.value.param == "metadata" + mock_logging.post_call_failure_hook.assert_awaited_once() diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 08e26125bd3..fc3b14592cd 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -1042,9 +1042,10 @@ async def test_route_request_override_enable_tag_filtering_beats_body_value(): [ ("acompletion", "messages", "/chat/completions"), ("aembedding", "input", "/embeddings"), + ("acreate_batch", "input_file_id", "/batches"), ], ) -@pytest.mark.parametrize("data_extra", [{}, {"messages": None, "input": None}]) +@pytest.mark.parametrize("data_extra", [{}, {"messages": None, "input": None, "input_file_id": None}]) def test_raise_if_required_body_param_missing_rejects_missing_param(route_type, param, route, data_extra): from litellm.proxy.route_llm_request import ( ProxyMissingRequiredParamError, @@ -1054,10 +1055,31 @@ def test_raise_if_required_body_param_missing_rejects_missing_param(route_type, with pytest.raises(ProxyMissingRequiredParamError) as exc_info: raise_if_required_body_param_missing(route_type=route_type, data={"model": "gpt-4o", **data_extra}) - assert exc_info.value.status_code == 400 + assert exc_info.value.code == "400" assert exc_info.value.param == param assert exc_info.value.type == "invalid_request_error" - assert exc_info.value.detail == {"error": f"{route}: Missing required parameter: '{param}'."} + assert exc_info.value.message == f"{route}: Missing required parameter: '{param}'." + + +@pytest.mark.parametrize( + "data, param", + [ + ({"endpoint": "/v1/chat/completions", "completion_window": "24h"}, "input_file_id"), + ({"input_file_id": "file-abc", "completion_window": "24h"}, "endpoint"), + ({"input_file_id": "file-abc", "endpoint": "/v1/chat/completions"}, "completion_window"), + ({}, "input_file_id"), + ], +) +def test_raise_if_required_body_param_missing_names_first_missing_batch_param(data, param): + from litellm.proxy.route_llm_request import ( + ProxyMissingRequiredParamError, + raise_if_required_body_param_missing, + ) + + with pytest.raises(ProxyMissingRequiredParamError) as exc_info: + raise_if_required_body_param_missing(route_type="acreate_batch", data=data) + + assert exc_info.value.param == param @pytest.mark.parametrize( @@ -1069,6 +1091,10 @@ def test_raise_if_required_body_param_missing_rejects_missing_param(route_type, ("aembedding", {"model": "text-embedding-3-small", "input": "hi"}), ("arerank", {"model": "rerank-model"}), ("aimage_generation", {"model": "dall-e-3"}), + ( + "acreate_batch", + {"input_file_id": "file-abc", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + ), ], ) def test_raise_if_required_body_param_missing_allows_valid_requests(route_type, data): @@ -1088,7 +1114,7 @@ async def test_route_request_rejects_chat_completion_without_messages(): with pytest.raises(ProxyMissingRequiredParamError) as exc_info: await route_request({"model": "gpt-4o"}, llm_router, None, "acompletion") - assert exc_info.value.status_code == 400 + assert exc_info.value.code == "400" assert exc_info.value.param == "messages" llm_router.acompletion.assert_not_called() diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 87525273911..b33ed3bb581 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -452,6 +452,42 @@ async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_functio assert captured.get("litellm_trace_id") == "tid" +@pytest.mark.asyncio +async def test_execute_tool_calls_threads_logging_obj_into_call_tool(monkeypatch): + """The Responses-API MCP path must hand the request's litellm_logging_obj to + global_mcp_server_manager.call_tool, otherwise pre_call_tool_check / + _create_during_hook_task get None and no guardrail evaluation is bridged onto + the request logger, so MCP tool calls made through the Responses API report zero + guardrail evaluations in the monitor. Drop the litellm_logging_obj kwarg on the + call_tool invocation and this fails.""" + _setup_proxy_logging(monkeypatch) + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + + sentinel_logging_obj = MagicMock() + sentinel_logging_obj.async_post_mcp_tool_call_hook = AsyncMock() + sentinel_logging_obj.async_success_handler = AsyncMock() + + handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + monkeypatch.setattr( + handler_module, + "function_setup", + lambda *_args, **_kwargs: (sentinel_logging_obj, None), + ) + + tool_name = "deepwiki-read_wiki_structure" + tool_calls = [{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}] + + await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=tool_calls, + user_api_key_auth=None, + ) + + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None + assert call_tool_mock.await_args.kwargs["litellm_logging_obj"] is sentinel_logging_obj + + @pytest.mark.asyncio async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch): """ diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ca848190a32..6d70a6aa5f4 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -30,7 +30,7 @@ "limit": 16713 }, "LIT011": { - "limit": 5591 + "limit": 5590 }, "LIT012": { "limit": 4519 diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 4a6d7006893..b1d9cf28503 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2947,9 +2947,6 @@ "max-lines": { "count": 1 }, - "no-nested-ternary": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.test.ts new file mode 100644 index 00000000000..5a79bf74ce3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.test.ts @@ -0,0 +1,40 @@ +import { renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const useInfiniteQuery = vi.fn(); +vi.mock("@/lib/http/api", () => ({ $api: { useInfiniteQuery: (...args: unknown[]) => useInfiniteQuery(...args) } })); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +import { useInfiniteSpendLogUsers } from "./useSpendLogUsers"; + +const WINDOW = { start_date: "2026-07-23 00:00:00", end_date: "2026-07-24 00:00:00" }; + +describe("useInfiniteSpendLogUsers", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token" }); + }); + + it("calls the scoped spend-log user facet with the visible window", () => { + renderHook(() => useInfiniteSpendLogUsers(WINDOW, 25, "alice")); + + const expectedQuery = { + "filter[startTime][gte]": "2026-07-23 00:00:00", + "filter[startTime][lte]": "2026-07-24 00:00:00", + page_size: 25, + q: "alice", + }; + expect(useInfiniteQuery.mock.calls[0][1]).toBe("/management/v1/spend_logs/users"); + expect(useInfiniteQuery.mock.calls[0][2].params.query).toEqual(expectedQuery); + }); + + it("omits q when the search box is empty", () => { + renderHook(() => useInfiniteSpendLogUsers(WINDOW, 50, "")); + + expect(useInfiniteQuery.mock.calls[0][2].params.query).not.toHaveProperty("q"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.ts new file mode 100644 index 00000000000..3a82c9e9d91 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.ts @@ -0,0 +1,21 @@ +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { $api } from "@/lib/http/api"; + +import { nextPageFromLinks, type SpendLogsWindow } from "./useSpendLogEndUsers"; + +export const useInfiniteSpendLogUsers = (window: SpendLogsWindow, pageSize: number = 50, q?: string) => { + const { accessToken } = useAuthorized(); + const query = { + "filter[startTime][gte]": window.start_date, + "filter[startTime][lte]": window.end_date, + page_size: pageSize, + ...(q !== undefined && q !== "" ? { q } : {}), + }; + const options = { + pageParamName: "page", + initialPageParam: 1, + getNextPageParam: nextPageFromLinks, + enabled: Boolean(accessToken), + }; + return $api.useInfiniteQuery("get", "/management/v1/spend_logs/users", { params: { query } }, options); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 928337010f7..c9faaf91238 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -24,7 +24,7 @@ import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip import React, { type ReactNode, useMemo, useState } from "react"; import TeamMultiSelect from "@/components/common_components/team_multi_select"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; -import { UsageExportHeader } from "@/components/EntityUsageExport"; +import { UsageExportHeader, type UsageFilterSelectProps } from "@/components/EntityUsageExport"; import type { EntityType } from "@/components/EntityUsageExport/types"; import { agentDailyActivityCall, @@ -84,6 +84,7 @@ interface EntityUsageProps { entityList: EntityList[] | null; premiumUser: boolean; dateValue: DateRangePickerValue; + filterSelectProps?: UsageFilterSelectProps; } const ENTITY_FETCH_FNS: Record Promise> = { @@ -107,6 +108,7 @@ const EntityUsage: React.FC = ({ entityList, userRole, dateValue, + filterSelectProps, }) => { const { teams } = useTeams(); const [selectedTags, setSelectedTags] = useState([]); @@ -621,6 +623,9 @@ const EntityUsage: React.FC = ({ { key: "endpoints", label: "Endpoint Activity", content: }, ]; + const hasEntityFilterOptions = entityList !== null && entityList.length > 0; + const showEntityFilters = entityType !== "team" && (filterSelectProps !== undefined || hasEntityFilterOptions); + return (
{isFetchingMore && ( @@ -679,7 +684,7 @@ const EntityUsage: React.FC = ({ dateValue={dateValue} entityType={entityType} spendData={spendData} - showFilters={entityType !== "team" && entityList !== null && entityList.length > 0} + showFilters={showEntityFilters} filterSlot={ entityType === "team" ? : undefined } @@ -689,6 +694,7 @@ const EntityUsage: React.FC = ({ onFiltersChange={setSelectedTags} filterOptions={getAllTags() || undefined} filterMode={entityType === "user" ? "single" : "multiple"} + filterSelectProps={filterSelectProps} teams={teams || []} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx index 9cd1e53c9b4..016d96cfa5c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx @@ -47,7 +47,18 @@ vi.mock("@/components/UsagePage/components/EntityUsage/TopKeyView", () => ({ })); vi.mock("./EntityUsage/EntityUsage", () => ({ - default: () =>
Entity Usage
, + default: ({ + entityType, + filterSelectProps, + }: { + entityType?: string; + filterSelectProps?: { onSearchChange?: (query: string) => void }; + }) => ( +
+ Entity Usage + {entityType === "user" && filterSelectProps !== undefined && Searchable user filter} +
+ ), EntityList: [], })); @@ -77,6 +88,7 @@ vi.mock("./UsageViewSelect/UsageViewSelect", async () => { React.createElement("option", { value: "customer" }, "Customer Usage"), tagOption, React.createElement("option", { value: "agent" }, "Agent Usage"), + React.createElement("option", { value: "user" }, "User Usage"), React.createElement("option", { value: "user-agent-activity" }, "User Agent Activity"), ); }; @@ -744,6 +756,18 @@ describe("UsagePage", () => { expect(mockUseInfiniteUsers).toHaveBeenCalledWith(50, undefined); }); + it("should reuse the searchable user filter in the user usage view", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + fireEvent.change(screen.getByTestId("usage-view-select"), { target: { value: "user" } }); + + expect(await screen.findByText("Searchable user filter")).toBeInTheDocument(); + }); + it("should deduplicate users across pages", async () => { mockUseInfiniteUsers.mockReturnValue({ data: { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index 96cb0098a3a..ea8f02a4e70 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -28,7 +28,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import { all_admin_roles, internalUserRoles } from "@/utils/roles"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; import CloudZeroExportModal from "@/components/cloudzero_export_modal"; -import EntityUsageExportModal from "@/components/EntityUsageExport"; +import EntityUsageExportModal, { type UsageFilterSelectProps } from "@/components/EntityUsageExport"; import { Team } from "@/components/key_team_helpers/key_list"; import { gatewayDailyActivityCall, @@ -135,6 +135,14 @@ const UsagePage: React.FC = ({ teams, organizations }) => { return result; }, [usersInfiniteData]); + const userFilterSelectProps: UsageFilterSelectProps = { + onSearchChange: setSettledUserSearch, + onLoadMore: () => void fetchNextUsersPage(), + hasNextPage: hasNextUsersPage, + isLoading: isLoadingUsers, + isFetchingNextPage: isFetchingNextUsersPage, + emptyText: "No users found", + }; // For admins: null means global view (all users), a string means filter by that user // For non-admins: always set to their own user ID const [selectedUserId, setSelectedUserId] = useState(isAdmin ? null : userID || null); @@ -1043,6 +1051,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { userID={userID} userRole={userRole} entityList={userOptions.length > 0 ? userOptions : null} + filterSelectProps={userFilterSelectProps} premiumUser={premiumUser} dateValue={dateValue} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx index 2632f40adbe..1a3d517df3e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx @@ -148,12 +148,25 @@ describe("ViewUserDashboard", () => { expect(screen.getByRole("textbox", { name: "Default setting" })).toHaveValue("unsaved change"); }); + it("renders invite and bulk invite as toolbar actions alongside the other admin controls", async () => { + renderDashboard(); + + const inviteButton = await screen.findByRole("button", { name: /\+ invite user/i }); + const bulkInviteButton = screen.getByRole("button", { name: /\+ bulk invite users/i }); + const toolbar = screen.getByTestId("toggle-user-selection").parentElement; + + expect(inviteButton.parentElement).toBe(toolbar); + expect(bulkInviteButton.parentElement).toBe(toolbar); + }); + it("shows the users table without admin controls for non-proxy admins", async () => { renderDashboard({ userRole: "Internal User" }); expect(await screen.findByText("test@example.com")).toBeInTheDocument(); expect(screen.queryByRole("tab")).not.toBeInTheDocument(); expect(screen.queryByTestId("toggle-user-selection")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /\+ invite user/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /\+ bulk invite users/i })).not.toBeInTheDocument(); }); it("keeps actions unavailable while the user list is loading", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index 1de01e88866..df2345b0af4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -2,6 +2,7 @@ import { parseAsString, useQueryState } from "nuqs"; import React, { useCallback, useEffect, useMemo, useState } from "react"; import BulkEditUserModal from "./BulkEditUsers"; +import BulkCreateUsersButton from "@/components/bulk_create_users_button"; import { CreateUserButton } from "@/components/CreateUserButton"; import { Button } from "@/components/ui/button"; import { Skeleton } from "@/components/ui/skeleton"; @@ -365,12 +366,11 @@ const ViewUserDashboard: React.FC = ({ {!userListQuery.isLoading && userID && accessToken && ( <> {isProxyAdmin && ( - + + )} + + {isProxyAdmin && ( + )} {isProxyAdmin && ( diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx index 246c01e9a9a..de1c3508dac 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx @@ -20,10 +20,6 @@ vi.mock("./networking", () => ({ getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost"), })); -vi.mock("./bulk_create_users_button", () => ({ - default: () =>
Bulk Create Users
, -})); - vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ useOrganizations: vi.fn().mockReturnValue({ data: [], isLoading: false }), })); @@ -42,7 +38,6 @@ const createQueryClient = () => const defaultProps = { userID: "123", accessToken: "token", - teams: [], possibleUIRoles: null as Record> | null, }; @@ -75,6 +70,14 @@ describe("CreateUserButton", () => { }); }); + it("should not render the bulk invite button", async () => { + renderWithProviders(); + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + expect(screen.queryByRole("button", { name: /bulk invite users/i })).not.toBeInTheDocument(); + }); + it("should open the invite modal when invite user button is clicked", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.tsx index 8de9f010fa3..8a9146d756f 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.tsx @@ -1,22 +1,11 @@ -import { InfoCircleOutlined, UserAddOutlined } from "@ant-design/icons"; +import { InfoCircleOutlined } from "@ant-design/icons"; import { useQueryClient } from "@tanstack/react-query"; import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { Button } from "@/components/ui/button"; import { Accordion, AccordionBody, AccordionHeader, SelectItem, TextInput } from "@tremor/react"; -import { - Alert, - Button, - Checkbox, - Form, - Input, - Modal, - Select, - Select as Select2, - Space, - Tooltip, - Typography, -} from "antd"; +import { Alert, Checkbox, Form, Input, Modal, Select, Select as Select2, Space, Tooltip, Typography } from "antd"; +import { UserPlus } from "lucide-react"; import React, { useEffect, useState } from "react"; -import BulkCreateUsers from "./bulk_create_users_button"; import TeamDropdown from "./common_components/team_dropdown"; import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key"; import NotificationsManager from "./molecules/notifications_manager"; @@ -46,7 +35,6 @@ const generateUUID = (): string => { interface CreateuserProps { userID: string; accessToken: string; - teams: any[] | null; possibleUIRoles: null | Record>; onUserCreated?: (userId: string) => void; isEmbedded?: boolean; @@ -63,7 +51,6 @@ interface UISettings { export const CreateUserButton: React.FC = ({ userID, accessToken, - teams, possibleUIRoles, onUserCreated, isEmbedded = false, @@ -79,8 +66,6 @@ export const CreateUserButton: React.FC = ({ const [baseUrl, setBaseUrl] = useState(null); const { data: organizations = [] } = useOrganizations(); - // Derive teams from the user's organizations, falling back to the teams prop - useEffect(() => { const fetchData = async () => { try { @@ -237,7 +222,7 @@ export const CreateUserButton: React.FC = ({
- +
); @@ -245,11 +230,10 @@ export const CreateUserButton: React.FC = ({ // Original return for standalone mode return ( -
- - = ({
-
@@ -391,6 +376,6 @@ export const CreateUserButton: React.FC = ({ invitationLinkData={invitationLinkData} /> )} -
+ ); }; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.test.tsx index 729d6fd3406..bfc667f60ba 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.test.tsx @@ -1,4 +1,4 @@ -import { renderWithProviders, screen } from "../../../tests/test-utils"; +import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import { vi } from "vitest"; import UsageExportHeader from "./UsageExportHeader"; @@ -70,4 +70,29 @@ describe("UsageExportHeader", () => { ); expect(screen.getByText("Team")).toBeInTheDocument(); }); + + it("should keep a searchable single filter usable when no options match", async () => { + const user = userEvent.setup(); + const onSearchChange = vi.fn(); + + renderWithProviders( + , + ); + + const userFilter = screen.getByRole("combobox"); + await user.click(userFilter); + await user.type(userFilter, "alice"); + + await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("alice")); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx index db386af1ffd..2716182f702 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx @@ -1,6 +1,7 @@ import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import { Download } from "lucide-react"; import React, { useState } from "react"; +import { PaginatedSearchSelect } from "@/components/shared/PaginatedSearchSelect"; import { Button } from "@/components/ui/button"; import { Combobox, @@ -20,6 +21,16 @@ import EntityUsageExportModal from "./EntityUsageExportModal"; import type { EntitySpendData, EntityType } from "./types"; import type { Team } from "@/components/key_team_helpers/key_list"; +export interface UsageFilterSelectProps { + onSearchChange: (query: string) => void; + onLoadMore: () => void; + hasNextPage?: boolean; + isLoading?: boolean; + isFetchingNextPage?: boolean; + emptyText?: string; + loadingText?: string; +} + interface UsageExportHeaderProps { dateValue: DateRangePickerValue; entityType: EntityType; @@ -32,6 +43,7 @@ interface UsageExportHeaderProps { onFiltersChange?: (filters: string[]) => void; filterOptions?: Array<{ label: string; value: string }>; filterMode?: "multiple" | "single"; + filterSelectProps?: UsageFilterSelectProps; filterSlot?: React.ReactNode; customTitle?: string; compactLayout?: boolean; @@ -49,6 +61,7 @@ const UsageExportHeader: React.FC = ({ onFiltersChange, filterOptions = [], filterMode = "multiple", + filterSelectProps, filterSlot, customTitle, compactLayout = false, @@ -57,7 +70,8 @@ const UsageExportHeader: React.FC = ({ const anchor = useComboboxAnchor(); const [isExportModalOpen, setIsExportModalOpen] = useState(false); - const hasFilters = filterSlot != null || (showFilters && filterOptions.length > 0); + const hasBuiltInFilter = filterOptions.length > 0 || filterSelectProps !== undefined; + const hasFilters = filterSlot != null || (showFilters && hasBuiltInFilter); const optionValues = filterOptions.map((option) => option.value); const labelOf = (value: string) => filterOptions.find((option) => option.value === value)?.label ?? value; @@ -74,45 +88,59 @@ const UsageExportHeader: React.FC = ({ ); - const builtInFilter = - filterMode === "single" ? ( - onFiltersChange?.(next ? [next] : [])} - itemToStringLabel={labelOf} - > - 0} - /> - {filterList} - - ) : ( - onFiltersChange?.(next)} - > - } className="w-full"> - - {(selected: string[]) => - selected.map((value) => ( - - {labelOf(value)} - - )) - } - - - {selectedFilters.length > 0 && } - - {filterList} - - ); + const searchableSingleFilter = + filterSelectProps !== undefined ? ( + onFiltersChange?.(next ? [next] : [])} + placeholder={filterPlaceholder} + {...filterSelectProps} + /> + ) : undefined; + + const singleFilter = searchableSingleFilter ?? ( + onFiltersChange?.(next ? [next] : [])} + itemToStringLabel={labelOf} + > + 0} + /> + {filterList} + + ); + + const multiFilter = ( + onFiltersChange?.(next)} + > + } className="w-full"> + + {(selected: string[]) => + selected.map((value) => ( + + {labelOf(value)} + + )) + } + + + {selectedFilters.length > 0 && } + + {filterList} + + ); + + const builtInFilter = filterMode === "single" ? singleFilter : multiFilter; return ( <> diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/index.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/index.ts index dca6d9e0a6a..2a5e94aabbe 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/index.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/index.ts @@ -1,3 +1,4 @@ export { default } from "./EntityUsageExportModal"; export { default as UsageExportHeader } from "./UsageExportHeader"; +export type { UsageFilterSelectProps } from "./UsageExportHeader"; export * from "./types"; diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index fdbc07ee020..48eb4dd08d5 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -5,6 +5,7 @@ import { ColumnDef } from "@tanstack/react-table"; import { Popover, Typography } from "antd"; import { DataTableMultiSortHeader, DataTableSortHeader, type DataTableSortField } from "@/components/shared/DataTable"; +import { inheritedBudgetGates } from "@/components/shared/InheritedBudgetHint"; import { Skeleton } from "@/components/ui/skeleton"; import { DateCell, @@ -304,13 +305,14 @@ export const getKeyTableColumns = ({ size: 180, enableSorting: true, cell: ({ row }) => { - const teamId = row.original.team_id; - const team = allTeams.find((t) => t.team_id === teamId); + const team = allTeams.find((t) => t.team_id === row.original.team_id); + const orgId = row.original.organization_id || row.original.org_id || team?.organization_id; + const organization = organizations.find((o) => o.organization_id === orgId); return ( ); }, 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 fc0dae1f912..8573c2f7415 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1718,7 +1718,6 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp { + it("returns team then org gates when both have budgets", () => { + expect(inheritedBudgetGates(team, organization)).toEqual([ + { scope: "Team", alias: "Platform", maxBudget: 1200, budgetDuration: "30d" }, + { scope: "Organization", alias: "Acme", maxBudget: 5000, budgetDuration: null }, + ]); + }); + + it("skips a team or org whose max_budget is null", () => { + expect(inheritedBudgetGates({ ...team, max_budget: null }, organization)).toEqual([ + { scope: "Organization", alias: "Acme", maxBudget: 5000, budgetDuration: null }, + ]); + expect(inheritedBudgetGates(team, { ...organization, litellm_budget_table: { max_budget: null } })).toEqual([ + { scope: "Team", alias: "Platform", maxBudget: 1200, budgetDuration: "30d" }, + ]); + }); + + it("returns nothing when team and org are missing or budgetless", () => { + expect(inheritedBudgetGates(null, undefined)).toEqual([]); + expect( + inheritedBudgetGates({ ...team, max_budget: null }, { ...organization, litellm_budget_table: null }), + ).toEqual([]); + }); + + it("falls back to ids when aliases are empty", () => { + expect( + inheritedBudgetGates({ ...team, team_alias: "" }, { ...organization, organization_alias: "" }).map( + (g) => g.alias, + ), + ).toEqual(["team-1", "org-1"]); + }); +}); + +describe("InheritedBudgetHint", () => { + it("renders nothing without gates", () => { + const { container } = render(); + expect(container).toBeEmptyDOMElement(); + }); + + it("shows each gate with its budget and duration on hover", async () => { + render(); + await userEvent.setup().hover(screen.getByLabelText("question-circle")); + expect(screen.getByTestId("inherited-budget-hint")).toHaveTextContent("Team Platform: $1,200.00 / 30d"); + expect(screen.getByTestId("inherited-budget-hint")).toHaveTextContent("Organization Acme: $5,000.00"); + expect(screen.getByTestId("inherited-budget-hint")).not.toHaveTextContent("Organization Acme: $5,000.00 /"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx new file mode 100644 index 00000000000..6be1dad564f --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx @@ -0,0 +1,67 @@ +"use client"; + +import { Tooltip } from "@/components/atoms/Tooltip"; +import type { Team } from "@/components/key_team_helpers/key_list"; +import type { Organization } from "@/components/networking"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; + +export interface InheritedBudgetGate { + scope: "Team" | "Organization"; + alias: string; + maxBudget: number; + budgetDuration: string | null; +} + +type TeamBudgetSource = Pick; +type OrganizationBudgetSource = Pick; + +const teamGate = (team: TeamBudgetSource | null | undefined): InheritedBudgetGate | null => + team && team.max_budget != null + ? { + scope: "Team", + alias: team.team_alias || team.team_id, + maxBudget: team.max_budget, + budgetDuration: team.budget_duration ?? null, + } + : null; + +const organizationGate = (organization: OrganizationBudgetSource | null | undefined): InheritedBudgetGate | null => { + const budgetTable: { max_budget?: number | null; budget_duration?: string | null } | null | undefined = + organization?.litellm_budget_table; + return organization && budgetTable?.max_budget != null + ? { + scope: "Organization", + alias: organization.organization_alias || organization.organization_id, + maxBudget: budgetTable.max_budget, + budgetDuration: budgetTable.budget_duration ?? null, + } + : null; +}; + +export const inheritedBudgetGates = ( + team: TeamBudgetSource | null | undefined, + organization: OrganizationBudgetSource | null | undefined, +): readonly InheritedBudgetGate[] => [teamGate(team), organizationGate(organization)].filter((gate) => gate !== null); + +const formatGate = (gate: InheritedBudgetGate): string => + `${gate.scope} ${gate.alias}: $${formatNumberWithCommas(gate.maxBudget, 2)}${gate.budgetDuration ? ` / ${gate.budgetDuration}` : ""}`; + +interface InheritedBudgetHintProps { + gates: readonly InheritedBudgetGate[]; +} + +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: + {gates.map((gate) => ( + {formatGate(gate)} + ))} +
+ } + /> + ); +} diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx index d4af8428d69..0bc14849040 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx @@ -53,9 +53,24 @@ describe("SpendBudgetCell", () => { expect(indicator(container)?.className).toContain("bg-destructive"); }); - it("falls back to the team budget and labels it", () => { - render(); - expect(screen.getByText("of $200 (Team)")).toBeInTheDocument(); - expect(screen.getByRole("meter")).toHaveAttribute("aria-valuemax", "200"); + it("never meters key spend against an inherited team/org budget", () => { + const gates = [{ scope: "Team" as const, alias: "Team A", maxBudget: 200, budgetDuration: "30d" }]; + render(); + expect(screen.getByText("· Unlimited")).toBeInTheDocument(); + expect(screen.queryByText(/\(Team\)/)).not.toBeInTheDocument(); + expect(screen.queryByRole("meter")).not.toBeInTheDocument(); + expect(screen.getByLabelText("question-circle")).toBeInTheDocument(); + }); + + it("shows no inherited-budget hint when there is nothing to inherit", () => { + render(); + expect(screen.queryByLabelText("question-circle")).not.toBeInTheDocument(); + }); + + it("shows no inherited-budget hint when the key has its own budget", () => { + const gates = [{ scope: "Team" as const, alias: "Team A", maxBudget: 200, budgetDuration: null }]; + render(); + expect(screen.getByText("of $50")).toBeInTheDocument(); + expect(screen.queryByLabelText("question-circle")).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx index 60b42615967..943b9aa766c 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx @@ -1,12 +1,13 @@ "use client"; +import { InheritedBudgetHint, type InheritedBudgetGate } from "@/components/shared/InheritedBudgetHint"; import { Meter, MeterIndicator, MeterTrack } from "@/components/ui/meter"; import { formatNumberWithCommas, getSpendString } from "@/utils/dataUtils"; interface SpendBudgetCellProps { spend: number | null | undefined; maxBudget: number | null | undefined; - teamMaxBudget?: number | null; + inheritedGates?: readonly InheritedBudgetGate[]; spendDecimals?: number; budgetDecimals?: number; } @@ -20,27 +21,24 @@ const meterTone = (pct: number): "default" | "warning" | "over" => { export function SpendBudgetCell({ spend, maxBudget, - teamMaxBudget, + inheritedGates = [], spendDecimals = 4, budgetDecimals = 0, }: SpendBudgetCellProps) { const spendValue = typeof spend === "number" && !Number.isNaN(spend) ? spend : 0; - const budget = maxBudget ?? teamMaxBudget ?? null; - const isTeamBudget = maxBudget == null && teamMaxBudget != null; + const budget = maxBudget ?? null; const hasBudget = typeof budget === "number" && budget > 0; const pct = hasBudget ? (spendValue / budget) * 100 : 0; const spendText = spendValue > 0 ? getSpendString(spendValue, spendDecimals) : "$0.00"; - const budgetLabel = - budget === null - ? "· Unlimited" - : `of $${formatNumberWithCommas(budget, budgetDecimals)}${isTeamBudget ? " (Team)" : ""}`; + const budgetLabel = budget === null ? "· Unlimited" : `of $${formatNumberWithCommas(budget, budgetDecimals)}`; return (
{spendText}{" "} {budgetLabel} + {budget === null && }
{hasBudget && ( ({ useRouter: () => ({ push: vi.fn() }) })); -vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ - useOrganizations: () => ({ data: [] }), -})); - vi.mock("./key_edit_view", () => ({ KeyEditView: () =>
, })); vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ default: vi.fn() })); +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ useOrganizations: vi.fn() })); vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: vi.fn() })); vi.mock("@/app/(dashboard)/hooks/projects/useProjects", () => ({ useProjects: vi.fn().mockReturnValue({ data: [], isLoading: false }), @@ -130,10 +130,34 @@ const makeTeam = (overrides: Partial): Team => ({ ...overrides, }); +const makeOrganization = (overrides: Partial): Organization => + ({ + organization_id: "org-1", + organization_alias: "Acme Org", + budget_id: "budget-1", + metadata: {}, + models: [], + spend: 0, + model_spend: {}, + created_at: "2026-01-01T00:00:00Z", + created_by: "admin", + updated_at: "2026-01-01T00:00:00Z", + updated_by: "admin", + litellm_budget_table: { max_budget: null, budget_duration: null }, + teams: null, + users: null, + members: null, + ...overrides, + }) as Organization; + +const mockOrganizations = (organizations: Organization[]) => + vi.mocked(useOrganizations).mockReturnValue({ data: organizations } as ReturnType); + describe("KeyInfoView overview budget display (LIT-2845)", () => { beforeEach(() => { vi.mocked(useTeams).mockReturnValue({ teams: [], setTeams: vi.fn() }); vi.mocked(useAuthorized).mockReturnValue(baseAuthorized); + mockOrganizations([]); }); it("renders a sub-dollar max_budget ($0.10) with 2-decimal precision in the overview Spend card", async () => { @@ -188,7 +212,7 @@ describe("KeyInfoView overview budget display (LIT-2845)", () => { }); }); - it("renders team budget with alias and duration when key has no own budget but team has one", async () => { + it("never pairs key spend with the team budget: shows Unlimited plus an inherited-budget hint", async () => { vi.mocked(useTeams).mockReturnValue({ teams: [makeTeam({ team_id: "team-123", team_alias: "Test Budget", max_budget: 1200, budget_duration: "30d" })], setTeams: vi.fn(), @@ -203,15 +227,20 @@ describe("KeyInfoView overview budget display (LIT-2845)", () => { />, ); await waitFor(() => { - expect(screen.getByText(/of \$1,200\.00 \(Team: Test Budget \/ 30d\)/)).toBeInTheDocument(); + expect(screen.getByText(/of Unlimited/)).toBeInTheDocument(); }); + expect(screen.queryByText(/of \$1,200\.00/)).not.toBeInTheDocument(); + expect(screen.queryByText(/\(Team: Test Budget/)).not.toBeInTheDocument(); + await userEvent.setup().hover(screen.getByLabelText("question-circle")); + expect(screen.getByTestId("inherited-budget-hint")).toHaveTextContent("Team Test Budget: $1,200.00 / 30d"); }); - it("renders team budget without duration when team has no budget_duration", async () => { + it("lists the organization budget in the hint when the team's org has one", async () => { vi.mocked(useTeams).mockReturnValue({ - teams: [makeTeam({ team_id: "team-456", team_alias: "No Duration Team", max_budget: 500 })], + teams: [makeTeam({ team_id: "team-456", team_alias: "Org Team", organization_id: "org-1" })], setTeams: vi.fn(), }); + mockOrganizations([makeOrganization({ litellm_budget_table: { max_budget: 5000, budget_duration: null } })]); renderWithProviders( { />, ); await waitFor(() => { - expect(screen.getByText(/of \$500\.00 \(Team: No Duration Team\)/)).toBeInTheDocument(); + expect(screen.getByText(/of Unlimited/)).toBeInTheDocument(); }); + await userEvent.setup().hover(screen.getByLabelText("question-circle")); + expect(screen.getByTestId("inherited-budget-hint")).toHaveTextContent("Organization Acme Org: $5,000.00"); + expect(screen.getByTestId("inherited-budget-hint")).not.toHaveTextContent("Team Org Team"); }); - it("renders 'Unlimited' when key has no budget and team also has no budget", async () => { + it("renders 'Unlimited' with no hint when neither key, team, nor org has a budget", async () => { vi.mocked(useTeams).mockReturnValue({ teams: [makeTeam({ team_id: "team-789", team_alias: "Free Team" })], setTeams: vi.fn(), @@ -243,6 +275,27 @@ describe("KeyInfoView overview budget display (LIT-2845)", () => { await waitFor(() => { expect(screen.getByText(/of Unlimited/)).toBeInTheDocument(); }); + expect(screen.queryByLabelText("question-circle")).not.toBeInTheDocument(); + }); + + it("shows no hint when the key has its own budget even if the team has one", async () => { + vi.mocked(useTeams).mockReturnValue({ + teams: [makeTeam({ team_id: "team-123", team_alias: "Test Budget", max_budget: 1200 })], + setTeams: vi.fn(), + }); + renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + await waitFor(() => { + expect(screen.getByText(/of \$25\.00/)).toBeInTheDocument(); + }); + expect(screen.queryByLabelText("question-circle")).not.toBeInTheDocument(); }); }); @@ -250,6 +303,7 @@ describe("KeyInfoView budget reset visibility", () => { beforeEach(() => { vi.mocked(useTeams).mockReturnValue({ teams: [], setTeams: vi.fn() }); vi.mocked(useAuthorized).mockReturnValue(baseAuthorized); + mockOrganizations([]); }); const KEY_WITH_RESET = { diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 3ee11325a65..99981edb1c8 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -36,6 +36,7 @@ import { extractMcpEntitlement } from "../mcp_server_management/mcpEntitlement"; import ObjectPermissionsView from "../object_permissions_view"; import { RegenerateKeyModal } from "../organisms/RegenerateKeyModal"; import { parseErrorMessage } from "../shared/errorUtils"; +import { InheritedBudgetHint, inheritedBudgetGates } from "../shared/InheritedBudgetHint"; import { KeyEditView } from "./key_edit_view"; interface KeyInfoViewProps { @@ -460,12 +461,9 @@ export default function KeyInfoView({ const orgId = currentKeyData.organization_id || currentKeyData.org_id || parentTeam?.organization_id || ""; const parentOrg = orgId ? organizations?.find((org) => org.organization_id === orgId) : null; - const budgetDisplay = - currentKeyData.max_budget !== null - ? `$${formatNumberWithCommas(currentKeyData.max_budget, 2)}` - : parentTeam?.max_budget != null - ? `$${formatNumberWithCommas(parentTeam.max_budget, 2)} (Team: ${parentTeam.team_alias || parentTeam.team_id}${parentTeam.budget_duration ? ` / ${parentTeam.budget_duration}` : ""})` - : "Unlimited"; + const hasOwnBudget = currentKeyData.max_budget !== null; + const budgetDisplay = hasOwnBudget ? `$${formatNumberWithCommas(currentKeyData.max_budget, 2)}` : "Unlimited"; + const inheritedGates = hasOwnBudget ? [] : inheritedBudgetGates(parentTeam, parentOrg); return (
@@ -616,7 +614,10 @@ export default function KeyInfoView({

Spend

${formatNumberWithCommas(currentKeyData.spend, 4)}

-

of {budgetDisplay}

+

+ of {budgetDisplay} + +

{currentKeyData.budget_reset_at && (

Resets {formatTimestamp(currentKeyData.budget_reset_at)}

)} diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx index 82b179b2654..1c94e6418a0 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx @@ -14,11 +14,16 @@ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useInfiniteModelInfo: vi.fn(), })); +vi.mock("@/app/(dashboard)/hooks/spendLogs/useSpendLogUsers", () => ({ + useInfiniteSpendLogUsers: vi.fn(), +})); + vi.mock("@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers", () => ({ useInfiniteSpendLogEndUsers: vi.fn(), })); import { useInfiniteSpendLogEndUsers } from "@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers"; +import { useInfiniteSpendLogUsers } from "@/app/(dashboard)/hooks/spendLogs/useSpendLogUsers"; import { useInfiniteKeyAliases } from "@/app/(dashboard)/hooks/keys/useKeyAliases"; import { useInfiniteModelInfo } from "@/app/(dashboard)/hooks/models/useModels"; @@ -50,6 +55,9 @@ describe("RequestLogsFilters", () => { vi.mocked(useInfiniteModelInfo).mockReturnValue( emptyInfiniteQuery as unknown as ReturnType, ); + vi.mocked(useInfiniteSpendLogUsers).mockReturnValue( + emptyInfiniteQuery as unknown as ReturnType, + ); vi.mocked(useInfiniteSpendLogEndUsers).mockReturnValue( emptyInfiniteQuery as unknown as ReturnType, ); @@ -62,6 +70,7 @@ describe("RequestLogsFilters", () => { "Team ID", "Status", "Key Alias", + "User ID", "End User", "Error Code", "Error Message", @@ -74,6 +83,78 @@ describe("RequestLogsFilters", () => { } }); + it("places User ID between Key Alias and End User", async () => { + renderFilters(); + + const labels = ["Key Alias", "User ID", "End User"].map((label) => screen.getByText(label)); + expect(labels[0].compareDocumentPosition(labels[1]) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + expect(labels[1].compareDocumentPosition(labels[2]) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + }); + + it("selects an internal user value from the caller's visible spend logs", async () => { + vi.mocked(useInfiniteSpendLogUsers).mockReturnValue({ + ...emptyInfiniteQuery, + data: { + pages: [ + { + data: ["alice@example.com"], + meta: { page: 1, page_size: 50, has_more: false }, + links: { self: "", next: null }, + }, + ], + pageParams: [1], + }, + } as unknown as ReturnType); + const user = userEvent.setup(); + const { set } = renderFilters(); + + await user.click(await screen.findByPlaceholderText("Search an internal user")); + await user.click(await screen.findByText("alice@example.com")); + + expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.USER_ID, "alice@example.com"); + }); + + it("pushes the User ID picker query to the paginated user lookup", async () => { + const user = userEvent.setup(); + renderFilters(); + + const input = await screen.findByPlaceholderText("Search an internal user"); + await user.click(input); + await user.type(input, "alice@example.com"); + + await waitFor(() => expect(useInfiniteSpendLogUsers).toHaveBeenCalledWith(LOGS_WINDOW, 50, "alice@example.com")); + }); + + it("loads the next page when the User ID list is scrolled near the end", async () => { + const fetchNextPage = vi.fn(); + vi.mocked(useInfiniteSpendLogUsers).mockReturnValue({ + ...emptyInfiniteQuery, + fetchNextPage, + hasNextPage: true, + data: { + pages: [ + { + data: ["alice@example.com"], + meta: { page: 1, page_size: 50, has_more: true }, + links: { self: "", next: "?page=2" }, + }, + ], + pageParams: [1], + }, + } as unknown as ReturnType); + const user = userEvent.setup(); + renderFilters(); + + await user.click(await screen.findByPlaceholderText("Search an internal user")); + const list = await screen.findByTestId("paginated-search-select-list"); + Object.defineProperty(list, "scrollTop", { value: 90, configurable: true }); + Object.defineProperty(list, "clientHeight", { value: 10, configurable: true }); + Object.defineProperty(list, "scrollHeight", { value: 100, configurable: true }); + list.dispatchEvent(new Event("scroll", { bubbles: true })); + + await waitFor(() => expect(fetchNextPage).toHaveBeenCalled()); + }); + it("scopes the Key Alias lookup to the selected team", async () => { renderFilters({ [LOG_FILTER_IDS.TEAM_ID]: "team-42" }); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx index 2005b868cd6..017260230dd 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx @@ -3,6 +3,7 @@ import { useMemo, useState } from "react"; import { useInfiniteSpendLogEndUsers } from "@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers"; +import { useInfiniteSpendLogUsers } from "@/app/(dashboard)/hooks/spendLogs/useSpendLogUsers"; import { useInfiniteKeyAliases } from "@/app/(dashboard)/hooks/keys/useKeyAliases"; import { useInfiniteModelInfo } from "@/app/(dashboard)/hooks/models/useModels"; import { DataTableFilterField } from "@/components/shared/DataTable"; @@ -144,6 +145,51 @@ function ModelFilterField({ value, onChange }: { value: string; onChange: (value ); } +function UserIdFilterField({ + value, + onChange, + logsWindow, +}: { + value: string; + onChange: (value: string | undefined) => void; + logsWindow: LogsWindow; +}) { + const [search, setSearch] = useState(""); + const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteSpendLogUsers( + logsWindow, + PAGE_SIZE, + emptyToUndefined(search), + ); + + const options = useMemo(() => { + const seen = new Set(); + return (data?.pages ?? []).flatMap((page) => + page.data.flatMap((userId) => { + if (!userId || seen.has(userId)) return []; + seen.add(userId); + return [{ label: userId, value: userId }]; + }), + ); + }, [data]); + + return ( + + onChange(emptyToUndefined(next))} + onSearchChange={setSearch} + onLoadMore={() => void fetchNextPage()} + hasNextPage={hasNextPage} + isLoading={isLoading} + isFetchingNextPage={isFetchingNextPage} + placeholder="Search an internal user" + emptyText="No users found" + /> + + ); +} + function EndUserFilterField({ value, onChange, @@ -279,6 +325,12 @@ export function RequestLogsFilters({ get, set, teams, logsWindow }: RequestLogsF teamId={valueOf(LOG_FILTER_IDS.TEAM_ID)} /> + + { }); }); - describe("filterByCurrentUser", () => { - it("scopes to the current user when no explicit user filter is set", async () => { - renderFilterHook({ filterByCurrentUser: true }); + describe("user scope", () => { + it("leaves an empty user filter for the backend to authorize", async () => { + renderFilterHook(); await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled()); - expect(lastCallParams()?.params).toMatchObject({ user_id: "user-1" }); + expect(lastCallParams()?.params?.user_id).toBeUndefined(); }); - it("lets an explicit user filter win over the current-user scope", async () => { + it("sends an explicit user filter for the backend to intersect with authorization", async () => { renderFilterHook({ - filterByCurrentUser: true, columnFilters: [{ id: LOG_FILTER_IDS.USER_ID, value: "someone-else" }], }); diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index e1089c6a16c..474f51e93b3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -36,6 +36,7 @@ export const LOG_FILTER_LABELS: Record = { [LOG_FILTER_IDS.TEAM_ID]: "Team ID", [LOG_FILTER_IDS.STATUS]: "Status", [LOG_FILTER_IDS.KEY_ALIAS]: "Key Alias", + [LOG_FILTER_IDS.USER_ID]: "User ID", [LOG_FILTER_IDS.END_USER]: "End User", [LOG_FILTER_IDS.ERROR_CODE]: "Error Code", [LOG_FILTER_IDS.ERROR_MESSAGE]: "Error Message", @@ -98,7 +99,6 @@ export function useLogFilterLogic({ userRole, userID, columnFilters, - filterByCurrentUser, activeTab, isLiveTail, startTime, @@ -112,7 +112,6 @@ export function useLogFilterLogic({ userRole: string | null; userID: string | null; columnFilters: ColumnFiltersState; - filterByCurrentUser: boolean | null; activeTab: string; isLiveTail: boolean; startTime: string; @@ -136,7 +135,6 @@ export function useLogFilterLogic({ endTime, isCustomDate, columnFilters, - filterByCurrentUser ? userID : null, sortBy, sortOrder, ], @@ -166,7 +164,7 @@ export function useLogFilterLogic({ team_id: getFilterValue(columnFilters, LOG_FILTER_IDS.TEAM_ID), request_id: getFilterValue(columnFilters, LOG_FILTER_IDS.REQUEST_ID), session_id: getFilterValue(columnFilters, LOG_FILTER_IDS.SESSION_ID), - user_id: userIdFilter ?? (filterByCurrentUser ? userID ?? undefined : undefined), + user_id: userIdFilter, end_user: getFilterValue(columnFilters, LOG_FILTER_IDS.END_USER), status_filter: getFilterValue(columnFilters, LOG_FILTER_IDS.STATUS), model_id: getFilterValue(columnFilters, LOG_FILTER_IDS.MODEL_ID), diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 095afc6110f..2bc21735248 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7622,6 +7622,26 @@ export interface paths { patch?: never; trace?: never; }; + "/management/v1/spend_logs/users": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List Spend Log Users + * @description The distinct internal users appearing in spend logs the caller can read. + */ + get: operations["list_spend_log_users_management_v1_spend_logs_users_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/mcp": { parameters: { query?: never; @@ -45939,6 +45959,46 @@ export interface operations { }; }; }; + list_spend_log_users_management_v1_spend_logs_users_get: { + parameters: { + query: { + /** @description Window start (UTC when no offset is given) */ + "filter[startTime][gte]": string; + /** @description Window end (UTC when no offset is given) */ + "filter[startTime][lte]": string; + /** @description Case-insensitive partial match on the internal user id */ + q?: string | null; + /** @description Page number */ + page?: number; + /** @description Page size */ + page_size?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["FacetListResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; aggregate_mcp_route_mcp_get: { parameters: { query?: never; diff --git a/uv.lock b/uv.lock index 1cc924e50fe..63e8e31223c 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-12T18:09:57.187615Z" +exclude-newer = "2026-08-14T18:59:56.524034Z" exclude-newer-span = "P3D" [manifest] @@ -9075,11 +9075,11 @@ wheels = [ [[package]] name = "sqlparse" -version = "0.5.5" +version = "0.6.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/90/76/437d71068094df0726366574cf3432a4ed754217b436eb7429415cf2d480/sqlparse-0.5.5.tar.gz", hash = "sha256:e20d4a9b0b8585fdf63b10d30066c7c94c5d7a7ec47c889a2d83a3caa93ff28e", size = 120815, upload-time = "2025-12-19T07:17:45.073Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5f/d3/3f06a1006f2261d1342aefb3c71eed02f5d4ca5bdbecd86ebc12ad38306e/sqlparse-0.6.0.tar.gz", hash = "sha256:113c35c75365ab9cc9c7231d68c6428fb11c085fc8e9eb1ad659b7ddbf6cd2b9", size = 178477, upload-time = "2026-08-13T19:16:06.396Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/49/4b/359f28a903c13438ef59ebeee215fb25da53066db67b305c125f1c6d2a25/sqlparse-0.5.5-py3-none-any.whl", hash = "sha256:12a08b3bf3eec877c519589833aed092e2444e68240a3577e8e26148acc7b1ba", size = 46138, upload-time = "2025-12-19T07:17:46.573Z" }, + { url = "https://files.pythonhosted.org/packages/d9/50/f00935da0ec7cbf325f8dc4f772ae46fbc7b672dd62876e73f0a94adda57/sqlparse-0.6.0-py3-none-any.whl", hash = "sha256:b861c0288ce2fa56209a9a6412d2e066ac664b3873b89c26c9d8415e8e32996f", size = 50070, upload-time = "2026-08-13T19:16:04.062Z" }, ] [[package]]