Merge branch 'litellm_internal_staging' into fix-vertex-typearray-null

This commit is contained in:
ErenAta16 2026-08-01 16:17:35 +03:00 • committed by GitHub
commit 2505fe5b78
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
114 changed files with 7972 additions and 936 deletions

View file

@ -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

View file

@ -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

View file

@ -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")

View file

@ -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

View file

@ -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)

View file

@ -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

View file

@ -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={},

View file

@ -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)

View file

@ -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(

View file

@ -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

View file

@ -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 = []

View file

@ -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,

View file

@ -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

View file

@ -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)

View file

@ -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

View file

@ -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)

View file

@ -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",

View file

@ -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,
)
#########################################################

View file

@ -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:

View file

@ -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"]

View file

@ -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.",
)
)

View file

@ -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),
)

View file

@ -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 <op> 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),
)

View file

@ -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(

View file

@ -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

View file

@ -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))

View file

@ -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,

View file

@ -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

View file

@ -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",
)

View file

@ -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.

View file

@ -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,

View file

@ -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)

View file

@ -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

View file

@ -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.

View file

@ -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:

View file

@ -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()

View file

@ -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

View file

@ -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"""

View file

@ -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",

View file

@ -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

View file

@ -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: <reason>`.
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: <reason>`.
LIT003 noqa suppression without rule codes or without a reason.
Required shape: `# noqa: TID251 # <reason>`
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:

View file

@ -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)"}

View file

@ -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(

View file

@ -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 ----------

View file

@ -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,

View file

@ -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"],

View file

@ -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,

View file

@ -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

View file

@ -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()

View file

@ -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"

View file

@ -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"}

View file

@ -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

View file

@ -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

View file

@ -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()

View file

@ -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)

View file

@ -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"})

View file

@ -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",

View file

@ -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"
)

View file

@ -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()

View file

@ -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
)

View file

@ -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"

View file

@ -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

View file

@ -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()

View file

@ -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"]}

View file

@ -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")

View file

@ -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

View file

@ -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 = {

View file

@ -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

View file

@ -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"])

View file

@ -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}"
)

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -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
}
}

View file

@ -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

View file

@ -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<typeof import("@/utils/dataUtils")>()),
copyToClipboard: copyToClipboardMock,
}));
const makeBudget = (overrides: Partial<budgetItem> = {}): 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<budgetItem>> = {}): ResourceListResult<budgetItem> => ({
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<typeof userEvent.setup>, 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(<BudgetTable {...defaultProps} />);
renderWithProviders(<BudgetTable {...defaultProps} list={makeList()} />);
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(<BudgetTable {...defaultProps} list={makeList()} />);
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(<BudgetTable {...defaultProps} list={makeList()} />);
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(<BudgetTable {...defaultProps} list={list} />);
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(<BudgetTable {...defaultProps} budgets={[makeBudget({ budget_id: budgetId })]} />);
const list = makeList({ rows: [makeBudget({ budget_id: budgetId })] });
renderWithProviders(<BudgetTable {...defaultProps} list={list} />);
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(<BudgetTable {...defaultProps} list={list} />);
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(<BudgetTable {...defaultProps} list={list} />);
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(<BudgetTable {...defaultProps} list={makeList()} />);
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(<BudgetTable {...defaultProps} list={makeList()} />);
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(<BudgetTable {...defaultProps} list={makeList({ onSortingChange })} />);
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(
<BudgetTable {...defaultProps} budgets={[makeBudget({ max_budget: null, tpm_limit: null, rpm_limit: null })]} />,
);
const list = makeList({ rows: [makeBudget({ max_budget: null, tpm_limit: null, rpm_limit: null })] });
renderWithProviders(<BudgetTable {...defaultProps} list={list} />);
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(<BudgetTable {...defaultProps} budgets={budgets} />);
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(<BudgetTable {...defaultProps} />);
const list = makeList();
renderWithProviders(<BudgetTable {...defaultProps} list={list} />);
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(<BudgetTable {...defaultProps} />);
const list = makeList();
renderWithProviders(<BudgetTable {...defaultProps} list={list} />);
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(<BudgetTable {...defaultProps} canModify={false} />);
renderWithProviders(<BudgetTable {...defaultProps} canModify={false} list={makeList()} />);
expect(screen.queryByTestId("budget-actions-budget-1")).not.toBeInTheDocument();
});
it("should show skeleton rows when loading", () => {
renderWithProviders(<BudgetTable {...defaultProps} budgets={[]} isLoading />);
renderWithProviders(<BudgetTable {...defaultProps} list={makeList({ rows: [], isLoading: true })} />);
expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0);
});
it("should show the empty state when there are no budgets", () => {
renderWithProviders(<BudgetTable {...defaultProps} budgets={[]} />);
renderWithProviders(<BudgetTable {...defaultProps} list={makeList({ rows: [], rowCount: 0 })} />);
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(<BudgetTable {...defaultProps} list={list} />);
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(<BudgetTable {...defaultProps} list={list} />);
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(<BudgetTable {...defaultProps} list={list} />);
expect(screen.getByText("Could not load budgets")).toBeInTheDocument();
expect(screen.getByText("budget store unavailable")).toBeInTheDocument();
});
});

View file

@ -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<budgetItem>;
canModify: boolean;
onEditClick: (budget: budgetItem) => void;
onDeleteClick: (budget: budgetItem) => void;
}
function EmptyState() {
const PAGE_SIZE_OPTIONS = [25, 50, 100];
const FILTER_LABELS: Record<string, string> = {
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 (
<div className="flex flex-col items-center gap-1 py-6">
<div className="mb-1 flex size-10 items-center justify-center rounded-lg bg-muted">
<Inbox className="size-5 text-muted-foreground" />
</div>
<div className="text-sm font-medium text-foreground">No budgets yet</div>
<div className="text-sm font-medium text-foreground">{hasQuery ? "No matching budgets" : "No budgets yet"}</div>
<div className="text-sm text-muted-foreground">
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."}
</div>
</div>
);
}
const BudgetTable: React.FC<BudgetTableProps> = ({ 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 (
<div className="flex flex-col items-center gap-1 py-6">
<div className="mb-1 flex size-10 items-center justify-center rounded-lg bg-muted">
<ShieldAlert className="size-5 text-muted-foreground" />
</div>
<div className="text-sm font-medium text-foreground">
{forbidden ? "You do not have access to budgets" : "Could not load budgets"}
</div>
<div className="text-sm text-muted-foreground">
{forbidden ? "Ask a proxy admin to grant you the admin viewer role." : error.message}
</div>
</div>
);
}
/** "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 (
<div className="flex flex-col gap-2">
{BUDGET_DURATION_FILTER_OPTIONS.map((option) => (
<Label key={option.value} className="font-normal">
<Checkbox
checked={selected.includes(option.value)}
onCheckedChange={(checked) => toggle(option.value, checked === true)}
data-testid={`budget-filter-duration-${option.value}`}
/>
{option.label}
</Label>
))}
</div>
);
}
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 (
<>
<DataTableFilterField label="Reset">
<DurationFilter
selected={(get("budget_duration") as string[] | undefined) ?? []}
onChange={(selected) => set("budget_duration", selected)}
/>
</DataTableFilterField>
<DataTableFilterField label="Max Budget (USD)">
<div className="flex items-center gap-2">
<Input
type="number"
min={0}
step="0.01"
value={maxBudget.min ?? ""}
disabled={unlimitedOnly}
onChange={(event) => set("max_budget", normalizeMaxBudget({ ...maxBudget, min: event.target.value }))}
placeholder="Min"
aria-label="Minimum max budget"
data-testid="budget-filter-max-budget-min"
/>
<Input
type="number"
min={0}
step="0.01"
value={maxBudget.max ?? ""}
disabled={unlimitedOnly}
onChange={(event) => set("max_budget", normalizeMaxBudget({ ...maxBudget, max: event.target.value }))}
placeholder="Max"
aria-label="Maximum max budget"
data-testid="budget-filter-max-budget-max"
/>
</div>
<Label className="mt-1 font-normal">
<Checkbox
checked={unlimitedOnly}
onCheckedChange={(checked) => set("max_budget", normalizeMaxBudget({ unlimitedOnly: checked === true }))}
data-testid="budget-filter-max-budget-unlimited"
/>
Unlimited only
</Label>
</DataTableFilterField>
<DataTableFilterField label="Created">
<div className="flex items-center gap-2">
<Input
type="date"
value={created.from ?? ""}
onChange={(event) => set("created_at", normalizeCreatedAt({ ...created, from: event.target.value }))}
aria-label="Created from"
data-testid="budget-filter-created-from"
/>
<Input
type="date"
value={created.to ?? ""}
onChange={(event) => set("created_at", normalizeCreatedAt({ ...created, to: event.target.value }))}
aria-label="Created to"
data-testid="budget-filter-created-to"
/>
</div>
</DataTableFilterField>
</>
);
}
const BudgetTable: React.FC<BudgetTableProps> = ({ 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 ? <EmptyState hasQuery={hasQuery} /> : <ErrorState error={list.error} />;
return (
<DataTable
data={rows}
data={list.rows}
columns={columns}
getRowId={(budget, index) => 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={<EmptyState />}
noDataMessage={emptyMessage}
size="compact"
toolbar={(table) => (
<>
<DataTableToolbar
table={table}
searchValue={list.searchValue}
onSearchChange={list.onSearchChange}
searchPlaceholder="Search by budget ID…"
onOpenFilters={() => setFiltersOpen(true)}
onRefresh={list.refetch}
isRefreshing={list.isFetching}
filterLabels={FILTER_LABELS}
formatFilterValue={formatFilterValue}
/>
<DataTableFilterDrawer
table={table}
open={filtersOpen}
onOpenChange={setFiltersOpen}
title="Filters"
description="Narrow down your budgets"
>
{(draft) => <BudgetFilterFields {...draft} />}
</DataTableFilterDrawer>
</>
)}
/>
);
};

View file

@ -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<budgetItem> = () => true;
serverFilter.autoRemove = () => false;
function RateLimitCell({ value }: { value: number | null | undefined }) {
if (value == null) {
return <span className="text-muted-foreground">n/a</span>;
}
return <span className="tabular-nums">{value}</span>;
}
function BudgetDurationCell({ value }: { value: string | null | undefined }) {
if (!value) {
return <span className="text-muted-foreground">Not set</span>;
}
return <span className="whitespace-nowrap">{getBudgetDurationLabel(value)}</span>;
}
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<string, boolean> = {
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 }) => <IdCell value={row.original.budget_id} variant="plain" className="max-w-full" />,
header: ({ column }) => <DataTableSortHeader column={column} title="Budget ID" />,
cell: ({ row }) => (
<IdCell value={row.original.budget_id} variant="plain" truncate={false} copyable className="whitespace-nowrap" />
),
},
{
id: "max_budget",
accessorKey: "max_budget",
filterFn: serverFilter,
meta: { title: "Max Budget", numeric: true },
header: "Max Budget",
header: ({ column }) => <DataTableSortHeader column={column} title="Max Budget" />,
size: 120,
enableSorting: false,
cell: ({ row }) => <MoneyCell value={row.original.max_budget} decimals={2} showZero emptyText="Unlimited" />,
},
{
id: "tpm_limit",
accessorKey: "tpm_limit",
meta: { title: "TPM", numeric: true },
header: "TPM",
header: ({ column }) => <DataTableSortHeader column={column} title="TPM" />,
size: 100,
enableSorting: false,
cell: ({ row }) => <RateLimitCell value={row.original.tpm_limit} />,
},
{
id: "rpm_limit",
accessorKey: "rpm_limit",
meta: { title: "RPM", numeric: true },
header: "RPM",
header: ({ column }) => <DataTableSortHeader column={column} title="RPM" />,
size: 100,
enableSorting: false,
cell: ({ row }) => <RateLimitCell value={row.original.rpm_limit} />,
},
{
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 }) => <DataTableSortHeader column={column} title="Reset" />,
size: 110,
cell: ({ row }) => <BudgetDurationCell value={row.original.budget_duration} />,
},
{
id: "created_at",
accessorKey: "created_at",
filterFn: serverFilter,
meta: { title: "Created" },
header: ({ column }) => <DataTableSortHeader column={column} title="Created" />,
size: 160,
cell: ({ row }) => <DateCell value={row.original.created_at} />,
},
...(canModify
? [
{

View file

@ -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(<QueryClientProvider client={qc}>{ui}</QueryClientProvider>);
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<typeof budgetRow>[], totalCount: number) => {
getMock.mockResolvedValue({
data: rows,
meta: { total_count: totalCount, page: 1, page_size: 50, total_pages: Math.ceil(totalCount / 50) },
});
};
type QueryRecord = Record<string, string | number>;
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(
<QueryClientProvider client={client}>
<BudgetPanel accessToken="sk-test" />
</QueryClientProvider>,
);
};
const openFilters = async (user: ReturnType<typeof userEvent.setup>) => {
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(<BudgetPanel accessToken="token-123" />);
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(<BudgetPanel accessToken="token-123" />);
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(<BudgetPanel accessToken="token-123" />);
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(<BudgetPanel accessToken="token-123" />);
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(<BudgetPanel accessToken="token-123" />);
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(<BudgetPanel accessToken="token-123" />);
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));
});
});

View file

@ -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<BudgetSettingsPageProps> = ({ 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<BudgetSettingsPageProps> = ({ accessToken }) => {
};
return (
<div className="w-full mx-auto flex-auto overflow-y-auto m-8 p-2">
{canModify && (
<Button size="sm" className="mb-2" onClick={() => setIsCreateModelVisible(true)}>
+ Create Budget
</Button>
)}
<Tabs defaultValue="budgets">
<TabsList variant="line" className="h-auto w-full justify-start rounded-none border-b p-0">
<TabsTrigger value="budgets" className="flex-none rounded-none px-4 py-2">
Budgets
</TabsTrigger>
<TabsTrigger value="examples" className="flex-none rounded-none px-4 py-2">
Examples
</TabsTrigger>
</TabsList>
<TabsContent value="budgets">
<div className="mt-6">
<div className="flex h-full flex-col gap-4 p-6 px-12">
<PageHeader
icon={<Wallet className="size-5" />}
title="Budgets"
subtitle="Spend, TPM and RPM limits you can assign to customers."
/>
<Tabs defaultValue="budgets" className="min-h-0 flex-1 gap-0">
<div className="flex items-center gap-4 border-b border-border">
{canModify && (
<>
<Button onClick={() => setIsCreateModelVisible(true)}>
<Plus className="size-4" />
Create Budget
</Button>
<ToolbarSeparator className="h-6" />
</>
)}
<TabsList variant="line">
<TabsTrigger value="budgets" className="flex-none px-4">
Budgets
</TabsTrigger>
<TabsTrigger value="examples" className="flex-none px-4">
Examples
</TabsTrigger>
</TabsList>
</div>
<TabsContent value="budgets" className="flex min-h-0 flex-1 flex-col">
<div className="flex min-h-0 flex-1 flex-col pt-6">
<BudgetModal isModalVisible={isCreateModelVisible} setIsModalVisible={setIsCreateModelVisible} />
{selectedBudget && (
<EditBudgetModal
@ -97,10 +115,8 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
existingBudget={selectedBudget}
/>
)}
<p className="mb-4 text-sm text-muted-foreground">Create a budget to assign to customers.</p>
<BudgetTable
budgets={budgetList}
isLoading={isLoading}
list={budgetList}
canModify={canModify}
onEditClick={handleEditCall}
onDeleteClick={handleDeleteClick}
@ -122,8 +138,8 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
/>
</div>
</TabsContent>
<TabsContent value="examples">
<div className="mt-6">
<TabsContent value="examples" className="min-h-0 flex-1 overflow-y-auto">
<div className="pt-6">
<p className="text-base text-muted-foreground">How to use budget id</p>
<Tabs defaultValue="assign-budget">
<TabsList variant="line" className="h-auto w-full justify-start rounded-none border-b p-0">

View file

@ -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({});
});
});

View file

@ -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<string, unknown> =>
typeof value === "object" && value !== null ? (value as Record<string, unknown>) : {};
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<Record<string, string>> =>
Object.fromEntries(filters.flatMap(filterParams));

View file

@ -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<budgetItem[]> => {
const DEFAULT_PAGE_SIZE = 50;
const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }];
export const useBudgetList = (): ResourceListResult<budgetItem> => {
const { accessToken } = useAuthorized();
return useQuery<budgetItem[]>({
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<BudgetListResponse> =>
apiClient.get<BudgetListResponse>(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<budgetItem>(listOptions);
};
export const useCreateBudget = () => {

View file

@ -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<Row> => ({
data: rows,
meta: { total_count: totalCount, page: 1, page_size: 50, total_pages: 1 },
});
const noFilters = (): Readonly<Record<string, string>> => ({});
const calls: ResourceListQuery[] = [];
const renderList = (overrides: Partial<UseResourceListOptions<Row>> = {}) => {
const client = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } });
const wrapper = ({ children }: PropsWithChildren) => (
<QueryClientProvider client={client}>{children}</QueryClientProvider>
);
const fetchPage = vi.fn((query: ResourceListQuery) => {
calls.push(query);
return Promise.resolve(page([{ id: "a" }], 3));
});
const options: UseResourceListOptions<Row> = {
queryKey: ["widgets", "list"],
fetchPage,
serializeFilters: noFilters,
defaultSorting: [{ id: "created_at", desc: true }],
defaultPageSize: 50,
enabled: true,
...overrides,
};
return renderHook(() => useResourceList<Row>(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<Record<string, string>> =>
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([]);
});
});

View file

@ -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<Record<string, string | number>>;
/** 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<TRow> {
data: TRow[];
meta: ResourceListMeta;
}
export interface UseResourceListOptions<TRow> {
/** 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<ResourceListPage<TRow>>;
/** Must be referentially stable; it feeds the query key. */
serializeFilters: (filters: ColumnFiltersState) => Readonly<Record<string, string>>;
defaultSorting: SortingState;
defaultPageSize: number;
enabled: boolean;
}
export interface ResourceListResult<TRow> {
rows: TRow[];
rowCount: number;
isLoading: boolean;
isFetching: boolean;
error: Error | null;
refetch: () => void;
sorting: SortingState;
onSortingChange: OnChangeFn<SortingState>;
pagination: PaginationState;
onPaginationChange: OnChangeFn<PaginationState>;
columnFilters: ColumnFiltersState;
onColumnFiltersChange: OnChangeFn<ColumnFiltersState>;
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<TRow>(options: UseResourceListOptions<TRow>): ResourceListResult<TRow> {
const { queryKey, fetchPage, serializeFilters, defaultSorting, defaultPageSize, enabled } = options;
const [sorting, setSorting] = useState<SortingState>(defaultSorting);
const [pagination, setPagination] = useState<PaginationState>({ pageIndex: 0, pageSize: defaultPageSize });
const [columnFilters, setColumnFilters] = useState<ColumnFiltersState>([]);
const [searchValue, setSearchValue] = useState("");
const [debouncedSearch] = useDebouncedValue(searchValue, { wait: DEBOUNCE_WAIT_MS });
const query = useMemo<ResourceListQuery>(() => {
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<ResourceListPage<TRow>, Error, ResourceListPage<TRow>, 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<OnChangeFn<SortingState>>(
(updater) => {
setSorting(updater);
toFirstPage();
},
[toFirstPage],
);
const onColumnFiltersChange = useCallback<OnChangeFn<ColumnFiltersState>>(
(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,
};
}

View file

@ -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,

View file

@ -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 }) => (
<span className="text-sm font-medium text-gray-700 flex items-center">
{label}
<Tooltip title={tooltip}>
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
</Tooltip>
</span>
);
const IdJagFormFields: React.FC<IdJagFormFieldsProps> = ({ isEditing = false }) => {
const placeholderSuffix = isEditing ? " (leave blank to keep existing)" : "";
return (
<>
<Form.Item
label={
<FieldLabel
label="Org Token Endpoint (leg 1)"
tooltip="Your IdP org authorization server's token endpoint. LiteLLM exchanges the user's identity assertion here for an ID-JAG assertion (RFC 8693 with requested_token_type=urn:ietf:params:oauth:token-type:id-jag)."
/>
}
name="token_exchange_endpoint"
rules={[{ required: !isEditing, message: "The org token endpoint is required for ID-JAG" }]}
>
<Input placeholder="https://your-org.okta.com/oauth2/v1/token" className={fieldClassName} />
</Form.Item>
<Form.Item
label={
<FieldLabel
label="Resource Token Endpoint (leg 2)"
tooltip="The upstream resource authorization server's token endpoint. LiteLLM posts the ID-JAG assertion here as an RFC 7523 jwt-bearer grant to get the access token the MCP server accepts."
/>
}
name={["credentials", "id_jag_resource_token_endpoint"]}
rules={[{ required: !isEditing, message: "The resource token endpoint is required for ID-JAG" }]}
>
<Input placeholder="https://upstream.example.com/oauth2/token" className={fieldClassName} />
</Form.Item>
<Form.Item
label={<FieldLabel label="Client ID" tooltip="OAuth2 client ID LiteLLM authenticates as on both legs." />}
name={["credentials", "client_id"]}
rules={[{ required: !isEditing, message: "Client ID is required for ID-JAG" }]}
>
<Input.Password placeholder={`Enter OAuth client ID${placeholderSuffix}`} className={fieldClassName} />
</Form.Item>
<Form.Item
label={
<FieldLabel
label="Client Secret"
tooltip="Authenticates LiteLLM as the OAuth client via client_secret_post. Leave blank when using a private key instead; a private key takes precedence over this secret."
/>
}
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"));
},
}),
]}
>
<Input.Password placeholder={`Enter OAuth client secret${placeholderSuffix}`} className={fieldClassName} />
</Form.Item>
<Form.Item
label={
<FieldLabel
label="Client Private Key (PEM)"
tooltip="PEM private key signing the RFC 7523 private_key_jwt client assertion. Okta Cross App Access normally requires this. When set it takes precedence over the client secret."
/>
}
name={["credentials", "client_private_key"]}
>
<Input.TextArea
rows={3}
placeholder={`-----BEGIN PRIVATE KEY-----${placeholderSuffix}`}
className={fieldClassName}
/>
</Form.Item>
<Form.Item
label={
<FieldLabel
label="Private Key ID (optional)"
tooltip="The kid advertised in the client assertion JWT header, so the IdP can select the right registered key."
/>
}
name={["credentials", "client_private_key_id"]}
>
<Input placeholder="my-signing-key-1" className={fieldClassName} />
</Form.Item>
<Form.Item
label={
<FieldLabel
label="Client Assertion Signing Algorithm (optional)"
tooltip="Algorithm signing the client assertion JWT. Defaults to RS256."
/>
}
name={["credentials", "client_assertion_signing_alg"]}
>
<Input placeholder="RS256" className={fieldClassName} />
</Form.Item>
<Form.Item
label={
<FieldLabel
label="Audience (optional)"
tooltip="RFC 8693 audience sent on leg 1, identifying the upstream the ID-JAG assertion is minted for."
/>
}
name="audience"
>
<Input placeholder="https://upstream.example.com" className={fieldClassName} />
</Form.Item>
<Form.Item
label={
<FieldLabel
label="Resource Indicator (optional)"
tooltip="RFC 8707 resource indicator sent on leg 1. Separate from Audience, which is the RFC 8693 parameter."
/>
}
name={["credentials", "id_jag_resource"]}
>
<Input placeholder="https://upstream.example.com/mcp" className={fieldClassName} />
</Form.Item>
<Form.Item
label={
<FieldLabel
label="Subject Token Type (optional)"
tooltip="Type of the identity assertion exchanged on leg 1. Defaults to urn:ietf:params:oauth:token-type:id_token."
/>
}
name="subject_token_type"
>
<Input placeholder="urn:ietf:params:oauth:token-type:id_token" className={fieldClassName} />
</Form.Item>
<Form.Item
label={<FieldLabel label="Scopes (optional)" tooltip="Scopes requested on leg 1 of the exchange." />}
name={["credentials", "scopes"]}
>
<Select mode="tags" tokenSeparators={[","]} placeholder="Add scopes" className="rounded-lg" size="large" />
</Form.Item>
</>
);
};
export default IdJagFormFields;

View file

@ -1058,6 +1058,127 @@ describe("CreateMCPServer", () => {
});
});
it("shows the ID-JAG fields only for the ID-JAG auth type", async () => {
await selectHttpTransport();
// The sibling OBO mode must not render the ID-JAG section.
await selectAntOption("Authentication", "OAuth Token Exchange (OBO)");
await waitFor(() => {
expect(screen.queryByText("Org Token Endpoint (leg 1)")).not.toBeInTheDocument();
});
expect(screen.queryByText("Resource Token Endpoint (leg 2)")).not.toBeInTheDocument();
await selectAntOption("Authentication", "ID-JAG (Okta Cross App Access)");
await waitFor(() => {
expect(screen.getByText("Org Token Endpoint (leg 1)")).toBeInTheDocument();
});
expect(screen.getByText("Resource Token Endpoint (leg 2)")).toBeInTheDocument();
expect(screen.getByText("Client Private Key (PEM)")).toBeInTheDocument();
await selectAntOption("Authentication", "API Key");
await waitFor(() => {
expect(screen.queryByText("Org Token Endpoint (leg 1)")).not.toBeInTheDocument();
});
expect(screen.queryByText("Resource Token Endpoint (leg 2)")).not.toBeInTheDocument();
});
it("routes ID-JAG config to the backend payload with both legs and the private key", async () => {
await selectHttpTransport();
fireEvent.change(getServerNameInput(), { target: { value: "IdJag_Server" } });
fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
target: { value: "https://upstream.example.com/mcp" },
});
await selectAntOption("Authentication", "ID-JAG (Okta Cross App Access)");
await waitFor(() => {
expect(screen.getByPlaceholderText("https://your-org.okta.com/oauth2/v1/token")).toBeInTheDocument();
});
fireEvent.change(screen.getByPlaceholderText("https://your-org.okta.com/oauth2/v1/token"), {
target: { value: "https://acme.okta.com/oauth2/v1/token" },
});
fireEvent.change(screen.getByPlaceholderText("https://upstream.example.com/oauth2/token"), {
target: { value: "https://jira.example.com/oauth2/token" },
});
fireEvent.change(screen.getByPlaceholderText("Enter OAuth client ID"), {
target: { value: "id-jag-client" },
});
fireEvent.change(screen.getByPlaceholderText("-----BEGIN PRIVATE KEY-----"), {
target: { value: "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----" },
});
fireEvent.change(screen.getByPlaceholderText("my-signing-key-1"), {
target: { value: "kid-1" },
});
vi.mocked(networking.createMCPServer).mockResolvedValue({
server_id: "new-server-id-jag",
server_name: "IdJag_Server",
alias: "IdJag_Server",
url: "https://upstream.example.com/mcp",
transport: "http",
auth_type: "oauth2_id_jag",
created_at: "2024-01-01T00:00:00Z",
created_by: "user-1",
updated_at: "2024-01-01T00:00:00Z",
updated_by: "user-1",
});
await act(async () => {
fireEvent.click(screen.getByRole("button", { name: "Add MCP Server" }));
});
await waitFor(() => {
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
});
const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0];
expect(payload.auth_type).toBe("oauth2_id_jag");
// Leg 1 rides the shared token_exchange_endpoint column; leg 2 is ID-JAG specific.
expect(payload.token_exchange_endpoint).toBe("https://acme.okta.com/oauth2/v1/token");
expect(payload.credentials).toMatchObject({
client_id: "id-jag-client",
id_jag_resource_token_endpoint: "https://jira.example.com/oauth2/token",
client_private_key: "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----",
client_private_key_id: "kid-1",
});
});
it("blocks an ID-JAG submit that provides neither a client secret nor a private key", async () => {
await selectHttpTransport();
fireEvent.change(getServerNameInput(), { target: { value: "IdJag_NoCreds" } });
fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
target: { value: "https://upstream.example.com/mcp" },
});
await selectAntOption("Authentication", "ID-JAG (Okta Cross App Access)");
await waitFor(() => {
expect(screen.getByPlaceholderText("https://your-org.okta.com/oauth2/v1/token")).toBeInTheDocument();
});
fireEvent.change(screen.getByPlaceholderText("https://your-org.okta.com/oauth2/v1/token"), {
target: { value: "https://acme.okta.com/oauth2/v1/token" },
});
fireEvent.change(screen.getByPlaceholderText("https://upstream.example.com/oauth2/token"), {
target: { value: "https://jira.example.com/oauth2/token" },
});
fireEvent.change(screen.getByPlaceholderText("Enter OAuth client ID"), {
target: { value: "id-jag-client" },
});
await act(async () => {
fireEvent.click(screen.getByRole("button", { name: "Add MCP Server" }));
});
await waitFor(() => {
expect(screen.getByText("Provide either a client secret or a client private key")).toBeInTheDocument();
});
expect(networking.createMCPServer).not.toHaveBeenCalled();
});
it("makes scope required when the Entra OBO profile is selected", async () => {
await selectHttpTransport();

View file

@ -26,6 +26,7 @@ import OAuthFormFields from "./OAuthFormFields";
import TruePassthroughWarning from "./TruePassthroughWarning";
import PassthroughAuthorizeSection from "./PassthroughAuthorizeSection";
import TokenExchangeFormFields from "./TokenExchangeFormFields";
import IdJagFormFields from "./IdJagFormFields";
import MCPServerCostConfig from "./mcp_server_cost_config";
import MCPConnectionStatus from "./mcp_connection_status";
import MCPToolConfiguration from "./mcp_tool_configuration";
@ -61,6 +62,7 @@ const AUTH_TYPES_REQUIRING_CREDENTIALS = [
...AUTH_TYPES_REQUIRING_AUTH_VALUE,
AUTH_TYPE.OAUTH2,
AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE,
AUTH_TYPE.OAUTH2_ID_JAG,
AUTH_TYPE.AWS_SIGV4,
AUTH_TYPE.TRUE_PASSTHROUGH,
AUTH_TYPE.OAUTH_DELEGATE,
@ -140,6 +142,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
const shouldShowAuthValueField = authType ? AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(authType) : false;
const isOAuthAuthType = authType === AUTH_TYPE.OAUTH2;
const isTokenExchangeAuthType = authType === AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE;
const isIdJagAuthType = authType === AUTH_TYPE.OAUTH2_ID_JAG;
const isAwsSigV4AuthType = authType === AUTH_TYPE.AWS_SIGV4;
const isM2MFlow = isOAuthAuthType && formValues.oauth_flow_type === OAUTH_FLOW.M2M;
@ -1071,7 +1074,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
children: (
<>
<Form.Item name="auth_type" rules={[{ required: true, message: "Please select an auth type" }]}>
<Select placeholder="Select auth type" className="rounded-lg" size="large">
<Select placeholder="Select auth type" className="rounded-lg" size="large" virtual={false}>
<Select.Option value="none">None</Select.Option>
<Select.Option value="api_key">API Key</Select.Option>
<Select.Option value="bearer_token">Bearer Token</Select.Option>
@ -1079,6 +1082,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
<Select.Option value="basic">Basic Auth</Select.Option>
<Select.Option value="oauth2">OAuth</Select.Option>
<Select.Option value="oauth2_token_exchange">OAuth Token Exchange (OBO)</Select.Option>
<Select.Option value="oauth2_id_jag">ID-JAG (Okta Cross App Access)</Select.Option>
<Select.Option value="aws_sigv4">AWS SigV4 (Bedrock AgentCore MCPs)</Select.Option>
<Select.Option value="true_passthrough">True Passthrough (no LiteLLM auth)</Select.Option>
<Select.Option value="oauth_delegate">
@ -1144,6 +1148,8 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
)}
{isTokenExchangeAuthType && <TokenExchangeFormFields />}
{isIdJagAuthType && <IdJagFormFields />}
</>
),
},

View file

@ -36,6 +36,7 @@ import PassthroughAuthorizeSection from "./PassthroughAuthorizeSection";
import MCPToolConfiguration from "./mcp_tool_configuration";
import StdioConfiguration from "./StdioConfiguration";
import TokenExchangeFormFields from "./TokenExchangeFormFields";
import IdJagFormFields from "./IdJagFormFields";
import OAuthFormFields from "./OAuthFormFields";
import MCPLogoSelector from "./MCPLogoSelector";
import EnvVarsSection from "./EnvVarsSection";
@ -64,6 +65,7 @@ const AUTH_TYPES_REQUIRING_CREDENTIALS = [
...AUTH_TYPES_REQUIRING_AUTH_VALUE,
AUTH_TYPE.OAUTH2,
AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE,
AUTH_TYPE.OAUTH2_ID_JAG,
AUTH_TYPE.AWS_SIGV4,
AUTH_TYPE.TRUE_PASSTHROUGH,
AUTH_TYPE.OAUTH_DELEGATE,
@ -103,6 +105,7 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
const shouldShowAuthValueField = authType ? AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(authType) : false;
const isOAuthAuthType = authType === AUTH_TYPE.OAUTH2;
const isTokenExchangeAuthType = authType === AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE;
const isIdJagAuthType = authType === AUTH_TYPE.OAUTH2_ID_JAG;
const isAwsSigV4AuthType = authType === AUTH_TYPE.AWS_SIGV4;
const oauthFlowTypeValue = Form.useWatch("oauth_flow_type", form) as string | undefined;
const isM2MFlow = isOAuthAuthType && oauthFlowTypeValue === OAUTH_FLOW.M2M;
@ -1111,7 +1114,7 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
{!isStdioTransport && (
<>
<Form.Item label="Authentication" name="auth_type" rules={[{ required: true }]}>
<Select>
<Select virtual={false}>
<Select.Option value="none">None</Select.Option>
<Select.Option value="api_key">API Key</Select.Option>
<Select.Option value="bearer_token">Bearer Token</Select.Option>
@ -1119,6 +1122,7 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
<Select.Option value="basic">Basic Auth</Select.Option>
<Select.Option value="oauth2">OAuth</Select.Option>
<Select.Option value="oauth2_token_exchange">OAuth Token Exchange (OBO)</Select.Option>
<Select.Option value="oauth2_id_jag">ID-JAG (Okta Cross App Access)</Select.Option>
<Select.Option value="aws_sigv4">AWS SigV4 (Bedrock AgentCore MCPs)</Select.Option>
<Select.Option value="true_passthrough">True Passthrough (no LiteLLM auth)</Select.Option>
<Select.Option value="oauth_delegate">
@ -1257,6 +1261,8 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
{!isStdioTransport && isTokenExchangeAuthType && <TokenExchangeFormFields isEditing />}
{!isStdioTransport && isIdJagAuthType && <IdJagFormFields isEditing />}
{!isStdioTransport && isAwsSigV4AuthType && (
<>
<p className="text-sm text-gray-500 mb-2">

View file

@ -957,25 +957,23 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
placeholder="Select vector stores (optional)"
/>
</Form.Item>
<Form.Item label="Allowed Pass Through Routes" name="allowed_passthrough_routes" className="mt-8">
<Tooltip
title={
!premiumUser
? "Premium feature - Upgrade to set allowed pass through routes"
: !isProxyAdminRole(userRole || "")
? "Only proxy admins can set allowed pass through routes"
: ""
}
placement="top"
>
<PassThroughRoutesSelector
onChange={(values: string[]) => form.setFieldValue("allowed_passthrough_routes", values)}
value={form.getFieldValue("allowed_passthrough_routes")}
accessToken={accessToken || ""}
placeholder="Select pass through routes (optional)"
disabled={!premiumUser || !isProxyAdminRole(userRole || "")}
/>
</Tooltip>
<Form.Item
label="Allowed Pass Through Routes"
name="allowed_passthrough_routes"
className="mt-8"
tooltip={
!premiumUser
? "Premium feature - Upgrade to set allowed pass through routes"
: !isProxyAdminRole(userRole || "")
? "Only proxy admins can set allowed pass through routes"
: undefined
}
>
<PassThroughRoutesSelector
accessToken={accessToken || ""}
placeholder="Select pass through routes (optional)"
disabled={!premiumUser || !isProxyAdminRole(userRole || "")}
/>
</Form.Item>
</AccordionBody>
</Accordion>

View file

@ -196,6 +196,19 @@ describe("server-side filtering maps controls to the right query params", () =>
expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: "platform" }));
});
});
it("opts into team id prefix matching so a partial id from a proxy error finds the team", async () => {
renderTable();
fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "66c432fa" } });
await waitFor(() => {
expect(mockUseTeamsTable).toHaveBeenLastCalledWith(
1,
50,
expect.objectContaining({ search: "66c432fa", searchTeamIdMatch: "prefix" }),
);
});
});
});
describe("non-admin scoping", () => {

View file

@ -66,6 +66,7 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet
team_alias: getFilterValue("alias"),
teamID: getFilterValue("team_id"),
search: searchQuery.trim() || undefined,
searchTeamIdMatch: "prefix" as const,
userID: isAdminView ? undefined : userID ?? undefined,
sortBy: sorting[0]?.id,
sortOrder: toSortOrder(sorting),

View file

@ -17,9 +17,25 @@ import {
formatKeywords,
parseSkillSource,
isValidSubPath,
buildMarketplaceSettingsSnippet,
} from "./helpers";
import { MarketplacePluginEntry, PluginSource } from "./types";
describe("buildMarketplaceSettingsSnippet", () => {
it("nests the url under a source object so Claude Code accepts the marketplace", () => {
expect(JSON.parse(buildMarketplaceSettingsSnippet("https://proxy.example.com"))).toEqual({
extraKnownMarketplaces: {
"my-org": {
source: {
source: "url",
url: "https://proxy.example.com/claude-code/marketplace.json",
},
},
},
});
});
});
describe("formatInstallCommand", () => {
it("formats github source with repo", () => {
const source: PluginSource = { source: "github", repo: "org/repo" };

View file

@ -176,6 +176,27 @@ export const parseSkillSource = (rawUrl: string, subPath?: string): SkillSourceP
return parseRawGitSource(url, subPath);
};
/**
* Build the `~/.claude/settings.json` snippet that registers the proxy as a marketplace.
* Claude Code expects `extraKnownMarketplaces.<name>.source` to be a source object, not a
* bare `"url"` string, so the url/source pair is nested one level deeper.
*/
export const buildMarketplaceSettingsSnippet = (proxyOrigin: string): string =>
JSON.stringify(
{
extraKnownMarketplaces: {
"my-org": {
source: {
source: "url",
url: `${proxyOrigin}/claude-code/marketplace.json`,
},
},
},
},
null,
2,
);
/**
* Generate install command for Claude Code CLI
* Format: /plugin marketplace add org/repo OR /plugin marketplace add url

View file

@ -1,6 +1,6 @@
import React, { useState } from "react";
import { ArrowLeftOutlined, CopyOutlined, CheckOutlined, LinkOutlined } from "@ant-design/icons";
import { formatInstallCommand } from "./helpers";
import { buildMarketplaceSettingsSnippet, formatInstallCommand } from "./helpers";
import { Plugin } from "./types";
interface SkillDetailProps {
@ -31,6 +31,10 @@ const SkillDetail: React.FC<SkillDetailProps> = ({ skill, onBack }) => {
const installCommand = formatInstallCommand(skill);
const settingsSnippet = buildMarketplaceSettingsSnippet(
typeof window !== "undefined" ? window.location.origin : "<proxy-url>",
);
const detailRows = [
...(skill.category ? [{ property: "Category", value: skill.category }] : []),
...(skill.domain ? [{ property: "Domain", value: skill.domain }] : []),
@ -298,21 +302,7 @@ const SkillDetail: React.FC<SkillDetailProps> = ({ skill, onBack }) => {
>
<span style={{ fontSize: 13, color: "#3c4043", fontWeight: 500 }}>~/.claude/settings.json</span>
<button
onClick={() => {
const snippet = JSON.stringify(
{
extraKnownMarketplaces: {
"my-org": {
source: "url",
url: `${typeof window !== "undefined" ? window.location.origin : ""}/claude-code/marketplace.json`,
},
},
},
null,
2,
);
copyToClipboard(snippet, "settings");
}}
onClick={() => copyToClipboard(settingsSnippet, "settings")}
style={{
display: "flex",
alignItems: "center",
@ -339,18 +329,7 @@ const SkillDetail: React.FC<SkillDetailProps> = ({ skill, onBack }) => {
backgroundColor: "#fff",
}}
>
{JSON.stringify(
{
extraKnownMarketplaces: {
"my-org": {
source: "url",
url: `${typeof window !== "undefined" ? window.location.origin : "<proxy-url>"}/claude-code/marketplace.json`,
},
},
},
null,
2,
)}
{settingsSnippet}
</pre>
</div>
</div>

View file

@ -3,7 +3,7 @@ import { Select } from "antd";
import { getPassThroughEndpointsCall } from "../networking";
interface PassThroughRoutesSelectorProps {
onChange: (selectedRoutes: string[]) => void;
onChange?: (selectedRoutes: string[]) => void;
value?: string[];
className?: string;
accessToken: string;

View file

@ -40,6 +40,7 @@ export const AUTH_TYPE = {
BASIC: "basic",
OAUTH2: "oauth2",
OAUTH2_TOKEN_EXCHANGE: "oauth2_token_exchange",
OAUTH2_ID_JAG: "oauth2_id_jag",
AWS_SIGV4: "aws_sigv4",
TRUE_PASSTHROUGH: "true_passthrough",
OAUTH_DELEGATE: "oauth_delegate",
@ -66,11 +67,12 @@ export const gatewayMintsClientFor = (server: { auth_type?: string | null; dcr_b
(server.auth_type === AUTH_TYPE.OAUTH_DELEGATE && !server.dcr_bridge);
// Auth modes that cannot be used through the gateway aggregate connect flow, where the client holds
// only an identity-only session bearer and upstream credentials are resolved server-side from the
// per-user vault. The vault is only populated by interactive authorization_code (oauth2). The
// client-forwarded modes need the caller to present the upstream Authorization per call, and
// only an identity-only session bearer and upstream credentials are resolved server-side per user.
// The client-forwarded modes need the caller to present the upstream Authorization per call, and
// oauth2_token_exchange (OBO) needs the caller's own IdP token as the subject to exchange; the
// session bearer is neither, so none of these can complete a tool call on this connection.
// oauth2_id_jag is deliberately NOT here: it falls back to the identity assertion captured for the
// session's user at SSO login, so the identity-only bearer is enough to resolve it server-side.
export const isUnsupportedOnGatewayConnect = (authType?: string | null): boolean =>
isClientForwardedTokenMode(authType) || authType === AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE;

View file

@ -1385,8 +1385,6 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
}
>
<PassThroughRoutesSelector
onChange={(values: string[]) => form.setFieldValue("allowed_passthrough_routes", values)}
value={form.getFieldValue("allowed_passthrough_routes")}
accessToken={accessToken}
placeholder={
!premiumUser

View file

@ -625,6 +625,44 @@ describe("DataTable layout", () => {
const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement;
expect(scroller.style.maxHeight).toBe("240px");
});
it("caps fillHeight at the parent's height instead of stretching to it, so a short table stays short", () => {
const { container } = render(<DataTable data={CHARLIE_ALICE_BOB} columns={nameEmailColumns} fillHeight />);
const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement;
const frame = scroller.parentElement as HTMLElement;
const outer = frame.parentElement as HTMLElement;
// A ceiling, not a stretch: flex-1 here would hold the footer at the bottom on a two-row table.
expect(outer.className).toContain("max-h-full");
expect(outer.className).not.toContain("flex-1");
expect(frame.className).not.toContain("flex-1");
expect(scroller.className).not.toContain("flex-1");
expect(outer.className).toContain("flex-col");
expect(frame.className).toContain("flex-col");
expect(scroller.className).toContain("min-h-0");
expect(scroller.className).toContain("overflow-auto");
expect(scroller.style.maxHeight).toBe("");
// Without this the Table primitive's own overflow container captures the sticky header.
expect(scroller.className).toContain("[&_[data-slot=table-container]]:overflow-visible");
const thead = container.querySelector("thead") as HTMLElement;
expect(thead.className).toContain("sticky");
// Rows pass under the header, so the semi-transparent row tint alone would let them show through.
expect(thead.className).toContain("bg-background");
});
it("leaves the default layout untouched when neither height mode is set", () => {
const { container } = render(<DataTable data={CHARLIE_ALICE_BOB} columns={nameEmailColumns} />);
const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement;
expect(scroller.className).toContain("overflow-x-auto");
expect(scroller.className).not.toContain("min-h-0");
expect(scroller.style.maxHeight).toBe("");
expect((scroller.parentElement as HTMLElement).className).not.toContain("flex-col");
expect(container.querySelector("thead")?.className).not.toContain("sticky");
expect(container.querySelector("thead")?.className).not.toContain("bg-background");
});
});
describe("DataTable misconfiguration guards", () => {

Some files were not shown because too many files have changed in this diff Show more