diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index 90ede5a653f..8d791ca5bc7 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -58,6 +58,8 @@ jobs: # free OSS, run as a pinned, checksum-verified binary; no GitHub Action # dependency and no vendor SaaS callout. - name: Scan image for fixable HIGH/CRITICAL CVEs + env: + GRYPE_MATCH_PYTHON_USING_CPES: "true" run: | "$RUNNER_TEMP/grype" litellm-image-scan:${{ github.sha }} \ --only-fixed \ diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 438ff5600ba..79036367652 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -7,8 +7,8 @@ duration_in_seconds is used in diff parts of the code base, example """ import re -import time -from datetime import datetime, timedelta, timezone, tzinfo +import time as time_module +from datetime import datetime, time, timedelta, timezone, tzinfo from typing import Optional, Tuple from zoneinfo import ZoneInfo @@ -61,7 +61,7 @@ def duration_in_seconds(duration: str) -> int: elif unit == "w": return value * 604800 elif unit == "mo": - now = time.time() + now = time_module.time() current_time = datetime.fromtimestamp(now) # Calculate target month and year, handling overflow past December @@ -94,12 +94,17 @@ def duration_in_seconds(duration: str) -> int: raise ValueError(f"Unsupported duration unit, passed duration: {duration}") -def get_next_standardized_reset_time(duration: str, current_time: datetime, timezone_str: str = "UTC") -> datetime: +def get_next_standardized_reset_time( + duration: str, + current_time: datetime, + timezone_str: str = "UTC", + reset_time_of_day: time = time(0, 0), +) -> datetime: """ Get the next standardized reset time based on the duration. All durations will reset at predictable intervals, aligned from the current time: - - Nd: If N=1, reset at next midnight; if N>1, reset every N days from now + - Nd: If N=1, reset at the next `reset_time_of_day`; if N>1, reset every N days from now - Nh: Every N hours, aligned to hour boundaries (e.g., 1:00, 2:00) - Nm: Every N minutes, aligned to minute boundaries (e.g., 1:05, 1:10) - Ns: Every N seconds, aligned to second boundaries @@ -108,12 +113,15 @@ def get_next_standardized_reset_time(duration: str, current_time: datetime, time - duration: Duration string (e.g. "30s", "30m", "30h", "30d") - current_time: Current datetime - timezone_str: Timezone string (e.g. "UTC", "US/Eastern", "Asia/Kolkata") + - reset_time_of_day: Wall-clock time the reset lands on for day/week/month + durations (defaults to midnight). Ignored for sub-day durations, where a + time-of-day is meaningless. Returns: - Next reset time at a standardized interval in the specified timezone """ # Set up timezone and normalize current time - current_time, tz = _setup_timezone(current_time, timezone_str) + current_time, _ = _setup_timezone(current_time, timezone_str) # Parse duration value, unit = _parse_duration(duration) @@ -126,9 +134,9 @@ def get_next_standardized_reset_time(duration: str, current_time: datetime, time # Handle different time units if unit == "d": - return _handle_day_reset(current_time, base_midnight, value, tz) + return _handle_day_reset(current_time, base_midnight, value, reset_time_of_day) elif unit == "w": - return _handle_day_reset(current_time, base_midnight, value * 7, tz) + return _handle_day_reset(current_time, base_midnight, value * 7, reset_time_of_day) elif unit == "h": return _handle_hour_reset(current_time, base_midnight, value) elif unit == "m": @@ -136,7 +144,7 @@ def get_next_standardized_reset_time(duration: str, current_time: datetime, time elif unit == "s": return _handle_second_reset(current_time, base_midnight, value) elif unit == "mo": - return _handle_month_reset(current_time, base_midnight, value) + return _handle_month_reset(current_time, base_midnight, value, reset_time_of_day) else: # Unrecognized unit, default to next midnight return base_midnight + timedelta(days=1) @@ -175,46 +183,58 @@ def _parse_duration(duration: str) -> Tuple[Optional[int], Optional[str]]: return int(value), unit -def _handle_day_reset(current_time: datetime, base_midnight: datetime, value: int, tz: tzinfo) -> datetime: +def _apply_time_of_day(dt: datetime, reset_time_of_day: time) -> datetime: + """Set the wall-clock time of `dt` to `reset_time_of_day`, keeping its date and tzinfo.""" + return dt.replace( + hour=reset_time_of_day.hour, + minute=reset_time_of_day.minute, + second=reset_time_of_day.second, + microsecond=reset_time_of_day.microsecond, + ) + + +def _next_occurrence( + boundary_midnight: datetime, + reset_time_of_day: time, + current_time: datetime, + period: timedelta, +) -> datetime: + """Place the reset at `reset_time_of_day` on the boundary day, rolling forward one + `period` if that instant has already passed (or is exactly now).""" + candidate = _apply_time_of_day(boundary_midnight, reset_time_of_day) + if candidate <= current_time: + return candidate + period + return candidate + + +def _first_of_next_month(first_of_month: datetime) -> datetime: + """Given the 1st of some month, return the 1st of the following month.""" + if first_of_month.month == 12: + return first_of_month.replace(year=first_of_month.year + 1, month=1) + return first_of_month.replace(month=first_of_month.month + 1) + + +def _handle_day_reset( + current_time: datetime, + base_midnight: datetime, + value: int, + reset_time_of_day: time, +) -> datetime: """Handle day-based reset times.""" # Handle zero value - immediate expiration if value == 0: return current_time - if value == 1: # Daily reset at midnight - return base_midnight + timedelta(days=1) - elif value == 7: # Weekly reset on Monday at midnight + if value == 1: # Daily reset at the configured time of day + return _next_occurrence(base_midnight, reset_time_of_day, current_time, timedelta(days=1)) + elif value == 7: # Weekly reset on Monday at the configured time of day days_until_monday = (7 - current_time.weekday()) % 7 - if days_until_monday == 0: # If today is Monday - days_until_monday = 7 - return base_midnight + timedelta(days=days_until_monday) - elif value == 30: # Monthly reset on 1st at midnight - # Get 1st of next month at midnight - if current_time.month == 12: - next_reset = datetime( - year=current_time.year + 1, - month=1, - day=1, - hour=0, - minute=0, - second=0, - microsecond=0, - tzinfo=tz, - ) - else: - next_reset = datetime( - year=current_time.year, - month=current_time.month + 1, - day=1, - hour=0, - minute=0, - second=0, - microsecond=0, - tzinfo=tz, - ) - return next_reset - else: # Custom day value - next interval is value days from current - return current_time.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=value) + upcoming_monday = base_midnight + timedelta(days=days_until_monday) + return _next_occurrence(upcoming_monday, reset_time_of_day, current_time, timedelta(days=7)) + elif value == 30: # Monthly reset on 1st at the configured time of day + return _handle_month_reset(current_time, base_midnight, 1, reset_time_of_day) + else: # Custom day value - next interval is value days from the start of today + return _apply_time_of_day(base_midnight + timedelta(days=value), reset_time_of_day) def _handle_hour_reset(current_time: datetime, base_midnight: datetime, value: int) -> datetime: @@ -316,36 +336,30 @@ def _handle_second_reset(current_time: datetime, base_midnight: datetime, value: return current_time.replace(hour=next_hour, minute=next_minute, second=next_second, microsecond=0) -def _handle_month_reset(current_time: datetime, base_midnight: datetime, value: int) -> datetime: +def _handle_month_reset( + current_time: datetime, + base_midnight: datetime, + value: int, + reset_time_of_day: time, +) -> datetime: """ - Handle monthly reset times. For monthly resets, we always reset at the start of the next month. + Handle monthly reset times. Resets land on the 1st at `reset_time_of_day`; if the + 1st of the current month at that time has already passed, roll to the 1st of next month. Args: current_time: Current datetime base_midnight: Midnight of current day value: Number of months (currently only supports 1 month resets) + reset_time_of_day: Wall-clock time the reset lands on Returns: - datetime: First day of next month at midnight + datetime: First day of the next reset month at `reset_time_of_day` """ if value != 1: raise ValueError("Monthly resets currently only support 1 month intervals") - # Get the first day of next month - if current_time.month == 12: - next_month = 1 - next_year = current_time.year + 1 - else: - next_month = current_time.month + 1 - next_year = current_time.year - - return datetime( - year=next_year, - month=next_month, - day=1, - hour=0, - minute=0, - second=0, - microsecond=0, - tzinfo=current_time.tzinfo, - ) + first_of_this_month = base_midnight.replace(day=1) + candidate = _apply_time_of_day(first_of_this_month, reset_time_of_day) + if candidate <= current_time: + return _apply_time_of_day(_first_of_next_month(first_of_this_month), reset_time_of_day) + return candidate diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8f30071eb5d..ff714b8de22 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -610,6 +610,34 @@ def _passthrough_token_from_mcp_auth_header( return None +async def _materialize_auth_headers(auth: httpx.Auth | None) -> dict[str, str] | None: + """Extract the header a resolved ``httpx.Auth`` would set, as a plain dict, or None. + + OpenAPI tool closures egress through ``AsyncHTTPHandler`` methods that accept headers but no + ``auth``, so a resolved credential must be materialized into a header value. Driving one step + of the auth's own flow (against a throwaway request that is never sent) keeps this generic + across every auth shape without per-class branching; ``header_name`` is the resolver-arm + convention for "this auth sets a header" (``NoOpAuth`` has none and yields nothing to apply). + The materialized value is point-in-time: flow behaviors past the first request, like the M2M + one-shot 401 refetch, do not apply on this arm. + """ + if auth is None: + return None + header_name = getattr(auth, "header_name", None) + if not isinstance(header_name, str) or not header_name: + return None + probe = httpx.Request("GET", "http://localhost/") + flow = auth.async_auth_flow(probe) + try: + first_request = await flow.__anext__() + except StopAsyncIteration: + return None + finally: + await flow.aclose() + header_value = first_request.headers.get(header_name) + return {header_name: header_value} if header_value else None + + def _consumes_caller_authorization(server: MCPServer) -> bool: """True when this server's egress forwards the caller's request-wide ``Authorization`` upstream: the client-forwarded token modes, legacy OAuth pass-through, and legacy upstream-delegated @@ -4705,6 +4733,61 @@ class MCPServerManager: ) return oauth2_headers + async def resolve_openapi_upstream_auth( + self, + *, + mcp_server: MCPServer, + oauth2_headers: dict[str, str] | None, + raw_headers: dict[str, str] | None, + mcp_auth_header: str | dict[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, + forwarded_headers: dict[str, str] | None, + ) -> tuple[dict[str, str] | None, dict[str, str] | None]: + """Resolve the gateway-owned upstream credential for a spec_path (OpenAPI) tool call. + + OpenAPI tools egress through a plain httpx call assembled from ContextVars, never through + ``_create_mcp_client``, so the v2 resolver graft there does not run for them and a resolved + credential (authorization_code's stored per-user token, client_credentials' minted M2M + token, token_exchange's exchanged token, passthrough's forwarded caller token) must be + materialized into headers here. Returns ``(resolved_auth_headers, forwarded_headers)``: + the resolved headers are authoritative over every other Authorization source (the same + rule ``_resolve_v2_auth`` applies on the MCPClient path) and ``forwarded_headers`` comes + back with any header the resolver claimed already dropped. Unmigrated (v1) servers resolve + through the stored-token lookup instead, and a missing per-user credential raises the same + discovery challenge the MCPClient path serves, rather than egressing unauthenticated. + + The resolved headers carry only credentials the gateway itself resolved (a stored per-user + token, a minted or exchanged token). Caller-supplied ``oauth2_headers`` are never promoted + into them: on the v2 arm they feed only subject-token extraction (the designed RFC 8693 + input), and on the v1 arm their presence disables the stored lookup entirely, so a + caller's gateway credential can never displace a per-server BYOK header or leak upstream + as the resolved credential. + """ + spec = to_server_spec(mcp_server) + if spec is None: + if oauth2_headers: + return None, forwarded_headers + stored_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) + return stored_headers, forwarded_headers + + subject_token: str | None = None + if isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)): + subject_token = self._extract_bearer_token(oauth2_headers, raw_headers) + elif isinstance(spec.config, PassthroughConfig): + inbound_token, forwarded_headers = _take_forwarded_authorization(forwarded_headers) + per_server_token = _passthrough_token_from_mcp_auth_header(mcp_auth_header) + subject_token = per_server_token if per_server_token is not None else inbound_token + + resolved_auth, forwarded_headers = await self._resolve_v2_auth( + server=mcp_server, + spec=spec, + provider=self._cred_provider, + subject_token=subject_token, + user_api_key_auth=user_api_key_auth, + extra_headers=forwarded_headers, + ) + return await _materialize_auth_headers(resolved_auth), forwarded_headers + async def _gather_openapi_tool_tasks( self, tasks: list[Any], @@ -4796,6 +4879,7 @@ class MCPServerManager: ) tasks.append(during_hook_task) + caller_oauth2_headers = oauth2_headers oauth2_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, oauth2_headers, user_api_key_auth) # For OpenAPI servers, call the tool handler directly instead of via MCP client @@ -4813,22 +4897,32 @@ class MCPServerManager: auth_header_value = ( _format_byok_openapi_auth_header(mcp_server, mcp_auth_header) if mcp_auth_header else None ) - forwarded_headers = _openapi_forwarded_extra_headers(mcp_server, raw_headers, user_api_key_auth) + resolved_auth_headers, forwarded_headers = await self.resolve_openapi_upstream_auth( + mcp_server=mcp_server, + oauth2_headers=caller_oauth2_headers, + raw_headers=raw_headers, + mcp_auth_header=mcp_auth_header, + user_api_key_auth=user_api_key_auth, + forwarded_headers=_openapi_forwarded_extra_headers(mcp_server, raw_headers, user_api_key_auth), + ) async def _call_openapi_via_handler(): from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, _request_extra_headers, + _request_resolved_auth_headers, ) auth_token = _request_auth_header.set(auth_header_value) extra_token = _request_extra_headers.set(forwarded_headers) + resolved_token = _request_resolved_auth_headers.set(resolved_auth_headers) try: async with self._limit_outbound_concurrency(mcp_server): return await self._call_openapi_tool_handler(mcp_server, name, arguments) finally: _request_auth_header.reset(auth_token) _request_extra_headers.reset(extra_token) + _request_resolved_auth_headers.reset(resolved_token) tasks.append(asyncio.create_task(_call_openapi_via_handler())) else: diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 1ee300be718..0b795057837 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -62,6 +62,14 @@ _request_extra_headers: contextvars.ContextVar[Optional[Dict[str, str]]] = conte "_request_extra_headers", default=None ) +# Per-request headers carrying the gateway-resolved upstream credential +# (stored per-user OAuth token, minted M2M token, exchanged OBO token). +# Set from MCPServerManager.resolve_openapi_upstream_auth; authoritative +# over every other Authorization source in _merge_openapi_tool_request_headers. +_request_resolved_auth_headers: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar( + "_request_resolved_auth_headers", default=None +) + def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str: """Ensure path params cannot introduce directory traversal.""" @@ -294,10 +302,15 @@ def _merge_openapi_tool_request_headers( """Merge static closure headers with per-request ContextVar overrides. Precedence (highest to lowest): - 1. ``_request_auth_header`` — BYOK override of ``Authorization`` - 2. ``static_headers`` — operator-configured headers baked into the + 1. ``_request_resolved_auth_headers`` — the gateway-resolved upstream + credential (stored per-user OAuth token, minted M2M token, + exchanged OBO token). The resolver is authoritative: a BYOK or + forwarded ``Authorization`` must not shadow it, mirroring + ``_resolve_v2_auth`` on the MCPClient path + 2. ``_request_auth_header`` — BYOK override of ``Authorization`` + 3. ``static_headers`` — operator-configured headers baked into the tool closure at registration time - 3. ``_request_extra_headers`` — per-request headers forwarded from + 4. ``_request_extra_headers`` — per-request headers forwarded from the MCP caller (allowlisted by ``MCPServer.extra_headers``) This matches the existing MCP invariant in @@ -323,6 +336,12 @@ def _merge_openapi_tool_request_headers( del effective_headers[existing] effective_headers["Authorization"] = override_auth + resolved_auth_headers = _request_resolved_auth_headers.get() or {} + for name, value in resolved_auth_headers.items(): + for existing in [k for k in effective_headers if k.lower() == name.lower()]: + del effective_headers[existing] + effective_headers[name] = value + return effective_headers diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a8ab0937124..396dd6c7dc7 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -376,6 +376,7 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, _request_extra_headers, + _request_resolved_auth_headers, ) from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport from litellm.proxy._experimental.mcp_server.tool_registry import ( @@ -2785,13 +2786,29 @@ if MCP_AVAILABLE: forwarded_headers = {} forwarded_headers[header_name] = value + resolved_auth_headers: dict[str, str] | None = None + if mcp_server: + ( + resolved_auth_headers, + forwarded_headers, + ) = await global_mcp_server_manager.resolve_openapi_upstream_auth( + mcp_server=mcp_server, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + mcp_auth_header=mcp_auth_header, + user_api_key_auth=user_api_key_auth, + forwarded_headers=forwarded_headers, + ) + _auth_token = _request_auth_header.set(auth_header_value) _extra_token = _request_extra_headers.set(forwarded_headers) + _resolved_token = _request_resolved_auth_headers.set(resolved_auth_headers) try: local_content = await _handle_local_mcp_tool(name, arguments) finally: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) + _request_resolved_auth_headers.reset(_resolved_token) response = CallToolResult(content=cast(Any, local_content), isError=False) # Try managed MCP server tool (pass the full prefixed name) diff --git a/litellm/proxy/a2a/agent_card.py b/litellm/proxy/a2a/agent_card.py index e97ab4a01ae..29a689a32de 100644 --- a/litellm/proxy/a2a/agent_card.py +++ b/litellm/proxy/a2a/agent_card.py @@ -7,23 +7,46 @@ the base; specific fields are replaced so all traffic flows through the proxy and uses LiteLLM auth. """ +import re from copy import deepcopy -from typing import Any, Dict, List, Mapping +from typing import Any, Dict, List, Literal, Mapping + +SupportedA2AVersion = Literal["0.3", "1.0"] # Protocol versions LiteLLM can serve to A2A clients. The admin pins one per agent; # responses are normalized to it regardless of the upstream agent's own version. -SUPPORTED_A2A_PROTOCOL_VERSIONS = ("0.3", "1.0") +SUPPORTED_A2A_PROTOCOL_VERSIONS: tuple[SupportedA2AVersion, ...] = ("0.3", "1.0") # Default served version when the agent card does not pin one. LITELLM_A2A_PROTOCOL_VERSION = "1.0" +_PROTOCOL_VERSION_PATTERN = re.compile( + r"^(\d+\.\d+)(?:\.\d+(?:-[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?(?:\+[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?)?$" +) + + +def normalize_protocol_version(version: object) -> SupportedA2AVersion | None: + """Map a raw ``protocolVersion`` value to the supported canonical major.minor version. + + Accepts the bare major.minor convention of the 1.0 spec (``"0.3"``, ``"1.0"``) and the + full semver forms older SDKs emit (``"0.3.0"``, ``"1.0.1"``, including prerelease and + build suffixes like ``"0.3.0-rc1"``). Malformed strings, versions outside the + supported set, and non-strings yield ``None``. + """ + if not isinstance(version, str): + return None + match = _PROTOCOL_VERSION_PATTERN.match(version) + if match is None: + return None + major_minor = match.group(1) + return next((supported for supported in SUPPORTED_A2A_PROTOCOL_VERSIONS if supported == major_minor), None) + + def resolve_served_protocol_version(card: Mapping[str, Any] | None) -> str: """Return the validated protocol version an agent card pins, else the default.""" - version = card.get("protocolVersion") if card else None - if version in SUPPORTED_A2A_PROTOCOL_VERSIONS: - return version - return LITELLM_A2A_PROTOCOL_VERSION + normalized = normalize_protocol_version(card.get("protocolVersion") if card else None) + return normalized if normalized is not None else LITELLM_A2A_PROTOCOL_VERSION # Security scheme exposed by the LiteLLM-fronted agent card. Always replaces diff --git a/litellm/proxy/a2a/version_convert.py b/litellm/proxy/a2a/version_convert.py index e8f49e6f6a9..9de33a0966a 100644 --- a/litellm/proxy/a2a/version_convert.py +++ b/litellm/proxy/a2a/version_convert.py @@ -30,6 +30,7 @@ from typing import Callable, Literal, Union from pydantic import BaseModel from litellm._logging import verbose_proxy_logger +from litellm.proxy.a2a.agent_card import normalize_protocol_version A2AVersion = Literal["0.3", "1.0"] RequestId = Union[str, int, None] @@ -103,16 +104,14 @@ def normalize_request_params(params: JsonDict, served: A2AVersion, *, method: st def _detect_card_version(card: JsonDict) -> A2AVersion: """Infer the wire version of an agent card dict. - ``protocolVersion`` is the authoritative indicator; fall back to presence of - ``supportedInterfaces`` (a 1.0-only field) only when the explicit field is absent. - Cards that set ``protocolVersion: "0.3"`` or carry neither signal are treated as 0.3. + ``protocolVersion`` is the authoritative indicator; semver values normalize to + their major.minor (``"0.3.0"`` -> ``"0.3"``). Fall back to presence of + ``supportedInterfaces`` (a 1.0-only field) only when the explicit field is + absent or unrecognized; cards carrying neither signal are treated as 0.3. """ - pv = card.get("protocolVersion") - if pv == "1.0": - return "1.0" - if pv == "0.3": - return "0.3" - # No protocolVersion field: use structural heuristic. + normalized = normalize_protocol_version(card.get("protocolVersion")) + if normalized is not None: + return normalized return "1.0" if "supportedInterfaces" in card else "0.3" diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index a7ceffed97b..2421f270974 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKey from litellm.proxy.a2a.agent_card import ( SUPPORTED_A2A_PROTOCOL_VERSIONS, merge_agent_card, + normalize_protocol_version, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user @@ -51,7 +52,7 @@ def _proxy_base_url(http_request: Request) -> str: def _validate_protocol_version(upstream_card: Mapping[str, Any] | None) -> None: """Reject an agent card pinning an unsupported A2A protocol version.""" version = upstream_card.get("protocolVersion") if upstream_card else None - if version is not None and version not in SUPPORTED_A2A_PROTOCOL_VERSIONS: + if version is not None and normalize_protocol_version(version) is None: raise HTTPException( status_code=400, detail=( diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index d48868960ce..0416a969569 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -14,6 +14,11 @@ from litellm.proxy._types import ( LiteLLM_UserTable, LiteLLM_VerificationToken, ) +from litellm.proxy.common_utils.timezone_utils import ( + BudgetResetSettings, + compute_budget_reset_at, + get_budget_reset_settings, +) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( @@ -33,9 +38,15 @@ class ResetBudgetJob: Resets the budget for all the keys, users, and teams that need it """ - def __init__(self, proxy_logging_obj: ProxyLogging, prisma_client: PrismaClient): + def __init__( + self, + proxy_logging_obj: ProxyLogging, + prisma_client: PrismaClient, + reset_settings: BudgetResetSettings | None = None, + ): self.proxy_logging_obj: ProxyLogging = proxy_logging_obj self.prisma_client: PrismaClient = prisma_client + self.reset_settings: BudgetResetSettings = reset_settings or get_budget_reset_settings() async def reset_budget( self, @@ -238,7 +249,7 @@ class ResetBudgetJob: if budgets_to_reset is not None and len(budgets_to_reset) > 0: for budget in budgets_to_reset: - budget = await ResetBudgetJob._reset_budget_reset_at_date(budget, now) + budget = await ResetBudgetJob._reset_budget_reset_at_date(budget, now, self.reset_settings) await self.prisma_client.update_data( query_type="update_many", @@ -443,7 +454,11 @@ class ResetBudgetJob: if keys_to_reset is not None and len(keys_to_reset) > 0: for key in keys_to_reset: try: - updated_key = await ResetBudgetJob._reset_budget_for_key(key=key, current_time=now) + updated_key = await ResetBudgetJob._reset_budget_for_key( + key=key, + current_time=now, + reset_settings=self.reset_settings, + ) if updated_key is not None: updated_keys.append(updated_key) else: @@ -514,7 +529,11 @@ class ResetBudgetJob: if users_to_reset is not None and len(users_to_reset) > 0: for user in users_to_reset: try: - updated_user = await ResetBudgetJob._reset_budget_for_user(user=user, current_time=now) + updated_user = await ResetBudgetJob._reset_budget_for_user( + user=user, + current_time=now, + reset_settings=self.reset_settings, + ) if updated_user is not None: updated_users.append(updated_user) else: @@ -589,7 +608,11 @@ class ResetBudgetJob: if teams_to_reset is not None and len(teams_to_reset) > 0: for team in teams_to_reset: try: - updated_team = await ResetBudgetJob._reset_budget_for_team(team=team, current_time=now) + updated_team = await ResetBudgetJob._reset_budget_for_team( + team=team, + current_time=now, + reset_settings=self.reset_settings, + ) if updated_team is not None: updated_teams.append(updated_team) else: @@ -656,10 +679,9 @@ class ResetBudgetJob: counter_key: str, spend_counter_cache: Any, now: datetime, + reset_settings: BudgetResetSettings, ) -> bool: """Reset a single budget window if expired. Returns True if the window was reset.""" - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - reset_at_str = window.get("reset_at") if not reset_at_str: return False @@ -671,7 +693,9 @@ class ResetBudgetJob: await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=0.0) except Exception as redis_err: verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err) - window["reset_at"] = get_budget_reset_time(budget_duration=window["budget_duration"]).isoformat() + window["reset_at"] = compute_budget_reset_at( + budget_duration=window["budget_duration"], settings=reset_settings + ).isoformat() return True async def reset_budget_windows(self) -> None: @@ -703,7 +727,13 @@ class ResetBudgetJob: changed = False for window in windows: counter_key = f"spend:key:{row['token']}:window:{window['budget_duration']}" - if await ResetBudgetJob._reset_expired_window(window, counter_key, spend_counter_cache, now): + if await ResetBudgetJob._reset_expired_window( + window, + counter_key, + spend_counter_cache, + now, + self.reset_settings, + ): changed = True if changed: await VerificationTokenRepository(self.prisma_client).table.update( @@ -726,7 +756,13 @@ class ResetBudgetJob: changed = False for window in windows: counter_key = f"spend:team:{row['team_id']}:window:{window['budget_duration']}" - if await ResetBudgetJob._reset_expired_window(window, counter_key, spend_counter_cache, now): + if await ResetBudgetJob._reset_expired_window( + window, + counter_key, + spend_counter_cache, + now, + self.reset_settings, + ): changed = True if changed: await TeamRepository(self.prisma_client).table.update( @@ -741,6 +777,7 @@ class ResetBudgetJob: item: Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken], current_time: datetime, item_type: Literal["key", "team", "user"], + reset_settings: BudgetResetSettings, ): """ In-place, updates spend=0, and sets budget_reset_at to current_time + budget_duration @@ -755,24 +792,40 @@ class ResetBudgetJob: try: item.spend = 0.0 if hasattr(item, "budget_duration") and item.budget_duration is not None: - from litellm.proxy.common_utils.timezone_utils import ( - get_budget_reset_time, + item.budget_reset_at = compute_budget_reset_at( + budget_duration=item.budget_duration, settings=reset_settings ) - - item.budget_reset_at = get_budget_reset_time(budget_duration=item.budget_duration) return item except Exception as e: verbose_proxy_logger.exception("Error resetting budget for %s: %s. Item: %s", item_type, e, item) raise e @staticmethod - async def _reset_budget_for_team(team: LiteLLM_TeamTable, current_time: datetime) -> Optional[LiteLLM_TeamTable]: - await ResetBudgetJob._reset_budget_common(item=team, current_time=current_time, item_type="team") + async def _reset_budget_for_team( + team: LiteLLM_TeamTable, + current_time: datetime, + reset_settings: BudgetResetSettings, + ) -> LiteLLM_TeamTable | None: + await ResetBudgetJob._reset_budget_common( + item=team, + current_time=current_time, + item_type="team", + reset_settings=reset_settings, + ) return team @staticmethod - async def _reset_budget_for_user(user: LiteLLM_UserTable, current_time: datetime) -> Optional[LiteLLM_UserTable]: - await ResetBudgetJob._reset_budget_common(item=user, current_time=current_time, item_type="user") + async def _reset_budget_for_user( + user: LiteLLM_UserTable, + current_time: datetime, + reset_settings: BudgetResetSettings, + ) -> LiteLLM_UserTable | None: + await ResetBudgetJob._reset_budget_common( + item=user, + current_time=current_time, + item_type="user", + reset_settings=reset_settings, + ) return user @staticmethod @@ -788,15 +841,15 @@ class ResetBudgetJob: @staticmethod async def _reset_budget_reset_at_date( - budget: LiteLLM_BudgetTableFull, current_time: datetime + budget: LiteLLM_BudgetTableFull, + current_time: datetime, + reset_settings: BudgetResetSettings, ) -> LiteLLM_BudgetTableFull: try: if budget.budget_duration is not None: - from litellm.proxy.common_utils.timezone_utils import ( - get_budget_reset_time, + budget.budget_reset_at = compute_budget_reset_at( + budget_duration=budget.budget_duration, settings=reset_settings ) - - budget.budget_reset_at = get_budget_reset_time(budget_duration=budget.budget_duration) except Exception as e: verbose_proxy_logger.exception("Error resetting budget_reset_at for budget: %s. Item: %s", e, budget) raise e @@ -804,7 +857,14 @@ class ResetBudgetJob: @staticmethod async def _reset_budget_for_key( - key: LiteLLM_VerificationToken, current_time: datetime - ) -> Optional[LiteLLM_VerificationToken]: - await ResetBudgetJob._reset_budget_common(item=key, current_time=current_time, item_type="key") + key: LiteLLM_VerificationToken, + current_time: datetime, + reset_settings: BudgetResetSettings, + ) -> LiteLLM_VerificationToken | None: + await ResetBudgetJob._reset_budget_common( + item=key, + current_time=current_time, + item_type="key", + reset_settings=reset_settings, + ) return key diff --git a/litellm/proxy/common_utils/timezone_utils.py b/litellm/proxy/common_utils/timezone_utils.py index 32f9f47d519..a50daf40144 100644 --- a/litellm/proxy/common_utils/timezone_utils.py +++ b/litellm/proxy/common_utils/timezone_utils.py @@ -1,10 +1,47 @@ -from datetime import datetime, timezone +from datetime import datetime, time, timezone + +from pydantic import BaseModel, ConfigDict import litellm from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time -def get_budget_reset_timezone(): +class BudgetResetSettings(BaseModel): + """Immutable, validated settings that govern when budgets reset. + + Parsed once from `litellm_settings` and injected into consumers (the reset + job, management endpoints) so reset times never depend on reaching into + module-level globals at call time. + """ + + model_config = ConfigDict(frozen=True) + + timezone: str = "UTC" + reset_time_of_day: time = time(0, 0) + + +def parse_budget_reset_time(raw: object) -> time: + """Parse a `budget_reset_time` config value (e.g. "12:00") into a `time`. + + Falls back to midnight when unset; raises a clear error on a malformed value + so a bad config fails loudly at startup instead of silently resetting at midnight. + """ + if raw is None or raw == "": + return time(0, 0) + if not isinstance(raw, str): + raise ValueError(f"Invalid budget_reset_time {raw!r}; must be a quoted 24-hour 'HH:MM' string, e.g. \"12:00\"") + for fmt in ("%H:%M", "%H:%M:%S"): + try: + parsed = datetime.strptime(raw, fmt) + return time(hour=parsed.hour, minute=parsed.minute, second=parsed.second) + except ValueError: + continue + raise ValueError( + f"Invalid budget_reset_time {raw!r}; expected a 24-hour 'HH:MM' or 'HH:MM:SS' string, e.g. \"12:00\"" + ) + + +def get_budget_reset_timezone() -> str: """ Get the budget reset timezone from litellm_settings. Falls back to UTC if not specified. @@ -15,15 +52,29 @@ def get_budget_reset_timezone(): return getattr(litellm, "timezone", None) or "UTC" -def get_budget_reset_time(budget_duration: str) -> datetime: - """ - Get the budget reset time based on the configured timezone. - Falls back to UTC if not specified. - """ +def get_budget_reset_settings() -> BudgetResetSettings: + """Build validated reset settings from litellm_settings. Raises on a malformed + `budget_reset_time`, which lets the proxy fail fast at startup.""" + return BudgetResetSettings( + timezone=get_budget_reset_timezone(), + reset_time_of_day=parse_budget_reset_time(getattr(litellm, "budget_reset_time", None)), + ) - reset_at = get_next_standardized_reset_time( + +def compute_budget_reset_at(budget_duration: str, settings: BudgetResetSettings) -> datetime: + """Compute the next reset time for a budget duration using injected settings.""" + return get_next_standardized_reset_time( duration=budget_duration, current_time=datetime.now(timezone.utc), - timezone_str=get_budget_reset_timezone(), + timezone_str=settings.timezone, + reset_time_of_day=settings.reset_time_of_day, ) - return reset_at + + +def get_budget_reset_time(budget_duration: str) -> datetime: + """Get the budget reset time using the globally-configured timezone and reset time. + + Thin wrapper over `compute_budget_reset_at` for callers that don't yet receive + `BudgetResetSettings` by injection (creation/update endpoints, startup backfill). + """ + return compute_budget_reset_at(budget_duration, get_budget_reset_settings()) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e3eab60bcb4..dc896a9ea1b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -320,7 +320,10 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.common_utils.timezone_utils import ( + get_budget_reset_settings, + get_budget_reset_time, +) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, get_management_object_ttl, @@ -4598,6 +4601,13 @@ class ProxyConfig: litellm.json_logs = True litellm._turn_on_json() verbose_proxy_logger.debug(f"{blue_color_code} Enabled JSON logging via config{reset_color_code}") + elif key == "budget_reset_time": + from litellm.proxy.common_utils.timezone_utils import ( + parse_budget_reset_time, + ) + + parse_budget_reset_time(value) + setattr(litellm, key, value) else: verbose_proxy_logger.debug( f"{blue_color_code} setting litellm.{key}={_redact_general_setting_value(key, value, is_full_admin=False)}{reset_color_code}" @@ -7869,6 +7879,7 @@ class ProxyStartupEvent: budget_reset_job = ResetBudgetJob( proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client, + reset_settings=get_budget_reset_settings(), ) scheduler.add_job( diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index 793cc02ff17..4a1ef5ed696 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -173,6 +173,7 @@ class Status1(Enum): cancelled = "cancelled" incomplete = "incomplete" budget_exceeded = "budget_exceeded" + queued = "queued" class InteractionStatusUpdate(BaseModel): @@ -341,6 +342,7 @@ class Status3(Enum): CANCELLED = "cancelled" INCOMPLETE = "incomplete" BUDGET_EXCEEDED = "budget_exceeded" + QUEUED = "queued" class ModelOption(RootModel[str]): diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 5c96eb619bf..44da3ea06a0 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -30,6 +30,7 @@ def _attrify(d: dict): None)` (et al), which returns None for plain dicts — that would silently skip the row. """ + class _AttrDict(dict): def __getattr__(self, k): try: @@ -120,9 +121,11 @@ async def test_reset_budget_keys_partial_failure(): key1, key2, key3, key4, key5, key6 = ( _attrify(k) for k in [key1, key2, key3, key4, key5, key6] ) - prisma_client.get_data = AsyncMock(return_value=[key1, key2, key3, key4, key5, key6]) + prisma_client.get_data = AsyncMock( + return_value=[key1, key2, key3, key4, key5, key6] + ) - async def fake_reset_key(key, current_time): + async def fake_reset_key(key, current_time, reset_settings=None): if key["id"] == "key1": # Simulate a failure on key1 (for example, this might be due to an invariant check) raise Exception("Simulated failure for key1") @@ -207,9 +210,11 @@ async def test_reset_budget_users_partial_failure(): user1, user2, user3, user4, user5, user6 = ( _attrify(u) for u in [user1, user2, user3, user4, user5, user6] ) - prisma_client.get_data = AsyncMock(return_value=[user1, user2, user3, user4, user5, user6]) + prisma_client.get_data = AsyncMock( + return_value=[user1, user2, user3, user4, user5, user6] + ) - async def fake_reset_user(user, current_time): + async def fake_reset_user(user, current_time, reset_settings=None): if user["id"] == "user1": raise Exception("Simulated failure for user1") else: @@ -397,7 +402,7 @@ async def test_reset_budget_teams_partial_failure(): team1, team2 = _attrify(team1), _attrify(team2) prisma_client.get_data = AsyncMock(return_value=[team1, team2]) - async def fake_reset_team(team, current_time): + async def fake_reset_team(team, current_time, reset_settings=None): if team["id"] == "team1": raise Exception("Simulated failure for team1") else: @@ -513,14 +518,14 @@ async def test_reset_budget_continues_other_categories_on_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_key(key, current_time): + async def fake_reset_key(key, current_time, reset_settings=None): key["spend"] = 0.0 key["budget_reset_at"] = ( current_time + timedelta(seconds=key["budget_duration"]) ).isoformat() return key - async def fake_reset_user(user, current_time): + async def fake_reset_user(user, current_time, reset_settings=None): if user["id"] == "user1": raise Exception("Simulated failure for user1") user["spend"] = 0.0 @@ -529,7 +534,7 @@ async def test_reset_budget_continues_other_categories_on_failure(): ).isoformat() return user - async def fake_reset_team(team, current_time): + async def fake_reset_team(team, current_time, reset_settings=None): team["spend"] = 0.0 team["budget_reset_at"] = ( current_time + timedelta(seconds=team["budget_duration"]) @@ -632,7 +637,7 @@ async def test_service_logger_keys_success(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_key(key, current_time): + async def fake_reset_key(key, current_time, reset_settings=None): key["spend"] = 0.0 key["budget_reset_at"] = ( current_time + timedelta(seconds=key["budget_duration"]) @@ -688,7 +693,7 @@ async def test_service_logger_keys_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_key(key, current_time): + async def fake_reset_key(key, current_time, reset_settings=None): if key["id"] == "key1": raise Exception("Simulated failure for key1") key["spend"] = 0.0 @@ -750,7 +755,7 @@ async def test_service_logger_users_success(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_user(user, current_time): + async def fake_reset_user(user, current_time, reset_settings=None): user["spend"] = 0.0 user["budget_reset_at"] = ( current_time + timedelta(seconds=user["budget_duration"]) @@ -802,7 +807,7 @@ async def test_service_logger_users_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_user(user, current_time): + async def fake_reset_user(user, current_time, reset_settings=None): if user["id"] == "user1": raise Exception("Simulated failure for user1") user["spend"] = 0.0 @@ -863,7 +868,7 @@ async def test_service_logger_teams_success(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_team(team, current_time): + async def fake_reset_team(team, current_time, reset_settings=None): team["spend"] = 0.0 team["budget_reset_at"] = ( current_time + timedelta(seconds=team["budget_duration"]) @@ -915,7 +920,7 @@ async def test_service_logger_teams_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_team(team, current_time): + async def fake_reset_team(team, current_time, reset_settings=None): if team["id"] == "team1": raise Exception("Simulated failure for team1") team["spend"] = 0.0 diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/test_litellm/interactions/test_openapi_compliance.py index 209e99895db..11b08fa45a8 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/test_litellm/interactions/test_openapi_compliance.py @@ -194,6 +194,7 @@ class TestResponseCompliance: "cancelled", "incomplete", "budget_exceeded", + "queued", ] assert status_prop["enum"] == expected_statuses print(f"✓ Status enum values: {expected_statuses}") diff --git a/tests/test_litellm/litellm_core_utils/test_duration_parser.py b/tests/test_litellm/litellm_core_utils/test_duration_parser.py index 3e4446c6672..b6b617610a8 100644 --- a/tests/test_litellm/litellm_core_utils/test_duration_parser.py +++ b/tests/test_litellm/litellm_core_utils/test_duration_parser.py @@ -1,5 +1,5 @@ import unittest -from datetime import datetime, timezone +from datetime import datetime, time, timezone from zoneinfo import ZoneInfo from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time @@ -199,5 +199,122 @@ class TestStandardizedResetTime(unittest.TestCase): self.assertEqual(result, expected) +class TestResetTimeOfDay(unittest.TestCase): + """A configurable reset_time_of_day shifts day/week/month resets off midnight.""" + + def test_daily_reset_before_offset_is_today(self): + now = datetime(2023, 5, 15, 8, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1d", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 15, 12, 0, 0, tzinfo=timezone.utc)) + + def test_daily_reset_after_offset_is_tomorrow(self): + now = datetime(2023, 5, 15, 14, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1d", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 16, 12, 0, 0, tzinfo=timezone.utc)) + + def test_daily_reset_exactly_at_offset_rolls_forward(self): + now = datetime(2023, 5, 15, 12, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1d", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 16, 12, 0, 0, tzinfo=timezone.utc)) + + def test_daily_reset_with_seconds_offset(self): + now = datetime(2023, 5, 15, 8, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1d", now, "UTC", reset_time_of_day=time(9, 30, 15) + ) + self.assertEqual(result, datetime(2023, 5, 15, 9, 30, 15, tzinfo=timezone.utc)) + + def test_offset_applies_in_configured_timezone(self): + # 2023-05-15 22:30 UTC == 2023-05-16 01:30 in Jerusalem (IDT, UTC+3), + # so the next noon-Jerusalem reset is 2023-05-16 12:00 IDT. + now = datetime(2023, 5, 15, 22, 30, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1d", now, "Asia/Jerusalem", reset_time_of_day=time(12, 0) + ) + jerusalem = result.astimezone(ZoneInfo("Asia/Jerusalem")) + self.assertEqual( + (jerusalem.year, jerusalem.month, jerusalem.day), (2023, 5, 16) + ) + self.assertEqual(jerusalem.hour, 12) + self.assertEqual(jerusalem.minute, 0) + + def test_weekly_reset_lands_on_monday_at_offset(self): + wednesday = datetime(2023, 5, 17, 15, 45, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "7d", wednesday, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 22, 12, 0, 0, tzinfo=timezone.utc)) + + def test_weekly_reset_today_is_monday_before_offset_is_today(self): + monday_morning = datetime(2023, 5, 22, 9, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "7d", monday_morning, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 22, 12, 0, 0, tzinfo=timezone.utc)) + + def test_weekly_reset_today_is_monday_after_offset_is_next_week(self): + monday_afternoon = datetime(2023, 5, 22, 15, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "7d", monday_afternoon, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 29, 12, 0, 0, tzinfo=timezone.utc)) + + def test_monthly_30d_lands_on_first_at_offset(self): + now = datetime(2023, 5, 15, 10, 30, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "30d", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 6, 1, 12, 0, 0, tzinfo=timezone.utc)) + + def test_monthly_1mo_today_is_first_before_offset_is_today(self): + now = datetime(2023, 5, 1, 9, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1mo", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 1, 12, 0, 0, tzinfo=timezone.utc)) + + def test_monthly_year_rollover_at_offset(self): + now = datetime(2023, 12, 15, 9, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1mo", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)) + + def test_custom_day_reset_applies_offset(self): + now = datetime(2023, 5, 15, 10, 30, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "3d", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 18, 12, 0, 0, tzinfo=timezone.utc)) + + def test_sub_day_durations_ignore_offset(self): + base = datetime(2023, 5, 15, 15, 20, 30, tzinfo=timezone.utc) + self.assertEqual( + get_next_standardized_reset_time( + "2h", base, "UTC", reset_time_of_day=time(12, 0) + ), + datetime(2023, 5, 15, 16, 0, 0, tzinfo=timezone.utc), + ) + self.assertEqual( + get_next_standardized_reset_time( + "30m", base, "UTC", reset_time_of_day=time(12, 0) + ), + datetime(2023, 5, 15, 15, 30, 0, tzinfo=timezone.utc), + ) + + def test_default_offset_is_midnight(self): + now = datetime(2023, 5, 15, 10, 30, 0, tzinfo=timezone.utc) + self.assertEqual( + get_next_standardized_reset_time("1d", now, "UTC"), + datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc), + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 73486fe0b6a..b56a12db5b1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1033,3 +1033,166 @@ class TestResolveByokMcpAuthHeader: check_mock.assert_awaited_once_with(server, user_auth) assert result == "caller-header" + + +class TestOpenApiResolvedUpstreamAuth: + """LIT-4629: spec_path servers egress through plain httpx, so the manager's OpenAPI arm must + materialize the v2-resolved credential into the `_request_resolved_auth_headers` ContextVar; + before the fix the resolved token never reached the upstream API.""" + + def _oauth_server(self, **overrides: Any) -> MCPServer: + fields: Dict[str, Any] = dict( + server_id="srv-sheets", + name="google_sheets", + server_name="google_sheets", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + spec_path="https://example.com/sheets-openapi.yaml", + ) + fields.update(overrides) + return MCPServer(**fields) + + @pytest.mark.asyncio + async def test_call_tool_openapi_injects_v2_resolved_token_contextvar(self): + """The managed spec_path arm resolves the v2 credential and sets the ContextVar; kills + the mutant that drops the resolve_openapi_upstream_auth call in call_tool.""" + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_resolved_auth_headers, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + manager = MCPServerManager() + server = self._oauth_server() + user_auth = UserAPIKeyAuth(user_id="alice", api_key="sk-user") + captured: Dict[str, Any] = {} + + async def fake_openapi_handler(_server, _name, _arguments): + captured["resolved"] = _request_resolved_auth_headers.get() + return MagicMock() + + with patch.object(manager, "_resolve_mcp_server_for_tool_call", return_value=server): + with patch.object( + manager._cred_provider, + "resolve_credentials", + new=AsyncMock(return_value=Ok(StaticHeaderAuth("Bearer stored-user-token"))), + ): + with patch.object(manager, "_call_openapi_tool_handler", side_effect=fake_openapi_handler): + await manager.call_tool( + server_name=server.server_name, + name="get_values", + arguments={}, + user_api_key_auth=user_auth, + ) + + assert captured["resolved"] == {"Authorization": "Bearer stored-user-token"} + assert _request_resolved_auth_headers.get() is None + + @pytest.mark.asyncio + async def test_call_tool_openapi_m2m_missing_token_url_fails_closed(self): + """A url-less M2M spec server with no token_url must fail with a typed error instead of + egressing unauthenticated (the pre-#32259 silent failure this arm previously preserved). + Drives the real adapter/resolver chain: ClientCredentialsConfig with missing grant fields + resolves to a misconfigured CredError, raised as an HTTPException.""" + from fastapi import HTTPException + + manager = MCPServerManager() + server = self._oauth_server( + oauth2_flow="client_credentials", + client_id="m2m-client", + client_secret="m2m-secret", + token_url=None, + ) + called = AsyncMock() + + with patch.object(manager, "_resolve_mcp_server_for_tool_call", return_value=server): + with patch.object(manager, "_call_openapi_tool_handler", new=called): + with pytest.raises(HTTPException): + await manager.call_tool( + server_name=server.server_name, + name="get_values", + arguments={}, + user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="sk-user"), + ) + + called.assert_not_awaited() + + @pytest.mark.asyncio + async def test_caller_oauth2_headers_never_become_resolved_for_byok_server(self): + """Greptile P1 regression: BYOK servers defer to v1 (to_server_spec None), and the v1 arm + must never promote caller-supplied oauth2 headers into the resolved-auth slot, where they + would override the per-server BYOK credential and leak the caller's gateway Authorization + upstream.""" + manager = MCPServerManager() + server = MCPServer( + server_id="byok-spec", + name="byok_spec", + server_name="byok_spec", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + spec_path="https://example.com/openapi.yaml", + is_byok=True, + ) + + resolved, forwarded = await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers={"Authorization": "Bearer sk-litellm-gateway-key"}, + raw_headers=None, + mcp_auth_header="user-byok-key", + user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="sk-user"), + forwarded_headers=None, + ) + + assert resolved is None + assert forwarded is None + + @pytest.mark.asyncio + async def test_v1_server_threads_stored_headers_only_without_caller_headers(self): + """The v1 (unmigrated) arm resolves the stored per-user token only when the caller sent no + oauth2 headers of their own; with caller headers present the stored lookup is skipped and + nothing is promoted to resolved.""" + manager = MCPServerManager() + server = MCPServer( + server_id="v1-spec", + name="v1_spec", + server_name="v1_spec", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + spec_path="https://example.com/openapi.yaml", + delegate_auth_to_upstream=True, + ) + stored = {"Authorization": "Bearer stored-v1-token"} + user_auth = UserAPIKeyAuth(user_id="alice", api_key="sk-user") + + with patch.object( + manager, "_resolve_oauth2_headers_for_tool_call", new=AsyncMock(return_value=stored) + ) as lookup: + resolved, _ = await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers=None, + raw_headers=None, + mcp_auth_header=None, + user_api_key_auth=user_auth, + forwarded_headers=None, + ) + assert resolved == stored + lookup.assert_awaited_once_with(server, None, user_auth) + + with patch.object( + manager, "_resolve_oauth2_headers_for_tool_call", new=AsyncMock(return_value=stored) + ) as lookup: + resolved, _ = await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers={"Authorization": "Bearer caller-supplied"}, + raw_headers=None, + mcp_auth_header=None, + user_api_key_auth=user_auth, + forwarded_headers=None, + ) + assert resolved is None + lookup.assert_not_awaited() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 03f91260955..576b9f4f139 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -8891,3 +8891,48 @@ async def test_resolve_toolset_tool_permissions_single_db_fetch_across_checks(): assert first == {"server-a": ["lookup_status"]} assert second == first list_toolsets_mock.assert_awaited_once() + + +class TestMaterializeAuthHeaders: + """_materialize_auth_headers drives one step of a resolved httpx.Auth's own flow to turn it + into a header dict for the OpenAPI egress arm, which sends plain headers and cannot carry an + httpx.Auth. Generic across auth shapes via the resolver-arm header_name convention.""" + + @pytest.mark.asyncio + async def test_static_header_auth_materializes_its_header(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _materialize_auth_headers, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + + headers = await _materialize_auth_headers(StaticHeaderAuth("Bearer stored-token")) + assert headers == {"Authorization": "Bearer stored-token"} + + @pytest.mark.asyncio + async def test_client_credentials_bearer_auth_materializes_bearer(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _materialize_auth_headers, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ) + + async def _refetch(_stale: str): + return None + + headers = await _materialize_auth_headers(ClientCredentialsBearerAuth("m2m-token", _refetch)) + assert headers == {"Authorization": "Bearer m2m-token"} + + @pytest.mark.asyncio + async def test_noop_and_none_materialize_to_none(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _materialize_auth_headers, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + NoOpAuth, + ) + + assert await _materialize_auth_headers(None) is None + assert await _materialize_auth_headers(NoOpAuth()) is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 39f3c767220..7bcacb3ff4a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -17,6 +17,7 @@ import pytest from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, _request_extra_headers, + _request_resolved_auth_headers, _resolve_param_list, _resolve_ref, build_input_schema, @@ -1207,3 +1208,61 @@ class TestRequestExtraHeaders: call_args = async_client.get.call_args headers_sent = call_args[1]["headers"] assert "X-TOKEN" not in headers_sent + + @pytest.mark.asyncio + async def test_resolved_auth_headers_win_over_every_other_authorization_source(self): + """The gateway-resolved credential (stored per-user OAuth / minted M2M token) is + authoritative: it must override the BYOK override, static headers, and forwarded caller + headers on the Authorization name, case-insensitively, mirroring _resolve_v2_auth's rule + on the MCPClient path. Without this, a spec_path oauth2 server's completed OAuth flow + stores a token that never reaches the upstream API (LIT-4629).""" + operation = {} + func = create_tool_function( + path="/secure", + method="get", + operation=operation, + base_url="https://api.example.com", + headers={"authorization": "Bearer static-operator"}, + ) + + with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + async_client = _create_mock_client("get", "secure-data") + mock_client.return_value = async_client + + extra_token = _request_extra_headers.set({"Authorization": "Bearer caller-forwarded"}) + auth_token = _request_auth_header.set("Bearer byok-credential") + resolved_token = _request_resolved_auth_headers.set({"Authorization": "Bearer resolved-oauth"}) + try: + result = await func() + finally: + _request_auth_header.reset(auth_token) + _request_extra_headers.reset(extra_token) + _request_resolved_auth_headers.reset(resolved_token) + + assert result == "secure-data" + headers_sent = async_client.get.call_args[1]["headers"] + authorization_values = [v for k, v in headers_sent.items() if k.lower() == "authorization"] + assert authorization_values == ["Bearer resolved-oauth"] + + @pytest.mark.asyncio + async def test_resolved_auth_headers_not_leaked_between_calls(self): + """After resetting the resolved-auth ContextVar, subsequent calls send no credential.""" + operation = {} + func = create_tool_function( + path="/data", + method="get", + operation=operation, + base_url="https://api.example.com", + ) + + with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + async_client = _create_mock_client("get", "ok") + mock_client.return_value = async_client + + token = _request_resolved_auth_headers.set({"Authorization": "Bearer resolved-oauth"}) + _request_resolved_auth_headers.reset(token) + + await func() + + headers_sent = async_client.get.call_args[1]["headers"] + assert "Authorization" not in headers_sent diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 3ad01e9c3ec..1e4349c3143 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -218,3 +218,86 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable(): assert exc.value.status_code == 503 pre_call.assert_not_awaited() handle_local.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_openapi_local_tool_injects_resolved_oauth_token(): + """LIT-4629: the local-registry (OpenAPI) dispatch is the primary egress for spec_path + tools, and before the fix it dropped the gateway-resolved OAuth credential entirely, so a + user's completed OAuth flow stored a token that never reached the upstream API. The resolved + credential must land in the `_request_resolved_auth_headers` ContextVar the tool closure + reads. Kills the mutant that deletes the resolve_openapi_upstream_auth call in server.py.""" + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_resolved_auth_headers, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + user = UserAPIKeyAuth( + api_key="sk-user", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + oauth_server = MCPServer( + server_id="srv-sheets", + name="google_sheets", + server_name="google_sheets", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + spec_path="https://example.com/sheets-openapi.yaml", + ) + + fake_tool = MagicMock() + fake_tool.name = "get_values" + captured: dict = {} + + async def handle_local(_name, _arguments): + captured["resolved"] = _request_resolved_auth_headers.get() + return [] + + with ( + patch.object( + mcp_module.global_mcp_server_manager, + "_get_mcp_server_from_tool_name", + return_value=oauth_server, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "pre_call_tool_check", + new=AsyncMock(return_value={}), + ), + patch.object( + mcp_module.global_mcp_tool_registry, + "get_tool", + return_value=fake_tool, + ), + patch.object( + mcp_module.global_mcp_server_manager._cred_provider, + "resolve_credentials", + new=AsyncMock(return_value=Ok(StaticHeaderAuth("Bearer stored-user-token"))), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + new=handle_local, + ), + patch( + "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed", + return_value=True, + ), + ): + await mcp_module.execute_mcp_tool( + name="get_values", + arguments={}, + allowed_mcp_servers=[oauth_server], + start_time=datetime.now(timezone.utc), + user_api_key_auth=user, + ) + + assert captured["resolved"] == {"Authorization": "Bearer stored-user-token"} + assert _request_resolved_auth_headers.get() is None diff --git a/tests/test_litellm/proxy/a2a/test_agent_card.py b/tests/test_litellm/proxy/a2a/test_agent_card.py index d302bde7895..dfa848e335e 100644 --- a/tests/test_litellm/proxy/a2a/test_agent_card.py +++ b/tests/test_litellm/proxy/a2a/test_agent_card.py @@ -1,10 +1,14 @@ """Unit tests for the pure merge logic in litellm/proxy/a2a/agent_card.py.""" +import pytest + from litellm.proxy.a2a.agent_card import ( LITELLM_A2A_PROTOCOL_VERSION, LITELLM_SECURITY_REQUIREMENTS, LITELLM_SECURITY_SCHEMES, merge_agent_card, + normalize_protocol_version, + resolve_served_protocol_version, ) PROXY_URL = "https://proxy.example/a2a/agent-xyz" @@ -205,3 +209,54 @@ def test_strips_additional_interfaces_to_prevent_backend_url_leak(): ] merged = merge_agent_card(upstream, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE) assert "additionalInterfaces" not in merged + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("0.3", "0.3"), + ("0.3.0", "0.3"), + ("1.0", "1.0"), + ("1.0.0", "1.0"), + ("1.0.1", "1.0"), + ("0.3.0-rc1", "0.3"), + ("1.0.0-rc.1+build.5", "1.0"), + ("0.2.6", None), + ("2.0", None), + ("0.30", None), + ("0.3.garbage", None), + ("0.3.", None), + ("1.0.not-semver", None), + ("0.3.0.0", None), + ("0.3-rc1", None), + ("garbage", None), + ("", None), + (None, None), + (1.0, None), + ], +) +def test_normalize_protocol_version(raw, expected): + assert normalize_protocol_version(raw) == expected + + +def test_resolve_served_protocol_version_canonicalizes_semver_pins(): + assert resolve_served_protocol_version({"protocolVersion": "0.3.0"}) == "0.3" + assert resolve_served_protocol_version({"protocolVersion": "1.0.0"}) == "1.0" + assert resolve_served_protocol_version({"protocolVersion": "0.3"}) == "0.3" + assert resolve_served_protocol_version({"protocolVersion": "1.0"}) == "1.0" + + +def test_resolve_served_protocol_version_falls_back_for_unsupported(): + assert ( + resolve_served_protocol_version({"protocolVersion": "0.2.6"}) + == LITELLM_A2A_PROTOCOL_VERSION + ) + assert resolve_served_protocol_version(None) == LITELLM_A2A_PROTOCOL_VERSION + + +def test_serves_semver_pinned_protocol_version_as_major_minor(): + card = _full_upstream_card() + card["protocolVersion"] = "0.3.0" + merged = merge_agent_card(card, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE) + assert merged["protocolVersion"] == "0.3" + assert merged["supportedInterfaces"][0]["protocolVersion"] == "0.3" diff --git a/tests/test_litellm/proxy/a2a/test_version_convert.py b/tests/test_litellm/proxy/a2a/test_version_convert.py index f3c51ca6b72..7eb5debb792 100644 --- a/tests/test_litellm/proxy/a2a/test_version_convert.py +++ b/tests/test_litellm/proxy/a2a/test_version_convert.py @@ -313,3 +313,13 @@ def test_agent_card_with_0_3_pin_and_supported_interfaces_is_lowered(): def test_agent_card_same_version_passthrough(): card = _extended_card_1_0() assert normalize_agent_card(card, "1.0") is card + + +def test_detect_card_version_normalizes_semver_protocol_version(): + from litellm.proxy.a2a.version_convert import _detect_card_version + + assert _detect_card_version({"protocolVersion": "1.0.0"}) == "1.0" + assert ( + _detect_card_version({"protocolVersion": "0.3.0", "supportedInterfaces": []}) + == "0.3" + ) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 3740c01b7fc..bcd3333baf9 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -540,6 +540,53 @@ class TestAgentRBACProxyAdmin: assert resp.status_code == 200 +class TestAgentProtocolVersionValidation: + """Registration accepts spec-default semver protocolVersion values and still + rejects genuinely unsupported versions.""" + + @pytest.fixture(autouse=True) + def _setup(self, monkeypatch): + self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN) + self.mock_registry = MagicMock() + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry) + + def _create_agent_with_protocol_version(self, protocol_version: str): + config = _sample_agent_config() + config["agent_card_params"]["protocolVersion"] = protocol_version + with patch("litellm.proxy.proxy_server.prisma_client"): + self.mock_registry.get_agent_by_name = MagicMock(return_value=None) + self.mock_registry.add_agent_to_db = AsyncMock( + return_value=_sample_agent_response() + ) + self.mock_registry.register_agent = MagicMock() + return self.admin_client.post( + "/v1/agents", + json=config, + headers={"Authorization": "Bearer k"}, + ) + + def test_semver_protocol_version_registers_and_stores_major_minor(self): + resp = self._create_agent_with_protocol_version("0.3.0") + assert resp.status_code == 200 + stored_card = self.mock_registry.add_agent_to_db.await_args.kwargs["agent"][ + "agent_card_params" + ] + assert stored_card["protocolVersion"] == "0.3" + assert stored_card["supportedInterfaces"][0]["protocolVersion"] == "0.3" + + def test_unsupported_protocol_version_is_rejected(self): + resp = self._create_agent_with_protocol_version("0.2.6") + assert resp.status_code == 400 + assert "Unsupported protocolVersion '0.2.6'" in resp.json()["detail"] + self.mock_registry.add_agent_to_db.assert_not_awaited() + + def test_malformed_protocol_version_is_rejected(self): + resp = self._create_agent_with_protocol_version("0.3.garbage") + assert resp.status_code == 400 + assert "Unsupported protocolVersion '0.3.garbage'" in resp.json()["detail"] + self.mock_registry.add_agent_to_db.assert_not_awaited() + + class TestCheckAgentManagementPermission: """Unit tests for the _check_agent_management_permission helper.""" diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 5e348b1bb7e..be5bc74c385 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -5,25 +5,23 @@ import sys import time import types from datetime import datetime, timedelta, timezone +from datetime import time as dt_time from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock import pytest -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob +from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings from litellm.proxy.utils import ProxyLogging # Mock classes for testing class MockLiteLLMTeamMembership: - async def update_many( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Dict[str, Any]: + async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: # Mock the update_many method for litellm_teammembership return {"count": 1} @@ -32,9 +30,7 @@ class MockLiteLLMVerificationToken: def __init__(self): self.update_many_calls: List[Dict[str, Any]] = [] - async def update_many( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Dict[str, Any]: + async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_many_calls.append({"where": where, "data": data}) return {"count": 1} @@ -52,9 +48,7 @@ class MockLiteLLMOrganizationTable: self.find_many_calls.append({"where": where}) return self._find_many_results - async def update_many( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Dict[str, Any]: + async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_many_calls.append({"where": where, "data": data}) return {"count": 1} @@ -72,9 +66,7 @@ class MockLiteLLMTagTable: self.find_many_calls.append({"where": where}) return self._find_many_results - async def update_many( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Dict[str, Any]: + async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_many_calls.append({"where": where, "data": data}) return {"count": 1} @@ -110,9 +102,7 @@ class MockBatcher: _self._outer = outer def update(_self, where, data): - _self._outer.calls.append( - {"table": _self._table_name, "where": where, "data": data} - ) + _self._outer.calls.append({"table": _self._table_name, "where": where, "data": data}) self.litellm_verificationtoken = _Table("key", self) self.litellm_usertable = _Table("user", self) @@ -172,11 +162,7 @@ class MockPrismaClient: return [item for item in data if hasattr(item, "budget_reset_at")] # Handle specific filtering for enduser table queries - if ( - table_name == "enduser" - and query_type == "find_all" - and "budget_id_list" in kwargs - ): + if table_name == "enduser" and query_type == "find_all" and "budget_id_list" in kwargs: budget_id_list = kwargs["budget_id_list"] # Return endusers that match the budget IDs return [ @@ -188,11 +174,7 @@ class MockPrismaClient: ] # Handle key queries with expires and reset_at - if ( - table_name == "key" - and query_type == "find_all" - and ("expires" in kwargs or "reset_at" in kwargs) - ): + if table_name == "key" and query_type == "find_all" and ("expires" in kwargs or "reset_at" in kwargs): return [item for item in data if hasattr(item, "budget_reset_at")] return data @@ -227,9 +209,7 @@ def mock_proxy_logging(): @pytest.fixture def reset_budget_job(mock_prisma_client, mock_proxy_logging): - return ResetBudgetJob( - proxy_logging_obj=mock_proxy_logging, prisma_client=mock_prisma_client - ) + return ResetBudgetJob(proxy_logging_obj=mock_proxy_logging, prisma_client=mock_prisma_client) # Helper function to run async tests @@ -270,6 +250,40 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): assert set(write["data"].keys()) == {"spend", "budget_reset_at"} +def test_reset_budget_for_key_honors_injected_reset_time(mock_prisma_client, mock_proxy_logging): + """Injected BudgetResetSettings drives the written reset time end to end (DI, no globals). + + Before the configurable-reset-time change this wrote a midnight reset_at (hour 0); + with noon injected it must write a noon reset_at. + """ + job = ResetBudgetJob( + proxy_logging_obj=mock_proxy_logging, + prisma_client=mock_prisma_client, + reset_settings=BudgetResetSettings(timezone="UTC", reset_time_of_day=dt_time(12, 0)), + ) + now = datetime.now(timezone.utc) + test_key = type( + "LiteLLM_VerificationToken", + (), + { + "spend": 100.0, + "budget_duration": "1d", + "budget_reset_at": now, + "id": "test-key-noon", + "token": "tok-noon", + }, + ) + mock_prisma_client.data["key"] = [test_key] + + asyncio.run(job.reset_budget_for_litellm_keys()) + + key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"] + assert len(key_writes) == 1 + reset_at = key_writes[0]["data"]["budget_reset_at"].astimezone(timezone.utc) + assert reset_at.hour == 12 + assert reset_at.minute == 0 + + def test_reset_budget_for_user(reset_budget_job, mock_prisma_client): # Setup test data with timezone-aware datetime now = datetime.now(timezone.utc) @@ -486,11 +500,7 @@ def test_reset_budget_for_keys_linked_to_budgets(reset_budget_job, mock_prisma_c budgets_to_reset = [test_budget] # Run the method - asyncio.run( - reset_budget_job.reset_budget_for_keys_linked_to_budgets( - budgets_to_reset=budgets_to_reset - ) - ) + asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset)) # Verify that update_many was called on litellm_verificationtoken calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls @@ -531,11 +541,7 @@ def test_reset_budget_for_keys_linked_to_budgets_excludes_keys_with_own_budget_d budgets_to_reset = [test_budget] - asyncio.run( - reset_budget_job.reset_budget_for_keys_linked_to_budgets( - budgets_to_reset=budgets_to_reset - ) - ) + asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset)) calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls assert len(calls) == 1 @@ -548,17 +554,13 @@ def test_reset_budget_for_keys_linked_to_budgets_excludes_keys_with_own_budget_d assert call["where"]["budget_id"] == {"in": ["7d-budget-tier"]} -def test_reset_budget_for_keys_linked_to_budgets_empty( - reset_budget_job, mock_prisma_client -): +def test_reset_budget_for_keys_linked_to_budgets_empty(reset_budget_job, mock_prisma_client): """ Test that when there are no budgets to reset, no update is performed on the verification token table. """ # Run with empty list - asyncio.run( - reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=[]) - ) + asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=[])) # Verify no update_many calls were made calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls @@ -584,11 +586,7 @@ def test_reset_budget_for_orgs_linked_to_budgets(reset_budget_job, mock_prisma_c }, ) - asyncio.run( - reset_budget_job.reset_budget_for_orgs_linked_to_budgets( - budgets_to_reset=[test_budget] - ) - ) + asyncio.run(reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[test_budget])) calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls assert len(calls) == 1 @@ -598,16 +596,12 @@ def test_reset_budget_for_orgs_linked_to_budgets(reset_budget_job, mock_prisma_c assert call["data"]["spend"] == 0 -def test_reset_budget_for_orgs_linked_to_budgets_empty( - reset_budget_job, mock_prisma_client -): +def test_reset_budget_for_orgs_linked_to_budgets_empty(reset_budget_job, mock_prisma_client): """ Test that when there are no budgets to reset, no update is performed on the organization table. """ - asyncio.run( - reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[]) - ) + asyncio.run(reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[])) calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls assert len(calls) == 0 @@ -631,11 +625,7 @@ def test_reset_budget_for_tags_linked_to_budgets(reset_budget_job, mock_prisma_c }, ) - asyncio.run( - reset_budget_job.reset_budget_for_tags_linked_to_budgets( - budgets_to_reset=[test_budget] - ) - ) + asyncio.run(reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[test_budget])) calls = mock_prisma_client.db.litellm_tagtable.update_many_calls assert len(calls) == 1 @@ -645,16 +635,12 @@ def test_reset_budget_for_tags_linked_to_budgets(reset_budget_job, mock_prisma_c assert call["data"]["spend"] == 0 -def test_reset_budget_for_tags_linked_to_budgets_empty( - reset_budget_job, mock_prisma_client -): +def test_reset_budget_for_tags_linked_to_budgets_empty(reset_budget_job, mock_prisma_client): """ Test that when there are no budgets to reset, no update is performed on the tag table. """ - asyncio.run( - reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[]) - ) + asyncio.run(reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[])) calls = mock_prisma_client.db.litellm_tagtable.update_many_calls assert len(calls) == 0 @@ -668,9 +654,7 @@ def test_reset_budget_for_tags_linked_to_budgets_empty( ], ids=["30d-calendar-month", "1mo-calendar-month", "1d-next-midnight"], ) -def test_reset_budget_reset_at_date_calendar_aligned( - budget_duration, expected_day, expected_month -): +def test_reset_budget_reset_at_date_calendar_aligned(budget_duration, expected_day, expected_month): """ Verify that _reset_budget_reset_at_date produces calendar-aligned reset times (matching get_budget_reset_time), not sliding-window offsets. @@ -694,7 +678,7 @@ def test_reset_budget_reset_at_date_calendar_aligned( with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt: mock_dt.now.return_value = fixed_now mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now)) + asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings())) assert test_budget.budget_reset_at.day == expected_day assert test_budget.budget_reset_at.month == expected_month @@ -724,7 +708,7 @@ def test_reset_budget_reset_at_date_7d_next_monday(): with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt: mock_dt.now.return_value = fixed_now mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now)) + asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings())) # Next Monday after Wednesday June 14 is June 19 assert test_budget.budget_reset_at.day == 19 @@ -749,7 +733,7 @@ def test_reset_budget_reset_at_date_none_duration(): }, ) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, now)) + asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, now, BudgetResetSettings())) assert test_budget.budget_reset_at == original_reset_at @@ -773,7 +757,7 @@ def test_reset_budget_reset_at_date_none_reset_at(): with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt: mock_dt.now.return_value = fixed_now mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now)) + asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings())) # Should be set to 1st of next month (July 1) assert test_budget.budget_reset_at is not None @@ -781,9 +765,7 @@ def test_reset_budget_reset_at_date_none_reset_at(): assert test_budget.budget_reset_at.month == 7 -def test_budget_table_reset_also_resets_linked_keys( - reset_budget_job, mock_prisma_client -): +def test_budget_table_reset_also_resets_linked_keys(reset_budget_job, mock_prisma_client): """ Integration-style test: when reset_budget_for_litellm_budget_table runs, it should also reset spend for keys linked to the expiring budget tiers @@ -818,9 +800,7 @@ def test_budget_table_reset_also_resets_linked_keys( assert calls[0]["data"]["spend"] == 0 -def test_budget_table_reset_also_resets_linked_orgs( - reset_budget_job, mock_prisma_client -): +def test_budget_table_reset_also_resets_linked_orgs(reset_budget_job, mock_prisma_client): """ Integration-style test: when reset_budget_for_litellm_budget_table runs, it should also reset spend for orgs linked to the expiring budget tiers @@ -853,9 +833,7 @@ def test_budget_table_reset_also_resets_linked_orgs( assert calls[0]["data"]["spend"] == 0 -def test_budget_table_reset_also_resets_linked_tags( - reset_budget_job, mock_prisma_client -): +def test_budget_table_reset_also_resets_linked_tags(reset_budget_job, mock_prisma_client): """ Integration-style test: when reset_budget_for_litellm_budget_table runs, it should also reset spend for tags linked to the expiring budget tiers. @@ -887,9 +865,7 @@ def test_budget_table_reset_also_resets_linked_tags( assert calls[0]["data"]["spend"] == 0 -def test_reset_budget_resets_endusers_with_null_budget_id( - reset_budget_job, mock_prisma_client -): +def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock_prisma_client): """ When litellm.max_end_user_budget_id is configured and that budget is being reset, end users with budget_id=NULL should also have their spend @@ -959,17 +935,13 @@ def test_reset_budget_resets_endusers_with_null_budget_id( mock_prisma_client.data["enduser"] = [enduser_with_budget] # Set up the DB mock for NULL-budget-id end users - mock_prisma_client.db.litellm_endusertable.set_find_many_results( - [enduser_no_budget_row] - ) + mock_prisma_client.db.litellm_endusertable.set_find_many_results([enduser_no_budget_row]) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) # Both end users should have been reset updated = mock_prisma_client.updated_data["enduser"] - assert ( - len(updated) == 2 - ), f"Expected 2 endusers reset (1 explicit + 1 implicit), got {len(updated)}" + assert len(updated) == 2, f"Expected 2 endusers reset (1 explicit + 1 implicit), got {len(updated)}" user_ids = {u.user_id for u in updated} assert "enduser-explicit" in user_ids @@ -986,9 +958,7 @@ def test_reset_budget_resets_endusers_with_null_budget_id( litellm.max_end_user_budget_id = None -def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured( - reset_budget_job, mock_prisma_client -): +def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured(reset_budget_job, mock_prisma_client): """ When litellm.max_end_user_budget_id is NOT configured, end users with budget_id=NULL should NOT be fetched or reset. @@ -1073,20 +1043,14 @@ def test_reset_budget_for_team_members_preserves_total_spend(): mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_teammembership.update_many = AsyncMock( - return_value={"count": 1} - ) + mock_prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1}) - job = ResetBudgetJob( - proxy_logging_obj=MagicMock(), prisma_client=mock_prisma_client - ) + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=mock_prisma_client) asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) mock_prisma_client.db.litellm_teammembership.update_many.assert_called_once() - call_kwargs = ( - mock_prisma_client.db.litellm_teammembership.update_many.call_args.kwargs - ) + call_kwargs = mock_prisma_client.db.litellm_teammembership.update_many.call_args.kwargs assert call_kwargs["where"]["budget_id"]["in"] == ["budget-1"] assert call_kwargs["data"] == {"spend": 0} assert "total_spend" not in call_kwargs["data"] @@ -1142,9 +1106,7 @@ def test_reset_budget_windows_uses_is_not_null_filter(monkeypatch): raises `MissingRequiredValueError`. We work around it by using `query_raw` with `IS NOT NULL`. If someone reverts to the ORM filter, this test fails. """ - job, prisma_client, _ = _make_reset_budget_windows_job( - monkeypatch, key_rows=[], team_rows=[] - ) + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=[], team_rows=[]) asyncio.run(job.reset_budget_windows()) @@ -1184,15 +1146,11 @@ def test_reset_budget_windows_resets_expired_key_window(monkeypatch): # The `budget_limits` payload is re-serialized JSON with a bumped reset_at. written_windows = json.loads(call_kwargs["data"]["budget_limits"]) assert len(written_windows) == 1 - new_reset_at = datetime.fromisoformat( - written_windows[0]["reset_at"].replace("Z", "+00:00") - ).replace(tzinfo=None) + new_reset_at = datetime.fromisoformat(written_windows[0]["reset_at"].replace("Z", "+00:00")).replace(tzinfo=None) assert new_reset_at > now # The spend counter for this key+window was cleared. - spend_counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:key:sk-expired:window:1d", value=0.0 - ) + spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0) def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch): @@ -1206,9 +1164,7 @@ def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch): "budget_limits": [{"budget_duration": "1d", "reset_at": future}], } ] - job, prisma_client, _ = _make_reset_budget_windows_job( - monkeypatch, key_rows=key_rows, team_rows=[] - ) + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) asyncio.run(job.reset_budget_windows()) @@ -1237,9 +1193,7 @@ def test_reset_budget_windows_resets_expired_team_window(monkeypatch): assert call_kwargs["where"] == {"team_id": "team-expired"} assert "budget_limits" in call_kwargs["data"] - spend_counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:team:team-expired:window:30d", value=0.0 - ) + spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team:team-expired:window:30d", value=0.0) def test_reset_budget_windows_handles_string_budget_limits(monkeypatch): @@ -1252,14 +1206,10 @@ def test_reset_budget_windows_handles_string_budget_limits(monkeypatch): key_rows = [ { "token": "sk-string-limits", - "budget_limits": json.dumps( - [{"budget_duration": "1d", "reset_at": expired}] - ), + "budget_limits": json.dumps([{"budget_duration": "1d", "reset_at": expired}]), } ] - job, prisma_client, _ = _make_reset_budget_windows_job( - monkeypatch, key_rows=key_rows, team_rows=[] - ) + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) asyncio.run(job.reset_budget_windows()) @@ -1274,9 +1224,7 @@ def test_reset_budget_windows_skips_row_with_empty_budget_limits(monkeypatch): {"token": "sk-empty-list", "budget_limits": []}, {"token": "sk-empty-str", "budget_limits": ""}, ] - job, prisma_client, _ = _make_reset_budget_windows_job( - monkeypatch, key_rows=key_rows, team_rows=[] - ) + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) asyncio.run(job.reset_budget_windows()) @@ -1361,27 +1309,17 @@ def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch): ) prisma_client = MagicMock() - prisma_client.db.litellm_teammembership.find_many = AsyncMock( - return_value=[membership] - ) - prisma_client.db.litellm_teammembership.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[membership]) + prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:team_member:alice:team-x", value=0.0, ttl=60 - ) - counter_cache.redis_cache.async_set_cache.assert_any_await( - key="spend:team_member:alice:team-x", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team_member:alice:team-x", value=0.0, ttl=60) + counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:team_member:alice:team-x", value=0.0, ttl=60) -def test_reset_budget_for_keys_invalidates_redis_counter( - reset_budget_job, mock_prisma_client, monkeypatch -): +def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): """Key budget reset must clear the Redis spend counter.""" counter_cache = _make_counter_invalidation_job(monkeypatch) @@ -1402,14 +1340,10 @@ def test_reset_budget_for_keys_invalidates_redis_counter( asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:key:sk-abc", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-abc", value=0.0, ttl=60) -def test_reset_budget_for_users_invalidates_redis_counter( - reset_budget_job, mock_prisma_client, monkeypatch -): +def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): """User budget reset must clear the Redis spend counter.""" counter_cache = _make_counter_invalidation_job(monkeypatch) @@ -1430,14 +1364,10 @@ def test_reset_budget_for_users_invalidates_redis_counter( asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:user:alice", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:user:alice", value=0.0, ttl=60) -def test_reset_budget_for_teams_invalidates_redis_counter( - reset_budget_job, mock_prisma_client, monkeypatch -): +def test_reset_budget_for_teams_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): """Team budget reset must clear the Redis spend counter.""" counter_cache = _make_counter_invalidation_job(monkeypatch) @@ -1458,9 +1388,7 @@ def test_reset_budget_for_teams_invalidates_redis_counter( asyncio.run(reset_budget_job.reset_budget_for_litellm_teams()) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:team:team-x", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team:team-x", value=0.0, ttl=60) def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch): @@ -1511,9 +1439,7 @@ def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch): batcher.commit = failing_commit prisma_client.db.batch_ = MagicMock(return_value=batcher) - job = ResetBudgetJob( - proxy_logging_obj=MockProxyLogging(), prisma_client=prisma_client - ) + job = ResetBudgetJob(proxy_logging_obj=MockProxyLogging(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_litellm_keys()) @@ -1543,8 +1469,8 @@ def test_reset_budget_for_keys_writes_only_spend_and_reset_at(reset_budget_job, "budget_duration": "30d", "budget_reset_at": now, "token": "sk-problematic", - "object_permission_id": "perm-abc", # would be rejected on update - "budget_limits": [{"max_budget": 5}], # would be rejected on update + "object_permission_id": "perm-abc", # would be rejected on update + "budget_limits": [{"max_budget": 5}], # would be rejected on update "metadata": {"some": "thing"}, }, ) @@ -1570,19 +1496,13 @@ def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monke linked_key = type("Key", (), {"token": "sk-linked"}) prisma_client = MagicMock() - prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[linked_key] - ) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[linked_key]) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget])) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:key:sk-linked", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-linked", value=0.0, ttl=60) def test_reset_budget_for_orgs_linked_to_budgets_invalidates_redis_counter(monkeypatch): @@ -1593,22 +1513,14 @@ def test_reset_budget_for_orgs_linked_to_budgets_invalidates_redis_counter(monke linked_org = type("Org", (), {"organization_id": "org-acme"}) prisma_client = MagicMock() - prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[linked_org] - ) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[linked_org]) + prisma_client.db.litellm_organizationtable.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget])) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:org:org-acme", value=0.0, ttl=60 - ) - counter_cache.redis_cache.async_set_cache.assert_any_await( - key="spend:org:org-acme", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:org:org-acme", value=0.0, ttl=60) + counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:org:org-acme", value=0.0, ttl=60) def test_reset_budget_for_tags_linked_to_budgets_invalidates_redis_counter(monkeypatch): @@ -1625,12 +1537,8 @@ def test_reset_budget_for_tags_linked_to_budgets_invalidates_redis_counter(monke job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:tag:tenant-42", value=0.0, ttl=60 - ) - counter_cache.redis_cache.async_set_cache.assert_any_await( - key="spend:tag:tenant-42", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:tag:tenant-42", value=0.0, ttl=60) + counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:tag:tenant-42", value=0.0, ttl=60) def test_reset_budget_for_tags_linked_to_budgets_invalidates_management_cache( @@ -1657,9 +1565,7 @@ def test_reset_budget_for_tags_linked_to_budgets_invalidates_management_cache( job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) - counter_cache.user_api_key_cache.async_delete_cache.assert_any_await( - key="tag:tenant-42" - ) + counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="tag:tenant-42") def test_reset_budget_for_tags_linked_to_budgets_invalidates_each_tag_management_cache( @@ -1684,8 +1590,7 @@ def test_reset_budget_for_tags_linked_to_budgets_invalidates_each_tag_management asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) deleted_keys = { - call.kwargs.get("key") - for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list + call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list } assert deleted_keys == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"} @@ -1711,19 +1616,13 @@ def test_reset_budget_for_keys_linked_to_budgets_invalidates_management_cache( linked_key = type("Key", (), {"token": "sk-linked"}) prisma_client = MagicMock() - prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[linked_key] - ) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[linked_key]) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget])) - counter_cache.user_api_key_cache.async_delete_cache.assert_any_await( - key="sk-linked" - ) + counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="sk-linked") def test_reset_budget_for_orgs_linked_to_budgets_invalidates_management_cache( @@ -1736,19 +1635,14 @@ def test_reset_budget_for_orgs_linked_to_budgets_invalidates_management_cache( linked_org = type("Org", (), {"organization_id": "org-acme"}) prisma_client = MagicMock() - prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[linked_org] - ) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[linked_org]) + prisma_client.db.litellm_organizationtable.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget])) deleted_keys = { - call.kwargs.get("key") - for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list + call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list } assert deleted_keys == { "org_id:org-acme", @@ -1768,19 +1662,13 @@ def test_reset_budget_for_team_members_invalidates_management_cache(monkeypatch) ) prisma_client = MagicMock() - prisma_client.db.litellm_teammembership.find_many = AsyncMock( - return_value=[membership] - ) - prisma_client.db.litellm_teammembership.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[membership]) + prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) - counter_cache.user_api_key_cache.async_delete_cache.assert_any_await( - key="team-x_alice" - ) + counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="team-x_alice") def test_reset_budget_for_tags_linked_to_budgets_management_cache_delete_failure_still_resets( @@ -1788,9 +1676,7 @@ def test_reset_budget_for_tags_linked_to_budgets_management_cache_delete_failure ): """If ``async_delete_cache`` raises, the DB cascade must still complete.""" counter_cache = _make_counter_invalidation_job(monkeypatch) - counter_cache.user_api_key_cache.async_delete_cache = AsyncMock( - side_effect=RuntimeError("cache unavailable") - ) + counter_cache.user_api_key_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("cache unavailable")) expired_budget = type("B", (), {"budget_id": "budget-1"}) linked_tag = type("Tag", (), {"tag_name": "tenant-42"}) diff --git a/tests/test_litellm/proxy/common_utils/test_timezone_utils.py b/tests/test_litellm/proxy/common_utils/test_timezone_utils.py index 80b813226df..7f686c53c95 100644 --- a/tests/test_litellm/proxy/common_utils/test_timezone_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_timezone_utils.py @@ -1,19 +1,33 @@ import os import sys -from datetime import datetime, timezone +from datetime import datetime, time, timezone from zoneinfo import ZoneInfo +import pytest + sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path import litellm from litellm.proxy.common_utils.timezone_utils import ( + BudgetResetSettings, + compute_budget_reset_at, + get_budget_reset_settings, get_budget_reset_time, get_budget_reset_timezone, + parse_budget_reset_time, ) +def _restore_attr(obj, name, original): + if original is None: + if hasattr(obj, name): + delattr(obj, name) + else: + setattr(obj, name, original) + + def test_get_budget_reset_time(): """ Test that the budget reset time is set to the first of the next month @@ -100,3 +114,69 @@ def test_get_budget_reset_time_respects_timezone(): delattr(litellm, "timezone") else: litellm.timezone = original + + +def test_parse_budget_reset_time_hh_mm(): + assert parse_budget_reset_time("12:00") == time(12, 0) + + +def test_parse_budget_reset_time_hh_mm_ss(): + assert parse_budget_reset_time("09:30:15") == time(9, 30, 15) + + +def test_parse_budget_reset_time_unset_defaults_to_midnight(): + assert parse_budget_reset_time(None) == time(0, 0) + assert parse_budget_reset_time("") == time(0, 0) + + +def test_parse_budget_reset_time_invalid_string_raises(): + with pytest.raises(ValueError): + parse_budget_reset_time("25:00") + with pytest.raises(ValueError): + parse_budget_reset_time("noon") + + +def test_parse_budget_reset_time_non_string_raises(): + # Unquoted "12:00" in YAML parses to the int 720; it must fail loudly, + # not silently fall back to midnight. + with pytest.raises(ValueError): + parse_budget_reset_time(720) + + +def test_get_budget_reset_settings_reads_globals(): + orig_tz = getattr(litellm, "timezone", None) + orig_rt = getattr(litellm, "budget_reset_time", None) + try: + litellm.timezone = "Asia/Jerusalem" + litellm.budget_reset_time = "12:00" + settings = get_budget_reset_settings() + assert settings.timezone == "Asia/Jerusalem" + assert settings.reset_time_of_day == time(12, 0) + finally: + _restore_attr(litellm, "timezone", orig_tz) + _restore_attr(litellm, "budget_reset_time", orig_rt) + + +def test_compute_budget_reset_at_applies_offset(): + settings = BudgetResetSettings( + timezone="Asia/Jerusalem", reset_time_of_day=time(12, 0) + ) + reset_at = compute_budget_reset_at("1d", settings) + jerusalem = reset_at.astimezone(ZoneInfo("Asia/Jerusalem")) + assert jerusalem.hour == 12 + assert jerusalem.minute == 0 + assert reset_at > datetime.now(timezone.utc) + + +def test_get_budget_reset_time_honors_global_budget_reset_time(): + orig_tz = getattr(litellm, "timezone", None) + orig_rt = getattr(litellm, "budget_reset_time", None) + try: + litellm.timezone = "UTC" + litellm.budget_reset_time = "12:00" + reset_at = get_budget_reset_time(budget_duration="1d") + assert reset_at.astimezone(timezone.utc).hour == 12 + assert reset_at.astimezone(timezone.utc).minute == 0 + finally: + _restore_attr(litellm, "timezone", orig_tz) + _restore_attr(litellm, "budget_reset_time", orig_rt) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.test.ts new file mode 100644 index 00000000000..5eb3bdc105d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.test.ts @@ -0,0 +1,103 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useSetKeyBlockedState, setKeyBlockedState } from "./useSetKeyBlockedState"; +import { apiClient } from "@/components/networking"; + +vi.mock("@/components/networking", () => ({ + apiClient: { post: vi.fn() }, +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +const mockPost = vi.mocked(apiClient.post); + +const createWrapper = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false }, mutations: { retry: false } } }); + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + return { queryClient, wrapper }; +}; + +describe("setKeyBlockedState", () => { + beforeEach(() => { + mockPost.mockReset(); + }); + + it("POSTs the key hash to /key/block when blocking", async () => { + mockPost.mockResolvedValueOnce({ blocked: true }); + + const result = await setKeyBlockedState("sk-access", { keyToken: "hashed-token", blocked: true }); + + expect(mockPost).toHaveBeenCalledWith("/key/block", { + accessToken: "sk-access", + body: { key: "hashed-token" }, + }); + expect(result).toEqual({ blocked: true }); + }); + + it("POSTs the key hash to /key/unblock when unblocking", async () => { + mockPost.mockResolvedValueOnce({ blocked: false }); + + const result = await setKeyBlockedState("sk-access", { keyToken: "hashed-token", blocked: false }); + + expect(mockPost).toHaveBeenCalledWith("/key/unblock", { + accessToken: "sk-access", + body: { key: "hashed-token" }, + }); + expect(result).toEqual({ blocked: false }); + }); + + it("falls back to the requested state when the response has no blocked field", async () => { + mockPost.mockResolvedValueOnce(null); + + const result = await setKeyBlockedState("sk-access", { keyToken: "hashed-token", blocked: true }); + + expect(result).toEqual({ blocked: true }); + }); +}); + +describe("useSetKeyBlockedState", () => { + beforeEach(() => { + mockPost.mockReset(); + mockUseAuthorized.mockReturnValue({ accessToken: "sk-access" }); + }); + + it("invalidates key queries after a successful mutation", async () => { + mockPost.mockResolvedValueOnce({ blocked: true }); + const { queryClient, wrapper } = createWrapper(); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + + const { result } = renderHook(() => useSetKeyBlockedState(), { wrapper }); + result.current.mutate({ keyToken: "hashed-token", blocked: true }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ["keys"] }); + }); + + it("surfaces request failures as mutation errors", async () => { + mockPost.mockRejectedValueOnce(new Error("Key not found.")); + const { wrapper } = createWrapper(); + + const { result } = renderHook(() => useSetKeyBlockedState(), { wrapper }); + result.current.mutate({ keyToken: "missing", blocked: true }); + + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(result.current.error?.message).toBe("Key not found."); + }); + + it("errors without an access token", async () => { + mockUseAuthorized.mockReturnValue({ accessToken: null }); + const { wrapper } = createWrapper(); + + const { result } = renderHook(() => useSetKeyBlockedState(), { wrapper }); + result.current.mutate({ keyToken: "hashed-token", blocked: true }); + + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(mockPost).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.ts new file mode 100644 index 00000000000..792ef567f99 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.ts @@ -0,0 +1,45 @@ +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { apiClient } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { keyKeys } from "./useKeys"; + +export interface SetKeyBlockedStateInput { + keyToken: string; + blocked: boolean; +} + +export interface SetKeyBlockedStateResult { + blocked: boolean; +} + +interface BlockKeyResponse { + blocked?: boolean | null; +} + +export const setKeyBlockedState = async ( + accessToken: string, + { keyToken, blocked }: SetKeyBlockedStateInput, +): Promise => { + const response = await apiClient.post(blocked ? "/key/block" : "/key/unblock", { + accessToken, + body: { key: keyToken }, + }); + return { blocked: response?.blocked ?? blocked }; +}; + +export const useSetKeyBlockedState = () => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async (input) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return setKeyBlockedState(accessToken, input); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: keyKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 513054aae7a..95f45ea199e 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -174,6 +174,13 @@ it("should render VirtualKeysTable component", () => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); +it("shows the Budget Reset column by default", async () => { + renderWithProviders(); + await waitFor(() => { + expect(screen.getByText("Budget Reset")).toBeInTheDocument(); + }); +}); + it("left-anchors the create-key CTA below the title, between the header and the table toolbar", () => { renderWithProviders(Create New Key} />); @@ -498,8 +505,13 @@ describe("Status column reflects blocked / expiry / scim metadata", () => { renderWithProviders(); + const tag = await screen.findByTestId(`key-status-${mockKey.token_id}`); + expect(tag).toHaveTextContent("Active"); + + const user = userEvent.setup(); + await user.hover(tag); await waitFor(() => { - expect(screen.getByTestId(`key-status-${mockKey.token_id}`)).toHaveTextContent("Active"); + expect(screen.getByText(/not blocked and has not expired/i)).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index 133ff89a898..fdbc07ee020 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -46,7 +46,11 @@ const getKeyStatus = (key: KeyResponse): KeyStatus => { if (!Number.isNaN(expiresAt) && expiresAt < Date.now()) { return { tone: "warning", label: "Expired", tooltip: "This key has passed its expiry date." }; } - return { tone: "success", label: "Active" }; + return { + tone: "success", + label: "Active", + tooltip: "This key is not blocked and has not expired.", + }; }; const UserPopoverCell = ({ @@ -359,6 +363,5 @@ export const KEY_TABLE_HIDDEN_COLUMNS: Record = { created_by: false, updated_at: false, expires: false, - budget_reset_at: false, rate_limits: false, }; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index d44a491b840..d6e9ba5665c 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -372,7 +372,7 @@ export function getGlobalLitellmHeaderName(): string { return globalLitellmHeaderName; } -const apiClient = createApiClient({ +export const apiClient = createApiClient({ getBaseUrl: getProxyBaseUrl, getAuthHeaderName: getGlobalLitellmHeaderName, onError: handleError, diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx index 14750787609..f67f70e7df9 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx @@ -60,22 +60,16 @@ describe("KeyInfoHeader", () => { }); describe("action buttons", () => { - it("should show Regenerate and Delete buttons by default", () => { + it("should show Regenerate button and actions dropdown by default", () => { render(); expect(screen.getByRole("button", { name: /regenerate key/i })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: /delete key/i })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /more key actions/i })).toBeInTheDocument(); }); - it("should show Regenerate and Delete buttons when canModifyKey is true", () => { - render(); - expect(screen.getByRole("button", { name: /regenerate key/i })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: /delete key/i })).toBeInTheDocument(); - }); - - it("should hide Regenerate and Delete buttons when canModifyKey is false", () => { + it("should hide Regenerate button and actions dropdown when canModifyKey is false", () => { render(); expect(screen.queryByRole("button", { name: /regenerate key/i })).not.toBeInTheDocument(); - expect(screen.queryByRole("button", { name: /delete key/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /more key actions/i })).not.toBeInTheDocument(); }); it("should call onRegenerate when Regenerate Key is clicked", async () => { @@ -85,13 +79,6 @@ describe("KeyInfoHeader", () => { expect(onRegenerate).toHaveBeenCalledTimes(1); }); - it("should call onDelete when Delete Key is clicked", async () => { - const onDelete = vi.fn(); - render(); - await userEvent.click(screen.getByRole("button", { name: /delete key/i })); - expect(onDelete).toHaveBeenCalledTimes(1); - }); - it("should disable Regenerate button when regenerateDisabled is true", () => { render(); expect(screen.getByRole("button", { name: /regenerate key/i })).toBeDisabled(); @@ -103,6 +90,79 @@ describe("KeyInfoHeader", () => { }); }); + describe("destructive actions dropdown", () => { + const openDropdown = async () => { + await userEvent.click(screen.getByRole("button", { name: /more key actions/i })); + }; + + it("should list Block Key, Reset Spend, and Delete Key when all handlers are provided", async () => { + render(); + await openDropdown(); + expect(await screen.findByRole("menuitem", { name: /block key/i })).toBeInTheDocument(); + expect(screen.getByRole("menuitem", { name: /reset spend/i })).toBeInTheDocument(); + expect(screen.getByRole("menuitem", { name: /delete key/i })).toBeInTheDocument(); + }); + + it("should omit Block Key and Reset Spend when their handlers are not provided", async () => { + render(); + await openDropdown(); + expect(await screen.findByRole("menuitem", { name: /delete key/i })).toBeInTheDocument(); + expect(screen.queryByRole("menuitem", { name: /block key/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("menuitem", { name: /reset spend/i })).not.toBeInTheDocument(); + }); + + it("should show Unblock Key instead of Block Key when the key is blocked", async () => { + render(); + await openDropdown(); + expect(await screen.findByRole("menuitem", { name: /unblock key/i })).toBeInTheDocument(); + expect(screen.queryByRole("menuitem", { name: /^block key/i })).not.toBeInTheDocument(); + }); + + it("should call onToggleBlocked when Block Key is clicked", async () => { + const onToggleBlocked = vi.fn(); + render(); + await openDropdown(); + await userEvent.click(await screen.findByRole("menuitem", { name: /block key/i })); + expect(onToggleBlocked).toHaveBeenCalledTimes(1); + }); + + it("should call onToggleBlocked when Unblock Key is clicked", async () => { + const onToggleBlocked = vi.fn(); + render(); + await openDropdown(); + await userEvent.click(await screen.findByRole("menuitem", { name: /unblock key/i })); + expect(onToggleBlocked).toHaveBeenCalledTimes(1); + }); + + it("should call onResetSpend when Reset Spend is clicked", async () => { + const onResetSpend = vi.fn(); + render(); + await openDropdown(); + await userEvent.click(await screen.findByRole("menuitem", { name: /reset spend/i })); + expect(onResetSpend).toHaveBeenCalledTimes(1); + }); + + it("should call onDelete when Delete Key is clicked", async () => { + const onDelete = vi.fn(); + render(); + await openDropdown(); + await userEvent.click(await screen.findByRole("menuitem", { name: /delete key/i })); + expect(onDelete).toHaveBeenCalledTimes(1); + }); + }); + + describe("blocked tag", () => { + it("should show a Blocked tag when isBlocked is true", () => { + render(); + expect(screen.getByText("Blocked")).toBeInTheDocument(); + }); + + it("should not show a Blocked tag by default", () => { + render(); + expect(screen.queryByText("Blocked")).not.toBeInTheDocument(); + }); + }); + describe("Create New Key button", () => { it("should show when onCreateNew is provided", () => { render(); diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx index c9c393352cc..d0dd782a697 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx @@ -1,5 +1,6 @@ import React from "react"; -import { Button, Typography, Tooltip, Space, Divider, Flex, Popover } from "antd"; +import { Button, Typography, Tooltip, Space, Divider, Flex, Popover, Dropdown, Tag } from "antd"; +import type { MenuProps } from "antd"; import { ArrowLeftOutlined, SyncOutlined, @@ -12,6 +13,9 @@ import { SafetyCertificateOutlined, TransactionOutlined, FieldTimeOutlined, + MoreOutlined, + StopOutlined, + CheckCircleOutlined, } from "@ant-design/icons"; import LabeledField from "../common_components/LabeledField"; import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; @@ -38,6 +42,8 @@ interface KeyInfoHeaderProps { onRegenerate?: () => void; onDelete?: () => void; onResetSpend?: () => void; + onToggleBlocked?: () => void; + isBlocked?: boolean; canModifyKey?: boolean; backButtonText?: string; regenerateDisabled?: boolean; @@ -133,11 +139,33 @@ export function KeyInfoHeader({ onRegenerate, onDelete, onResetSpend, + onToggleBlocked, + isBlocked = false, canModifyKey = true, backButtonText = "Back to Keys", regenerateDisabled = false, regenerateTooltip, }: KeyInfoHeaderProps) { + const destructiveActionItems: MenuProps["items"] = [ + ...(onToggleBlocked + ? [ + isBlocked + ? { key: "unblock", label: "Unblock Key", icon: } + : { key: "block", label: "Block Key", icon: , danger: true }, + ] + : []), + ...(onResetSpend + ? [{ key: "reset-spend", label: "Reset Spend", icon: , danger: true }] + : []), + { key: "delete", label: "Delete Key", icon: , danger: true }, + ]; + + const handleDestructiveActionClick: MenuProps["onClick"] = ({ key }) => { + if (key === "block" || key === "unblock") onToggleBlocked?.(); + if (key === "reset-spend") onResetSpend?.(); + if (key === "delete") onDelete?.(); + }; + return (
{onCreateNew && ( @@ -156,9 +184,16 @@ export function KeyInfoHeader({
- - {data.keyName} - + + + {data.keyName} + + {isBlocked && ( + }> + Blocked + + )} + Key ID: {data.keyId} @@ -172,14 +207,12 @@ export function KeyInfoHeader({ - {onResetSpend && ( - - )} - + +
+
+ Budget Reset + + {currentKeyData.budget_reset_at + ? `${currentKeyData.budget_duration ? `Every ${currentKeyData.budget_duration}, next ` : ""}${formatTimestamp(currentKeyData.budget_reset_at)}` + : "Never"} + +
+ {currentKeyData.budget_fallbacks && Object.keys(currentKeyData.budget_fallbacks).length > 0 && (
Budget Fallbacks diff --git a/uv.lock b/uv.lock index 1dfa2c1201c..cee24aca330 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-15T21:54:47.972166Z" +exclude-newer = "2026-07-18T19:44:23.519632Z" exclude-newer-span = "P3D" [manifest] @@ -6816,11 +6816,11 @@ wheels = [ [[package]] name = "pyasn1" -version = "0.6.3" +version = "0.6.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5c/5f/6583902b6f79b399c9c40674ac384fd9cd77805f9e6205075f828ef11fb2/pyasn1-0.6.3.tar.gz", hash = "sha256:697a8ecd6d98891189184ca1fa05d1bb00e2f84b5977c481452050549c8a72cf", size = 148685, upload-time = "2026-03-17T01:06:53.382Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a4/9a/23310166d960def5897e91fe20e5b724601b02a22e84ba1f94232c0b7f67/pyasn1-0.6.4.tar.gz", hash = "sha256:9c447d8431c947fe4c8febc4ed9e760bc29011a5b01e5c74b67025bd9fb8ce81", size = 151262, upload-time = "2026-07-09T01:12:33.988Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5d/a0/7d793dce3fa811fe047d6ae2431c672364b462850c6235ae306c0efd025f/pyasn1-0.6.3-py3-none-any.whl", hash = "sha256:a80184d120f0864a52a073acc6fc642847d0be408e7c7252f31390c0f4eadcde", size = 83997, upload-time = "2026-03-17T01:06:52.036Z" }, + { url = "https://files.pythonhosted.org/packages/9a/3b/6163796d69c3977d1e4287bea4a6979161cbbdd170ebb430511e8e1999ce/pyasn1-0.6.4-py3-none-any.whl", hash = "sha256:deda9277cfd454080ec40b207fb6df82206a3a2688735233cdcd8d3d565f088b", size = 84410, upload-time = "2026-07-09T01:12:32.92Z" }, ] [[package]] @@ -7140,14 +7140,14 @@ wheels = [ [[package]] name = "pypdf" -version = "6.13.3" +version = "6.14.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/17/18/9947cc201af9ccf76720fd3347bf4f70eb882ce3fcf4cb05f7443e4cf871/pypdf-6.13.3.tar.gz", hash = "sha256:f3cb822769725f1bac658c406cfc9460399043f3750c2d3e4650e0a85eacabd7", size = 6484063, upload-time = "2026-06-17T15:22:00.898Z" } +sdist = { url = "https://files.pythonhosted.org/packages/03/72/7dfd5ff1c9c37de97a731701f51af091325f123d9d4270361c9c69e4431f/pypdf-6.14.2.tar.gz", hash = "sha256:7873f502fe4385e79539b21d872392dc0c4e3714327c15881cbc7fbfd1f95b25", size = 6491182, upload-time = "2026-06-23T14:18:30.859Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/94/56/2967e621598987905fb8cdfadd8f8de6b5c68c9351f0523c4df8409f28f1/pypdf-6.13.3-py3-none-any.whl", hash = "sha256:c6e3f86afb625791510b02ad5480e94b63970bb957df75d44657c282ecc52224", size = 347288, upload-time = "2026-06-17T15:21:59.512Z" }, + { url = "https://files.pythonhosted.org/packages/49/e6/136aa8993a2ae7214e0b0ef2edaa0d2e08d1d4e4982635b08a835ff31ec8/pypdf-6.14.2-py3-none-any.whl", hash = "sha256:3f07891af76dc002657e04993ab9b4de81de29f9013b9761d0b7968bff12e946", size = 349514, upload-time = "2026-06-23T14:18:28.867Z" }, ] [[package]]