mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
Merge branch 'litellm_internal_staging' into fix-vertex-typearray-null
This commit is contained in:
commit
2505fe5b78
114 changed files with 7972 additions and 936 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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={},
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
#########################################################
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
205
litellm/proxy/management_endpoints/management_v1/budgets.py
Normal file
205
litellm/proxy/management_endpoints/management_v1/budgets.py
Normal 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.",
|
||||
)
|
||||
)
|
||||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)"}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 ----------
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"})
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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"
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]}
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
</>
|
||||
)}
|
||||
/>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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
|
||||
? [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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">
|
||||
|
|
|
|||
|
|
@ -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({});
|
||||
});
|
||||
});
|
||||
|
|
@ -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));
|
||||
|
|
@ -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 = () => {
|
||||
|
|
|
|||
|
|
@ -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([]);
|
||||
});
|
||||
});
|
||||
|
|
@ -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,
|
||||
};
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
@ -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();
|
||||
|
||||
|
|
|
|||
|
|
@ -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 />}
|
||||
</>
|
||||
),
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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">
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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", () => {
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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" };
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Reference in a new issue