diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 65142091712..43df27ea2e2 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 31903 + "limit": 31256 }, "reportArgumentType": { "limit": 2645 @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 10214 + "limit": 10208 }, "reportFunctionMemberAccess": { "limit": 11 @@ -33,7 +33,7 @@ "limit": 227 }, "reportIncompatibleMethodOverride": { - "limit": 78 + "limit": 77 }, "reportIncompatibleVariableOverride": { "limit": 12 @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45366 + "limit": 45357 }, "reportUnknownLambdaType": { "limit": 113 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 8821736d0ff..d57c1a78f3d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -215,7 +215,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) if result: - return LiteLLM_ManagedFileTable(**result) + return LiteLLM_ManagedFileTable.model_validate(result) ## CHECK DB db_object = await self.prisma_client.db.litellm_managedfiletable.find_first( @@ -223,7 +223,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) if db_object: - return LiteLLM_ManagedFileTable(**db_object.model_dump()) + return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump()) return None async def delete_unified_file_id( @@ -349,7 +349,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if isinstance(batch.file_object, str) else batch.file_object ) - batch_obj = LiteLLMBatch(**batch_data) + batch_obj = LiteLLMBatch.model_validate(batch_data) batch_obj.id = batch.unified_object_id batch_objects.append(batch_obj) @@ -382,7 +382,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "flat_model_file_ids": {"hasSome": model_object_ids}, } ) - return [OpenAIFileObject(**file_object.file_object) for file_object in file_ids] + return [OpenAIFileObject.model_validate(file_object.file_object) for file_object in file_ids] async def check_managed_file_id_access( self, data: Dict, user_api_key_dict: UserAPIKeyAuth diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9c7bbbd3b4c..fbe70c03d74 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1182,6 +1182,7 @@ class CustomGuardrail(CustomLogger): call_type == CallTypes.completion.value or call_type == CallTypes.acompletion.value or call_type == CallTypes.anthropic_messages.value + or call_type == CallTypes.call_mcp_tool.value ): return data.get("messages") diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 60dbf7c644a..0ac22b5bb1e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -131,8 +131,8 @@ class CustomStreamWrapper: self.sent_last_chunk = False self._stream_created_time: float = time.time() - litellm_params: GenericLiteLLMParams = GenericLiteLLMParams( - **self.logging_obj.model_call_details.get("litellm_params", {}) + litellm_params: GenericLiteLLMParams = GenericLiteLLMParams.model_validate( + dict(**self.logging_obj.model_call_details.get("litellm_params", {})) ) self.merge_reasoning_content_in_choices: bool = litellm_params.merge_reasoning_content_in_choices or False self.sent_first_thinking_block = False diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index d0b0dbb070d..1a860cca5a9 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -66,7 +66,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): ) # Create ResponseReasoningItem object from the item data - reasoning_item = ResponseReasoningItem(**item_data) + reasoning_item = ResponseReasoningItem.model_validate(item_data) # Convert back to dict with exclude_none=True to exclude None fields dict_reasoning_item = reasoning_item.model_dump(exclude_none=True) @@ -346,4 +346,4 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIError raise AzureOpenAIError(message=raw_response.text, status_code=raw_response.status_code) - return ResponsesAPIResponse(**raw_response_json) + return ResponsesAPIResponse.model_validate(raw_response_json) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index df5b10b3bdc..2c5f455692c 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -1,10 +1,11 @@ import asyncio +import concurrent.futures import contextlib import os import ssl import typing import urllib.request -from typing import Any, Callable, Dict, Optional, Union +from typing import Any, Callable, ClassVar, Dict, Optional, Union import aiohttp import aiohttp.client_exceptions @@ -138,6 +139,11 @@ class LiteLLMAiohttpTransport(AiohttpTransport): Credit to: https://github.com/karpetrosyan/httpx-aiohttp for this implementation """ + # Strong references to scheduled session-close tasks. A bare + # asyncio.create_task() result may be garbage-collected before it runs, + # leaving the recycled session unclosed ("Unclosed client session"). + _background_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes + def __init__( self, client: Union[ClientSession, Callable[[], ClientSession]], @@ -164,6 +170,92 @@ class LiteLLMAiohttpTransport(AiohttpTransport): self._owns_session = True return session + @classmethod + def _on_close_task_done(cls, task: "asyncio.Task[None]") -> None: + cls._background_close_tasks.discard(task) + if task.cancelled(): + return + exc = task.exception() + if exc is not None: + verbose_logger.debug("Error closing recycled aiohttp session: %s", exc) + + @staticmethod + def _on_threadsafe_close_done(future: "concurrent.futures.Future[None]") -> None: + if future.cancelled(): + return + exc = future.exception() + if exc is not None: + verbose_logger.debug("Error closing recycled aiohttp session on its own loop: %s", exc) + + @staticmethod + def _mark_connector_closed(session: ClientSession) -> None: + """Synchronously dispose a session whose event loop is gone. + + An async close can no longer run on a closed loop. BaseConnector._close + is the same synchronous teardown aiohttp's own finalizer (__del__) + uses: it is guarded for closed loops, releases pooled connections, and + flips the flags that ClientSession.closed / BaseConnector.closed read - + so no "Unclosed client session" / "Unclosed connector" warnings reach + the event-loop exception handler at garbage collection. + """ + connector = getattr(session, "_connector", None) + close_sync = getattr(connector, "_close", None) + if not callable(close_sync): + return + try: + close_sync() + except (RuntimeError, AttributeError, OSError) as e: + verbose_logger.debug("Best-effort connector close failed: %s", e) + + def _close_recycled_session(self, session: ClientSession) -> None: + """Deterministically dispose a ClientSession this transport is replacing. + + Covers the three lifecycles a recycled session can be in: + - its loop is the current running loop: schedule an async close and keep + a strong reference to the task until it completes; + - its loop is still running elsewhere (e.g. another thread): hand the + close to that loop thread-safely; + - its loop is stopped or closed, or there is no running loop: fall + back to the synchronous finalizer-safe teardown. + """ + if session.closed: + return + + session_loop = getattr(session, "_loop", None) + try: + current_loop: Optional[asyncio.AbstractEventLoop] = asyncio.get_running_loop() + except RuntimeError: + current_loop = None + + if session_loop is not None and session_loop is not current_loop: + if not session_loop.is_closed() and session_loop.is_running(): + # The session's loop is running somewhere else (e.g. another + # thread): closing from here would touch that loop's internals + # unsafely; hand the close to its own loop. + try: + future = asyncio.run_coroutine_threadsafe(session.close(), session_loop) + except RuntimeError as e: # loop shut down between the checks + verbose_logger.debug("Threadsafe session close failed: %s", e) + self._mark_connector_closed(session) + else: + future.add_done_callback(self._on_threadsafe_close_done) + return + + # Foreign loop that is stopped or closed: an async close can no + # longer run there, and running it on the current loop would touch + # another loop's internals. Dispose synchronously instead. + self._mark_connector_closed(session) + return + + if current_loop is None: + self._mark_connector_closed(session) + return + + task = current_loop.create_task(session.close()) + cls = type(self) + cls._background_close_tasks.add(task) + task.add_done_callback(cls._on_close_task_done) + def _get_valid_client_session(self) -> ClientSession: """ Helper to get a valid ClientSession for the current event loop. @@ -193,21 +285,25 @@ class LiteLLMAiohttpTransport(AiohttpTransport): # Close old session to prevent leaks old_session = self.client try: - if self._owns_session and not old_session.closed: - try: - asyncio.create_task(old_session.close()) - except RuntimeError: - # Different event loop - can't schedule task, rely on GC - verbose_logger.debug("Old session from different loop, relying on GC") + if self._owns_session: + self._close_recycled_session(old_session) except Exception as e: verbose_logger.debug(f"Error closing old session: {e}") # Create a new session in the current event loop self.client = self._rebuild_session() - except (RuntimeError, AttributeError): - # If we can't check the loop or session is invalid, recreate it + except (RuntimeError, AttributeError) as e: + # If we can't check the loop or session is invalid, recreate it, + # but still dispose of the session being replaced. + old_session = self.client + if self._owns_session: + try: + self._close_recycled_session(old_session) + except (RuntimeError, AttributeError, OSError) as close_error: + verbose_logger.debug(f"Error closing old session: {close_error}") self.client = self._rebuild_session() + verbose_logger.debug(f"Error checking session loop, created new session: {e}") return self.client @@ -301,7 +397,14 @@ class LiteLLMAiohttpTransport(AiohttpTransport): # Handle the case where session was closed between our check and actual use if "Session is closed" in str(e): verbose_logger.debug(f"Session closed during request, retrying with new session: {e}") - # Force creation of a new session + # Dispose of the session that actually faulted. Do NOT read + # self.client here: a concurrent task may already have + # replaced it with a healthy session that must stay open. + # Guarded by isinstance: factory-injected sessions may be + # duck-typed test doubles without a close() coroutine. + # Read _owns_session before _rebuild_session() claims ownership. + if self._owns_session and isinstance(client_session, ClientSession): + self._close_recycled_session(client_session) self.client = self._rebuild_session() client_session = self.client diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ec1301e5923..d6acbaae434 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6184,10 +6184,12 @@ class BaseLLMHTTPHandler: import websockets from websockets.asyncio.client import ClientConnection - litellm_params = GenericLiteLLMParams( - api_base=api_base, - api_key=api_key, - **kwargs, + litellm_params = GenericLiteLLMParams.model_validate( + { + "api_base": api_base, + "api_key": api_key, + **kwargs, + } ) headers = responses_api_provider_config.validate_environment( headers={}, diff --git a/litellm/llms/manus/responses/transformation.py b/litellm/llms/manus/responses/transformation.py index 0db53f90330..e6fbbac0563 100644 --- a/litellm/llms/manus/responses/transformation.py +++ b/litellm/llms/manus/responses/transformation.py @@ -217,7 +217,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): raw_response_json["id"] = f"unknown-{uuid.uuid4().hex[:8]}" try: - response = ResponsesAPIResponse(**raw_response_json) + response = ResponsesAPIResponse.model_validate(raw_response_json) except Exception: verbose_logger.debug(f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct") response = ResponsesAPIResponse.model_construct(**raw_response_json) @@ -305,7 +305,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): raw_response_json["id"] = f"unknown-{uuid.uuid4().hex[:8]}" try: - response = ResponsesAPIResponse(**raw_response_json) + response = ResponsesAPIResponse.model_validate(raw_response_json) except Exception: verbose_logger.debug(f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct") response = ResponsesAPIResponse.model_construct(**raw_response_json) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index cbca57c5e62..c0da1536c42 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( _get_image_mime_type_from_url, ) from litellm.litellm_core_utils.prompt_templates.factory import ( + _get_thought_signature_from_tool, convert_generic_image_chunk_to_openai_image_obj, convert_to_anthropic_image_obj, convert_to_gemini_tool_call_invoke, @@ -635,6 +636,58 @@ def check_if_part_exists_in_parts(parts: List[PartType], part: PartType, exclude return False +def _collect_tool_call_thought_signatures( + assistant_msg: ChatCompletionAssistantMessage, +) -> frozenset[str]: + """Thought signatures already carried by this message's tool-call parts. + + Gemini returns each thoughtSignature on exactly one part. When the signed + part is a function call, the signature is replayed on that tool-call part + by convert_to_gemini_tool_call_invoke, so attaching the same signature to + the text part as well would send two copies and double-bill the previous + turn's reasoning tokens on gemini-3 and newer models. + + Detection deliberately calls _get_thought_signature_from_tool without the + model argument: with a gemini-3 model that helper synthesizes a dummy + signature for unsigned tool calls, which must not suppress a real + text-part signature (e.g. replaying gemini-2.5 history to a newer model). + """ + signatures: tuple[str, ...] = () + + tool_calls = assistant_msg.get("tool_calls") + if isinstance(tool_calls, list): + for tool in tool_calls: + if not isinstance(tool, dict): + continue + signature = _get_thought_signature_from_tool(tool) + if signature: + signatures += (signature,) + + function_call = assistant_msg.get("function_call") + if isinstance(function_call, dict): + signature = _get_thought_signature_from_tool({"function": function_call}) + if signature: + signatures += (signature,) + + provider_specific_fields = assistant_msg.get("provider_specific_fields") + if not isinstance(provider_specific_fields, dict): + return frozenset(signatures) + + invocations = provider_specific_fields.get("server_side_tool_invocations") + if not isinstance(invocations, list): + return frozenset(signatures) + + for invocation in invocations: + if not isinstance(invocation, dict): + continue + for key in ("thought_signature", "response_thought_signature"): + invocation_signature = invocation.get(key) + if isinstance(invocation_signature, str) and invocation_signature: + signatures += (invocation_signature,) + + return frozenset(signatures) + + def _gemini_convert_messages_with_history( messages: List[AllMessageValues], model: Optional[str] = None, @@ -868,8 +921,18 @@ def _gemini_convert_messages_with_history( if provider_specific_fields and isinstance(provider_specific_fields, dict): thought_signatures = provider_specific_fields.get("thought_signatures") - # If we have thought signatures, add them to the part - if thought_signatures and isinstance(thought_signatures, list) and len(thought_signatures) > 0: + # A signature that is already carried by one of this message's + # tool-call parts must not be attached to the text part too: + # Gemini bills every replayed copy as the previous turn's full + # reasoning token count on gemini-3 and newer models + tool_call_signatures = _collect_tool_call_thought_signatures(assistant_msg) + + if ( + thought_signatures + and isinstance(thought_signatures, list) + and len(thought_signatures) > 0 + and thought_signatures[0] not in tool_call_signatures + ): # Use the first signature for the text part (Gemini expects one signature per part) assistant_content.append( PartType( diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index dc796837277..81983cc62fd 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -67,6 +67,18 @@ def _as_list(values: Sequence[str] | None) -> list[str] | None: # mutable-ok: r return None if values is None else list(values) +class UnloadableEntitlementError(Exception): + """A principal's row NAMES an ``object_permission_id`` whose contents could not be read. + + Raised only where there is POSITIVE evidence an entitlement exists, so every caller must DENY + rather than fall back to "this level places no restriction": a ceiling we know exists but cannot + read would otherwise silently widen the caller for as long as the fault lasts. + + Deliberately distinct from a lookup that fails before the principal's entitlement is known at + all. Not knowing whether someone is entitled is the state that existed before the level did, so + it places no ceiling; denying there would refuse MCP to every caller during a cold-cache fault.""" + + def _parse_mcp_server_names_from_path(path: str, mcp_servers_header: Optional[List[str]] = None) -> Optional[List[str]]: """Resolve the single MCP server name a cold-start passthrough bypass may target. Delegates parsing to @@ -292,6 +304,22 @@ class MCPRequestHandler: 3. Header extraction and validation Utilizes the main `user_api_key_auth` function to validate authentication + + Entitlement-fault contract (``get_allowed_mcp_servers`` / ``get_allowed_tools_for_server``) + ------------------------------------------------------------------------------------------ + Every level (key, team, end user, agent, org) answers "which servers/tools does this level + permit", and a level that answers nothing places no restriction. A lookup FAULT is not that + answer, and the two callers resolve it differently on purpose: + + - A keyless gateway-admitted subject fails CLOSED on any fault at any level. Each of its grant + sources is resolved independently and unioned, so a fault that returned "no restriction" would + win the union as allow-all, and its per-source org ceiling is the ONLY org bound it has. + - Key auth fails closed only where there is POSITIVE evidence an entitlement exists: a principal + row that NAMES an ``object_permission_id`` we cannot load is a known entitlement with unknown + contents (``UnloadableEntitlementError`` -> deny). A fault so early we cannot tell whether the + principal is entitled at all leaves no ceiling, because that is the state that existed before + the level did; denying there would refuse MCP to every caller, most of whom have no entitlement + configured, for the duration of a cold-cache or DB fault. """ LITELLM_API_KEY_HEADER_NAME_PRIMARY = SpecialHeaders.custom_litellm_api_key.value @@ -1348,6 +1376,9 @@ class MCPRequestHandler: has an explicit MCP server list, the combined key/team/end_user/agent result is capped to that list. If the org has no list, no extra restriction is applied. + A level that cannot answer is NOT a level that permits everything; see the class docstring + for how each caller shape resolves an entitlement fault. + Returns: List[str]: List of allowed MCP servers by server id """ @@ -1478,7 +1509,12 @@ class MCPRequestHandler: return list(set(allowed_mcp_servers)) except Exception as e: - verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}") + if isinstance(e, UnloadableEntitlementError): + # A ceiling we KNOW exists and cannot read. Denying is the only answer that does not + # widen this caller past what an operator configured, for both caller shapes. + verbose_logger.warning(f"Denying MCP access, entitlement unreadable: {str(e)}") + else: + verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}") return [] @staticmethod @@ -1491,11 +1527,15 @@ class MCPRequestHandler: """Cap the resolved server list by this caller's org ceiling: an explicit org list intersects lower-level restrictions (else becomes the ceiling); no org or an empty list leaves it unchanged. - ``keyless_source`` governs both divergences for a keyless admitted source. An UNRESOLVABLE ceiling - fails CLOSED for it (its only org bound is this ceiling, so dropping it on a fault would escalate a - cross-org user) while a key stays fail-open. And an org list may only ever INTERSECT a source (the - admitted model unions grants, so a ceiling must not become one), whereas for a key it may - substitute, that being the key ceiling model.""" + ``keyless_source`` governs both divergences for a keyless admitted source. An INDETERMINATE ceiling + (we cannot tell whether the org restricts at all) fails CLOSED for it (its only org bound is this + ceiling, so dropping it on a fault would escalate a cross-org user) while a key stays fail-open. And + an org list may only ever INTERSECT a source (the admitted model unions grants, so a ceiling must not + become one), whereas for a key it may substitute, that being the key ceiling model. + + The fail-open arm is reached only for an INDETERMINATE fault: a ceiling the org NAMES but that + cannot be read raises out of ``_get_allowed_mcp_servers_for_org`` and never arrives here as + ``None``, so key auth cannot silently shed a ceiling an operator did configure.""" if not (user_api_key_auth and user_api_key_auth.org_id): return allowed_mcp_servers allowed_mcp_servers_for_org = await MCPRequestHandler._get_allowed_mcp_servers_for_org(user_api_key_auth) @@ -1900,12 +1940,19 @@ class MCPRequestHandler: ) except Exception as e: - verbose_logger.warning(f"Failed to get allowed tools for server: {str(e)}") + # An entitlement known to exist but unreadable denies for BOTH caller shapes, so [] rather + # than the None (allow-all) key auth gets for an indeterminate fault. + unreadable_entitlement = isinstance(e, UnloadableEntitlementError) + if unreadable_entitlement: + verbose_logger.warning(f"Denying MCP tools, entitlement unreadable: {str(e)}") + else: + verbose_logger.warning(f"Failed to get allowed tools for server: {str(e)}") # Fail CLOSED for a keyless admitted subject: ANY error must deny the server's tools ([]), # not collapse to allow-all (None); key/JWT auth keeps its prior allow-all-on-error. Both # keyless_source AND the marker are needed: each source resolves through an UNMARKED auth, so # without keyless_source a fault under a source returns None and wins the union as allow-all. - return [] if (keyless_source or _is_mcp_admitted_user_subject(user_api_key_auth)) else None + deny_all = unreadable_entitlement or keyless_source or _is_mcp_admitted_user_subject(user_api_key_auth) + return [] if deny_all else None @staticmethod async def _apply_agent_and_org_tool_ceilings( @@ -1944,7 +1991,9 @@ class MCPRequestHandler: try: org_obj_perm = await MCPRequestHandler._get_org_object_permission(user_api_key_auth) except Exception as e: # noqa: BLE001 # unresolvable org ceiling, decided per caller shape - if keyless_source: + # A ceiling the org NAMES but that cannot be read denies at every caller shape; only an + # INDETERMINATE fault (we cannot tell whether a ceiling exists) keeps key auth open. + if keyless_source or isinstance(e, UnloadableEntitlementError): raise verbose_logger.warning( f"MCP org tool ceiling unresolvable for org_id={user_api_key_auth.org_id!r}; " @@ -2275,18 +2324,54 @@ class MCPRequestHandler: verbose_logger.warning(f"Failed to get allowed MCP servers for team: {str(e)}") return [] + @staticmethod + async def _load_named_object_permission( + principal: str, + object_permission_id: str, + prisma_client: "PrismaClient", + user_api_key_auth: UserAPIKeyAuth, + ) -> LiteLLM_ObjectPermissionTable: + """Load the object permission a principal's row NAMES, or raise ``UnloadableEntitlementError``. + + The single place that fault is minted, so end user, agent and org cannot drift on what counts + as "known entitlement, unknown contents". ``get_object_permission`` answers None for both an + absent row and a failed read, and neither is evidence the principal is unrestricted: the link + proves an entitlement was configured, so both must deny.""" + from litellm.proxy.auth.auth_checks import get_object_permission + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + + unloadable = UnloadableEntitlementError( + f"{principal} names object_permission_id {object_permission_id!r} which could not be loaded" + ) + try: + object_permission = await get_object_permission( + object_permission_id=object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # a named entitlement we cannot read denies, whatever the read failed with + raise unloadable from e + if object_permission is None: + raise unloadable + return object_permission + @staticmethod async def _get_org_object_permission( user_api_key_auth: Optional[UserAPIKeyAuth] = None, - ): + ) -> LiteLLM_ObjectPermissionTable | None: """ Get org object_permission via the established ``get_org_object`` / ``get_object_permission`` helpers so MCP requests share the same ``user_api_key_cache`` entries as the rest of the proxy. + + ``None`` means the org places NO ceiling: no ``org_id``, no DB, or an org row naming no + permission. A row that NAMES one it cannot load raises ``UnloadableEntitlementError``; + every other lookup failure propagates as itself, leaving the ceiling merely unresolved. """ from litellm.proxy.auth.auth_checks import ( OrganizationNotFoundError, - get_object_permission, get_org_object, ) from litellm.proxy.proxy_server import ( @@ -2322,31 +2407,29 @@ class MCPRequestHandler: if org_obj is None or not org_obj.object_permission_id: return None - # The org NAMES a permission; failing to read it is INDETERMINATE and must not collapse into the - # None that means "no ceiling". Raise and let each caller pick fail-open or fail-closed. - object_permission = await get_object_permission( + # The org NAMES a permission; failing to read it is a KNOWN ceiling with unknown contents and + # must not collapse into the None that means "no ceiling". Raising denies at every caller shape. + return await MCPRequestHandler._load_named_object_permission( + principal=f"org {user_api_key_auth.org_id!r}", object_permission_id=org_obj.object_permission_id, prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + user_api_key_auth=user_api_key_auth, ) - if object_permission is None: - raise ValueError( - f"org {user_api_key_auth.org_id!r} names object_permission_id " - f"{org_obj.object_permission_id!r} which could not be loaded" - ) - return object_permission @staticmethod async def _get_allowed_mcp_servers_for_org( user_api_key_auth: Optional[UserAPIKeyAuth] = None, - ) -> List[str]: + ) -> list[str] | None: """ Get allowed MCP servers for an organization. Returns the MCP servers from the org's object_permission. - An empty result means the org places no restriction (allow-all from this level). + An empty result means the org places no restriction (allow-all from this level), ``None`` + that the ceiling could not be resolved, which the caller decides per shape. + + A ceiling the org NAMES but we cannot read is neither: it raises out of here so both caller + shapes deny, because dropping a ceiling known to exist is exactly the silent widening the + level is there to prevent. """ try: object_permissions = await MCPRequestHandler._get_org_object_permission(user_api_key_auth) @@ -2374,34 +2457,28 @@ class MCPRequestHandler: except Exception as e: # None = ceiling UNRESOLVED, distinct from [] = org places no restriction. Collapsing them # let a DB fault silently drop a ceiling; the caller picks fail-open/closed from this signal. + # A NAMED-but-unreadable ceiling is a stronger fact than "unresolved" and denies everywhere. + if isinstance(e, UnloadableEntitlementError): + raise verbose_logger.warning(f"Failed to get allowed MCP servers for org: {str(e)}") return None @staticmethod - async def _get_allowed_mcp_servers_for_end_user( - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - ) -> List[str]: - """ - Get allowed MCP servers for an end user. + async def _get_end_user_object_permission( + user_api_key_auth: UserAPIKeyAuth, + prisma_client: "PrismaClient", + ) -> LiteLLM_ObjectPermissionTable | None: + """The end user's own object_permission, or ``None`` when this level places no restriction. - Returns the MCP servers from the end_user's object_permission. - """ + ``None`` covers an end user row that is absent or names no permission, and an end user we + could not resolve at all (``get_end_user_object`` answers None for an absent row AND for a + failed read, so this level genuinely cannot tell those apart). A row that DOES name a + permission we cannot load raises ``UnloadableEntitlementError``: the link is positive + evidence of an entitlement, so its contents may not be assumed empty.""" from litellm.proxy.auth.auth_checks import get_end_user_object - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - if not user_api_key_auth or not user_api_key_auth.end_user_id: - return [] - - if prisma_client is None: - verbose_logger.debug("prisma_client is None") - return [] + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache try: - # Use optimized get_end_user_object function with caching end_user_obj = await get_end_user_object( end_user_id=user_api_key_auth.end_user_id, prisma_client=prisma_client, @@ -2410,29 +2487,65 @@ class MCPRequestHandler: proxy_logging_obj=proxy_logging_obj, route="/mcp", ) + except Exception as e: # noqa: BLE001 # entitlement unknown, not known-absent: no ceiling, as before this level + verbose_logger.warning(f"Failed to resolve end_user for MCP permissions: {str(e)}") + return None - if end_user_obj is None or end_user_obj.object_permission is None: - return [] + if end_user_obj is None: + return None + if end_user_obj.object_permission is not None: + return end_user_obj.object_permission + if not end_user_obj.object_permission_id: + return None + # The row NAMES a permission the relation did not carry. One shared (cached) lookup decides + # whether it is readable; an unreadable one denies rather than reading as "no restriction". + return await MCPRequestHandler._load_named_object_permission( + principal=f"end user {user_api_key_auth.end_user_id!r}", + object_permission_id=end_user_obj.object_permission_id, + prisma_client=prisma_client, + user_api_key_auth=user_api_key_auth, + ) + @staticmethod + async def _get_allowed_mcp_servers_for_end_user( + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + ) -> List[str]: + """ + Get allowed MCP servers for an end user. + + Returns the MCP servers from the end_user's object_permission; an empty result means this + level places no restriction. An entitlement the end user row NAMES but that cannot be read + raises ``UnloadableEntitlementError`` out of here so the resolver denies. + """ + from litellm.proxy.proxy_server import prisma_client + + if not user_api_key_auth or not user_api_key_auth.end_user_id: + return [] + + if prisma_client is None: + verbose_logger.debug("prisma_client is None") + return [] + + object_permission = await MCPRequestHandler._get_end_user_object_permission(user_api_key_auth, prisma_client) + if object_permission is None: + return [] + + try: # Permission entries may be server_ids OR names/aliases — expand to ids. from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - direct_mcp_servers = global_mcp_server_manager.expand_permission_list( - end_user_obj.object_permission.mcp_servers or [] - ) + direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permission.mcp_servers or []) # Get MCP servers from access groups access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups( - end_user_obj.object_permission.mcp_access_groups or [] + object_permission.mcp_access_groups or [] ) # servers referenced in tool permissions should also be accessible tool_perm_servers = list( - global_mcp_server_manager.expand_tool_permissions( - end_user_obj.object_permission.mcp_tool_permissions - ).keys() + global_mcp_server_manager.expand_tool_permissions(object_permission.mcp_tool_permissions).keys() ) # Combine all lists @@ -2643,22 +2756,51 @@ class MCPRequestHandler: # don't re-query the DB on every MCP request for that agent. _AGENT_NO_PERMISSION_SENTINEL = "__agent_no_mcp_permission__" + @staticmethod + async def _agent_object_permission_id(agent_id: str, prisma_client: "PrismaClient") -> str | None: + """The permission row this agent's row links to, or ``None`` when it links none. + + Caches the link (with a sentinel for "links none") so an agent without an entitlement costs + no DB read per MCP request. A read that fails also answers ``None``: not knowing whether the + agent is entitled is the state that existed before this level, so it places no ceiling. Only + a link we DID resolve can make the caller deny.""" + from litellm.proxy.proxy_server import user_api_key_cache + + cache_key = f"agent_object_permission_id:{agent_id}" + try: + cached: object = await user_api_key_cache.async_get_cache(key=cache_key) + if cached == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL: + return None + if isinstance(cached, str) and cached: + return cached + agent_row = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) + linked: object = getattr(agent_row, "object_permission_id", None) if agent_row is not None else None + object_permission_id = linked if isinstance(linked, str) and linked else None + await user_api_key_cache.async_set_cache( + key=cache_key, + value=object_permission_id or MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL, + ttl=get_management_object_ttl(user_api_key_cache), + ) + return object_permission_id + except Exception as e: # noqa: BLE001 # entitlement unknown, not known-absent: no ceiling, as before this level + verbose_logger.warning(f"Failed to resolve object_permission_id for agent {agent_id!r}: {str(e)}") + return None + @staticmethod async def _get_agent_object_permission( user_api_key_auth: Optional[UserAPIKeyAuth] = None, - ): + ) -> LiteLLM_ObjectPermissionTable | None: """ Get agent object_permission via the established ``get_object_permission`` helper. Caches the ``agent_id -> object_permission_id`` mapping so we avoid re-reading the agent row on every request, and reuses the shared ``object_permission_id`` cache populated by the org / team / key paths. + + ``None`` means the agent places NO restriction: no ``agent_id``, no DB, or an agent linking + no permission. An agent that LINKS one we cannot load raises ``UnloadableEntitlementError``, + since a known entitlement with unknown contents must deny rather than read as unrestricted. """ - from litellm.proxy.auth.auth_checks import get_object_permission - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.proxy_server import prisma_client if not user_api_key_auth or not user_api_key_auth.agent_id: return None @@ -2668,40 +2810,17 @@ class MCPRequestHandler: return None agent_id = user_api_key_auth.agent_id - cache_key = f"agent_object_permission_id:{agent_id}" - - try: - object_permission_id: Optional[str] = await user_api_key_cache.async_get_cache(key=cache_key) - - if object_permission_id == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL: - return None - - if object_permission_id is None: - agent_row = await AgentsRepository(prisma_client).table.find_unique( - where={"agent_id": agent_id}, - ) - object_permission_id = ( - getattr(agent_row, "object_permission_id", None) if agent_row is not None else None - ) - await user_api_key_cache.async_set_cache( - key=cache_key, - value=object_permission_id or MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL, - ttl=get_management_object_ttl(user_api_key_cache), - ) - if not object_permission_id: - return None - - return await get_object_permission( - object_permission_id=object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: - verbose_logger.warning(f"Failed to get agent object permission: {str(e)}") + object_permission_id = await MCPRequestHandler._agent_object_permission_id(agent_id, prisma_client) + if object_permission_id is None: return None + return await MCPRequestHandler._load_named_object_permission( + principal=f"agent {agent_id!r}", + object_permission_id=object_permission_id, + prisma_client=prisma_client, + user_api_key_auth=user_api_key_auth, + ) + @staticmethod async def _get_allowed_mcp_servers_for_agent( user_api_key_auth: Optional[UserAPIKeyAuth] = None, @@ -2711,7 +2830,9 @@ class MCPRequestHandler: Get allowed MCP servers for an agent (from the agent's object_permission). Returns the MCP servers from the agent's object_permission. - If agent has no object_permission, returns [] (no extra restriction). + If agent has no object_permission, returns [] (no extra restriction). An entitlement the + agent LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here so the + resolver denies. Args: user_api_key_auth: User auth with agent_id @@ -2721,13 +2842,13 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.agent_id: return [] - try: - obj_perm = agent_object_permission - if obj_perm is None: - obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) - if obj_perm is None: - return [] + obj_perm = agent_object_permission + if obj_perm is None: + obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + if obj_perm is None: + return [] + try: direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or [] if isinstance(direct_mcp_servers, str): direct_mcp_servers = [] @@ -2757,7 +2878,9 @@ class MCPRequestHandler: ) -> Optional[List[str]]: """ Get allowed tool names for a server from the agent's object_permission. - Returns None if agent has no tool restrictions for this server. + Returns None if agent has no tool restrictions for this server. An entitlement the agent + LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here, which the + tool resolver turns into deny-all for the server rather than an unrestricted tool list. Args: server_id: Server ID to check permissions for @@ -2768,13 +2891,13 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.agent_id: return None - try: - obj_perm = agent_object_permission - if obj_perm is None: - obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) - if obj_perm is None: - return None + obj_perm = agent_object_permission + if obj_perm is None: + obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + if obj_perm is None: + return None + try: mcp_tool_permissions = getattr(obj_perm, "mcp_tool_permissions", None) if not mcp_tool_permissions or not isinstance(mcp_tool_permissions, dict): return None diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 89c83459524..e3f7352e8ca 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -906,6 +906,20 @@ def _extract_upstream_auth_failure( return upstream_auth_challenge(exc) +def _obo_retry_applies(server: MCPServer, subject_token: str | None) -> bool: + """Whether an upstream 401/403 should invalidate the minted credential and retry once. + + ``oauth2_token_exchange`` can only mint from an inbound subject token, so with no token there is + nothing to re-mint and the plain single call is correct. ``oauth2_id_jag`` also sources its + subject from the identity assertion stored for the user at SSO login, so it qualifies whether or + not the caller presented a token of its own; gating it on the inbound token would leave a + store-sourced bearer un-invalidated and replayed until its TTL. + """ + if server.auth_type == MCPAuth.oauth2_id_jag: + return True + return server.auth_type == MCPAuth.oauth2_token_exchange and bool(subject_token) + + def _warn_on_server_name_fields( *, server_id: str, @@ -4422,13 +4436,18 @@ class MCPServerManager: arguments: dict[str, Any], server_name: str, user_api_key_auth: Optional[UserAPIKeyAuth], - proxy_logging_obj: ProxyLogging, + proxy_logging_obj: ProxyLogging | None, server: MCPServer, raw_headers: Optional[dict[str, str]] = None, ) -> dict[str, Any]: """ Run pre-call checks and guardrail hooks for an MCP tool call. + Authorization runs unconditionally; only the guardrail hooks, which are + dispatched through ``proxy_logging_obj``, depend on a logger being + present. An absent logger must never be able to turn an authorization + decision into a no-op. + Returns a dict that may contain: - "arguments": hook-modified tool arguments (only if changed) - "extra_headers": headers injected by pre_mcp_call guardrail hooks @@ -4456,6 +4475,10 @@ class MCPServerManager: server=server, ) + hook_result: dict[str, Any] = {} + if proxy_logging_obj is None: + return hook_result + # Extract incoming Bearer token from raw request headers so # guardrails like MCPJWTSigner can verify + re-sign it (FR-5). normalized_raw = {k.lower(): v for k, v in (raw_headers or {}).items()} @@ -4485,7 +4508,6 @@ class MCPServerManager: # Convert to LLM format for existing guardrail compatibility synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(mcp_request_obj, pre_hook_kwargs) - hook_result: dict[str, Any] = {} try: # Use standard pre_call_hook modified_data = await proxy_logging_obj.pre_call_hook( @@ -4778,7 +4800,7 @@ class MCPServerManager: arguments=arguments, ) - if mcp_server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag) and subject_token: + if _obo_retry_applies(mcp_server, subject_token): # OBO / ID-JAG: the exchanged token may have been revoked/rotated upstream since it was # cached, so an upstream 401 gets one invalidate + re-mint + retry. Gated to these modes; # all others keep the plain single call below. @@ -5111,19 +5133,17 @@ class MCPServerManager: # Allow validation and modification of tool calls before execution # Using standard pre_call_hook ######################################################### - hook_result: dict[str, Any] = {} - if proxy_logging_obj: - hook_result = await self.pre_call_tool_check( - name=name, - arguments=arguments, - server_name=server_name, - user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, - server=mcp_server, - raw_headers=raw_headers, - ) - if "arguments" in hook_result: - arguments = hook_result["arguments"] + hook_result: dict[str, Any] = await self.pre_call_tool_check( + name=name, + arguments=arguments, + server_name=server_name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=mcp_server, + raw_headers=raw_headers, + ) + if "arguments" in hook_result: + arguments = hook_result["arguments"] # Prepare tasks for during hooks tasks = [] diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 69984a56311..b70db64ba94 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -10,19 +10,24 @@ at runtime instead of returning `None`. `none`, `api_key` (shared-key source), and `passthrough` (forwards the caller's own inbound token) are live, as is `authorization_code`, which reads the user's token from the injected `OAuthTokenStore`, `token_exchange`, which swaps the caller's inbound token through the injected -`TokenExchanger`, and `client_credentials`, which mints and caches the gateway's M2M token through -the injected `ClientCredentialsTokenSource`. The remaining arms are `not_implemented` stubs that -each land in a follow-up PR with their seam. Pure v2: no imports from v1. +`TokenExchanger`, `client_credentials`, which mints and caches the gateway's M2M token through the +injected `ClientCredentialsTokenSource`, and `id_jag`, which runs the two-leg identity-assertion +grant against a subject token taken from the request or from the injected `SSOAssertionStore`. The +remaining arms are `not_implemented` stubs that each land in a follow-up PR with their seam. Pure +v2: no imports from v1. """ from __future__ import annotations import hashlib +from datetime import datetime, timezone from functools import partial import httpx from typing_extensions import assert_never +from litellm._logging import verbose_proxy_logger + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( ClientCredentialsBearerAuth, ClientCredentialsTokenSource, @@ -41,6 +46,12 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( Ok, Result, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + AssertionStoreUnavailable, + DbSSOAssertionStore, + SSOAssertionStore, + SSOIdentityAssertion, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( ExchangedToken, ExchangedTokenCache, @@ -111,12 +122,14 @@ class UpstreamCredentialProvider: token_endpoint: TokenEndpointClient | None = None, exchanged_tokens: ExchangedTokenCache | None = None, client_credentials_source: ClientCredentialsTokenSource | None = None, + sso_assertion_store: SSOAssertionStore | None = None, ) -> None: self._oauth_token_store: OAuthTokenStore = oauth_token_store or _NullOAuthTokenStore() self._token_exchanger: TokenExchanger = token_exchanger or _NullTokenExchanger() self._token_endpoint: TokenEndpointClient = token_endpoint or TokenEndpointClient() self._exchanged_tokens: ExchangedTokenCache = exchanged_tokens or ExchangedTokenCache() self._client_credentials_source = client_credentials_source or ClientCredentialsTokenSource() + self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or DbSSOAssertionStore() async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: @@ -171,15 +184,73 @@ class UpstreamCredentialProvider: assert_never(config.key_source) async def _id_jag(self, subject: Subject, server: ServerSpec, config: IdJagConfig) -> Result[httpx.Auth, CredError]: - if subject.inbound_token is None: + match await self._id_jag_subject_token(subject): + case Error(err): + return Error(err) + case Ok(subject_token): + return await self._id_jag_exchange(subject, subject_token, server, config) + + async def _id_jag_subject_token(self, subject: Subject) -> Result[str, CredError]: + """The identity token ID-JAG leg 1 asserts, from the request or from the SSO login it was captured at. + + A caller that presents its own IdP identity token wins: that is the strongest available + assertion of who is calling. Otherwise the subject is the assertion captured for this user + at LiteLLM SSO login, which is what lets an agent holding a brokered LiteLLM credential + reach an upstream as the user it was issued for. The user is always taken from the + authenticated principal, never from a caller-supplied field, so no caller can select whose + identity is asserted upstream. + + Every miss is ``precondition_required`` (412) rather than a fall-through to a weaker + credential: ID-JAG exists to assert a specific user, so a missing subject has no safe + substitute. A store outage is the one exception: it is ``upstream_unavailable`` (503), not + 412, because the user has nothing to fix by signing in again, and it is a value rather than + a raised error so a DB blip cannot 500 the egress or the upstream-401 retry. + """ + if subject.inbound_token is not None: + return Ok(subject.inbound_token.get_secret_value()) + if not subject.subject_id: return Error( CredError.of_precondition_required( - "ID-JAG requires a caller identity token; it asserts the calling " - "user's identity upstream and cannot use a static credential." + "ID-JAG requires an identified caller; this request carries neither an " + "identity token nor a resolved LiteLLM user." ) ) - token = subject.inbound_token.get_secret_value() - cache_key = _id_jag_cache_key(token, server.server_id, config) + try: + assertion = await self._sso_assertion_store.fetch(subject.subject_id) + except AssertionStoreUnavailable as exc: + # The driver's message can name hosts, schemas or connection details, and this summary + # is returned to the caller verbatim as a 503 body. Operators get it from the log. + verbose_proxy_logger.warning( + "ID-JAG: the IdP identity assertion store is unreachable for user_id=%s: %s", + subject.subject_id, + exc, + ) + return Error( + CredError.of_upstream_unavailable( + "The IdP identity assertion store is unreachable, so ID-JAG cannot resolve a subject." + ) + ) + if assertion is None: + return Error( + CredError.of_precondition_required( + "ID-JAG requires an IdP identity assertion for this user and none is stored. " + "Sign in through LiteLLM SSO so the gateway captures one." + ) + ) + if _assertion_expired(assertion, datetime.now(timezone.utc)): + return Error( + CredError.of_precondition_required( + "The stored IdP identity assertion for this user has expired. Sign in through " + "LiteLLM SSO again to capture a current one." + ) + ) + return Ok(assertion.id_token.get_secret_value()) + + async def _id_jag_exchange( + self, subject: Subject, token: str, server: ServerSpec, config: IdJagConfig + ) -> Result[httpx.Auth, CredError]: + slot = _id_jag_slot_key(subject, server) + fingerprint = _id_jag_fingerprint(token, server.server_id, config) async def _exchange() -> Result[ExchangedToken, CredError]: leg1_params = { @@ -211,7 +282,7 @@ class UpstreamCredentialProvider: config.client_auth, ) - match await self._exchanged_tokens.get_or_compute(cache_key, _exchange): + match await self._exchanged_tokens.get_or_compute(slot, _exchange, fingerprint=fingerprint): case Ok(access_token): return Ok(StaticHeaderAuth(f"Bearer {access_token}")) case Error(err): @@ -273,17 +344,27 @@ class UpstreamCredentialProvider: re-mintable cached credential here; `client_credentials` recovers inside its own auth flow (`ClientCredentialsBearerAuth` retries the 401'd request once with a fresh token), and other modes are a no-op. + + `id_jag` evicts by a slot key derived from the principal, so it needs no lookup against the + assertion store on this path; the fingerprint stored beside the entry is what keeps a slot + shared between callers safe. """ - if subject.inbound_token is None: - return - if isinstance(server.config, TokenExchangeConfig): + if isinstance(server.config, IdJagConfig): + self._invalidate_id_jag(subject, server) + elif isinstance(server.config, TokenExchangeConfig) and subject.inbound_token is not None: await self._token_exchanger.invalidate( subject.inbound_token.get_secret_value(), server, server.config, tenant_id=subject.tenant_id ) - if isinstance(server.config, IdJagConfig): - self._exchanged_tokens.invalidate( - _id_jag_cache_key(subject.inbound_token.get_secret_value(), server.server_id, server.config) - ) + + def _invalidate_id_jag(self, subject: Subject, server: ServerSpec) -> None: + """Evict the bearer this `(subject, server)` last resolved, without depending on the store. + + The slot is addressed by the principal (plus the caller's own token when it presented one), + never by the credential material, so it stays computable when the assertion store is down. + The fingerprint stored with the entry is what keeps that safe: an entry minted for different + inputs reads as a miss rather than being served. + """ + self._exchanged_tokens.invalidate(_id_jag_slot_key(subject, server)) async def _authz_token(self, subject: Subject, server: ServerSpec) -> OAuthToken | None: """The user's authorization_code token, or None when absent or the store is unreachable. @@ -297,8 +378,37 @@ class UpstreamCredentialProvider: return None -def _id_jag_cache_key(subject_token: str, server_id: str, config: IdJagConfig) -> str: - """Bind the cached leg-2 bearer to the caller token, the server, AND the config that minted it. +def _id_jag_slot_key(subject: Subject, server: ServerSpec) -> str: + """Which cache slot this caller's bearer for this upstream lives in. + + Addressed by the principal, plus the caller's own token when it presented one so two callers + sharing an empty principal do not contend for one slot. Deliberately free of the stored + assertion, which is what lets invalidation compute this while the assertion store is down. The + entry's fingerprint, not this key, is what guarantees a cached bearer matches current inputs. + """ + inbound = subject.inbound_token.get_secret_value() if subject.inbound_token is not None else "" + material = "\x00".join((subject.tenant_id, subject.subject_id, server.server_id, inbound)) + return hashlib.sha256(material.encode()).hexdigest() + + +def _assertion_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool: + """Whether the stored assertion's ``exp`` has passed. An assertion carrying no expiry is + treated as usable and left for the IdP to reject, since the store records what the id_token + claimed rather than imposing a lifetime of its own. A naive ``expires_at`` is read as UTC so a + stored value that lost its offset compares instead of raising. + """ + expires_at = assertion.expires_at + if expires_at is None: + return False + normalized = expires_at if expires_at.tzinfo is not None else expires_at.replace(tzinfo=timezone.utc) + return normalized <= now + + +def _id_jag_fingerprint(subject_token: str, server_id: str, config: IdJagConfig) -> str: + """What the cached leg-2 bearer was minted from: the subject token, the server, and the config. + + Stored beside the bearer and compared on every read, so a rotated assertion or an edited server + config reads as a miss and re-mints instead of serving a bearer authorized under the old policy. Every exchange parameter derives from the config (endpoints, audience, resource, scopes, client auth), so a server update that changes any of them must change the key; otherwise the old bearer, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index e0927cc4f64..d52c718c0b8 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -18,7 +18,7 @@ from __future__ import annotations import json from datetime import datetime, timezone -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Protocol import jwt from pydantic import BaseModel, ConfigDict, SecretStr, TypeAdapter, ValidationError @@ -160,6 +160,42 @@ async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | N ) +class AssertionStoreUnavailable(Exception): + """Raised by ``fetch`` when the backing store is unreachable (e.g. the DB is down). + + Distinct from returning ``None`` for "this user has no captured assertion": an outage must not + read as a definite absence, which would tell the user to sign in again over a transient failure, + and it must not escape as an unhandled error on the egress or retry path. Mirrors + ``TokenStoreUnavailable`` on the sibling per-user OAuth store. + """ + + +class SSOAssertionStore(Protocol): + """The read seam the ``id_jag`` egress arm depends on, so the arm takes a collaborator + rather than reaching for a module-level function and a proxy global at call time. + + Returns the user's captured assertion, or ``None`` when they have never signed in. Raises + ``AssertionStoreUnavailable`` when the backing store is unreachable. + """ + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: ... + + +class DbSSOAssertionStore: + """The live store: the row the SSO callback wrote, read back by ``user_id``. + + A storage failure is re-raised as ``AssertionStoreUnavailable`` so the resolver can map it to a + typed fail-closed result; letting the raw driver error escape would surface a DB blip as a 500 + from credential resolution and from the upstream-401 retry. + """ + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + try: + return await fetch_sso_identity_assertion(user_id) + except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence + raise AssertionStoreUnavailable(str(exc)) from exc + + async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient, new_master_key: str) -> None: """Re-encrypt every stored assertion under ``new_master_key`` during a salt-key rotation, mirroring the sibling per-user credential tables; an unreadable row is skipped so one diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py index 4bc5732ec0e..3ed22732c90 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py @@ -23,7 +23,7 @@ from dataclasses import dataclass import httpx import jwt -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import assert_never from litellm._logging import verbose_proxy_logger @@ -51,6 +51,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ) from litellm.types.llms.custom_http import httpxSpecialProvider +# The cache stores (fingerprint, token); anything else in the slot is treated as absent. +_CACHED_ENTRY_ADAPTER: TypeAdapter[tuple[str, str]] = TypeAdapter(tuple[str, str]) + CLIENT_ASSERTION_TYPE = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer" CLIENT_ASSERTION_LIFETIME_SECONDS = 60 @@ -134,19 +137,28 @@ class ExchangedTokenCache: self, cache_key: str, compute: Callable[[], Awaitable[Result[ExchangedToken, CredError]]], + *, + fingerprint: str = "", ) -> Result[str, CredError]: - cached = self._get(cache_key) + """The cached token for `cache_key`, minting one when absent. + + `fingerprint` lets a caller address a slot by something stable (a principal) while still + guaranteeing the token it gets back was minted for the *current* inputs: a stored entry + whose fingerprint differs reads as a miss and is re-minted over. That keeps eviction + addressable without the key having to encode the credential material it protects. + """ + cached = self._get(cache_key, fingerprint) if cached is not None: return Ok(cached) async with self._lock(cache_key): - cached = self._get(cache_key) + cached = self._get(cache_key, fingerprint) if cached is not None: return Ok(cached) match await compute(): case Ok(token): self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped cache_key, - token.access_token, + (fingerprint, token.access_token), ttl=_cache_ttl_seconds(token.expires_in), ) return Ok(token.access_token) @@ -157,9 +169,18 @@ class ExchangedTokenCache: """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401).""" self._cache.delete_cache(cache_key) # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped - def _get(self, cache_key: str) -> str | None: - value = self._cache.get_cache(cache_key) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # InMemoryCache is untyped; narrowed by isinstance below - return value if isinstance(value, str) else None + def _get(self, cache_key: str, fingerprint: str) -> str | None: + """The stored token, or None when absent or minted for different inputs. + + The fingerprint comparison is what makes a shared slot safe: a mismatch never returns the + other party's token, it just reads as a miss. + """ + value = self._cache.get_cache(cache_key) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # InMemoryCache is untyped; the adapter below is the type gate + try: + stored_fingerprint, token = _CACHED_ENTRY_ADAPTER.validate_python(value) + except ValidationError: + return None + return token if stored_fingerprint == fingerprint else None def _lock(self, cache_key: str) -> asyncio.Lock: lock = self._locks.get(cache_key) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 0f276cb8e5c..f80954986e5 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -391,7 +391,8 @@ class Subject(BaseModel): tenant_id: str subject_id: str - # Opaque, already-validated inbound identity. Only `token_exchange` / `passthrough` read it. + # Opaque, already-validated inbound identity. Read by `token_exchange`, `passthrough`, and + # `id_jag` (which falls back to the user's stored SSO assertion when it is absent). inbound_token: SecretStr | None = None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 48effdb0f6e..d8431b9b3bb 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2902,6 +2902,54 @@ if MCP_AVAILABLE: # Deprecated: Local MCP Server Tool ######################################################### else: + # Gate only what can actually dispatch. When the unprefixed name is + # not in the registry either, `_handle_local_mcp_tool` below reports + # 404 and nothing runs, so demanding a server here would turn every + # unknown tool name into a misleading 503. + if global_mcp_tool_registry.get_tool(original_tool_name) is not None: + # `mcp_server` is None here because the tool name is not in the + # tool -> server mapping, but the name still carries a prefix + # that the server-level check above compared against the + # caller's `allowed_mcp_servers` by exact `name`. So the named + # server is in that list and can carry the tool-level checks, + # even with the mapping cold. Resolve it from + # `allowed_mcp_servers` rather than the registry: the registry + # would happily return a server the caller holds no grant for, + # and matching anything other than `name` would accept a server + # the check never validated. + prefix_server = next( + (candidate for candidate in allowed_mcp_servers if candidate.name == server_name), + None, + ) + if prefix_server is None: + # A non-empty prefix that passed the server-level check + # always matches here, so this arm only fires when the + # prefix was empty, which is exactly the case that check + # skips. Fail closed rather than dispatch with no server to + # evaluate a tool ceiling against. + raise HTTPException( + status_code=503, + detail=( + f"MCP server for tool '{original_tool_name}' is not available; " + "refusing to dispatch without authorization checks. " + "Retry once the server is registered." + ), + ) + + from litellm.proxy.proxy_server import proxy_logging_obj + + hook_result = await global_mcp_server_manager.pre_call_tool_check( + name=original_tool_name, + arguments=arguments, + server_name=server_name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=prefix_server, + raw_headers=raw_headers, + ) + if "arguments" in hook_result: + arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args + local_content = await _handle_local_mcp_tool(original_tool_name, arguments) response = CallToolResult(content=cast(Any, local_content), isError=False) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d85ad173434..23fe7af9994 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -841,6 +841,7 @@ class LiteLLMRoutes(enum.Enum): "/config/list", "/config/field/info", "/budget/list", + "/management/v1/budgets", "/budget/settings", # Invitation viewing (admin viewer cannot create/delete; can read). "/invitation/info", diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index cec682d772a..bdbe3f193d4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -1634,6 +1634,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) event_type: GuardrailEventHooks = GuardrailEventHooks.during_call + if call_type == CallTypes.call_mcp_tool.value: + event_type = GuardrailEventHooks.during_mcp_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: return @@ -1666,7 +1668,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source="INPUT", messages=filtered_messages, request_data=data, - logging_event_type=GuardrailEventHooks.during_call, + logging_event_type=event_type, ) ######################################################### diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index 6122f0594e8..e40db7f3f1f 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -6,8 +6,6 @@ import asyncio from datetime import datetime, timezone from typing import Optional -from pydantic import BaseModel - import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -16,7 +14,6 @@ from litellm.proxy._types import ( CommonProxyErrors, LiteLLM_AuditLogs, Litellm_EntityType, - LiteLLM_UserTable, LitellmTableNames, NewUserRequest, NewUserResponse, @@ -58,20 +55,20 @@ class UserManagementEventHooks: try: if prisma_client is None: raise Exception(CommonProxyErrors.db_not_connected_error.value) - user_row: BaseModel = await UserRepository(prisma_client).table.find_first( - where={"user_id": response.user_id} - ) - - user_row_litellm_typed = LiteLLM_UserTable(**user_row.model_dump(exclude_none=True)) + if response.user_id is None: + raise Exception("no user_id returned for the newly created user") + user_row = await UserRepository(prisma_client).find_by_id(response.user_id) + if user_row is None: + raise Exception(f"no user row found for user_id={response.user_id}") asyncio.create_task( UserManagementEventHooks.create_internal_user_audit_log( - user_id=user_row_litellm_typed.user_id, + user_id=user_row.user_id, action="created", litellm_changed_by=user_api_key_dict.user_id, user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=litellm_proxy_admin_name, before_value=None, - after_value=user_row_litellm_typed.model_dump_json(exclude_none=True), + after_value=user_row.model_dump_json(exclude_none=True), ) ) except Exception as e: diff --git a/litellm/proxy/management_endpoints/management_v1/__init__.py b/litellm/proxy/management_endpoints/management_v1/__init__.py index 257de66130b..a06c6b2591c 100644 --- a/litellm/proxy/management_endpoints/management_v1/__init__.py +++ b/litellm/proxy/management_endpoints/management_v1/__init__.py @@ -2,11 +2,15 @@ from fastapi import APIRouter +from litellm.proxy.management_endpoints.management_v1.budgets import ( + router as budgets_router, +) from litellm.proxy.management_endpoints.management_v1.spend_logs import ( router as spend_logs_router, ) router = APIRouter() +router.include_router(budgets_router) router.include_router(spend_logs_router) __all__ = ["router"] diff --git a/litellm/proxy/management_endpoints/management_v1/budgets.py b/litellm/proxy/management_endpoints/management_v1/budgets.py new file mode 100644 index 00000000000..bc1521caf0b --- /dev/null +++ b/litellm/proxy/management_endpoints/management_v1/budgets.py @@ -0,0 +1,205 @@ +"""`GET /management/v1/budgets`.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType +from typing import Annotated + +from fastapi import APIRouter, Depends, Request +from pydantic import BaseModel, TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + CommonProxyErrors, + UserAPIKeyAuth, + user_api_key_has_admin_view, +) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.management_v1.common import ( + MANAGEMENT_V1_PREFIX, + PROBLEM_TYPE_BASE, + ManagementProblem, +) +from litellm.proxy.management_endpoints.management_v1.list_framework import ( + FilterSpec, + ListSpec, + Predicate, + QueryPlan, + Scope, + ScopeAll, + ScopeDenied, + SortKey, + handle_list, + order_by_sql, + where_sql, +) +from litellm.proxy.utils import PrismaClient +from litellm.types.proxy.management_endpoints.management_v1 import ( + ListResponse, + ProblemDetail, +) + +router = APIRouter(prefix=MANAGEMENT_V1_PREFIX) + +BUDGET_TABLE = '"LiteLLM_BudgetTable"' + + +class BudgetListItem(BaseModel): + """One budget as the Budgets page reads it, and as it comes back off the table. + + Validating the raw row through here is what makes `tpm_limit` / `rpm_limit` + numbers: they are `BigInt?` in the schema, which the query engine hands back as + decimal strings, and a quoted "60000" breaks arithmetic in the dashboard. + """ + + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + budget_duration: str | None = None + budget_reset_at: datetime | None = None + created_at: datetime + updated_at: datetime + + +class _RowCount(BaseModel): + count: int + + +_BUDGET_ROWS = TypeAdapter(tuple[BudgetListItem, ...]) +_ROW_COUNTS = TypeAdapter(tuple[_RowCount, ...]) + +SELECTED_COLUMNS = ", ".join(f'"{name}"' for name in BudgetListItem.model_fields) + + +@dataclass(frozen=True, slots=True) +class PrismaBudgetListExecutor: + """The database half of the budgets list. Every caller-supplied value is bound to a + placeholder by `where_sql`; only the spec's own column names reach the SQL text.""" + + prisma_client: PrismaClient + + async def count(self, where: tuple[Predicate, ...]) -> int: + clauses, params = where_sql(where) + sql = f"SELECT COUNT(*) AS count FROM {BUDGET_TABLE}" + (f" WHERE {clauses}" if clauses else "") + rows = await self.prisma_client.db.query_raw(sql, *params) + counted = _ROW_COUNTS.validate_python(rows) + return counted[0].count if counted else 0 + + async def find_many(self, plan: QueryPlan) -> Sequence[BudgetListItem]: + clauses, params = where_sql(plan.where) + sql = ( + f"SELECT {SELECTED_COLUMNS} FROM {BUDGET_TABLE}" + + (f" WHERE {clauses}" if clauses else "") + + f" ORDER BY {order_by_sql(plan.order)}" + + f" LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" + ) + rows = await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip) + return _BUDGET_ROWS.validate_python(rows) + + +def _serialize(row: BudgetListItem) -> BudgetListItem: + """The row shape is the wire shape: the query selects exactly the columns served.""" + return row + + +def _scope(caller: UserAPIKeyAuth) -> Scope: + if user_api_key_has_admin_view(caller): + return ScopeAll() + return ScopeDenied(reason="Only proxy admins can list budgets, your role={}".format(caller.user_role)) + + +# budget_duration is deliberately absent from `sortable`: the column holds strings +# like "7d" and "30d", so a lexicographic ORDER BY puts "30d" ahead of "7d". +BUDGET_FILTERS: Mapping[str, FilterSpec] = MappingProxyType( + { # mutable-ok: an immutable mapping has no literal form; MappingProxyType freezes this one and it never escapes + "budget_duration": FilterSpec(type=str, ops=frozenset(("in", "is_null"))), + "max_budget": FilterSpec(type=float, ops=frozenset(("gte", "lte", "is_null"))), + "created_at": FilterSpec(type=datetime, ops=frozenset(("gte", "lte"))), + } +) + +BUDGETS_LIST_SPEC: ListSpec[BudgetListItem, BudgetListItem] = ListSpec( + resource="budgets", + sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "created_at")), + searchable=frozenset(("budget_id",)), + filters=BUDGET_FILTERS, + default_sort=(SortKey(field="created_at", descending=True),), + default_page_size=50, + max_page_size=100, + scope=_scope, + serialize=_serialize, + tiebreaker="budget_id", +) + + +@router.get( + "/budgets", + tags=("budget management",), + dependencies=(Depends(user_api_key_auth),), + response_model=ListResponse[BudgetListItem], +) +async def list_budgets( + request: Request, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> ListResponse[BudgetListItem]: + """ + The budgets defined on this proxy, paged, sortable and filterable, for the + Budgets page. + + Readable by a proxy admin or an admin viewer; anyone else is refused 403. The + older `/budget/list` answers with the whole table as a bare array and has no + way to page, sort or filter it. + + `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`, + `rpm_limit` or `created_at`, each optionally prefixed with `-` for descending, + and defaults to `-created_at`. `budget_id` is appended to every sort as the + tiebreaker. `q` is a case-insensitive substring match on `budget_id`. + `page_size` defaults to 50 and is capped at 100. Filters are + `filter[budget_duration][in|is_null]`, `filter[max_budget][gte|lte|is_null]` + and `filter[created_at][gte|lte]`. + + Example curl: + ``` + curl --location --globoff 'http://0.0.0.0:4000/management/v1/budgets?sort=-max_budget&filter[budget_duration][in]=7d,30d&page_size=25' \ + --header 'Authorization: Bearer sk-1234' + ``` + """ + 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, + ) + ) + + return await handle_list( + spec=BUDGETS_LIST_SPEC, + executor=PrismaBudgetListExecutor(prisma_client=prisma_client), + request=request, + caller=user_api_key_dict, + ) + + except ManagementProblem: + raise + except Exception as e: # noqa: BLE001 # a driver error answers as a problem document, not the OpenAI error shape + verbose_proxy_logger.exception( + "litellm.proxy.management_endpoints.management_v1.budgets.list_budgets(): Exception occured - {}".format( + str(e) + ) + ) + raise ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}internal-server-error", + title="Internal server error", + status=500, + detail="Failed to list budgets.", + ) + ) diff --git a/litellm/proxy/management_endpoints/management_v1/common.py b/litellm/proxy/management_endpoints/management_v1/common.py index c0e7f49f2e9..daa2c60ac5e 100644 --- a/litellm/proxy/management_endpoints/management_v1/common.py +++ b/litellm/proxy/management_endpoints/management_v1/common.py @@ -7,6 +7,7 @@ from fastapi.dependencies.utils import get_flat_dependant from fastapi.responses import JSONResponse from litellm.types.proxy.management_endpoints.management_v1 import ( + ListLinks, PageLinks, ProblemDetail, ) @@ -43,6 +44,21 @@ def _declared_query_params(request: Request) -> frozenset[str]: return frozenset(field.alias for field in get_flat_dependant(dependant, skip_repeats=True).query_params) +def escape_like(value: str) -> str: + """Escape LIKE/ILIKE metacharacters. Ids routinely contain `_`, which is a wildcard unescaped.""" + return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + + +def unknown_query_param_problem(unknown: tuple[str, ...], allowed: tuple[str, ...]) -> ProblemDetail: + return ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}unknown-query-parameter", + title="Unknown query parameter", + status=400, + detail=f"Unrecognized query parameter(s): {', '.join(unknown)}.", + allowed=sorted(allowed), + ) + + async def reject_unknown_query_params(request: Request) -> None: """Reject any query param the route did not declare. @@ -53,15 +69,7 @@ async def reject_unknown_query_params(request: Request) -> None: unknown: tuple[str, ...] = tuple(sorted(name for name in request.query_params if name not in declared)) if not unknown: return - raise ManagementProblem( - ProblemDetail( - type=f"{PROBLEM_TYPE_BASE}unknown-query-parameter", - title="Unknown query parameter", - status=400, - detail=f"Unrecognized query parameter(s): {', '.join(unknown)}.", - allowed=sorted(declared), - ) - ) + raise ManagementProblem(unknown_query_param_problem(unknown=unknown, allowed=tuple(sorted(declared)))) def _page_url(request: Request, page: int) -> str: @@ -75,3 +83,15 @@ def build_page_links(request: Request, page: int, has_more: bool) -> PageLinks: prev=_page_url(request, page - 1) if page > 1 else None, next=_page_url(request, page + 1) if has_more else None, ) + + +def build_list_links(request: Request, page: int, total_pages: int) -> ListLinks: + """Page-mode links. `last` clamps to page 1 on an empty result set so every link still resolves.""" + last = max(total_pages, 1) + return ListLinks( + self_link=_page_url(request, page), + first=_page_url(request, 1), + prev=_page_url(request, page - 1) if page > 1 else None, + next=_page_url(request, page + 1) if page < last else None, + last=_page_url(request, last), + ) diff --git a/litellm/proxy/management_endpoints/management_v1/list_framework.py b/litellm/proxy/management_endpoints/management_v1/list_framework.py new file mode 100644 index 00000000000..8f25c45016b --- /dev/null +++ b/litellm/proxy/management_endpoints/management_v1/list_framework.py @@ -0,0 +1,548 @@ +"""Generic list handling for `/management/v1` collection routes. + +A resource declares a `ListSpec`; `build_query_plan` turns query parameters into a +`QueryPlan` or an RFC 9457 problem without touching a database, and `handle_list` +runs that plan through an injected `ListExecutor`. Keeping the planning pure is what +lets a caller assert the plan as a value instead of asserting against a live Prisma +client, and it keeps this module free of any database dependency. + +A plan's `where` is a tuple of frozen `Predicate`s rather than a backend-shaped +mapping, so the framework never has to know which query builder executes it and a +planned predicate cannot be rewritten afterwards. `where_sql` renders one for a +raw-SQL executor with every caller-supplied value bound to a placeholder. +""" + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from functools import partial, reduce +from math import ceil +from typing import Generic, Literal, Protocol, TypeVar + +from fastapi import Request +from pydantic import TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.management_endpoints.management_v1.common import ( + PROBLEM_TYPE_BASE, + ManagementProblem, + build_list_links, + escape_like, + unknown_query_param_problem, +) +from litellm.types.proxy.management_endpoints.management_v1 import ( + ListMeta, + ListResponse, + ProblemDetail, +) + +ComparisonOp = Literal["eq", "gte", "lte", "gt", "lt", "contains", "not"] +# `is_null` is not in the design doc's operator set. It is here because there is no +# other way to ask for "max_budget IS NULL", and a table that renders nulls as +# "Unlimited" has to be able to filter on them. +FilterOp = ComparisonOp | Literal["in", "is_null"] + +FilterType = type[str] | type[int] | type[float] | type[datetime] +FilterValue = str | int | float | datetime + +PAGE_PARAM = "page" +PAGE_SIZE_PARAM = "page_size" +SORT_PARAM = "sort" +SEARCH_PARAM = "q" + +TRow = TypeVar("TRow") +TRow_co = TypeVar("TRow_co", covariant=True) +TOut = TypeVar("TOut") + +_FILTER_OP_ADAPTER: TypeAdapter[FilterOp] = TypeAdapter(FilterOp) + + +@dataclass(frozen=True, slots=True) +class Compare: + """`field value`.""" + + field: str + op: ComparisonOp + value: FilterValue + + +@dataclass(frozen=True, slots=True) +class Within: + """`field IN (values)`.""" + + field: str + values: tuple[FilterValue, ...] + + +@dataclass(frozen=True, slots=True) +class IsNull: + """`field IS NULL`, or `IS NOT NULL` when negated.""" + + field: str + negated: bool + + +@dataclass(frozen=True, slots=True) +class AnyOf: + """Disjunction of its clauses. `?q=` is the only producer today.""" + + clauses: tuple["Predicate", ...] + + +Predicate = Compare | Within | IsNull | AnyOf + + +@dataclass(frozen=True, slots=True) +class FilterSpec: + type: FilterType + ops: frozenset[FilterOp] + + +@dataclass(frozen=True, slots=True) +class SortKey: + field: str + descending: bool + + +@dataclass(frozen=True, slots=True) +class ScopeAll: + """The caller may read every row of the resource.""" + + +@dataclass(frozen=True, slots=True) +class ScopeWhere: + """The caller may read the rows matching every predicate in `where`.""" + + where: tuple[Predicate, ...] + + +@dataclass(frozen=True, slots=True) +class ScopeDenied: + """The caller may read no rows at all, and should be told so rather than shown an empty page.""" + + reason: str + + +Scope = ScopeAll | ScopeWhere | ScopeDenied + + +@dataclass(frozen=True, slots=True) +class ListSpec(Generic[TRow, TOut]): + resource: str + sortable: frozenset[str] + searchable: frozenset[str] + filters: Mapping[str, FilterSpec] + default_sort: tuple[SortKey, ...] + default_page_size: int + max_page_size: int + scope: Callable[[UserAPIKeyAuth], Scope] + serialize: Callable[[TRow], TOut] + tiebreaker: str + + def __post_init__(self) -> None: + """A malformed spec is a programming error at import time, so this raises rather than + returning a problem: there is no request in flight and no caller to answer.""" + if not 1 <= self.default_page_size <= self.max_page_size: + raise ValueError( + f"{self.resource}: default_page_size must be between 1 and max_page_size " + f"({self.max_page_size}), got {self.default_page_size}. A default above the cap " + f"would serve more rows than the resource allows whenever page_size is omitted." + ) + if not self.tiebreaker: + raise ValueError(f"{self.resource}: tiebreaker is required; it is the final sort key on every query.") + undeclared = tuple(sorted(frozenset(key.field for key in self.default_sort) - self.sortable)) + if undeclared: + raise ValueError(f"{self.resource}: default_sort orders by non-sortable field(s): {', '.join(undeclared)}.") + non_text = tuple( + sorted(field for field, spec in self.filters.items() if "contains" in spec.ops and spec.type is not str) + ) + if non_text: + raise ValueError( + f"{self.resource}: contains renders as ILIKE and is only meaningful on text columns, " + f"but is declared on: {', '.join(non_text)}." + ) + + +@dataclass(frozen=True, slots=True) +class QueryPlan: + """`where` is an implicit AND, ordered scope-first; `order` always ends with the spec's tiebreaker.""" + + where: tuple[Predicate, ...] + order: tuple[SortKey, ...] + skip: int + take: int + + +class ListExecutor(Protocol[TRow_co]): + """The database half of a list, injected so this module never imports Prisma.""" + + async def count(self, where: tuple[Predicate, ...]) -> int: ... + + async def find_many(self, plan: QueryPlan) -> Sequence[TRow_co]: ... + + +def order_by_sql(order: tuple[SortKey, ...]) -> str: + """`ORDER BY` body for a plan, NULLS LAST in both directions. + + Postgres sorts nulls last ascending but first descending, so an unqualified flip of + the sort direction drags every "Unlimited" row to the top of the table. Every field + reaching here is either a member of `ListSpec.sortable` (caller-supplied sort is + checked against it, `default_sort` at construction) or the spec's `tiebreaker`, so + these are developer-declared column names, never caller-controlled text. + """ + return ", ".join(f'"{key.field}" {"DESC" if key.descending else "ASC"} NULLS LAST' for key in order) + + +def _sql_operator(op: ComparisonOp) -> str: + match op: + case "eq": + return "=" + case "not": + return "<>" + case "gte": + return ">=" + case "lte": + return "<=" + case "gt": + return ">" + case "lt": + return "<" + case "contains": + return "ILIKE" + case _: + assert_never(op) + + +def _placeholder(index: int, value: FilterValue) -> str: + """`$n`, cast when the bind is a datetime. + + Binds cross into the query engine as JSON, so a datetime arrives as text and + Postgres refuses `timestamp >= text` outright. Prisma stores DateTime as a naive + `TIMESTAMP(3)` holding UTC, so the bind is read as an instant and then dropped to + naive UTC to match the column, the same cast `/spend/logs/ui` applies. + """ + return f"${index}::timestamptz AT TIME ZONE 'UTC'" if isinstance(value, datetime) else f"${index}" + + +def _render(predicate: Predicate, index: int) -> tuple[str, tuple[object, ...]]: + match predicate: + case IsNull(field=field, negated=negated): + return f'"{field}" IS {"NOT NULL" if negated else "NULL"}', () + case Within(field=field, values=values): + placeholders = ", ".join(_placeholder(index + offset, value) for offset, value in enumerate(values)) + return f'"{field}" IN ({placeholders})', values + case AnyOf(clauses=clauses): + rendered, params = _render_all(clauses, index) + return f"({' OR '.join(rendered)})", params + case Compare(field=field, op="contains", value=value): + return f"\"{field}\" ILIKE ${index} ESCAPE '\\'", (f"%{escape_like(str(value))}%",) + case Compare(field=field, op=op, value=value): + return f'"{field}" {_sql_operator(op)} {_placeholder(index, value)}', (value,) + case _: + assert_never(predicate) + + +def _render_one( + rendered: tuple[tuple[str, ...], tuple[object, ...]], + predicate: Predicate, + first_index: int, +) -> tuple[tuple[str, ...], tuple[object, ...]]: + """Append one predicate, numbering it after the binds already consumed.""" + clauses, params = rendered + clause, clause_params = _render(predicate, first_index + len(params)) + return (*clauses, clause), (*params, *clause_params) + + +def _render_all(predicates: tuple[Predicate, ...], index: int) -> tuple[tuple[str, ...], tuple[object, ...]]: + """Render every predicate, numbering placeholders continuously across them. + + Folded rather than self-recursive: walking a predicate list is a running index, and + recursing per predicate grew the stack with the filter count for nothing. `_render` + still re-enters here for `AnyOf`, whose clauses are plain `Compare`s from `?q=`, so + that nesting is one level deep and cannot be driven deeper by a caller. + """ + return reduce(partial(_render_one, first_index=index), predicates, ((), ())) + + +def where_sql(where: tuple[Predicate, ...], first_index: int = 1) -> tuple[str, tuple[object, ...]]: + """`WHERE` body and its bind parameters, numbered from `first_index`. + + Returns `("", ())` when there is nothing to filter on. Every caller-supplied value + becomes a `$n` placeholder rather than being written into the SQL text; only column + names reach the text, and those come from the spec's own declarations. + """ + clauses, params = _render_all(where, first_index) + return " AND ".join(clauses), params + + +def _problem(slug: str, title: str, status: int, detail: str, allowed: tuple[str, ...] | None = None) -> ProblemDetail: + return ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}{slug}", + title=title, + status=status, + detail=detail, + allowed=sorted(allowed) if allowed is not None else None, + ) + + +def _invalid(detail: str) -> ProblemDetail: + return _problem("invalid-query-parameter", "Invalid query parameter", 400, detail) + + +def _parse_filter_key(name: str) -> tuple[str, FilterOp] | None: + """`filter[max_budget][gte]` -> `("max_budget", "gte")`; bare `filter[status]` -> `("status", "eq")`.""" + if not name.startswith("filter[") or not name.endswith("]"): + return None + field, separator, raw_op = name[len("filter[") : -1].partition("][") + if not separator: + return field, "eq" + try: + return field, _FILTER_OP_ADAPTER.validate_python(raw_op) + except ValidationError: + return None + + +def _is_known_param(spec: ListSpec[TRow, TOut], name: str) -> bool: + if name in (PAGE_PARAM, PAGE_SIZE_PARAM): + return True + if name == SORT_PARAM: + return bool(spec.sortable) + if name == SEARCH_PARAM: + return bool(spec.searchable) + parsed = _parse_filter_key(name) + return parsed is not None and parsed[0] in spec.filters + + +def _allowed_params(spec: ListSpec[TRow, TOut]) -> tuple[str, ...]: + return tuple( + sorted( + (PAGE_PARAM, PAGE_SIZE_PARAM) + + ((SORT_PARAM,) if spec.sortable else ()) + + ((SEARCH_PARAM,) if spec.searchable else ()) + + tuple( + f"filter[{field}]" if op == "eq" else f"filter[{field}][{op}]" + for field, filter_spec in spec.filters.items() + for op in filter_spec.ops + ) + ) + ) + + +def _parse_positive_int(name: str, raw: str) -> int | ProblemDetail: + try: + value = int(raw) + except ValueError: + return _invalid(f"'{name}' must be an integer.") + if value < 1: + return _invalid(f"'{name}' must be 1 or greater.") + return value + + +def _parse_page(params: Mapping[str, str]) -> int | ProblemDetail: + raw = params.get(PAGE_PARAM) + return 1 if raw is None else _parse_positive_int(PAGE_PARAM, raw) + + +def _parse_page_size(spec: ListSpec[TRow, TOut], params: Mapping[str, str]) -> int | ProblemDetail: + raw = params.get(PAGE_SIZE_PARAM) + if raw is None: + return spec.default_page_size + value = _parse_positive_int(PAGE_SIZE_PARAM, raw) + if isinstance(value, ProblemDetail): + return value + return min(value, spec.max_page_size) + + +def _parse_sort(spec: ListSpec[TRow, TOut], params: Mapping[str, str]) -> tuple[SortKey, ...] | ProblemDetail: + raw = params.get(SORT_PARAM) + if raw is None: + return spec.default_sort + segments = tuple(segment.strip() for segment in raw.split(",")) + keys = tuple( + SortKey(field=segment[1:] if segment.startswith("-") else segment, descending=segment.startswith("-")) + for segment in segments + ) + rejected = tuple(sorted(frozenset(key.field for key in keys) - spec.sortable)) + if rejected: + return _problem( + "invalid-sort-field", + "Invalid sort field", + 400, + f"Cannot sort {spec.resource} by: {', '.join(repr(field) for field in rejected)}.", + tuple(spec.sortable), + ) + return keys + + +def _to_utc(value: datetime) -> datetime: + return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc) + + +def _coerce(field: str, op: FilterOp, raw: str, target: FilterType) -> FilterValue | ProblemDetail: + try: + if target is str: + return raw + if target is int: + return int(raw) + if target is float: + return float(raw) + return _to_utc(datetime.fromisoformat(raw[:-1] + "+00:00" if raw.endswith("Z") else raw)) + except ValueError: + return _invalid(f"'filter[{field}][{op}]' is not a valid {target.__name__}: {raw!r}.") + + +def _null_predicate(field: str, raw: str) -> Predicate | ProblemDetail: + if raw.lower() == "true": + return IsNull(field=field, negated=False) + if raw.lower() == "false": + return IsNull(field=field, negated=True) + return _invalid(f"'filter[{field}][is_null]' must be 'true' or 'false'.") + + +def _within_predicate(field: str, raw: str, target: FilterType) -> Predicate | ProblemDetail: + coerced = tuple(_coerce(field, "in", item.strip(), target) for item in raw.split(",")) + problems = tuple(item for item in coerced if isinstance(item, ProblemDetail)) + if problems: + return problems[0] + return Within(field=field, values=tuple(item for item in coerced if not isinstance(item, ProblemDetail))) + + +def _parse_filter(field: str, op: FilterOp, raw: str, filter_spec: FilterSpec) -> Predicate | ProblemDetail: + if op not in filter_spec.ops: + return _problem( + "unsupported-filter-operator", + "Unsupported filter operator", + 400, + f"Operator '{op}' is not supported on '{field}'.", + tuple(filter_spec.ops), + ) + if op == "is_null": + return _null_predicate(field, raw) + if op == "in": + return _within_predicate(field, raw, filter_spec.type) + value = _coerce(field, op, raw, filter_spec.type) + if isinstance(value, ProblemDetail): + return value + return Compare(field=field, op=op, value=value) + + +def _parse_filters(spec: ListSpec[TRow, TOut], params: Mapping[str, str]) -> tuple[Predicate, ...] | ProblemDetail: + keys = tuple( + (name, parsed) + for name in sorted(params) + if (parsed := _parse_filter_key(name)) is not None and parsed[0] in spec.filters + ) + parsed = tuple(_parse_filter(field, op, params[name], spec.filters[field]) for name, (field, op) in keys) + problems = tuple(item for item in parsed if isinstance(item, ProblemDetail)) + if problems: + return problems[0] + return tuple(item for item in parsed if not isinstance(item, ProblemDetail)) + + +def _search_predicate(spec: ListSpec[TRow, TOut], params: Mapping[str, str]) -> Predicate | None: + raw = params.get(SEARCH_PARAM) + if not raw: + return None + return AnyOf(clauses=tuple(Compare(field=field, op="contains", value=raw) for field in sorted(spec.searchable))) + + +def _scope_predicates(scope: Scope) -> tuple[Predicate, ...] | ProblemDetail: + match scope: + case ScopeAll(): + return () + case ScopeWhere(where=where): + return where + case ScopeDenied(reason=reason): + return _problem("forbidden", "Forbidden", 403, reason) + case _: + assert_never(scope) + + +def build_query_plan( + spec: ListSpec[TRow, TOut], + params: Mapping[str, str], + caller: UserAPIKeyAuth, +) -> QueryPlan | ProblemDetail: + """Turn query parameters into a plan, or into the problem that explains why they are not one.""" + scope_predicates = _scope_predicates(spec.scope(caller)) + if isinstance(scope_predicates, ProblemDetail): + return scope_predicates + + unknown = tuple(sorted(name for name in params if not _is_known_param(spec, name))) + if unknown: + return unknown_query_param_problem(unknown=unknown, allowed=_allowed_params(spec)) + + page = _parse_page(params) + if isinstance(page, ProblemDetail): + return page + + page_size = _parse_page_size(spec, params) + if isinstance(page_size, ProblemDetail): + return page_size + + sort = _parse_sort(spec, params) + if isinstance(sort, ProblemDetail): + return sort + + filters = _parse_filters(spec, params) + if isinstance(filters, ProblemDetail): + return filters + + search = _search_predicate(spec, params) + return QueryPlan( + # Scope first: conjuncts a caller filter sits behind and cannot replace. + where=scope_predicates + filters + ((search,) if search is not None else ()), + # Ordering by an all-null column without a unique final key lets Postgres return + # the same row on two different pages. + order=sort + (SortKey(field=spec.tiebreaker, descending=False),), + skip=(page - 1) * page_size, + take=page_size, + ) + + +def _duplicate_params(request: Request) -> tuple[str, ...]: + names = tuple(name for name, _ in request.query_params.multi_items()) + return tuple(sorted(frozenset(name for name in names if names.count(name) > 1))) + + +async def handle_list( + spec: ListSpec[TRow, TOut], + executor: ListExecutor[TRow], + request: Request, + caller: UserAPIKeyAuth, +) -> ListResponse[TOut]: + """Plan, execute, count, serialize, envelope. Failures reach the client as RFC 9457 problems.""" + plan = build_query_plan(spec=spec, params=request.query_params, caller=caller) + if isinstance(plan, ProblemDetail): + raise ManagementProblem(plan) + + # Checked here rather than in build_query_plan because a Mapping[str, str] cannot + # represent a repeat: query_params.get() silently keeps the last one, so ?page=1&page=999 + # would page from 999 without the caller ever being told which value won. + duplicates = _duplicate_params(request) + if duplicates: + raise ManagementProblem( + _problem( + "duplicate-query-parameter", + "Duplicate query parameter", + 400, + f"Repeated query parameter(s): {', '.join(duplicates)}. Each may appear once; " + f"use a comma-separated list for multiple sort keys or filter values.", + ) + ) + + total_count = await executor.count(plan.where) + rows = await executor.find_many(plan) + total_pages = ceil(total_count / plan.take) + page = plan.skip // plan.take + 1 + return ListResponse[TOut]( + data=tuple(spec.serialize(row) for row in rows), + meta=ListMeta( + total_count=total_count, + page=page, + page_size=plan.take, + total_pages=total_pages, + ), + links=build_list_links(request=request, page=page, total_pages=total_pages), + ) diff --git a/litellm/proxy/management_endpoints/management_v1/spend_logs.py b/litellm/proxy/management_endpoints/management_v1/spend_logs.py index c11a14bbfea..ccde3c4112c 100644 --- a/litellm/proxy/management_endpoints/management_v1/spend_logs.py +++ b/litellm/proxy/management_endpoints/management_v1/spend_logs.py @@ -13,6 +13,7 @@ from litellm.proxy.management_endpoints.management_v1.common import ( PROBLEM_TYPE_BASE, ManagementProblem, build_page_links, + escape_like, reject_unknown_query_params, ) from litellm.proxy.utils import PrismaClient @@ -34,10 +35,6 @@ def _as_utc(value: datetime) -> datetime: return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc) -def _escape_like(value: str) -> str: - return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") - - async def _end_user_scope_clause( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, @@ -133,7 +130,7 @@ async def list_spend_log_end_users( ) window_params: tuple[Any, ...] = (_as_utc(start_time), _as_utc(end_time)) - search_params: tuple[Any, ...] = (f"%{_escape_like(q)}%",) if q else () + search_params: tuple[Any, ...] = (f"%{escape_like(q)}%",) if q else () search_clause = (f"end_user ILIKE ${len(window_params) + 1} ESCAPE '\\'",) if q else () scope_clause, scope_params = await _end_user_scope_clause( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b6422d7f5ae..41904e3883b 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -324,7 +324,7 @@ async def patch_model( # Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates) live_before_reload = live_model_ids_snapshot() - await clear_cache() + still_desired_ids = await clear_cache() ## CREATE AUDIT LOG ## asyncio.create_task( @@ -344,6 +344,7 @@ async def patch_model( before=live_before_reload, written_models=[(model_id, getattr(updated_model, "model_info", None))], action="update", + still_desired=still_desired_ids, ) return updated_model @@ -429,7 +430,7 @@ async def _set_model_blocked_status( ) live_before_reload = live_model_ids_snapshot() - await clear_cache() + still_desired_ids = await clear_cache() asyncio.create_task( create_object_audit_log( @@ -450,6 +451,7 @@ async def _set_model_blocked_status( before=live_before_reload, written_models=[(data.model_id, getattr(updated_model, "model_info", None))], action=action, + still_desired=still_desired_ids, ) return updated_model @@ -1355,6 +1357,7 @@ async def add_new_model( """ live_before_reload = live_model_ids_snapshot() + still_desired_ids: frozenset[str] | None = None try: _original_litellm_model_name = model_params.model_name if model_params.model_info.team_id is None: @@ -1369,7 +1372,9 @@ async def add_new_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) - await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) + still_desired_ids = await proxy_config.add_deployment( + prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) # don't let failed slack alert block the /model/new response _alerting = general_settings.get("alerting", []) or [] if "slack" in _alerting: @@ -1414,6 +1419,7 @@ async def add_new_model( before=live_before_reload, written_models=[(model_response.model_id, getattr(model_response, "model_info", None))], action="create", + still_desired=still_desired_ids, ) return model_response @@ -1542,7 +1548,7 @@ async def update_model( # Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates) live_before_reload = live_model_ids_snapshot() - await clear_cache() + still_desired_ids = await clear_cache() ## CREATE AUDIT LOG ## asyncio.create_task( create_object_audit_log( @@ -1569,6 +1575,7 @@ async def update_model( before=live_before_reload, written_models=[(_model_id, getattr(model_response, "model_info", None))], action="update", + still_desired=still_desired_ids, ) return model_response @@ -1814,6 +1821,7 @@ def reload_serving_verdict( before: frozenset[str], written_models: Sequence[tuple[str, object]], written_must_serve: bool, + still_desired: frozenset[str] | None = None, ) -> tuple[tuple[str, ...], tuple[str, ...]]: """Judge a write-triggered reload by diffing the router's serving state instead of trusting any layer of the reload stack to report its own failure. @@ -1828,10 +1836,14 @@ def reload_serving_verdict( this write and blaming it would block unrelated metadata fixes - not written but live before and gone now: collateral degradation of this pod caused by the reload this request triggered (a wholesale re-add failure, or a - newly introduced conflict), always reported + newly introduced conflict), reported only when the db still wants that id + + ``still_desired`` is the db + config id set the reload just reconciled against. An + id absent from it was deleted on purpose, most often by another pod this one had not + yet polled, so the reload dropping it is the reconcile working rather than damage. + Without it (no reconcile ran) every drop is reported, which is the safe direction. Returns (written ids violating their obligation, collateral ids no longer served). - Best effort under concurrent admin writes: the snapshot spans only this request. """ now = live_model_ids_snapshot() written_ids = frozenset(model_id for model_id, _ in written_models) @@ -1843,7 +1855,8 @@ def reload_serving_verdict( ) else: missing = tuple(model_id for model_id, _ in written_models if model_id in before and model_id not in now) - collateral = tuple(sorted(before - now - written_ids)) + dropped = before - now - written_ids + collateral = tuple(sorted(dropped if still_desired is None else dropped & still_desired)) return (missing, collateral) @@ -1851,12 +1864,18 @@ def raise_if_reload_degraded_serving( before: frozenset[str], written_models: Sequence[tuple[str, object]], action: str, + still_desired: frozenset[str] | None = None, ) -> None: """The caller-visible error this pod's model-write endpoints owe their caller when the model they wrote is not being served after the reload they triggered. The DB write is durable either way and every other pod reloads on its own interval; this speaks only for the handling pod.""" - missing, collateral = reload_serving_verdict(before=before, written_models=written_models, written_must_serve=True) + missing, collateral = reload_serving_verdict( + before=before, + written_models=written_models, + written_must_serve=True, + still_desired=still_desired, + ) if not missing and not collateral: return missing_clause = ( @@ -1882,9 +1901,12 @@ def raise_if_reload_degraded_serving( ) -async def clear_cache(): +async def clear_cache() -> frozenset[str] | None: """ Clear router caches and reload models. + + Returns the db + config id set the reload reconciled against, or None when no + reload ran, so callers can pass it to raise_if_reload_degraded_serving. """ from litellm.proxy.proxy_server import ( llm_router, @@ -1896,7 +1918,7 @@ async def clear_cache(): if llm_router is None or prisma_client is None: verbose_proxy_logger.debug("llm_router or prisma_client is None, skipping cache clear") - return + return None try: # Only clear DB models, preserve config models @@ -1943,10 +1965,14 @@ async def clear_cache(): llm_router.quality_routers.pop(model_name, None) # Reload only DB models - await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) + still_desired_ids = await proxy_config.add_deployment( + prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) verbose_proxy_logger.debug( f"Cleared {len(db_model_ids)} DB models, preserved {len(config_models)} config models" ) + return still_desired_ids except Exception as e: verbose_proxy_logger.exception(f"Failed to clear cache and reload models. Due to error - {str(e)}") + return None diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 47a8670e26f..ac53f254981 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -96,7 +96,6 @@ class _VerificationTokenTableClient(Protocol): async def find_many( self, where: Mapping[str, object] | None = None, - select: Mapping[str, object] | None = None, ) -> "Sequence[PrismaVerificationToken]": ... @@ -157,7 +156,6 @@ async def _get_internal_user_api_keys( key_records = await _table(VerificationTokenRepository(prisma_client)).find_many( where={"user_id": user_id}, - select={"token": True}, ) user_api_keys.update(key_record.token for key_record in key_records if getattr(key_record, "token", None)) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c35c17aa359..fa01f43d049 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -131,6 +131,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkUpdateTeamMemberPermissionsRequest, BulkUpdateTeamMemberPermissionsResponse, GetTeamMemberPermissionsResponse, + TeamIdSearchMatch, TeamListItem, TeamListResponse, TeamMemberAddResult, @@ -3973,6 +3974,7 @@ async def _build_team_list_where_conditions( user_id: Optional[str], use_deleted_table: bool, search: Optional[str] = None, + search_team_id_match: TeamIdSearchMatch = "exact", org_admin_org_ids: Optional[List[str]] = None, user_api_key_cache: Optional[Any] = None, proxy_logging_obj: Optional[Any] = None, @@ -3996,7 +3998,7 @@ async def _build_team_list_where_conditions( if search: where_conditions["OR"] = [ - {"team_id": search}, + ({"team_id": {"startsWith": search}} if search_team_id_match == "prefix" else {"team_id": search}), {"team_alias": {"contains": search, "mode": "insensitive"}}, ] @@ -4230,8 +4232,14 @@ async def list_team_v2( ), search: Optional[str] = fastapi.Query( default=None, - description="Combined search: matches teams whose 'team_id' equals the value OR whose 'team_alias' contains it (case-insensitive).", + description="Combined search: matches teams whose 'team_id' matches the value OR whose 'team_alias' contains it (case-insensitive).", ), + search_team_id_match: Annotated[ + TeamIdSearchMatch, + fastapi.Query( + description="How 'search' matches 'team_id': 'exact' (default) or 'prefix' for a case-sensitive prefix match." + ), + ] = "exact", page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), page_size: int = fastapi.Query(default=10, description="Number of teams per page", ge=1, le=100), sort_by: Optional[str] = fastapi.Query( @@ -4308,6 +4316,7 @@ async def list_team_v2( user_id=user_id, use_deleted_table=use_deleted_table, search=search, + search_team_id_match=search_team_id_match, org_admin_org_ids=org_admin_org_ids, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index ad0b4f7444c..11ec19a8581 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -478,10 +478,7 @@ async def _resolve_team_id_to_object_permission_id( if not team_id or not team_id.strip(): return None team_id_clean = team_id.strip() - row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id_clean}, - select={"object_permission_id": True}, - ) + row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean}) if row is None: return None op_id = getattr(row, "object_permission_id", None) @@ -497,10 +494,7 @@ async def _resolve_team_id_to_object_permission_id( ) if updated_count == 0: await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id}) - row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id_clean}, - select={"object_permission_id": True}, - ) + row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean}) return getattr(row, "object_permission_id", None) if row else None return new_id diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 8682b61f910..110e5883485 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -3137,7 +3137,7 @@ class SSOAuthenticationHandler: _default_team_params = deepcopy(litellm.default_team_params) _new_team_request = team_request.model_dump() _new_team_request.update(_default_team_params) - team_request = NewTeamRequest(**_new_team_request) + team_request = NewTeamRequest.model_validate(_new_team_request) return team_request @staticmethod @@ -3271,14 +3271,6 @@ class SSOAuthenticationHandler: # User might not be already created on first generation of key # But if it is, we want their models preferences - default_ui_key_values: Dict[str, Any] = { - "duration": LITELLM_UI_SESSION_DURATION, - "key_max_budget": litellm.max_ui_session_budget, - "aliases": {}, - "config": {}, - "spend": 0, - "team_id": "litellm-dashboard", - } user_defined_values: Optional[SSOUserDefinedValues] = None if user_custom_sso is not None: @@ -3338,10 +3330,20 @@ class SSOAuthenticationHandler: verbose_proxy_logger.info(f"user_defined_values for creating ui key: {user_defined_values}") - default_ui_key_values.update(user_defined_values) - default_ui_key_values["request_type"] = "key" response = await generate_key_helper_fn( - **default_ui_key_values, # type: ignore + request_type="key", + duration=LITELLM_UI_SESSION_DURATION, + key_max_budget=litellm.max_ui_session_budget, + aliases={}, + config={}, + spend=0, + team_id="litellm-dashboard", + models=user_defined_values["models"], + user_id=user_defined_values["user_id"], + user_email=user_defined_values["user_email"], + user_role=user_defined_values["user_role"], + max_budget=user_defined_values["max_budget"], + budget_duration=user_defined_values["budget_duration"], table_name="key", ) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b2e36188681..2960b031cd3 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1035,7 +1035,7 @@ async def get_batch_from_database( if isinstance(db_batch_object.file_object, str) else db_batch_object.file_object ) - response = LiteLLMBatch(**batch_data) + response = LiteLLMBatch.model_validate(batch_data) response.id = batch_id # The stored batch object has the raw provider input_file_id. Resolve to unified ID. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a60ea2da019..9023f5fb23f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1041,6 +1041,8 @@ async def proxy_startup_event(app: FastAPI): redis_usage_cache=transaction_buffer_redis_cache, ) + ProxyStartupEvent._warn_if_mock_testing_params_enabled(general_settings=general_settings) + ## SEMANTIC TOOL FILTER ## # Read litellm_settings from config for semantic filter initialization try: @@ -5300,7 +5302,7 @@ class ProxyConfig: _model_info = RouterModelInfo(id=model.model_id, db_model=db_model) return _model_info - async def _delete_deployment(self, db_models: list) -> int: + async def _delete_deployment(self, db_models: list) -> frozenset[str] | None: """ (Helper function of add deployment) -> combined to reduce prisma db calls @@ -5309,14 +5311,16 @@ class ProxyConfig: - Remove any that are missing Return: - - int - returns number of deleted deployments + - frozenset[str] - the ids the db + config say should be served after this + reconcile, so a caller can tell an id this evicted on purpose from one that + went missing. None when no reconcile ran and that set is therefore unknown. """ global user_config_file_path, llm_router combined_id_list = [] ## BASE CASES ## if llm_router is None: - return 0 + return None # NOTE: db_models may be legitimately empty when all DB models have been deleted. # Do NOT short-circuit on len(db_models) == 0 — we must still evict any # DB-sourced deployments that are no longer in the DB. The caller @@ -5337,7 +5341,7 @@ class ProxyConfig: "Skipping deployment cleanup to avoid removing valid models.", str(e), ) - return 0 + return None model_list = config.get("model_list", None) if model_list: for model in model_list: @@ -5361,13 +5365,10 @@ class ProxyConfig: router_model_ids = llm_router.get_model_ids() # Check for model IDs in llm_router not present in combined_id_list and delete them - deleted_deployments = 0 for model_id in router_model_ids: if model_id not in combined_id_list: - is_deleted = llm_router.delete_deployment(id=model_id) - if is_deleted is not None: - deleted_deployments += 1 - return deleted_deployments + llm_router.delete_deployment(id=model_id) + return frozenset(combined_id_list) def _resolve_db_litellm_param(self, key: str, value: object) -> object: if not isinstance(value, str): @@ -5452,9 +5453,11 @@ class ProxyConfig: self, new_models: Optional[Json], proxy_logging_obj: ProxyLogging, - ): + ) -> frozenset[str] | None: global llm_router, llm_model_list, master_key, general_settings + still_desired_ids: frozenset[str] | None = None + # Load config separately so a timeout here doesn't block model loading config_data: dict = {} search_tools = None @@ -5501,7 +5504,7 @@ class ProxyConfig: if search_tools is not None and llm_router is not None: llm_router.search_tools = search_tools ## DELETE MODEL LOGIC - await self._delete_deployment(db_models=models_list) + still_desired_ids = await self._delete_deployment(db_models=models_list) ## ADD MODEL LOGIC self._add_deployment(db_models=models_list) @@ -5527,6 +5530,8 @@ class ProxyConfig: proxy_logging_obj=proxy_logging_obj, ) + return still_desired_ids + def _add_callback_from_db_to_in_memory_litellm_callbacks( self, callback: str, @@ -6159,14 +6164,20 @@ class ProxyConfig: self, prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging, - ): + ) -> frozenset[str] | None: """ - Check db for new models - Check if model id's in router already - If not, add to router + + Returns the ids the db + config say should be served after the reconcile, or + None when no reconcile ran. Callers that judge their own reload need it to tell + a deliberate eviction from a deployment that went missing. """ global llm_router, llm_model_list, master_key, general_settings + still_desired_ids: frozenset[str] | None = None + try: # warm the config cache so the per-param reads below all hit await prefetch_config_params( @@ -6186,7 +6197,9 @@ class ProxyConfig: new_models = await self._get_models_from_db(prisma_client=prisma_client) # update llm router - await self._update_llm_router(new_models=new_models, proxy_logging_obj=proxy_logging_obj) + still_desired_ids = await self._update_llm_router( + new_models=new_models, proxy_logging_obj=proxy_logging_obj + ) db_general_settings = await get_config_param(prisma_client, "general_settings") @@ -6204,6 +6217,8 @@ class ProxyConfig: "litellm.proxy.proxy_server.py::ProxyConfig:add_deployment - {}".format(str(e)) ) + return still_desired_ids + async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient): """ Use this to read non-llm objects from the db and initialize them @@ -7712,6 +7727,37 @@ class ProxyStartupEvent: proxy_logging_obj.startup_event(llm_router=llm_router, redis_usage_cache=redis_usage_cache) + @staticmethod + def _warn_if_mock_testing_params_enabled(general_settings: dict) -> None: + """Announce, loudly, that any caller may inject synthetic failures.""" + from litellm.proxy.route_llm_request import ( + GATED_MOCK_PARAM_NAMES, + MOCK_TESTING_CONFIG_KEY, + ) + + if general_settings.get(MOCK_TESTING_CONFIG_KEY, False) is not True: + return + + verbose_proxy_logger.warning( + "\n%s\n" + " DANGEROUS SETTING ENABLED\n" + " general_settings.%s = true\n" + "\n" + " Any caller with a valid key on this proxy can now inject synthetic\n" + " failures and latency into their own requests using these body params:\n" + "%s\n" + "\n" + " A request using them consumes a connection and a concurrency slot\n" + " without reaching a provider, and returns an error the caller chose.\n" + "\n" + " Intended for testing fallback chains. Do not leave enabled.\n" + "%s", + "=" * 72, + MOCK_TESTING_CONFIG_KEY, + "\n".join(f" {name}" for name in GATED_MOCK_PARAM_NAMES), + "=" * 72, + ) + @staticmethod def _validate_redis_transaction_buffer_config( general_settings: dict, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 1f5aacc2115..12e9857219e 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -8,18 +8,24 @@ import litellm from litellm.proxy._types import UserAPIKeyAuth from litellm.router_utils.common_utils import _is_proxy_admin_request -# Router-internal mock_testing_* flag names — kept in sync with -# ``litellm.types.router.MockRouterTestingParams`` by the test -# ``test_mock_testing_kwarg_names_matches_dataclass``. Hardcoding (rather -# than deriving via ``dataclasses.fields(MockRouterTestingParams)`` at +# Client-supplied params that make the router or the call path fabricate a +# failure or a delay instead of calling the provider. The ``mock_testing_*`` +# names are kept in sync with ``litellm.types.router.MockRouterTestingParams`` +# by ``test_gated_mock_params_cover_mock_router_testing_params``. Hardcoding +# (rather than deriving via ``dataclasses.fields(MockRouterTestingParams)`` at # import time) avoids a cyclic import: ``litellm.types.router`` imports # back into proxy modules before this module finishes loading. -_MOCK_TESTING_KWARG_NAMES: tuple = ( +GATED_MOCK_PARAM_NAMES: tuple[str, ...] = ( "mock_testing_fallbacks", "mock_testing_context_fallbacks", "mock_testing_content_policy_fallbacks", + "mock_testing_rate_limit_error", + "mock_timeout", + "mock_delay", ) +MOCK_TESTING_CONFIG_KEY = "dangerously_allow_mock_testing_request_params" + if TYPE_CHECKING: from litellm.router import Router as _Router @@ -169,6 +175,41 @@ def raise_if_required_body_param_missing(route_type: str, data: Mapping[str, obj ) +class MockTestingParamsDisabledError(HTTPException): + def __init__(self, params: tuple[str, ...]): + super().__init__( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: HTTPException.detail has no immutable form; same shape as the sibling errors here + "error": ( + f"Mock testing request params are disabled on this proxy: {', '.join(params)}. " + f"An admin can enable them by setting `general_settings.{MOCK_TESTING_CONFIG_KEY}: true` " + "in config.yaml. This setting cannot be changed from the Admin UI or the API." + ) + }, + ) + + +def raise_if_mock_testing_params_disallowed(data: Mapping[str, object], *, allowed: bool) -> None: + """Reject client-supplied mock testing params unless an admin opted in. + + Rejecting (rather than silently dropping) keeps a request that asked for a + synthetic failure from returning a normal success, which reads as a passing + fallback test that never ran. + """ + if allowed: + return + present = tuple(name for name in GATED_MOCK_PARAM_NAMES if name in data) + if present: + raise MockTestingParamsDisabledError(params=present) + + +def mock_testing_params_allowed() -> bool: + """Read the opt-in from the running proxy's ``general_settings``.""" + import litellm.proxy.proxy_server as proxy_server + + return proxy_server.general_settings.get(MOCK_TESTING_CONFIG_KEY, False) is True + + def get_team_id_from_data(data: dict) -> Optional[str]: """ Get the team id from the data's metadata or litellm_metadata params. @@ -381,12 +422,7 @@ async def route_request( await add_shared_session_to_data(data) - # Strip router-internal mock_testing_* flags. Combined with an - # unauthorized fallback in ``router_settings_override`` they let a - # caller deterministically execute requests against restricted - # models. VERIA-44. - for _key in _MOCK_TESTING_KWARG_NAMES: - data.pop(_key, None) + raise_if_mock_testing_params_disallowed(data, allowed=mock_testing_params_allowed()) data.pop("enable_tag_filtering", None) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 42788227acc..f7e429ddb31 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -11,6 +11,7 @@ from typing import ( Literal, Mapping, NamedTuple, + Sequence, Union, ) @@ -2025,19 +2026,7 @@ async def ui_view_spend_logs( data = await prisma_client.db.query_raw(sql_query, *sql_params) - # query_raw returns the JSONB `metadata` column as a string (the Prisma - # serialiser bypasses the model-layer JSON hydration we get on the ORM - # path). The UI reads `metadata.status` / `metadata.error_information` - # as object fields, so failure rows looked like successes (#29674). - # Re-hydrate to dict here. - for row in data: - if isinstance(row, dict): - md = row.get("metadata") - if isinstance(md, str): - try: - row["metadata"] = json.loads(md) - except (ValueError, TypeError): - row["metadata"] = {} + _hydrate_spend_log_metadata(data) # Calculate total pages total_pages = (total_records + page_size - 1) // page_size @@ -2078,6 +2067,27 @@ def _spend_log_field_has_content(value: Union[str, list, dict] | None) -> bool: return True +def _hydrate_spend_log_metadata(rows: Sequence[Any]) -> None: + """Re-hydrate the JSONB ``metadata`` column returned by ``query_raw`` as a string. + + The Prisma serialiser bypasses the model-layer JSON hydration we get on the ORM + path, while the UI reads ``metadata.status`` / ``metadata.error_information`` / + ``metadata.internal_call_origin`` as object fields. Property access on a string + is silently undefined, so failure rows looked like successes (#29674). Every + ``query_raw`` reader of this column goes through here so a new one cannot + reintroduce that. + """ + for row in rows: + if not isinstance(row, dict): + continue + md = row.get("metadata") + if isinstance(md, str): + try: + row["metadata"] = json.loads(md) + except (ValueError, TypeError): + row["metadata"] = {} + + def _cold_storage_object_key_from_metadata( metadata: Union[str, dict] | None, ) -> str | None: @@ -3361,6 +3371,7 @@ async def ui_view_session_spend_logs( LIMIT $2 OFFSET $3 """ result = await prisma_client.db.query_raw(sql_query, session_id, page_size, skip, *scope_params) + _hydrate_spend_log_metadata(result) total_pages = (total_records + page_size - 1) // page_size diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index f0b5dfd8bc5..e93991463f8 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -3,61 +3,56 @@ User repository for database operations on LiteLLM_UserTable. """ import json -from typing import Any, Dict, List, Optional, Type +from typing import Any, Dict, List, Mapping, Optional, Type from litellm.models.user import LiteLLM_UserTable -from litellm.repositories.base_repository import BaseRepository +from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict + +_JSON_ENCODED_COLUMNS = frozenset({"metadata", "model_spend", "model_max_budget"}) class UserRepository(BaseRepository[LiteLLM_UserTable]): """Repository for user database operations.""" @property - def table(self) -> Any: + def table(self) -> Any: # any-ok: Prisma table actions are reached through the untyped client wrapper return self.prisma_client.db.litellm_usertable @property def model_class(self) -> Type[LiteLLM_UserTable]: return LiteLLM_UserTable - def _to_model(self, record: Any) -> Optional[LiteLLM_UserTable]: + def _to_model(self, record: Optional[DbRecord]) -> Optional[LiteLLM_UserTable]: """Convert a database record to a User model.""" if record is None: return None - data = record.dict() if hasattr(record, "dict") else dict(record) + return LiteLLM_UserTable.model_validate( + { + column: json.loads(value) if column in _JSON_ENCODED_COLUMNS and isinstance(value, str) else value + for column, value in record_to_dict(record).items() + } + ) - json_fields = ["metadata", "model_spend", "model_max_budget"] - for field in json_fields: - if isinstance(data.get(field), str): - data[field] = json.loads(data[field]) - - return LiteLLM_UserTable(**data) - - async def find_by_id(self, user_id: str, id_field: str = "user_id") -> Optional[LiteLLM_UserTable]: - return await super().find_by_id(user_id, id_field) + async def find_by_id(self, id_value: str, id_field: str = "user_id") -> Optional[LiteLLM_UserTable]: + return await super().find_by_id(id_value, id_field) async def find_by_email(self, user_email: str) -> Optional[LiteLLM_UserTable]: """Find a user by email.""" - records = await self.table.find_many(where={"user_email": user_email}) - if records: - return self._to_model(records[0]) - return None + records = await self.find_many(where={"user_email": user_email}) + return records[0] if records else None async def find_by_sso_id(self, sso_user_id: str) -> Optional[LiteLLM_UserTable]: """Find a user by SSO ID.""" - record = await self.table.find_unique(where={"sso_user_id": sso_user_id}) - return self._to_model(record) + return await self.find_by_id(sso_user_id, id_field="sso_user_id") async def find_by_organization_id(self, organization_id: str) -> List[LiteLLM_UserTable]: """Find all users in an organization.""" - records = await self.table.find_many(where={"organization_id": organization_id}) - return self._to_model_list(records) + return await self.find_many(where={"organization_id": organization_id}) async def find_by_team_id(self, team_id: str) -> List[LiteLLM_UserTable]: """Find all users in a team.""" - records = await self.table.find_many(where={"teams": {"has": team_id}}) - return self._to_model_list(records) + return await self.find_many(where={"teams": {"has": team_id}}) async def count_billable_users(self) -> int: """Number of users that count toward the license seat limit. @@ -86,7 +81,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): max_budget: Optional[float] = None, user_email: Optional[str] = None, models: Optional[List[str]] = None, - metadata: Optional[Dict[str, Any]] = None, + metadata: Optional[Mapping[str, object]] = None, max_parallel_requests: Optional[int] = None, tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, @@ -96,7 +91,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): object_permission_id: Optional[str] = None, ) -> LiteLLM_UserTable: """Create a new user.""" - data: Dict[str, Any] = {"user_id": user_id} + data: Dict[str, object] = {"user_id": user_id} if user_alias is not None: data["user_alias"] = user_alias if team_id is not None: @@ -149,7 +144,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): max_budget: Optional[float] = None, user_email: Optional[str] = None, models: Optional[List[str]] = None, - metadata: Optional[Dict[str, Any]] = None, + metadata: Optional[Mapping[str, object]] = None, max_parallel_requests: Optional[int] = None, tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, @@ -159,7 +154,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): object_permission_id: Optional[str] = None, ) -> Optional[LiteLLM_UserTable]: """Update a user.""" - data: Dict[str, Any] = {} + data: Dict[str, object] = {} if user_alias is not None: data["user_alias"] = user_alias if team_id is not None: @@ -212,11 +207,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): if not await self.exists(user_id, id_field="user_id"): return None - record = await self.table.update( - where={"user_id": user_id}, - data={"teams": {"push": team_id}}, - ) - return self._to_model(record) + return await self.update(user_id, {"teams": {"push": team_id}}, id_field="user_id") async def remove_from_team(self, user_id: str, team_id: str) -> Optional[LiteLLM_UserTable]: """Remove a user from a team. diff --git a/litellm/router.py b/litellm/router.py index ac00cdbc2b0..399bfcdf3e1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8613,9 +8613,9 @@ class Router: deployment = self.get_deployment(model_id=model_id) if deployment is None or self._is_deployment_blocked(deployment): return None - return CredentialLiteLLMParams(**deployment.litellm_params.model_dump(exclude_none=True)).model_dump( - exclude_none=True - ) + return CredentialLiteLLMParams.model_validate( + deployment.litellm_params.model_dump(exclude_none=True) + ).model_dump(exclude_none=True) def get_deployment_by_model_group_name(self, model_group_name: str) -> Optional[Deployment]: """ @@ -8755,9 +8755,9 @@ class Router: return None # Get basic credentials - credentials = CredentialLiteLLMParams(**deployment.litellm_params.model_dump(exclude_none=True)).model_dump( - exclude_none=True - ) + credentials = CredentialLiteLLMParams.model_validate( + deployment.litellm_params.model_dump(exclude_none=True) + ).model_dump(exclude_none=True) # Resolve litellm_credential_name to actual credentials if deployment.litellm_params.litellm_credential_name is not None: diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index b43fe0da4ca..1c6a83c6cd9 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -329,7 +329,7 @@ class ComplexityRouter(CustomLogger): # Parse config - always create a new instance to avoid singleton mutation if complexity_router_config: - self.config = ComplexityRouterConfig(**complexity_router_config) + self.config = ComplexityRouterConfig.model_validate(complexity_router_config) else: self.config = ComplexityRouterConfig() diff --git a/litellm/types/proxy/management_endpoints/management_v1.py b/litellm/types/proxy/management_endpoints/management_v1.py index 2aecc54f114..b2244f6eb9b 100644 --- a/litellm/types/proxy/management_endpoints/management_v1.py +++ b/litellm/types/proxy/management_endpoints/management_v1.py @@ -1,7 +1,11 @@ """Shared response shapes for the `/management/v1` control-plane surface.""" +from typing import Generic, TypeVar + from pydantic import BaseModel, ConfigDict, Field +TOut = TypeVar("TOut") + class ProblemDetail(BaseModel): """RFC 9457 problem details, served as `application/problem+json`.""" @@ -37,3 +41,33 @@ class FacetListResponse(BaseModel): data: list[str] meta: PageMeta links: PageLinks + + +class ListMeta(BaseModel): + """Page-mode counterpart to `PageMeta`: an entity list pays for the COUNT(*) so the table can show a page count.""" + + total_count: int + page: int + page_size: int + total_pages: int + + +class ListLinks(BaseModel): + """Page-mode counterpart to `PageLinks`. `first`/`last` are knowable here because the total count is.""" + + model_config = ConfigDict(populate_by_name=True) + + self_link: str = Field(alias="self") + first: str + prev: str | None = None + next: str | None = None + last: str + + +class ListResponse(BaseModel, Generic[TOut]): + """Rows stay flat: JSON:API's `{type, id, attributes}` wrapper is a deliberate deviation, so every + dashboard column accessor would otherwise have to go through `.attributes`.""" + + data: list[TOut] + meta: ListMeta + links: ListLinks diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 0e555535874..2d9387da956 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Union from pydantic import BaseModel @@ -11,6 +11,8 @@ from litellm.proxy._types import ( Member, ) +TeamIdSearchMatch = Literal["exact", "prefix"] + class GetTeamMemberPermissionsRequest(BaseModel): """Request to get the team member permissions for a team""" diff --git a/pyproject.toml b/pyproject.toml index 93fb32da464..678e7384a05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ dependencies = [ "tokenizers>=0.21.0,<1.0", "click>=8.0.0,<9.0", "jinja2>=3.1.6,<4.0", - "aiohttp>=3.10,<4.0", + "aiohttp>=3.14.2,<4.0", "pydantic>=2.10.0,<3.0.0", "jsonschema>=4.0.0,<5.0", ] @@ -277,7 +277,7 @@ exclude = [ [tool.uv] constraint-dependencies = [ "tornado>=6.5.6", - "aiohttp>=3.14.1,<4.0", + "aiohttp>=3.14.2,<4.0", "packaging>=24.0", "soupsieve>=2.8.4", "httplib2>=0.32.0", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index dfdc4efe800..6fb3ed748b6 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 130 }, "ANN401": { - "limit": 2010 + "limit": 2009 }, "ASYNC230": { "limit": 14 @@ -324,7 +324,7 @@ "limit": 879 }, "UP006": { - "limit": 12138 + "limit": 12135 }, "UP007": { "limit": 2526 diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index 809dc141eb8..43a4cb66484 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -20,7 +20,10 @@ LIT002 Mutable-collection *construction*: a list/dict/set literal or comprehens generator (`tuple(f(x) for x in xs)`), a tuple literal, or a frozen dataclass / NamedTuple / ReadOnly TypedDict. Generator expressions and `tuple`/`frozenset` calls are not construction and pass. Annotation-internal lists (`Callable[[int], - str]`) are exempt. Suppress with `# mutable-ok: `. + str]`) are exempt, as is a value passed directly to a freezing wrapper + (`tuple(...)`, `frozenset(...)`, `MappingProxyType(...)`): it is frozen before + it can escape, though anything mutable nested inside it still counts. + Suppress with `# mutable-ok: `. LIT003 noqa suppression without rule codes or without a reason. Required shape: `# noqa: TID251 # ` LIT004 pyright/mypy ignore without bracketed codes or without a reason. @@ -90,6 +93,7 @@ MUTABLE_CONSTRUCTORS = frozenset(( # are common methods (e.g. pydantic's `model.dict()`), not collection construction. A # qualified `collections.deque(...)` still counts. QUALIFIED_CONSTRUCTORS = MUTABLE_CONSTRUCTORS - frozenset(("dict", "list", "set")) +FREEZING_WRAPPERS = frozenset(("tuple", "frozenset", "MappingProxyType")) UNSAFE_GUARDS = frozenset(("TypeGuard", "TypeIs")) MIN_REASON_LEN = 3 @@ -382,6 +386,34 @@ def _annotation_node_ids(tree: ast.AST) -> frozenset[int]: ) +def _is_freezing_wrapper(func: ast.expr) -> bool: + if isinstance(func, ast.Name): + return func.id in FREEZING_WRAPPERS + return ( + isinstance(func, ast.Attribute) + and func.attr == "MappingProxyType" + and isinstance(func.value, ast.Name) + and func.value.id == "types" + ) + + +def _frozen_argument_ids(tree: ast.AST) -> frozenset[int]: + """ids() of every expression passed directly to a freezing wrapper. + + `MappingProxyType({...})`, `frozenset({...})`, and `tuple([...])` freeze their + argument before it can escape, so the literal inside is a one-shot build, not a + mutable value anyone can grow later. Only the argument itself is exempt; a + mutable collection nested inside it still trips LIT002. Only bare names (plus + `types.MappingProxyType`) qualify, so an unrelated method that happens to share + a wrapper's name cannot exempt its argument. + """ + return frozenset( + id(node.args[0]) + for node in ast.walk(tree) + if isinstance(node, ast.Call) and len(node.args) == 1 and _is_freezing_wrapper(node.func) + ) + + def _construction_kind(node: ast.expr) -> str | None: """Human label if `node` builds a mutable collection, else None.""" if isinstance(node, ast.List): @@ -407,8 +439,9 @@ def _construction_kind(node: ast.expr) -> str | None: def iter_construction_violations(path: Path, tree: ast.AST, comments: Comments) -> Iterator[Violation]: in_annotation = _annotation_node_ids(tree) + frozen_arguments = _frozen_argument_ids(tree) for node in ast.walk(tree): - if not isinstance(node, ast.expr) or id(node) in in_annotation: + if not isinstance(node, ast.expr) or id(node) in in_annotation or id(node) in frozen_arguments: continue kind = _construction_kind(node) if kind is None or node.lineno in comments.mutable_ok_lines: diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 68c5ef6b31d..2a0fc5c9f29 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -57,6 +57,8 @@ - {id: mgmt.budget.update.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "budget_management_endpoints.py:155", rationale: "Limit changes apply"} - {id: mgmt.budget.delete.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "budget_management_endpoints.py:280", rationale: "Clears limits"} - {id: mgmt.budget.list.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "budget_management_endpoints.py:215", rationale: "Budget enumeration"} +- {id: mgmt.budget.list_v1.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "management_v1/budgets.py:129", rationale: "Budget enumeration the Budgets page can page, sort and filter"} +- {id: mgmt.budget.list_v1.admin_only, module: mgmt, tier: P1, surface: api, assertions: [admin_only], source: "management_v1/budgets.py:129", rationale: "A caller without admin view is refused, not served an empty page"} - {id: mgmt.callback.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "callback_management_endpoints.py", rationale: "Callback config (smoke)"} - {id: mgmt.cache_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cache_settings_endpoints.py", rationale: "Cache config (smoke). Deliberately uncovered: the previous test read the live settings and wrote them back, which proves nothing (identical values in, so a no-op POST still passes) while being able to break the deployment. /cache/settings persists what it receives and that row outranks YAML cache_params, re-applied on a timer, so a write that omits ssl or redis_startup_nodes turns a TLS cluster into a plaintext standalone node and every later Redis call hangs. That took out 60 of 72 tests on 2026-07-25. GET cannot round-trip it either: it resolves the stored row overlaid with REDIS_* env and never reads YAML, so on a fresh deploy it cannot see YAML ssl to echo back. A safe test needs an isolated proxy, or LIT-4816 fixed so a partial write cannot downgrade transport. Do not re-add a read-then-write-back test against a shared proxy."} - {id: mgmt.cost_tracking.estimate.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cost_tracking_settings.py", rationale: "Cost estimate (smoke)"} diff --git a/tests/e2e/load/test_chat_completions_throughput_e2e.py b/tests/e2e/load/test_chat_completions_throughput_e2e.py index 6dd5fff971a..dd9ca370c6f 100644 --- a/tests/e2e/load/test_chat_completions_throughput_e2e.py +++ b/tests/e2e/load/test_chat_completions_throughput_e2e.py @@ -16,6 +16,17 @@ pytestmark = [pytest.mark.e2e, pytest.mark.load] class TestChatCompletionsThroughput: + @pytest.mark.skip( + reason=( + "LIT-5054: the SLO measures how many gateway replicas happen to be warm, not the " + "request path. Clearing the floor needs roughly 5-7 replicas at ~10-14 RPS each, " + "stage idles at one, and reactive HPA scale-up lands minutes into a ~3 minute " + "test. It has failed both assertions on consecutive days: 93.3% errors at an " + "inflated 264 RPS (closed-loop RPS rises when requests fail fast, and those " + "requests never reached a pod), then 16.7 RPS with zero failures. Unskip once the " + "assertion is independent of fleet size." + ) + ) @pytest.mark.covers("reliability.perf.throughput.under_slo") def test_sustains_throughput_slo_under_load(self, client: LoadClient, load_key: str) -> None: result = run_chat_load( diff --git a/tests/e2e/management/test_budget_customer_user_org_e2e.py b/tests/e2e/management/test_budget_customer_user_org_e2e.py index 54cc18b228b..12372bb7cc1 100644 --- a/tests/e2e/management/test_budget_customer_user_org_e2e.py +++ b/tests/e2e/management/test_budget_customer_user_org_e2e.py @@ -19,10 +19,10 @@ import time from collections.abc import Callable import pytest -from pydantic import BaseModel, RootModel +from pydantic import BaseModel, Field, RootModel from e2e_config import unique_marker -from e2e_http import NoBody, unwrap +from e2e_http import NoBody, Success, UnauthorizedError, UnknownApiError, unwrap from lifecycle import ResourceManager from management_client import ManagementClient from models import KeyGenerateBody, OrgInfoParams, OrgNewBody, UserNewBody @@ -44,9 +44,11 @@ def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: class BudgetNewBody(BaseModel): - max_budget: float + max_budget: float | None = None soft_budget: float | None = None budget_duration: str | None = None + budget_id: str | None = None + tpm_limit: int | None = None class BudgetNewResponse(BaseModel): @@ -204,6 +206,164 @@ class TestBudgetManagement: ) +# ---------- /management/v1/budgets ---------- + +_BUDGETS_V1 = "/management/v1/budgets" + + +class BudgetPageParams(BaseModel): + """Query for GET /management/v1/budgets. The filter fields serialize to the + bracketed keys the route reads them under, so nothing here is a raw dict.""" + + q: str | None = None + sort: str | None = None + page: int | None = None + page_size: int | None = None + duration_in: str | None = Field(default=None, serialization_alias="filter[budget_duration][in]") + max_budget_is_null: bool | None = Field(default=None, serialization_alias="filter[max_budget][is_null]") + not_a_parameter: str | None = Field(default=None, serialization_alias="filter[budget_id][eq]") + + +class BudgetPageMeta(BaseModel): + page: int + page_size: int + total_count: int + total_pages: int + + +class BudgetPageLinks(BaseModel): + first: str + prev: str | None = None + next: str | None = None + last: str + + +class BudgetPageRow(BaseModel): + budget_id: str + max_budget: float | None = None + tpm_limit: int | None = None + budget_duration: str | None = None + + +class BudgetPageResponse(BaseModel): + data: list[BudgetPageRow] + meta: BudgetPageMeta + links: BudgetPageLinks + + +def _list_budgets(client: ManagementClient, params: BudgetPageParams) -> BudgetPageResponse: + return unwrap( + client.proxy.transport.get( + _BUDGETS_V1, + headers=client.proxy.transport.master, + params=params, + response_type=BudgetPageResponse, + ) + ) + + +def _list_budget_ids(client: ManagementClient, params: BudgetPageParams) -> tuple[str, ...]: + return tuple(row.budget_id for row in _list_budgets(client, params).data) + + +def _list_status(client: ManagementClient, params: BudgetPageParams, key: str | None = None) -> int: + headers = client.proxy.transport.master if key is None else client.proxy.transport.bearer(key) + outcome = client.proxy.transport.get( + _BUDGETS_V1, headers=headers, params=params, response_type=BudgetPageResponse + ) + match outcome: + case Success(status_code=status_code): + return status_code + case UnauthorizedError(): + return 401 + case UnknownApiError(status_code=status_code): + return status_code + case _: + raise AssertionError(outcome) + + +class TestBudgetListV1: + """The paged, sorted, filtered budget list the Budgets page reads. + + Every test tags its own budgets with a marker in the budget_id and searches on + it, so budgets left behind by other suites cannot move the assertions. + """ + + @pytest.mark.covers("mgmt.budget.list_v1.happy_path") + def test_sorts_pages_and_filters_the_budgets_it_created( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + marker = unique_marker() + small, medium, large = (f"{marker}-small", f"{marker}-medium", f"{marker}-large") + for budget_id, max_budget, duration in ( + (small, 1.0, "7d"), + (medium, 2.0, "30d"), + (large, 3.0, "30d"), + ): + _create_budget( + client, + resources, + BudgetNewBody( + budget_id=budget_id, max_budget=max_budget, budget_duration=duration, tpm_limit=60000 + ), + ) + + mine = BudgetPageParams(q=marker, sort="-max_budget") + _ = _poll( + client, + lambda: mine if len(_list_budget_ids(client, mine)) == 3 else None, + f"{_BUDGETS_V1} never listed all three budgets tagged {marker}", + ) + + assert _list_budget_ids(client, mine) == (large, medium, small) + + page_two = _list_budgets(client, BudgetPageParams(q=marker, sort="-max_budget", page=2, page_size=1)) + assert [row.budget_id for row in page_two.data] == [medium] + assert page_two.meta.total_count == 3 + assert page_two.meta.total_pages == 3 + assert page_two.meta.page_size == 1 + assert page_two.links.prev is not None and page_two.links.next is not None + + assert set(_list_budget_ids(client, BudgetPageParams(q=marker, duration_in="30d"))) == {medium, large} + + limits = _list_budgets(client, BudgetPageParams(q=marker, sort="budget_id")).data + assert [row.tpm_limit for row in limits] == [60000, 60000, 60000] + + @pytest.mark.covers("mgmt.budget.list_v1.happy_path") + def test_is_null_finds_the_budget_left_uncapped( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + marker = unique_marker() + uncapped = _create_budget(client, resources, BudgetNewBody(budget_id=f"{marker}-uncapped")) + _ = _create_budget(client, resources, BudgetNewBody(budget_id=f"{marker}-capped", max_budget=4.0)) + + params = BudgetPageParams(q=marker, max_budget_is_null=True) + found = _poll( + client, + lambda: params if _list_budget_ids(client, params) == (uncapped,) else None, + f"{_BUDGETS_V1} never isolated the uncapped budget {uncapped}", + ) + + assert [row.max_budget for row in _list_budgets(client, found).data] == [None] + + @pytest.mark.covers("mgmt.budget.list_v1.happy_path") + def test_refuses_a_sort_field_and_a_parameter_it_does_not_support(self, client: ManagementClient) -> None: + assert _list_status(client, BudgetPageParams(sort="budget_duration")) == 400 + assert _list_status(client, BudgetPageParams(not_a_parameter="b-1")) == 400 + + @pytest.mark.covers("mgmt.budget.list_v1.admin_only") + def test_is_refused_for_a_non_admin_key(self, client: ManagementClient, resources: ResourceManager) -> None: + key = client.proxy.generate_key(KeyGenerateBody()) + resources.defer(lambda: client.proxy.delete_key(key)) + + status = _list_status(client, BudgetPageParams(), key=key) + + assert status in (401, 403), ( + f"a non-admin key listing budgets must be refused 401/403, got {status}. Serving 200 with an " + f"empty page would read as 'this proxy has no budgets'" + ) + + # ---------- customer / end-user ---------- diff --git a/tests/e2e/mcp/test_mcp_datadog_e2e.py b/tests/e2e/mcp/test_mcp_datadog_e2e.py index d093e307f99..138f654272d 100644 --- a/tests/e2e/mcp/test_mcp_datadog_e2e.py +++ b/tests/e2e/mcp/test_mcp_datadog_e2e.py @@ -49,6 +49,16 @@ def _seed_completion(proxy: ProxyClient, *, key: str, marker: str) -> None: class TestDatadogMcpRoundTrip: + @pytest.mark.skip( + reason=( + "LIT-5052: this test sends a `telemetry` argument that Datadog's " + "search_datadog_logs tool now rejects, so every tool call fails validation with " + "'unexpected additional properties [\"telemetry\"]' before the round-trip " + "assertion is reached. `telemetry` was never a documented Datadog parameter; the " + "test relied on the server ignoring unknown properties. Unskip once the argument " + "is dropped." + ) + ) @pytest.mark.covers("mcp.list_tools.api_key.succeeds", "mcp.call_tool.api_key.succeeds") def test_search_logs_finds_seeded_completion( self, diff --git a/tests/e2e/mcp/test_mcp_guardrail_e2e.py b/tests/e2e/mcp/test_mcp_guardrail_e2e.py index 9e3a8c48395..60a349ddc5e 100644 --- a/tests/e2e/mcp/test_mcp_guardrail_e2e.py +++ b/tests/e2e/mcp/test_mcp_guardrail_e2e.py @@ -78,6 +78,16 @@ def _search_on_synced_pod( class TestMcpToolCallGuardrail: + @pytest.mark.skip( + reason=( + "LIT-5052: the control call sends a `telemetry` argument that Datadog's " + "search_datadog_logs tool now rejects, so the clean-argument half of this test " + "errors with 'unexpected additional properties [\"telemetry\"]' and the guardrail " + "block it exists to prove is never exercised. `telemetry` was never a documented " + "Datadog parameter; the test relied on the server ignoring unknown properties. " + "Unskip once the argument is dropped." + ) + ) @pytest.mark.covers( "guardrail.litellm_content_filter.pre_mcp_call.blocks", exercised_on=["mcp_operations"], diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 678424e36d1..788a0a3f45c 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -51,6 +51,16 @@ class TestMcpKeyWithoutAccessIsDenied: f"boundary: {denied_tools}" ) + @pytest.mark.skip( + reason=( + "LIT-5052: the control call proving a granted key CAN invoke the tool sends a " + "`telemetry` argument that Datadog's search_datadog_logs tool now rejects, so it " + "errors with 'unexpected additional properties [\"telemetry\"]' and the denial " + "assertion is never reached. `telemetry` was never a documented Datadog " + "parameter; the test relied on the server ignoring unknown properties. Unskip " + "once the argument is dropped." + ) + ) @pytest.mark.covers("mcp.call_tool.api_key.denied_without_permission") def test_call_tool_denied_without_permission( self, diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index e31c3953714..1f260f86eeb 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -142,6 +142,77 @@ def test_cli_extra_is_a_thin_client_install(): assert not leaked, f"`cli` extra leaks proxy-server deps onto laptops: {leaked}" +AIOHTTP_POOL_POISONING_RANGE = ">=3.14.0,<3.14.2" +AIOHTTP_POOL_POISONING_RELEASES = ("3.14.0", "3.14.1") + + +def _load_toml(path): + try: + import tomllib as tomli + except ImportError: + try: + import tomli + except ImportError: + pytest.skip("tomli/tomllib not available - skipping dependency check") + + with open(path, "rb") as f: + return tomli.load(f) + + +def _declared_aiohttp_specifier(): + from packaging.requirements import Requirement + + pyproject = _load_toml(os.path.join(PROJECT_ROOT, "pyproject.toml")) + for requirement in pyproject["project"]["dependencies"]: + parsed = Requirement(requirement) + if parsed.name.lower() == "aiohttp": + return parsed.specifier + pytest.fail("aiohttp is no longer a declared runtime dependency of litellm") + + +def _locked_aiohttp_version(): + lock = _load_toml(os.path.join(PROJECT_ROOT, "uv.lock")) + for package in lock["package"]: + if package["name"].lower() == "aiohttp": + return package["version"] + pytest.fail("aiohttp is missing from uv.lock") + + +def test_declared_aiohttp_floor_excludes_pool_poisoning_releases(): + """aiohttp 3.14.0/3.14.1 re-arm the sock_read timer on a keep-alive connection + after it is back in the idle pool, so the next request to reuse it fails + instantly with a bogus timeout (aio-libs/aiohttp#12953, fixed in 3.14.2). + + The wheel's own metadata is what pip resolves against, so the floor declared + here - not just the lockfile - has to exclude that range. + """ + specifier = _declared_aiohttp_specifier() + + admitted = [v for v in AIOHTTP_POOL_POISONING_RELEASES if specifier.contains(v)] + assert not admitted, ( + f"litellm declares aiohttp{specifier}, which still admits {admitted}. " + "Those releases poison pooled keep-alive connections and cause " + "cross-provider sub-millisecond 'Connection timed out' failures; " + "keep the floor at >=3.14.2." + ) + + +def test_locked_aiohttp_version_is_not_pool_poisoning(): + """uv.lock is what the published Docker images install (uv sync --frozen), so a + lock that drifts back onto 3.14.0/3.14.1 ships the regression regardless of + what pyproject.toml declares. + """ + from packaging.specifiers import SpecifierSet + + locked = _locked_aiohttp_version() + + assert not SpecifierSet(AIOHTTP_POOL_POISONING_RANGE).contains(locked), ( + f"uv.lock resolves aiohttp {locked}, which is inside the pool-poisoning " + f"range {AIOHTTP_POOL_POISONING_RANGE} (aio-libs/aiohttp#12953). " + "Re-run `uv lock` against an aiohttp>=3.14.2 floor." + ) + + import os import subprocess import time diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py index 0550898c73d..b0b092a541f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py @@ -1,4 +1,5 @@ import asyncio +import concurrent.futures import os import sys @@ -827,3 +828,305 @@ async def test_stale_loop_rebuild_does_not_close_unowned_session(): shared_session._loop = running_loop other_loop.close() await shared_session.close() + + +# --------------------------------------------------------------------------- +# Recycled-session leak tests (#24230) +# --------------------------------------------------------------------------- + + +async def _new_session() -> aiohttp.ClientSession: + return aiohttp.ClientSession() + + +def _make_session_on_dead_loop() -> aiohttp.ClientSession: + """Create a ClientSession bound to an event loop that is then closed. + + Runs in a worker thread: the caller may already be inside a running + event loop, where a nested run_until_complete is forbidden. + """ + import threading + + result: dict = {} + + def build() -> None: + loop = asyncio.new_event_loop() + try: + result["session"] = loop.run_until_complete(_new_session()) + finally: + loop.close() + + thread = threading.Thread(target=build) + thread.start() + thread.join(5) + return result["session"] + + +def _flaky_get_running_loop_factory(): + """get_running_loop stand-in that fails once, then delegates. + + Reproduces #24230: a transient loop-inspection failure sends + _get_valid_client_session into its (RuntimeError, AttributeError) + fallback branch. + """ + real_get_running_loop = asyncio.get_running_loop + calls = {"count": 0} + + def flaky(): + calls["count"] += 1 + if calls["count"] == 1: + raise RuntimeError("simulated loop inspection failure") + return real_get_running_loop() + + return flaky + + +@pytest.mark.asyncio +async def test_fallback_recreate_closes_previous_session(): + """ + Regression test for #24230: when loop inspection fails and the fallback + branch recreates the session, the replaced session must still be closed - + not silently abandoned to the garbage collector. + """ + from unittest.mock import patch + + old_session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = old_session + + with patch( + "litellm.llms.custom_httpx.aiohttp_transport.asyncio.get_running_loop", + side_effect=_flaky_get_running_loop_factory(), + ): + new_session = transport._get_valid_client_session() + + try: + assert new_session is not old_session + for _ in range(3): + await asyncio.sleep(0) + assert old_session.closed, "replaced session must be closed, not leaked" + finally: + await new_session.close() + if not old_session.closed: + await old_session.close() + + +@pytest.mark.asyncio +async def test_replaced_session_emits_no_unclosed_warnings(): + """ + Regression test for #24230: a session replaced by the fallback branch must + not surface "Unclosed client session" / "Unclosed connector" warnings when + the garbage collector finalizes it. + """ + import gc + import warnings as warnings_mod + from unittest.mock import patch + + old_session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = old_session + + with patch( + "litellm.llms.custom_httpx.aiohttp_transport.asyncio.get_running_loop", + side_effect=_flaky_get_running_loop_factory(), + ): + new_session = transport._get_valid_client_session() + + try: + for _ in range(3): + await asyncio.sleep(0) + + del old_session + with warnings_mod.catch_warnings(record=True) as caught: + warnings_mod.simplefilter("always") + gc.collect() + + unclosed = [ + str(w.message) + for w in caught + if "Unclosed client session" in str(w.message) or "Unclosed connector" in str(w.message) + ] + assert not unclosed, f"leaked session warnings: {unclosed}" + finally: + await new_session.close() + + +@pytest.mark.asyncio +async def test_dead_loop_session_closed_synchronously_on_recycle(): + """ + Regression test for #24230: a session whose event loop is already closed + cannot run an async close anywhere. Recycling it must dispose of it + deterministically, the session reads closed as soon as the recycle + returns, so no finalizer warning window remains. + """ + old_session = _make_session_on_dead_loop() + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = old_session + + new_session = transport._get_valid_client_session() + + try: + assert new_session is not old_session + assert old_session.closed, "session from a closed loop must be disposed synchronously at recycle" + finally: + await new_session.close() + + +@pytest.mark.asyncio +async def test_close_task_strongly_referenced_until_done(): + """ + Regression test for #24230: scheduled session-close tasks must be strongly + referenced (and pruned on completion) so they cannot be garbage-collected + before they run. + """ + old_session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + + transport._close_recycled_session(old_session) + + assert LiteLLMAiohttpTransport._background_close_tasks, "close task must be strongly referenced while pending" + for _ in range(5): + await asyncio.sleep(0) + assert old_session.closed + assert not LiteLLMAiohttpTransport._background_close_tasks, "completed close tasks must be pruned from the registry" + + +@pytest.mark.asyncio +async def test_session_from_other_running_loop_closed_threadsafe(): + """ + Regression test for #24230: a session that belongs to a loop still running + in another thread must be closed on its own loop (thread-safe), not driven + from the current loop. + """ + import threading + import time + + ready = threading.Event() + holder: dict = {} + + def worker() -> None: + loop = asyncio.new_event_loop() + holder["loop"] = loop + + async def make() -> None: + holder["session"] = aiohttp.ClientSession() + + loop.run_until_complete(make()) + ready.set() + loop.run_forever() + loop.close() + + thread = threading.Thread(target=worker, daemon=True) + thread.start() + assert ready.wait(5), "worker loop failed to start" + + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = holder["session"] + + new_session = transport._get_valid_client_session() + + try: + deadline = time.monotonic() + 5 + while not holder["session"].closed and time.monotonic() < deadline: + await asyncio.sleep(0.01) + assert holder["session"].closed, "foreign-loop session was never closed" + finally: + holder["loop"].call_soon_threadsafe(holder["loop"].stop) + thread.join(5) + await new_session.close() + + +def test_threadsafe_close_done_callback_tolerates_cancelled_future(): + """ + Regression test for #24230 (review finding): when the foreign loop stops + before the handed-off close coroutine runs, asyncio cancels the + concurrent.futures.Future. The done-callback must return quietly instead + of letting future.exception() raise CancelledError (a BaseException that + escapes _invoke_callbacks and crashes the foreign loop's thread). + """ + future: "concurrent.futures.Future[None]" = concurrent.futures.Future() + future.cancel() + + LiteLLMAiohttpTransport._on_threadsafe_close_done(future) + + +@pytest.mark.asyncio +async def test_session_closed_retry_does_not_close_concurrent_replacement(): + """ + Regression test for #24230 (review finding): when the "Session is closed" + retry fires, the handler must dispose the session that actually faulted, + not self.client - a concurrent task may already have replaced self.client + with a healthy session, which must stay open. + """ + from unittest.mock import patch + + faulted_session = aiohttp.ClientSession() + healthy_replacement = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = faulted_session + + calls = {"n": 0} + + async def fake_make_request(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + # simulate a concurrent task replacing the shared session between + # the failed await and the exception handler + transport.client = healthy_replacement + raise RuntimeError("Session is closed") + raise StopAsyncIteration("stop after retry dispatch") + + with patch.object(transport, "_make_aiohttp_request", side_effect=fake_make_request): + with pytest.raises(Exception): + await transport.handle_async_request(httpx.Request("GET", "http://example.com")) + + try: + assert not healthy_replacement.closed, "concurrent replacement session must not be closed by the retry handler" + for _ in range(3): + await asyncio.sleep(0) + assert faulted_session.closed, "the faulted session must be disposed" + finally: + await faulted_session.close() + await healthy_replacement.close() + new_session = transport.client + if isinstance(new_session, aiohttp.ClientSession): + await new_session.close() + + +@pytest.mark.asyncio +async def test_stopped_loop_session_disposed_synchronously_on_recycle(): + """ + Regression test for #24230 (review finding): a session whose loop is + stopped but not yet closed cannot safely run an async close on another + loop, and nothing will ever process a close handed to the stopped loop. + Recycling must dispose it synchronously, like the closed-loop case. + """ + import threading + + result: dict = {} + + def build() -> None: + loop = asyncio.new_event_loop() + + async def make() -> None: + result["session"] = aiohttp.ClientSession() + + loop.run_until_complete(make()) + result["loop"] = loop # stopped, deliberately NOT closed + + thread = threading.Thread(target=build) + thread.start() + thread.join(5) + + old_session = result["session"] + transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession()) + transport.client = old_session + + new_session = transport._get_valid_client_session() + + try: + assert new_session is not old_session + assert old_session.closed, "session from a stopped (not yet closed) loop must be disposed synchronously" + finally: + await new_session.close() + result["loop"].close() diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index d99c190c6e5..8ee8186f6bb 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -2097,3 +2097,319 @@ def test_multi_turn_function_calling_roles(): assert ( content["role"] == "user" ), f"Content block {i} with function_response has role='{content['role']}', expected 'user'" + + +def test_gemini_thought_signature_preservation_real_response(): + """Test that thought signatures are preserved on the text part if originally there, without dropping or duplicating (real response case).""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + real_candidate = { + "content": { + "parts": [ + { + "text": "I will explain and then list files.", + "thoughtSignature": "mock_signature_from_text_part", + }, + { + "functionCall": { + "name": "list_files", + "args": {}, + } + }, + ] + } + } + + parts = real_candidate["content"]["parts"] + + content, reasoning_content = ( + VertexGeminiConfig().get_assistant_content_message(parts=parts) + ) + thought_signatures = ( + VertexGeminiConfig()._extract_thought_signatures_from_parts( + parts=parts + ) + ) + functions, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts, + cumulative_tool_call_idx=0, + is_function_call=False, + ) + + msg: dict = {"role": "assistant"} + if content is not None: + msg["content"] = content + if tools: + msg["tool_calls"] = tools + if functions is not None: + msg["function_call"] = functions + if thought_signatures is not None: + msg["provider_specific_fields"] = { + "thought_signatures": thought_signatures + } + + converted_real = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted_real) == 1 + assert "parts" in converted_real[0] + parts_out = converted_real[0]["parts"] + assert len(parts_out) == 2 + assert "text" in parts_out[0] + assert ( + parts_out[0]["thoughtSignature"] == "mock_signature_from_text_part" + ) + assert "function_call" in parts_out[1] + assert "thoughtSignature" not in parts_out[1] + + +def test_gemini_thought_signature_deduplication_assumed_response(): + """Test that thought signatures are deduplicated and not attached to the text part if already present in the tool call (assumed response case).""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + pr_assumed_msg = { + "role": "assistant", + "content": "I will list the directory.", + "provider_specific_fields": { + "thought_signatures": ["mock_signature_63k"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": { + "thought_signature": "mock_signature_63k" + }, + } + ], + } + + converted_pr = _gemini_convert_messages_with_history( + messages=[pr_assumed_msg], + model="gemini-2.5-pro", + ) + + assert len(converted_pr) == 1 + assert "parts" in converted_pr[0] + parts_out = converted_pr[0]["parts"] + assert len(parts_out) == 2 + assert "text" in parts_out[0] + assert "thoughtSignature" not in parts_out[0] + assert "function_call" in parts_out[1] + assert parts_out[1]["thoughtSignature"] == "mock_signature_63k" + + +def test_gemini_thought_signature_pure_text(): + """Test that thought signatures are preserved on the text part for responses with no tool calls.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Hello, I am a model.", + "provider_specific_fields": { + "thought_signatures": ["pure_text_signature"] + }, + } + + converted = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted) == 1 + assert "parts" in converted[0] + parts_out = converted[0]["parts"] + assert len(parts_out) == 1 + assert "text" in parts_out[0] + assert parts_out[0]["thoughtSignature"] == "pure_text_signature" + + +def test_gemini_thought_signature_pure_tool_call(): + """Test that thought signatures are preserved on the tool call for responses with no intermediate text.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "thought_signatures": ["pure_tool_signature"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": { + "thought_signature": "pure_tool_signature" + }, + } + ], + } + + converted = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted) == 1 + assert "parts" in converted[0] + parts_out = converted[0]["parts"] + assert len(parts_out) == 1 + assert "function_call" in parts_out[0] + assert parts_out[0]["thoughtSignature"] == "pure_tool_signature" + + +def test_gemini_distinct_text_and_tool_signatures_are_both_preserved(): + """A text-part signature that differs from the tool-call signature must stay on the text part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Some analysis.", + "provider_specific_fields": { + "thought_signatures": ["text_signature", "tool_signature"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": {"thought_signature": "tool_signature"}, + } + ], + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-2.5-pro" + )[0]["parts"] + + assert parts[0]["text"] == "Some analysis." + assert parts[0]["thoughtSignature"] == "text_signature" + assert "function_call" in parts[1] + assert parts[1]["thoughtSignature"] == "tool_signature" + + +def test_gemini_25_text_signature_survives_replay_to_gemini_3(): + """gemini-2.5 history (signed text, unsigned tool call) replayed to gemini-3 keeps the real + text signature; the dummy signature synthesized for the unsigned tool call must not suppress it.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + _get_dummy_thought_signature, + ) + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "I will list the directory.", + "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + } + ], + } + + parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ + 0 + ]["parts"] + + assert parts[0]["text"] == "I will list the directory." + assert parts[0]["thoughtSignature"] == "real_25_signature" + assert "function_call" in parts[1] + assert parts[1]["thoughtSignature"] == _get_dummy_thought_signature() + + +def test_gemini_function_call_signature_round_trip_no_duplicate(): + """End to end: a gemini-3-style response (unsigned text + signed functionCall) parsed and + re-serialized sends the signature exactly once, on the function-call part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + response_parts = [ + {"text": "I will calculate the result for you."}, + { + "functionCall": {"name": "add_numbers", "args": {"a": 17, "b": 25}}, + "thoughtSignature": "signature_from_function_call", + }, + ] + + config = VertexGeminiConfig() + content, _ = config.get_assistant_content_message(parts=response_parts) + thought_signatures = config._extract_thought_signatures_from_parts( + parts=response_parts + ) + _, tools, _ = VertexGeminiConfig._transform_parts( + parts=response_parts, cumulative_tool_call_idx=0, is_function_call=False + ) + + msg = { + "role": "assistant", + "content": content, + "tool_calls": tools, + "provider_specific_fields": {"thought_signatures": thought_signatures}, + } + + parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ + 0 + ]["parts"] + + signatures = [p["thoughtSignature"] for p in parts if "thoughtSignature" in p] + assert signatures == ["signature_from_function_call"] + assert "thoughtSignature" not in parts[0] + assert "function_call" in parts[1] + + +def test_gemini_server_side_tool_signature_not_duplicated_on_text(): + """A signature already re-injected on a server-side toolCall part is not attached to the text part again.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "The weather in Buenos Aires is sunny.", + "provider_specific_fields": { + "thought_signatures": ["server_side_signature"], + "server_side_tool_invocations": [ + { + "tool_type": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "args": {"queries": ["weather Buenos Aires"]}, + "response": {"weather": "Sunny"}, + "thought_signature": "server_side_signature", + } + ], + }, + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-2.5-pro" + )[0]["parts"] + + text_part = next(p for p in parts if "text" in p) + assert "thoughtSignature" not in text_part + tool_call_part = next(p for p in parts if "toolCall" in p) + assert tool_call_part["thoughtSignature"] == "server_side_signature" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 89b8f018e5c..0b95a497882 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -8191,3 +8191,177 @@ class TestGetUserObjectPermission: async def test_no_user_id_places_no_ceiling(self): assert await MCPRequestHandler._get_user_object_permission(UserAPIKeyAuth(api_key="sk-test")) is None assert await MCPRequestHandler._get_user_object_permission(None) is None + + +def _key_auth_reaching(server, *, tools=None, **fields): + """A key-authenticated caller whose OWN key grant reaches ``server`` (and optionally its ``tools``). + + The key grant is the thing an upper-level entitlement fault must not silently hand back: every + test below asserts against what this key reaches when the level under test cannot be resolved. + """ + return UserAPIKeyAuth( + api_key="sk-hash", + user_id="u1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", + mcp_servers=[server], + mcp_tool_permissions={server: tools} if tools else None, + ), + **fields, + ) + + +def _agent_prisma(object_permission_id=None, side_effect=None): + prisma_client = MagicMock() + prisma_client.db.litellm_agentstable.find_unique = AsyncMock( + return_value=MagicMock(object_permission_id=object_permission_id), + side_effect=side_effect, + ) + return prisma_client + + +@contextlib.contextmanager +def _entitlement_fault_globals(prisma_client=None): + from litellm.caching.dual_cache import DualCache + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client or MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + ): + yield + + +@pytest.mark.asyncio +class TestEntitlementFaultSemantics: + """Each entitlement level distinguishes two fault classes for a KEY-authenticated caller. + + A principal row that NAMES an object_permission we cannot load is a known entitlement with + unknown contents, so the level denies rather than handing back the wider key scope. A lookup + that fails before we can tell whether the principal is entitled at all leaves no ceiling, which + is the state that existed before the level did; denying there would refuse MCP to the majority + of callers, who have no such entitlement configured, for the duration of a cold-cache fault. + """ + + async def test_end_user_named_but_unloadable_permission_denies(self): + end_user = MagicMock(object_permission=None, object_permission_id="op-eu") + auth = _key_auth_reaching("srv1", end_user_id="eu-1") + with _entitlement_fault_globals(): + with ( + patch("litellm.proxy.auth.auth_checks.get_end_user_object", AsyncMock(return_value=end_user)), + patch("litellm.proxy.auth.auth_checks.get_object_permission", AsyncMock(return_value=None)), + ): + allowed = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed == [], "an end-user entitlement we know exists but cannot read must deny" + + async def test_end_user_without_an_entitlement_places_no_ceiling(self): + """The three shapes that are NOT evidence of an entitlement: an end user row linking no + permission, no end user row at all, and a lookup that blew up before answering either.""" + auth = _key_auth_reaching("srv1", end_user_id="eu-1") + linked_none = MagicMock(object_permission=None, object_permission_id=None) + for lookup, shape in ( + (AsyncMock(return_value=linked_none), "row links no permission"), + (AsyncMock(return_value=None), "no end user row"), + (AsyncMock(side_effect=RuntimeError("connection reset by peer")), "lookup failed"), + ): + with _entitlement_fault_globals(): + with patch("litellm.proxy.auth.auth_checks.get_end_user_object", lookup): + allowed = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(allowed) == {"srv1"}, f"{shape}: no evidence of an entitlement, so no ceiling" + + async def test_agent_named_but_unloadable_permission_denies(self): + auth = _key_auth_reaching("srv1", agent_id="agent-unloadable") + with _entitlement_fault_globals(_agent_prisma(object_permission_id="op-agent")): + with patch("litellm.proxy.auth.auth_checks.get_object_permission", AsyncMock(return_value=None)): + allowed = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed == [], "an agent entitlement we know exists but cannot read must deny" + + async def test_agent_without_an_entitlement_places_no_ceiling(self): + """An agent row linking no permission, and an agent row we could not read at all.""" + for prisma_client, agent_id, shape in ( + (_agent_prisma(object_permission_id=None), "agent-unlinked", "agent links no permission"), + (_agent_prisma(side_effect=RuntimeError("connection reset by peer")), "agent-unread", "row read failed"), + ): + auth = _key_auth_reaching("srv1", agent_id=agent_id) + with _entitlement_fault_globals(prisma_client): + allowed = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(allowed) == {"srv1"}, f"{shape}: no evidence of an entitlement, so no ceiling" + + async def test_agent_named_but_unloadable_permission_denies_tools(self): + """The tools axis denies with [] rather than the None (allow-all) key auth gets for an + indeterminate fault, so an unreadable agent entitlement cannot widen the key's tool scope.""" + auth = _key_auth_reaching("srv1", tools=["tool_a"], agent_id="agent-tools-unloadable") + with _entitlement_fault_globals(_agent_prisma(object_permission_id="op-agent")): + with patch("litellm.proxy.auth.auth_checks.get_object_permission", AsyncMock(return_value=None)): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert tools == [], "an agent entitlement we know exists but cannot read must deny its tools" + + async def test_org_named_but_unloadable_ceiling_denies(self): + auth = _key_auth_reaching("srv1", org_id="org-a") + org = MagicMock(object_permission_id="op-org") + with _entitlement_fault_globals(): + with ( + patch("litellm.proxy.auth.auth_checks.get_org_object", AsyncMock(return_value=org)), + patch("litellm.proxy.auth.auth_checks.get_object_permission", AsyncMock(return_value=None)), + ): + allowed = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed == [], "an org ceiling we know exists but cannot read must deny, key auth included" + + async def test_org_named_but_unloadable_ceiling_denies_tools(self): + auth = _key_auth_reaching("srv1", tools=["tool_a"], org_id="org-a") + org = MagicMock(object_permission_id="op-org") + with _entitlement_fault_globals(): + with ( + patch("litellm.proxy.auth.auth_checks.get_org_object", AsyncMock(return_value=org)), + patch("litellm.proxy.auth.auth_checks.get_object_permission", AsyncMock(return_value=None)), + ): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert tools == [], "an org tool ceiling we know exists but cannot read must deny its tools" + + async def test_org_without_a_resolvable_entitlement_places_no_ceiling(self): + """A deleted org and an org lookup that failed are both cases where we cannot point at a + ceiling; key auth keeps its long-standing fail-open behavior for them.""" + from litellm.proxy.auth.auth_checks import OrganizationNotFoundError + + auth = _key_auth_reaching("srv1", org_id="org-a") + for lookup, shape in ( + (AsyncMock(return_value=MagicMock(object_permission_id=None)), "org names no permission"), + (AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db.")), "org deleted"), + (AsyncMock(side_effect=RuntimeError("connection reset by peer")), "org lookup failed"), + ): + with _entitlement_fault_globals(): + with patch("litellm.proxy.auth.auth_checks.get_org_object", lookup): + allowed = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(allowed) == {"srv1"}, f"{shape}: no ceiling we can point at, so key auth stays open" + + async def test_keyless_org_ceiling_denies_on_either_fault_class(self): + """The keyless gateway-admitted path is untouched: it already denied on ANY org-ceiling + fault, and still denies on both classes, because a per-source org ceiling is the only org + bound a keyless subject has and an unbounded source would win the union.""" + auth = _make_admitted_subject("sso-user", org_id="org-a", own_servers=["srv1"]) + org = MagicMock(object_permission_id="op-org") + with _entitlement_fault_globals(): + with patch("litellm.proxy.auth.auth_checks.get_org_object", AsyncMock(return_value=org)): + with patch("litellm.proxy.auth.auth_checks.get_object_permission", AsyncMock(return_value=None)): + named_unloadable = await MCPRequestHandler.get_allowed_mcp_servers(auth) + with patch( + "litellm.proxy.auth.auth_checks.get_org_object", + AsyncMock(side_effect=RuntimeError("connection reset by peer")), + ): + indeterminate = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert named_unloadable == [] and indeterminate == [] + + async def test_keyless_source_never_consults_the_end_user_or_agent_levels(self): + """A keyless subject's grant sources carry neither end_user_id nor agent_id, so neither + level runs for it and neither new deny can reach its union. Pinned because a source that + DID consult them would fail closed on a fault and silently drop a team's grants.""" + auth = _make_admitted_subject("sso-user", own_servers=["srv1"]) + auth.end_user_id = "eu-1" + auth.agent_id = "agent-unloadable" + with _entitlement_fault_globals(_agent_prisma(object_permission_id="op-agent")): + with ( + patch("litellm.proxy.auth.auth_checks.get_end_user_object", AsyncMock(side_effect=AssertionError)), + patch("litellm.proxy.auth.auth_checks.get_object_permission", AsyncMock(return_value=None)), + ): + allowed = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(allowed) == {"srv1"} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index a710da81962..0d130767bd5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -7,6 +7,10 @@ also guards reachability: a dropped `case` would hit `assert_never` and raise in returning the stub. """ +import asyncio +import logging +from datetime import datetime, timedelta, timezone + import httpx import pytest from pydantic import SecretStr @@ -25,6 +29,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( NoOpAuth, Ok, PassthroughConfig, + PrivateKeyJwtAuth, Result, ServerSpec, SharedKey, @@ -37,6 +42,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto OAuthToken, TokenStoreUnavailable, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + AssertionStoreUnavailable, + SSOIdentityAssertion, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( ExchangedToken, ) @@ -71,6 +80,23 @@ def _with_inbound(token: str) -> Subject: return Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr(token)) +class _FakeAssertionStore: + """The SSO assertion read seam, canned per user_id and recording every lookup.""" + + def __init__(self, assertions: dict[str, SSOIdentityAssertion] | None = None) -> None: + self._assertions = dict(assertions or {}) + self.lookups: list[str] = [] + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + self.lookups.append(user_id) + return self._assertions.get(user_id) + + +def _assertion(id_token: str, expires_in: timedelta | None = timedelta(minutes=30)) -> SSOIdentityAssertion: + expires_at = datetime.now(timezone.utc) + expires_in if expires_in is not None else None + return SSOIdentityAssertion(id_token=SecretStr(id_token), expires_at=expires_at) + + def _spec(config): return ServerSpec(server_id="s", resource="https://upstream.example.com", config=config) @@ -471,9 +497,10 @@ async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer(): @pytest.mark.asyncio -async def test_id_jag_without_inbound_token_is_precondition_required_no_http(): +async def test_id_jag_without_inbound_token_or_stored_assertion_is_precondition_required_no_http(): endpoint = _FakeTokenEndpoint([]) - provider = UpstreamCredentialProvider(token_endpoint=endpoint) + store = _FakeAssertionStore() + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) result = await provider.resolve_credentials( Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) ) @@ -481,6 +508,397 @@ async def test_id_jag_without_inbound_token_is_precondition_required_no_http(): assert isinstance(result, Error) assert result.error.tag == "precondition_required" assert endpoint.calls == [] + assert store.lookups == ["alice"] + + +@pytest.mark.asyncio +async def test_id_jag_exchanges_the_stored_sso_assertion_when_the_caller_presents_no_token(): + """The agent-triggered flow: a brokered LiteLLM credential carries no IdP token, so leg 1's + subject is the assertion captured for that user at SSO login.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + store = _FakeAssertionStore({"alice": _assertion("alice-id-token")}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Ok) + assert _emitted(result.ok)["Authorization"] == "Bearer final-access" + assert store.lookups == ["alice"] + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "alice-id-token" + assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" + + +@pytest.mark.asyncio +async def test_id_jag_prefers_the_callers_own_token_over_the_stored_assertion(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + store = _FakeAssertionStore({"alice": _assertion("stored-id-token")}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials(_with_inbound("inbound-id-token"), _spec(_id_jag_config())) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "inbound-id-token" + assert store.lookups == [] + + +@pytest.mark.asyncio +async def test_id_jag_refuses_an_expired_stored_assertion_without_calling_the_idp(): + endpoint = _FakeTokenEndpoint([]) + store = _FakeAssertionStore({"alice": _assertion("stale-id-token", expires_in=-timedelta(seconds=1))}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_id_jag_accepts_a_stored_assertion_that_declares_no_expiry(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + store = _FakeAssertionStore({"alice": _assertion("undated-id-token", expires_in=None)}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "undated-id-token" + + +@pytest.mark.asyncio +async def test_id_jag_never_reads_the_store_for_an_unidentified_caller(): + """An empty subject_id must not select a credential; otherwise every anonymous caller would + share one store slot.""" + endpoint = _FakeTokenEndpoint([]) + store = _FakeAssertionStore({"": _assertion("anonymous-slot")}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id=""), _spec(_id_jag_config())) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert store.lookups == [] + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_id_jag_keeps_store_sourced_bearers_partitioned_per_user(): + endpoint = _FakeTokenEndpoint( + [ + Ok(ExchangedToken(access_token="alice-id-jag", expires_in=300)), + Ok(ExchangedToken(access_token="alice-bearer", expires_in=3600)), + Ok(ExchangedToken(access_token="bob-id-jag", expires_in=300)), + Ok(ExchangedToken(access_token="bob-bearer", expires_in=3600)), + ] + ) + store = _FakeAssertionStore( + {"alice": _assertion("alice-id-token"), "bob": _assertion("bob-id-token")} + ) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + alice = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + bob = await provider.resolve_credentials(Subject(tenant_id="", subject_id="bob"), _spec(_id_jag_config())) + + assert isinstance(alice, Ok) and isinstance(bob, Ok) + assert _emitted(alice.ok)["Authorization"] == "Bearer alice-bearer" + assert _emitted(bob.ok)["Authorization"] == "Bearer bob-bearer" + + +_DRIVER_DETAIL = "could not connect to host=pg-primary.internal port=5432 user=litellm" + + +class _OutageAssertionStore: + """A store whose backing DB is down, failing with a driver message full of internals.""" + + def __init__(self) -> None: + self.lookups: list[str] = [] + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + self.lookups.append(user_id) + raise AssertionStoreUnavailable(_DRIVER_DETAIL) + + +@pytest.mark.asyncio +async def test_id_jag_maps_an_assertion_store_outage_to_upstream_unavailable(): + """A store outage must not escape as an unhandled error, and must not be reported as a missing + assertion: telling the user to sign in again does not fix a database that is down.""" + endpoint = _FakeTokenEndpoint([]) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=_OutageAssertionStore()) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_id_jag_store_outage_does_not_leak_driver_detail_to_the_caller(caplog): + """`upstream_unavailable` is rendered into the 503 body verbatim, so the driver's message, which + can name hosts, ports and users, must stay out of the summary and go to the log instead.""" + provider = UpstreamCredentialProvider( + token_endpoint=_FakeTokenEndpoint([]), sso_assertion_store=_OutageAssertionStore() + ) + + with caplog.at_level(logging.WARNING): + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert _DRIVER_DETAIL not in result.error.summary + assert "pg-primary.internal" not in result.error.summary + # The operator still needs it, so it must be in the log. + assert _DRIVER_DETAIL in caplog.text + + +@pytest.mark.asyncio +async def test_id_jag_invalidation_survives_an_assertion_store_outage(): + """invalidate_credentials runs on the upstream-401 retry path, so a store outage there must be + swallowed rather than turning a recoverable 401 into a 500.""" + provider = UpstreamCredentialProvider( + token_endpoint=_FakeTokenEndpoint([]), sso_assertion_store=_OutageAssertionStore() + ) + + await provider.invalidate_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + +class _FlakyAssertionStore: + """Serves an assertion, but fails while ``down`` is set.""" + + def __init__(self, assertion: SSOIdentityAssertion) -> None: + self._assertion = assertion + self.down = False + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + if self.down: + raise AssertionStoreUnavailable("connection refused") + return self._assertion + + +@pytest.mark.asyncio +async def test_id_jag_evicts_the_rejected_bearer_even_if_the_store_is_down_during_invalidation(): + """The upstream-401 recovery sequence with a transient store blip. + + Invalidation runs while the store is unreachable and the store recovers before the retry + resolves. Deriving the eviction key from a fresh lookup would evict nothing and then recompute + the identical key, handing the retry the very bearer the upstream just rejected. + """ + endpoint = _FakeTokenEndpoint(_two_leg_ok("rejected-bearer") + _two_leg_ok("reminted-bearer")) + store = _FlakyAssertionStore(_assertion("alice-id-token")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + subject = Subject(tenant_id="", subject_id="alice") + spec = _spec(_id_jag_config()) + + first = await provider.resolve_credentials(subject, spec) + assert isinstance(first, Ok) + assert _emitted(first.ok)["Authorization"] == "Bearer rejected-bearer" + + store.down = True + await provider.invalidate_credentials(subject, spec) + store.down = False + + second = await provider.resolve_credentials(subject, spec) + assert isinstance(second, Ok) + assert _emitted(second.ok)["Authorization"] == "Bearer reminted-bearer" + assert len(endpoint.calls) == 4 + + +class _SwitchableAssertionStore: + """Serves whichever assertion the test currently points it at, as a re-login would.""" + + def __init__(self, id_token: str) -> None: + self.id_token = id_token + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + return _assertion(self.id_token) + + +@pytest.mark.asyncio +async def test_id_jag_invalidation_clears_every_live_bearer_for_the_principal(): + """Overlapping store-sourced requests for one principal can hold different keys (a re-login + between them mints a different subject token). Invalidation must clear all of them: keeping + only the newest would let one request's 401 recovery evict the other's entry and leave its own + rejected bearer cached to be replayed on the retry.""" + endpoint = _FakeTokenEndpoint( + _two_leg_ok("bearer-from-first") + _two_leg_ok("bearer-from-second") + _two_leg_ok("reminted") + ) + store = _SwitchableAssertionStore("id-token-first") + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + subject = Subject(tenant_id="", subject_id="alice") + spec = _spec(_id_jag_config()) + + first = await provider.resolve_credentials(subject, spec) + store.id_token = "id-token-second" + second = await provider.resolve_credentials(subject, spec) + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(first.ok)["Authorization"] == "Bearer bearer-from-first" + assert _emitted(second.ok)["Authorization"] == "Bearer bearer-from-second" + + await provider.invalidate_credentials(subject, spec) + + # Point the store back at the first token. If that entry had survived the invalidation this + # would replay "bearer-from-first", which is the bearer an upstream may already have rejected. + store.id_token = "id-token-first" + third = await provider.resolve_credentials(subject, spec) + assert isinstance(third, Ok) + assert _emitted(third.ok)["Authorization"] == "Bearer reminted" + + +class _SequentialAssertionStore: + """Issues a distinct assertion per call unless pinned, so concurrent resolutions genuinely + mint distinct credentials rather than collapsing onto one through single-flight.""" + + def __init__(self) -> None: + self.pinned: str | None = None + self.issued: list[str] = [] + self._n = 0 + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + await asyncio.sleep(0) + if self.pinned is not None: + return _assertion(self.pinned) + self._n += 1 + token = f"id-token-{self._n}" + self.issued.append(token) + return _assertion(token) + + +class _CountingTokenEndpoint: + """Mints a unique bearer per exchange and yields, so exchanges interleave.""" + + def __init__(self) -> None: + self._n = 0 + + async def fetch(self, endpoint, client_id, grant_params, client_auth): + await asyncio.sleep(0) + self._n += 1 + return Ok(ExchangedToken(access_token=f"tok-{self._n}", expires_in=3600)) + + +@pytest.mark.asyncio +async def test_id_jag_invalidation_leaves_no_bearer_behind_under_concurrency(): + """After invalidation, no bearer minted before it may ever be served again. + + Drives many overlapping resolutions that each mint a distinct credential, invalidates once, + then replays every subject token that was issued. Any credential the eviction could not reach + would show up here as a replayed pre-invalidation bearer. + """ + concurrency = 20 + endpoint = _CountingTokenEndpoint() + store = _SequentialAssertionStore() + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + subject = Subject(tenant_id="t", subject_id="alice") + spec = _spec(_id_jag_config()) + + results = await asyncio.gather(*(provider.resolve_credentials(subject, spec) for _ in range(concurrency))) + before = {_emitted(r.ok)["Authorization"] for r in results if isinstance(r, Ok)} + issued = list(store.issued) + # Guard the guard: if these collapsed onto one credential the test would prove nothing. + assert len(before) > 1 + + await provider.invalidate_credentials(subject, spec) + + for token in issued: + store.pinned = token + replayed = await provider.resolve_credentials(subject, spec) + assert isinstance(replayed, Ok) + assert _emitted(replayed.ok)["Authorization"] not in before + + +@pytest.mark.asyncio +async def test_id_jag_never_serves_a_bearer_minted_for_a_different_caller(): + """Two unidentified-principal callers share a slot, so the fingerprint, not the key, is what + keeps them apart: a mismatch must read as a miss rather than hand over the other's bearer.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("first-callers-bearer") + _two_leg_ok("second-callers-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + spec = _spec(_id_jag_config()) + + first = await provider.resolve_credentials(_with_inbound("caller-one-token"), spec) + second = await provider.resolve_credentials(_with_inbound("caller-two-token"), spec) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(first.ok)["Authorization"] == "Bearer first-callers-bearer" + assert _emitted(second.ok)["Authorization"] == "Bearer second-callers-bearer" + + +@pytest.mark.asyncio +async def test_id_jag_rotating_the_signing_key_does_not_reuse_the_cached_bearer(): + """The cache key fingerprints the private-key-JWT client auth, so a rotated signing key + re-mints instead of serving a bearer authorized under the retired key.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("old-key-bearer") + _two_leg_ok("new-key-bearer")) + store = _FakeAssertionStore({"alice": _assertion("alice-id-token")}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + subject = Subject(tenant_id="", subject_id="alice") + + def _with_key(pem: str) -> IdJagConfig: + return _id_jag_config().model_copy( + update={"client_auth": PrivateKeyJwtAuth(private_key=SecretStr(pem), key_id="kid-1")} + ) + + first = await provider.resolve_credentials(subject, _spec(_with_key("-----OLD KEY-----"))) + second = await provider.resolve_credentials(subject, _spec(_with_key("-----NEW KEY-----"))) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(first.ok)["Authorization"] == "Bearer old-key-bearer" + assert _emitted(second.ok)["Authorization"] == "Bearer new-key-bearer" + assert len(endpoint.calls) == 4 + + +@pytest.mark.asyncio +async def test_id_jag_reads_a_naive_stored_expiry_as_utc(): + """A stored expires_at that lost its offset must still compare rather than raise: an aware/naive + comparison would be a TypeError on the egress path, turning a 412 into a 500.""" + endpoint = _FakeTokenEndpoint([]) + naive_past = datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(hours=1) + store = _FakeAssertionStore( + {"alice": SSOIdentityAssertion(id_token=SecretStr("stale"), expires_at=naive_past)} + ) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_invalidate_evicts_a_store_sourced_id_jag_bearer(): + """The upstream-401 recovery path. Keyed off the request alone the eviction would miss, and the + rejected bearer would be replayed until its TTL.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("first-bearer") + _two_leg_ok("second-bearer")) + store = _FakeAssertionStore({"alice": _assertion("alice-id-token")}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + subject = Subject(tenant_id="", subject_id="alice") + spec = _spec(_id_jag_config()) + + first = await provider.resolve_credentials(subject, spec) + await provider.invalidate_credentials(subject, spec) + second = await provider.resolve_credentials(subject, spec) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(first.ok)["Authorization"] == "Bearer first-bearer" + assert _emitted(second.ok)["Authorization"] == "Bearer second-bearer" + assert len(endpoint.calls) == 4 @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py index a3f46a49ba9..7b82e004f37 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py @@ -8,6 +8,7 @@ rotation re-encrypts stored rows like the sibling per-user credential tables. """ import json +import os import time from unittest.mock import AsyncMock, MagicMock, patch @@ -15,6 +16,8 @@ import jwt as pyjwt import pytest from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + AssertionStoreUnavailable, + DbSSOAssertionStore, assertion_from_sso_login, ema_assertion_retention_enabled, fetch_sso_identity_assertion, @@ -341,3 +344,23 @@ async def test_rotation_skips_unreadable_rows_but_rotates_readable_ones(): await rotate_sso_identity_assertions_master_key(prisma_client=prisma, new_master_key="another-new-salt-key-0000") assert stored["bad"] == "garbage-blob" assert stored["good"] != good_blob_before + + +@pytest.mark.asyncio +async def test_db_store_converts_a_driver_failure_into_assertion_store_unavailable(): + """The live store must not let a raw driver error escape: the resolver distinguishes an outage + from an absent assertion, and only a typed failure lets it do that.""" + prisma = MagicMock() + prisma.db.litellm_ssoidentityassertion.find_unique = AsyncMock(side_effect=RuntimeError("connection refused")) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(AssertionStoreUnavailable): + await DbSSOAssertionStore().fetch("alice") + + +@pytest.mark.asyncio +async def test_db_store_returns_none_for_a_user_with_no_stored_assertion(): + """An absent row stays an absence, not an outage, so a user who never signed in still gets the + 412 that tells them to.""" + with patch.dict(os.environ, {"LITELLM_SALT_KEY": SALT_KEY}): + with patch("litellm.proxy.proxy_server.prisma_client", _make_prisma({})): + assert await DbSSOAssertionStore().fetch("nobody") is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index db5b64ef131..54fd5242d5f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -40,17 +40,21 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _oauth_endpoints_unresolved, _deserialize_json_list, _normalize_mcp_server_cost_info, + _obo_retry_applies, _should_strip_caller_authorization, _without_authorization, ) from litellm.proxy._types import ( LiteLLM_MCPServerTable, + LiteLLM_ObjectPermissionTable, + LitellmUserRoles, MCPApprovalStatus, MCPEnvVar, MCPEnvVarScope, MCPTransport, + UserAPIKeyAuth, ) -from litellm.types.mcp import MCPAuth +from litellm.types.mcp import MCPAuth, MCPAuthType from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer @@ -8232,6 +8236,35 @@ def test_should_strip_caller_authorization_for_token_exchange(): assert _should_strip_caller_authorization(mcp_server=server, raw_headers=None, user_api_key_auth=None) is True +def _retry_gate_server(auth_type: MCPAuthType) -> MCPServer: + return MCPServer( + server_id="retry-gate", + name="retry-gate-server", + url="https://up.example.com", + transport=MCPTransport.http, + auth_type=auth_type, + ) + + +def test_obo_retry_applies_to_id_jag_without_an_inbound_subject_token(): + """ID-JAG can source its subject from the user's stored SSO assertion, so the upstream-401 + invalidate-and-retry path must engage even when the caller presented no token of its own; + otherwise a store-sourced bearer is replayed until its TTL after being rejected.""" + assert _obo_retry_applies(_retry_gate_server(MCPAuth.oauth2_id_jag), None) is True + assert _obo_retry_applies(_retry_gate_server(MCPAuth.oauth2_id_jag), "inbound-id-token") is True + + +def test_obo_retry_still_requires_a_subject_token_for_token_exchange(): + """token_exchange can only mint from an inbound token, so with none there is nothing to re-mint.""" + assert _obo_retry_applies(_retry_gate_server(MCPAuth.oauth2_token_exchange), None) is False + assert _obo_retry_applies(_retry_gate_server(MCPAuth.oauth2_token_exchange), "inbound-token") is True + + +def test_obo_retry_does_not_apply_to_other_auth_modes(): + for auth_type in (MCPAuth.none, MCPAuth.api_key, MCPAuth.oauth2, MCPAuth.true_passthrough): + assert _obo_retry_applies(_retry_gate_server(auth_type), "some-token") is False + + class _UpstreamAuthError(Exception): """Mimics a wrapped upstream 401 the way _extract_upstream_auth_failure detects it.""" @@ -9688,3 +9721,75 @@ class TestOpenAPIRegistryKeyMatchesRegistration: assert result.isError is True assert "not found in registry" in result.content[0].text + + +class TestToolAuthorizationIsNotConditionalOnLogging: + """`call_tool` used to run `pre_call_tool_check` — the only place tool-level + MCP entitlements are enforced — inside `if proxy_logging_obj:`, so a caller + reached with no logging object got no authorization decision at all. Both + production call sites pass the module-level `ProxyLogging` singleton, which + is never None, so this was not a live hole; the invariant being restored is + that an authorization decision cannot be skipped by an absent logger. + """ + + @staticmethod + def _manager_with_scoped_server() -> tuple[MCPServerManager, UserAPIKeyAuth]: + manager = MCPServerManager() + manager.registry["srv-gated"] = MCPServer( + server_id="srv-gated", + name="gated_server", + server_name="gated_server", + alias="gated_server", + url="http://127.0.0.1:1/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + for tool_name in ("read_only_tool", "delete_everything"): + manager.tool_name_to_mcp_server_name_mapping[tool_name] = "gated_server" + user = UserAPIKeyAuth( + api_key="sk-caller", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-gated", + mcp_servers=["srv-gated"], + mcp_tool_permissions={"srv-gated": ["read_only_tool"]}, + ), + ) + return manager, user + + @pytest.mark.asyncio + async def test_unentitled_tool_refused_without_proxy_logging_obj(self): + manager, user = self._manager_with_scoped_server() + upstream = AsyncMock(return_value=CallToolResult(content=[], isError=False)) + + with patch.object(manager, "_call_regular_mcp_tool", new=upstream): + with pytest.raises(HTTPException) as exc: + await manager.call_tool( + server_name="gated_server", + name="delete_everything", + arguments={}, + user_api_key_auth=user, + proxy_logging_obj=None, + ) + + assert exc.value.status_code == 403 + upstream.assert_not_awaited() + + @pytest.mark.asyncio + async def test_entitled_tool_still_dispatches_without_proxy_logging_obj(self): + """The gate must refuse only what the entitlement excludes; an allowed + tool still reaches the upstream when there is no logging object.""" + manager, user = self._manager_with_scoped_server() + upstream = AsyncMock(return_value=CallToolResult(content=[], isError=False)) + + with patch.object(manager, "_call_regular_mcp_tool", new=upstream): + await manager.call_tool( + server_name="gated_server", + name="read_only_tool", + arguments={}, + user_api_key_auth=user, + proxy_logging_obj=None, + ) + + upstream.assert_awaited_once() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index c4b3c7f5f67..63473bf3cf5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -9,7 +9,13 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LitellmUserRoles, + UserAPIKeyAuth, +) +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer @pytest.mark.asyncio @@ -305,3 +311,241 @@ async def test_openapi_local_tool_injects_resolved_oauth_token(): assert captured["resolved"] == {"Authorization": "Bearer stored-user-token"} assert _request_resolved_auth_headers.get() is None + + + +LEGACY_SERVER_ID = "srv-legacy-petstore" +LEGACY_SERVER_NAME = "legacy_petstore" +LEGACY_TOOL = "dump_secrets" + + +@pytest.fixture +def legacy_local_tool(): + """A bare `mcp_tools`-style handler plus a registered server whose tools were + never listed, which is what leaves `tool_name_to_mcp_server_name_mapping` + cold and routes `{server}-{tool}` into `execute_mcp_tool`'s legacy fallback. + + Yields the server and the list the handler appends to, so a test can tell + "refused" from "dispatched" by whether the handler actually ran. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.tool_registry import ( + global_mcp_tool_registry, + ) + + executed: list[dict] = [] + server = MCPServer( + server_id=LEGACY_SERVER_ID, + name=LEGACY_SERVER_NAME, + server_name=LEGACY_SERVER_NAME, + alias=LEGACY_SERVER_NAME, + url="http://127.0.0.1:1/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + global_mcp_tool_registry.register_tool( + name=LEGACY_TOOL, + description="bare tool registered from the mcp_tools config block", + input_schema={"type": "object", "properties": {}}, + handler=lambda **kwargs: executed.append(kwargs) or "legacy local tool ran", + ) + global_mcp_server_manager.registry[LEGACY_SERVER_ID] = server + assert ( + global_mcp_server_manager._get_mcp_server_from_tool_name( + f"{LEGACY_SERVER_NAME}-{LEGACY_TOOL}" + ) + is None + ), "fixture precondition: the prefixed name must resolve to no server" + try: + yield server, executed + finally: + global_mcp_tool_registry.tools.pop(LEGACY_TOOL, None) + global_mcp_server_manager.registry.pop(LEGACY_SERVER_ID, None) + + +def _caller_entitled_to(tools: list[str]) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-caller", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-legacy-fallback", + mcp_servers=[LEGACY_SERVER_ID], + mcp_tool_permissions={LEGACY_SERVER_ID: tools}, + ), + ) + + +@pytest.mark.asyncio +async def test_legacy_local_tool_fallback_refuses_unentitled_caller(legacy_local_tool): + """The legacy fallback dispatched into the local tool registry with no + tool-level authorization at all: no allowed/banned check, no key/team/org + tool permissions, no parameter validation. It must now run the same gate, + so a caller whose entitlement excludes the tool is refused and the handler + never runs. + + Nothing is mocked: the real registries and the real entitlement gate decide. + """ + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + server, executed = legacy_local_tool + user = _caller_entitled_to(["list_pets"]) + + # The gate answers "no" for this caller/tool pair, so a dispatch below would + # be an entitlement bypass rather than a routing quirk. + assert ( + await MCPRequestHandler.is_tool_allowed_for_server( + tool_name=LEGACY_TOOL, + server_id=LEGACY_SERVER_ID, + user_api_key_auth=user, + ) + is False + ) + + with pytest.raises(HTTPException) as exc: + await mcp_module.execute_mcp_tool( + name=f"{LEGACY_SERVER_NAME}-{LEGACY_TOOL}", + arguments={}, + allowed_mcp_servers=[server], + start_time=datetime.now(timezone.utc), + user_api_key_auth=user, + ) + + assert exc.value.status_code == 403 + # Pin the refusal to the ENTITLEMENT gate. The server-level check earlier in + # execute_mcp_tool also raises 403 (with a plain-string detail), and the + # allowed/banned-tools check raises a dict naming the server rather than the + # key/team, so asserting on the status alone would pass for the wrong reason. + detail = exc.value.detail + assert isinstance(detail, dict), detail + assert "not allowed for your key/team" in detail["error"], detail + assert executed == [] + + +@pytest.mark.asyncio +async def test_legacy_local_tool_fallback_still_dispatches_entitled_caller( + legacy_local_tool, +): + """The gate must do per-tool work rather than disabling the fallback: the + same shape of call, from a caller entitled to the tool, still dispatches. + + This is the backwards-compatibility half. Refusing this call would trade an + authorization hole for an outage on a configuration that worked before. + """ + from litellm.proxy._experimental.mcp_server import server as mcp_module + + server, executed = legacy_local_tool + user = _caller_entitled_to([LEGACY_TOOL]) + + result = await mcp_module.execute_mcp_tool( + name=f"{LEGACY_SERVER_NAME}-{LEGACY_TOOL}", + arguments={}, + allowed_mcp_servers=[server], + start_time=datetime.now(timezone.utc), + user_api_key_auth=user, + ) + + assert result.isError is False + assert executed == [{}] + assert "legacy local tool ran" in result.content[0].text + + +@pytest.mark.asyncio +async def test_legacy_local_tool_fallback_fails_closed_on_empty_prefix( + legacy_local_tool, +): + """An empty prefix segment skips the server-level check outright: + `split_server_prefix_from_name` yields an empty `server_name`, and `execute_mcp_tool` + only runs `is_tool_allowed` `if server_name`. The legacy fallback then dispatched for a + caller holding no server grant at all, so this arm of the guard is reachable rather than + defensive. Nothing is patched here; the empty prefix segment is the whole of it. + """ + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + _server, executed = legacy_local_tool + + with pytest.raises(HTTPException) as exc: + await mcp_module.execute_mcp_tool( + name=f"-{LEGACY_TOOL}", + arguments={}, + allowed_mcp_servers=[], + start_time=datetime.now(timezone.utc), + user_api_key_auth=_caller_entitled_to([LEGACY_TOOL]), + ) + + assert exc.value.status_code == 503 + assert executed == [] + + +@pytest.mark.asyncio +async def test_legacy_local_tool_fallback_fails_closed_when_prefix_names_no_server( + legacy_local_tool, +): + """Second arm of the same guard: a non-empty prefix that named a server the caller does + hold, but which is absent from `allowed_mcp_servers` by the time dispatch runs. Patching + the server-level check (which would otherwise refuse first) is what makes the arm + observable, so a later refactor cannot make the branch dispatch with no server to + evaluate a tool ceiling against. + """ + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + _server, executed = legacy_local_tool + other_server = MCPServer( + server_id="srv-unrelated", + name="unrelated_server", + server_name="unrelated_server", + alias="unrelated_server", + url="http://127.0.0.1:1/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed", + return_value=True, + ): + with pytest.raises(HTTPException) as exc: + await mcp_module.execute_mcp_tool( + name=f"{LEGACY_SERVER_NAME}-{LEGACY_TOOL}", + arguments={}, + allowed_mcp_servers=[other_server], + start_time=datetime.now(timezone.utc), + user_api_key_auth=_caller_entitled_to([LEGACY_TOOL]), + ) + + assert exc.value.status_code == 503 + assert executed == [] + + +@pytest.mark.asyncio +async def test_unknown_tool_name_still_reports_not_found(): + """The guard must gate dispatch, not existence. An unprefixed name that no registry + knows cannot dispatch anything, so it has to keep reporting 404 rather than collapsing + into the guard's 503; every typo'd tool name takes this branch. + """ + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + with pytest.raises(HTTPException) as exc: + await mcp_module.execute_mcp_tool( + name="tool_no_registry_knows", + arguments={}, + allowed_mcp_servers=[], + start_time=datetime.now(timezone.utc), + user_api_key_auth=_caller_entitled_to([LEGACY_TOOL]), + ) + + assert exc.value.status_code == 404 + assert "not found" in str(exc.value.detail) diff --git a/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py b/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py index b0d2595e48c..36f656a7adc 100644 --- a/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py +++ b/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py @@ -58,10 +58,14 @@ def admin_viewer_client(monkeypatch): mock_config_table = MagicMock() mock_config_table.find_first = AsyncMock(return_value=None) + # /management/v1/budgets reads through query_raw: count first, then the page. + mock_query_raw = AsyncMock(side_effect=[[{"count": 0}], []]) + mock_prisma.db = types.SimpleNamespace( litellm_budgettable=mock_budget_table, litellm_invitationlink=mock_invitation_table, litellm_config=mock_config_table, + query_raw=mock_query_raw, ) monkeypatch.setattr(ps, "prisma_client", mock_prisma) @@ -106,6 +110,14 @@ def test_budget_list_allows_admin_viewer(admin_viewer_client): assert resp.status_code == 200, resp.text +def test_management_v1_budgets_allows_admin_viewer(admin_viewer_client): + """`/management/v1/budgets` is the paged/sortable budget list; same read tier as + `/budget/list`, and it answers 403 rather than an empty page when it refuses.""" + resp = admin_viewer_client.get("/management/v1/budgets") + _assert_not_role_blocked(resp) + assert resp.status_code == 200, resp.text + + def test_budget_settings_allows_admin_viewer(admin_viewer_client): """`/budget/settings` describes a budget's fields; read-only.""" resp = admin_viewer_client.get("/budget/settings", params={"budget_id": "b1"}) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index a6d4dc63697..06764139eda 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -1960,6 +1960,7 @@ ADMIN_VIEWER_SETTINGS_ROUTES = [ "/config/field/info", # Budgets page "/budget/list", + "/management/v1/budgets", "/budget/settings", # Invitation viewing (admin viewer cannot create/delete; can read) "/invitation/info", diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 53f32fd96fb..65d6e33588f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -21,7 +21,7 @@ from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( ) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import ModelResponse +from litellm.types.utils import CallTypes, ModelResponse @pytest.mark.asyncio @@ -3614,3 +3614,59 @@ class TestBedrockIncrementalFlagInteractions: ) scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] assert scanned == ["q1"], "latest-role selection must exclude the system prompt" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode, call_type, should_scan", + [ + ("during_mcp_call", "call_mcp_tool", True), + ("during_call", "completion", True), + ("during_mcp_call", "completion", False), + ("during_call", "call_mcp_tool", False), + ], +) +async def test_moderation_hook_honors_the_mcp_event_type(mode, call_type, should_scan): + """A guardrail configured mode: during_mcp_call must actually scan MCP tool calls. + + ProxyLogging.during_call_hook remaps call_mcp_tool to during_mcp_call before + dispatching, so re-checking during_call here would reject the very requests the + guardrail was configured for and let the tool call through unscanned. + """ + guardrail = BedrockGuardrail( + guardrail_name="bedrock-mcp", + guardrailIdentifier="gid", + guardrailVersion="1", + event_hook=mode, + default_on=True, + ) + data = { + "messages": [{"role": "user", "content": "scan me"}], + "mcp_tool_name": "search", + "mcp_arguments": {"query": "scan me"}, + } + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = MagicMock( + action="NONE", output=[], outputs=[], assessments=[] + ) + await guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u"), + call_type=call_type, + ) + + assert (mock_api.call_count == 1) is should_scan, ( + f"mode={mode} call_type={call_type}: expected scan={should_scan}, " + f"bedrock api called {mock_api.call_count} times" + ) + if should_scan: + expected_event = ( + GuardrailEventHooks.during_mcp_call + if call_type == CallTypes.call_mcp_tool.value + else GuardrailEventHooks.during_call + ) + assert mock_api.call_args.kwargs["logging_event_type"] == expected_event, ( + "the scan must be logged under the event it actually ran for, so guardrail logs, " + "OTel spans, and Langfuse metadata do not misclassify MCP enforcement as an LLM call" + ) diff --git a/tests/test_litellm/proxy/hooks/test_user_management_event_hooks.py b/tests/test_litellm/proxy/hooks/test_user_management_event_hooks.py new file mode 100644 index 00000000000..102430ab985 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_user_management_event_hooks.py @@ -0,0 +1,99 @@ +import asyncio +import json +import sys +from types import SimpleNamespace +from typing import Any, Dict, List, Optional +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm.proxy._types import NewUserRequest, NewUserResponse, UserAPIKeyAuth +from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks + + +class FakeUserTable: + def __init__(self, rows: List[Dict[str, Any]]): + self._rows = rows + + async def find_unique(self, where: Dict[str, Any]) -> Optional[Dict[str, Any]]: + return next( + (row for row in self._rows if all(row.get(key) == value for key, value in where.items())), + None, + ) + + +class FakePrismaClient: + def __init__(self, rows: List[Dict[str, Any]]): + self.db = SimpleNamespace(litellm_usertable=FakeUserTable(rows)) + + +async def _run_created_hook(prisma_client: FakePrismaClient, audit_log: AsyncMock) -> None: + proxy_server = SimpleNamespace( + prisma_client=prisma_client, + litellm_proxy_admin_name="admin-user", + ) + with ( + patch.dict(sys.modules, {"litellm.proxy.proxy_server": proxy_server}), + patch.object(litellm, "store_audit_logs", True), + patch( + "litellm.proxy.hooks.user_management_event_hooks.create_audit_log_for_update", + audit_log, + ), + patch.object( + UserManagementEventHooks, + "async_send_user_invitation_email", + AsyncMock(), + ), + ): + await UserManagementEventHooks.async_user_created_hook( + data=NewUserRequest(user_email="new@example.com", send_invite_email=False), + response=NewUserResponse( + user_id="user-1", + user_email="new@example.com", + key="sk-test", + ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-user", api_key="sk-admin"), + ) + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_created_hook_audit_logs_the_row_read_back_from_the_database(): + """The audit entry must describe the persisted user row, not the /user/new response.""" + prisma_client = FakePrismaClient( + [ + { + "user_id": "user-1", + "user_email": "new@example.com", + "user_role": "proxy_admin", + "models": ["gpt-4"], + "teams": ["team-a"], + "metadata": '{"source": "api"}', + } + ] + ) + audit_log = AsyncMock() + + await _run_created_hook(prisma_client, audit_log) + + audit_log.assert_awaited_once() + request_data = audit_log.await_args.kwargs["request_data"] + assert request_data.object_id == "user-1" + assert request_data.action == "created" + + updated_values = json.loads(request_data.updated_values) + assert updated_values["user_role"] == "proxy_admin" + assert updated_values["models"] == ["gpt-4"] + assert updated_values["teams"] == ["team-a"] + assert updated_values["metadata"] == {"source": "api"} + + +@pytest.mark.asyncio +async def test_created_hook_skips_the_audit_log_when_no_user_row_exists(): + """A user id that resolves to nothing must not produce an audit entry.""" + audit_log = AsyncMock() + + await _run_created_hook(FakePrismaClient([]), audit_log) + + audit_log.assert_not_awaited() diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py new file mode 100644 index 00000000000..40473f1a25a --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py @@ -0,0 +1,525 @@ +from dataclasses import replace +from datetime import datetime, timezone +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI, Request +from fastapi.exceptions import RequestValidationError +from fastapi.testclient import TestClient + +from litellm.proxy._types import LiteLLMRoutes, LitellmUserRoles +from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth +from litellm.proxy.management_endpoints.management_v1 import router +from litellm.proxy.management_endpoints.management_v1.budgets import ( + BUDGETS_LIST_SPEC, + BudgetListItem, +) +from litellm.proxy.management_endpoints.management_v1.common import ( + MANAGEMENT_V1_PREFIX, + PROBLEM_TYPE_BASE, + ManagementProblem, + problem_response, +) +from litellm.proxy.management_endpoints.management_v1.list_framework import ( + Compare, + ScopeWhere, + build_query_plan, +) +from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail + +app = FastAPI() + + +@app.exception_handler(ManagementProblem) +async def management_problem_exception_handler(request: Request, exc: ManagementProblem): + return problem_response(exc.problem) + + +@app.exception_handler(RequestValidationError) +async def validation_exception_handler(request: Request, exc: RequestValidationError): + return problem_response( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}invalid-query-parameter", + title="Invalid query parameter", + status=400, + detail="The request query parameters are invalid.", + ) + ) + + +app.include_router(router) +client = TestClient(app) + +BUDGETS_PATH = f"{MANAGEMENT_V1_PREFIX}/budgets" +SORTABLE = ["budget_id", "created_at", "max_budget", "rpm_limit", "tpm_limit"] + + +def _row(budget_id: str, **overrides: Any) -> dict[str, Any]: + return { + "budget_id": budget_id, + "max_budget": 10.0, + "soft_budget": None, + "tpm_limit": None, + "rpm_limit": None, + "budget_duration": "30d", + "budget_reset_at": None, + "created_at": "2026-07-20T12:00:00+00:00", + "updated_at": "2026-07-21T12:00:00+00:00", + **overrides, + } + + +@pytest.fixture +def query_raw(monkeypatch): + """Mocks the one call the executor makes. `count` reads the first result, `find_many` the second.""" + mock = AsyncMock(side_effect=[[{"count": 0}], []]) + prisma_client = MagicMock() + prisma_client.db.query_raw = mock + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + return mock + + +def _serve(query_raw, rows: list[dict[str, Any]], total: int | None = None) -> None: + query_raw.side_effect = [[{"count": len(rows) if total is None else total}], rows] + + +@pytest.fixture +def as_proxy_admin(): + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + yield + app.dependency_overrides.clear() + + +def _as_role(role: LitellmUserRoles): + original = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="u", user_role=role) + return original + + +def _get(query: str = ""): + suffix = f"?{query}" if query else "" + return client.get(f"{BUDGETS_PATH}{suffix}", headers={"Authorization": "Bearer k"}) + + +def _select_call(query_raw): + """The find_many call: (sql, *params). The count call comes first.""" + return query_raw.call_args_list[1].args + + +def test_returns_flat_rows_in_the_control_plane_envelope(query_raw, as_proxy_admin): + """`{data, meta, links}` with flat rows; no JSON:API `{type, id, attributes}` wrapper.""" + _serve(query_raw, [_row("b-1")]) + + response = _get() + + assert response.status_code == 200 + body = response.json() + assert set(body) == {"data", "meta", "links"} + assert body["data"][0]["budget_id"] == "b-1" + assert "attributes" not in body["data"][0] + + +def test_serves_the_columns_the_budgets_page_renders(query_raw, as_proxy_admin): + _serve(query_raw, [_row("b-1", soft_budget=5.0, budget_reset_at="2026-08-01T00:00:00+00:00")]) + + row = _get().json()["data"][0] + + assert set(row) == { + "budget_id", + "max_budget", + "soft_budget", + "tpm_limit", + "rpm_limit", + "budget_duration", + "budget_reset_at", + "created_at", + "updated_at", + } + assert row["soft_budget"] == 5.0 + assert row["budget_reset_at"].startswith("2026-08-01T00:00:00") + + +def test_selects_only_the_columns_it_serves(query_raw, as_proxy_admin): + """A `SELECT *` would ship created_by/updated_by and model_max_budget to the browser.""" + _serve(query_raw, []) + + _get() + + sql = _select_call(query_raw)[0] + assert "SELECT *" not in sql + assert '"budget_id"' in sql and '"max_budget"' in sql + assert "created_by" not in sql and "model_max_budget" not in sql + + +def test_defaults_to_newest_first_with_budget_id_breaking_ties(query_raw, as_proxy_admin): + """Two budgets created in the same transaction share a created_at; without the + tiebreaker their relative order is undefined and pages can repeat or drop rows.""" + _serve(query_raw, []) + + _get() + + assert 'ORDER BY "created_at" DESC NULLS LAST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0] + + +def test_appends_the_tiebreaker_to_an_explicit_sort(query_raw, as_proxy_admin): + _serve(query_raw, []) + + _get("sort=-max_budget") + + assert 'ORDER BY "max_budget" DESC NULLS LAST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0] + + +def test_refuses_to_sort_on_budget_duration(query_raw, as_proxy_admin): + """The column holds "7d"/"30d", so a lexicographic ORDER BY would put "30d" + before "7d" and silently mis-order the page.""" + _serve(query_raw, []) + + response = _get("sort=budget_duration") + + assert response.status_code == 400 + assert response.headers["content-type"].startswith("application/problem+json") + body = response.json() + assert "budget_duration" in body["detail"] + assert body["allowed"] == SORTABLE + query_raw.assert_not_called() + + +def test_the_advertised_sort_fields_are_the_ones_that_work(query_raw, as_proxy_admin): + """Guards the rejection above against drifting from what the spec actually accepts.""" + for field in SORTABLE: + _serve(query_raw, []) + assert _get(f"sort={field}").status_code == 200, field + assert sorted(BUDGETS_LIST_SPEC.sortable) == SORTABLE + + +def test_rejects_an_unknown_query_parameter(query_raw, as_proxy_admin): + """A silently ignored filter over-returns budgets, which is worse than a rejected request.""" + _serve(query_raw, []) + + response = _get("filtre[max_budget][gte]=5") + + assert response.status_code == 400 + assert response.headers["content-type"].startswith("application/problem+json") + body = response.json() + assert "filtre[max_budget][gte]" in body["detail"] + assert "filter[max_budget][gte]" in body["allowed"] + query_raw.assert_not_called() + + +def test_rejects_an_operator_the_filter_does_not_declare(query_raw, as_proxy_admin): + """`max_budget` takes ranges, not `in`; accepting an undeclared operator is how a + filter starts meaning something the query planner never checked.""" + for query in ("filter[max_budget][in]=5,10", "filter[created_at][is_null]=true"): + _serve(query_raw, []) + assert _get(query).status_code == 400, query + + +def test_omitted_page_size_serves_fifty(query_raw, as_proxy_admin): + _serve(query_raw, []) + + body = _get().json() + + assert body["meta"]["page_size"] == 50 + assert _select_call(query_raw)[-2] == 50 + + +def test_clamps_an_oversized_page_size_to_a_hundred(query_raw, as_proxy_admin): + """Unclamped, one request can ask the proxy to serialize the whole budget table.""" + _serve(query_raw, []) + + body = _get("page_size=500").json() + + assert body["meta"]["page_size"] == 100 + assert _select_call(query_raw)[-2] == 100 + + +def test_offsets_by_page(query_raw, as_proxy_admin): + _serve(query_raw, []) + + _get("page=3&page_size=25") + + assert _select_call(query_raw)[-2:] == (25, 50) + + +@pytest.mark.parametrize( + "role", + [ + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + LitellmUserRoles.TEAM, + ], +) +def test_refuses_a_caller_without_admin_view(query_raw, role): + """Budgets are proxy-wide, so a caller who cannot read all of them must be told + so. Answering 200 with an empty list would read as "there are no budgets".""" + _serve(query_raw, [_row("b-1")]) + original = _as_role(role) + try: + response = _get() + finally: + app.dependency_overrides = original + + assert response.status_code == 403 + assert response.headers["content-type"].startswith("application/problem+json") + assert response.json()["status"] == 403 + query_raw.assert_not_called() + + +@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]) +def test_admins_and_admin_viewers_may_read_every_budget(query_raw, role): + _serve(query_raw, [_row("b-1")]) + original = _as_role(role) + try: + response = _get() + finally: + app.dependency_overrides = original + + assert response.status_code == 200 + assert [row["budget_id"] for row in response.json()["data"]] == ["b-1"] + assert "WHERE" not in _select_call(query_raw)[0] + + +def test_a_denied_caller_stays_denied_whatever_they_filter_on(query_raw): + """The scope decision reads the caller, never the query string.""" + _serve(query_raw, [_row("b-1")]) + original = _as_role(LitellmUserRoles.INTERNAL_USER) + try: + response = _get("filter[max_budget][gte]=0&q=b-") + finally: + app.dependency_overrides = original + + assert response.status_code == 403 + + +def test_a_filter_sits_behind_the_scope_predicate_instead_of_replacing_it(): + """Same spec, planned for a row-scoped caller: the scope clause has to survive, and + lead, whatever the caller filtered on. Replacing it would let a filter widen a read.""" + scoped = replace( + BUDGETS_LIST_SPEC, + scope=lambda _caller: ScopeWhere(where=(Compare(field="budget_id", op="eq", value="b-1"),)), + ) + + plan = build_query_plan( + spec=scoped, + params={"filter[max_budget][gte]": "5"}, + caller=UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert plan.where[0] == Compare(field="budget_id", op="eq", value="b-1") + assert Compare(field="max_budget", op="gte", value=5.0) in plan.where + + +def test_q_matches_budget_id_case_insensitively(query_raw, as_proxy_admin): + """budget_id is the only text identity on the row; matching anything else would + return budgets whose ids do not contain what the user typed.""" + _serve(query_raw, []) + + _get("q=Prod") + + sql, *params = _select_call(query_raw) + assert '"budget_id" ILIKE $1' in sql + assert params[0] == "%Prod%" + + +def test_q_escapes_like_metacharacters(query_raw, as_proxy_admin): + """Budget ids routinely contain '_'; unescaped it is a single-character wildcard.""" + _serve(query_raw, []) + + _get("q=team_a%25") + + assert _select_call(query_raw)[1] == r"%team\_a\%%" + + +def test_q_does_not_search_any_other_column(query_raw, as_proxy_admin): + _serve(query_raw, []) + + _get("q=30d") + + assert "budget_duration" not in _select_call(query_raw)[0].split("WHERE")[1] + assert BUDGETS_LIST_SPEC.searchable == frozenset({"budget_id"}) + + +def test_is_null_selects_the_unlimited_budgets(query_raw, as_proxy_admin): + """"Unlimited" is max_budget IS NULL; `max_budget = 0` would be a hard zero cap.""" + _serve(query_raw, [_row("b-unlimited", max_budget=None)]) + + body = _get("filter[max_budget][is_null]=true").json() + + assert '"max_budget" IS NULL' in _select_call(query_raw)[0] + assert body["data"][0]["max_budget"] is None + + +def test_is_null_false_selects_the_capped_budgets(query_raw, as_proxy_admin): + _serve(query_raw, []) + + _get("filter[max_budget][is_null]=false") + + assert '"max_budget" IS NOT NULL' in _select_call(query_raw)[0] + + +def test_in_filter_binds_each_requested_duration(query_raw, as_proxy_admin): + _serve(query_raw, []) + + _get("filter[budget_duration][in]=7d,30d") + + sql, *params = _select_call(query_raw) + assert '"budget_duration" IN ($1, $2)' in sql + assert params[:2] == ["7d", "30d"] + + +def test_created_at_range_is_bound_as_a_timestamp(query_raw, as_proxy_admin): + """The bind crosses into the query engine as JSON, so an uncast placeholder reaches + Postgres as text and `timestamp >= text` is a hard error, not a wrong answer.""" + _serve(query_raw, []) + + _get("filter[created_at][gte]=2026-07-01T00:00:00Z") + + sql, *params = _select_call(query_raw) + assert "\"created_at\" >= $1::timestamptz AT TIME ZONE 'UTC'" in sql + assert params[0] == datetime(2026, 7, 1, tzinfo=timezone.utc) + + +def test_numbers_placeholders_continuously_across_predicates(query_raw, as_proxy_admin): + """Each predicate is numbered after the binds the ones before it consumed. Restart + the count and `$1` gets read as the duration while the search string goes unbound.""" + _serve(query_raw, []) + + _get("filter[budget_duration][in]=7d,30d&filter[max_budget][gte]=5&q=prod") + + sql, *params = _select_call(query_raw) + assert '"budget_duration" IN ($1, $2)' in sql + assert '"max_budget" >= $3' in sql + assert '"budget_id" ILIKE $4' in sql + assert params[:4] == ["7d", "30d", 5.0, "%prod%"] + + +def test_a_non_datetime_bind_is_not_cast(query_raw, as_proxy_admin): + """Guards the cast above from being applied to every placeholder.""" + _serve(query_raw, []) + + _get("filter[max_budget][gte]=5") + + sql = _select_call(query_raw)[0] + assert '"max_budget" >= $1' in sql + assert "timestamptz" not in sql + + +def test_an_offsetless_created_at_bound_is_read_as_utc(query_raw, as_proxy_admin): + """The dashboard sends 'YYYY-MM-DDTHH:MM:SS' with no offset. Left naive, Postgres + would compare it in the session timezone and shift the window off the rows shown.""" + _serve(query_raw, []) + + _get("filter[created_at][gte]=2026-07-01T00:00:00") + + bound = _select_call(query_raw)[1] + assert bound == datetime(2026, 7, 1, tzinfo=timezone.utc) + assert bound.tzinfo is not None + + +def test_rejects_a_filter_value_that_is_not_of_the_declared_type(query_raw, as_proxy_admin): + for query in ("filter[max_budget][gte]=lots", "filter[created_at][gte]=yesterday"): + _serve(query_raw, []) + assert _get(query).status_code == 400, query + + +def test_reports_the_total_and_links_every_page_on_a_middle_page(query_raw, as_proxy_admin): + """The Budgets page renders a page count, so the total has to be the match total, + not the length of the page it just received.""" + _serve(query_raw, [_row("b-3"), _row("b-4")], total=7) + + body = _get("page=2&page_size=2").json() + + assert body["meta"] == {"page": 2, "page_size": 2, "total_count": 7, "total_pages": 4} + links = body["links"] + assert "page=1" in links["first"] and "page_size=2" in links["first"] + assert "page=1" in links["prev"] + assert "page=3" in links["next"] + assert "page=4" in links["last"] + assert "page=2" in links["self"] + + +def test_the_last_page_has_no_next(query_raw, as_proxy_admin): + _serve(query_raw, [_row("b-5")], total=5) + + links = _get("page=3&page_size=2").json()["links"] + + assert links["next"] is None + assert "page=2" in links["prev"] + + +def test_the_first_page_has_no_prev(query_raw, as_proxy_admin): + _serve(query_raw, [_row("b-1")], total=5) + + links = _get("page_size=2").json()["links"] + + assert links["prev"] is None + assert "page=2" in links["next"] + + +def test_an_empty_table_still_links_a_first_and_last_page(query_raw, as_proxy_admin): + _serve(query_raw, [], total=0) + + body = _get().json() + + assert body["meta"]["total_count"] == 0 + assert body["meta"]["total_pages"] == 0 + assert "page=1" in body["links"]["first"] and "page=1" in body["links"]["last"] + + +def test_counts_over_the_same_predicate_it_pages(query_raw, as_proxy_admin): + """A total counted without the caller's filter would page through rows the + filter excluded.""" + _serve(query_raw, [], total=0) + + _get("filter[budget_duration][in]=30d") + + count_sql, *count_params = query_raw.call_args_list[0].args + assert '"budget_duration" IN ($1)' in count_sql + assert count_params == ["30d"] + assert "COUNT(*)" in count_sql + + +def test_bigint_limits_serialize_as_json_numbers(query_raw, as_proxy_admin): + """tpm_limit/rpm_limit are BigInt? in Prisma; the query engine can hand them back + as decimal strings, and a quoted "60000" breaks arithmetic in the dashboard.""" + _serve(query_raw, [_row("b-1", tpm_limit="60000", rpm_limit=1200)]) + + response = _get() + row = response.json()["data"][0] + + assert row["tpm_limit"] == 60000 + assert row["rpm_limit"] == 1200 + assert isinstance(row["tpm_limit"], int) and not isinstance(row["tpm_limit"], bool) + assert '"tpm_limit": "60000"' not in response.text + + +def test_reports_a_missing_database_as_a_problem_document(monkeypatch, as_proxy_admin): + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + + response = _get() + + assert response.status_code == 503 + assert response.headers["content-type"].startswith("application/problem+json") + + +def test_the_spec_serves_what_the_row_model_declares(): + """SELECTED_COLUMNS is built off the model, so a field added to one cannot go + missing from the other and produce a row the validator rejects.""" + from litellm.proxy.management_endpoints.management_v1.budgets import SELECTED_COLUMNS + + assert SELECTED_COLUMNS == ", ".join(f'"{name}"' for name in BudgetListItem.model_fields) + assert BUDGETS_LIST_SPEC.tiebreaker in BudgetListItem.model_fields + assert BUDGETS_LIST_SPEC.sortable <= frozenset(BudgetListItem.model_fields) + assert frozenset(BUDGETS_LIST_SPEC.filters) <= frozenset(BudgetListItem.model_fields) + + +def test_is_reachable_by_the_roles_that_can_open_the_budgets_page(): + """Route-level auth gate, which the dependency_overrides above bypass. The handler's + admin-view check is dead code if RouteChecks rejects the role first.""" + assert BUDGETS_PATH in LiteLLMRoutes.admin_viewer_routes.value + assert ("/budget/list" in LiteLLMRoutes.admin_viewer_routes.value) == ( + BUDGETS_PATH in LiteLLMRoutes.admin_viewer_routes.value + ) diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_list_framework.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_list_framework.py new file mode 100644 index 00000000000..35bd5517361 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_list_framework.py @@ -0,0 +1,871 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from datetime import datetime, timezone + +import pytest +from fastapi import Request +from pydantic import BaseModel + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.management_endpoints.management_v1.common import ( + MANAGEMENT_V1_PREFIX, + PROBLEM_TYPE_BASE, + ManagementProblem, + build_page_links, +) +from litellm.proxy.management_endpoints.management_v1.list_framework import ( + AnyOf, + Compare, + FilterSpec, + IsNull, + ListSpec, + QueryPlan, + ScopeAll, + ScopeDenied, + ScopeWhere, + SortKey, + Within, + build_query_plan, + handle_list, + order_by_sql, + where_sql, +) +from litellm.types.proxy.management_endpoints.management_v1 import ( + PageLinks, + PageMeta, + ProblemDetail, +) + +BUDGETS_PATH = f"{MANAGEMENT_V1_PREFIX}/budgets" +CALLER = UserAPIKeyAuth(user_id="caller-1") + + +@dataclass(frozen=True, slots=True) +class BudgetRow: + budget_id: str + max_budget: float | None + created_by: str + + +class BudgetOut(BaseModel): + budget_id: str + max_budget: float | None + + +def _serialize(row: BudgetRow) -> BudgetOut: + return BudgetOut(budget_id=row.budget_id, max_budget=row.max_budget) + + +def _spec( + scope=lambda caller: ScopeAll(), + searchable=frozenset({"budget_id", "created_by"}), + sortable=frozenset({"max_budget", "created_at", "budget_id"}), +) -> ListSpec[BudgetRow, BudgetOut]: + return ListSpec( + resource="budgets", + sortable=sortable, + searchable=searchable, + filters={ + "max_budget": FilterSpec(type=float, ops=frozenset({"eq", "gte", "lte", "is_null"})), + "created_at": FilterSpec(type=datetime, ops=frozenset({"gte", "lte"})), + "created_by": FilterSpec(type=str, ops=frozenset({"eq", "in", "contains"})), + "tpm_limit": FilterSpec(type=int, ops=frozenset({"eq"})), + }, + default_sort=(SortKey(field="created_at", descending=True),), + default_page_size=25, + max_page_size=100, + scope=scope, + serialize=_serialize, + tiebreaker="budget_id", + ) + + +def _spec_with(**overrides) -> ListSpec[BudgetRow, BudgetOut]: + """`replace` re-runs `__init__`, so the spec's own validation applies to the override.""" + return replace(_spec(), **overrides) + + +class RecordingExecutor: + """In-memory stand-in for the Prisma-backed executor PR 2 supplies.""" + + def __init__(self, rows: tuple[BudgetRow, ...], total_count: int | None = None) -> None: + self.rows = rows + self.total_count = len(rows) if total_count is None else total_count + self.plan: QueryPlan | None = None + self.count_where: tuple[object, ...] | None = None + + async def count(self, where: tuple[object, ...]) -> int: + self.count_where = where + return self.total_count + + async def find_many(self, plan: QueryPlan) -> Sequence[BudgetRow]: + self.plan = plan + return self.rows[plan.skip : plan.skip + plan.take] + + +def _request(query: str = "") -> Request: + return Request( + { + "type": "http", + "method": "GET", + "scheme": "http", + "root_path": "", + "path": BUDGETS_PATH, + "query_string": query.encode(), + "headers": [(b"host", b"testserver")], + } + ) + + +def _plan(query_params: Mapping[str, str], spec: ListSpec[BudgetRow, BudgetOut] | None = None) -> QueryPlan: + result = build_query_plan(spec=spec or _spec(), params=query_params, caller=CALLER) + assert isinstance(result, QueryPlan), result + return result + + +def _problem(query_params: Mapping[str, str], spec: ListSpec[BudgetRow, BudgetOut] | None = None) -> ProblemDetail: + result = build_query_plan(spec=spec or _spec(), params=query_params, caller=CALLER) + assert isinstance(result, ProblemDetail), result + return result + + +def _conjuncts(plan: QueryPlan) -> tuple[object, ...]: + return plan.where + + +# ---------------------------------------------------------------- invariant 1 + + +def test_appends_the_tiebreaker_to_the_default_sort(): + """Without a unique final key, ordering by an all-null column lets Postgres hand the + same row back on two different pages.""" + assert _plan({}).order == (SortKey(field="created_at", descending=True), SortKey(field="budget_id", descending=False)) + + +def test_appends_the_tiebreaker_to_an_explicit_multi_key_sort(): + order = _plan({"sort": "-max_budget,created_at"}).order + + assert len(order) == 3 + assert order[-1] == SortKey(field="budget_id", descending=False) + + +def test_appends_the_tiebreaker_even_when_the_caller_already_sorts_by_it(): + """Deduplicating it away is the tempting simplification, and it is the one that + reintroduces a non-total order the moment the leading key stops being unique.""" + assert _plan({"sort": "-budget_id"}).order == ( + SortKey(field="budget_id", descending=True), + SortKey(field="budget_id", descending=False), + ) + + +# ---------------------------------------------------------------- invariant 2 + + +def test_orders_nulls_last_in_both_directions(): + """Postgres sorts nulls last ascending but first descending, so flipping the sort + direction on max_budget would otherwise float every "Unlimited" row to the top.""" + sql = order_by_sql((SortKey(field="max_budget", descending=True), SortKey(field="budget_id", descending=False))) + + assert sql == '"max_budget" DESC NULLS LAST, "budget_id" ASC NULLS LAST' + + +def test_order_sql_covers_every_key_in_the_plan(): + sql = order_by_sql(_plan({"sort": "-max_budget,created_at"}).order) + + assert sql.count("NULLS LAST") == 3 + assert sql == '"max_budget" DESC NULLS LAST, "created_at" ASC NULLS LAST, "budget_id" ASC NULLS LAST' + + +# ------------------------------------------------------------ where rendering + + +def test_every_caller_value_is_bound_not_interpolated(): + """The one property that keeps a filter value from reaching the SQL text. A value that + looks like SQL has to come back as a parameter, never as part of the statement.""" + sql, params = where_sql((Compare(field="budget_id", op="eq", value="'; DROP TABLE x --"),)) + + assert sql == '"budget_id" = $1' + assert params == ("'; DROP TABLE x --",) + assert "DROP" not in sql + + +def test_placeholders_are_numbered_across_the_whole_plan(): + """A predicate that binds several values has to advance the counter by that many, or + every later predicate reads the wrong parameter.""" + sql, params = where_sql( + ( + Compare(field="created_by", op="eq", value="alice"), + Within(field="budget_id", values=("a", "b", "c")), + Compare(field="max_budget", op="gte", value=5.0), + ) + ) + + assert sql == '"created_by" = $1 AND "budget_id" IN ($2, $3, $4) AND "max_budget" >= $5' + assert params == ("alice", "a", "b", "c", 5.0) + + +def test_placeholder_numbering_can_start_past_earlier_parameters(): + sql, params = where_sql((Compare(field="created_by", op="eq", value="alice"),), first_index=4) + + assert sql == '"created_by" = $4' + assert params == ("alice",) + + +def test_is_null_binds_no_parameter_and_does_not_consume_a_placeholder(): + sql, params = where_sql( + (IsNull(field="max_budget", negated=False), Compare(field="created_by", op="eq", value="alice")) + ) + + assert sql == '"max_budget" IS NULL AND "created_by" = $1' + assert params == ("alice",) + + +def test_is_null_negated_renders_is_not_null(): + assert where_sql((IsNull(field="max_budget", negated=True),))[0] == '"max_budget" IS NOT NULL' + + +def test_a_search_renders_as_a_parenthesised_or(): + """Without the parentheses the OR would bind looser than the surrounding ANDs and the + scope predicate would stop constraining the search branch.""" + sql, params = where_sql( + ( + Compare(field="created_by", op="eq", value="alice"), + AnyOf( + clauses=( + Compare(field="budget_id", op="contains", value="prod"), + Compare(field="created_by", op="contains", value="prod"), + ) + ), + ) + ) + + assert sql == ( + '"created_by" = $1 AND (' + "\"budget_id\" ILIKE $2 ESCAPE '\\'" + " OR " + "\"created_by\" ILIKE $3 ESCAPE '\\'" + ")" + ) + assert params == ("alice", "%prod%", "%prod%") + + +def test_contains_escapes_like_metacharacters(): + """Budget ids routinely contain '_', which is a single-character wildcard unescaped.""" + _, params = where_sql((Compare(field="budget_id", op="contains", value="device_id%"),)) + + assert params == (r"%device\_id\%%",) + + +@pytest.mark.parametrize( + ("op", "operator"), + [("eq", "="), ("not", "<>"), ("gte", ">="), ("lte", "<="), ("gt", ">"), ("lt", "<")], +) +def test_each_comparison_operator_renders_its_sql_spelling(op, operator): + assert where_sql((Compare(field="max_budget", op=op, value=1),))[0] == f'"max_budget" {operator} $1' + + +def test_an_empty_plan_renders_no_where_body(): + assert where_sql(()) == ("", ()) + + +def test_a_planned_filter_renders_end_to_end(): + """Ties the parser to the renderer: what build_query_plan produces is what executes.""" + sql, params = where_sql(_plan({"filter[max_budget][is_null]": "true", "q": "prod"}).where) + + assert sql == ( + '"max_budget" IS NULL AND (' + "\"budget_id\" ILIKE $1 ESCAPE '\\'" + " OR " + "\"created_by\" ILIKE $2 ESCAPE '\\'" + ")" + ) + assert params == ("%prod%", "%prod%") + + +# ---------------------------------------------------------------- invariant 3 + + +def test_the_scope_predicate_is_the_first_conjunct(): + spec = _spec(scope=lambda caller: ScopeWhere(where=(Compare(field="created_by", op="eq", value=caller.user_id),))) + + conjuncts = _conjuncts(_plan({"filter[max_budget][gte]": "5"}, spec=spec)) + + assert conjuncts[0] == Compare(field="created_by", op="eq", value="caller-1") + + +def test_a_caller_filter_cannot_replace_the_scope_predicate(): + """The failure this guards is a `{**scope, **filters}` merge: a caller filtering on + the scoped column would silently overwrite the scope and read another user's rows.""" + spec = _spec(scope=lambda caller: ScopeWhere(where=(Compare(field="created_by", op="eq", value=caller.user_id),))) + + conjuncts = _conjuncts(_plan({"filter[created_by][eq]": "someone-else"}, spec=spec)) + + assert conjuncts[0] == Compare(field="created_by", op="eq", value="caller-1") + assert Compare(field="created_by", op="eq", value="someone-else") in conjuncts + assert len(conjuncts) == 2 + + +def test_the_scope_predicate_survives_a_search(): + spec = _spec(scope=lambda caller: ScopeWhere(where=(Compare(field="created_by", op="eq", value=caller.user_id),))) + + conjuncts = _conjuncts(_plan({"q": "prod"}, spec=spec)) + + assert conjuncts[0] == Compare(field="created_by", op="eq", value="caller-1") + assert any(isinstance(conjunct, AnyOf) for conjunct in conjuncts) + + +def test_an_unscoped_caller_gets_no_scope_conjunct(): + assert _plan({"filter[max_budget][gte]": "5"}).where == (Compare(field="max_budget", op="gte", value=5.0),) + + +def test_an_unfiltered_unscoped_list_has_an_empty_where(): + assert _plan({}).where == () + + +# ---------------------------------------------------------------- invariant 4 + + +def test_a_denied_scope_is_a_403_problem(): + spec = _spec(scope=lambda caller: ScopeDenied(reason="Only a proxy admin can list budgets.")) + + problem = _problem({}, spec=spec) + + assert problem.status == 403 + assert problem.type == f"{PROBLEM_TYPE_BASE}forbidden" + assert problem.detail == "Only a proxy admin can list budgets." + + +@pytest.mark.asyncio +async def test_a_denied_scope_never_reaches_the_database(): + """A 200 with an empty list would tell the caller the resource is empty rather than + that they cannot read it, and would still pay for the query.""" + spec = _spec(scope=lambda caller: ScopeDenied(reason="nope")) + executor = RecordingExecutor(rows=(BudgetRow(budget_id="b1", max_budget=None, created_by="x"),)) + + with pytest.raises(ManagementProblem) as raised: + await handle_list(spec=spec, executor=executor, request=_request(), caller=CALLER) + + assert raised.value.problem.status == 403 + assert executor.plan is None + assert executor.count_where is None + + +# ---------------------------------------------------------------- invariant 5 + + +def test_page_size_falls_back_to_the_spec_default(): + assert _plan({}).take == 25 + + +def test_page_size_is_clamped_to_the_spec_maximum(): + """Clamped rather than rejected: an over-large page is a UI bug, not a caller error, + but serving it would let one request read the whole table.""" + assert _plan({"page_size": "100000"}).take == 100 + + +def test_page_offsets_by_page_size(): + plan = _plan({"page": "3", "page_size": "10"}) + + assert (plan.skip, plan.take) == (20, 10) + + +@pytest.mark.parametrize("page", ["0", "-1"], ids=["zero", "negative"]) +def test_page_below_one_is_rejected(page): + problem = _problem({"page": page}) + + assert problem.status == 400 + assert problem.type == f"{PROBLEM_TYPE_BASE}invalid-query-parameter" + + +@pytest.mark.parametrize( + "params", + [{"page": "one"}, {"page_size": "many"}, {"page_size": "0"}], + ids=["page-not-an-int", "page-size-not-an-int", "page-size-zero"], +) +def test_unusable_paging_values_are_rejected(params): + assert _problem(params).status == 400 + + +# ---------------------------------------------------------------- invariant 6 + + +def test_an_unknown_query_parameter_is_rejected_with_the_allowed_set(): + problem = _problem({"page_sizee": "10"}) + + assert problem.status == 400 + assert problem.type == f"{PROBLEM_TYPE_BASE}unknown-query-parameter" + assert "page_sizee" in problem.detail + assert problem.allowed is not None + assert "page_size" in problem.allowed + assert "filter[max_budget][gte]" in problem.allowed + + +def test_the_allowed_set_enumerates_only_operators_the_field_declares(): + problem = _problem({"nope": "1"}) + + assert problem.allowed is not None + assert "filter[max_budget][is_null]" in problem.allowed + assert "filter[tpm_limit][gte]" not in problem.allowed + assert "filter[created_by][in]" in problem.allowed + + +def test_a_filter_on_an_undeclared_field_is_an_unknown_parameter(): + problem = _problem({"filter[secret_column][eq]": "x"}) + + assert problem.type == f"{PROBLEM_TYPE_BASE}unknown-query-parameter" + assert "filter[secret_column][eq]" in problem.detail + + +def test_every_declared_parameter_is_accepted(): + """Guards the unknown-param check against rejecting the spec's own contract.""" + plan = _plan( + { + "page": "2", + "page_size": "10", + "sort": "-max_budget", + "q": "prod", + "filter[max_budget][gte]": "5", + "filter[created_by][in]": "a,b", + } + ) + + assert plan.take == 10 + + +# ---------------------------------------------------------------- invariant 7 + + +def test_sorting_by_an_undeclared_field_is_rejected(): + problem = _problem({"sort": "api_key"}) + + assert problem.status == 400 + assert problem.type == f"{PROBLEM_TYPE_BASE}invalid-sort-field" + assert problem.allowed == ["budget_id", "created_at", "max_budget"] + assert "api_key" in problem.detail + + +def test_one_bad_key_rejects_the_whole_multi_key_sort(): + """Dropping the unknown key and sorting by the rest would silently return a + differently-ordered page than the one asked for.""" + assert _problem({"sort": "-created_at,api_key"}).type == f"{PROBLEM_TYPE_BASE}invalid-sort-field" + + +def test_a_double_dash_prefix_is_not_a_descending_sort(): + assert _problem({"sort": "--created_at"}).type == f"{PROBLEM_TYPE_BASE}invalid-sort-field" + + +# ---------------------------------------------------------------- invariant 8 + + +def test_an_operator_the_field_does_not_declare_is_rejected(): + problem = _problem({"filter[max_budget][contains]": "5"}) + + assert problem.status == 400 + assert problem.type == f"{PROBLEM_TYPE_BASE}unsupported-filter-operator" + assert problem.allowed == ["eq", "gte", "is_null", "lte"] + assert "contains" in problem.detail + + +def test_the_same_operator_is_accepted_on_a_field_that_declares_it(): + """Pins the rejection to the field's own operator set rather than a global denylist.""" + conjuncts = _conjuncts(_plan({"filter[created_by][contains]": "ops"})) + + assert conjuncts == (Compare(field="created_by", op="contains", value="ops"),) + + +def test_a_string_that_is_not_an_operator_at_all_is_an_unknown_parameter(): + """`gt3` is a typo, not an operator the field withheld, so the useful reply is the + parameter list rather than this field's operator set.""" + problem = _problem({"filter[max_budget][gt3]": "5"}) + + assert problem.type == f"{PROBLEM_TYPE_BASE}unknown-query-parameter" + assert problem.allowed is not None + assert "filter[max_budget][gte]" in problem.allowed + + +# ---------------------------------------------------------------- invariant 9 + + +def test_search_against_a_spec_with_nothing_searchable_is_rejected(): + """A silently-empty search filter returns the unfiltered table, which reads as + "no results were filtered out" rather than "this resource cannot be searched".""" + problem = _problem({"q": "prod"}, spec=_spec(searchable=frozenset())) + + assert problem.status == 400 + assert problem.type == f"{PROBLEM_TYPE_BASE}unknown-query-parameter" + assert problem.allowed is not None + assert "q" not in problem.allowed + + +def test_search_is_a_case_insensitive_or_across_every_searchable_field(): + conjuncts = _conjuncts(_plan({"q": "Prod"})) + + assert conjuncts == ( + AnyOf( + clauses=( + Compare(field="budget_id", op="contains", value="Prod"), + Compare(field="created_by", op="contains", value="Prod"), + ) + ), + ) + + +def test_an_empty_search_string_adds_no_filter(): + assert _plan({"q": ""}).where == () + + +# --------------------------------------------------------------- invariant 10 + + +def test_multi_key_sort_parses_the_json_api_grammar(): + order = _plan({"sort": "-created_at,budget_id,-max_budget"}).order + + assert order[:3] == ( + SortKey(field="created_at", descending=True), + SortKey(field="budget_id", descending=False), + SortKey(field="max_budget", descending=True), + ) + + +def test_sort_segments_tolerate_surrounding_whitespace(): + assert _plan({"sort": "-created_at, budget_id"}).order[:2] == ( + SortKey(field="created_at", descending=True), + SortKey(field="budget_id", descending=False), + ) + + +# ------------------------------------------------------------- filter parsing + + +def test_comparison_operators_become_prisma_range_fragments(): + conjuncts = _conjuncts(_plan({"filter[max_budget][gte]": "5", "filter[max_budget][lte]": "50"})) + + assert conjuncts == ( + Compare(field="max_budget", op="gte", value=5.0), + Compare(field="max_budget", op="lte", value=50.0), + ) + + +def test_eq_is_a_bare_value_not_a_wrapped_one(): + assert _conjuncts(_plan({"filter[tpm_limit][eq]": "100"})) == (Compare(field="tpm_limit", op="eq", value=100),) + + +def test_a_filter_with_no_operator_bracket_means_eq(): + """`filter[status]=active` is the design doc's canonical spelling for equality; + only the non-eq operators carry a second bracket.""" + assert _conjuncts(_plan({"filter[tpm_limit]": "100"})) == (Compare(field="tpm_limit", op="eq", value=100),) + + +def test_the_bare_form_and_the_explicit_eq_form_agree(): + assert _plan({"filter[created_by]": "alice"}) == _plan({"filter[created_by][eq]": "alice"}) + + +def test_the_bare_form_still_coerces_to_the_declared_type(): + assert _problem({"filter[tpm_limit]": "1.5"}).status == 400 + + +def test_the_bare_form_is_rejected_on_a_field_that_does_not_declare_eq(): + """The shorthand is sugar for the eq operator, not a bypass around the operator set.""" + problem = _problem({"filter[created_at]": "2026-07-23T00:00:00Z"}) + + assert problem.type == f"{PROBLEM_TYPE_BASE}unsupported-filter-operator" + assert problem.allowed == ["gte", "lte"] + + +def test_the_allowed_set_advertises_the_bare_spelling_for_eq(): + allowed = _problem({"nope": "1"}).allowed + + assert allowed is not None + assert "filter[max_budget]" in allowed + assert "filter[max_budget][eq]" not in allowed + assert "filter[created_at][gte]" in allowed + assert "filter[created_at]" not in allowed + + +@pytest.mark.parametrize( + "name", + ["filter[]", "filter[a][b][c]", "filter[a][", "filter", "filter[a][gte", "filter[max_budget]]["], + ids=["empty", "triple", "unbalanced", "bare-word", "unterminated", "bracketed-field"], +) +def test_malformed_filter_keys_are_unknown_parameters_not_eq_filters(name): + """A malformed key must not fall through to the bare-eq branch and silently filter + on a field nobody declared. `field in spec.filters` is the gate that makes this hold, + which is also why the parser needs no separate well-formedness guard.""" + assert _problem({name: "x"}).type == f"{PROBLEM_TYPE_BASE}unknown-query-parameter" + + +def test_in_splits_on_commas_and_coerces_every_member(): + assert _conjuncts(_plan({"filter[created_by][in]": "alice, bob"})) == ( + Within(field="created_by", values=("alice", "bob")), + ) + + +def test_is_null_true_matches_rows_with_no_budget(): + """Budgets renders a null max_budget as "Unlimited"; without is_null there is no way + to ask for those rows.""" + assert _conjuncts(_plan({"filter[max_budget][is_null]": "true"})) == ( + IsNull(field="max_budget", negated=False), + ) + + +def test_is_null_false_matches_rows_that_have_one(): + assert _conjuncts(_plan({"filter[max_budget][is_null]": "false"})) == ( + IsNull(field="max_budget", negated=True), + ) + + +def test_is_null_rejects_a_non_boolean(): + assert _problem({"filter[max_budget][is_null]": "maybe"}).status == 400 + + +@pytest.mark.parametrize( + "params", + [ + {"filter[max_budget][gte]": "lots"}, + {"filter[tpm_limit][eq]": "1.5"}, + {"filter[created_at][gte]": "yesterday"}, + {"filter[created_by][in]": "alice,"}, + ], + ids=["float", "int", "datetime", "in-member"], +) +def test_a_value_that_does_not_match_the_declared_type_is_rejected(params): + numeric_in = _spec_with(filters={**_spec().filters, "created_by": FilterSpec(type=int, ops=frozenset({"in"}))}) + target = numeric_in if "filter[created_by][in]" in params else _spec() + + assert _problem(params, spec=target).status == 400 + + +def test_a_datetime_filter_is_normalised_to_utc(): + """The dashboard sends both offset-bearing and naive timestamps; reading a naive one + as server-local time would shift the window off what the table is showing.""" + with_offset = _conjuncts(_plan({"filter[created_at][gte]": "2026-07-23T02:00:00+02:00"})) + naive = _conjuncts(_plan({"filter[created_at][gte]": "2026-07-23 00:00:00"})) + + assert with_offset == (Compare(field="created_at", op="gte", value=datetime(2026, 7, 23, tzinfo=timezone.utc)),) + assert naive == with_offset + + +def test_filters_are_ordered_deterministically(): + """Two requests differing only in query-string order must plan identically, or the + plan stops being a comparable value.""" + forwards = _plan({"filter[created_by][eq]": "a", "filter[max_budget][gte]": "5"}) + backwards = _plan({"filter[max_budget][gte]": "5", "filter[created_by][eq]": "a"}) + + assert forwards == backwards + + +# ------------------------------------------------------------------- envelope + + +@pytest.mark.asyncio +async def test_returns_the_page_mode_envelope(): + executor = RecordingExecutor( + rows=tuple(BudgetRow(budget_id=f"b{i}", max_budget=float(i), created_by="u") for i in range(10)), + total_count=42, + ) + + response = await handle_list( + spec=_spec(), executor=executor, request=_request("page=2&page_size=5"), caller=CALLER + ) + body = response.model_dump(by_alias=True) + + assert body["meta"] == {"total_count": 42, "page": 2, "page_size": 5, "total_pages": 9} + assert set(body) == {"data", "meta", "links"} + assert "has_more" not in body["meta"] + + +@pytest.mark.asyncio +async def test_serializes_rows_flat_without_a_json_api_resource_wrapper(): + executor = RecordingExecutor(rows=(BudgetRow(budget_id="b1", max_budget=None, created_by="u"),)) + + response = await handle_list(spec=_spec(), executor=executor, request=_request(), caller=CALLER) + body = response.model_dump(by_alias=True) + + assert body["data"] == [{"budget_id": "b1", "max_budget": None}] + assert "attributes" not in body["data"][0] + assert "created_by" not in body["data"][0] + + +@pytest.mark.asyncio +async def test_links_let_a_client_page_without_building_urls(): + executor = RecordingExecutor(rows=(), total_count=42) + + response = await handle_list( + spec=_spec(), executor=executor, request=_request("page=2&page_size=5"), caller=CALLER + ) + links = response.model_dump(by_alias=True)["links"] + + assert links["self"] == f"{BUDGETS_PATH}?page_size=5&page=2" + assert links["first"] == f"{BUDGETS_PATH}?page_size=5&page=1" + assert links["prev"] == f"{BUDGETS_PATH}?page_size=5&page=1" + assert links["next"] == f"{BUDGETS_PATH}?page_size=5&page=3" + assert links["last"] == f"{BUDGETS_PATH}?page_size=5&page=9" + + +@pytest.mark.asyncio +async def test_the_last_page_has_no_next_link(): + executor = RecordingExecutor(rows=(), total_count=10) + + response = await handle_list( + spec=_spec(), executor=executor, request=_request("page=2&page_size=5"), caller=CALLER + ) + links = response.model_dump(by_alias=True)["links"] + + assert links["next"] is None + assert links["prev"] == f"{BUDGETS_PATH}?page_size=5&page=1" + + +@pytest.mark.asyncio +async def test_an_empty_result_set_still_resolves_every_link(): + executor = RecordingExecutor(rows=(), total_count=0) + + response = await handle_list(spec=_spec(), executor=executor, request=_request(), caller=CALLER) + body = response.model_dump(by_alias=True) + + assert body["data"] == [] + assert body["meta"]["total_pages"] == 0 + assert body["links"]["first"] == body["links"]["last"] == f"{BUDGETS_PATH}?page=1" + assert body["links"]["next"] is None + assert body["links"]["prev"] is None + + +@pytest.mark.asyncio +async def test_the_executor_counts_the_same_predicate_it_reads(): + """Counting a wider predicate than the read inflates total_pages and hands the UI + pages that are always empty.""" + executor = RecordingExecutor(rows=(), total_count=3) + spec = _spec(scope=lambda caller: ScopeWhere(where=(Compare(field="created_by", op="eq", value=caller.user_id),))) + + await handle_list(spec=spec, executor=executor, request=_request("filter[max_budget][gte]=5"), caller=CALLER) + + assert executor.plan is not None + assert executor.count_where == executor.plan.where + + +@pytest.mark.asyncio +async def test_a_rejected_request_is_raised_as_a_problem_before_any_query(): + executor = RecordingExecutor(rows=()) + + with pytest.raises(ManagementProblem) as raised: + await handle_list(spec=_spec(), executor=executor, request=_request("sort=api_key"), caller=CALLER) + + assert raised.value.problem.status == 400 + assert executor.count_where is None + + +# ------------------------------------------------------- spec construction + + +def test_a_default_page_size_above_the_cap_is_rejected_at_construction(): + """The cap is only enforced on a supplied page_size, so a default above it would serve + more rows than the resource allows on exactly the request that omits page_size.""" + with pytest.raises(ValueError, match="default_page_size"): + _spec_with(default_page_size=200, max_page_size=100) + + +@pytest.mark.parametrize( + "overrides", + [{"default_page_size": 0}, {"default_page_size": -5}, {"max_page_size": 0}], + ids=["zero-default", "negative-default", "zero-cap"], +) +def test_a_non_positive_page_size_is_rejected_at_construction(overrides): + """take=0 divides by zero when handle_list computes total_pages, so the resource would + 500 on every request instead of failing when it is registered.""" + with pytest.raises(ValueError, match="default_page_size"): + _spec_with(**overrides) + + +def test_a_default_sort_on_a_non_sortable_field_is_rejected_at_construction(): + """Caller-supplied sort is validated against `sortable`; default_sort is not read from + the request, so without this it reaches order_by_sql and yields invalid SQL.""" + with pytest.raises(ValueError, match="default_sort"): + _spec_with(default_sort=(SortKey(field="not_a_column", descending=True),)) + + +def test_an_empty_tiebreaker_is_rejected_at_construction(): + with pytest.raises(ValueError, match="tiebreaker"): + _spec_with(tiebreaker="") + + +def test_a_page_size_equal_to_the_cap_is_a_valid_spec(): + """Guards the bound against being tightened into an off-by-one that bans max==default.""" + assert _spec_with(default_page_size=100, max_page_size=100).default_page_size == 100 + + +# ------------------------------------------------------ repeated parameters + + +@pytest.mark.asyncio +async def test_a_repeated_query_parameter_is_rejected(): + """Starlette keeps the last value, so ?page=1&page=999 would page from 999 with nothing + telling the caller which one won. The doc rejects silently-altered params for this reason.""" + executor = RecordingExecutor(rows=()) + + with pytest.raises(ManagementProblem) as raised: + await handle_list(spec=_spec(), executor=executor, request=_request("page=1&page=999"), caller=CALLER) + + assert raised.value.problem.status == 400 + assert raised.value.problem.type == f"{PROBLEM_TYPE_BASE}duplicate-query-parameter" + assert "page" in raised.value.problem.detail + assert executor.count_where is None + + +@pytest.mark.asyncio +async def test_a_repeated_filter_parameter_is_rejected(): + executor = RecordingExecutor(rows=()) + + with pytest.raises(ManagementProblem) as raised: + await handle_list( + spec=_spec(), + executor=executor, + request=_request("filter[created_by][eq]=alice&filter[created_by][eq]=bob"), + caller=CALLER, + ) + + assert raised.value.problem.type == f"{PROBLEM_TYPE_BASE}duplicate-query-parameter" + assert "filter[created_by][eq]" in raised.value.problem.detail + + +@pytest.mark.asyncio +async def test_distinct_parameters_are_not_treated_as_duplicates(): + """Guards the check against rejecting two different operators on one field, which is + how a range filter is expressed.""" + executor = RecordingExecutor(rows=(), total_count=0) + + response = await handle_list( + spec=_spec(), + executor=executor, + request=_request("filter[max_budget][gte]=5&filter[max_budget][lte]=50&page=2"), + caller=CALLER, + ) + + assert response.meta.page == 2 + assert executor.count_where is not None + + +@pytest.mark.asyncio +async def test_a_denied_scope_outranks_a_duplicate_parameter(): + """Permission is the stronger statement about the caller, so it is answered first.""" + spec = _spec(scope=lambda caller: ScopeDenied(reason="nope")) + executor = RecordingExecutor(rows=()) + + with pytest.raises(ManagementProblem) as raised: + await handle_list(spec=spec, executor=executor, request=_request("page=1&page=2"), caller=CALLER) + + assert raised.value.problem.status == 403 + + +# --------------------------------------------------- facet-mode regression + + +def test_the_facet_page_shapes_are_untouched_by_page_mode(): + """The live facet endpoint reports `has_more` and has no first/last, because it + deliberately skips the COUNT(*). Folding it into the page-mode shapes would either + break its response or make every keystroke pay for a full-table count.""" + assert set(PageMeta.model_fields) == {"page", "page_size", "has_more"} + assert set(PageLinks.model_fields) == {"self_link", "prev", "next"} + + links = build_page_links(request=_request("q=ac&page=2"), page=2, has_more=True).model_dump(by_alias=True) + + assert set(links) == {"self", "prev", "next"} + assert links["next"] == "/management/v1/budgets?q=ac&page=3" diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 1e8add52f74..8dbf38555b3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -839,7 +839,7 @@ class TestUpdateModel: ), patch( "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", - new=AsyncMock(return_value=True), + new=AsyncMock(return_value=None), ) as mock_clear_cache, ): await update_model( @@ -3223,7 +3223,7 @@ class TestPatchModelBlockedAuthGate: ), patch( "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", - new=AsyncMock(return_value=True), + new=AsyncMock(return_value=None), ), ): result = await patch_model( @@ -3314,6 +3314,72 @@ class TestWriteSurfacesReloadDrop: before=frozenset({"m-live", "m-collateral"}), written_models=[("m-live", None)], action="update" ) + def test_a_model_the_db_no_longer_has_is_not_collateral(self, monkeypatch): + """Another pod deleting a model is not this pod's reload breaking. + + A pod that has not yet polled the delete still lists the id when the write + snapshots `before`; the reload it triggers then evicts the id because the db no + longer has it. That eviction is the reconcile working, so it must not fail the + write. `still_desired` is the db + config set the reload reconciled against, so + an id missing from it drops out of the collateral diff. + + The cases below, in order: an id the db no longer wants is not collateral and the + write succeeds; an id the db still wants that stopped serving is real degradation + and still raises, so a genuinely broken reload is caught; and with no reconcile at + all the desired set is unknown, so every drop is reported. + """ + import litellm + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + raise_if_reload_degraded_serving, + reload_serving_verdict, + ) + + live_router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "m-live", "db_model": True}, + } + ] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", live_router) + + _, collateral = reload_serving_verdict( + before=frozenset({"m-live", "m-deleted-elsewhere"}), + written_models=[("m-live", None)], + written_must_serve=True, + still_desired=frozenset({"m-live"}), + ) + assert collateral == () + + assert ( + raise_if_reload_degraded_serving( + before=frozenset({"m-live", "m-deleted-elsewhere"}), + written_models=[("m-live", None)], + action="create", + still_desired=frozenset({"m-live"}), + ) + is None + ) + + with pytest.raises(ProxyException, match="m-should-be-serving"): + raise_if_reload_degraded_serving( + before=frozenset({"m-live", "m-should-be-serving"}), + written_models=[("m-live", None)], + action="create", + still_desired=frozenset({"m-live", "m-should-be-serving"}), + ) + + with pytest.raises(ProxyException, match="m-deleted-elsewhere"): + raise_if_reload_degraded_serving( + before=frozenset({"m-live", "m-deleted-elsewhere"}), + written_models=[("m-live", None)], + action="create", + still_desired=None, + ) + class TestModelInfoAsMapping: """The model_info column reaches consumers as a dict or as its JSON string; this is 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 76ba0e3dc67..4fe1b54694f 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 @@ -1,17 +1,20 @@ +import inspect import json import os import sys -from typing import Any, Dict, Optional +from collections.abc import Sequence +from typing import Optional import pytest from fastapi import HTTPException from fastapi.testclient import TestClient +from prisma.actions import LiteLLM_VerificationTokenActions sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from unittest.mock import patch +from unittest.mock import Mock, patch import litellm from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -21,6 +24,28 @@ from litellm.types.tag_management import TagDeleteRequest, TagInfoRequest, TagNe client = TestClient(app) +class FakeVerificationTokenTable: + """Stand-in for ``prisma_client.db.litellm_verificationtoken``. + + ``AsyncMock`` swallows any keyword argument, so a plain mock cannot catch a + call that the generated prisma client would reject at runtime. This double + binds every call against the real ``find_many`` signature, so passing an + unsupported kwarg (e.g. ``select``) raises the same ``TypeError`` the proxy + surfaces as an HTTP 500. + """ + + def __init__(self, records: Sequence[Mock]): + self._records = tuple(records) + self.calls: list[dict[str, object]] = [] + + async def find_many(self, **kwargs: object) -> tuple[Mock, ...]: + inspect.signature(LiteLLM_VerificationTokenActions.find_many).bind( + self, **kwargs + ) + self.calls.append(kwargs) + return self._records + + @pytest.mark.asyncio async def test_create_and_get_tag(): """ @@ -380,6 +405,7 @@ async def test_list_tags_no_dynamic_tags(): app.dependency_overrides.clear() +@pytest.mark.asyncio async def test_internal_user_list_tags_only_returns_tags_used_by_their_keys(): """ Internal users can view tag usage, but the tag list must be scoped to tags @@ -404,9 +430,8 @@ async def test_internal_user_list_tags_only_returns_tags_used_by_their_keys(): owned_key_record = Mock() owned_key_record.token = "owned-key" - mock_db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[owned_key_record] - ) + fake_token_table = FakeVerificationTokenTable([owned_key_record]) + mock_db.litellm_verificationtoken = fake_token_table mock_db.litellm_dailytagspend.group_by = AsyncMock( return_value=[ @@ -446,10 +471,9 @@ async def test_internal_user_list_tags_only_returns_tags_used_by_their_keys(): "stored-owned-tag", "dynamic-owned-tag", ] - mock_db.litellm_verificationtoken.find_many.assert_awaited_once_with( - where={"user_id": "internal-user-123"}, - select={"token": True}, - ) + assert fake_token_table.calls == [ + {"where": {"user_id": "internal-user-123"}} + ] mock_db.litellm_dailytagspend.group_by.assert_awaited_once_with( by=["tag"], where={ @@ -468,6 +492,54 @@ async def test_internal_user_list_tags_only_returns_tags_used_by_their_keys(): app.dependency_overrides.clear() +@pytest.mark.asyncio +async def test_internal_user_list_tags_does_not_500_on_unsupported_prisma_kwarg(): + """ + Regression: /tag/list returned 500 for every internal user because the + non-admin branch looked up the caller's keys with + ``find_many(select={"token": True})``, and the generated prisma client has no + ``select`` kwarg. This reproduces the reported case exactly: a freshly created + internal user with no tag spend yet, which must get an empty 200 rather than + "LiteLLM_VerificationTokenActions.find_many() got an unexpected keyword + argument 'select'". + """ + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + api_key="new-user-key", + user_id="brand-new-internal-user", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_db = Mock() + mock_prisma.db = mock_db + + key_record = Mock() + key_record.token = "new-user-key" + fake_token_table = FakeVerificationTokenTable([key_record]) + mock_db.litellm_verificationtoken = fake_token_table + + mock_db.litellm_dailytagspend.group_by = AsyncMock(return_value=[]) + mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[]) + + response = client.get( + "/tag/list", headers={"Authorization": "Bearer new-user-key"} + ) + + assert response.status_code == 200, response.text + assert response.json() == [] + assert fake_token_table.calls == [ + {"where": {"user_id": "brand-new-internal-user"}} + ] + finally: + app.dependency_overrides.clear() + + @pytest.mark.asyncio async def test_list_tags_with_date_range_filters_dynamic_tags(): """ @@ -537,9 +609,8 @@ async def test_internal_user_tag_daily_activity_is_scoped_to_their_keys(): owned_key_record = Mock() owned_key_record.token = "owned-key" - mock_db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[owned_key_record] - ) + fake_token_table = FakeVerificationTokenTable([owned_key_record]) + mock_db.litellm_verificationtoken = fake_token_table mock_get_daily_activity.return_value = "daily-activity-response" result = await get_tag_daily_activity( @@ -549,6 +620,7 @@ async def test_internal_user_tag_daily_activity_is_scoped_to_their_keys(): ) assert result == "daily-activity-response" + assert fake_token_table.calls == [{"where": {"user_id": "internal-user-123"}}] mock_get_daily_activity.assert_awaited_once() assert mock_get_daily_activity.await_args.kwargs["api_key"] == ["owned-key"] @@ -583,9 +655,8 @@ async def test_internal_user_tag_daily_activity_rejects_unowned_api_key_filter() owned_key_record = Mock() owned_key_record.token = "owned-key" - mock_db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[owned_key_record] - ) + fake_token_table = FakeVerificationTokenTable([owned_key_record]) + mock_db.litellm_verificationtoken = fake_token_table result = await get_tag_daily_activity( start_date="2025-01-01", end_date="2025-01-31", @@ -593,6 +664,7 @@ async def test_internal_user_tag_daily_activity_rejects_unowned_api_key_filter() user_api_key_dict=mock_user_auth, ) + assert fake_token_table.calls == [{"where": {"user_id": "internal-user-123"}}] assert result.results == [] assert result.metadata.total_spend == 0 assert result.metadata.total_api_requests == 0 @@ -626,7 +698,8 @@ async def test_internal_user_tag_daily_activity_scopes_to_current_key_without_us ): mock_db = Mock() mock_prisma.db = mock_db - mock_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + fake_token_table = FakeVerificationTokenTable([]) + mock_db.litellm_verificationtoken = fake_token_table mock_get_daily_activity.return_value = "daily-activity-response" result = await get_tag_daily_activity( @@ -636,7 +709,7 @@ async def test_internal_user_tag_daily_activity_scopes_to_current_key_without_us ) assert result == "daily-activity-response" - mock_db.litellm_verificationtoken.find_many.assert_not_awaited() + assert fake_token_table.calls == [] mock_get_daily_activity.assert_awaited_once() assert mock_get_daily_activity.await_args.kwargs["api_key"] == [ "current-owned-key" @@ -669,7 +742,8 @@ async def test_internal_user_tag_daily_activity_without_any_scoped_keys_returns_ ): mock_db = Mock() mock_prisma.db = mock_db - mock_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + fake_token_table = FakeVerificationTokenTable([]) + mock_db.litellm_verificationtoken = fake_token_table result = await get_tag_daily_activity( start_date="2025-01-01", @@ -680,7 +754,7 @@ async def test_internal_user_tag_daily_activity_without_any_scoped_keys_returns_ assert result.results == [] assert result.metadata.total_spend == 0 assert result.metadata.total_api_requests == 0 - mock_db.litellm_verificationtoken.find_many.assert_not_awaited() + assert fake_token_table.calls == [] mock_get_daily_activity.assert_not_awaited() diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a658a4b7353..25a0ff644f8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -3627,9 +3627,9 @@ async def test_list_team_v2_with_invalid_status(): @pytest.mark.asyncio async def test_list_team_v2_search_builds_or_clause(): """ - `search` should be passed as a Prisma OR across team_id (exact) and - team_alias (case-insensitive contains), so the UI can hit a single - backend filter with either a UUID or a name fragment. + `search` should be passed as a Prisma OR across an exact team_id match and a + case-insensitive team_alias contains, so the UI needs one backend filter. + Exact id matching is the documented default and must not change. """ from unittest.mock import AsyncMock, Mock, patch @@ -3671,6 +3671,54 @@ async def test_list_team_v2_search_builds_or_clause(): } +@pytest.mark.asyncio +async def test_list_team_v2_search_team_id_match_prefix(): + """ + Opting into `search_team_id_match="prefix"` should widen the team_id side of + the search OR to an index-friendly prefix match, so the first characters of a + team id quoted in a proxy error find the team. + """ + from unittest.mock import AsyncMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 + + mock_request = Mock(spec=Request) + mock_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: + mock_db = Mock() + mock_prisma_client.db = mock_db + mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + mock_db.litellm_teamtable.count = AsyncMock(return_value=0) + + await list_team_v2( + http_request=mock_request, + user_id=None, + organization_id=None, + team_id=None, + team_alias=None, + search="66c432fa", + search_team_id_match="prefix", + user_api_key_dict=mock_admin, + page=1, + page_size=10, + status=None, + ) + + find_many_kwargs = mock_db.litellm_teamtable.find_many.call_args.kwargs + assert find_many_kwargs["where"] == { + "OR": [ + {"team_id": {"startsWith": "66c432fa"}}, + {"team_alias": {"contains": "66c432fa", "mode": "insensitive"}}, + ] + } + + @pytest.mark.asyncio async def test_list_team_v2_search_composes_with_user_id_filter(): """ @@ -3724,6 +3772,7 @@ async def test_list_team_v2_search_composes_with_user_id_filter(): team_id=None, team_alias=None, search="team_a", + search_team_id_match="prefix", user_api_key_dict=mock_user_api_key_dict, page=1, page_size=10, @@ -3733,7 +3782,7 @@ async def test_list_team_v2_search_composes_with_user_id_filter(): find_many_kwargs = mock_db.litellm_teamtable.find_many.call_args.kwargs where = find_many_kwargs["where"] assert where["OR"] == [ - {"team_id": "team_a"}, + {"team_id": {"startsWith": "team_a"}}, {"team_alias": {"contains": "team_a", "mode": "insensitive"}}, ] assert where["team_id"] == {"in": ["team_a", "team_b"]} diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py index 45c3c6c2466..18ea5c3f27d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -7,8 +7,10 @@ and litellm.proxy.proxy_server.prisma_client) because the endpoint code imports these inside function bodies to avoid circular imports. """ +import inspect import os import sys +from collections.abc import Sequence from datetime import datetime, timedelta, timezone from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -16,6 +18,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import FastAPI from fastapi.testclient import TestClient +from prisma.actions import LiteLLM_TeamTableActions sys.path.insert(0, os.path.abspath("../../..")) @@ -74,6 +77,47 @@ def _group_row(tool_name: str, spend: float, request_count: int, total_tokens: i return {"tool_name": tool_name, "_sum": {"spend": spend, "total_tokens": total_tokens, "request_count": request_count}} +class FakeTeamTable: + """Stand-in for ``prisma_client.db.litellm_teamtable``. + + ``AsyncMock`` accepts any keyword argument, so a plain mock cannot catch a call + the generated prisma client would reject at runtime. This double binds every + call against the real action signature, so an unsupported kwarg (e.g. ``select``) + raises the same ``TypeError`` the proxy surfaces as an HTTP 500. + """ + + def __init__(self, rows: Sequence[MagicMock], updated_count: int = 1): + self._rows = tuple(rows) + self._updated_count = updated_count + self.find_unique_calls: list[dict[str, object]] = [] + self.update_many_calls: list[dict[str, object]] = [] + + async def find_unique(self, **kwargs: object) -> Optional[MagicMock]: + inspect.signature(LiteLLM_TeamTableActions.find_unique).bind(self, **kwargs) + index = len(self.find_unique_calls) + self.find_unique_calls.append(kwargs) + return self._rows[index] if index < len(self._rows) else None + + async def update_many(self, **kwargs: object) -> int: + inspect.signature(LiteLLM_TeamTableActions.update_many).bind(self, **kwargs) + self.update_many_calls.append(kwargs) + return self._updated_count + + +def _team_row(object_permission_id: Optional[str]) -> MagicMock: + row = MagicMock() + row.object_permission_id = object_permission_id + return row + + +def _team_policy_prisma(team_table: FakeTeamTable) -> MagicMock: + prisma = MagicMock() + prisma.db.litellm_teamtable = team_table + prisma.db.litellm_objectpermissiontable.create = AsyncMock() + prisma.db.litellm_objectpermissiontable.delete = AsyncMock() + return prisma + + def _rollup_prisma(group_rows: list, daily_rows: list | None = None) -> MagicMock: prisma = MagicMock() prisma.db.query_raw = AsyncMock(return_value=[]) @@ -161,6 +205,72 @@ class TestToolManagementEndpoints: assert body["input_policy"] == "blocked" assert body["updated"] is True + @patch( + "litellm.proxy.db.tool_registry_writer.add_tool_to_object_permission_blocked", + new_callable=AsyncMock, + ) + def test_update_tool_policy_for_team_does_not_500_on_unsupported_prisma_kwarg(self, mock_block): + """ + Regression: POST /v1/tool/policy with a team_id returned 500 with + "LiteLLM_TeamTableActions.find_unique() got an unexpected keyword argument + 'select'". The team lookup passed select={"object_permission_id": True}, and + the generated prisma client has no select kwarg, so the call raised TypeError + that the handler turned into a 500. + """ + mock_block.return_value = True + team_table = FakeTeamTable([_team_row("existing-op-id")]) + + with patch("litellm.proxy.proxy_server.prisma_client", _team_policy_prisma(team_table)), patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry" + ) as mock_registry: + mock_registry.return_value.is_initialized.return_value = False + resp = self.client.post( + "/v1/tool/policy", + json={"tool_name": "my_tool", "input_policy": "blocked", "team_id": "team-123"}, + ) + + assert resp.status_code == 200, resp.text + assert resp.json()["team_id"] == "team-123" + assert team_table.find_unique_calls == [{"where": {"team_id": "team-123"}}] + mock_block.assert_awaited_once_with( + prisma_client=mock_block.await_args.kwargs["prisma_client"], + object_permission_id="existing-op-id", + tool_name="my_tool", + ) + + @patch( + "litellm.proxy.db.tool_registry_writer.add_tool_to_object_permission_blocked", + new_callable=AsyncMock, + ) + def test_update_tool_policy_for_team_reresolves_when_permission_race_is_lost(self, mock_block): + """ + Same regression on the second team lookup: when the team has no object + permission yet and a concurrent writer wins the update_many (0 rows updated), + the handler re-reads the team to pick up the winner's id. That fallback read + passed the same unsupported select kwarg. + """ + mock_block.return_value = True + team_table = FakeTeamTable( + [_team_row(None), _team_row("winner-op-id")], + updated_count=0, + ) + + with patch("litellm.proxy.proxy_server.prisma_client", _team_policy_prisma(team_table)), patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry" + ) as mock_registry: + mock_registry.return_value.is_initialized.return_value = False + resp = self.client.post( + "/v1/tool/policy", + json={"tool_name": "my_tool", "input_policy": "blocked", "team_id": "team-123"}, + ) + + assert resp.status_code == 200, resp.text + assert team_table.find_unique_calls == [ + {"where": {"team_id": "team-123"}}, + {"where": {"team_id": "team-123"}}, + ] + assert mock_block.await_args.kwargs["object_permission_id"] == "winner-op-id" + @patch("litellm.proxy.proxy_server.prisma_client", None) def test_list_tools_no_db_returns_500(self): resp = self.client.get("/v1/tool/list") diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 8d1d8185e4d..99287b92b8f 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -1404,12 +1404,12 @@ def test_ProxyConfig_get_model_info_with_id_missing_model_id_raises(monkeypatch) @pytest.mark.asyncio -async def test_ProxyConfig__delete_deployment_empty_returns_zero(monkeypatch): +async def test_ProxyConfig__delete_deployment_no_router_returns_none(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) pc = ProxyConfig() result = await pc._delete_deployment(db_models=[]) - snapshot = {"deleted": result, "router_was": "none", "empty_db_models": True} - assert snapshot == {"deleted": 0, "router_was": "none", "empty_db_models": True} + snapshot = {"still_desired": result, "router_was": "none", "empty_db_models": True} + assert snapshot == {"still_desired": None, "router_was": "none", "empty_db_models": True} @pytest.mark.asyncio 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 aa20c3f6ed4..f6216d1646e 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 @@ -1582,6 +1582,47 @@ async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_ui_view_session_spend_logs_rehydrates_metadata_jsonb_text(client, monkeypatch): + """The session sidebar reads metadata fields as object properties, so this endpoint + must re-hydrate the JSONB column that query_raw hands back as a string, exactly as + /spend/logs/ui does (#29674). Property access on a string is silently undefined, + so a row's origin, status and error information all read as absent without this. + """ + raw_row = { + "request_id": "req-classifier-1", + "session_id": "session-123", + "startTime": "2024-01-01T00:00:00Z", + "metadata": json.dumps({"internal_call_origin": "autorouter_classifier", "status": "success"}), + } + + class MockDB: + async def count(self, *args, **kwargs): + return 1 + + async def query_raw(self, sql_query, session_id, page_size, skip, *scope_params): + return [dict(raw_row)] + + class MockPrismaClient: + def __init__(self): + self.db = MockDB() + self.db.litellm_spendlogs = self.db + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient()) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + response = client.get("/spend/logs/session/ui", params={"session_id": "session-123"}) + assert response.status_code == 200 + row = response.json()["data"][0] + assert isinstance(row["metadata"], dict) + assert row["metadata"]["internal_call_origin"] == "autorouter_classifier" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_session_spend_logs_scopes_non_admin_to_own_logs(client, monkeypatch): own_log = { diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b9a33bd2cef..ab64224fdc3 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2459,13 +2459,11 @@ async def test_delete_deployment_type_mismatch(): patch("litellm.proxy.proxy_server.user_config_file_path", "test_config.yaml"), ): # Call the function under test - deleted_count = await pc._delete_deployment(db_models=[]) + still_desired = await pc._delete_deployment(db_models=[]) # The two SHA-hash models have no corresponding entry in combined_id_list # and must be evicted. - assert ( - deleted_count == 2 - ), f"Expected 2 deletions (SHA-hash models), got {deleted_count}" + assert len(deleted_ids) == 2, f"Expected 2 deletions (SHA-hash models), got {deleted_ids}" assert ( "a96e12e76b36a57cfae57a41288eb41567629cac89b4828c6f7074afc3534695" in deleted_ids @@ -2485,6 +2483,12 @@ async def test_delete_deployment_type_mismatch(): "12345679" not in deleted_ids ), f"Model 12345679 should NOT be deleted. Deleted IDs: {deleted_ids}" + assert still_desired is not None + assert {"12345678", "12345679"} <= still_desired, ( + "the int-keyed config models must come back as strings in the desired set, so a " + f"caller judging its own reload reads them as wanted rather than evicted; got {still_desired}" + ) + @pytest.mark.asyncio async def test_get_config_from_file(tmp_path, monkeypatch): @@ -9119,10 +9123,13 @@ class TestDeleteDeploymentSync: with patch.object( proxy_config, "get_config", AsyncMock(return_value={"model_list": []}) ): - count = await proxy_config._delete_deployment(db_models=[]) + still_desired = await proxy_config._delete_deployment(db_models=[]) mock_router.delete_deployment.assert_called_once_with(id="model-id-to-evict") - assert count == 1 + assert still_desired == frozenset(), ( + "an empty db and an empty config want nothing, which must stay distinct from " + f"the None returned when no reconcile ran at all; got {still_desired}" + ) @pytest.mark.asyncio async def test_update_llm_router_skips_update_on_db_fetch_failure(self): @@ -10678,3 +10685,83 @@ async def test_async_data_generator_forwards_usage_chunk_without_strip_marker(): assert len(data_frames) == 4 assert any('"usage"' in frame and '"completion_tokens":188' in frame.replace(" ", "") for frame in data_frames) assert frames[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_config_field_update_rejects_mock_testing_flag(): + """The mock-testing opt-in is deliberately absent from + ``ConfigGeneralSettings`` so that ``/config/field/update`` refuses it. If + someone later adds the field for tidiness, this test fails and tells them + they have just opened an API write path into a config-file-only setting.""" + from fastapi import HTTPException + + from litellm.proxy._types import ConfigFieldUpdate + from litellm.proxy.proxy_server import update_config_general_settings + from litellm.proxy.route_llm_request import MOCK_TESTING_CONFIG_KEY + + admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-test", + ) + + with patch.object(proxy_server_module, "prisma_client", MagicMock()): + with pytest.raises(HTTPException) as exc_info: + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name=MOCK_TESTING_CONFIG_KEY, + field_value=True, + config_type="general_settings", + ), + user_api_key_dict=admin, + ) + + assert exc_info.value.status_code == 400 + + +def test_config_update_body_drops_mock_testing_flag(): + """``/config/update`` parses its body as ``ConfigYAML``, whose + ``general_settings`` is a ``ConfigGeneralSettings``. Undeclared keys are + dropped on parse, so the flag never reaches the DB by that route either.""" + from litellm.proxy._types import ConfigYAML + from litellm.proxy.route_llm_request import MOCK_TESTING_CONFIG_KEY + + parsed = ConfigYAML.model_validate({"general_settings": {MOCK_TESTING_CONFIG_KEY: True}}) + + assert parsed.general_settings is not None + assert MOCK_TESTING_CONFIG_KEY not in parsed.general_settings.model_dump(exclude_none=True) + + +def test_startup_warns_when_mock_testing_params_enabled(caplog): + """Enabling the opt-in must announce itself, naming every param it + unlocks — the config key says ``mock_testing`` but the gate also covers + ``mock_timeout`` and ``mock_delay``, so coverage cannot be inferred from + the name alone.""" + import logging + + from litellm.proxy.proxy_server import ProxyStartupEvent + from litellm.proxy.route_llm_request import ( + GATED_MOCK_PARAM_NAMES, + MOCK_TESTING_CONFIG_KEY, + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + ProxyStartupEvent._warn_if_mock_testing_params_enabled( + general_settings={MOCK_TESTING_CONFIG_KEY: True} + ) + + assert MOCK_TESTING_CONFIG_KEY in caplog.text + for param_name in GATED_MOCK_PARAM_NAMES: + assert param_name in caplog.text + + +def test_startup_is_silent_when_mock_testing_params_disabled(caplog): + """A proxy that never set the opt-in must not emit the warning.""" + import logging + + from litellm.proxy.proxy_server import ProxyStartupEvent + from litellm.proxy.route_llm_request import MOCK_TESTING_CONFIG_KEY + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + ProxyStartupEvent._warn_if_mock_testing_params_enabled(general_settings={}) + + assert MOCK_TESTING_CONFIG_KEY not in caplog.text diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 93b3ef1cce8..e74693a4f6c 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -8,6 +8,8 @@ sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to from unittest.mock import MagicMock +from fastapi import HTTPException + from litellm.proxy.route_llm_request import ProxyModelNotFoundError, route_request @@ -471,38 +473,147 @@ async def test_route_request_with_router_settings_override_preserves_existing(): assert call_kwargs["timeout"] == 30 -def test_mock_testing_kwarg_names_matches_dataclass(): - """``_MOCK_TESTING_KWARG_NAMES`` is hardcoded to avoid a cyclic import - against ``litellm.types.router``. This test guards against drift — - if a new ``mock_testing_*`` field is added to ``MockRouterTestingParams`` - the strip list must be updated to keep covering it.""" +def test_gated_mock_params_cover_mock_router_testing_params(): + """``GATED_MOCK_PARAM_NAMES`` is hardcoded to avoid a cyclic import + against ``litellm.types.router``. This test guards against drift — if a + new ``mock_testing_*`` field is added to ``MockRouterTestingParams`` the + gate must be updated to keep covering it. The gate is a superset: it also + covers params consumed outside that dataclass.""" from dataclasses import fields - from litellm.proxy.route_llm_request import _MOCK_TESTING_KWARG_NAMES + from litellm.proxy.route_llm_request import GATED_MOCK_PARAM_NAMES from litellm.types.router import MockRouterTestingParams - assert set(_MOCK_TESTING_KWARG_NAMES) == {f.name for f in fields(MockRouterTestingParams)} + assert {f.name for f in fields(MockRouterTestingParams)} <= set(GATED_MOCK_PARAM_NAMES) + assert {"mock_testing_rate_limit_error", "mock_timeout", "mock_delay"} <= set(GATED_MOCK_PARAM_NAMES) -@pytest.mark.asyncio @pytest.mark.parametrize( - "mock_flag", + "mock_param", [ "mock_testing_fallbacks", "mock_testing_context_fallbacks", "mock_testing_content_policy_fallbacks", + "mock_testing_rate_limit_error", + "mock_timeout", + "mock_delay", ], ) -async def test_route_request_strips_mock_testing_flags(mock_flag): - """VERIA-44: router-internal testing flags must not survive a - user-supplied request body. Without this strip, an attacker can - combine ``mock_testing_fallbacks=true`` with an unauthorized fallback - in ``router_settings_override`` to deterministically execute requests - against restricted models.""" +def test_mock_params_rejected_when_not_allowed(mock_param): + """Every gated param must be rejected by name when the proxy has not + opted in, and the error must point the caller at the config key.""" + from litellm.proxy.route_llm_request import ( + MOCK_TESTING_CONFIG_KEY, + raise_if_mock_testing_params_disallowed, + ) + + data = {"model": "gpt-3.5-turbo", mock_param: True} + + with pytest.raises(HTTPException) as exc_info: + raise_if_mock_testing_params_disallowed(data, allowed=False) + + assert exc_info.value.status_code == 400 + error_message = exc_info.value.detail["error"] + assert mock_param in error_message + assert MOCK_TESTING_CONFIG_KEY in error_message + + +@pytest.mark.parametrize( + "mock_param", + [ + "mock_testing_fallbacks", + "mock_testing_context_fallbacks", + "mock_testing_content_policy_fallbacks", + "mock_testing_rate_limit_error", + "mock_timeout", + "mock_delay", + ], +) +def test_mock_params_pass_through_when_allowed(mock_param): + """With the opt-in set, gated params must survive untouched — a gate that + rejects correctly but strips anyway would leave the feature unusable.""" + from litellm.proxy.route_llm_request import raise_if_mock_testing_params_disallowed + + data = {"model": "gpt-3.5-turbo", mock_param: True} + + raise_if_mock_testing_params_disallowed(data, allowed=True) + + assert data[mock_param] is True + + +def test_mock_param_gate_reports_every_param_present(): + """A request carrying several gated params must name all of them, so a + caller fixing one is not surprised by the next.""" + from litellm.proxy.route_llm_request import raise_if_mock_testing_params_disallowed + + data = { + "model": "gpt-3.5-turbo", + "mock_testing_fallbacks": True, + "mock_delay": 30, + } + + with pytest.raises(HTTPException) as exc_info: + raise_if_mock_testing_params_disallowed(data, allowed=False) + + error_message = exc_info.value.detail["error"] + assert "mock_testing_fallbacks" in error_message + assert "mock_delay" in error_message + + +def test_ordinary_request_is_not_rejected_by_the_mock_param_gate(): + """The gate must not fire on a request that carries no gated param.""" + from litellm.proxy.route_llm_request import raise_if_mock_testing_params_disallowed + data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}], - mock_flag: True, + "mock_response": "hi", + } + + raise_if_mock_testing_params_disallowed(data, allowed=False) + + +@pytest.mark.asyncio +async def test_route_request_rejects_mock_params_by_default(monkeypatch): + """End-to-end through ``route_request``: with no opt-in configured the + request is rejected before it ever reaches the router.""" + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) + + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + "mock_testing_fallbacks": True, + } + llm_router = MagicMock() + + with pytest.raises(HTTPException) as exc_info: + await route_request(data, llm_router, None, "acompletion") + + assert exc_info.value.status_code == 400 + llm_router.acompletion.assert_not_called() + + +@pytest.mark.asyncio +async def test_route_request_forwards_mock_params_when_opted_in(monkeypatch): + """End-to-end through ``route_request``: with the opt-in set the param + reaches the router, which is what makes a fallback drill possible.""" + import litellm.proxy.proxy_server as proxy_server + + from litellm.proxy.route_llm_request import MOCK_TESTING_CONFIG_KEY + + monkeypatch.setattr( + proxy_server, + "general_settings", + {MOCK_TESTING_CONFIG_KEY: True}, + raising=False, + ) + + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + "mock_testing_fallbacks": True, } llm_router = MagicMock() llm_router.acompletion.return_value = "ok" @@ -510,10 +621,7 @@ async def test_route_request_strips_mock_testing_flags(mock_flag): await route_request(data, llm_router, None, "acompletion") call_kwargs = llm_router.acompletion.call_args[1] - assert mock_flag not in call_kwargs - # The flag is also gone from the original data dict so any subsequent - # processing (e.g. logging) doesn't see it either. - assert mock_flag not in data + assert call_kwargs["mock_testing_fallbacks"] is True @pytest.mark.parametrize("route_type", ["agenerate_content", "agenerate_content_stream"]) diff --git a/tests/test_litellm/proxy/test_update_llm_router_resilience.py b/tests/test_litellm/proxy/test_update_llm_router_resilience.py index fd0df4805e6..a7dc9c1783e 100644 --- a/tests/test_litellm/proxy/test_update_llm_router_resilience.py +++ b/tests/test_litellm/proxy/test_update_llm_router_resilience.py @@ -115,8 +115,10 @@ class TestDeleteDeploymentResilience: """Test _delete_deployment handles get_config failures gracefully.""" @pytest.mark.asyncio - async def test_returns_zero_when_get_config_times_out(self): - """Should return 0 (no deletions) when get_config fails, not raise.""" + async def test_returns_none_when_get_config_times_out(self): + """Should return None (no reconcile ran, desired set unknown) when get_config + fails, not raise. A caller judging its own reload must not read that as "the db + wants nothing" and blame the reload for every model it serves.""" proxy_config = ProxyConfig() db_models = [_make_db_model("gpt-5.1", "db-id-1")] @@ -136,8 +138,8 @@ class TestDeleteDeploymentResilience: ): result = await proxy_config._delete_deployment(db_models=db_models) - # Should safely return 0 instead of raising - assert result == 0 + # Should safely return None instead of raising + assert result is None # Should NOT have deleted any deployments mock_router.delete_deployment.assert_not_called() @@ -175,5 +177,8 @@ class TestDeleteDeploymentResilience: result = await proxy_config._delete_deployment(db_models=db_models) # "stale-id" should have been deleted (not in db_models or config) - assert result == 1 mock_router.delete_deployment.assert_called_once_with(id="stale-id") + assert result == frozenset({"db-id-1", "config-id-1"}), ( + "the returned set must be what the db + config still want, so a caller can " + f"tell that eviction apart from a deployment that went missing; got {result}" + ) diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index c923b722991..8253308f393 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -1905,6 +1905,28 @@ class TestUserRepositoryExtended: ) assert updated.user_email == "new@example.com" + @pytest.mark.asyncio + async def test_find_by_id_decodes_only_the_json_encoded_columns(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-json"] = { + "user_id": "user-json", + "user_email": "json@example.com", + "user_role": '{"not": "json"}', + "teams": [], + "models": [], + "metadata": '{"department": "engineering"}', + "model_spend": '{"gpt-4": 10.5}', + "model_max_budget": '{"gpt-4": 100.0}', + } + + user = await repo.find_by_id("user-json") + + assert user is not None + assert user.metadata == {"department": "engineering"} + assert user.model_spend == {"gpt-4": 10.5} + assert user.model_max_budget == {"gpt-4": 100.0} + assert user.user_email == "json@example.com" + assert user.user_role == '{"not": "json"}' + class TestProjectRepositoryExtended: @pytest.fixture diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py index 13edf6d1a95..53d672fc4a8 100644 --- a/tests/test_litellm/test_check_type_discipline.py +++ b/tests/test_litellm/test_check_type_discipline.py @@ -152,6 +152,28 @@ def test_qualified_collections_constructors_still_count(tmp_path): assert "LIT002" in _codes(tmp_path, "import collections\nm = collections.defaultdict(list)\n") +def test_value_frozen_by_wrapper_is_exempt(tmp_path): + assert "LIT002" not in _codes(tmp_path, "from types import MappingProxyType\nm = MappingProxyType({'a': 1})\n") + assert "LIT002" not in _codes(tmp_path, "import types\nm = types.MappingProxyType({'a': 1})\n") + assert "LIT002" not in _codes(tmp_path, "from types import MappingProxyType\nm = MappingProxyType(dict(a=1))\n") + assert "LIT002" not in _codes(tmp_path, "f = frozenset({1, 2})\n") + assert "LIT002" not in _codes(tmp_path, "t = tuple([1, 2])\n") + + +def test_same_named_method_does_not_exempt_its_argument(tmp_path): + assert "LIT002" in _codes(tmp_path, "t = obj.tuple([1, 2])\n") + assert "LIT002" in _codes(tmp_path, "f = obj.frozenset({1, 2})\n") + assert "LIT002" in _codes(tmp_path, "m = obj.MappingProxyType({'a': 1})\n") + + +def test_mutable_nested_inside_frozen_wrapper_still_counts(tmp_path): + assert "LIT002" in _codes(tmp_path, "from types import MappingProxyType\nm = MappingProxyType({'a': []})\n") + + +def test_unfrozen_literal_still_counts(tmp_path): + assert "LIT002" in _codes(tmp_path, "from types import MappingProxyType\nd = {'a': 1}\nm = MappingProxyType(d)\n") + + def test_mutable_ok_with_reason_suppresses_both_rules(tmp_path): codes = _codes(tmp_path, "x: dict[str, int] = {} # mutable-ok: in-place buffer mutated hot path\n") assert "LIT001" not in codes diff --git a/tests/test_litellm/test_model_block_unblock.py b/tests/test_litellm/test_model_block_unblock.py index dc0098e405e..ff66bedf0dc 100644 --- a/tests/test_litellm/test_model_block_unblock.py +++ b/tests/test_litellm/test_model_block_unblock.py @@ -36,7 +36,7 @@ def _setup_model_block_mocks(monkeypatch, *, updated_blocked: bool): mock_router = MagicMock() mock_router.get_model_ids.return_value = [model_id] - mock_clear_cache = AsyncMock(return_value=True) + mock_clear_cache = AsyncMock(return_value=None) mock_audit_log = AsyncMock(return_value=None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index c9a1b59cc06..05499e83c42 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23253 + "limit": 23250 }, "LIT002": { - "limit": 27427 + "limit": 27277 }, "LIT003": { "limit": 292 @@ -24,6 +24,6 @@ "limit": 1004 }, "LIT009": { - "limit": 2474 + "limit": 2473 } } diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 4cb3ebe7e16..8c6d940cfea 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -126,11 +126,6 @@ "count": 2 } }, - "src/app/(dashboard)/budgets/_components/budget_panel.test.tsx": { - "unused-imports/no-unused-imports": { - "count": 2 - } - }, "src/app/(dashboard)/budgets/_components/budget_panel.tsx": { "local/filename-pascal-case": { "count": 1 @@ -794,6 +789,11 @@ "count": 1 } }, + "src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.test.tsx": { "unused-imports/no-unused-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx index 2b97bcbc072..f78ba34770d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx @@ -1,22 +1,63 @@ -import { screen, within } from "@testing-library/react"; +import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders } from "@/../tests/test-utils"; +import { renderWithProviders, testQueryClient } from "@/../tests/test-utils"; import BudgetTable from "./BudgetTable"; -import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import type { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import type { ResourceListResult } from "@/app/(dashboard)/hooks/common/useResourceList"; +import { ApiError } from "@/lib/http/client"; + +const { copyToClipboardMock } = vi.hoisted(() => ({ copyToClipboardMock: vi.fn() })); + +vi.mock("@/utils/dataUtils", async (importOriginal) => ({ + ...(await importOriginal()), + copyToClipboard: copyToClipboardMock, +})); const makeBudget = (overrides: Partial = {}): budgetItem => ({ budget_id: "budget-1", max_budget: 100, + soft_budget: null, tpm_limit: 1000, rpm_limit: 10, + budget_duration: "30d", + budget_reset_at: null, + created_at: "2024-01-01T00:00:00Z", updated_at: "2024-01-01T00:00:00Z", ...overrides, }); -const defaultProps = { - budgets: [makeBudget()], +const makeList = (overrides: Partial> = {}): ResourceListResult => ({ + rows: [makeBudget()], + rowCount: 1, isLoading: false, + isFetching: false, + error: null, + refetch: vi.fn(), + sorting: [{ id: "created_at", desc: true }], + onSortingChange: vi.fn(), + pagination: { pageIndex: 0, pageSize: 50 }, + onPaginationChange: vi.fn(), + columnFilters: [], + onColumnFiltersChange: vi.fn(), + searchValue: "", + onSearchChange: vi.fn(), + ...overrides, +}); + +const FORBIDDEN_PROBLEM = { + type: "about:blank", + title: "Forbidden", + status: 403, + detail: "Only proxy admins can view budgets", +}; + +const showColumn = async (user: ReturnType, columnId: string) => { + await user.click(screen.getByTestId("view-options-trigger")); + await user.click(await screen.findByTestId(`view-option-${columnId}`)); +}; + +const defaultProps = { canModify: true, onEditClick: vi.fn(), onDeleteClick: vi.fn(), @@ -25,72 +66,151 @@ const defaultProps = { describe("BudgetTable", () => { beforeEach(() => { vi.clearAllMocks(); + testQueryClient.clear(); }); it("should display budget information", () => { - renderWithProviders(); + renderWithProviders(); expect(screen.getByText("budget-1")).toBeInTheDocument(); expect(screen.getByText("$100.00")).toBeInTheDocument(); expect(screen.getByText("1000")).toBeInTheDocument(); expect(screen.getByText("10")).toBeInTheDocument(); }); - it("should render the budget id without a fixed character-count clamp", () => { + it("should open on the four columns the page has always shown, with reset and created off", () => { + renderWithProviders(); + const headers = screen.getAllByRole("columnheader").map((header) => header.textContent); + expect(headers).toEqual(expect.arrayContaining(["Budget ID", "Max Budget", "TPM", "RPM"])); + expect(headers).not.toContain("Reset"); + expect(headers).not.toContain("Created"); + }); + + it("should render the reset column with the friendly duration label once it is turned on", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await showColumn(user, "budget_duration"); + expect(screen.getByText("monthly")).toBeInTheDocument(); + }); + + it("should render 'Not set' when a budget has no reset duration", async () => { + const user = userEvent.setup(); + const list = makeList({ rows: [makeBudget({ budget_duration: null })] }); + renderWithProviders(); + await showColumn(user, "budget_duration"); + expect(screen.getByText("Not set")).toBeInTheDocument(); + }); + + it("should render the budget id in full, with no truncation", () => { const budgetId = "ecc1869c-6231-4380-a56d-1a0be457477d"; - renderWithProviders(); + const list = makeList({ rows: [makeBudget({ budget_id: budgetId })] }); + renderWithProviders(); const idCell = screen.getByText(budgetId); + expect(idCell.className).not.toContain("truncate"); expect(idCell.className).not.toMatch(/max-w-\[\d+(ch|rem|px)\]/); - expect(idCell.className).toContain("max-w-full"); - expect(idCell.className).toContain("truncate"); + }); + + it("should keep the budget id on a single line", () => { + const budgetId = "ecc1869c-6231-4380-a56d-1a0be457477d"; + const list = makeList({ rows: [makeBudget({ budget_id: budgetId })] }); + renderWithProviders(); + expect(screen.getByText(budgetId).className).toContain("whitespace-nowrap"); + }); + + it("should copy the budget id from the cell's copy button", async () => { + const user = userEvent.setup(); + const budgetId = "ecc1869c-6231-4380-a56d-1a0be457477d"; + const list = makeList({ rows: [makeBudget({ budget_id: budgetId })] }); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: "Copy ID" })); + expect(copyToClipboardMock).toHaveBeenCalledWith(budgetId); + }); + + it("should offer sorting on every backend-sortable column", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await showColumn(user, "created_at"); + for (const field of ["budget_id", "max_budget", "tpm_limit", "rpm_limit", "created_at"]) { + expect(screen.getByTestId(`sort-header-${field}`)).toBeInTheDocument(); + } + }); + + it("should not make the reset column sortable", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await showColumn(user, "budget_duration"); + const headers = screen.getAllByRole("columnheader").map((header) => header.textContent); + expect(headers).toContain("Reset"); + expect(screen.queryByTestId("sort-header-budget_duration")).not.toBeInTheDocument(); + }); + + it("should ask the list for a new sort when a sortable header is clicked", async () => { + const user = userEvent.setup(); + const onSortingChange = vi.fn(); + renderWithProviders(); + await user.click(screen.getByTestId("sort-header-max_budget")); + expect(onSortingChange).toHaveBeenCalled(); }); it("should show n/a for missing rate limits and Unlimited for a missing max budget", () => { - renderWithProviders( - , - ); + const list = makeList({ rows: [makeBudget({ max_budget: null, tpm_limit: null, rpm_limit: null })] }); + renderWithProviders(); expect(screen.getAllByText("n/a")).toHaveLength(2); expect(screen.getByText("Unlimited")).toBeInTheDocument(); }); - it("should sort budgets by updated_at descending", () => { - const budgets = [ - makeBudget({ budget_id: "budget-old", updated_at: "2024-01-01T00:00:00Z" }), - makeBudget({ budget_id: "budget-new", updated_at: "2024-06-01T00:00:00Z" }), - ]; - renderWithProviders(); - const rows = screen.getAllByRole("row").slice(1); - expect(within(rows[0]).getByText("budget-new")).toBeInTheDocument(); - expect(within(rows[1]).getByText("budget-old")).toBeInTheDocument(); - }); - it("should call onEditClick from the actions menu", async () => { const user = userEvent.setup(); - renderWithProviders(); + const list = makeList(); + renderWithProviders(); await user.click(screen.getByTestId("budget-actions-budget-1")); await user.click(await screen.findByTestId("budget-action-edit")); - expect(defaultProps.onEditClick).toHaveBeenCalledWith(defaultProps.budgets[0]); + expect(defaultProps.onEditClick).toHaveBeenCalledWith(list.rows[0]); }); it("should call onDeleteClick from the actions menu", async () => { const user = userEvent.setup(); - renderWithProviders(); + const list = makeList(); + renderWithProviders(); await user.click(screen.getByTestId("budget-actions-budget-1")); await user.click(await screen.findByTestId("budget-action-delete")); - expect(defaultProps.onDeleteClick).toHaveBeenCalledWith(defaultProps.budgets[0]); + expect(defaultProps.onDeleteClick).toHaveBeenCalledWith(list.rows[0]); }); it("should not render the actions menu when the user cannot modify budgets", () => { - renderWithProviders(); + renderWithProviders(); expect(screen.queryByTestId("budget-actions-budget-1")).not.toBeInTheDocument(); }); it("should show skeleton rows when loading", () => { - renderWithProviders(); + renderWithProviders(); expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); }); it("should show the empty state when there are no budgets", () => { - renderWithProviders(); + renderWithProviders(); expect(screen.getByText("No budgets yet")).toBeInTheDocument(); }); + + it("should tell the user their search matched nothing rather than that no budgets exist", () => { + const list = makeList({ rows: [], rowCount: 0, searchValue: "nope" }); + renderWithProviders(); + expect(screen.getByText("No matching budgets")).toBeInTheDocument(); + }); + + it("should render an access-denied state for a 403 instead of an empty table", () => { + const error = new ApiError("Only proxy admins can view budgets", 403, FORBIDDEN_PROBLEM); + const list = makeList({ rows: [], rowCount: 0, error }); + const { container } = renderWithProviders(); + expect(screen.getByText("You do not have access to budgets")).toBeInTheDocument(); + expect(screen.queryByText("No budgets yet")).not.toBeInTheDocument(); + expect(container.querySelector(".lucide-shield-alert")).not.toBeNull(); + }); + + it("should surface the problem detail for a non-403 failure", () => { + const error = new ApiError("budget store unavailable", 500, null); + const list = makeList({ rows: [], rowCount: 0, error }); + renderWithProviders(); + expect(screen.getByText("Could not load budgets")).toBeInTheDocument(); + expect(screen.getByText("budget store unavailable")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx index 4bc06425f80..80872c58d28 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx @@ -1,55 +1,271 @@ "use client"; -import { Inbox } from "lucide-react"; -import React, { useMemo } from "react"; +import { Inbox, ShieldAlert } from "lucide-react"; +import React, { useMemo, useState } from "react"; -import { DataTable } from "@/components/shared/DataTable"; -import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import { + BUDGET_DURATION_FILTER_OPTIONS, + BUDGET_DURATION_UNSET, + type CreatedAtFilterValue, + type MaxBudgetFilterValue, +} from "@/app/(dashboard)/hooks/budgets/budgetFilters"; +import type { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import type { ResourceListResult } from "@/app/(dashboard)/hooks/common/useResourceList"; +import { + DataTable, + DataTableFilterDrawer, + DataTableFilterField, + DataTableToolbar, + type FilterDraft, +} from "@/components/shared/DataTable"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { ApiError } from "@/lib/http/client"; -import { getBudgetTableColumns } from "./BudgetTableColumns"; +import { BUDGET_TABLE_HIDDEN_COLUMNS, getBudgetTableColumns } from "./BudgetTableColumns"; interface BudgetTableProps { - budgets: budgetItem[]; - isLoading: boolean; + list: ResourceListResult; canModify: boolean; onEditClick: (budget: budgetItem) => void; onDeleteClick: (budget: budgetItem) => void; } -function EmptyState() { +const PAGE_SIZE_OPTIONS = [25, 50, 100]; + +const FILTER_LABELS: Record = { + budget_duration: "Reset", + max_budget: "Max Budget", + created_at: "Created", +}; + +const durationLabel = (value: string): string => + BUDGET_DURATION_FILTER_OPTIONS.find((option) => option.value === value)?.label ?? value; + +const formatFilterValue = (columnId: string, value: unknown): string => { + if (columnId === "budget_duration") { + return (Array.isArray(value) ? value : []).map((entry) => durationLabel(String(entry))).join(", "); + } + if (columnId === "max_budget") { + const { min, max, unlimitedOnly } = (value ?? {}) as MaxBudgetFilterValue; + return unlimitedOnly === true ? "Unlimited only" : `${min ? `$${min}` : "any"} to ${max ? `$${max}` : "any"}`; + } + if (columnId === "created_at") { + const { from, to } = (value ?? {}) as CreatedAtFilterValue; + return `${from || "any"} to ${to || "any"}`; + } + return String(value); +}; + +/** The drawer keeps any non-empty object as an active filter, so collapse a blank draft to nothing. */ +const normalizeMaxBudget = (draft: MaxBudgetFilterValue): MaxBudgetFilterValue | undefined => { + if (draft.unlimitedOnly === true) { + return { unlimitedOnly: true }; + } + const min = draft.min?.trim() ?? ""; + const max = draft.max?.trim() ?? ""; + if (min === "" && max === "") { + return undefined; + } + return { ...(min === "" ? {} : { min }), ...(max === "" ? {} : { max }) }; +}; + +const normalizeCreatedAt = (draft: CreatedAtFilterValue): CreatedAtFilterValue | undefined => { + const from = draft.from ?? ""; + const to = draft.to ?? ""; + if (from === "" && to === "") { + return undefined; + } + return { ...(from === "" ? {} : { from }), ...(to === "" ? {} : { to }) }; +}; + +function EmptyState({ hasQuery }: { hasQuery: boolean }) { return (
-
No budgets yet
+
{hasQuery ? "No matching budgets" : "No budgets yet"}
- Create a budget to set spend, TPM and RPM limits for customers. + {hasQuery + ? "No budget matches your search or filters." + : "Create a budget to set spend, TPM and RPM limits for customers."}
); } -const BudgetTable: React.FC = ({ budgets, isLoading, canModify, onEditClick, onDeleteClick }) => { - const rows = useMemo( - () => [...budgets].sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()), - [budgets], +function ErrorState({ error }: { error: Error }) { + const forbidden = error instanceof ApiError && error.status === 403; + return ( +
+
+ +
+
+ {forbidden ? "You do not have access to budgets" : "Could not load budgets"} +
+
+ {forbidden ? "Ask a proxy admin to grant you the admin viewer role." : error.message} +
+
); +} + +/** "Not set" and the concrete durations are exclusive; see serializeBudgetFilters for why. */ +function DurationFilter({ selected, onChange }: { selected: string[]; onChange: (selected: string[]) => void }) { + const toggle = (value: string, checked: boolean): void => { + if (!checked) { + onChange(selected.filter((entry) => entry !== value)); + return; + } + const kept = value === BUDGET_DURATION_UNSET ? [] : selected.filter((entry) => entry !== BUDGET_DURATION_UNSET); + onChange([...kept, value]); + }; + + return ( +
+ {BUDGET_DURATION_FILTER_OPTIONS.map((option) => ( + + ))} +
+ ); +} + +function BudgetFilterFields({ get, set }: FilterDraft) { + const maxBudget = (get("max_budget") as MaxBudgetFilterValue | undefined) ?? {}; + const created = (get("created_at") as CreatedAtFilterValue | undefined) ?? {}; + const unlimitedOnly = maxBudget.unlimitedOnly === true; + + return ( + <> + + set("budget_duration", selected)} + /> + + +
+ set("max_budget", normalizeMaxBudget({ ...maxBudget, min: event.target.value }))} + placeholder="Min" + aria-label="Minimum max budget" + data-testid="budget-filter-max-budget-min" + /> + set("max_budget", normalizeMaxBudget({ ...maxBudget, max: event.target.value }))} + placeholder="Max" + aria-label="Maximum max budget" + data-testid="budget-filter-max-budget-max" + /> +
+ + + +
+ set("created_at", normalizeCreatedAt({ ...created, from: event.target.value }))} + aria-label="Created from" + data-testid="budget-filter-created-from" + /> + set("created_at", normalizeCreatedAt({ ...created, to: event.target.value }))} + aria-label="Created to" + data-testid="budget-filter-created-to" + /> +
+ + + ); +} + +const BudgetTable: React.FC = ({ list, canModify, onEditClick, onDeleteClick }) => { + const [filtersOpen, setFiltersOpen] = useState(false); const columns = useMemo( () => getBudgetTableColumns({ canModify, onEditClick, onDeleteClick }), [canModify, onEditClick, onDeleteClick], ); + const hasQuery = list.searchValue.trim() !== "" || list.columnFilters.length > 0; + const emptyMessage = list.error === null ? : ; + return ( budget.budget_id || String(index)} - isLoading={isLoading} + defaultColumnVisibility={BUDGET_TABLE_HIDDEN_COLUMNS} + fillHeight + sortingMode="server" + sorting={list.sorting} + onSortingChange={list.onSortingChange} + paginationMode="server" + pagination={list.pagination} + onPaginationChange={list.onPaginationChange} + rowCount={list.rowCount} + pageSizeOptions={PAGE_SIZE_OPTIONS} + filterMode="server" + columnFilters={list.columnFilters} + onColumnFiltersChange={list.onColumnFiltersChange} + isLoading={list.isLoading} loadingMessage="Loading budgets…" - noDataMessage={} + noDataMessage={emptyMessage} size="compact" + toolbar={(table) => ( + <> + setFiltersOpen(true)} + onRefresh={list.refetch} + isRefreshing={list.isFetching} + filterLabels={FILTER_LABELS} + formatFilterValue={formatFilterValue} + /> + + {(draft) => } + + + )} /> ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx index e3fbc9dba08..e5cd9043492 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx @@ -1,11 +1,13 @@ "use client"; -import { ColumnDef } from "@tanstack/react-table"; +import { ColumnDef, FilterFn } from "@tanstack/react-table"; import { MoreHorizontal, Pencil, Trash2 } from "lucide-react"; -import { IdCell, MoneyCell } from "@/components/shared/table_cells"; -import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells"; +import type { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { buttonVariants } from "@/components/ui/button"; +import { getBudgetDurationLabel } from "@/components/common_components/budget_duration_dropdown"; import { DropdownMenu, DropdownMenuContent, @@ -15,13 +17,29 @@ import { } from "@/components/ui/dropdown-menu"; import { cn } from "@/lib/cva.config"; -function RateLimitCell({ value }: { value: number | null }) { +/** + * Filtering happens on the server, so this never runs as a predicate. It exists to override + * TanStack's auto-remove heuristic, which infers a filter shape from the column's first cell + * and silently discards a filter whose value is not that shape (a range object on a numeric + * column, for instance). + */ +const serverFilter: FilterFn = () => true; +serverFilter.autoRemove = () => false; + +function RateLimitCell({ value }: { value: number | null | undefined }) { if (value == null) { return n/a; } return {value}; } +function BudgetDurationCell({ value }: { value: string | null | undefined }) { + if (!value) { + return Not set; + } + return {getBudgetDurationLabel(value)}; +} + interface BudgetRowActionsProps { budget: budgetItem; onEditClick: (budget: budgetItem) => void; @@ -57,6 +75,12 @@ function BudgetRowActions({ budget, onEditClick, onDeleteClick }: BudgetRowActio ); } +/** Off by default so the table opens on the four columns it has always shown; the Columns menu turns them on. */ +export const BUDGET_TABLE_HIDDEN_COLUMNS: Record = { + budget_duration: false, + created_at: false, +}; + interface BudgetTableColumnsDeps { canModify: boolean; onEditClick: (budget: budgetItem) => void; @@ -72,38 +96,56 @@ export const getBudgetTableColumns = ({ id: "budget_id", accessorKey: "budget_id", meta: { title: "Budget ID" }, - header: "Budget ID", - size: 220, - enableSorting: false, - cell: ({ row }) => , + header: ({ column }) => , + cell: ({ row }) => ( + + ), }, { id: "max_budget", accessorKey: "max_budget", + filterFn: serverFilter, meta: { title: "Max Budget", numeric: true }, - header: "Max Budget", + header: ({ column }) => , size: 120, - enableSorting: false, cell: ({ row }) => , }, { id: "tpm_limit", accessorKey: "tpm_limit", meta: { title: "TPM", numeric: true }, - header: "TPM", + header: ({ column }) => , size: 100, - enableSorting: false, cell: ({ row }) => , }, { id: "rpm_limit", accessorKey: "rpm_limit", meta: { title: "RPM", numeric: true }, - header: "RPM", + header: ({ column }) => , size: 100, - enableSorting: false, cell: ({ row }) => , }, + { + id: "budget_duration", + accessorKey: "budget_duration", + filterFn: serverFilter, + meta: { title: "Reset" }, + // "7d"/"30d" sort lexicographically, not chronologically, so the route does not offer it. + enableSorting: false, + header: ({ column }) => , + size: 110, + cell: ({ row }) => , + }, + { + id: "created_at", + accessorKey: "created_at", + filterFn: serverFilter, + meta: { title: "Created" }, + header: ({ column }) => , + size: 160, + cell: ({ row }) => , + }, ...(canModify ? [ { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx index 392616f1935..46f72cd8886 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx @@ -1,217 +1,254 @@ -import { fireEvent, render, waitFor, screen } from "@testing-library/react"; -import { act } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { afterEach, describe, expect, it, vi } from "vitest"; +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { ApiError } from "@/lib/http/client"; + import BudgetPanel from "./budget_panel"; -const mockBudgets = [ - { - budget_id: "budget-1", - max_budget: 100, - rpm_limit: 10, - tpm_limit: 1000, - updated_at: "2024-01-01T00:00:00Z", - }, -]; - -vi.mock("@/app/(dashboard)/hooks/budgets/useBudgets", () => ({ - useBudgets: vi.fn().mockReturnValue({ data: [], isLoading: false }), - useDeleteBudget: vi.fn().mockReturnValue({ mutateAsync: vi.fn(), isPending: false }), - useCreateBudget: vi.fn().mockReturnValue({ mutateAsync: vi.fn() }), - useUpdateBudget: vi.fn().mockReturnValue({ mutateAsync: vi.fn() }), +const { getMock, budgetDeleteMock } = vi.hoisted(() => ({ + getMock: vi.fn(), + budgetDeleteMock: vi.fn(), })); -import { - useBudgets, - useDeleteBudget, - useCreateBudget, - useUpdateBudget, -} from "@/app/(dashboard)/hooks/budgets/useBudgets"; +vi.mock("@/components/networking", () => ({ + apiClient: { get: getMock }, + budgetCreateCall: vi.fn(), + budgetUpdateCall: vi.fn(), + budgetDeleteCall: budgetDeleteMock, + getProxyBaseUrl: () => "", +})); -const createQueryClient = () => - new QueryClient({ - defaultOptions: { queries: { retry: false, gcTime: 0 } }, - }); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "sk-test", userRole: "Admin", userId: "u1" }), +})); -function renderWithProviders(ui: React.ReactElement) { - const qc = createQueryClient(); - return render({ui}); +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { success: vi.fn(), info: vi.fn(), fromBackend: vi.fn() }, +})); + +interface BudgetSeed { + budget_id: string; + max_budget: number | null; + budget_duration: string | null; } +const budgetRow = (seed: BudgetSeed) => ({ + soft_budget: null, + tpm_limit: 1000, + rpm_limit: 10, + budget_reset_at: null, + created_at: "2026-01-01T00:00:00Z", + updated_at: "2026-01-01T00:00:00Z", + ...seed, +}); + +const FORBIDDEN_PROBLEM = { + type: "about:blank", + title: "Forbidden", + status: 403, + detail: "Only proxy admins can view budgets", +}; + +const DEFAULT_ROWS = [ + budgetRow({ budget_id: "ecc1869c-6231-4380-a56d-1a0be457477d", max_budget: 100, budget_duration: "30d" }), +]; + +const respondWith = (rows: ReturnType[], totalCount: number) => { + getMock.mockResolvedValue({ + data: rows, + meta: { total_count: totalCount, page: 1, page_size: 50, total_pages: Math.ceil(totalCount / 50) }, + }); +}; + +type QueryRecord = Record; + +const queries = (): QueryRecord[] => getMock.mock.calls.map((call) => (call[1] as { query: QueryRecord }).query); +const lastQuery = (): QueryRecord => queries()[queries().length - 1]; +const paths = (): string[] => getMock.mock.calls.map((call) => String(call[0])); + +const renderPanel = () => { + const client = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + return render( + + + , + ); +}; + +const openFilters = async (user: ReturnType) => { + await user.click(screen.getByTestId("datatable-filters-trigger")); + await screen.findByTestId("filter-drawer-body"); +}; + describe("Budget Panel", () => { - afterEach(() => { + beforeEach(() => { vi.clearAllMocks(); + respondWith(DEFAULT_ROWS, 1); }); - it("should render the budget panel and load budgets", async () => { - vi.mocked(useBudgets).mockReturnValue({ - data: mockBudgets, - isLoading: false, - } as any); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("Create a budget to assign to customers.")).toBeInTheDocument(); - expect(screen.getByText("budget-1")).toBeInTheDocument(); - }); + it("loads the first page of budgets, newest first", async () => { + renderPanel(); + await waitFor(() => expect(getMock).toHaveBeenCalled()); + expect(paths()[0]).toBe("/management/v1/budgets"); + expect(queries()[0]).toEqual({ page: 1, page_size: 50, sort: "-created_at" }); + expect(await screen.findByText("ecc1869c-6231-4380-a56d-1a0be457477d")).toBeInTheDocument(); }); - it("should open delete modal from the actions menu", async () => { + it("asks the server to sort when a sortable header is clicked", async () => { const user = userEvent.setup(); - vi.mocked(useBudgets).mockReturnValue({ - data: [ - { - budget_id: "budget-to-delete", - max_budget: 200, - rpm_limit: 20, - tpm_limit: 2000, - updated_at: "2024-01-02T00:00:00Z", - }, - ], - isLoading: false, - } as any); + renderPanel(); + await waitFor(() => expect(getMock).toHaveBeenCalled()); - renderWithProviders(); + await user.click(screen.getByTestId("sort-header-max_budget")); + await waitFor(() => expect(lastQuery().sort).toBe("-max_budget")); - await waitFor(() => { - expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); - }); + await user.click(screen.getByTestId("sort-header-max_budget")); + await waitFor(() => expect(lastQuery().sort).toBe("max_budget")); - await user.click(screen.getByTestId("budget-actions-budget-to-delete")); + await user.click(screen.getByTestId("sort-header-budget_id")); + await waitFor(() => expect(lastQuery().sort).toBe("budget_id")); + }); + + it("searches on budget_id with a debounced q", async () => { + const user = userEvent.setup(); + renderPanel(); + await waitFor(() => expect(getMock).toHaveBeenCalled()); + + await user.type(screen.getByTestId("datatable-search"), "ecc"); + await waitFor(() => expect(lastQuery().q).toBe("ecc")); + expect(queries().some((query) => query.q === "e" || query.q === "ec")).toBe(false); + }); + + it("filters by reset duration and clears it again", async () => { + const user = userEvent.setup(); + renderPanel(); + await waitFor(() => expect(getMock).toHaveBeenCalled()); + + await openFilters(user); + await user.click(screen.getByTestId("budget-filter-duration-7d")); + await user.click(screen.getByTestId("budget-filter-duration-30d")); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(lastQuery()["filter[budget_duration][in]"]).toBe("7d,30d")); + + await user.click(screen.getByTestId("filter-chip-remove-budget_duration")); + await waitFor(() => expect(lastQuery()).not.toHaveProperty("filter[budget_duration][in]")); + }); + + it("filters by budgets with no reset duration", async () => { + const user = userEvent.setup(); + renderPanel(); + await waitFor(() => expect(getMock).toHaveBeenCalled()); + + await openFilters(user); + await user.click(screen.getByTestId("budget-filter-duration-__unset__")); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(lastQuery()["filter[budget_duration][is_null]"]).toBe("true")); + expect(lastQuery()).not.toHaveProperty("filter[budget_duration][in]"); + }); + + it("filters by a max budget range and clears it again", async () => { + const user = userEvent.setup(); + renderPanel(); + await waitFor(() => expect(getMock).toHaveBeenCalled()); + + await openFilters(user); + await user.type(screen.getByTestId("budget-filter-max-budget-min"), "10"); + await user.type(screen.getByTestId("budget-filter-max-budget-max"), "500"); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(lastQuery()["filter[max_budget][gte]"]).toBe("10")); + expect(lastQuery()["filter[max_budget][lte]"]).toBe("500"); + + await user.click(screen.getByTestId("datatable-clear-filters")); + await waitFor(() => expect(lastQuery()).not.toHaveProperty("filter[max_budget][gte]")); + expect(lastQuery()).not.toHaveProperty("filter[max_budget][lte]"); + }); + + it("filters to unlimited budgets only", async () => { + const user = userEvent.setup(); + renderPanel(); + await waitFor(() => expect(getMock).toHaveBeenCalled()); + + await openFilters(user); + await user.type(screen.getByTestId("budget-filter-max-budget-min"), "10"); + await user.click(screen.getByTestId("budget-filter-max-budget-unlimited")); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(lastQuery()["filter[max_budget][is_null]"]).toBe("true")); + expect(lastQuery()).not.toHaveProperty("filter[max_budget][gte]"); + }); + + it("filters by a created date range covering whole local days", async () => { + const user = userEvent.setup(); + renderPanel(); + await waitFor(() => expect(getMock).toHaveBeenCalled()); + + await openFilters(user); + await user.type(screen.getByTestId("budget-filter-created-from"), "2026-01-05"); + await user.type(screen.getByTestId("budget-filter-created-to"), "2026-01-06"); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => + expect(lastQuery()["filter[created_at][gte]"]).toBe(new Date("2026-01-05T00:00:00.000").toISOString()), + ); + expect(lastQuery()["filter[created_at][lte]"]).toBe(new Date("2026-01-06T23:59:59.999").toISOString()); + }); + + it("pages through the results and changes page size", async () => { + const user = userEvent.setup(); + respondWith(DEFAULT_ROWS, 400); + renderPanel(); + await waitFor(() => expect(getMock).toHaveBeenCalled()); + + await user.click(screen.getByTestId("pagination-next")); + await waitFor(() => expect(lastQuery().page).toBe(2)); + expect(lastQuery().page_size).toBe(50); + + await user.click(screen.getByTestId("pagination-page-size")); + await user.click(await screen.findByRole("option", { name: "25" })); + await waitFor(() => expect(lastQuery().page_size).toBe(25)); + }); + + it("renders an access-denied state when the route rejects the caller", async () => { + getMock.mockRejectedValue(new ApiError("Only proxy admins can view budgets", 403, FORBIDDEN_PROBLEM)); + renderPanel(); + expect(await screen.findByText("You do not have access to budgets")).toBeInTheDocument(); + expect(screen.queryByText("No budgets yet")).not.toBeInTheDocument(); + }); + + it("deletes a budget from the actions menu", async () => { + const user = userEvent.setup(); + budgetDeleteMock.mockResolvedValue(undefined); + renderPanel(); + await screen.findByText("ecc1869c-6231-4380-a56d-1a0be457477d"); + + await user.click(screen.getByTestId("budget-actions-ecc1869c-6231-4380-a56d-1a0be457477d")); await user.click(await screen.findByTestId("budget-action-delete")); + await screen.findByText("Delete Budget?"); + await user.click(screen.getByRole("button", { name: /^delete$/i })); - await waitFor(() => { - expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); - }); + await waitFor(() => + expect(budgetDeleteMock).toHaveBeenCalledWith("sk-test", "ecc1869c-6231-4380-a56d-1a0be457477d"), + ); }); - it("should successfully delete a budget", async () => { + it("refetches the current page after a delete", async () => { const user = userEvent.setup(); - const deleteMutateAsync = vi.fn().mockResolvedValue(undefined); - vi.mocked(useBudgets).mockReturnValue({ - data: [ - { - budget_id: "budget-to-delete", - max_budget: 200, - rpm_limit: 20, - tpm_limit: 2000, - updated_at: "2024-01-02T00:00:00Z", - }, - ], - isLoading: false, - } as any); - vi.mocked(useDeleteBudget).mockReturnValue({ - mutateAsync: deleteMutateAsync, - isPending: false, - } as any); + budgetDeleteMock.mockResolvedValue(undefined); + renderPanel(); + await screen.findByText("ecc1869c-6231-4380-a56d-1a0be457477d"); + const before = getMock.mock.calls.length; - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); - }); - - await user.click(screen.getByTestId("budget-actions-budget-to-delete")); + await user.click(screen.getByTestId("budget-actions-ecc1869c-6231-4380-a56d-1a0be457477d")); await user.click(await screen.findByTestId("budget-action-delete")); + await screen.findByText("Delete Budget?"); + await user.click(screen.getByRole("button", { name: /^delete$/i })); - await waitFor(() => { - expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); - }); - - const confirmButton = screen.getByRole("button", { name: /delete/i }); - act(() => { - fireEvent.click(confirmButton); - }); - - await waitFor(() => { - expect(deleteMutateAsync).toHaveBeenCalledWith("budget-to-delete"); - }); - }); - - it("should render empty state without crashing", async () => { - vi.mocked(useBudgets).mockReturnValue({ - data: [], - isLoading: false, - } as any); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("Create a budget to assign to customers.")).toBeInTheDocument(); - }); - }); - - it("should handle delete error", async () => { - const user = userEvent.setup(); - const deleteMutateAsync = vi.fn().mockRejectedValue(new Error("Delete failed")); - vi.mocked(useBudgets).mockReturnValue({ - data: [ - { - budget_id: "budget-to-delete", - max_budget: 200, - rpm_limit: 20, - tpm_limit: 2000, - updated_at: "2024-01-02T00:00:00Z", - }, - ], - isLoading: false, - } as any); - vi.mocked(useDeleteBudget).mockReturnValue({ - mutateAsync: deleteMutateAsync, - isPending: false, - } as any); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); - }); - - await user.click(screen.getByTestId("budget-actions-budget-to-delete")); - await user.click(await screen.findByTestId("budget-action-delete")); - - await waitFor(() => { - expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); - }); - - const confirmButton = screen.getByRole("button", { name: /delete/i }); - act(() => { - fireEvent.click(confirmButton); - }); - - await waitFor(() => { - expect(deleteMutateAsync).toHaveBeenCalledWith("budget-to-delete"); - }); - }); - - it("should open edit modal from the actions menu", async () => { - const user = userEvent.setup(); - vi.mocked(useBudgets).mockReturnValue({ - data: [ - { - budget_id: "budget-to-edit", - max_budget: 300, - rpm_limit: 30, - tpm_limit: 3000, - updated_at: "2024-01-03T00:00:00Z", - }, - ], - isLoading: false, - } as any); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("budget-to-edit")).toBeInTheDocument(); - }); - - await user.click(screen.getByTestId("budget-actions-budget-to-edit")); - await user.click(await screen.findByTestId("budget-action-edit")); - - await waitFor(() => { - expect(screen.getByText("Edit Budget")).toBeInTheDocument(); - }); + await waitFor(() => expect(getMock.mock.calls.length).toBeGreaterThan(before)); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx index 2cf2a4c06ec..e49f0ffc722 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx @@ -3,13 +3,16 @@ * */ -import React, { useState } from "react"; +import { Plus, Wallet } from "lucide-react"; +import React, { useCallback, useState } from "react"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; +import { PageHeader } from "@/components/shared/PageHeader"; +import { ToolbarSeparator } from "@/components/shared/ToolbarSeparator"; import { Button } from "@/components/ui/button"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; import NotificationsManager from "@/components/molecules/notifications_manager"; -import { useBudgets, useDeleteBudget, budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import { useBudgetList, useDeleteBudget, budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import BudgetModal from "./budget_modal"; import BudgetTable from "./BudgetTable"; import EditBudgetModal from "./edit_budget_modal"; @@ -31,21 +34,25 @@ const BudgetPanel: React.FC = ({ accessToken }) => { // Admin Viewer follows the read-parity rule: see budgets, no writes. const canModify = isProxyAdminRole(userRole ?? ""); - const { data: budgetList = [], isLoading } = useBudgets(); + const budgetList = useBudgetList(); const deleteBudget = useDeleteBudget(); - const handleEditCall = async (budget: budgetItem) => { - if (accessToken == null) { - return; - } - setSelectedBudget(budget); - setIsEditModalVisible(true); - }; + // Stable identities keep the memoized column defs stable; new ones remount every header and cell. + const handleEditCall = useCallback( + (budget: budgetItem) => { + if (accessToken == null) { + return; + } + setSelectedBudget(budget); + setIsEditModalVisible(true); + }, + [accessToken], + ); - const handleDeleteClick = (budget: budgetItem) => { + const handleDeleteClick = useCallback((budget: budgetItem) => { setSelectedBudget(budget); setIsDeleteModalVisible(true); - }; + }, []); const handleDeleteConfirm = async () => { if (!selectedBudget || accessToken == null) { @@ -72,23 +79,34 @@ const BudgetPanel: React.FC = ({ accessToken }) => { }; return ( -
- {canModify && ( - - )} - - - - Budgets - - - Examples - - - -
+
+ } + title="Budgets" + subtitle="Spend, TPM and RPM limits you can assign to customers." + /> + +
+ {canModify && ( + <> + + + + )} + + + Budgets + + + Examples + + +
+ +
{selectedBudget && ( = ({ accessToken }) => { existingBudget={selectedBudget} /> )} -

Create a budget to assign to customers.

= ({ accessToken }) => { />
- -
+ +

How to use budget id

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/budgetFilters.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/budgetFilters.test.ts new file mode 100644 index 00000000000..af24e6f4ef2 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/budgetFilters.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, it } from "vitest"; + +import { BUDGET_DURATION_UNSET, serializeBudgetFilters } from "./budgetFilters"; + +describe("serializeBudgetFilters", () => { + it("sends nothing when no filter is active", () => { + expect(serializeBudgetFilters([])).toEqual({}); + }); + + it("maps selected durations onto the in operator", () => { + expect(serializeBudgetFilters([{ id: "budget_duration", value: ["7d", "30d"] }])).toEqual({ + "filter[budget_duration][in]": "7d,30d", + }); + }); + + it("maps 'Not set' onto is_null instead of in", () => { + expect(serializeBudgetFilters([{ id: "budget_duration", value: [BUDGET_DURATION_UNSET] }])).toEqual({ + "filter[budget_duration][is_null]": "true", + }); + }); + + it("never sends in alongside is_null for the same field", () => { + const params = serializeBudgetFilters([{ id: "budget_duration", value: ["7d", BUDGET_DURATION_UNSET] }]); + expect(params["filter[budget_duration][in]"]).toBeUndefined(); + expect(params["filter[budget_duration][is_null]"]).toBe("true"); + }); + + it("maps a max budget range onto gte and lte", () => { + expect(serializeBudgetFilters([{ id: "max_budget", value: { min: "10", max: "250.5" } }])).toEqual({ + "filter[max_budget][gte]": "10", + "filter[max_budget][lte]": "250.5", + }); + }); + + it("sends only the bound that was filled in", () => { + expect(serializeBudgetFilters([{ id: "max_budget", value: { min: "10", max: "" } }])).toEqual({ + "filter[max_budget][gte]": "10", + }); + }); + + it("maps 'Unlimited only' onto is_null and drops the range", () => { + const params = serializeBudgetFilters([{ id: "max_budget", value: { min: "10", unlimitedOnly: true } }]); + expect(params).toEqual({ "filter[max_budget][is_null]": "true" }); + }); + + it("widens a created-at day range to cover the whole local days", () => { + const params = serializeBudgetFilters([{ id: "created_at", value: { from: "2026-01-05", to: "2026-01-06" } }]); + expect(params["filter[created_at][gte]"]).toBe(new Date("2026-01-05T00:00:00.000").toISOString()); + expect(params["filter[created_at][lte]"]).toBe(new Date("2026-01-06T23:59:59.999").toISOString()); + }); + + it("ignores an unparseable date rather than sending a broken bound", () => { + expect(serializeBudgetFilters([{ id: "created_at", value: { from: "not-a-date" } }])).toEqual({}); + }); + + it("ignores filter ids the route does not declare", () => { + expect(serializeBudgetFilters([{ id: "spend", value: "5" }])).toEqual({}); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/budgetFilters.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/budgetFilters.ts new file mode 100644 index 00000000000..f54eddb3913 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/budgetFilters.ts @@ -0,0 +1,91 @@ +import type { ColumnFilter, ColumnFiltersState } from "@tanstack/react-table"; + +export const BUDGET_DURATION_UNSET = "__unset__"; + +export const BUDGET_DURATION_FILTER_OPTIONS: readonly { value: string; label: string }[] = [ + { value: "1h", label: "hourly" }, + { value: "24h", label: "daily" }, + { value: "7d", label: "weekly" }, + { value: "30d", label: "monthly" }, + { value: BUDGET_DURATION_UNSET, label: "Not set" }, +]; + +export interface MaxBudgetFilterValue { + min?: string; + max?: string; + unlimitedOnly?: boolean; +} + +export interface CreatedAtFilterValue { + from?: string; + to?: string; +} + +type QueryEntry = readonly [string, string]; + +const entries = (key: string, value: string): QueryEntry[] => (value === "" ? [] : [[key, value]]); + +const asStringArray = (value: unknown): string[] => + Array.isArray(value) ? value.filter((item): item is string => typeof item === "string") : []; + +const asRecord = (value: unknown): Record => + typeof value === "object" && value !== null ? (value as Record) : {}; + +const asTrimmed = (value: unknown): string => (typeof value === "string" ? value.trim() : ""); + +/** The date inputs give a calendar day; the route wants an instant, so widen to the viewer's whole local day. */ +const isoAt = (day: string, time: string): string => { + if (day === "") { + return ""; + } + const parsed = new Date(`${day}T${time}`); + return Number.isNaN(parsed.getTime()) ? "" : parsed.toISOString(); +}; + +/** + * "Not set" is exclusive with the concrete durations. The route's contract does not say how it + * combines `in` with `is_null` on one field, and under AND semantics that pair can only match + * nothing, so we never send both. + */ +const durationParams = (value: unknown): QueryEntry[] => { + const selected = asStringArray(value); + if (selected.includes(BUDGET_DURATION_UNSET)) { + return [["filter[budget_duration][is_null]", "true"]]; + } + return entries("filter[budget_duration][in]", selected.join(",")); +}; + +const maxBudgetParams = (value: unknown): QueryEntry[] => { + const draft = asRecord(value); + if (draft.unlimitedOnly === true) { + return [["filter[max_budget][is_null]", "true"]]; + } + return [ + ...entries("filter[max_budget][gte]", asTrimmed(draft.min)), + ...entries("filter[max_budget][lte]", asTrimmed(draft.max)), + ]; +}; + +const createdAtParams = (value: unknown): QueryEntry[] => { + const draft = asRecord(value); + return [ + ...entries("filter[created_at][gte]", isoAt(asTrimmed(draft.from), "00:00:00.000")), + ...entries("filter[created_at][lte]", isoAt(asTrimmed(draft.to), "23:59:59.999")), + ]; +}; + +const filterParams = (filter: ColumnFilter): QueryEntry[] => { + switch (filter.id) { + case "budget_duration": + return durationParams(filter.value); + case "max_budget": + return maxBudgetParams(filter.value); + case "created_at": + return createdAtParams(filter.value); + default: + return []; + } +}; + +export const serializeBudgetFilters = (filters: ColumnFiltersState): Readonly> => + Object.fromEntries(filters.flatMap(filterParams)); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/useBudgets.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/useBudgets.ts index 0d8d94f2369..50151e59d3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/useBudgets.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/useBudgets.ts @@ -1,28 +1,46 @@ -import { useQuery, useMutation, useQueryClient, UseQueryResult } from "@tanstack/react-query"; -import { createQueryKeys } from "../common/queryKeysFactory"; -import { getBudgetList, budgetCreateCall, budgetUpdateCall, budgetDeleteCall } from "@/components/networking"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +"use client"; -export interface budgetItem { - budget_id: string; - max_budget: number | null; - rpm_limit: number | null; - tpm_limit: number | null; - updated_at: string; -} +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import type { SortingState } from "@tanstack/react-table"; +import { useCallback } from "react"; + +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { apiClient, budgetCreateCall, budgetUpdateCall, budgetDeleteCall } from "@/components/networking"; +import type { components } from "@/lib/http/schema"; + +import { createQueryKeys } from "../common/queryKeysFactory"; +import { useResourceList, type ResourceListQuery, type ResourceListResult } from "../common/useResourceList"; +import { serializeBudgetFilters } from "./budgetFilters"; + +export type budgetItem = components["schemas"]["BudgetListItem"]; + +type BudgetListResponse = components["schemas"]["ListResponse_BudgetListItem_"]; + +export const BUDGET_LIST_PATH = "/management/v1/budgets"; export const budgetKeys = createQueryKeys("budgets"); -export const useBudgets = (): UseQueryResult => { +const DEFAULT_PAGE_SIZE = 50; +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +export const useBudgetList = (): ResourceListResult => { const { accessToken } = useAuthorized(); - return useQuery({ - queryKey: budgetKeys.list({}), - queryFn: async () => { - const data = await getBudgetList(accessToken!); - return (data ?? []).filter((item: budgetItem | null): item is budgetItem => item != null); - }, + + const fetchPage = useCallback( + (query: ResourceListQuery, signal: AbortSignal): Promise => + apiClient.get(BUDGET_LIST_PATH, { accessToken, query, signal }), + [accessToken], + ); + + const listOptions = { + queryKey: budgetKeys.lists(), + fetchPage, + serializeFilters: serializeBudgetFilters, + defaultSorting: DEFAULT_SORTING, + defaultPageSize: DEFAULT_PAGE_SIZE, enabled: Boolean(accessToken), - }); + }; + return useResourceList(listOptions); }; export const useCreateBudget = () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx new file mode 100644 index 00000000000..3ca67b082ec --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx @@ -0,0 +1,161 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import type { ColumnFiltersState } from "@tanstack/react-table"; +import { act, renderHook, waitFor } from "@testing-library/react"; +import React, { type PropsWithChildren } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { + toSortParam, + useResourceList, + type ResourceListPage, + type ResourceListQuery, + type UseResourceListOptions, +} from "./useResourceList"; + +interface Row { + id: string; +} + +const page = (rows: Row[], totalCount: number): ResourceListPage => ({ + data: rows, + meta: { total_count: totalCount, page: 1, page_size: 50, total_pages: 1 }, +}); + +const noFilters = (): Readonly> => ({}); + +const calls: ResourceListQuery[] = []; + +const renderList = (overrides: Partial> = {}) => { + const client = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const fetchPage = vi.fn((query: ResourceListQuery) => { + calls.push(query); + return Promise.resolve(page([{ id: "a" }], 3)); + }); + const options: UseResourceListOptions = { + queryKey: ["widgets", "list"], + fetchPage, + serializeFilters: noFilters, + defaultSorting: [{ id: "created_at", desc: true }], + defaultPageSize: 50, + enabled: true, + ...overrides, + }; + return renderHook(() => useResourceList(options), { wrapper }); +}; + +const lastCall = (): ResourceListQuery => calls[calls.length - 1]; + +describe("toSortParam", () => { + it("prefixes descending fields with a minus and joins with commas", () => { + expect(toSortParam([{ id: "created_at", desc: true }])).toBe("-created_at"); + expect(toSortParam([{ id: "max_budget", desc: false }])).toBe("max_budget"); + expect( + toSortParam([ + { id: "a", desc: false }, + { id: "b", desc: true }, + ]), + ).toBe("a,-b"); + }); +}); + +describe("useResourceList", () => { + beforeEach(() => { + calls.length = 0; + }); + + it("requests the first page with the default sort", async () => { + const { result } = renderList(); + await waitFor(() => expect(result.current.rowCount).toBe(3)); + expect(lastCall()).toEqual({ page: 1, page_size: 50, sort: "-created_at" }); + }); + + it("exposes the returned rows and total count", async () => { + const { result } = renderList(); + await waitFor(() => expect(result.current.rows).toEqual([{ id: "a" }])); + expect(result.current.rowCount).toBe(3); + }); + + it("does not fetch while disabled", async () => { + const { result } = renderList({ enabled: false }); + await waitFor(() => expect(result.current.isLoading).toBe(false)); + expect(calls).toHaveLength(0); + }); + + it("sends the new sort and returns to the first page", async () => { + const { result } = renderList(); + await waitFor(() => expect(calls).toHaveLength(1)); + + act(() => result.current.onPaginationChange({ pageIndex: 2, pageSize: 50 })); + await waitFor(() => expect(lastCall().page).toBe(3)); + + act(() => result.current.onSortingChange([{ id: "max_budget", desc: false }])); + await waitFor(() => expect(lastCall().sort).toBe("max_budget")); + expect(lastCall().page).toBe(1); + }); + + it("omits sort entirely when nothing is sorted", async () => { + const { result } = renderList({ defaultSorting: [] }); + await waitFor(() => expect(calls).toHaveLength(1)); + expect(result.current.sorting).toEqual([]); + expect(lastCall()).not.toHaveProperty("sort"); + }); + + it("debounces the search into a single trimmed q and returns to the first page", async () => { + const { result } = renderList(); + await waitFor(() => expect(calls).toHaveLength(1)); + + act(() => result.current.onPaginationChange({ pageIndex: 1, pageSize: 50 })); + await waitFor(() => expect(lastCall().page).toBe(2)); + + act(() => result.current.onSearchChange("bud")); + act(() => result.current.onSearchChange("budg ")); + + await waitFor(() => expect(lastCall().q).toBe("budg")); + expect(lastCall().page).toBe(1); + expect(calls.some((call) => call.q === "bud")).toBe(false); + }); + + it("stops sending q once the search box is cleared", async () => { + const { result } = renderList(); + act(() => result.current.onSearchChange("budget")); + await waitFor(() => expect(lastCall().q).toBe("budget")); + + act(() => result.current.onSearchChange("")); + await waitFor(() => expect(lastCall()).not.toHaveProperty("q")); + }); + + it("merges serialized filters into the request and returns to the first page", async () => { + const serializeFilters = (filters: ColumnFiltersState): Readonly> => + filters.length === 0 ? {} : { "filter[colour][in]": String(filters[0].value) }; + const { result } = renderList({ serializeFilters }); + await waitFor(() => expect(calls).toHaveLength(1)); + + act(() => result.current.onPaginationChange({ pageIndex: 3, pageSize: 50 })); + await waitFor(() => expect(lastCall().page).toBe(4)); + + act(() => result.current.onColumnFiltersChange([{ id: "colour", value: "red" }])); + await waitFor(() => expect(lastCall()["filter[colour][in]"]).toBe("red")); + expect(lastCall().page).toBe(1); + + act(() => result.current.onColumnFiltersChange([])); + await waitFor(() => expect(lastCall()).not.toHaveProperty("filter[colour][in]")); + }); + + it("sends the requested page size", async () => { + const { result } = renderList(); + await waitFor(() => expect(calls).toHaveLength(1)); + + act(() => result.current.onPaginationChange({ pageIndex: 0, pageSize: 25 })); + await waitFor(() => expect(lastCall().page_size).toBe(25)); + }); + + it("surfaces a failed page as an error instead of empty rows", async () => { + const fetchPage = vi.fn(() => Promise.reject(new Error("boom"))); + const { result } = renderList({ fetchPage }); + await waitFor(() => expect(result.current.error?.message).toBe("boom")); + expect(result.current.rows).toEqual([]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts new file mode 100644 index 00000000000..8a6376b2248 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts @@ -0,0 +1,139 @@ +"use client"; + +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { useQuery, type UseQueryOptions } from "@tanstack/react-query"; +import type { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import { useCallback, useMemo, useState } from "react"; + +import type { components } from "@/lib/http/schema"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; + +export type ResourceListQuery = Readonly>; + +/** The management list envelope. The generated response models are monomorphic, so only `data` is generic here. */ +export type ResourceListMeta = components["schemas"]["ListMeta"]; + +export interface ResourceListPage { + data: TRow[]; + meta: ResourceListMeta; +} + +export interface UseResourceListOptions { + /** Prefix every list variant hangs off, so invalidating the resource root refetches whichever page is on screen. */ + queryKey: readonly unknown[]; + fetchPage: (query: ResourceListQuery, signal: AbortSignal) => Promise>; + /** Must be referentially stable; it feeds the query key. */ + serializeFilters: (filters: ColumnFiltersState) => Readonly>; + defaultSorting: SortingState; + defaultPageSize: number; + enabled: boolean; +} + +export interface ResourceListResult { + rows: TRow[]; + rowCount: number; + isLoading: boolean; + isFetching: boolean; + error: Error | null; + refetch: () => void; + + sorting: SortingState; + onSortingChange: OnChangeFn; + pagination: PaginationState; + onPaginationChange: OnChangeFn; + columnFilters: ColumnFiltersState; + onColumnFiltersChange: OnChangeFn; + searchValue: string; + onSearchChange: (value: string) => void; +} + +/** JSON:API sort form: comma separated fields, `-` prefix for descending. */ +export const toSortParam = (sorting: SortingState): string => + sorting.map((entry) => (entry.desc ? `-${entry.id}` : entry.id)).join(","); + +/** + * State container for a table whose sorting, paging, search and filtering all run + * on the server. It owns those four pieces of state, folds them into one JSON:API + * query, and returns the exact props DataTable's server modes want. + * + * Empty parameters are dropped rather than sent blank because the management + * routes reject query params they do not declare. + */ +export function useResourceList(options: UseResourceListOptions): ResourceListResult { + const { queryKey, fetchPage, serializeFilters, defaultSorting, defaultPageSize, enabled } = options; + + const [sorting, setSorting] = useState(defaultSorting); + const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: defaultPageSize }); + const [columnFilters, setColumnFilters] = useState([]); + const [searchValue, setSearchValue] = useState(""); + const [debouncedSearch] = useDebouncedValue(searchValue, { wait: DEBOUNCE_WAIT_MS }); + + const query = useMemo(() => { + const sort = toSortParam(sorting); + const search = debouncedSearch.trim(); + return { + page: pagination.pageIndex + 1, + page_size: pagination.pageSize, + ...(sort === "" ? {} : { sort }), + ...(search === "" ? {} : { q: search }), + ...serializeFilters(columnFilters), + }; + }, [sorting, pagination.pageIndex, pagination.pageSize, debouncedSearch, columnFilters, serializeFilters]); + + const queryOptions: UseQueryOptions, Error, ResourceListPage, readonly unknown[]> = { + queryKey: [...queryKey, query], + queryFn: ({ signal }) => fetchPage(query, signal), + enabled, + placeholderData: (previous) => previous, + }; + const { data, isLoading, isFetching, error, refetch: refetchQuery } = useQuery(queryOptions); + + const toFirstPage = useCallback(() => setPagination((previous) => ({ ...previous, pageIndex: 0 })), []); + + const onSortingChange = useCallback>( + (updater) => { + setSorting(updater); + toFirstPage(); + }, + [toFirstPage], + ); + + const onColumnFiltersChange = useCallback>( + (updater) => { + setColumnFilters(updater); + toFirstPage(); + }, + [toFirstPage], + ); + + const onSearchChange = useCallback( + (value: string) => { + setSearchValue(value); + toFirstPage(); + }, + [toFirstPage], + ); + + const refetch = useCallback(() => { + void refetchQuery(); + }, [refetchQuery]); + + const rows = useMemo(() => data?.data ?? [], [data]); + + return { + rows, + rowCount: data?.meta.total_count ?? 0, + isLoading, + isFetching, + error, + refetch, + sorting, + onSortingChange, + pagination, + onPaginationChange: setPagination, + columnFilters, + onColumnFiltersChange, + searchValue, + onSearchChange, + }; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index f532c44ffd7..4061026b94d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -24,6 +24,7 @@ export interface TeamListCallOptions { teamID?: string | null; team_alias?: string | null; search?: string | null; + searchTeamIdMatch?: "exact" | "prefix" | null; userID?: string | null; sortBy?: string | null; sortOrder?: string | null; @@ -48,6 +49,7 @@ export const teamListCall = async ( organization_id: options.organizationID, team_alias: options.team_alias, search: options.search, + search_team_id_match: options.searchTeamIdMatch, user_id: options.userID, page, page_size: pageSize, @@ -213,6 +215,7 @@ const deletedTeamListCall = async ( organization_id: options.organizationID, team_alias: options.team_alias, search: options.search, + search_team_id_match: options.searchTeamIdMatch, user_id: options.userID, page, page_size: pageSize, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx new file mode 100644 index 00000000000..e8730a5b974 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx @@ -0,0 +1,158 @@ +import React from "react"; +import { Form, Input, Select, Tooltip } from "antd"; +import { InfoCircleOutlined } from "@ant-design/icons"; + +interface IdJagFormFieldsProps { + isEditing?: boolean; +} + +const fieldClassName = "rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"; + +const FieldLabel: React.FC<{ label: string; tooltip: string }> = ({ label, tooltip }) => ( + + {label} + + + + +); + +const IdJagFormFields: React.FC = ({ isEditing = false }) => { + const placeholderSuffix = isEditing ? " (leave blank to keep existing)" : ""; + + return ( + <> + + } + name="token_exchange_endpoint" + rules={[{ required: !isEditing, message: "The org token endpoint is required for ID-JAG" }]} + > + + + + } + name={["credentials", "id_jag_resource_token_endpoint"]} + rules={[{ required: !isEditing, message: "The resource token endpoint is required for ID-JAG" }]} + > + + + } + name={["credentials", "client_id"]} + rules={[{ required: !isEditing, message: "Client ID is required for ID-JAG" }]} + > + + + + } + name={["credentials", "client_secret"]} + dependencies={[["credentials", "client_private_key"]]} + rules={[ + ({ getFieldValue }) => ({ + validator: (_, value) => { + if (isEditing || value || getFieldValue(["credentials", "client_private_key"])) { + return Promise.resolve(); + } + return Promise.reject(new Error("Provide either a client secret or a client private key")); + }, + }), + ]} + > + + + + } + name={["credentials", "client_private_key"]} + > + + + + } + name={["credentials", "client_private_key_id"]} + > + + + + } + name={["credentials", "client_assertion_signing_alg"]} + > + + + + } + name="audience" + > + + + + } + name={["credentials", "id_jag_resource"]} + > + + + + } + name="subject_token_type" + > + + + } + name={["credentials", "scopes"]} + > + + +