mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-30 01:52:18 +00:00
Merge remote-tracking branch 'origin/main' into litellm_lit8140_flux2_input_image_billing
This commit is contained in:
shrey kharbanda
commit
3d7b85bdd0
19 changed files with 2364 additions and 273 deletions
|
|
@ -28398,6 +28398,8 @@
|
|||
"input_cost_per_image": 0.0011,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"input_cost_per_token_batches": 1e-06,
|
||||
"input_cost_per_token_flex": 1e-06,
|
||||
"input_cost_per_token_priority": 3.6e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 32768,
|
||||
|
|
@ -28409,6 +28411,8 @@
|
|||
"rpm": 1000,
|
||||
"tpm": 4000000,
|
||||
"output_cost_per_token_batches": 6e-06,
|
||||
"output_cost_per_token_flex": 6e-06,
|
||||
"output_cost_per_token_priority": 2.16e-05,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
|
|
@ -41395,20 +41399,20 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/deepseek/deepseek-v4-pro": {
|
||||
"input_cost_per_token": 9.1263e-07,
|
||||
"input_cost_per_token": 8.44944e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 384000,
|
||||
"max_tokens": 384000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.82526e-06,
|
||||
"output_cost_per_token": 1.689888e-06,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 7.60525e-08,
|
||||
"cache_read_input_token_cost": 7.0412e-08,
|
||||
"supports_audio_input": false,
|
||||
"supports_pdf_input": false,
|
||||
"supports_vision": false,
|
||||
|
|
@ -56211,6 +56215,7 @@
|
|||
"input_cost_per_image_token": 1e-06,
|
||||
"input_cost_per_token": 7.5e-07,
|
||||
"input_cost_per_video_per_second": 3.3333333333333335e-05,
|
||||
"input_cost_per_video_token": 1e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 65536,
|
||||
|
|
@ -56403,6 +56408,7 @@
|
|||
"input_cost_per_image_token": 1e-06,
|
||||
"input_cost_per_token": 7.5e-07,
|
||||
"input_cost_per_video_per_second": 3.3333333333333335e-05,
|
||||
"input_cost_per_video_token": 1e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 65536,
|
||||
|
|
@ -60144,6 +60150,23 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": {
|
||||
"cache_read_input_token_cost": 9e-09,
|
||||
"input_cost_per_token": 4.5e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 393216,
|
||||
"max_tokens": 393216,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-06,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp": {
|
||||
"cache_read_input_token_cost": 7e-09,
|
||||
"deprecation_date": "2026-09-25",
|
||||
|
|
@ -60223,6 +60246,23 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/deepseek-v4p1-flash-us": {
|
||||
"cache_read_input_token_cost": 9e-09,
|
||||
"input_cost_per_token": 4.5e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 393216,
|
||||
"max_tokens": 393216,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-06,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/deepseek-v4-flash-vision-exp": {
|
||||
"cache_read_input_token_cost": 7e-09,
|
||||
"deprecation_date": "2026-09-25",
|
||||
|
|
@ -60319,13 +60359,16 @@
|
|||
},
|
||||
"fireworks_ai/kimi-k3-us": {
|
||||
"cache_read_input_token_cost": 4.5e-07,
|
||||
"cache_read_input_token_cost_priority": 5.625e-07,
|
||||
"input_cost_per_token": 4.5e-06,
|
||||
"input_cost_per_token_priority": 5.625e-06,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.25e-05,
|
||||
"output_cost_per_token_priority": 2.8125e-05,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
|
|
@ -60527,13 +60570,16 @@
|
|||
},
|
||||
"fireworks_ai/accounts/fireworks/routers/kimi-k3-us": {
|
||||
"cache_read_input_token_cost": 4.5e-07,
|
||||
"cache_read_input_token_cost_priority": 5.625e-07,
|
||||
"input_cost_per_token": 4.5e-06,
|
||||
"input_cost_per_token_priority": 5.625e-06,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.25e-05,
|
||||
"output_cost_per_token_priority": 2.8125e-05,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
|
|
@ -63292,6 +63338,22 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/routers/glm-5p3-us": {
|
||||
"cache_read_input_token_cost": 3.9e-07,
|
||||
"input_cost_per_token": 2.1e-06,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"fireworks_ai/glm-5p3": {
|
||||
"cache_read_input_token_cost": 2.6e-07,
|
||||
"cache_read_input_token_cost_priority": 3.25e-07,
|
||||
|
|
@ -63311,6 +63373,22 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"fireworks_ai/glm-5p3-us": {
|
||||
"cache_read_input_token_cost": 3.9e-07,
|
||||
"input_cost_per_token": 2.1e-06,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": {
|
||||
"cache_read_input_token_cost": 3.9e-07,
|
||||
"input_cost_per_token": 2.1e-06,
|
||||
|
|
@ -63358,6 +63436,20 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/routers/glm-5p3-flash-us": {
|
||||
"cache_read_input_token_cost": 4.5e-08,
|
||||
"input_cost_per_token": 2.25e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-07,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/glm-5p3-flash": {
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
"cache_read_input_token_cost_priority": 3.75e-08,
|
||||
|
|
@ -63375,6 +63467,20 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/glm-5p3-flash-us": {
|
||||
"cache_read_input_token_cost": 4.5e-08,
|
||||
"input_cost_per_token": 2.25e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-07,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/models/inkling": {
|
||||
"cache_read_input_token_cost": 1.7e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
|
|
@ -65906,9 +66012,9 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/z-ai/glm-5.3-flash": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"output_cost_per_token": 5e-07,
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"input_cost_per_token": 4.5e-08,
|
||||
"output_cost_per_token": 6e-07,
|
||||
"cache_read_input_token_cost": 2.85e-08,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1310720,
|
||||
"max_output_tokens": 943718,
|
||||
|
|
@ -66657,9 +66763,9 @@
|
|||
"supports_web_search": true
|
||||
},
|
||||
"openrouter/deepseek/deepseek-v4-flash": {
|
||||
"input_cost_per_token": 8.4e-08,
|
||||
"output_cost_per_token": 1.68e-07,
|
||||
"cache_read_input_token_cost": 1.68e-08,
|
||||
"input_cost_per_token": 4.9e-08,
|
||||
"output_cost_per_token": 9.8e-08,
|
||||
"cache_read_input_token_cost": 9.8e-09,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 384000,
|
||||
|
|
@ -73050,6 +73156,26 @@
|
|||
"supports_vision": true,
|
||||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/mistralai/mistral-large-2512": {
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"input_cost_per_token": 5e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 209715,
|
||||
"max_tokens": 209715,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-06,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/mistralai/mistral-large-2512:batch": {
|
||||
"cache_read_input_token_cost": 2.5e-08,
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
|
|
|
|||
|
|
@ -227,6 +227,9 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
|||
DeploymentAffinityCheck,
|
||||
warn_on_unknown_model_group_affinity_flags,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
|
||||
build_io_token_rate_limit_headers,
|
||||
deployment_has_io_token_limits,
|
||||
|
|
@ -438,6 +441,7 @@ _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = M
|
|||
{
|
||||
"prompt_caching": PromptCachingDeploymentCheck,
|
||||
"enforce_model_rate_limits": ModelRateLimitingCheck,
|
||||
"encrypted_content_affinity": EncryptedContentAffinityCheck,
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -2207,10 +2211,6 @@ class Router:
|
|||
)
|
||||
|
||||
def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None:
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
def _move_before_deployment_affinity(
|
||||
callback_list: list[Any],
|
||||
callback_to_move: EncryptedContentAffinityCheck,
|
||||
|
|
|
|||
|
|
@ -36,17 +36,10 @@ Safe to enable globally:
|
|||
- No cache required.
|
||||
"""
|
||||
|
||||
import time
|
||||
from collections.abc import Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Final, Optional, Protocol, cast
|
||||
|
||||
import httpx
|
||||
from typing import TYPE_CHECKING, Final, Optional, cast
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.exceptions import (
|
||||
RateLimitError,
|
||||
ServiceUnavailableError,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger, Span
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
|
|
@ -55,7 +48,6 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
strip_encrypted_reasoning_from_messages,
|
||||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.router_utils.cooldown_cache import CooldownCacheValue
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
|
|
@ -63,14 +55,6 @@ if TYPE_CHECKING:
|
|||
from litellm.router import Router
|
||||
|
||||
|
||||
class _SupportsActiveCooldowns(Protocol):
|
||||
"""Cooldown-cache handle: this check only reads back the currently active cooldowns."""
|
||||
|
||||
async def async_get_active_cooldowns(
|
||||
self, model_ids: list[str], parent_otel_span: Span | None
|
||||
) -> list[tuple[str, CooldownCacheValue]]: ...
|
||||
|
||||
|
||||
class EncryptedContentAffinityCheck(CustomLogger):
|
||||
"""
|
||||
Routes follow-up Responses API requests to the deployment that produced
|
||||
|
|
@ -194,23 +178,6 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
return deployment
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _request_team_id(request_kwargs: Mapping[str, object]) -> str | None:
|
||||
containers: Final = (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata"))
|
||||
team_ids: Final = (c.get("user_api_key_team_id") for c in containers if isinstance(c, Mapping))
|
||||
return next((tid for tid in team_ids if isinstance(tid, str)), None)
|
||||
|
||||
def _routed_group_candidate_model_ids(self, request_kwargs: Mapping[str, object], model: str) -> frozenset[str]:
|
||||
"""
|
||||
Deployment ids that could serve this turn's routed ``model``, as the router
|
||||
resolves a route (model_group_alias / routing group / model_name / team /
|
||||
pattern). Delegates to the router so the full precedence is not re-derived here
|
||||
and no deployment ids are written into request kwargs bound for the provider.
|
||||
"""
|
||||
if self.router is None:
|
||||
return frozenset()
|
||||
return self.router.get_candidate_model_ids_for_route(model=model, team_id=self._request_team_id(request_kwargs))
|
||||
|
||||
@staticmethod
|
||||
def _encryption_boundary_key(
|
||||
litellm_params: object,
|
||||
|
|
@ -262,9 +229,6 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
Deployments in ``healthy_deployments`` sharing the originating
|
||||
deployment's ``(api_base, api_key)``, alongside the originating
|
||||
deployment object (or ``None`` if it was removed / router unavailable).
|
||||
Returns ``([], originating_or_None)`` when no boundary match exists,
|
||||
so the caller can reuse the looked-up ``originating`` rather than
|
||||
re-querying the router.
|
||||
"""
|
||||
if self.router is None:
|
||||
return [], None
|
||||
|
|
@ -294,18 +258,12 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
"""
|
||||
If the request ``input`` contains litellm-encoded item IDs, or its Anthropic
|
||||
``messages`` replay a bridge-tagged thinking block, decode the embedded
|
||||
``model_id`` and pin the request to that deployment. Raises
|
||||
``RateLimitError`` / ``ServiceUnavailableError`` when the originating
|
||||
deployment is a member of the routed model group but currently unavailable
|
||||
and no encryption-boundary peer exists, rather than dispatching a doomed
|
||||
request to a non-peer deployment. When the origin is not a member of the
|
||||
routed group (an auto-router tier change, a model switch with no peer, a
|
||||
removed deployment, or an unknown/forged marker), the encrypted reasoning is
|
||||
stripped and the request dispatches with its readable history instead. The
|
||||
429/503 split mirrors the originating cooldown's status:
|
||||
a 429-induced cooldown surfaces as 429 (with ``Retry-After`` set to the
|
||||
remaining cooldown window) so OpenAI-compatible clients back off and
|
||||
retry after the deployment is eligible again.
|
||||
``model_id`` and pin the request to that deployment. When the origin cannot
|
||||
serve this turn and no encryption-boundary peer is configured (it is
|
||||
unhealthy, the request was routed to a different group by an auto-router tier
|
||||
change or model switch, or the marker is removed/unknown/forged), the
|
||||
encrypted reasoning is stripped and the request dispatches to the healthy
|
||||
pool with its readable history instead of failing.
|
||||
"""
|
||||
request_kwargs = request_kwargs or {}
|
||||
typed_healthy_deployments: Final = cast(list[dict], healthy_deployments)
|
||||
|
|
@ -348,7 +306,7 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
return [deployment]
|
||||
|
||||
# Follow-up switched model_name (LIT-2531): pin by Azure resource instead.
|
||||
boundary_matches, originating = self._find_deployments_on_same_encryption_boundary(
|
||||
boundary_matches, _originating = self._find_deployments_on_same_encryption_boundary(
|
||||
healthy_deployments=typed_healthy_deployments,
|
||||
model_id=model_id,
|
||||
)
|
||||
|
|
@ -362,101 +320,17 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
request_kwargs["_encrypted_content_affinity_pinned"] = True
|
||||
return boundary_matches
|
||||
|
||||
# The origin cannot serve this turn's routed group and no peer shares the boundary, so its
|
||||
# The origin cannot serve this turn and no peer shares its encryption boundary, so its
|
||||
# encrypted reasoning can never decrypt here. Strip it, keep the readable history, and dispatch
|
||||
# to the routed group instead of failing. Membership is tested by deployment id against the set
|
||||
# the router actually resolved for this route, not by model-group name, so an alias, a
|
||||
# provider-qualified spelling, a team-public name, or a pattern route of the same group is not
|
||||
# mistaken for a tier change. An unknown origin (a removed deployment, or a forged marker) is
|
||||
# treated the same as a cross-group one, which also denies an authenticated caller a
|
||||
# deployment-id existence oracle: a real cross-group id and a nonexistent id both strip and
|
||||
# dispatch rather than returning distinguishable responses. Only a genuine same-group member
|
||||
# that is currently unavailable falls through to the fail-fast, preserving the cooldown contract.
|
||||
routed_group_model_ids: Final = (
|
||||
self._routed_group_candidate_model_ids(request_kwargs, model) if originating is not None else frozenset()
|
||||
# to the healthy pool instead of failing the request. This also denies an authenticated caller a
|
||||
# deployment-id existence oracle: a same-group id, a cross-group id, a removed id and a forged
|
||||
# marker all strip and dispatch rather than returning distinguishable responses.
|
||||
verbose_router_logger.warning(
|
||||
"EncryptedContentAffinityCheck: model_id=%s cannot serve group %s and no deployment on the same "
|
||||
"encryption boundary is configured; forwarding without its encrypted reasoning",
|
||||
model_id[:64],
|
||||
model,
|
||||
)
|
||||
if str(model_id) not in routed_group_model_ids:
|
||||
verbose_router_logger.debug(
|
||||
"EncryptedContentAffinityCheck: model_id=%s is not a candidate for the routed group %s; "
|
||||
"forwarding without its encrypted reasoning",
|
||||
model_id,
|
||||
model,
|
||||
)
|
||||
ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input)
|
||||
strip_encrypted_reasoning_from_messages(anthropic_messages)
|
||||
return typed_healthy_deployments
|
||||
|
||||
# The origin is a member of the routed group but currently unavailable (cooled down); fail fast
|
||||
# rather than dispatching to a non-peer, which would guarantee an upstream 400.
|
||||
raise await self._unavailable_origin_error(
|
||||
model=model,
|
||||
model_id=model_id,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
||||
async def _unavailable_origin_error(
|
||||
self,
|
||||
model: str,
|
||||
model_id: str,
|
||||
parent_otel_span: Span | None,
|
||||
) -> Exception:
|
||||
# Public error messages intentionally omit the originating ``model_id`` so
|
||||
# an authenticated caller forging encrypted-content markers cannot use the
|
||||
# error surface to enumerate which deployment IDs exist on this router.
|
||||
cooldown: Final = await self._get_origin_cooldown(model_id=model_id, parent_otel_span=parent_otel_span)
|
||||
|
||||
if cooldown is not None and str(cooldown.get("status_code")) == "429":
|
||||
retry_after: Final = self._cooldown_seconds_remaining(cooldown)
|
||||
return RateLimitError(
|
||||
message=(
|
||||
"The deployment that produced this encrypted_content is "
|
||||
f"rate-limited (cooling down for ~{retry_after}s), and no "
|
||||
"deployment on the same encryption boundary is configured. "
|
||||
"Retry after the Retry-After window or configure a deployment "
|
||||
"with the same (api_base, api_key)."
|
||||
),
|
||||
llm_provider="",
|
||||
model=model,
|
||||
response=httpx.Response(
|
||||
status_code=429,
|
||||
headers={"retry-after": str(retry_after)},
|
||||
request=httpx.Request("POST", "https://litellm.ai/"),
|
||||
),
|
||||
)
|
||||
|
||||
return ServiceUnavailableError(
|
||||
message=(
|
||||
"The deployment that produced this encrypted_content is "
|
||||
"currently unavailable (likely cooled down), and no deployment "
|
||||
"on the same encryption boundary is configured. Retry later or "
|
||||
"configure a deployment with the same (api_base, api_key)."
|
||||
),
|
||||
llm_provider="",
|
||||
model=model,
|
||||
)
|
||||
|
||||
async def _get_origin_cooldown(
|
||||
self,
|
||||
model_id: str,
|
||||
parent_otel_span: Span | None,
|
||||
) -> CooldownCacheValue | None:
|
||||
if self.router is None:
|
||||
return None
|
||||
cooldown_cache: Final[_SupportsActiveCooldowns | None] = getattr(self.router, "cooldown_cache", None)
|
||||
if cooldown_cache is None:
|
||||
return None
|
||||
try:
|
||||
active: Final = await cooldown_cache.async_get_active_cooldowns(
|
||||
model_ids=[model_id], parent_otel_span=parent_otel_span
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
for cached_model_id, value in active:
|
||||
if cached_model_id == model_id:
|
||||
return value
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _cooldown_seconds_remaining(cooldown: CooldownCacheValue) -> int:
|
||||
remaining = float(cooldown.get("timestamp", 0.0)) + float(cooldown.get("cooldown_time", 0.0)) - time.time()
|
||||
return max(1, int(remaining))
|
||||
ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input)
|
||||
strip_encrypted_reasoning_from_messages(anthropic_messages)
|
||||
return typed_healthy_deployments
|
||||
|
|
|
|||
|
|
@ -28398,6 +28398,8 @@
|
|||
"input_cost_per_image": 0.0011,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"input_cost_per_token_batches": 1e-06,
|
||||
"input_cost_per_token_flex": 1e-06,
|
||||
"input_cost_per_token_priority": 3.6e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 32768,
|
||||
|
|
@ -28409,6 +28411,8 @@
|
|||
"rpm": 1000,
|
||||
"tpm": 4000000,
|
||||
"output_cost_per_token_batches": 6e-06,
|
||||
"output_cost_per_token_flex": 6e-06,
|
||||
"output_cost_per_token_priority": 2.16e-05,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
|
|
@ -41395,20 +41399,20 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/deepseek/deepseek-v4-pro": {
|
||||
"input_cost_per_token": 9.1263e-07,
|
||||
"input_cost_per_token": 8.44944e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 384000,
|
||||
"max_tokens": 384000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.82526e-06,
|
||||
"output_cost_per_token": 1.689888e-06,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 7.60525e-08,
|
||||
"cache_read_input_token_cost": 7.0412e-08,
|
||||
"supports_audio_input": false,
|
||||
"supports_pdf_input": false,
|
||||
"supports_vision": false,
|
||||
|
|
@ -56211,6 +56215,7 @@
|
|||
"input_cost_per_image_token": 1e-06,
|
||||
"input_cost_per_token": 7.5e-07,
|
||||
"input_cost_per_video_per_second": 3.3333333333333335e-05,
|
||||
"input_cost_per_video_token": 1e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 65536,
|
||||
|
|
@ -56403,6 +56408,7 @@
|
|||
"input_cost_per_image_token": 1e-06,
|
||||
"input_cost_per_token": 7.5e-07,
|
||||
"input_cost_per_video_per_second": 3.3333333333333335e-05,
|
||||
"input_cost_per_video_token": 1e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 65536,
|
||||
|
|
@ -60144,6 +60150,23 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": {
|
||||
"cache_read_input_token_cost": 9e-09,
|
||||
"input_cost_per_token": 4.5e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 393216,
|
||||
"max_tokens": 393216,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-06,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp": {
|
||||
"cache_read_input_token_cost": 7e-09,
|
||||
"deprecation_date": "2026-09-25",
|
||||
|
|
@ -60223,6 +60246,23 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/deepseek-v4p1-flash-us": {
|
||||
"cache_read_input_token_cost": 9e-09,
|
||||
"input_cost_per_token": 4.5e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 393216,
|
||||
"max_tokens": 393216,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-06,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/deepseek-v4-flash-vision-exp": {
|
||||
"cache_read_input_token_cost": 7e-09,
|
||||
"deprecation_date": "2026-09-25",
|
||||
|
|
@ -60319,13 +60359,16 @@
|
|||
},
|
||||
"fireworks_ai/kimi-k3-us": {
|
||||
"cache_read_input_token_cost": 4.5e-07,
|
||||
"cache_read_input_token_cost_priority": 5.625e-07,
|
||||
"input_cost_per_token": 4.5e-06,
|
||||
"input_cost_per_token_priority": 5.625e-06,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.25e-05,
|
||||
"output_cost_per_token_priority": 2.8125e-05,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
|
|
@ -60527,13 +60570,16 @@
|
|||
},
|
||||
"fireworks_ai/accounts/fireworks/routers/kimi-k3-us": {
|
||||
"cache_read_input_token_cost": 4.5e-07,
|
||||
"cache_read_input_token_cost_priority": 5.625e-07,
|
||||
"input_cost_per_token": 4.5e-06,
|
||||
"input_cost_per_token_priority": 5.625e-06,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.25e-05,
|
||||
"output_cost_per_token_priority": 2.8125e-05,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
|
|
@ -63292,6 +63338,22 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/routers/glm-5p3-us": {
|
||||
"cache_read_input_token_cost": 3.9e-07,
|
||||
"input_cost_per_token": 2.1e-06,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"fireworks_ai/glm-5p3": {
|
||||
"cache_read_input_token_cost": 2.6e-07,
|
||||
"cache_read_input_token_cost_priority": 3.25e-07,
|
||||
|
|
@ -63311,6 +63373,22 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"fireworks_ai/glm-5p3-us": {
|
||||
"cache_read_input_token_cost": 3.9e-07,
|
||||
"input_cost_per_token": 2.1e-06,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": {
|
||||
"cache_read_input_token_cost": 3.9e-07,
|
||||
"input_cost_per_token": 2.1e-06,
|
||||
|
|
@ -63358,6 +63436,20 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/routers/glm-5p3-flash-us": {
|
||||
"cache_read_input_token_cost": 4.5e-08,
|
||||
"input_cost_per_token": 2.25e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-07,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/glm-5p3-flash": {
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
"cache_read_input_token_cost_priority": 3.75e-08,
|
||||
|
|
@ -63375,6 +63467,20 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/glm-5p3-flash-us": {
|
||||
"cache_read_input_token_cost": 4.5e-08,
|
||||
"input_cost_per_token": 2.25e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-07,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/models/inkling": {
|
||||
"cache_read_input_token_cost": 1.7e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
|
|
@ -65906,9 +66012,9 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/z-ai/glm-5.3-flash": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"output_cost_per_token": 5e-07,
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"input_cost_per_token": 4.5e-08,
|
||||
"output_cost_per_token": 6e-07,
|
||||
"cache_read_input_token_cost": 2.85e-08,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1310720,
|
||||
"max_output_tokens": 943718,
|
||||
|
|
@ -66657,9 +66763,9 @@
|
|||
"supports_web_search": true
|
||||
},
|
||||
"openrouter/deepseek/deepseek-v4-flash": {
|
||||
"input_cost_per_token": 8.4e-08,
|
||||
"output_cost_per_token": 1.68e-07,
|
||||
"cache_read_input_token_cost": 1.68e-08,
|
||||
"input_cost_per_token": 4.9e-08,
|
||||
"output_cost_per_token": 9.8e-08,
|
||||
"cache_read_input_token_cost": 9.8e-09,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 384000,
|
||||
|
|
@ -73050,6 +73156,26 @@
|
|||
"supports_vision": true,
|
||||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/mistralai/mistral-large-2512": {
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"input_cost_per_token": 5e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 209715,
|
||||
"max_tokens": 209715,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-06,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/mistralai/mistral-large-2512:batch": {
|
||||
"cache_read_input_token_cost": 2.5e-08,
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
|
|
|
|||
|
|
@ -239,6 +239,14 @@ class Provider:
|
|||
response: Final = self.scenario_store.get(scenario_id)
|
||||
if response is None:
|
||||
return JSONResponse({"error": "Unknown scenario"}, status_code=404)
|
||||
if request.method == "POST" and "json" in request.headers.get("content-type", ""):
|
||||
raw_body: Final = await request.body()
|
||||
if raw_body:
|
||||
body: Final = JSON_OBJECT.validate_json(raw_body)
|
||||
if isinstance(body, dict):
|
||||
self.observations.put(
|
||||
Observation(request.url.path, request.headers.get("authorization", ""), body)
|
||||
)
|
||||
if isinstance(response, RoutedResponse):
|
||||
route_key: Final = f"{request.method} /{'/'.join(segments[1:])}"
|
||||
route: Final = next(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,240 @@
|
|||
"""Encrypted-content affinity when the origin deployment has no encryption-boundary peer.
|
||||
|
||||
A multi-region model group has several deployments sharing one upstream api_key but a
|
||||
distinct api_base each, with ``optional_pre_call_checks: [encrypted_content_affinity]``
|
||||
and ``disable_cooldowns: true`` (the integration proxy config). A follow-up
|
||||
``POST /v1/responses`` that replays a reasoning item must never fail at the proxy while
|
||||
sibling deployments in the same group are healthy: when the origin cannot serve the turn
|
||||
its encrypted reasoning should be stripped and the request dispatched to a sibling.
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.upstream import delete_scenario, register_scenario
|
||||
from integration.cost_calculation.cost_tracking_case import JsonResponse
|
||||
from pydantic import JsonValue
|
||||
|
||||
AFFINITY_CHECK: Final = "encrypted_content_affinity"
|
||||
PROVIDER_MODEL: Final = "openai/gpt-5"
|
||||
PROVIDER_KEY: Final = "integration-provider-key"
|
||||
DEPLOYMENT_COUNT: Final = 3
|
||||
|
||||
|
||||
def _responses_payload() -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": "resp_$REQUEST_ID",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-5-scripted",
|
||||
"output": [
|
||||
{
|
||||
"type": "reasoning",
|
||||
"id": "rs_$REQUEST_ID",
|
||||
"summary": [],
|
||||
"encrypted_content": "ZHNra2RrZA==",
|
||||
},
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_$REQUEST_ID",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "scripted answer"}],
|
||||
},
|
||||
],
|
||||
"usage": {"input_tokens": 5, "output_tokens": 7, "total_tokens": 12},
|
||||
}
|
||||
|
||||
|
||||
def _enable_affinity_check(scenario: Scenario) -> None:
|
||||
gateway: Final = scenario.gateway
|
||||
settings: Final = object_value(gateway.get("/router/settings")["current_values"])
|
||||
current: Final = settings.get("optional_pre_call_checks")
|
||||
original: Final = list(current) if isinstance(current, list) else []
|
||||
original_retries: Final = settings.get("num_retries")
|
||||
scenario.cleanups.callback(
|
||||
lambda: gateway.post(
|
||||
"/config/update",
|
||||
{"router_settings": {"optional_pre_call_checks": original, "num_retries": original_retries}},
|
||||
)
|
||||
)
|
||||
gateway.post(
|
||||
"/config/update",
|
||||
{"router_settings": {"optional_pre_call_checks": [*original, AFFINITY_CHECK], "num_retries": 0}},
|
||||
)
|
||||
|
||||
|
||||
def _last_group_request_body(gateway: Gateway, group: str) -> dict[str, JsonValue]:
|
||||
"""The request body the upstream last saw for this group, to prove the strip reached the wire."""
|
||||
with httpx.Client(base_url=gateway.upstream_url, trust_env=False) as upstream:
|
||||
requests: Final = object_value(upstream.get("/__observations").json()).get("requests")
|
||||
assert isinstance(requests, list)
|
||||
bodies: Final = [
|
||||
object_value(request).get("body")
|
||||
for request in requests
|
||||
if str(object_value(request).get("path")).startswith(f"/{group}-")
|
||||
]
|
||||
assert bodies, f"upstream saw no requests for {group}: {requests}"
|
||||
body: Final = bodies[-1]
|
||||
assert isinstance(body, dict), f"upstream request body is not an object: {body}"
|
||||
return body
|
||||
|
||||
|
||||
def _assert_no_encrypted_reasoning_reached_upstream(gateway: Gateway, group: str) -> None:
|
||||
body: Final = _last_group_request_body(gateway, group)
|
||||
items: Final = body.get("input")
|
||||
assert isinstance(items, list), f"upstream request carried no input list: {body}"
|
||||
assert not any(
|
||||
isinstance(item, dict) and ("encrypted_content" in item or "litellm_enc" in json.dumps(item)) for item in items
|
||||
), f"undecryptable reasoning reached the sibling: {body}"
|
||||
|
||||
|
||||
def _multi_region_group(scenario: Scenario) -> tuple[str, tuple[str, ...]]:
|
||||
"""Three deployments in one model group: one shared api_key, a distinct api_base each."""
|
||||
gateway: Final = scenario.gateway
|
||||
group: Final = f"enc-affinity-{uuid.uuid4().hex}"
|
||||
handles: Final = tuple(
|
||||
register_scenario(
|
||||
f"{group}-{index}",
|
||||
JsonResponse(content_type="application/json", body=_responses_payload()),
|
||||
)
|
||||
for index in range(DEPLOYMENT_COUNT)
|
||||
)
|
||||
for handle in handles:
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
|
||||
def delete_model_if_present(model_id: str) -> None:
|
||||
entries: Final = gateway.get("/model/info")["data"]
|
||||
assert isinstance(entries, list)
|
||||
if any(object_value(object_value(entry)["model_info"])["id"] == model_id for entry in entries):
|
||||
scenario.delete_model(model_id)
|
||||
|
||||
deployment_ids: Final = tuple(
|
||||
string_value(
|
||||
object_value(
|
||||
gateway.post(
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": group,
|
||||
"litellm_params": {
|
||||
"model": PROVIDER_MODEL,
|
||||
"api_key": PROVIDER_KEY,
|
||||
"api_base": handle.api_base(),
|
||||
},
|
||||
"model_info": {},
|
||||
},
|
||||
)["model_info"]
|
||||
)["id"]
|
||||
)
|
||||
for handle in handles
|
||||
)
|
||||
for model_id in deployment_ids:
|
||||
scenario.cleanups.callback(delete_model_if_present, model_id)
|
||||
return group, deployment_ids
|
||||
|
||||
|
||||
def _user_message(text: str) -> dict[str, JsonValue]:
|
||||
return {"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}
|
||||
|
||||
|
||||
def _responses_turn(gateway: Gateway, group: str, request_input: JsonValue) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": group, "input": request_input, "store": False, "include": ["reasoning.encrypted_content"]},
|
||||
)
|
||||
|
||||
|
||||
def _turn_one(gateway: Gateway, group: str, deployment_ids: tuple[str, ...]) -> tuple[str, list[JsonValue]]:
|
||||
# the proxy caches responses, so the turn-one prompt needs a unique marker or a
|
||||
# stale body carrying another run's encoded model_id would replay instead
|
||||
response: Final = _responses_turn(gateway, group, f"hello affinity {uuid.uuid4().hex}")
|
||||
assert response.status_code == 200, response.text
|
||||
origin: Final = str(response.headers["x-litellm-model-id"])
|
||||
assert origin in deployment_ids, f"turn one served by unknown deployment {origin}: {response.text}"
|
||||
output: Final = object_value(response.json()).get("output")
|
||||
assert isinstance(output, list), f"turn one returned no output items: {response.text}"
|
||||
reasoning: Final = next((item for item in output if object_value(item).get("type") == "reasoning"), None)
|
||||
assert isinstance(reasoning, dict), f"turn one returned no reasoning item: {response.text}"
|
||||
assert str(object_value(reasoning)["id"]).startswith("encitem_"), (
|
||||
f"affinity encoding did not run on turn one: {reasoning}"
|
||||
)
|
||||
assert isinstance(object_value(reasoning).get("encrypted_content"), str), (
|
||||
f"turn one reasoning item has no encrypted_content: {reasoning}"
|
||||
)
|
||||
message: Final = next((item for item in output if object_value(item).get("type") == "message"), None)
|
||||
assert isinstance(message, dict), f"turn one returned no message item: {response.text}"
|
||||
return origin, [reasoning, message]
|
||||
|
||||
|
||||
def _replay(gateway: Gateway, group: str, items: list[JsonValue]) -> httpx.Response:
|
||||
return _responses_turn(
|
||||
gateway,
|
||||
group,
|
||||
[_user_message("hello affinity"), *items, _user_message("continue the conversation")],
|
||||
)
|
||||
|
||||
|
||||
def _model_blocked(gateway: Gateway, model_id: str) -> bool:
|
||||
entries: Final = gateway.get("/model/info")["data"]
|
||||
assert isinstance(entries, list)
|
||||
entry: Final = next(
|
||||
(entry for entry in entries if object_value(object_value(entry)["model_info"])["id"] == model_id),
|
||||
None,
|
||||
)
|
||||
return entry is not None and object_value(object_value(entry)["model_info"]).get("blocked") is True
|
||||
|
||||
|
||||
def test_replayed_encrypted_content_serves_from_sibling_when_origin_blocked(gateway: Gateway) -> None:
|
||||
"""Origin excluded from healthy deployments (admin-blocked, no cooldown) must not 503.
|
||||
|
||||
On unfixed code the origin is a routed-group candidate with no (api_base, api_key)
|
||||
peer, so the affinity check raises a proxy-level 503 instead of stripping the
|
||||
encrypted reasoning and dispatching to a healthy sibling.
|
||||
"""
|
||||
with gateway.scenario() as scenario:
|
||||
_enable_affinity_check(scenario)
|
||||
group, deployment_ids = _multi_region_group(scenario)
|
||||
origin, items = _turn_one(gateway, group, deployment_ids)
|
||||
|
||||
gateway.post("/model/block", {"model_id": origin})
|
||||
eventually(lambda: _model_blocked(gateway, origin), lambda blocked: blocked)
|
||||
|
||||
response: Final = _replay(gateway, group, items)
|
||||
assert response.status_code == 200, response.text
|
||||
siblings: Final = tuple(model_id for model_id in deployment_ids if model_id != origin)
|
||||
assert response.headers.get("x-litellm-model-id") in siblings, (
|
||||
f"turn two served by {response.headers.get('x-litellm-model-id')}, "
|
||||
f"expected a sibling of blocked origin {origin}: {response.text}"
|
||||
)
|
||||
_assert_no_encrypted_reasoning_reached_upstream(gateway, group)
|
||||
|
||||
|
||||
def test_replayed_encrypted_content_serves_from_sibling_when_origin_deleted(gateway: Gateway) -> None:
|
||||
"""Origin permanently removed: strip-and-dispatch to a sibling, the shipped behavior."""
|
||||
with gateway.scenario() as scenario:
|
||||
_enable_affinity_check(scenario)
|
||||
group, deployment_ids = _multi_region_group(scenario)
|
||||
origin, items = _turn_one(gateway, group, deployment_ids)
|
||||
|
||||
gateway.post("/model/delete", {"id": origin})
|
||||
eventually(
|
||||
lambda: gateway.get("/model/info")["data"],
|
||||
lambda entries: (
|
||||
isinstance(entries, list)
|
||||
and all(object_value(object_value(entry)["model_info"])["id"] != origin for entry in entries)
|
||||
),
|
||||
)
|
||||
|
||||
response: Final = _replay(gateway, group, items)
|
||||
assert response.status_code == 200, response.text
|
||||
siblings: Final = tuple(model_id for model_id in deployment_ids if model_id != origin)
|
||||
assert response.headers.get("x-litellm-model-id") in siblings, (
|
||||
f"turn two served by {response.headers.get('x-litellm-model-id')}, "
|
||||
f"expected a sibling of deleted origin {origin}: {response.text}"
|
||||
)
|
||||
_assert_no_encrypted_reasoning_reached_upstream(gateway, group)
|
||||
|
|
@ -30,6 +30,7 @@ from pydantic import ValidationError
|
|||
|
||||
import litellm
|
||||
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck
|
||||
from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
|
||||
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck
|
||||
from litellm.types.router import RetryPolicy, UpdateRouterConfig
|
||||
|
|
@ -203,6 +204,63 @@ def test_update_settings_replaces_toggleable_pre_call_checks():
|
|||
assert any(isinstance(callback, ModelRateLimitingCheck) for callback in (router.optional_callbacks or []))
|
||||
|
||||
|
||||
def test_update_settings_clears_omitted_encrypted_content_affinity_check():
|
||||
router = _build_router()
|
||||
|
||||
router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"])
|
||||
router.update_settings(optional_pre_call_checks=[])
|
||||
|
||||
assert not any(
|
||||
isinstance(callback, EncryptedContentAffinityCheck) for callback in (router.optional_callbacks or [])
|
||||
)
|
||||
assert not any(isinstance(callback, EncryptedContentAffinityCheck) for callback in litellm.callbacks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_settings_turning_off_encrypted_content_affinity_stops_flagging_requests():
|
||||
router = _build_router()
|
||||
|
||||
router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"])
|
||||
enabled_request: Final = {"litellm_metadata": {}, "input": "hello"}
|
||||
await router.async_get_available_deployment(model="test-model", request_kwargs=enabled_request)
|
||||
assert enabled_request["litellm_metadata"]["encrypted_content_affinity_enabled"] is True
|
||||
|
||||
router.update_settings(optional_pre_call_checks=[])
|
||||
disabled_request: Final = {"litellm_metadata": {}, "input": "hello"}
|
||||
await router.async_get_available_deployment(model="test-model", request_kwargs=disabled_request)
|
||||
assert "encrypted_content_affinity_enabled" not in disabled_request["litellm_metadata"]
|
||||
|
||||
router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"])
|
||||
reenabled_request: Final = {"litellm_metadata": {}, "input": "hello"}
|
||||
await router.async_get_available_deployment(model="test-model", request_kwargs=reenabled_request)
|
||||
assert reenabled_request["litellm_metadata"]["encrypted_content_affinity_enabled"] is True
|
||||
|
||||
|
||||
def test_update_settings_keeps_per_group_encrypted_content_affinity_when_global_toggle_is_omitted():
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4",
|
||||
"api_key": "sk-fake",
|
||||
"api_base": "http://localhost:9999",
|
||||
},
|
||||
}
|
||||
],
|
||||
model_group_affinity_config={"test-model": ["encrypted_content_affinity"]},
|
||||
)
|
||||
|
||||
router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"])
|
||||
router.update_settings(optional_pre_call_checks=[])
|
||||
|
||||
affinity_checks: Final = [
|
||||
callback for callback in (router.optional_callbacks or []) if isinstance(callback, EncryptedContentAffinityCheck)
|
||||
]
|
||||
assert len(affinity_checks) == 1
|
||||
assert affinity_checks[0].enable_global_affinity is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_settings_preserves_router_budget_limiting_when_omitted(monkeypatch):
|
||||
async def _disable_periodic_sync(*args, **kwargs):
|
||||
|
|
|
|||
|
|
@ -1285,7 +1285,7 @@ def test_boundary_key_rejects_non_dict_like_inputs():
|
|||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fail-fast when originating deployment is unavailable and no boundary peer
|
||||
# Degraded dispatch when the originating deployment is unavailable and no boundary peer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
|
|
@ -1323,14 +1323,13 @@ def _make_router_mock_with_cooldown(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_affinity_raises_service_unavailable_when_origin_cooled_for_non_429():
|
||||
async def test_affinity_strips_and_dispatches_when_origin_cooled_for_non_429():
|
||||
"""
|
||||
Originating deployment is in the router config, in cooldown for a non-429
|
||||
cause (e.g. a 500), and no boundary peer is configured. The check must
|
||||
surface this as a 503 (transient, but not rate-limit-specific) rather than
|
||||
dispatching to a non-peer deployment.
|
||||
degrade: strip the encrypted reasoning and dispatch to the healthy pool
|
||||
rather than failing the request.
|
||||
"""
|
||||
from litellm.exceptions import ServiceUnavailableError
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
|
@ -1366,32 +1365,36 @@ async def test_affinity_raises_service_unavailable_when_origin_cooled_for_non_42
|
|||
}
|
||||
]
|
||||
request_kwargs = {
|
||||
"input": [{"id": encoded_id, "type": "reasoning"}],
|
||||
"input": [
|
||||
{
|
||||
"id": encoded_id,
|
||||
"type": "reasoning",
|
||||
"encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(
|
||||
"gAAAAA-blob", "deployment-a-cooled"
|
||||
),
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
with pytest.raises(ServiceUnavailableError) as excinfo:
|
||||
await check.async_filter_deployments(
|
||||
model="gpt-5.4",
|
||||
healthy_deployments=healthy_only_b,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
result = await check.async_filter_deployments(
|
||||
model="gpt-5.4",
|
||||
healthy_deployments=healthy_only_b,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
# Public error message intentionally omits the originating model_id to
|
||||
# avoid an authenticated-caller probing oracle.
|
||||
assert "deployment-a-cooled" not in str(excinfo.value)
|
||||
assert excinfo.value.status_code == 503
|
||||
assert result is healthy_only_b
|
||||
assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in request_kwargs["input"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_affinity_raises_rate_limit_with_retry_after_when_origin_cooled_for_429():
|
||||
async def test_affinity_strips_and_dispatches_when_origin_cooled_for_429():
|
||||
"""
|
||||
Originating deployment is in cooldown specifically because of a 429.
|
||||
The check must surface this as a 429 RateLimitError with a Retry-After
|
||||
header derived from the cooldown's remaining window, so OpenAI-compatible
|
||||
clients respect the backoff instead of giving up on a 503.
|
||||
Originating deployment is in cooldown specifically because of a 429 and no
|
||||
boundary peer is configured. The check degrades the same way: strip the
|
||||
encrypted reasoning and dispatch to the healthy pool rather than surfacing
|
||||
a rate-limit error to the caller.
|
||||
"""
|
||||
from litellm.exceptions import RateLimitError
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
|
@ -1431,29 +1434,23 @@ async def test_affinity_raises_rate_limit_with_retry_after_when_origin_cooled_fo
|
|||
"input": [{"id": encoded_id, "type": "reasoning"}],
|
||||
}
|
||||
|
||||
with pytest.raises(RateLimitError) as excinfo:
|
||||
await check.async_filter_deployments(
|
||||
model="gpt-5.4",
|
||||
healthy_deployments=healthy_only_b,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
result = await check.async_filter_deployments(
|
||||
model="gpt-5.4",
|
||||
healthy_deployments=healthy_only_b,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
assert "deployment-a-cooled-429" not in str(excinfo.value)
|
||||
assert excinfo.value.status_code == 429
|
||||
retry_after = excinfo.value.response.headers.get("retry-after")
|
||||
assert retry_after is not None
|
||||
assert 1 <= int(retry_after) <= 60
|
||||
assert result is healthy_only_b
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_affinity_raises_service_unavailable_when_origin_filtered_without_cooldown_entry():
|
||||
async def test_affinity_strips_and_dispatches_when_origin_filtered_without_cooldown_entry():
|
||||
"""
|
||||
Originating deployment is configured but absent from healthy_deployments
|
||||
with no active cooldown entry. Surface as 503 (we cannot prove the cause
|
||||
was rate-limiting) rather than guessing 429.
|
||||
with no active cooldown entry and no boundary peer. The check degrades:
|
||||
strip the encrypted reasoning and dispatch to the healthy pool.
|
||||
"""
|
||||
from litellm.exceptions import ServiceUnavailableError
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
|
@ -1480,15 +1477,72 @@ async def test_affinity_raises_service_unavailable_when_origin_filtered_without_
|
|||
"input": [{"id": encoded_id, "type": "reasoning"}],
|
||||
}
|
||||
|
||||
with pytest.raises(ServiceUnavailableError) as excinfo:
|
||||
await check.async_filter_deployments(
|
||||
model="gpt-5.4",
|
||||
healthy_deployments=healthy_only_b,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
result = await check.async_filter_deployments(
|
||||
model="gpt-5.4",
|
||||
healthy_deployments=healthy_only_b,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
assert excinfo.value.status_code == 503
|
||||
assert result is healthy_only_b
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_affinity_serves_sibling_when_candidate_origin_has_no_boundary_peer():
|
||||
"""
|
||||
Regression for the multi-region group where every deployment has a distinct
|
||||
api_base: the origin is still a routed-group candidate but absent from
|
||||
healthy_deployments, and no (api_base, api_key) peer exists. The turn must
|
||||
degrade on a sibling with its encrypted reasoning stripped, not fail.
|
||||
"""
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
originating = _make_originating_mock("https://region-a.example.com/v1", "shared-key")
|
||||
mock_router = _make_router_mock_with_cooldown(
|
||||
originating, cooldown_entries=[], routed_group_model_ids=["region-a", "region-b", "region-c"]
|
||||
)
|
||||
check = EncryptedContentAffinityCheck(router=mock_router)
|
||||
wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "region-a")
|
||||
siblings = [
|
||||
{
|
||||
"model_info": {"id": "region-b"},
|
||||
"model_name": "gpt-5.4",
|
||||
"litellm_params": {"api_base": "https://region-b.example.com/v1", "api_key": "shared-key"},
|
||||
},
|
||||
{
|
||||
"model_info": {"id": "region-c"},
|
||||
"model_name": "gpt-5.4",
|
||||
"litellm_params": {"api_base": "https://region-c.example.com/v1", "api_key": "shared-key"},
|
||||
},
|
||||
]
|
||||
request_kwargs = {
|
||||
"litellm_metadata": {},
|
||||
"input": [
|
||||
{"role": "user", "content": "why is the sky blue?"},
|
||||
{
|
||||
"type": "reasoning",
|
||||
"id": ResponsesAPIRequestUtils._build_encrypted_item_id("region-a", "rs_test"),
|
||||
"encrypted_content": wrapped,
|
||||
"summary": [{"type": "summary_text", "text": "scattering"}],
|
||||
},
|
||||
{"role": "user", "content": "and sunsets?"},
|
||||
],
|
||||
}
|
||||
|
||||
result = await check.async_filter_deployments(
|
||||
model="gpt-5.4",
|
||||
healthy_deployments=siblings,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
assert result is siblings
|
||||
assert request_kwargs["input"][1] == {
|
||||
"type": "reasoning",
|
||||
"summary": [{"type": "summary_text", "text": "scattering"}],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1981,15 +2035,14 @@ async def test_affinity_strips_encrypted_reasoning_when_routed_to_another_model_
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_affinity_fails_fast_within_the_origins_own_group():
|
||||
async def test_affinity_degrades_within_the_origins_own_group():
|
||||
"""
|
||||
Negative class for the tier-change discriminator: the routed group IS the
|
||||
origin's group (a same-group cooldown, not a tier change), so even with a
|
||||
healthy non-origin sibling that cannot decrypt the content, the request
|
||||
still fails fast and the encrypted reasoning is left intact rather than
|
||||
stripped. Preserves the LIT-3051 cooldown contract.
|
||||
The routed group IS the origin's group (a same-group cooldown), and the
|
||||
healthy sibling sits on a different encryption boundary, so it cannot
|
||||
decrypt the replayed reasoning. The check still degrades instead of
|
||||
failing: the encrypted reasoning is stripped and the request dispatches
|
||||
to the sibling.
|
||||
"""
|
||||
from litellm.exceptions import ServiceUnavailableError
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
|
@ -2014,29 +2067,26 @@ async def test_affinity_fails_fast_within_the_origins_own_group():
|
|||
]
|
||||
request_kwargs = _cross_group_request_kwargs()
|
||||
|
||||
with pytest.raises(ServiceUnavailableError):
|
||||
await check.async_filter_deployments(
|
||||
model="gpt-reasoning-tier",
|
||||
healthy_deployments=sibling_pool,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
result = await check.async_filter_deployments(
|
||||
model="gpt-reasoning-tier",
|
||||
healthy_deployments=sibling_pool,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
assert request_kwargs["input"][1].get("encrypted_content")
|
||||
assert result is sibling_pool
|
||||
assert not request_kwargs["input"][1].get("encrypted_content")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_affinity_does_not_strip_when_group_is_spelled_differently_but_same_by_id():
|
||||
async def test_affinity_strips_when_group_is_spelled_differently_but_same_by_id():
|
||||
"""
|
||||
The discriminator must key on deployment-id membership, not on the model-group
|
||||
name string. Here the origin's configured group is spelled ``openai/gpt-5.4-mini``
|
||||
while the routed group is the canonical ``gpt-5.4-mini``: same group, different
|
||||
spelling. A name compare (``originating.model_name != model``) would read this as
|
||||
a tier change and strip the reasoning it did not have to. Because the origin's id
|
||||
is a member of the routed group, this is a same-group cooldown instead: the request
|
||||
fails fast and the encrypted reasoning is left intact.
|
||||
The origin's configured group is spelled ``openai/gpt-5.4-mini`` while the
|
||||
routed group is the canonical ``gpt-5.4-mini``: same group, different
|
||||
spelling, with the origin a member but currently unavailable. There is no
|
||||
boundary peer, so the check degrades: strip the encrypted reasoning and
|
||||
dispatch to the sibling pool.
|
||||
"""
|
||||
from litellm.exceptions import ServiceUnavailableError
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
|
@ -2071,28 +2121,24 @@ async def test_affinity_does_not_strip_when_group_is_spelled_differently_but_sam
|
|||
],
|
||||
}
|
||||
|
||||
with pytest.raises(ServiceUnavailableError):
|
||||
await check.async_filter_deployments(
|
||||
model="gpt-5.4-mini",
|
||||
healthy_deployments=sibling_pool,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
result = await check.async_filter_deployments(
|
||||
model="gpt-5.4-mini",
|
||||
healthy_deployments=sibling_pool,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
assert request_kwargs["input"][1].get("encrypted_content")
|
||||
assert result is sibling_pool
|
||||
assert not request_kwargs["input"][1].get("encrypted_content")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_affinity_honors_router_candidate_ids_for_team_and_pattern_routes():
|
||||
async def test_affinity_strips_for_team_and_pattern_routes():
|
||||
"""
|
||||
The exact `model_name` index does not include team-public or pattern routes, so a
|
||||
same-group cooldown reached only through one of those would be misread as a tier change
|
||||
and stripped. The check asks the router for the candidate ids it resolves for the route
|
||||
(`get_candidate_model_ids_for_route`), which covers those paths, rather than the bare
|
||||
index. Here that set marks the origin as a candidate, so the request fails fast with its
|
||||
reasoning intact, and the routed group and team are passed through to the router.
|
||||
Team-public or pattern routes resolve the same routed group as the origin's,
|
||||
so an unavailable origin there also degrades rather than failing: the
|
||||
encrypted reasoning is stripped and the request dispatches to the sibling.
|
||||
"""
|
||||
from litellm.exceptions import ServiceUnavailableError
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
|
@ -2127,13 +2173,12 @@ async def test_affinity_honors_router_candidate_ids_for_team_and_pattern_routes(
|
|||
],
|
||||
}
|
||||
|
||||
with pytest.raises(ServiceUnavailableError):
|
||||
await check.async_filter_deployments(
|
||||
model="team-public-model",
|
||||
healthy_deployments=sibling_pool,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
result = await check.async_filter_deployments(
|
||||
model="team-public-model",
|
||||
healthy_deployments=sibling_pool,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
assert request_kwargs["input"][1].get("encrypted_content")
|
||||
mock_router.get_candidate_model_ids_for_route.assert_called_once_with(model="team-public-model", team_id="teamA")
|
||||
assert result is sibling_pool
|
||||
assert not request_kwargs["input"][1].get("encrypted_content")
|
||||
|
|
|
|||
|
|
@ -0,0 +1,103 @@
|
|||
import { render, screen } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import ResetMemberBudgetsDialog from "./ResetMemberBudgetsDialog";
|
||||
import type { MemberBudgetResetState } from "./useMemberBudgetReset";
|
||||
|
||||
const pendingFor = (memberCount: number, newBudget = 10) => ({
|
||||
teamId: "team-123",
|
||||
updateData: {},
|
||||
userIds: Array.from({ length: memberCount }, (_, i) => `u-${i}`),
|
||||
newBudget,
|
||||
});
|
||||
|
||||
const promptingWith = (memberCount: number, newBudget = 10): MemberBudgetResetState => ({
|
||||
phase: "prompting",
|
||||
pending: pendingFor(memberCount, newBudget),
|
||||
});
|
||||
|
||||
const defaultHandlers = () => ({
|
||||
onReset: vi.fn(),
|
||||
onRetry: vi.fn(),
|
||||
onKeep: vi.fn(),
|
||||
onDismiss: vi.fn(),
|
||||
});
|
||||
|
||||
describe("ResetMemberBudgetsDialog", () => {
|
||||
it("names the members and the new default so the prompt is self-explanatory", () => {
|
||||
render(<ResetMemberBudgetsDialog state={promptingWith(3)} {...defaultHandlers()} />);
|
||||
|
||||
expect(
|
||||
screen.getByText(
|
||||
"3 members have a custom budget, so the new team default of $10.00 will not apply to them. " +
|
||||
"Reset them to the default, or keep the custom budgets?",
|
||||
),
|
||||
).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("switches to singular for one member", () => {
|
||||
render(<ResetMemberBudgetsDialog state={promptingWith(1)} {...defaultHandlers()} />);
|
||||
|
||||
expect(
|
||||
screen.getByText(
|
||||
"1 member has a custom budget, so the new team default of $10.00 will not apply to that member. " +
|
||||
"Reset it to the default, or keep the custom budget?",
|
||||
),
|
||||
).toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Reset to $10.00" })).toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Keep custom budget" })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("formats large budgets like the rest of the team page", () => {
|
||||
render(<ResetMemberBudgetsDialog state={promptingWith(2, 1000)} {...defaultHandlers()} />);
|
||||
|
||||
expect(screen.getByRole("button", { name: "Reset all to $1,000.00" })).toBeInTheDocument();
|
||||
expect(screen.getByText(/new team default of \$1,000\.00/)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("routes each choice to the matching action", async () => {
|
||||
const handlers = defaultHandlers();
|
||||
const user = userEvent.setup();
|
||||
render(<ResetMemberBudgetsDialog state={promptingWith(2)} {...handlers} />);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Reset all to $10.00" }));
|
||||
expect(handlers.onReset).toHaveBeenCalledTimes(1);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Keep custom budgets" }));
|
||||
expect(handlers.onKeep).toHaveBeenCalledTimes(1);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Cancel" }));
|
||||
expect(handlers.onDismiss).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("locks every button while the reset is in flight", () => {
|
||||
const state: MemberBudgetResetState = { phase: "resetting", pending: pendingFor(2), attempted: 0 };
|
||||
render(<ResetMemberBudgetsDialog state={state} {...defaultHandlers()} />);
|
||||
|
||||
expect(screen.getByRole("button", { name: "Reset all to $10.00" })).toBeDisabled();
|
||||
expect(screen.getByRole("button", { name: "Keep custom budgets" })).toBeDisabled();
|
||||
expect(screen.getByRole("button", { name: "Cancel" })).toBeDisabled();
|
||||
});
|
||||
|
||||
it("offers retry or cancel after a failed reset", async () => {
|
||||
const handlers = defaultHandlers();
|
||||
const user = userEvent.setup();
|
||||
const state: MemberBudgetResetState = { phase: "resetFailed", pending: pendingFor(2), attempted: 0 };
|
||||
render(<ResetMemberBudgetsDialog state={state} {...handlers} />);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Retry reset" }));
|
||||
expect(handlers.onRetry).toHaveBeenCalledTimes(1);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Cancel" }));
|
||||
expect(handlers.onDismiss).toHaveBeenCalledTimes(1);
|
||||
|
||||
expect(screen.queryByRole("button", { name: "Keep custom budgets" })).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("button", { name: "Reset all to $10.00" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("renders nothing while idle", () => {
|
||||
render(<ResetMemberBudgetsDialog state={{ phase: "idle" }} {...defaultHandlers()} />);
|
||||
|
||||
expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,71 @@
|
|||
import React from "react";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
|
||||
import { formatNumberWithCommas } from "@/utils/dataUtils";
|
||||
import { pluralize } from "./memberBudgetReset";
|
||||
import type { MemberBudgetResetState } from "./useMemberBudgetReset";
|
||||
|
||||
interface ResetMemberBudgetsDialogProps {
|
||||
state: MemberBudgetResetState;
|
||||
onReset: () => void;
|
||||
onRetry: () => void;
|
||||
onKeep: () => void;
|
||||
onDismiss: () => void;
|
||||
}
|
||||
|
||||
export default function ResetMemberBudgetsDialog({
|
||||
state,
|
||||
onReset,
|
||||
onRetry,
|
||||
onKeep,
|
||||
onDismiss,
|
||||
}: ResetMemberBudgetsDialogProps) {
|
||||
const open = state.phase !== "idle";
|
||||
const busy = state.phase === "resetting";
|
||||
const failed = state.phase === "resetFailed";
|
||||
const memberCount = state.phase === "idle" ? 0 : state.pending.userIds.length;
|
||||
const newBudget = state.phase === "idle" ? 0 : state.pending.newBudget;
|
||||
|
||||
return (
|
||||
<Dialog
|
||||
open={open}
|
||||
onOpenChange={(nextOpen) => {
|
||||
if (!nextOpen && !busy) onDismiss();
|
||||
}}
|
||||
>
|
||||
<DialogContent>
|
||||
<DialogHeader>
|
||||
<DialogTitle>Reset member budgets?</DialogTitle>
|
||||
</DialogHeader>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
{memberCount} {pluralize(memberCount, "member has", "members have")} a custom budget, so the new team default
|
||||
of ${formatNumberWithCommas(newBudget, 2)} will not apply to {pluralize(memberCount, "that member", "them")}.
|
||||
Reset {pluralize(memberCount, "it", "them")} to the default, or keep the custom{" "}
|
||||
{pluralize(memberCount, "budget", "budgets")}?
|
||||
</p>
|
||||
<DialogFooter>
|
||||
{failed ? (
|
||||
<>
|
||||
<Button variant="outline" onClick={onDismiss}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button onClick={onRetry}>Retry reset</Button>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<Button variant="ghost" onClick={onDismiss} disabled={busy}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button variant="outline" onClick={onKeep} disabled={busy}>
|
||||
Keep custom {pluralize(memberCount, "budget", "budgets")}
|
||||
</Button>
|
||||
<Button onClick={onReset} disabled={busy}>
|
||||
{pluralize(memberCount, "Reset", "Reset all")} to ${formatNumberWithCommas(newBudget, 2)}
|
||||
</Button>
|
||||
</>
|
||||
)}
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
|
|
@ -0,0 +1,722 @@
|
|||
import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema";
|
||||
import * as networking from "@/components/networking";
|
||||
import { fireEvent, screen, waitFor } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { renderWithProviders } from "../../../tests/test-utils";
|
||||
import { toast } from "@/lib/toast";
|
||||
import type { Team } from "../key_team_helpers/key_list";
|
||||
import { MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES } from "./memberBudgetReset";
|
||||
import TeamInfoView, { type TeamData } from "./TeamInfo";
|
||||
|
||||
const authState = vi.hoisted(() => ({ userRole: "Admin" }));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
|
||||
default: () => ({
|
||||
token: "123",
|
||||
accessToken: "123",
|
||||
userId: "user-1",
|
||||
userEmail: "user@example.com",
|
||||
userRole: authState.userRole,
|
||||
premiumUser: false,
|
||||
disabledPersonalKeyCreation: null,
|
||||
showSSOBanner: false,
|
||||
}),
|
||||
}));
|
||||
|
||||
vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) }));
|
||||
|
||||
vi.mock("@/components/networking", () => ({
|
||||
serverRootPath: "",
|
||||
teamInfoCall: vi.fn(),
|
||||
teamMemberDeleteCall: vi.fn(),
|
||||
teamMemberAddCall: vi.fn(),
|
||||
teamMemberUpdateCall: vi.fn(),
|
||||
teamUpdateCall: vi.fn(),
|
||||
getGuardrailsList: vi.fn(),
|
||||
getPoliciesList: vi.fn(),
|
||||
getPolicyInfoWithGuardrails: vi.fn(),
|
||||
fetchMCPAccessGroups: vi.fn(),
|
||||
getTeamPermissionsCall: vi.fn(),
|
||||
organizationInfoCall: vi.fn(),
|
||||
getRouterSettingsCall: vi.fn().mockResolvedValue({ fields: [] }),
|
||||
getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: [] }),
|
||||
fetchMCPServers: vi.fn().mockResolvedValue([]),
|
||||
fetchMCPToolsets: vi.fn().mockResolvedValue([]),
|
||||
listMCPTools: vi.fn().mockResolvedValue({ tools: [] }),
|
||||
vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }),
|
||||
getAgentsList: vi.fn().mockResolvedValue({ agents: [] }),
|
||||
getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [], count: 0 }),
|
||||
}));
|
||||
|
||||
const { bulkUpdatePOST } = vi.hoisted(() => ({ bulkUpdatePOST: vi.fn() }));
|
||||
vi.mock("@/lib/http/api", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("@/lib/http/api")>();
|
||||
return { ...actual, fetchClient: { ...actual.fetchClient, POST: bulkUpdatePOST } };
|
||||
});
|
||||
|
||||
const can = vi.fn();
|
||||
vi.mock("@/app/(dashboard)/hooks/useCan", () => ({
|
||||
default: (...args: unknown[]) => can(...args),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/utils/dataUtils", () => ({
|
||||
copyToClipboard: vi.fn().mockResolvedValue(true),
|
||||
formatNumberWithCommas: vi.fn((value: number) => value.toLocaleString()),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/teams/useTeamMetadataSchema", () => ({
|
||||
useTeamMetadataSchema: vi.fn(() => ({ data: [], isLoading: false })),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({
|
||||
useUISettings: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({
|
||||
useAllProxyModels: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
|
||||
useTeam: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({
|
||||
organizationKeys: { all: ["organizations"] },
|
||||
useOrganization: vi.fn(),
|
||||
useOrganizations: vi.fn().mockReturnValue({ data: [], isLoading: false }),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({
|
||||
useCurrentUser: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({
|
||||
useMCPServers: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({
|
||||
useMCPToolsets: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({
|
||||
default: ({
|
||||
value,
|
||||
onChange,
|
||||
}: {
|
||||
value?: { servers: string[]; accessGroups: string[]; toolsets?: string[] };
|
||||
onChange: (next: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void;
|
||||
}) => (
|
||||
<>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() =>
|
||||
onChange({ servers: [], accessGroups: value?.accessGroups ?? [], toolsets: value?.toolsets ?? [] })
|
||||
}
|
||||
>
|
||||
deselect all mcp servers
|
||||
</button>
|
||||
<button type="button" onClick={() => onChange({ servers: value?.servers ?? [], accessGroups: [], toolsets: [] })}>
|
||||
remove all access groups
|
||||
</button>
|
||||
</>
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/team/TeamMemberTab", () => ({
|
||||
default: vi.fn(({ setIsAddMemberModalVisible }) => (
|
||||
<div>
|
||||
<button onClick={() => setIsAddMemberModalVisible(true)}>Add Member</button>
|
||||
</div>
|
||||
)),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/common_components/user_search_modal", () => ({
|
||||
default: vi.fn(({ isVisible, onCancel, onSubmit }) =>
|
||||
isVisible ? (
|
||||
<div>
|
||||
<button onClick={onCancel}>Cancel</button>
|
||||
<button onClick={() => onSubmit({ user_email: "new@test.com", user_id: "new-user", role: "user" })}>
|
||||
Submit
|
||||
</button>
|
||||
</div>
|
||||
) : null,
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/team/EditMembership", () => ({
|
||||
default: vi.fn(({ visible, onCancel, onSubmit }) =>
|
||||
visible ? (
|
||||
<div>
|
||||
<button onClick={onCancel}>Cancel</button>
|
||||
<button onClick={() => onSubmit({ user_email: "edit@test.com", user_id: "edit-user", role: "admin" })}>
|
||||
Submit
|
||||
</button>
|
||||
</div>
|
||||
) : null,
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/common_components/DeleteResourceModal", () => ({
|
||||
default: vi.fn(({ isOpen, onCancel, onOk }) =>
|
||||
isOpen ? (
|
||||
<div>
|
||||
<button onClick={onCancel}>Cancel</button>
|
||||
<button onClick={onOk}>Confirm Delete</button>
|
||||
</div>
|
||||
) : null,
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/team/member_permissions", () => ({
|
||||
default: vi.fn(() => <div>Member Permissions</div>),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/common_components/ModelAliasManager", () => ({
|
||||
default: vi.fn(({ initialModelAliases, onAliasUpdate }) => (
|
||||
<div>
|
||||
<div data-testid="alias-editor-initial">{JSON.stringify(initialModelAliases)}</div>
|
||||
<button type="button" onClick={() => onAliasUpdate({ "gpt-4o": "gpt-4" })}>
|
||||
Set Alias
|
||||
</button>
|
||||
<button type="button" onClick={() => onAliasUpdate({})}>
|
||||
Clear Aliases
|
||||
</button>
|
||||
</div>
|
||||
)),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({
|
||||
useAccessGroups: vi.fn().mockReturnValue({
|
||||
data: [
|
||||
{ access_group_id: "ag-1", access_group_name: "Group 1", access_mcp_server_ids: [] },
|
||||
{ access_group_id: "ag-2", access_group_name: "Group 2", access_mcp_server_ids: [] },
|
||||
],
|
||||
isLoading: false,
|
||||
isError: false,
|
||||
}),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/common_components/AccessGroupSelector", () => ({
|
||||
default: ({ value, onChange }: { value?: string[]; onChange?: (next: string[]) => void }) => (
|
||||
<button type="button" onClick={() => onChange?.((value ?? []).slice(1))}>
|
||||
remove first unified access group
|
||||
</button>
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => {
|
||||
const useKeysResult = {
|
||||
data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 },
|
||||
isPending: false,
|
||||
isFetching: false,
|
||||
refetch: vi.fn(),
|
||||
};
|
||||
return { useKeys: vi.fn().mockReturnValue(useKeysResult) };
|
||||
});
|
||||
|
||||
vi.mock("../key_team_helpers/filter_helpers", () => ({
|
||||
fetchTeamFilterOptions: vi.fn().mockResolvedValue({
|
||||
keyAliases: [],
|
||||
organizationIds: [],
|
||||
userIds: [],
|
||||
}),
|
||||
fetchAllKeyAliases: vi.fn().mockResolvedValue([]),
|
||||
fetchAllOrganizations: vi.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels";
|
||||
import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
|
||||
import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
|
||||
import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams";
|
||||
import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser";
|
||||
import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
|
||||
import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets";
|
||||
import { useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups";
|
||||
import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
|
||||
|
||||
const mockUseAllProxyModels = vi.mocked(useAllProxyModels);
|
||||
const mockUseKeys = vi.mocked(useKeys);
|
||||
const mockUseTeam = vi.mocked(useTeam);
|
||||
const mockUseOrganization = vi.mocked(useOrganization);
|
||||
const mockUseCurrentUser = vi.mocked(useCurrentUser);
|
||||
const mockUseMCPServers = vi.mocked(useMCPServers);
|
||||
const mockUseMCPToolsets = vi.mocked(useMCPToolsets);
|
||||
const mockUseAccessGroups = vi.mocked(useAccessGroups);
|
||||
const mockUseUISettings = vi.mocked(useUISettings);
|
||||
|
||||
const createMockTeamData = (overrides = {}) => ({
|
||||
team_id: "123",
|
||||
team_info: {
|
||||
team_alias: "Test Team",
|
||||
team_id: "123",
|
||||
organization_id: null,
|
||||
admins: ["admin@test.com"],
|
||||
members: ["user1@test.com"],
|
||||
members_with_roles: [
|
||||
{
|
||||
user_id: "user1@test.com",
|
||||
user_email: "user1@test.com",
|
||||
role: "member",
|
||||
spend: 0,
|
||||
budget_id: "budget1",
|
||||
},
|
||||
],
|
||||
metadata: {},
|
||||
tpm_limit: null,
|
||||
rpm_limit: null,
|
||||
max_budget: null,
|
||||
budget_duration: null,
|
||||
models: [],
|
||||
blocked: false,
|
||||
spend: 0,
|
||||
max_parallel_requests: null,
|
||||
budget_reset_at: null,
|
||||
model_id: null,
|
||||
litellm_model_table: null,
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
team_member_budget_table: null,
|
||||
guardrails: [],
|
||||
policies: [],
|
||||
object_permission: null,
|
||||
...overrides,
|
||||
},
|
||||
keys: [],
|
||||
team_memberships: [],
|
||||
});
|
||||
|
||||
const seedDefaultMocks = () => {
|
||||
mockUseAllProxyModels.mockReturnValue({
|
||||
data: { data: [] },
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useAllProxyModels>);
|
||||
mockUseTeam.mockReturnValue({
|
||||
data: undefined,
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useTeam>);
|
||||
mockUseOrganization.mockReturnValue({
|
||||
data: undefined,
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useOrganization>);
|
||||
mockUseCurrentUser.mockReturnValue({
|
||||
data: { models: [] },
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useCurrentUser>);
|
||||
mockUseMCPServers.mockReturnValue({ data: [], isLoading: false, isError: false } as unknown as ReturnType<
|
||||
typeof useMCPServers
|
||||
>);
|
||||
mockUseMCPToolsets.mockReturnValue({ data: [], isLoading: false, isError: false } as unknown as ReturnType<
|
||||
typeof useMCPToolsets
|
||||
>);
|
||||
mockUseAccessGroups.mockReturnValue({
|
||||
data: [
|
||||
{ access_group_id: "ag-1", access_group_name: "Group 1", access_mcp_server_ids: [] },
|
||||
{ access_group_id: "ag-2", access_group_name: "Group 2", access_mcp_server_ids: [] },
|
||||
],
|
||||
isLoading: false,
|
||||
isError: false,
|
||||
} as unknown as ReturnType<typeof useAccessGroups>);
|
||||
mockUseUISettings.mockReturnValue({
|
||||
data: { values: {} },
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useUISettings>);
|
||||
mockUseKeys.mockReturnValue({
|
||||
data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 },
|
||||
isPending: false,
|
||||
isFetching: false,
|
||||
refetch: vi.fn(),
|
||||
} as unknown as ReturnType<typeof useKeys>);
|
||||
vi.mocked(useTeamMetadataSchema).mockReturnValue({
|
||||
data: [],
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useTeamMetadataSchema>);
|
||||
|
||||
can.mockReturnValue(true);
|
||||
vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] });
|
||||
vi.mocked(networking.getPoliciesList).mockResolvedValue({ policies: [] });
|
||||
vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]);
|
||||
vi.mocked(networking.getTeamPermissionsCall).mockResolvedValue({
|
||||
all_available_permissions: [],
|
||||
team_member_permissions: [],
|
||||
});
|
||||
};
|
||||
describe("TeamInfoView - member budget reset prompt", () => {
|
||||
const props = {
|
||||
teamId: "123",
|
||||
onUpdate: vi.fn(),
|
||||
onClose: vi.fn(),
|
||||
accessToken: "test-token",
|
||||
is_team_admin: true,
|
||||
is_proxy_admin: true,
|
||||
userModels: ["gpt-4", "gpt-3.5-turbo"],
|
||||
editTeam: false,
|
||||
premiumUser: false,
|
||||
};
|
||||
|
||||
const customBudgetMembership = (
|
||||
userId: string,
|
||||
maxBudget: number | null = 50,
|
||||
): TeamData["team_memberships"][number] => ({
|
||||
user_id: userId,
|
||||
team_id: "123",
|
||||
budget_id: `budget-${userId}`,
|
||||
budget_source: "custom",
|
||||
spend: 0,
|
||||
total_spend: 0,
|
||||
litellm_budget_table: {
|
||||
budget_id: `budget-${userId}`,
|
||||
soft_budget: null,
|
||||
max_budget: maxBudget,
|
||||
max_parallel_requests: null,
|
||||
tpm_limit: null,
|
||||
rpm_limit: null,
|
||||
model_max_budget: null,
|
||||
budget_duration: null,
|
||||
budget_reset_at: null,
|
||||
},
|
||||
});
|
||||
|
||||
const savedTeam: Team = {
|
||||
team_id: "123",
|
||||
team_alias: "Test Team",
|
||||
models: [],
|
||||
max_budget: null,
|
||||
budget_duration: null,
|
||||
tpm_limit: null,
|
||||
rpm_limit: null,
|
||||
organization_id: "org-1",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
keys: [],
|
||||
members_with_roles: [],
|
||||
spend: 0,
|
||||
};
|
||||
|
||||
const openEditorWithCustomMembers = async (
|
||||
user: ReturnType<typeof userEvent.setup>,
|
||||
userIds: string[] = ["user-custom"],
|
||||
maxBudget: number | null = 50,
|
||||
) => {
|
||||
const data = {
|
||||
...createMockTeamData({
|
||||
team_member_budget_table: { max_budget: 10, budget_duration: null, tpm_limit: null, rpm_limit: null },
|
||||
}),
|
||||
team_memberships: userIds.map((id) => customBudgetMembership(id, maxBudget)),
|
||||
} as TeamData;
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(data);
|
||||
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" });
|
||||
|
||||
renderWithProviders(<TeamInfoView {...props} />);
|
||||
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
|
||||
await user.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
await user.click(screen.getByText("Team Member Settings"));
|
||||
return await screen.findByLabelText("Default Budget (USD)");
|
||||
};
|
||||
|
||||
const submitNewDefault = async (user: ReturnType<typeof userEvent.setup>, input: HTMLElement, value: string) => {
|
||||
fireEvent.change(input, { target: { value } });
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
seedDefaultMocks();
|
||||
bulkUpdatePOST.mockResolvedValue({ data: { data: [] } });
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("holds the save behind the prompt when the default changes while a member has a custom budget", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const input = await openEditorWithCustomMembers(user);
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
|
||||
expect(await screen.findByText("Reset member budgets?")).toBeInTheDocument();
|
||||
expect(screen.getByText(/1 member has a custom budget/)).toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Reset to $20.00" })).toBeInTheDocument();
|
||||
expect(networking.teamUpdateCall).not.toHaveBeenCalled();
|
||||
expect(bulkUpdatePOST).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("saves the new default and leaves custom budgets alone on keep", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const input = await openEditorWithCustomMembers(user);
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
await user.click(await screen.findByRole("button", { name: "Keep custom budget" }));
|
||||
|
||||
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20);
|
||||
expect(bulkUpdatePOST).not.toHaveBeenCalled();
|
||||
await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument());
|
||||
});
|
||||
|
||||
it("saves the new default then clears each custom budget through the bulk endpoint on reset all", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const input = await openEditorWithCustomMembers(user, ["user-a", "user-b"]);
|
||||
bulkUpdatePOST.mockResolvedValue({
|
||||
data: {
|
||||
data: [
|
||||
{ success: true, user_id: "user-a" },
|
||||
{ success: true, user_id: "user-b" },
|
||||
],
|
||||
},
|
||||
});
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
await user.click(await screen.findByRole("button", { name: "Reset all to $20.00" }));
|
||||
|
||||
await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalledTimes(1));
|
||||
expect(networking.teamUpdateCall).toHaveBeenCalled();
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20);
|
||||
expect(bulkUpdatePOST).toHaveBeenCalledWith("/management/v1/teams/{team_id}/members/bulk_update", {
|
||||
params: { path: { team_id: "123" } },
|
||||
body: {
|
||||
members: [
|
||||
{ user_id: "user-a", max_budget_in_team: null },
|
||||
{ user_id: "user-b", max_budget_in_team: null },
|
||||
],
|
||||
},
|
||||
});
|
||||
await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Reset 2 member budgets to the team default"));
|
||||
});
|
||||
|
||||
it("keeps the prompt mounted while the member-budget reset is in flight", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const input = await openEditorWithCustomMembers(user);
|
||||
const bulk = Promise.withResolvers<{ data: { data: { success: boolean; user_id: string }[] } }>();
|
||||
vi.mocked(networking.teamInfoCall).mockImplementationOnce(() => new Promise(() => {}));
|
||||
bulkUpdatePOST.mockImplementationOnce(() => bulk.promise);
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
await user.click(await screen.findByRole("button", { name: "Reset to $20.00" }));
|
||||
|
||||
await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalled());
|
||||
expect(screen.getByText("Reset member budgets?")).toBeInTheDocument();
|
||||
expect(screen.queryByText("Loading...")).not.toBeInTheDocument();
|
||||
|
||||
bulk.resolve({ data: { data: [{ success: true, user_id: "user-custom" }] } });
|
||||
await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument());
|
||||
});
|
||||
|
||||
it("offers a retry after a failed reset and never re-saves the team", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const input = await openEditorWithCustomMembers(user);
|
||||
bulkUpdatePOST
|
||||
.mockRejectedValueOnce(new Error("bulk update down"))
|
||||
.mockResolvedValueOnce({ data: { data: [{ success: true, user_id: "user-custom" }] } });
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
await user.click(await screen.findByRole("button", { name: "Reset to $20.00" }));
|
||||
|
||||
await waitFor(() =>
|
||||
expect(toast.fromError).toHaveBeenCalledWith("Team updated, but member budgets could not be reset"),
|
||||
);
|
||||
expect(screen.getByText("Reset member budgets?")).toBeInTheDocument();
|
||||
expect(screen.queryByText("Loading...")).not.toBeInTheDocument();
|
||||
await waitFor(() => expect(vi.mocked(networking.teamInfoCall).mock.calls.length).toBeGreaterThan(1));
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Retry reset" }));
|
||||
|
||||
await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalledTimes(2));
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls).toHaveLength(1);
|
||||
await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Reset 1 member budget to the team default"));
|
||||
await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument());
|
||||
});
|
||||
|
||||
it("returns to the prompt without its own toast when the team save fails", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const input = await openEditorWithCustomMembers(user);
|
||||
vi.mocked(networking.teamUpdateCall).mockRejectedValueOnce(new Error("save failed"));
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
await user.click(await screen.findByRole("button", { name: "Reset to $20.00" }));
|
||||
|
||||
expect(await screen.findByRole("button", { name: "Reset to $20.00" })).toBeEnabled();
|
||||
expect(bulkUpdatePOST).not.toHaveBeenCalled();
|
||||
expect(toast.fromError).not.toHaveBeenCalled();
|
||||
expect(toast.error).not.toHaveBeenCalled();
|
||||
expect(screen.queryByText("Loading...")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("reports how many budgets were already reset when a later batch fails", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const userIds = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `user-${i}`);
|
||||
const input = await openEditorWithCustomMembers(user, userIds);
|
||||
bulkUpdatePOST
|
||||
.mockResolvedValueOnce({
|
||||
data: {
|
||||
data: userIds.slice(0, MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES).map((user_id) => ({ success: true, user_id })),
|
||||
},
|
||||
})
|
||||
.mockRejectedValueOnce(new Error("second batch down"))
|
||||
.mockResolvedValueOnce({
|
||||
data: { data: [{ success: true, user_id: `user-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}` }] },
|
||||
});
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
await user.click(await screen.findByRole("button", { name: "Reset all to $20.00" }));
|
||||
|
||||
await waitFor(() =>
|
||||
expect(toast.error).toHaveBeenCalledWith(
|
||||
`Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES} of ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets; the rest could not be reset`,
|
||||
),
|
||||
);
|
||||
expect(bulkUpdatePOST).toHaveBeenCalledTimes(2);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Retry reset" }));
|
||||
|
||||
await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalledTimes(3));
|
||||
expect(vi.mocked(bulkUpdatePOST).mock.calls[2][1]).toEqual({
|
||||
params: { path: { team_id: "123" } },
|
||||
body: {
|
||||
members: [{ user_id: `user-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`, max_budget_in_team: null }],
|
||||
},
|
||||
});
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls).toHaveLength(1);
|
||||
await waitFor(() =>
|
||||
expect(toast.success).toHaveBeenCalledWith(
|
||||
`Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets to the team default`,
|
||||
),
|
||||
);
|
||||
});
|
||||
|
||||
it("surfaces a failure toast when some members cannot be reset", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const input = await openEditorWithCustomMembers(user, ["user-a", "user-b"]);
|
||||
bulkUpdatePOST.mockResolvedValue({
|
||||
data: {
|
||||
data: [
|
||||
{ success: true, user_id: "user-a" },
|
||||
{ success: false, user_id: "user-b", error: "no such member" },
|
||||
],
|
||||
},
|
||||
});
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
await user.click(await screen.findByRole("button", { name: "Reset all to $20.00" }));
|
||||
|
||||
await waitFor(() =>
|
||||
expect(toast.error).toHaveBeenCalledWith("Team updated, but 1 member budget could not be reset"),
|
||||
);
|
||||
});
|
||||
|
||||
it("aborts the save entirely on cancel", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const input = await openEditorWithCustomMembers(user);
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
await user.click(await screen.findByRole("button", { name: "Cancel" }));
|
||||
|
||||
await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument());
|
||||
expect(networking.teamUpdateCall).not.toHaveBeenCalled();
|
||||
expect(bulkUpdatePOST).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("saves directly when no member carries a custom budget", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const input = await openEditorWithCustomMembers(user, []);
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
|
||||
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20);
|
||||
expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument();
|
||||
expect(bulkUpdatePOST).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("saves directly when the default is resubmitted unchanged", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
await openEditorWithCustomMembers(user);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(10);
|
||||
expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument();
|
||||
expect(bulkUpdatePOST).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("dismisses a pending prompt when the selected team changes", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const data = {
|
||||
...createMockTeamData({
|
||||
team_member_budget_table: { max_budget: 10, budget_duration: null, tpm_limit: null, rpm_limit: null },
|
||||
}),
|
||||
team_memberships: [customBudgetMembership("user-custom")],
|
||||
} as TeamData;
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(data);
|
||||
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" });
|
||||
|
||||
const { rerender } = renderWithProviders(<TeamInfoView {...props} />);
|
||||
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
|
||||
await user.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
await user.click(screen.getByText("Team Member Settings"));
|
||||
const input = await screen.findByLabelText("Default Budget (USD)");
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
expect(await screen.findByText("Reset member budgets?")).toBeInTheDocument();
|
||||
|
||||
rerender(<TeamInfoView {...props} teamId="456" />);
|
||||
|
||||
await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument());
|
||||
expect(bulkUpdatePOST).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("does not reopen the dialog or refetch the old team when the team changes mid-reset", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const data = {
|
||||
...createMockTeamData({
|
||||
team_member_budget_table: { max_budget: 10, budget_duration: null, tpm_limit: null, rpm_limit: null },
|
||||
}),
|
||||
team_memberships: [customBudgetMembership("user-custom")],
|
||||
} as TeamData;
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(data);
|
||||
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" });
|
||||
const bulkDone = Promise.withResolvers<{ data: { data: { success: boolean; user_id: string }[] } }>();
|
||||
bulkUpdatePOST.mockReturnValue(bulkDone.promise);
|
||||
|
||||
const { rerender } = renderWithProviders(<TeamInfoView {...props} />);
|
||||
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
|
||||
await user.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
await user.click(screen.getByText("Team Member Settings"));
|
||||
const input = await screen.findByLabelText("Default Budget (USD)");
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
await user.click(await screen.findByRole("button", { name: "Reset to $20.00" }));
|
||||
await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalled());
|
||||
|
||||
const infoCallsBeforeSwitch = vi.mocked(networking.teamInfoCall).mock.calls.length;
|
||||
rerender(<TeamInfoView {...props} teamId="456" />);
|
||||
bulkDone.resolve({ data: { data: [{ success: true, user_id: "user-custom" }] } });
|
||||
|
||||
await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument());
|
||||
await waitFor(() =>
|
||||
expect(
|
||||
vi
|
||||
.mocked(networking.teamInfoCall)
|
||||
.mock.calls.slice(infoCallsBeforeSwitch)
|
||||
.some((call) => call[1] === "456"),
|
||||
).toBe(true),
|
||||
);
|
||||
const infoCallsAfterSwitch = vi.mocked(networking.teamInfoCall).mock.calls.slice(infoCallsBeforeSwitch);
|
||||
expect(infoCallsAfterSwitch.every((call) => call[1] === "456")).toBe(true);
|
||||
expect(screen.queryByRole("button", { name: "Retry reset" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("saves directly when a custom member's cap is null and already inherits the default", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const input = await openEditorWithCustomMembers(user, ["user-limits-only"], null);
|
||||
|
||||
await submitNewDefault(user, input, "20");
|
||||
|
||||
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20);
|
||||
expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument();
|
||||
expect(bulkUpdatePOST).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
@ -55,6 +55,12 @@ vi.mock("@/components/networking", () => ({
|
|||
getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [], count: 0 }),
|
||||
}));
|
||||
|
||||
const { bulkUpdatePOST } = vi.hoisted(() => ({ bulkUpdatePOST: vi.fn() }));
|
||||
vi.mock("@/lib/http/api", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("@/lib/http/api")>();
|
||||
return { ...actual, fetchClient: { ...actual.fetchClient, POST: bulkUpdatePOST } };
|
||||
});
|
||||
|
||||
const can = vi.fn();
|
||||
vi.mock("@/app/(dashboard)/hooks/useCan", () => ({
|
||||
default: (...args: unknown[]) => can(...args),
|
||||
|
|
|
|||
|
|
@ -119,6 +119,10 @@ import {
|
|||
} from "./tabVisibilityUtils";
|
||||
import TeamMembersComponent from "./TeamMemberTab";
|
||||
import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable";
|
||||
import ResetMemberBudgetsDialog from "./ResetMemberBudgetsDialog";
|
||||
import { customBudgetMemberUserIds, shouldPromptMemberBudgetReset } from "./memberBudgetReset";
|
||||
import { useMemberBudgetReset } from "./useMemberBudgetReset";
|
||||
import { fetchClient } from "@/lib/http/api";
|
||||
|
||||
const UI_MANAGED_METADATA_KEYS: ReadonlySet<string> = new Set([
|
||||
"logging",
|
||||
|
|
@ -884,17 +888,41 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
const persistTeamUpdate = async (token: string, updateData: Record<string, unknown>) => {
|
||||
await teamUpdateCall(token, updateData);
|
||||
queryClient.invalidateQueries({ queryKey: organizationKeys.all });
|
||||
|
||||
toast.success("Team settings updated successfully");
|
||||
setIsEditing(false);
|
||||
fetchTeamInfo();
|
||||
};
|
||||
|
||||
const memberBudgetReset = useMemberBudgetReset({
|
||||
saveTeam: async (updateData) => {
|
||||
if (!accessToken) return;
|
||||
setIsTeamSaving(true);
|
||||
try {
|
||||
await persistTeamUpdate(accessToken, updateData);
|
||||
} finally {
|
||||
setIsTeamSaving(false);
|
||||
}
|
||||
},
|
||||
resetMemberBudgets: async (bulkTeamId, userIds) => {
|
||||
const { data } = await fetchClient.POST("/management/v1/teams/{team_id}/members/bulk_update", {
|
||||
params: { path: { team_id: bulkTeamId } },
|
||||
body: { members: userIds.map((user_id) => ({ user_id, max_budget_in_team: null })) },
|
||||
});
|
||||
return data?.data ?? [];
|
||||
},
|
||||
refreshTeamData,
|
||||
});
|
||||
|
||||
const { dismiss: dismissMemberBudgetReset } = memberBudgetReset;
|
||||
useEffect(() => {
|
||||
dismissMemberBudgetReset();
|
||||
}, [teamId, dismissMemberBudgetReset]);
|
||||
|
||||
const saveTeamAdminSettings = async (changes: TeamAdminSettingsChanges) => {
|
||||
if (!accessToken) return;
|
||||
setIsTeamSaving(true);
|
||||
try {
|
||||
await persistTeamUpdate(accessToken, { team_id: teamId, ...changes });
|
||||
toast.success("Team settings updated successfully");
|
||||
await fetchTeamInfo();
|
||||
} catch (error) {
|
||||
console.error("Error updating team:", error);
|
||||
} finally {
|
||||
|
|
@ -1005,8 +1033,10 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
updateData.max_budget = mapEmptyStringToNull(updateData.max_budget);
|
||||
updateData.team_member_budget_duration = values.team_member_budget_duration;
|
||||
|
||||
if (values.team_member_budget !== undefined) {
|
||||
updateData.team_member_budget = Number(values.team_member_budget);
|
||||
const newTeamMemberBudget =
|
||||
values.team_member_budget !== undefined ? Number(values.team_member_budget) : undefined;
|
||||
if (newTeamMemberBudget !== undefined) {
|
||||
updateData.team_member_budget = newTeamMemberBudget;
|
||||
}
|
||||
|
||||
if (values.team_member_key_duration !== undefined) {
|
||||
|
|
@ -1152,7 +1182,28 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
}
|
||||
}
|
||||
|
||||
const customBudgetUserIds = customBudgetMemberUserIds(teamData?.team_memberships ?? []);
|
||||
if (
|
||||
newTeamMemberBudget !== undefined &&
|
||||
shouldPromptMemberBudgetReset(
|
||||
newTeamMemberBudget,
|
||||
info.team_member_budget_table?.max_budget,
|
||||
customBudgetUserIds,
|
||||
)
|
||||
) {
|
||||
const pendingReset = {
|
||||
teamId,
|
||||
updateData,
|
||||
userIds: customBudgetUserIds,
|
||||
newBudget: newTeamMemberBudget,
|
||||
};
|
||||
memberBudgetReset.prompt(pendingReset);
|
||||
return;
|
||||
}
|
||||
|
||||
await persistTeamUpdate(accessToken, updateData);
|
||||
toast.success("Team settings updated successfully");
|
||||
await fetchTeamInfo();
|
||||
} catch (error) {
|
||||
console.error("Error updating team:", error);
|
||||
} finally {
|
||||
|
|
@ -2426,6 +2477,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
onOk={handleDeleteConfirm}
|
||||
confirmLoading={isDeleting}
|
||||
/>
|
||||
|
||||
<ResetMemberBudgetsDialog
|
||||
state={memberBudgetReset.state}
|
||||
onReset={memberBudgetReset.reset}
|
||||
onRetry={memberBudgetReset.retry}
|
||||
onKeep={memberBudgetReset.keepCustom}
|
||||
onDismiss={memberBudgetReset.dismiss}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -682,6 +682,44 @@ describe("TeamMembersComponent", () => {
|
|||
expect(inheritedRow).toHaveTextContent("$25.00");
|
||||
});
|
||||
|
||||
it("caps a Custom member at the team default when the private row has no budget limit", () => {
|
||||
const base = createMockTeamData();
|
||||
renderTab(
|
||||
createMockTeamData({
|
||||
team_info: {
|
||||
...base.team_info,
|
||||
team_member_budget_table: { max_budget: 20, budget_duration: null, tpm_limit: null, rpm_limit: null },
|
||||
},
|
||||
team_memberships: [
|
||||
{
|
||||
user_id: "user2@test.com",
|
||||
team_id: "team-123",
|
||||
budget_id: "budget2",
|
||||
budget_source: "custom",
|
||||
spend: 0,
|
||||
total_spend: 0,
|
||||
litellm_budget_table: {
|
||||
budget_id: "budget3",
|
||||
soft_budget: null,
|
||||
max_budget: null,
|
||||
max_parallel_requests: null,
|
||||
tpm_limit: null,
|
||||
rpm_limit: 100,
|
||||
model_max_budget: null,
|
||||
budget_duration: null,
|
||||
budget_reset_at: null,
|
||||
},
|
||||
},
|
||||
],
|
||||
}),
|
||||
);
|
||||
|
||||
const row = screen.getByRole("row", { name: /user2@test\.com/ });
|
||||
expect(within(row).getByTestId("member-budget-source")).toHaveTextContent("Custom");
|
||||
expect(row).toHaveTextContent("$20.00");
|
||||
expect(row).not.toHaveTextContent("Unlimited");
|
||||
});
|
||||
|
||||
it("shows no source label for a member with neither a custom nor a team budget", () => {
|
||||
renderTab(createMockTeamData({ team_memberships: [] }));
|
||||
|
||||
|
|
|
|||
|
|
@ -106,10 +106,7 @@ export default function TeamMemberTab({
|
|||
const getUserBudget = (userId: string | null): number | null => {
|
||||
if (!userId) return null;
|
||||
const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
|
||||
return (
|
||||
membership?.litellm_budget_table?.max_budget ??
|
||||
(membership?.budget_source === "team_default" ? teamDefaultBudget : null)
|
||||
);
|
||||
return membership?.litellm_budget_table?.max_budget ?? teamDefaultBudget;
|
||||
};
|
||||
|
||||
// Helper function to get rate limits for a user
|
||||
|
|
|
|||
|
|
@ -0,0 +1,101 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
chunk,
|
||||
customBudgetMemberUserIds,
|
||||
MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES,
|
||||
pluralize,
|
||||
shouldPromptMemberBudgetReset,
|
||||
} from "./memberBudgetReset";
|
||||
|
||||
describe("customBudgetMemberUserIds", () => {
|
||||
const customRow = (user_id: string, max_budget: number | null = 50) => ({
|
||||
user_id,
|
||||
budget_source: "custom" as const,
|
||||
litellm_budget_table: { max_budget },
|
||||
});
|
||||
|
||||
it("returns only members whose budget_source is custom with a private cap", () => {
|
||||
const memberships = [
|
||||
customRow("u-custom"),
|
||||
{ user_id: "u-default", budget_source: "team_default" as const },
|
||||
{ user_id: "u-none", budget_source: "none" as const },
|
||||
];
|
||||
|
||||
expect(customBudgetMemberUserIds(memberships)).toEqual(["u-custom"]);
|
||||
});
|
||||
|
||||
it("skips a custom row whose cap is null since it already inherits the default", () => {
|
||||
const memberships = [customRow("u-rate-limits-only", null), customRow("u-custom")];
|
||||
|
||||
expect(customBudgetMemberUserIds(memberships)).toEqual(["u-custom"]);
|
||||
});
|
||||
|
||||
it("skips a custom row that has no budget table at all", () => {
|
||||
const memberships = [
|
||||
{ user_id: "u-no-table", budget_source: "custom" as const, litellm_budget_table: null },
|
||||
customRow("u-custom"),
|
||||
];
|
||||
|
||||
expect(customBudgetMemberUserIds(memberships)).toEqual(["u-custom"]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("shouldPromptMemberBudgetReset", () => {
|
||||
it("prompts when the default changes and custom-budget members exist", () => {
|
||||
expect(shouldPromptMemberBudgetReset(10, 5, ["u-1"])).toBe(true);
|
||||
});
|
||||
|
||||
it("prompts when a default is set for the first time", () => {
|
||||
expect(shouldPromptMemberBudgetReset(10, null, ["u-1"])).toBe(true);
|
||||
expect(shouldPromptMemberBudgetReset(10, undefined, ["u-1"])).toBe(true);
|
||||
});
|
||||
|
||||
it("does not prompt when the submitted budget is unchanged", () => {
|
||||
expect(shouldPromptMemberBudgetReset(10, 10, ["u-1"])).toBe(false);
|
||||
});
|
||||
|
||||
it("does not prompt when the budget is cleared or zero", () => {
|
||||
expect(shouldPromptMemberBudgetReset(undefined, 10, ["u-1"])).toBe(false);
|
||||
expect(shouldPromptMemberBudgetReset(0, 10, ["u-1"])).toBe(false);
|
||||
});
|
||||
|
||||
it("does not prompt when no member has a custom budget", () => {
|
||||
expect(shouldPromptMemberBudgetReset(10, 5, [])).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("chunk", () => {
|
||||
it("splits selections larger than the bulk endpoint limit", () => {
|
||||
const ids = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `u-${i}`);
|
||||
|
||||
const chunks = chunk(ids, MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES);
|
||||
|
||||
expect(chunks).toHaveLength(2);
|
||||
expect(chunks[0]).toHaveLength(MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES);
|
||||
expect(chunks[1]).toEqual([`u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`]);
|
||||
});
|
||||
|
||||
it("keeps a selection under the limit in a single chunk", () => {
|
||||
expect(chunk(["u-1", "u-2"], MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES)).toEqual([["u-1", "u-2"]]);
|
||||
});
|
||||
|
||||
it("returns no chunks for an empty selection", () => {
|
||||
expect(chunk([], MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES)).toEqual([]);
|
||||
});
|
||||
|
||||
it("returns no chunks for a non-positive size", () => {
|
||||
expect(chunk(["u-1"], 0)).toEqual([]);
|
||||
expect(chunk(["u-1"], -3)).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("pluralize", () => {
|
||||
it("uses the singular form for exactly one", () => {
|
||||
expect(pluralize(1, "budget", "budgets")).toBe("budget");
|
||||
});
|
||||
|
||||
it("uses the plural form for zero and for more than one", () => {
|
||||
expect(pluralize(0, "budget", "budgets")).toBe("budgets");
|
||||
expect(pluralize(3, "budget", "budgets")).toBe("budgets");
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,38 @@
|
|||
import type { components } from "@/lib/http/schema";
|
||||
import type { TeamMembership } from "./TeamInfo";
|
||||
|
||||
export const MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES = 500;
|
||||
|
||||
export type TeamUpdatePayload = components["schemas"]["UpdateTeamRequest"];
|
||||
|
||||
export interface MemberBudgetResetPending {
|
||||
readonly teamId: string;
|
||||
readonly updateData: TeamUpdatePayload;
|
||||
readonly userIds: readonly string[];
|
||||
readonly newBudget: number;
|
||||
}
|
||||
|
||||
type MembershipBudgetRow = Pick<TeamMembership, "user_id" | "budget_source"> & {
|
||||
litellm_budget_table?: Pick<TeamMembership["litellm_budget_table"], "max_budget"> | null;
|
||||
};
|
||||
|
||||
export const customBudgetMemberUserIds = (memberships: readonly MembershipBudgetRow[]): string[] =>
|
||||
memberships
|
||||
.filter((m) => m.budget_source === "custom" && m.litellm_budget_table?.max_budget != null)
|
||||
.map((m) => m.user_id);
|
||||
|
||||
export const shouldPromptMemberBudgetReset = (
|
||||
nextBudget: number | undefined,
|
||||
previousBudget: number | null | undefined,
|
||||
customBudgetUserIds: string[],
|
||||
): boolean => {
|
||||
const budgetChanged = typeof nextBudget === "number" && nextBudget > 0 && nextBudget !== previousBudget;
|
||||
return budgetChanged && customBudgetUserIds.length > 0;
|
||||
};
|
||||
|
||||
export const chunk = <T>(items: readonly T[], size: number): T[][] => {
|
||||
if (size <= 0) return [];
|
||||
return Array.from({ length: Math.ceil(items.length / size) }, (_, i) => items.slice(i * size, i * size + size));
|
||||
};
|
||||
|
||||
export const pluralize = (count: number, singular: string, plural: string): string => (count === 1 ? singular : plural);
|
||||
|
|
@ -0,0 +1,345 @@
|
|||
import { act, renderHook } from "@testing-library/react";
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES } from "./memberBudgetReset";
|
||||
import { useMemberBudgetReset, type MemberBudgetResetGateway } from "./useMemberBudgetReset";
|
||||
|
||||
const buildGateway = () => ({
|
||||
saveTeam: vi.fn(async (_updateData: Record<string, unknown>) => {}),
|
||||
resetMemberBudgets: vi.fn(async (_teamId: string, userIds: readonly string[]) =>
|
||||
userIds.map((user_id) => ({ success: true, user_id })),
|
||||
),
|
||||
refreshTeamData: vi.fn(async () => {}),
|
||||
});
|
||||
|
||||
const pendingFor = (userIds: string[]) => ({
|
||||
teamId: "team-123",
|
||||
updateData: { team_id: "team-123", team_member_budget: 20 },
|
||||
userIds,
|
||||
newBudget: 20,
|
||||
});
|
||||
|
||||
const renderReset = (gateway: MemberBudgetResetGateway) => renderHook(() => useMemberBudgetReset(gateway));
|
||||
|
||||
describe("useMemberBudgetReset", () => {
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("starts idle and prompts with the pending update", async () => {
|
||||
const gateway = buildGateway();
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1"]));
|
||||
});
|
||||
|
||||
expect(result.current.state).toEqual({
|
||||
phase: "prompting",
|
||||
pending: pendingFor(["u-1"]),
|
||||
});
|
||||
expect(gateway.saveTeam).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("saves the team once, resets every member, and refreshes on success", async () => {
|
||||
const gateway = buildGateway();
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1", "u-2"]));
|
||||
});
|
||||
await act(async () => {
|
||||
await result.current.reset();
|
||||
});
|
||||
|
||||
expect(gateway.saveTeam).toHaveBeenCalledTimes(1);
|
||||
expect(gateway.saveTeam).toHaveBeenCalledWith({ team_id: "team-123", team_member_budget: 20 });
|
||||
expect(gateway.resetMemberBudgets).toHaveBeenCalledWith("team-123", ["u-1", "u-2"]);
|
||||
expect(toast.success).toHaveBeenCalledWith("Reset 2 member budgets to the team default");
|
||||
expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1);
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
});
|
||||
|
||||
it("resets members in batches no larger than the bulk endpoint limit", async () => {
|
||||
const userIds = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `u-${i}`);
|
||||
const gateway = buildGateway();
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(userIds));
|
||||
});
|
||||
await act(async () => {
|
||||
await result.current.reset();
|
||||
});
|
||||
|
||||
expect(gateway.resetMemberBudgets).toHaveBeenCalledTimes(2);
|
||||
expect(gateway.resetMemberBudgets.mock.calls[0][0]).toBe("team-123");
|
||||
expect(gateway.resetMemberBudgets.mock.calls[0][1]).toHaveLength(MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES);
|
||||
expect(gateway.resetMemberBudgets.mock.calls[1][1]).toEqual([`u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`]);
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
});
|
||||
|
||||
it("resets against the team carried by the pending update", async () => {
|
||||
const gateway = buildGateway();
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
const pending = {
|
||||
teamId: "team-999",
|
||||
updateData: { team_id: "team-999", team_member_budget: 20 },
|
||||
userIds: ["u-1"],
|
||||
newBudget: 20,
|
||||
};
|
||||
await act(async () => {
|
||||
result.current.prompt(pending);
|
||||
});
|
||||
await act(async () => {
|
||||
await result.current.reset();
|
||||
});
|
||||
|
||||
expect(gateway.resetMemberBudgets).toHaveBeenCalledWith("team-999", ["u-1"]);
|
||||
expect(gateway.resetMemberBudgets).not.toHaveBeenCalledWith("team-123", expect.anything());
|
||||
});
|
||||
|
||||
it("returns to prompting without its own toast when the team save fails", async () => {
|
||||
const gateway = buildGateway();
|
||||
gateway.saveTeam.mockRejectedValueOnce(new Error("team update failed"));
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1"]));
|
||||
});
|
||||
await act(async () => {
|
||||
await result.current.reset();
|
||||
});
|
||||
|
||||
expect(result.current.state.phase).toBe("prompting");
|
||||
expect(gateway.resetMemberBudgets).not.toHaveBeenCalled();
|
||||
expect(gateway.refreshTeamData).not.toHaveBeenCalled();
|
||||
expect(toast.error).not.toHaveBeenCalled();
|
||||
expect(toast.fromError).not.toHaveBeenCalled();
|
||||
expect(toast.success).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("retries only the unsent members when a later batch fails, without re-saving the team", async () => {
|
||||
const userIds = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `u-${i}`);
|
||||
const gateway = buildGateway();
|
||||
gateway.resetMemberBudgets
|
||||
.mockResolvedValueOnce(
|
||||
userIds.slice(0, MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES).map((user_id) => ({ success: true, user_id })),
|
||||
)
|
||||
.mockRejectedValueOnce(new Error("second batch failed"));
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(userIds));
|
||||
});
|
||||
await act(async () => {
|
||||
await result.current.reset();
|
||||
});
|
||||
|
||||
expect(result.current.state.phase).toBe("resetFailed");
|
||||
expect(toast.error).toHaveBeenCalledWith(
|
||||
`Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES} of ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets; the rest could not be reset`,
|
||||
);
|
||||
|
||||
gateway.resetMemberBudgets.mockResolvedValueOnce([
|
||||
{ success: true, user_id: `u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}` },
|
||||
]);
|
||||
await act(async () => {
|
||||
await result.current.retry();
|
||||
});
|
||||
|
||||
expect(gateway.saveTeam).toHaveBeenCalledTimes(1);
|
||||
expect(gateway.resetMemberBudgets).toHaveBeenCalledTimes(3);
|
||||
expect(gateway.resetMemberBudgets.mock.calls[2][0]).toBe("team-123");
|
||||
expect(gateway.resetMemberBudgets.mock.calls[2][1]).toEqual([`u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`]);
|
||||
expect(toast.success).toHaveBeenCalledWith(
|
||||
`Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets to the team default`,
|
||||
);
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
});
|
||||
|
||||
it("reports partial progress when the first batch fails", async () => {
|
||||
const gateway = buildGateway();
|
||||
gateway.resetMemberBudgets.mockRejectedValueOnce(new Error("batch failed"));
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1", "u-2"]));
|
||||
});
|
||||
await act(async () => {
|
||||
await result.current.reset();
|
||||
});
|
||||
|
||||
expect(result.current.state.phase).toBe("resetFailed");
|
||||
expect(toast.fromError).toHaveBeenCalledWith("Team updated, but member budgets could not be reset");
|
||||
expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("reports members the backend could not reset and closes", async () => {
|
||||
const gateway = buildGateway();
|
||||
gateway.resetMemberBudgets.mockResolvedValueOnce([
|
||||
{ success: true, user_id: "u-1" },
|
||||
{ success: false, user_id: "u-2" },
|
||||
]);
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1", "u-2"]));
|
||||
});
|
||||
await act(async () => {
|
||||
await result.current.reset();
|
||||
});
|
||||
|
||||
expect(toast.error).toHaveBeenCalledWith("Team updated, but 1 member budget could not be reset");
|
||||
expect(toast.success).not.toHaveBeenCalled();
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("saves the team once and closes on keep custom", async () => {
|
||||
const gateway = buildGateway();
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1"]));
|
||||
});
|
||||
await act(async () => {
|
||||
await result.current.keepCustom();
|
||||
});
|
||||
|
||||
expect(gateway.saveTeam).toHaveBeenCalledTimes(1);
|
||||
expect(toast.success).toHaveBeenCalledWith("Team settings updated successfully");
|
||||
expect(gateway.resetMemberBudgets).not.toHaveBeenCalled();
|
||||
expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1);
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
});
|
||||
|
||||
it("abandons the pending update on dismiss", async () => {
|
||||
const gateway = buildGateway();
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1"]));
|
||||
});
|
||||
await act(async () => {
|
||||
result.current.dismiss();
|
||||
});
|
||||
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
expect(gateway.saveTeam).not.toHaveBeenCalled();
|
||||
expect(gateway.resetMemberBudgets).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("skips the refresh and stays idle when dismissed while a reset is in flight", async () => {
|
||||
const gateway = buildGateway();
|
||||
const bulkDone = Promise.withResolvers<{ success: boolean; user_id: string }[]>();
|
||||
gateway.resetMemberBudgets.mockReturnValueOnce(bulkDone.promise);
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1"]));
|
||||
});
|
||||
let resetPromise = Promise.resolve();
|
||||
await act(async () => {
|
||||
resetPromise = result.current.reset();
|
||||
});
|
||||
expect(result.current.state.phase).toBe("resetting");
|
||||
|
||||
act(() => {
|
||||
result.current.dismiss();
|
||||
});
|
||||
await act(async () => {
|
||||
bulkDone.resolve([{ success: true, user_id: "u-1" }]);
|
||||
await resetPromise;
|
||||
});
|
||||
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
expect(gateway.refreshTeamData).not.toHaveBeenCalled();
|
||||
expect(toast.success).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("does not return to prompting when the team save fails after a dismiss", async () => {
|
||||
const gateway = buildGateway();
|
||||
const saveDone = Promise.withResolvers<void>();
|
||||
gateway.saveTeam.mockReturnValueOnce(saveDone.promise);
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1"]));
|
||||
});
|
||||
let resetPromise = Promise.resolve();
|
||||
await act(async () => {
|
||||
resetPromise = result.current.reset();
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.dismiss();
|
||||
});
|
||||
await act(async () => {
|
||||
saveDone.reject(new Error("team update failed"));
|
||||
await resetPromise;
|
||||
});
|
||||
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
expect(gateway.resetMemberBudgets).not.toHaveBeenCalled();
|
||||
expect(gateway.refreshTeamData).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("stays idle without a refresh when the reset fails after a dismiss", async () => {
|
||||
const gateway = buildGateway();
|
||||
const bulkDone = Promise.withResolvers<{ success: boolean; user_id: string }[]>();
|
||||
gateway.resetMemberBudgets.mockReturnValueOnce(bulkDone.promise);
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1"]));
|
||||
});
|
||||
let resetPromise = Promise.resolve();
|
||||
await act(async () => {
|
||||
resetPromise = result.current.reset();
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.dismiss();
|
||||
});
|
||||
await act(async () => {
|
||||
bulkDone.reject(new Error("bulk failed"));
|
||||
await resetPromise;
|
||||
});
|
||||
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
expect(gateway.refreshTeamData).not.toHaveBeenCalled();
|
||||
expect(toast.error).not.toHaveBeenCalled();
|
||||
expect(toast.fromError).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("skips the refresh when dismissed while keep-custom is saving", async () => {
|
||||
const gateway = buildGateway();
|
||||
const saveDone = Promise.withResolvers<void>();
|
||||
gateway.saveTeam.mockReturnValueOnce(saveDone.promise);
|
||||
const { result } = renderReset(gateway);
|
||||
|
||||
await act(async () => {
|
||||
result.current.prompt(pendingFor(["u-1"]));
|
||||
});
|
||||
let keepPromise = Promise.resolve();
|
||||
await act(async () => {
|
||||
keepPromise = result.current.keepCustom();
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.dismiss();
|
||||
});
|
||||
await act(async () => {
|
||||
saveDone.resolve();
|
||||
await keepPromise;
|
||||
});
|
||||
|
||||
expect(result.current.state.phase).toBe("idle");
|
||||
expect(gateway.refreshTeamData).not.toHaveBeenCalled();
|
||||
expect(toast.success).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
134
ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts
Normal file
134
ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
import { useCallback, useRef, useState } from "react";
|
||||
import type { components } from "@/lib/http/schema";
|
||||
import { toast } from "@/lib/toast";
|
||||
import {
|
||||
chunk,
|
||||
MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES,
|
||||
pluralize,
|
||||
type MemberBudgetResetPending,
|
||||
type TeamUpdatePayload,
|
||||
} from "./memberBudgetReset";
|
||||
|
||||
export type MemberBudgetBulkResult = components["schemas"]["TeamMemberBudgetUpdateResult"];
|
||||
|
||||
export type MemberBudgetResetState =
|
||||
| { phase: "idle" }
|
||||
| { phase: "prompting"; pending: MemberBudgetResetPending }
|
||||
| { phase: "resetting"; pending: MemberBudgetResetPending; attempted: number }
|
||||
| { phase: "resetFailed"; pending: MemberBudgetResetPending; attempted: number };
|
||||
|
||||
export interface MemberBudgetResetGateway {
|
||||
saveTeam: (updateData: TeamUpdatePayload) => Promise<void>;
|
||||
resetMemberBudgets: (teamId: string, userIds: readonly string[]) => Promise<MemberBudgetBulkResult[]>;
|
||||
refreshTeamData: () => Promise<void>;
|
||||
}
|
||||
|
||||
type ResetRun =
|
||||
| { ok: true; attempted: number; results: readonly MemberBudgetBulkResult[] }
|
||||
| { ok: false; attempted: number; error: unknown };
|
||||
|
||||
export const useMemberBudgetReset = (gateway: MemberBudgetResetGateway) => {
|
||||
const [state, setState] = useState<MemberBudgetResetState>({ phase: "idle" });
|
||||
const activeRun = useRef<object | null>(null);
|
||||
const isCurrent = (run: object) => activeRun.current === run;
|
||||
|
||||
const runReset = async (pending: MemberBudgetResetPending, attempted: number, run: object) => {
|
||||
const { resetMemberBudgets, refreshTeamData } = gateway;
|
||||
|
||||
const runChunks = (
|
||||
chunks: readonly (readonly string[])[],
|
||||
attemptedSoFar: number,
|
||||
results: readonly MemberBudgetBulkResult[],
|
||||
): Promise<ResetRun> => {
|
||||
const [ids, ...rest] = chunks;
|
||||
if (ids === undefined) return Promise.resolve({ ok: true, attempted: attemptedSoFar, results });
|
||||
return resetMemberBudgets(pending.teamId, ids).then(
|
||||
(batch) => runChunks(rest, attemptedSoFar + ids.length, [...results, ...batch]),
|
||||
(error: unknown) => ({ ok: false as const, attempted: attemptedSoFar, error }),
|
||||
);
|
||||
};
|
||||
|
||||
const outcome = await runChunks(
|
||||
chunk(pending.userIds.slice(attempted), MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES),
|
||||
attempted,
|
||||
[],
|
||||
);
|
||||
|
||||
if (!isCurrent(run)) return;
|
||||
if (!outcome.ok) {
|
||||
console.error("Error resetting member budgets:", outcome.error);
|
||||
const total = pending.userIds.length;
|
||||
if (outcome.attempted > 0) {
|
||||
toast.error(
|
||||
`Reset ${outcome.attempted} of ${total} member ${pluralize(total, "budget", "budgets")}; the rest could not be reset`,
|
||||
);
|
||||
} else {
|
||||
toast.fromError("Team updated, but member budgets could not be reset");
|
||||
}
|
||||
setState({ phase: "resetFailed", pending, attempted: outcome.attempted });
|
||||
await refreshTeamData();
|
||||
return;
|
||||
}
|
||||
const failed = outcome.results.filter((r) => !r.success);
|
||||
if (failed.length > 0) {
|
||||
toast.error(
|
||||
`Team updated, but ${failed.length} member ${pluralize(failed.length, "budget", "budgets")} could not be reset`,
|
||||
);
|
||||
} else {
|
||||
toast.success(
|
||||
`Reset ${outcome.attempted} member ${pluralize(outcome.attempted, "budget", "budgets")} to the team default`,
|
||||
);
|
||||
}
|
||||
setState({ phase: "idle" });
|
||||
await refreshTeamData();
|
||||
};
|
||||
|
||||
const prompt = (pending: MemberBudgetResetPending) => setState({ phase: "prompting", pending });
|
||||
|
||||
const reset = async () => {
|
||||
if (state.phase !== "prompting") return;
|
||||
const { pending } = state;
|
||||
const run = {};
|
||||
activeRun.current = run;
|
||||
setState({ phase: "resetting", pending, attempted: 0 });
|
||||
try {
|
||||
await gateway.saveTeam(pending.updateData);
|
||||
} catch (error) {
|
||||
console.error("Error updating team:", error);
|
||||
if (isCurrent(run)) setState({ phase: "prompting", pending });
|
||||
return;
|
||||
}
|
||||
await runReset(pending, 0, run);
|
||||
};
|
||||
|
||||
const retry = async () => {
|
||||
if (state.phase !== "resetFailed") return;
|
||||
const { pending, attempted } = state;
|
||||
const run = {};
|
||||
activeRun.current = run;
|
||||
setState({ phase: "resetting", pending, attempted });
|
||||
await runReset(pending, attempted, run);
|
||||
};
|
||||
|
||||
const keepCustom = async () => {
|
||||
if (state.phase !== "prompting") return;
|
||||
const { pending } = state;
|
||||
const run = {};
|
||||
activeRun.current = run;
|
||||
setState({ phase: "idle" });
|
||||
try {
|
||||
await gateway.saveTeam(pending.updateData);
|
||||
if (isCurrent(run)) toast.success("Team settings updated successfully");
|
||||
} catch (error) {
|
||||
console.error("Error updating team:", error);
|
||||
}
|
||||
if (isCurrent(run)) await gateway.refreshTeamData();
|
||||
};
|
||||
|
||||
const dismiss = useCallback(() => {
|
||||
activeRun.current = null;
|
||||
setState({ phase: "idle" });
|
||||
}, []);
|
||||
|
||||
return { state, prompt, reset, retry, keepCustom, dismiss };
|
||||
};
|
||||
Loading…
Add table
Reference in a new issue