diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json
index 664e1669834..f4d4e25859a 100644
--- a/basedpyright-code-budget.json
+++ b/basedpyright-code-budget.json
@@ -1,6 +1,6 @@
{
"reportAny": {
- "limit": 19955
+ "limit": 19949
},
"reportArgumentType": {
"limit": 2566
@@ -54,7 +54,7 @@
"limit": 0
},
"reportMissingParameterType": {
- "limit": 5663
+ "limit": 5661
},
"reportMissingTypeArgument": {
"limit": 15555
@@ -105,10 +105,10 @@
"limit": 109
},
"reportUnknownMemberType": {
- "limit": 39011
+ "limit": 39009
},
"reportUnknownParameterType": {
- "limit": 19885
+ "limit": 19883
},
"reportUnknownVariableType": {
"limit": 30569
diff --git a/litellm/__init__.py b/litellm/__init__.py
index e95b553c5d4..ee2c551481c 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -1801,6 +1801,9 @@ if TYPE_CHECKING:
from .llms.gemini.interactions.transformation import (
GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
)
+ from .llms.vertex_ai.interactions.transformation import (
+ VertexAIInteractionsConfig as VertexAIInteractionsConfig,
+ )
from .llms.openai.chat.o_series_transformation import (
OpenAIOSeriesConfig as OpenAIOSeriesConfig,
OpenAIOSeriesConfig as OpenAIO1Config,
diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py
index 89c72acc06d..c34c9eefe85 100644
--- a/litellm/_lazy_imports_registry.py
+++ b/litellm/_lazy_imports_registry.py
@@ -242,6 +242,7 @@ LLM_CONFIG_NAMES: Final = (
"OpenRouterResponsesAPIConfig",
"BedrockMantleResponsesAPIConfig",
"GoogleAIStudioInteractionsConfig",
+ "VertexAIInteractionsConfig",
"OpenAIOSeriesConfig",
"AnthropicSkillsConfig",
"BaseSkillsAPIConfig",
@@ -977,6 +978,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.gemini.interactions.transformation",
"GoogleAIStudioInteractionsConfig",
),
+ "VertexAIInteractionsConfig": (
+ ".llms.vertex_ai.interactions.transformation",
+ "VertexAIInteractionsConfig",
+ ),
"OpenAIOSeriesConfig": (
".llms.openai.chat.o_series_transformation",
"OpenAIOSeriesConfig",
diff --git a/litellm/constants.py b/litellm/constants.py
index 0a1ada3bab2..78aba30f9c0 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -750,6 +750,7 @@ openai_compatible_endpoints: Final[list] = [
"api.groq.com/openai/v1",
"https://integrate.api.nvidia.com/v1",
"api.deepseek.com/v1",
+ "api.together.ai/v1",
"api.together.xyz/v1",
"app.empower.dev/api/v1",
"https://api.friendli.ai/serverless/v1",
diff --git a/litellm/interactions/utils.py b/litellm/interactions/utils.py
index 8a1e8836894..3895a85061d 100644
--- a/litellm/interactions/utils.py
+++ b/litellm/interactions/utils.py
@@ -47,6 +47,13 @@ def get_provider_interactions_api_config(
return GoogleAIStudioInteractionsConfig()
+ if provider in (LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value):
+ from litellm.llms.vertex_ai.interactions.transformation import (
+ VertexAIInteractionsConfig,
+ )
+
+ return VertexAIInteractionsConfig()
+
return None
diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py
index e674fc37673..005e94ebe82 100644
--- a/litellm/litellm_core_utils/get_llm_provider_logic.py
+++ b/litellm/litellm_core_utils/get_llm_provider_logic.py
@@ -272,6 +272,14 @@ def get_llm_provider(
elif endpoint == "api.deepseek.com/v1":
custom_llm_provider = "deepseek"
dynamic_api_key = get_secret_str("DEEPSEEK_API_KEY")
+ elif endpoint == "api.together.ai/v1" or endpoint == "api.together.xyz/v1":
+ custom_llm_provider = "together_ai"
+ dynamic_api_key = api_key or (
+ get_secret_str("TOGETHER_API_KEY")
+ or get_secret_str("TOGETHER_AI_API_KEY")
+ or get_secret_str("TOGETHERAI_API_KEY")
+ or get_secret_str("TOGETHER_AI_TOKEN")
+ )
elif endpoint == "ollama.com":
custom_llm_provider = "ollama"
dynamic_api_key = get_secret_str("OLLAMA_API_KEY")
@@ -707,7 +715,7 @@ def _get_openai_compatible_provider_info(
dynamic_api_key,
) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key)
elif custom_llm_provider == "together_ai":
- api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.xyz/v1"
+ api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.ai/v1"
dynamic_api_key = api_key or (
get_secret_str("TOGETHER_API_KEY")
or get_secret_str("TOGETHER_AI_API_KEY")
diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py
index 826a890eca9..86cfbf70255 100644
--- a/litellm/litellm_core_utils/prompt_templates/factory.py
+++ b/litellm/litellm_core_utils/prompt_templates/factory.py
@@ -643,49 +643,6 @@ def claude_2_1_pt(
return prompt
-### TOGETHER AI
-
-
-def get_model_info(token, model):
- try:
- headers: Final = {"Authorization": f"Bearer {token}"}
- client: Final = HTTPHandler(concurrent_limit=1)
- response: Final = client.get("https://api.together.xyz/models/info", headers=headers)
- if response.status_code == 200:
- model_info: Final = response.json()
- for m in model_info:
- if m["name"].lower().strip() == model.strip():
- return m["config"].get("prompt_format", None), m["config"].get("chat_template", None)
- return None, None
- else:
- return None, None
- except Exception: # safely fail a prompt template request
- return None, None
-
-
-## OLD TOGETHER AI FLOW
-# def format_prompt_togetherai(messages, prompt_format, chat_template):
-# if prompt_format is None:
-# return default_pt(messages)
-
-# human_prompt, assistant_prompt = prompt_format.split("{prompt}")
-
-# if chat_template is not None:
-# prompt = hf_chat_template(
-# model=None, messages=messages, chat_template=chat_template
-# )
-# elif prompt_format is not None:
-# prompt = custom_prompt(
-# role_dict={},
-# messages=messages,
-# initial_prompt_value=human_prompt,
-# final_prompt_value=assistant_prompt,
-# )
-# else:
-# prompt = default_pt(messages)
-# return prompt
-
-
### IBM Granite
diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py
index 23abca7d5f2..26380ad0af8 100644
--- a/litellm/llms/anthropic/chat/transformation.py
+++ b/litellm/llms/anthropic/chat/transformation.py
@@ -1215,8 +1215,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
if reasoning_effort is None or reasoning_effort == "none":
return None
if AnthropicConfig._is_adaptive_thinking_model(model, custom_llm_provider):
+ # without display, Anthropic defaults adaptive thinking to
+ # display="omitted" and returns a blank thinking block
return AnthropicThinkingParam(
type="adaptive",
+ display="summarized",
)
elif reasoning_effort == "low":
return AnthropicThinkingParam(
@@ -2144,7 +2147,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
@staticmethod
- def _thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None:
+ def thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None:
details: Final = usage_object.get("output_tokens_details")
if not isinstance(details, Mapping):
return None
@@ -2176,7 +2179,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
reported_thinking_tokens: Final = (
iteration_thinking_tokens
if iteration_thinking_tokens is not None
- else self._thinking_tokens_from_usage(usage_object)
+ else self.thinking_tokens_from_usage(usage_object)
)
if reported_thinking_tokens is not None:
capped_reported: Final = min(max(0, reported_thinking_tokens), completion_tokens)
@@ -2199,7 +2202,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None:
per_iteration: Final = tuple(
- self._thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None
+ self.thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None
for iteration in iterations
)
reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None)
diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py
index b437e25d24b..767677cbcbf 100644
--- a/litellm/llms/bedrock/chat/converse_transformation.py
+++ b/litellm/llms/bedrock/chat/converse_transformation.py
@@ -1617,6 +1617,8 @@ class AmazonConverseConfig(BaseConfig):
}
if additional_request_params:
data["additionalModelRequestFields"] = additional_request_params
+ if "thinking" in additional_request_params:
+ data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",)
if system_content_blocks:
data["system"] = system_content_blocks
@@ -1801,6 +1803,17 @@ class AmazonConverseConfig(BaseConfig):
thinking_blocks_list.append(_redacted_block)
return thinking_blocks_list
+ @staticmethod
+ def thinking_tokens_from_additional_fields(additional_fields: object) -> int | None:
+ """Converse omits thinking tokens from its usage block; they only arrive under
+ ``additionalModelResponseFields`` when ``/usage/output_tokens_details`` is requested."""
+ if not isinstance(additional_fields, Mapping):
+ return None
+ usage: Final = additional_fields.get("usage")
+ if not isinstance(usage, Mapping):
+ return None
+ return AnthropicConfig.thinking_tokens_from_usage(usage)
+
@staticmethod
def is_converse_usage_shape(usage_object: Mapping[str, object]) -> bool:
"""Converse-family models report camelCase token counts, not Anthropic's snake_case."""
@@ -1842,6 +1855,7 @@ class AmazonConverseConfig(BaseConfig):
usage: ConverseTokenUsageBlock,
reasoning_content: str | None = None,
thinking_ran: bool = False,
+ provider_reasoning_tokens: int | None = None,
) -> Usage:
input_tokens = usage["inputTokens"]
output_tokens: Final = usage["outputTokens"]
@@ -1862,9 +1876,14 @@ class AmazonConverseConfig(BaseConfig):
cache_creation_tokens=cache_creation_input_tokens,
text_tokens=raw_input_tokens,
)
- reasoning_tokens: Final = (
+ estimated_reasoning_tokens: Final = (
token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0
)
+ reasoning_tokens: Final = (
+ min(max(0, provider_reasoning_tokens), output_tokens)
+ if provider_reasoning_tokens is not None
+ else estimated_reasoning_tokens
+ )
completion_tokens_details: Final = (
CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens,
@@ -2272,6 +2291,9 @@ class AmazonConverseConfig(BaseConfig):
completion_response["usage"],
reasoning_content=chat_completion_message.get("reasoning_content"),
thinking_ran=reasoningContentBlocks is not None,
+ provider_reasoning_tokens=self.thinking_tokens_from_additional_fields(
+ completion_response.get("additionalModelResponseFields")
+ ),
)
## HANDLE TOOL CALLS
diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py
index 35e82f1961a..fc34e403beb 100644
--- a/litellm/llms/bedrock/chat/invoke_handler.py
+++ b/litellm/llms/bedrock/chat/invoke_handler.py
@@ -331,6 +331,7 @@ class AWSEventStreamDecoder:
self.json_mode = json_mode
self._current_tool_name: str | None = None
self._thinking_ran = False
+ self._provider_reasoning_tokens: int | None = None
def check_empty_tool_call_args(self) -> bool:
"""
@@ -559,10 +560,14 @@ class AWSEventStreamDecoder:
tool_use = self._handle_converse_stop_event(content_block_index)
elif "stopReason" in chunk_data:
finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop"))
+ self._provider_reasoning_tokens = AmazonConverseConfig.thinking_tokens_from_additional_fields(
+ chunk_data.get("additionalModelResponseFields")
+ )
elif "usage" in chunk_data:
usage = converse_config.transform_usage(
chunk_data.get("usage", {}),
thinking_ran=self._thinking_ran,
+ provider_reasoning_tokens=self._provider_reasoning_tokens,
)
if thinking_blocks:
self._thinking_ran = True
diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py
index 0ce2e6f60d3..d0a3c37ffb3 100644
--- a/litellm/llms/bedrock/passthrough/transformation.py
+++ b/litellm/llms/bedrock/passthrough/transformation.py
@@ -1,4 +1,5 @@
import json
+from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Optional, cast
from httpx import Response
@@ -93,6 +94,9 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
endpoint_url,
)
+ def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None:
+ return None
+
def sign_request(
self,
headers: dict,
@@ -109,6 +113,7 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
request_data=request_data or {},
api_base=api_base,
model=model,
+ api_key=self.get_bedrock_bearer_token(optional_params),
)
def logging_non_streaming_response(
diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py
index 889361cd808..d877fbb4e09 100644
--- a/litellm/llms/bedrock_mantle/common_utils.py
+++ b/litellm/llms/bedrock_mantle/common_utils.py
@@ -13,6 +13,7 @@ global state.
"""
import re
+from collections.abc import Mapping
from typing import Final
from botocore.exceptions import (
@@ -31,30 +32,39 @@ BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1"
MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE)
+def resolve_mantle_bearer_token(api_key: str | None) -> str | None:
+ return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
+
+
+def resolve_mantle_region(params: Mapping[str, object]) -> str:
+ region: Final = params.get("aws_region_name")
+ if isinstance(region, str) and region:
+ BaseAWSLLM._validate_aws_region_name(region)
+ return region
+ api_base: Final = params.get("api_base")
+ base: Final = (api_base if isinstance(api_base, str) else None) or get_secret_str("BEDROCK_MANTLE_API_BASE")
+ if base:
+ match: Final = MANTLE_HOST_RE.match(base.rstrip("/"))
+ if match:
+ return match.group(1)
+ return (
+ get_secret_str("BEDROCK_MANTLE_REGION")
+ or get_secret_str("AWS_REGION_NAME")
+ or get_secret_str("AWS_REGION")
+ or BEDROCK_MANTLE_DEFAULT_REGION
+ )
+
+
class BedrockMantleAuthMixin:
_aws_signer: BaseAWSLLM
@staticmethod
def _resolve_bearer_token(api_key: str | None) -> str | None:
- return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
+ return resolve_mantle_bearer_token(api_key)
@staticmethod
def _resolve_region(params: dict) -> str:
- region: Final = params.get("aws_region_name")
- if region:
- BaseAWSLLM._validate_aws_region_name(region)
- return region
- base: Final = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE")
- if base:
- match: Final = MANTLE_HOST_RE.match(base.rstrip("/"))
- if match:
- return match.group(1)
- return (
- get_secret_str("BEDROCK_MANTLE_REGION")
- or get_secret_str("AWS_REGION_NAME")
- or get_secret_str("AWS_REGION")
- or BEDROCK_MANTLE_DEFAULT_REGION
- )
+ return resolve_mantle_region(params)
def sign_request(
self,
diff --git a/litellm/llms/bedrock_mantle/passthrough/transformation.py b/litellm/llms/bedrock_mantle/passthrough/transformation.py
new file mode 100644
index 00000000000..e6b831efa57
--- /dev/null
+++ b/litellm/llms/bedrock_mantle/passthrough/transformation.py
@@ -0,0 +1,71 @@
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Final, Literal, Optional
+
+from httpx import Response
+
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig
+from litellm.llms.bedrock_mantle.common_utils import (
+ MANTLE_HOST_RE,
+ resolve_mantle_bearer_token,
+ resolve_mantle_region,
+)
+from litellm.types.utils import LlmProviders
+
+if TYPE_CHECKING:
+ from litellm.types.utils import CostResponseTypes
+
+
+class BedrockMantlePassthroughConfig(BedrockPassthroughConfig):
+ """Native Bedrock runtime passthrough (InvokeModel, Converse) for deployments declared as bedrock_mantle.
+
+ The Mantle host only serves the OpenAI-compatible surface, so a Mantle api_base lends its region and the
+ request itself goes to bedrock-runtime, signed with the deployment's Bearer token or SigV4 credentials.
+ """
+
+ def _get_aws_region_name(
+ self,
+ optional_params: Mapping[str, object],
+ model: str | None = None,
+ model_id: str | None = None,
+ ) -> str:
+ return resolve_mantle_region(optional_params)
+
+ def get_runtime_endpoint(
+ self,
+ api_base: str | None,
+ aws_bedrock_runtime_endpoint: str | None,
+ aws_region_name: str,
+ endpoint_type: Literal["runtime", "agent", "agentcore"] | None = "runtime",
+ ) -> tuple[str, str]:
+ is_mantle_host: Final = api_base is not None and MANTLE_HOST_RE.match(api_base.rstrip("/")) is not None
+ return super().get_runtime_endpoint(
+ api_base=None if is_mantle_host else api_base,
+ aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
+ aws_region_name=aws_region_name,
+ endpoint_type=endpoint_type,
+ )
+
+ def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None:
+ api_key: Final = litellm_params.get("api_key")
+ return resolve_mantle_bearer_token(api_key if isinstance(api_key, str) else None)
+
+ def logging_non_streaming_response(
+ self,
+ model: str,
+ custom_llm_provider: str,
+ httpx_response: Response,
+ request_data: dict, # mutable-ok: mirrors the inherited BedrockPassthroughConfig signature
+ logging_obj: Logging,
+ endpoint: str,
+ ) -> Optional["CostResponseTypes"]:
+ is_converse: Final = "invoke" not in endpoint and "converse" in endpoint
+ shape_provider: Final = LlmProviders.BEDROCK.value if is_converse else custom_llm_provider
+ return super().logging_non_streaming_response(
+ model=model,
+ custom_llm_provider=shape_provider,
+ httpx_response=httpx_response,
+ request_data=request_data,
+ logging_obj=logging_obj,
+ endpoint=endpoint,
+ )
diff --git a/litellm/llms/together_ai/rerank/handler.py b/litellm/llms/together_ai/rerank/handler.py
index 10246451a9d..b8079e52c97 100644
--- a/litellm/llms/together_ai/rerank/handler.py
+++ b/litellm/llms/together_ai/rerank/handler.py
@@ -16,11 +16,16 @@ from litellm.llms.together_ai.rerank.transformation import TogetherAIRerankConfi
from litellm.types.rerank import RerankRequest, RerankResponse
+def _rerank_url(api_base: str) -> str:
+ return f"{api_base.rstrip('/')}/rerank"
+
+
class TogetherAIRerank(BaseLLM):
def rerank(
self,
model: str,
api_key: str,
+ api_base: str,
query: str,
documents: list[str | dict[str, Any]],
top_n: int | None = None,
@@ -46,10 +51,10 @@ class TogetherAIRerank(BaseLLM):
raise ValueError("TogetherAI does not support max_chunks_per_doc")
if _is_async:
- return self.async_rerank(request_data_dict, api_key) # Call async method
+ return self.async_rerank(request_data_dict, api_key, api_base)
response: Final = client.post(
- "https://api.together.xyz/v1/rerank",
+ _rerank_url(api_base),
headers={
"accept": "application/json",
"content-type": "application/json",
@@ -69,11 +74,12 @@ class TogetherAIRerank(BaseLLM):
self,
request_data_dict: dict[str, Any],
api_key: str,
+ api_base: str,
) -> RerankResponse:
client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.TOGETHER_AI) # Use async client
response: Final = await client.post(
- "https://api.together.xyz/v1/rerank",
+ _rerank_url(api_base),
headers={
"accept": "application/json",
"content-type": "application/json",
diff --git a/litellm/llms/vertex_ai/interactions/__init__.py b/litellm/llms/vertex_ai/interactions/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py
new file mode 100644
index 00000000000..0764a8bea62
--- /dev/null
+++ b/litellm/llms/vertex_ai/interactions/transformation.py
@@ -0,0 +1,149 @@
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Final
+
+from litellm.litellm_core_utils.url_utils import encode_url_path_segment
+from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig
+from litellm.llms.vertex_ai.common_utils import validate_vertex_location
+from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
+from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+VERTEX_INTERACTIONS_API_VERSION: Final = "v1beta1"
+VERTEX_INTERACTIONS_DEFAULT_LOCATION: Final = "global"
+
+
+@dataclass(frozen=True, slots=True)
+class VertexInteractionsTarget:
+ base_url: str
+ project_id: str
+ location: str
+
+ @property
+ def collection_url(self) -> str:
+ return (
+ f"{self.base_url}/{VERTEX_INTERACTIONS_API_VERSION}"
+ f"/projects/{self.project_id}/locations/{self.location}/interactions"
+ )
+
+ def interaction_url(self, interaction_id: str) -> str:
+ encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id")
+ return f"{self.collection_url}/{encoded_interaction_id}"
+
+
+class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig):
+ def __init__(
+ self,
+ mint_access_token: Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]] | None = None,
+ ) -> None:
+ super().__init__()
+ self._mint_access_token: Final[Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]]] = (
+ mint_access_token or self._mint_access_token_with_vertex_base
+ )
+
+ def _mint_access_token_with_vertex_base(
+ self,
+ credentials: VERTEX_CREDENTIALS_TYPES | None,
+ project_id: str | None,
+ ) -> tuple[str, str]:
+ return self._ensure_access_token(
+ credentials=credentials, project_id=project_id, custom_llm_provider="vertex_ai"
+ )
+
+ @property
+ def custom_llm_provider(self) -> LlmProviders:
+ return LlmProviders.VERTEX_AI
+
+ @property
+ def api_version(self) -> str:
+ return VERTEX_INTERACTIONS_API_VERSION
+
+ def get_default_vertex_location(self) -> str:
+ return VERTEX_INTERACTIONS_DEFAULT_LOCATION
+
+ def _mint(self, litellm_params: GenericLiteLLMParams) -> tuple[str, str]:
+ raw_params: Final = litellm_params.model_dump()
+ return self._mint_access_token(
+ self.safe_get_vertex_ai_credentials(raw_params),
+ self.safe_get_vertex_ai_project(raw_params),
+ )
+
+ def _target(self, api_base: str | None, litellm_params: GenericLiteLLMParams) -> VertexInteractionsTarget:
+ _, project_id = self._mint(litellm_params)
+ if not project_id:
+ raise ValueError(
+ "Vertex AI project is required. Set vertex_project, litellm.vertex_project, or VERTEXAI_PROJECT"
+ )
+ location: Final = validate_vertex_location(
+ self.explicit_vertex_ai_location(litellm_params.model_dump()) or VERTEX_INTERACTIONS_DEFAULT_LOCATION
+ )
+ return VertexInteractionsTarget(
+ base_url=self.get_api_base(api_base or None, location),
+ project_id=project_id,
+ location=location,
+ )
+
+ def validate_environment(
+ self,
+ headers: Mapping[str, str],
+ model: str,
+ litellm_params: GenericLiteLLMParams | None,
+ ) -> dict: # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers
+ access_token, _ = self._mint(litellm_params or GenericLiteLLMParams())
+ return { # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers
+ "Content-Type": "application/json",
+ "Authorization": f"Bearer {access_token}",
+ **headers,
+ }
+
+ def get_complete_url(
+ self,
+ api_base: str | None,
+ model: str | None,
+ agent: str | None = None,
+ litellm_params: Mapping[str, object] | None = None,
+ stream: bool | None = None,
+ ) -> str:
+ params: Final = (
+ GenericLiteLLMParams.model_validate(litellm_params) if litellm_params else GenericLiteLLMParams()
+ )
+ collection_url: Final = self._target(api_base, params).collection_url
+ return f"{collection_url}?alt=sse" if stream else collection_url
+
+ def _interaction_by_id_request(
+ self,
+ interaction_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ url_suffix: str = "",
+ ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+ target: Final = self._target(api_base or None, litellm_params)
+ return f"{target.interaction_url(interaction_id)}{url_suffix}", {} # mutable-ok: same base contract
+
+ def transform_get_interaction_request(
+ self,
+ interaction_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: Mapping[str, str],
+ ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+ return self._interaction_by_id_request(interaction_id, api_base, litellm_params)
+
+ def transform_delete_interaction_request(
+ self,
+ interaction_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: Mapping[str, str],
+ ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+ return self._interaction_by_id_request(interaction_id, api_base, litellm_params)
+
+ def transform_cancel_interaction_request(
+ self,
+ interaction_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: Mapping[str, str],
+ ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+ return self._interaction_by_id_request(interaction_id, api_base, litellm_params, url_suffix=":cancel")
diff --git a/litellm/main.py b/litellm/main.py
index 191b1514650..fcf33653cb3 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -602,7 +602,7 @@ async def acompletion(
_, custom_llm_provider, _, _ = get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
- api_base=base_url,
+ api_base=kwargs.get("api_base") or base_url,
)
fallbacks = fallbacks or litellm.model_fallbacks
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 9f953e11df1..1aa4c7cd060 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -49016,12 +49016,13 @@
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@@ -49048,12 +49049,13 @@
"output_cost_per_token": 1.32e-05,
"output_cost_per_token_above_272k_tokens": 1.98e-05,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@@ -49080,12 +49082,13 @@
"output_cost_per_token": 1.32e-06,
"output_cost_per_token_above_272k_tokens": 1.98e-06,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py
index 8a2ee2a3af8..4b30afb2f98 100644
--- a/litellm/passthrough/main.py
+++ b/litellm/passthrough/main.py
@@ -199,7 +199,7 @@ def llm_passthrough_route(
api_key=api_key,
)
- litellm_params_dict: Final = get_litellm_params(**kwargs)
+ litellm_params_dict: Final = get_litellm_params(api_key=api_key, api_base=api_base, **kwargs)
if client is None:
from litellm.llms.custom_httpx.http_handler import (
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 0840d37ffa1..628e569e1b8 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -815,6 +815,7 @@ class LiteLLMRoutes(enum.Enum):
"/team/member_add",
"/team/member_delete",
"/team/member_update",
+ "/team/{team_id}/member/{user_id}/reset_spend",
"/team/permissions_list",
"/team/permissions_update",
"/team/daily/activity",
@@ -1287,6 +1288,16 @@ class RegenerateKeyRequest(GenerateKeyRequest):
class ResetSpendRequest(LiteLLMPydanticObjectBase):
reset_to: float
+ @field_validator("reset_to", mode="before")
+ @classmethod
+ def reject_bool_reset_to(cls, v):
+ # bool is a subclass of int, so pydantic silently coerces True/False into
+ # 1.0/0.0 for a `float` field: a caller who accidentally sends a boolean
+ # would otherwise get an unintended spend reset instead of a 422.
+ if isinstance(v, bool):
+ raise ValueError("reset_to must be a number, not a boolean") # noqa: TRY004 # pydantic needs ValueError
+ return v
+
class KeyRequest(LiteLLMPydanticObjectBase):
keys: list[str] | None = None
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index e7b98b3cc7f..4af942a357e 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -87,6 +87,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
object_permission_cache_key,
tag_cache_key,
tag_registry_cache_key,
+ team_membership_auth_cache_key,
+ team_membership_reservation_cache_key,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.guardrails.tool_name_extraction import (
@@ -1967,7 +1969,7 @@ async def get_team_membership(
if user_id is None or team_id is None:
return None
- _key: Final = f"team_membership:{user_id}:{team_id}"
+ _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
# check if in cache
cached_membership_obj: Final = await user_api_key_cache.async_get_cache(
@@ -2402,6 +2404,116 @@ async def _cache_team_object(
)
+async def invalidate_team_member_spend_state(
+ user_id: str,
+ team_id: str,
+ user_api_key_cache: UserApiKeyCache,
+ new_spend: float | None = None,
+) -> None:
+ """
+ Clear every cached read path for one team member's budget so a spend
+ reset or a raised cap takes effect on the next request instead of
+ waiting on the membership cache's TTL.
+
+ Two independently-keyed cache entries hold the same LiteLLM_TeamMembership
+ row: user_api_key_auth.py's admission check writes ``{team_id}_{user_id}``,
+ while budget_reservation.py's pre-call reservation and auth_checks.py's own
+ get_team_membership() (used by _check_team_member_budget) both write
+ ``team_membership:{user_id}:{team_id}``. Both formats must be invalidated
+ explicitly; writing one does not refresh the other. All keys are also
+ broadcast (LIT-3803): each worker's own in-memory copy (membership object,
+ spend counter, or the counter's own short-TTL DB-floor marker) survives
+ eviction elsewhere until its TTL, so the handling worker alone clearing its
+ copy leaves every other worker still enforcing the pre-reset budget.
+
+ ``new_spend`` is only passed by reset_team_member_spend_fn, which knows the
+ exact post-reset value: it is SET everywhere (matching /key/{key}/reset_spend's
+ own precedent) rather than deleted, so a worker's next read reflects it
+ directly instead of re-deriving it through a DB reseed. team_member_update
+ only changes the budget cap, not the tracked spend, so it passes no
+ new_spend; the live spend counter is untouched in that case (deleting it
+ would force a reseed from the DB's own spend column, which lags the live
+ counter via periodic batch writes, briefly under-enforcing the raised cap
+ against a spend value lower than what was actually tracked) and only the
+ membership caches carrying the new cap are invalidated.
+
+ The floor marker (``spend_db_floor:``, proxy_server.py's
+ _authoritative_floor_spend) caches the pre-reset DB spend for
+ SPEND_DB_FLOOR_CACHE_TTL_SECONDS; left stale after a real reset, a request
+ landing on the pod that cached it can read that higher floor and raise the
+ counter right back above the just-reset spend. It is overwritten here with
+ the post-reset floor (not merely deleted) and _authoritative_floor_spend
+ re-checks the marker after its DB read, so a floor read already in flight
+ on this pod when the reset commits cannot clobber it with the pre-reset
+ value. Both keys are broadcast as SETs carrying new_spend, not deletes:
+ every subscriber (remote pods AND this pod's own, which receives its own
+ message) writes the post-reset value, so the self-delivered message cannot
+ erase the guard just written here.
+
+ Raises HTTPException(503) if Redis still holds the stale pre-reset counter
+ after both the SET and the fallback DELETE fail: budget checks read Redis
+ first, so returning success would leave the old value authoritative for
+ every worker despite the DB write having committed.
+ """
+ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
+ evict_and_broadcast,
+ publish_auth_cache_invalidation,
+ )
+
+ if new_spend is not None:
+ from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
+
+ spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}"
+ spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}"
+
+ spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60)
+ if spend_counter_cache.redis_cache is not None:
+ try:
+ await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60)
+ except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
+ verbose_proxy_logger.warning(
+ "Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "
+ "read reseeds from the DB rather than keeping the stale pre-reset value authoritative",
+ spend_counter_key,
+ e,
+ )
+ try:
+ await spend_counter_cache.redis_cache.async_delete_cache(key=spend_counter_key)
+ except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success
+ verbose_proxy_logger.warning(
+ "Failed to delete stale spend counter %s in Redis after a failed reset write",
+ spend_counter_key,
+ exc_info=True,
+ )
+ raise HTTPException(
+ status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+ detail={ # mutable-ok: HTTPException.detail takes a dict
+ "error": "Spend was reset in the database, but Redis is unreachable and still "
+ "holds the pre-reset counter. Retry once Redis is reachable."
+ },
+ ) from e
+
+ spend_counter_cache.in_memory_cache.set_cache(
+ key=spend_db_floor_key,
+ value=new_spend,
+ ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
+ )
+ await publish_auth_cache_invalidation(cache_key=spend_counter_key, new_value=new_spend, ttl=60)
+ await publish_auth_cache_invalidation(
+ cache_key=spend_db_floor_key,
+ new_value=new_spend,
+ ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
+ )
+
+ await evict_and_broadcast(
+ cache_keys=(
+ team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
+ team_membership_reservation_cache_key(user_id=user_id, team_id=team_id),
+ ),
+ user_api_key_cache=user_api_key_cache,
+ )
+
+
async def delete_cache_team_object(
team_id: str,
team_alias: str | None,
diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index 658d176f6a7..28d76e6799c 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -87,7 +87,10 @@ from litellm.proxy.common_utils.http_parsing_utils import (
populate_request_with_path_params,
)
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
-from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+from litellm.proxy.common_utils.user_api_key_cache import (
+ UserApiKeyCache,
+ team_membership_auth_cache_key,
+)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.utils import (
@@ -1970,8 +1973,10 @@ async def _user_api_key_auth_builder(
# Check 3. Check if user is in their team budget
if not skip_budget_checks and valid_token.team_member_spend is not None:
- if prisma_client is not None:
- _cache_key: Final = f"{valid_token.team_id}_{valid_token.user_id}"
+ _user_id: Final = valid_token.user_id
+ _team_id: Final = valid_token.team_id
+ if prisma_client is not None and _user_id is not None and _team_id is not None:
+ _cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id)
team_member_info = await user_api_key_cache.async_get_cache(
key=_cache_key,
@@ -1979,25 +1984,21 @@ async def _user_api_key_auth_builder(
)
if team_member_info is None:
# read from DB
- _user_id: Final = valid_token.user_id
- _team_id: Final = valid_token.team_id
-
- if _user_id is not None and _team_id is not None:
- _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first(
- where={
- "user_id": _user_id,
- "team_id": _team_id,
- },
- include={"litellm_budget_table": True},
+ _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first(
+ where={
+ "user_id": _user_id,
+ "team_id": _team_id,
+ },
+ include={"litellm_budget_table": True},
+ )
+ if _db_member is not None:
+ team_member_info = LiteLLM_TeamMembership(**_db_member.dict())
+ await user_api_key_cache.async_set_cache(
+ key=_cache_key,
+ value=team_member_info,
+ model_type=LiteLLM_TeamMembership,
+ ttl=5,
)
- if _db_member is not None:
- team_member_info = LiteLLM_TeamMembership(**_db_member.dict())
- await user_api_key_cache.async_set_cache(
- key=_cache_key,
- value=team_member_info,
- model_type=LiteLLM_TeamMembership,
- ttl=5,
- )
if team_member_info is not None and team_member_info.litellm_budget_table is not None:
team_member_budget: Final = team_member_info.litellm_budget_table.max_budget
@@ -2013,11 +2014,16 @@ async def _user_api_key_auth_builder(
max_budget=team_member_budget,
)
if team_member_spend > team_member_budget:
+ _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}"
raise litellm.BudgetExceededError(
current_cost=team_member_spend,
max_budget=team_member_budget,
+ message=(
+ f"Budget has been exceeded! TeamMember={_entity_id} "
+ f"Current cost: {team_member_spend}, Max budget: {team_member_budget}"
+ ),
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
- entity_id=f"{valid_token.user_id}:{valid_token.team_id}",
+ entity_id=_entity_id,
)
# Check 3. If token is expired
diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
index acdc9728390..fb2ca6372c0 100644
--- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
+++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
@@ -12,6 +12,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
)
if TYPE_CHECKING:
+ from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@@ -30,15 +31,24 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str:
@dataclass(frozen=True, slots=True)
class _CacheInvalidationMessage:
cache_key: str
+ new_value: float | None = None
+ ttl: float | None = None
-def _cache_invalidation_message_json(cache_key: str) -> str:
- return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key)))
+def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str:
+ message: Final = asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl))
+ return json.dumps({field: value for field, value in message.items() if value is not None})
-def _cache_key_from_message_data(data: object) -> str | None:
+def _finite_number_or_none(value: object) -> float | None:
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ return None
+ return float(value)
+
+
+def _message_from_data(data: object) -> _CacheInvalidationMessage | None:
if isinstance(data, bytes):
- data = data.decode("utf-8", errors="replace")
+ data = data.decode("utf-8", errors="replace") # rebind-ok: normalizing the wire payload to str
if not isinstance(data, str):
return None
try:
@@ -48,14 +58,28 @@ def _cache_key_from_message_data(data: object) -> str | None:
if not isinstance(parsed, dict):
return None
cache_key: Final = parsed.get("cache_key")
- return cache_key if isinstance(cache_key, str) else None
+ if not isinstance(cache_key, str):
+ return None
+ return _CacheInvalidationMessage(
+ cache_key=cache_key,
+ new_value=_finite_number_or_none(parsed.get("new_value")),
+ ttl=_finite_number_or_none(parsed.get("ttl")),
+ )
-async def publish_auth_cache_invalidation(cache_key: str) -> None:
+async def publish_auth_cache_invalidation(
+ cache_key: str, new_value: float | None = None, ttl: float | None = None
+) -> None:
"""
Best-effort broadcast so every worker drops its local in-memory copy of a
mutated management object; without this, only the handling worker and Redis
are evicted and other workers keep serving the stale object until its TTL.
+
+ Passing ``new_value`` broadcasts a SET instead of a delete: every subscriber
+ (including the publishing worker's own, which receives its own message)
+ writes the value into its additional in-memory caches rather than deleting
+ the key. A spend reset uses this so the handler's self-delivered message
+ cannot erase the freshly-written post-reset counter or floor marker.
"""
redis_cache: Final = coordination_redis_cache()
if redis_cache is None:
@@ -68,7 +92,10 @@ async def publish_auth_cache_invalidation(cache_key: str) -> None:
cache_key,
)
return
- await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key))
+ await client.publish(
+ auth_cache_invalidation_channel(redis_cache),
+ _cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl),
+ )
except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors
verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e)
@@ -95,15 +122,17 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us
class AuthCacheInvalidationSubscriber:
- __slots__ = ("_redis_cache", "_task", "_user_api_key_cache")
+ __slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache")
def __init__(
self,
redis_cache: "RedisCache",
user_api_key_cache: "UserApiKeyCache",
+ additional_in_memory_caches: Sequence["InMemoryCache"] = (),
) -> None:
self._redis_cache = redis_cache
self._user_api_key_cache = user_api_key_cache
+ self._additional_in_memory_caches = tuple(additional_in_memory_caches)
self._task: asyncio.Task[None] | None = None
def start(self) -> None:
@@ -160,12 +189,18 @@ class AuthCacheInvalidationSubscriber:
def _apply_message(self, message: object) -> None:
data: Final = message.get("data") if isinstance(message, dict) else None
- cache_key: Final = _cache_key_from_message_data(data)
- if cache_key is None:
+ parsed: Final = _message_from_data(data)
+ if parsed is None:
+ return
+ if parsed.new_value is not None:
+ for additional_cache in self._additional_in_memory_caches:
+ additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl)
return
in_memory_cache: Final = self._user_api_key_cache.in_memory_cache
if in_memory_cache is not None:
- in_memory_cache.delete_cache(cache_key)
+ in_memory_cache.delete_cache(parsed.cache_key)
+ for additional_cache in self._additional_in_memory_caches:
+ additional_cache.delete_cache(parsed.cache_key)
@staticmethod
async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None:
diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py
index 93d51bdd461..b8df0105b7b 100644
--- a/litellm/proxy/common_utils/user_api_key_cache.py
+++ b/litellm/proxy/common_utils/user_api_key_cache.py
@@ -200,6 +200,21 @@ def end_user_restricted_registry_cache_key() -> str:
return "end_user_restricted_registry"
+def team_membership_auth_cache_key(team_id: str, user_id: str) -> str:
+ """Cache key one team member's ``LiteLLM_TeamMembership`` row is stored under for the admission check."""
+ return f"{team_id}_{user_id}"
+
+
+def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str:
+ """Cache key the pre-call budget reservation stores the same ``LiteLLM_TeamMembership`` row under.
+
+ Deliberately not unified with ``team_membership_auth_cache_key``: the two readers wrote independent
+ keys before this file existed, so a fix that invalidates one must invalidate both explicitly rather
+ than assume a single write is visible to both.
+ """
+ return f"team_membership:{user_id}:{team_id}"
+
+
def get_management_object_ttl(cache: DualCache) -> float:
"""
In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...).
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index 01254d5c064..49461d7841d 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -16,7 +16,7 @@ import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
-from typing import Annotated, Final, NamedTuple, Protocol, TypedDict, TypeVar, cast
+from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypedDict, TypeVar, cast
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@@ -56,6 +56,7 @@ from litellm.proxy._types import (
PatchTeamRequest,
ProxyErrorTypes,
ProxyException,
+ ResetSpendRequest,
SpecialManagementEndpointEnums,
SpecialModelNames,
SpecialProxyStrings,
@@ -84,6 +85,7 @@ from litellm.proxy.auth.auth_checks import (
get_team_membership,
get_team_object,
get_user_object,
+ invalidate_team_member_spend_state,
)
from litellm.proxy.auth.auth_utils import (
enforce_batch_enqueued_token_limit_is_admin_only,
@@ -3392,7 +3394,7 @@ async def team_member_update(
Update team member budgets and team member role
"""
- from litellm.proxy.proxy_server import premium_user, prisma_client
+ from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
@@ -3491,6 +3493,12 @@ async def team_member_update(
budget_patch=budget_patch,
team_default_budget_id=team_default_budget_id,
)
+ if budget_patch:
+ await invalidate_team_member_spend_state(
+ user_id=received_user_id,
+ team_id=data.team_id,
+ user_api_key_cache=user_api_key_cache,
+ )
### update team member role
if data.role is not None:
@@ -3527,6 +3535,125 @@ async def team_member_update(
)
+def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAuth) -> None:
+ """
+ _verify_team_access authorizes a team admin (or org admin) over their own
+ team, with no check that the target user_id differs from the caller. Left
+ unchecked, that admin could target their own LiteLLM_TeamMembership row and
+ repeatedly reset it to 0 right before it crosses their per-member cap,
+ consuming the shared team budget without the configured limit ever binding.
+ Only a proxy admin may reset an admin's own spend.
+ """
+ if user_id == user_api_key_dict.user_id and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
+ _raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.")
+
+
+def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn:
+ detail: Final = {"error": message} # mutable-ok: HTTPException.detail takes a dict
+ raise HTTPException(status_code=status_code, detail=detail)
+
+
+def _validate_team_member_reset_spend_value(
+ reset_to: object,
+ membership: LiteLLM_TeamMembership,
+) -> float:
+ if not isinstance(reset_to, (int, float)):
+ _raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a float")
+
+ reset_to_float: Final = float(reset_to)
+ if not math.isfinite(reset_to_float) or reset_to_float < 0:
+ _raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a finite number >= 0")
+
+ current_spend: Final = membership.spend or 0.0
+ if reset_to_float > current_spend:
+ _raise_reset_spend_error(
+ status.HTTP_400_BAD_REQUEST,
+ f"reset_to ({reset_to_float}) must be <= current spend ({current_spend})",
+ )
+
+ max_budget: Final = membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None
+ if max_budget is not None and reset_to_float > max_budget:
+ _raise_reset_spend_error(
+ status.HTTP_400_BAD_REQUEST,
+ f"reset_to ({reset_to_float}) must be <= budget ({max_budget})",
+ )
+
+ return reset_to_float
+
+
+@router.post(
+ "/team/{team_id}/member/{user_id}/reset_spend",
+ tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence
+ dependencies=(Depends(user_api_key_auth),),
+)
+@management_endpoint_wrapper
+async def reset_team_member_spend_fn(
+ team_id: str,
+ user_id: str,
+ data: ResetSpendRequest,
+ user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+):
+ """
+ Reset a team member's tracked spend against their per-member budget.
+
+ A member's spend is tracked separately from both their own personal
+ budget and the team's own budget (LiteLLM_TeamMembership.spend), so
+ neither /user/update nor /team/update can clear it: this is the only
+ endpoint that does. The cross-pod spend counter and cached membership
+ reads are invalidated so the reset takes effect on the member's next
+ request rather than waiting on the membership cache's TTL.
+ """
+ from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
+
+ if prisma_client is None:
+ _raise_reset_spend_error(status.HTTP_500_INTERNAL_SERVER_ERROR, "DB not connected. prisma_client is None")
+
+ team_obj: Final = await get_team_object(
+ team_id=team_id,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=proxy_logging_obj,
+ check_db_only=True,
+ )
+ await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict)
+ _check_not_resetting_own_spend(user_id=user_id, user_api_key_dict=user_api_key_dict)
+
+ membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument
+ "user_id_team_id": {"user_id": user_id, "team_id": team_id} # mutable-ok: same prisma where= argument
+ }
+ _membership_row: Final = await _team_membership_db(prisma_client).find_unique(
+ where=membership_where,
+ include={"litellm_budget_table": True}, # mutable-ok: prisma client requires a plain dict include= argument
+ )
+ if _membership_row is None:
+ _raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.")
+ membership: Final = LiteLLM_TeamMembership.model_validate(_membership_row.model_dump())
+
+ current_spend: Final = membership.spend or 0.0
+ reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership)
+
+ await _team_membership_db(prisma_client).update(
+ where=membership_where,
+ data={"spend": reset_to}, # mutable-ok: prisma client requires a plain dict data= argument
+ )
+
+ await invalidate_team_member_spend_state(
+ user_id=user_id,
+ team_id=team_id,
+ user_api_key_cache=user_api_key_cache,
+ new_spend=reset_to,
+ )
+
+ return { # mutable-ok: matches this router's established untyped-response-dict convention
+ "team_id": team_id,
+ "user_id": user_id,
+ "spend": reset_to,
+ "previous_spend": current_spend,
+ "max_budget": membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None,
+ }
+
+
def _create_results_from_response(
members: list[Member],
response: TeamAddMemberResponse,
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 7dced4e26b6..0abcdeaf3f6 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -2555,6 +2555,12 @@ async def _authoritative_floor_spend(
if db_spend is None:
return None
+ # a spend reset that committed during the DB read above wrote the post-reset
+ # floor to the marker; keep it over this read's now-stale pre-commit value
+ rechecked: Final = spend_counter_cache.in_memory_cache.get_cache(key=marker_key)
+ if rechecked is not None:
+ return float(rechecked)
+
spend_counter_cache.in_memory_cache.set_cache(
key=marker_key,
value=db_spend,
@@ -6798,6 +6804,7 @@ class ProxyConfig:
subscriber: Final = AuthCacheInvalidationSubscriber(
redis_cache=redis_cache,
user_api_key_cache=user_api_key_cache,
+ additional_in_memory_caches=(spend_counter_cache.in_memory_cache,),
)
self.auth_cache_invalidation_subscriber = subscriber
subscriber.start()
diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py
index ce6c9330620..149f9b960a1 100644
--- a/litellm/proxy/spend_tracking/budget_reservation.py
+++ b/litellm/proxy/spend_tracking/budget_reservation.py
@@ -25,7 +25,11 @@ from litellm.proxy._types import (
from litellm.proxy.auth.auth_utils import get_model_from_request
from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded
from litellm.proxy.auth.route_checks import RouteChecks
-from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key, tag_cache_key
+from litellm.proxy.common_utils.user_api_key_cache import (
+ end_user_cache_key,
+ tag_cache_key,
+ team_membership_reservation_cache_key,
+)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@@ -546,7 +550,9 @@ async def _get_team_member_budget_counter(
if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None:
return None
- membership_cache_key: Final = f"team_membership:{valid_token.user_id}:{team_object.team_id}"
+ membership_cache_key: Final = team_membership_reservation_cache_key(
+ user_id=valid_token.user_id, team_id=team_object.team_id
+ )
cached_team_membership: Final = await user_api_key_cache.async_get_cache(key=membership_cache_key)
team_membership: LiteLLM_TeamMembership | None = None
if isinstance(cached_team_membership, LiteLLM_TeamMembership):
diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py
index 15a6f18a6bb..c8f7842aebf 100644
--- a/litellm/rerank_api/main.py
+++ b/litellm/rerank_api/main.py
@@ -277,6 +277,8 @@ def rerank(
if api_key is None:
raise ValueError("TogetherAI API key is required, please set 'TOGETHERAI_API_KEY' in your environment")
+ api_base = dynamic_api_base or optional_params.api_base or litellm.api_base or "https://api.together.ai/v1"
+
response = together_rerank.rerank(
model=model,
query=query,
@@ -286,6 +288,7 @@ def rerank(
return_documents=return_documents,
max_chunks_per_doc=max_chunks_per_doc,
api_key=api_key,
+ api_base=api_base,
_is_async=_is_async,
)
elif _custom_llm_provider == litellm.LlmProviders.JINA_AI:
diff --git a/litellm/router.py b/litellm/router.py
index 324ecdba444..c658bc441ba 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -8346,6 +8346,7 @@ class Router:
) = litellm.get_llm_provider(
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None),
+ api_base=deployment.litellm_params.api_base,
)
# done reading model["litellm_params"]
# Check if provider is supported: either in enum or JSON-configured
diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py
index cc6eccbf3e0..d3b0f334163 100644
--- a/litellm/types/llms/anthropic.py
+++ b/litellm/types/llms/anthropic.py
@@ -685,6 +685,7 @@ ANTHROPIC_API_ONLY_HEADERS: Final = { # fails if calling anthropic on vertex ai
class AnthropicThinkingParam(TypedDict, total=False):
type: ReadOnly[Literal["enabled", "adaptive", "disabled"]]
budget_tokens: int
+ display: ReadOnly[Literal["summarized", "omitted"]]
class ANTHROPIC_HOSTED_TOOLS(str, Enum):
diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py
index 5665aa3277a..6ae2e31fe60 100644
--- a/litellm/types/llms/bedrock.py
+++ b/litellm/types/llms/bedrock.py
@@ -1,4 +1,5 @@
import json
+from collections.abc import Sequence
from enum import Enum
from typing import TYPE_CHECKING, Any, Final, Literal
@@ -396,7 +397,7 @@ class OutputConfigBlock(TypedDict, total=False):
class CommonRequestObject(TypedDict, total=False): # common request object across sync + async flows
additionalModelRequestFields: dict
- additionalModelResponseFieldPaths: list[str]
+ additionalModelResponseFieldPaths: Sequence[str]
inferenceConfig: InferenceConfig
system: list[SystemContentBlock]
toolConfig: ToolConfigBlock
diff --git a/litellm/utils.py b/litellm/utils.py
index 012e8785321..5cbd0519032 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -8610,6 +8610,12 @@ class ProviderConfigManager:
)
return BedrockPassthroughConfig()
+ elif LlmProviders.BEDROCK_MANTLE == provider:
+ from litellm.llms.bedrock_mantle.passthrough.transformation import (
+ BedrockMantlePassthroughConfig,
+ )
+
+ return BedrockMantlePassthroughConfig()
elif LlmProviders.VLLM == provider or LlmProviders.HOSTED_VLLM == provider:
from litellm.llms.vllm.passthrough.transformation import (
VLLMPassthroughConfig,
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 9f953e11df1..1aa4c7cd060 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -49016,12 +49016,13 @@
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@@ -49048,12 +49049,13 @@
"output_cost_per_token": 1.32e-05,
"output_cost_per_token_above_272k_tokens": 1.98e-05,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@@ -49080,12 +49082,13 @@
"output_cost_per_token": 1.32e-06,
"output_cost_per_token_above_272k_tokens": 1.98e-06,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json
index 03318718fb5..1ca152985f9 100644
--- a/ruff-strict-budget.json
+++ b/ruff-strict-budget.json
@@ -1,6 +1,6 @@
{
"ANN001": {
- "limit": 3018
+ "limit": 3016
},
"ANN002": {
"limit": 71
@@ -9,7 +9,7 @@
"limit": 827
},
"ANN201": {
- "limit": 2016
+ "limit": 2015
},
"ANN202": {
"limit": 852
@@ -57,7 +57,7 @@
"limit": 3
},
"BLE001": {
- "limit": 2919
+ "limit": 2918
},
"C401": {
"limit": 8
diff --git a/tests/mcp_tests/conftest.py b/tests/mcp_tests/conftest.py
index d1dc3ec7216..5823893afc0 100644
--- a/tests/mcp_tests/conftest.py
+++ b/tests/mcp_tests/conftest.py
@@ -7,6 +7,7 @@ import pytest
import litellm
import asyncio
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
@pytest.fixture(scope="session")
@@ -38,6 +39,8 @@ def setup_and_teardown():
yield
# Teardown code (executes after the yield point)
+ # LoggingWorker carries still-queued coroutines onto the next test's loop, where they'd log into that test's callbacks
+ asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
loop.close() # Close the loop created earlier
asyncio.set_event_loop(None) # Remove the reference to the loop
diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py
index 1903f29001f..fc9f675f837 100644
--- a/tests/mcp_tests/test_mcp_logging.py
+++ b/tests/mcp_tests/test_mcp_logging.py
@@ -1,6 +1,9 @@
import os
import pytest
import asyncio
+import subprocess
+import sys
+from pathlib import Path
from typing import Optional
from unittest.mock import AsyncMock, patch
@@ -458,3 +461,45 @@ async def test_mcp_tool_call_hook():
logged_standard_logging_payload is not None
), "Standard logging payload should not be None"
assert logged_standard_logging_payload["response_cost"] == 1.42
+
+
+_QUEUED_LOGGING_OUTLIVES_TEST = '''
+import time
+
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+ran_at = []
+
+
+async def _record_run():
+ ran_at.append(time.monotonic())
+
+
+async def test_1_leaves_logging_queued_behind_a_stopped_worker():
+ GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run())
+ await GLOBAL_LOGGING_WORKER.stop()
+ assert ran_at == []
+
+
+async def test_2_starts_after_the_previous_tests_logging_ran():
+ started_at = time.monotonic()
+ GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run())
+ await GLOBAL_LOGGING_WORKER.flush()
+ assert [t < started_at for t in ran_at] == [True, False]
+'''
+
+
+def test_logging_queued_by_one_test_is_drained_before_the_next(tmp_path: Path):
+ """Regression: a logging coroutine queued by one test must not run inside a later test (it would log into that
+ test's callbacks, which is how test_mcp_tool_call_hook captured a gpt-4o-mini payload under xdist)."""
+ (tmp_path / "conftest.py").write_text((Path(__file__).parent / "conftest.py").read_text())
+ (tmp_path / "pyproject.toml").write_text('[tool.pytest.ini_options]\nasyncio_mode = "auto"\n')
+ (tmp_path / "test_queued_logging.py").write_text(_QUEUED_LOGGING_OUTLIVES_TEST)
+ result = subprocess.run(
+ [sys.executable, "-m", "pytest", "-q", "-p", "no:cacheprovider", "test_queued_logging.py"],
+ cwd=tmp_path,
+ capture_output=True,
+ text=True,
+ timeout=120,
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
diff --git a/tests/proxy_behavior/management/test_team_member_reset_spend.py b/tests/proxy_behavior/management/test_team_member_reset_spend.py
new file mode 100644
index 00000000000..ec2c78139fe
--- /dev/null
+++ b/tests/proxy_behavior/management/test_team_member_reset_spend.py
@@ -0,0 +1,152 @@
+import uuid
+
+import pytest
+
+from .actors import Actor
+from .conftest import create_scratch_team
+
+pytestmark = pytest.mark.asyncio(loop_scope="session")
+
+_SEED_SPEND = 5.0
+_RESET_TO = 2.0
+
+
+# POST /team/{team_id}/member/{user_id}/reset_spend. The handler gate is
+# _verify_team_access (proxy admin / team admin of this team / org admin of
+# the team's org) — the same gate /team/member_update uses, so this mirrors
+# that file's matrix exactly.
+_MATRIX = [
+ ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200),
+ ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200),
+ ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200),
+ ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403),
+ ("alpha/owner", Actor.OWNER, "alpha", 403),
+ ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403),
+ ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403),
+ ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403),
+ ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403),
+ ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200),
+ ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403),
+ ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403),
+ ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200),
+]
+
+
+async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None:
+ if shape == "alpha":
+ await create_scratch_team(
+ prisma,
+ team_id,
+ organization_id=world.org_a_id,
+ admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
+ )
+ elif shape == "beta":
+ await create_scratch_team(prisma, team_id, organization_id=world.org_b_id)
+ else: # pragma: no cover - guard
+ pytest.fail(f"unknown shape={shape}")
+ await prisma.db.litellm_teammembership.create(
+ data={"user_id": member_id, "team_id": team_id, "spend": _SEED_SPEND}
+ )
+
+
+@pytest.mark.parametrize(
+ "actor,shape,expected_status",
+ [(a, sh, s) for (_id, a, sh, s) in _MATRIX],
+ ids=[s[0] for s in _MATRIX],
+)
+async def test_team_member_reset_spend_authz_matrix(
+ actor: Actor,
+ shape: str,
+ expected_status: int,
+ proxy_client,
+ prisma,
+ scratch,
+ world,
+):
+ member_id = scratch.tag("member")
+ await _seed_target(prisma, world, shape, scratch.prefix, member_id)
+ caller = world.keys[actor]
+
+ resp = await proxy_client.post(
+ f"/team/{scratch.prefix}/member/{member_id}/reset_spend",
+ headers={"Authorization": f"Bearer {caller.cleartext}"},
+ json={"reset_to": _RESET_TO},
+ )
+ assert (
+ resp.status_code == expected_status
+ ), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
+
+ row = await prisma.db.litellm_teammembership.find_unique(
+ where={"user_id_team_id": {"user_id": member_id, "team_id": scratch.prefix}}
+ )
+ assert row is not None
+ if expected_status == 200:
+ assert row.spend == _RESET_TO
+ else:
+ assert row.spend == _SEED_SPEND, "denied but spend reset"
+
+
+async def test_team_member_reset_spend_missing_team_is_404(proxy_client, world):
+ resp = await proxy_client.post(
+ f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_spend",
+ headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
+ json={"reset_to": 0.0},
+ )
+ assert resp.status_code == 404, resp.text
+
+
+async def test_team_member_reset_spend_missing_membership_is_404(
+ proxy_client, prisma, scratch, world
+):
+ """A well-formed team but a user_id with no LiteLLM_TeamMembership row is 404."""
+ await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
+ resp = await proxy_client.post(
+ f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_spend",
+ headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
+ json={"reset_to": 0.0},
+ )
+ assert resp.status_code == 404, resp.text
+
+
+async def test_team_member_reset_spend_above_current_spend_is_400(
+ proxy_client, prisma, scratch, world
+):
+ member_id = scratch.tag("member")
+ await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
+ await prisma.db.litellm_teammembership.create(
+ data={"user_id": member_id, "team_id": scratch.prefix, "spend": 1.0}
+ )
+ resp = await proxy_client.post(
+ f"/team/{scratch.prefix}/member/{member_id}/reset_spend",
+ headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
+ json={"reset_to": 5.0},
+ )
+ assert resp.status_code == 400, resp.text
+
+
+async def test_team_member_reset_spend_team_admin_cannot_reset_own_spend(
+ proxy_client, prisma, scratch, world
+):
+ """A team admin targeting their own LiteLLM_TeamMembership row is 403: unchecked, an
+ admin could repeatedly zero their own spend right before it crosses their per-member
+ cap, consuming the shared team budget without the configured limit ever binding."""
+ team_admin = world.keys[Actor.TEAM_ADMIN]
+ await create_scratch_team(
+ prisma,
+ scratch.prefix,
+ organization_id=world.org_a_id,
+ admin_user_ids=[team_admin.user_id],
+ )
+ await prisma.db.litellm_teammembership.create(
+ data={"user_id": team_admin.user_id, "team_id": scratch.prefix, "spend": _SEED_SPEND}
+ )
+ resp = await proxy_client.post(
+ f"/team/{scratch.prefix}/member/{team_admin.user_id}/reset_spend",
+ headers={"Authorization": f"Bearer {team_admin.cleartext}"},
+ json={"reset_to": 0.0},
+ )
+ assert resp.status_code == 403, resp.text
+ row = await prisma.db.litellm_teammembership.find_unique(
+ where={"user_id_team_id": {"user_id": team_admin.user_id, "team_id": scratch.prefix}}
+ )
+ assert row is not None and row.spend == _SEED_SPEND, "denied but spend reset"
diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
index b358e6fa28c..6ca48ce63b8 100644
--- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
+++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
@@ -2,7 +2,7 @@ import datetime
import json
import os
import unittest
-from typing import TYPE_CHECKING, List, Literal, Optional, Tuple
+from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple
from unittest.mock import ANY, MagicMock, Mock, patch
import httpx
@@ -2444,6 +2444,32 @@ def test_map_optional_params_preserves_reasoning_summary():
assert responses_api_request["reasoning"]["summary"] == "detailed"
+@pytest.mark.parametrize("reasoning_effort", ["max", "high"])
+def test_transform_request_bedrock_mantle_tools_keeps_reasoning_effort(monkeypatch, reasoning_effort):
+ """Regression for reasoning_effort=max being dropped on the chat -> Responses bridge (issue #38084)."""
+ from litellm.completion_extras.litellm_responses_transformation.transformation import (
+ LiteLLMResponsesTransformationHandler,
+ )
+
+ monkeypatch.setattr(litellm, "reasoning_auto_summary", False)
+ monkeypatch.delenv("LITELLM_REASONING_AUTO_SUMMARY", raising=False)
+ handler: Final = LiteLLMResponsesTransformationHandler()
+
+ result: Final = handler.transform_request(
+ model="openai.gpt-5.6-sol",
+ messages=[{"role": "user", "content": "Say pong"}],
+ optional_params={
+ "reasoning_effort": reasoning_effort,
+ "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}],
+ },
+ litellm_params={"custom_llm_provider": "bedrock_mantle"},
+ headers={},
+ litellm_logging_obj=Mock(),
+ )
+
+ assert result["reasoning"] == {"effort": reasoning_effort}
+
+
def test_map_optional_params_tool_choice_chat_nested_to_responses_api():
"""Chat tool_choice must become Responses ToolChoiceFunction (top-level name)."""
from litellm.completion_extras.litellm_responses_transformation.transformation import (
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
index c8c36032793..6f513ce1bd4 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
@@ -478,10 +478,10 @@ def test_generic_cost_per_token_minimax_m3_above_512k_tokens(_local_model_cost_m
],
)
def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(_local_model_cost_map, model):
- """Bedrock GPT-5.6 supports a 1M context window, billed at the long-context rates above 272K."""
+ """Bedrock GPT-5.6 enforces a 1,050,000-token context window, billed at the long-context rates above 272K."""
model_cost_map = litellm.model_cost[model]
- assert model_cost_map["max_input_tokens"] == 1000000
+ assert model_cost_map["max_input_tokens"] == 1050000
cached_tokens = 100000
completion_tokens = 1000
diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
index bda7ab4afc6..6cacd119030 100644
--- a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
+++ b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
@@ -133,3 +133,54 @@ class TestGetLlmProviderRejectsAttackerSmuggledApiBase:
assert provider == "groq"
assert dynamic_api_key == "server-real-groq-key"
+
+
+class TestTogetherApiBaseResolvesProvider:
+ """
+ Regression for the Together host migration: both the current
+ ``api.together.ai`` host and the legacy ``api.together.xyz`` host must
+ resolve to ``together_ai`` when passed as ``api_base``. Before the fix
+ the endpoint list carried the legacy host but the provider-mapping
+ chain had no branch for it, so the match fell through with a None
+ provider and the deployment failed with "LLM Provider NOT provided".
+ """
+
+ @pytest.mark.parametrize(
+ "api_base",
+ [
+ "https://api.together.ai/v1",
+ "https://api.together.xyz/v1",
+ ],
+ )
+ def test_together_api_base_resolves_to_together_ai(self, api_base, monkeypatch):
+ monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env")
+
+ model, provider, dynamic_api_key, returned_api_base = get_llm_provider(
+ model="some-model",
+ api_base=api_base,
+ )
+
+ assert provider == "together_ai"
+ assert dynamic_api_key == "together-key-from-env"
+ assert returned_api_base == api_base
+ assert model == "some-model"
+
+ def test_explicit_api_key_beats_together_env_key(self, monkeypatch):
+ monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env")
+
+ _, provider, dynamic_api_key, _ = get_llm_provider(
+ model="some-model",
+ api_base="https://api.together.ai/v1",
+ api_key="explicit-caller-key",
+ )
+
+ assert provider == "together_ai"
+ assert dynamic_api_key == "explicit-caller-key"
+
+ def test_together_default_api_base_is_together_ai(self, monkeypatch):
+ monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False)
+
+ _, provider, _, api_base = get_llm_provider(model="together_ai/some-model")
+
+ assert provider == "together_ai"
+ assert api_base == "https://api.together.ai/v1"
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py
index 12ab536ed45..c9170efd18a 100644
--- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py
@@ -42,7 +42,7 @@ def test_reasoning_effort_maps_to_output_config_for_adaptive_model(
)
assert "reasoning_effort" not in result
- assert result.get("thinking") == {"type": "adaptive"}
+ assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert result.get("output_config") == {"effort": expected_effort}
diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py
index ef74249ca8e..288817dff07 100644
--- a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py
+++ b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py
@@ -5,6 +5,8 @@ Verifies that reasoning_effort=None returns None for all models,
including Claude Opus 4.6.
"""
+import pytest
+
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
@@ -35,6 +37,16 @@ class TestMapReasoningEffort:
)
assert result["type"] == "adaptive"
+ @pytest.mark.parametrize("effort", ["low", "medium", "high"])
+ def test_adaptive_mapping_requests_summarized_display(self, effort):
+ """Regression LIT-5714: adaptive thinking without ``display`` makes Anthropic
+ return a blank thinking block, so reasoning_effort callers always got
+ ``reasoning_content: ""``."""
+ result = AnthropicConfig._map_reasoning_effort(
+ reasoning_effort=effort, model="claude-opus-4-6", custom_llm_provider="anthropic"
+ )
+ assert result["display"] == "summarized"
+
def test_other_model_low_returns_enabled_with_budget(self):
result = AnthropicConfig._map_reasoning_effort(
reasoning_effort="low", model="claude-4-sonnet-20250514", custom_llm_provider="anthropic"
diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py
index 53a432427d3..326edde743d 100644
--- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py
+++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py
@@ -341,7 +341,7 @@ def test_messages_thinking_shape_follows_exact_azure_entry_flag(local_model_cost
)
result = transform()
- assert result.get("thinking") == {"type": "adaptive"}
+ assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert result.get("output_config") == {"effort": "medium"}
monkeypatch.setitem(
diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
index 604f3414775..4d2c077b548 100644
--- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
+++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
@@ -366,6 +366,96 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model):
assert additional.get("output_config") == {"effort": "high"}
+def test_reasoning_effort_requests_summarized_display_converse():
+ """Regression LIT-5714: adaptive thinking synthesized from reasoning_effort must
+ request the summarized display, otherwise the provider returns a blank thinking
+ block and reasoning_content is always empty."""
+ config = AmazonConverseConfig()
+
+ optional_params = config.map_openai_params(
+ non_default_params={"reasoning_effort": "high"},
+ optional_params={},
+ model="bedrock/converse/us.anthropic.claude-opus-4-7",
+ drop_params=False,
+ )
+
+ assert optional_params["thinking"]["type"] == "adaptive"
+ assert optional_params["thinking"]["display"] == "summarized"
+
+
+def test_thinking_request_adds_output_tokens_details_response_path():
+ """Regression LIT-5714: the Converse usage block has no thinking-token field, so
+ thinking requests must ask for ``/usage/output_tokens_details`` via
+ ``additionalModelResponseFieldPaths``."""
+ config = AmazonConverseConfig()
+
+ result = config._transform_request(
+ model="bedrock/converse/us.anthropic.claude-opus-4-7",
+ messages=[{"role": "user", "content": "hi"}],
+ optional_params={
+ "maxTokens": 256,
+ "thinking": {"type": "adaptive", "display": "summarized"},
+ "output_config": {"effort": "high"},
+ },
+ litellm_params={},
+ headers={},
+ )
+
+ assert result["additionalModelResponseFieldPaths"] == ("/usage/output_tokens_details",)
+
+
+def test_request_without_thinking_omits_response_field_paths():
+ config = AmazonConverseConfig()
+
+ result = config._transform_request(
+ model="bedrock/converse/us.anthropic.claude-opus-4-7",
+ messages=[{"role": "user", "content": "hi"}],
+ optional_params={"maxTokens": 256},
+ litellm_params={},
+ headers={},
+ )
+
+ assert "additionalModelResponseFieldPaths" not in result
+
+
+def test_transform_usage_prefers_provider_reasoning_tokens():
+ """Regression LIT-5714: provider-reported thinking tokens must win over the
+ token_counter estimate derived from visible reasoning text."""
+ config = AmazonConverseConfig()
+
+ usage = config.transform_usage(
+ {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042},
+ reasoning_content="a short reasoning summary",
+ thinking_ran=True,
+ provider_reasoning_tokens=1033,
+ )
+
+ assert usage.completion_tokens_details.reasoning_tokens == 1033
+ assert usage.completion_tokens_details.text_tokens == 3002 - 1033
+
+
+def test_transform_usage_falls_back_to_estimate_without_provider_tokens():
+ config = AmazonConverseConfig()
+
+ usage = config.transform_usage(
+ {"inputTokens": 40, "outputTokens": 300, "totalTokens": 340},
+ reasoning_content="a short reasoning summary",
+ thinking_ran=True,
+ )
+
+ assert usage.completion_tokens_details.reasoning_tokens > 0
+ assert usage.completion_tokens_details.reasoning_tokens < 300
+
+
+def test_thinking_tokens_parsed_from_additional_model_response_fields():
+ parsed = AmazonConverseConfig.thinking_tokens_from_additional_fields(
+ {"usage": {"output_tokens_details": {"thinking_tokens": 92}}}
+ )
+ assert parsed == 92
+ assert AmazonConverseConfig.thinking_tokens_from_additional_fields(None) is None
+ assert AmazonConverseConfig.thinking_tokens_from_additional_fields({"usage": {}}) is None
+
+
@pytest.mark.parametrize(
"model,effort,expected_effort",
[
diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
index cd305d8ed26..4bef59842f1 100644
--- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
+++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
@@ -208,6 +208,29 @@ def test_bedrock_converse_streaming_consistent_id():
), "All chunk IDs must match the one captured from the messageStart event"
+def test_converse_streaming_usage_uses_provider_thinking_tokens():
+ """Regression LIT-5714: the messageStop event carries provider thinking tokens
+ under ``additionalModelResponseFields``; the usage chunk must report them instead
+ of a token_counter estimate."""
+ chunks = [
+ {
+ "contentBlockIndex": 0,
+ "delta": {"reasoningContent": {"text": "thinking about it"}},
+ },
+ {
+ "stopReason": "end_turn",
+ "additionalModelResponseFields": {"usage": {"output_tokens_details": {"thinking_tokens": 1033}}},
+ },
+ {"usage": {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}},
+ ]
+
+ decoder = AWSEventStreamDecoder(model="bedrock/anthropic.claude-opus-4-7")
+ parsed = [decoder.converse_chunk_parser(chunk) for chunk in chunks]
+
+ usage = parsed[-1].usage
+ assert usage.completion_tokens_details.reasoning_tokens == 1033
+
+
@pytest.mark.asyncio
async def test_make_call_does_not_rechunk_stream_by_default():
"""Re-chunking the event stream into fixed 1024-byte blocks holds small
diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
index d3c28302bf9..1e09afd6919 100644
--- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
+++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
@@ -1387,7 +1387,7 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model(
)
assert "reasoning_effort" not in result
- assert result.get("thinking") == {"type": "adaptive"}
+ assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert result.get("output_config") == {"effort": expected_effort}
@@ -2935,7 +2935,7 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag(
)
result = transform()
- assert result.get("thinking") == {"type": "adaptive"}
+ assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert result.get("output_config") == {"effort": "medium"}
monkeypatch.setitem(litellm.model_cost[model], "supports_adaptive_thinking", False)
diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py
new file mode 100644
index 00000000000..8c6eda605ca
--- /dev/null
+++ b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py
@@ -0,0 +1,197 @@
+import json
+from unittest.mock import MagicMock, patch
+
+import httpx
+import pytest
+from botocore.credentials import Credentials
+
+from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig
+from litellm.llms.bedrock_mantle.passthrough.transformation import BedrockMantlePassthroughConfig
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
+from litellm.passthrough.main import llm_passthrough_route
+from litellm.types.utils import LlmProviders
+from litellm.utils import ProviderConfigManager
+
+MANTLE_API_BASE = "https://bedrock-mantle.us-east-2.api.aws"
+INVOKE_ENDPOINT = "model/us.openai.gpt-5.6-sol/invoke"
+CONVERSE_ENDPOINT = "model/us.openai.gpt-5.6-sol/converse"
+REQUEST_BODY = {"messages": [{"role": "user", "content": "say pong"}], "max_completion_tokens": 64}
+
+
+@pytest.fixture
+def no_ambient_aws(monkeypatch):
+ for name in (
+ "AWS_BEARER_TOKEN_BEDROCK",
+ "BEDROCK_MANTLE_API_KEY",
+ "BEDROCK_MANTLE_API_BASE",
+ "BEDROCK_MANTLE_REGION",
+ "AWS_BEDROCK_RUNTIME_ENDPOINT",
+ "AWS_REGION_NAME",
+ "AWS_REGION",
+ "AWS_DEFAULT_REGION",
+ ):
+ monkeypatch.delenv(name, raising=False)
+
+
+def test_bedrock_mantle_registers_its_own_bedrock_passthrough_config():
+ config = ProviderConfigManager.get_provider_passthrough_config(
+ model="us.openai.gpt-5.6-sol", provider=LlmProviders.BEDROCK_MANTLE
+ )
+ assert isinstance(config, BedrockMantlePassthroughConfig)
+ assert isinstance(config, BedrockPassthroughConfig)
+
+
+def test_mantle_api_base_only_lends_its_region_to_the_runtime_url(no_ambient_aws):
+ url, base_url = BedrockMantlePassthroughConfig().get_complete_url(
+ api_base=MANTLE_API_BASE,
+ api_key=None,
+ model="us.openai.gpt-5.6-sol",
+ endpoint=INVOKE_ENDPOINT,
+ request_query_params=None,
+ litellm_params={"api_base": MANTLE_API_BASE},
+ )
+ assert str(url) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}"
+ assert base_url == "https://bedrock-runtime.us-east-2.amazonaws.com"
+
+
+def test_explicit_region_and_non_mantle_api_base_are_kept(no_ambient_aws):
+ vpc_endpoint = "https://vpce-0123.bedrock-runtime.us-east-1.vpce.amazonaws.com"
+ url, base_url = BedrockMantlePassthroughConfig().get_complete_url(
+ api_base=vpc_endpoint,
+ api_key=None,
+ model="us.openai.gpt-5.6-sol",
+ endpoint=INVOKE_ENDPOINT,
+ request_query_params=None,
+ litellm_params={"api_base": vpc_endpoint, "aws_region_name": "us-east-1"},
+ )
+ assert str(url) == f"{vpc_endpoint}/{INVOKE_ENDPOINT}"
+ assert base_url == vpc_endpoint
+
+
+def test_region_falls_back_to_the_mantle_default_without_any_hint(no_ambient_aws):
+ url, _ = BedrockMantlePassthroughConfig().get_complete_url(
+ api_base=None,
+ api_key=None,
+ model="us.openai.gpt-5.6-sol",
+ endpoint=INVOKE_ENDPOINT,
+ request_query_params=None,
+ litellm_params={},
+ )
+ assert str(url) == f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}"
+
+
+@pytest.mark.parametrize(
+ ("litellm_params", "env", "expected_bearer"),
+ [
+ ({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"),
+ ({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"),
+ ({}, {"AWS_BEARER_TOKEN_BEDROCK": "aws-env-key"}, "aws-env-key"),
+ ],
+)
+def test_sign_request_uses_the_deployment_bearer_token(no_ambient_aws, monkeypatch, litellm_params, env, expected_bearer):
+ for name, value in env.items():
+ monkeypatch.setenv(name, value)
+ headers, body = BedrockMantlePassthroughConfig().sign_request(
+ headers={},
+ litellm_params=litellm_params,
+ request_data=REQUEST_BODY,
+ api_base=f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}",
+ model="us.openai.gpt-5.6-sol",
+ )
+ assert headers["Authorization"] == f"Bearer {expected_bearer}"
+ assert body is not None
+ assert json.loads(body) == REQUEST_BODY
+
+
+def test_sign_request_falls_back_to_sigv4_scoped_to_the_mantle_region(no_ambient_aws):
+ config = BedrockMantlePassthroughConfig()
+ with patch.object(config, "get_credentials", return_value=Credentials("AKIA", "secret")):
+ headers, body = config.sign_request(
+ headers={},
+ litellm_params={"api_base": MANTLE_API_BASE},
+ request_data=REQUEST_BODY,
+ api_base=f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}",
+ model="us.openai.gpt-5.6-sol",
+ )
+ assert headers["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIA/")
+ assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
+ assert body is not None
+ assert json.loads(body) == REQUEST_BODY
+
+
+@pytest.mark.parametrize(
+ ("route_kwargs", "env", "expected_bearer"),
+ [
+ ({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"),
+ ({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"),
+ ],
+)
+def test_invoke_passthrough_route_reaches_bedrock_runtime_for_a_mantle_deployment(
+ no_ambient_aws, monkeypatch, route_kwargs, env, expected_bearer
+):
+ for name, value in env.items():
+ monkeypatch.setenv(name, value)
+ client = HTTPHandler()
+ with (
+ patch.object(client.client, "send", return_value=MagicMock(status_code=200)),
+ patch.object(client.client, "build_request", wraps=client.client.build_request) as build_request,
+ ):
+ response = llm_passthrough_route(
+ model="bedrock_mantle/us.openai.gpt-5.6-sol",
+ endpoint=INVOKE_ENDPOINT,
+ method="POST",
+ api_base=MANTLE_API_BASE,
+ json=dict(REQUEST_BODY),
+ client=client,
+ litellm_logging_obj=MagicMock(),
+ **route_kwargs,
+ )
+ assert response.status_code == 200
+ sent = build_request.call_args.kwargs
+ assert str(sent["url"]) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}"
+ assert sent["headers"]["Authorization"] == f"Bearer {expected_bearer}"
+ assert json.loads(sent["content"]) == REQUEST_BODY
+
+
+def _logged_model_response(endpoint, body):
+ request = httpx.Request("POST", f"https://bedrock-runtime.us-east-1.amazonaws.com/{endpoint}")
+ return BedrockMantlePassthroughConfig().logging_non_streaming_response(
+ model="us.openai.gpt-5.6-sol",
+ custom_llm_provider="bedrock_mantle",
+ httpx_response=httpx.Response(200, json=body, request=request),
+ request_data={"messages": [{"role": "user", "content": [{"text": "say pong"}]}]},
+ logging_obj=MagicMock(),
+ endpoint=endpoint,
+ )
+
+
+def test_converse_logging_parses_the_converse_response_shape():
+ result = _logged_model_response(
+ CONVERSE_ENDPOINT,
+ {
+ "metrics": {"latencyMs": 800.0},
+ "output": {"message": {"content": [{"text": "pong"}], "role": "assistant"}},
+ "stopReason": "end_turn",
+ "usage": {"inputTokens": 8, "outputTokens": 5, "totalTokens": 13},
+ },
+ )
+ assert result.choices[0].message.content == "pong"
+ assert result.usage.prompt_tokens == 8
+ assert result.usage.completion_tokens == 5
+
+
+def test_invoke_logging_parses_the_openai_chat_response_shape():
+ result = _logged_model_response(
+ INVOKE_ENDPOINT,
+ {
+ "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "pong", "role": "assistant"}}],
+ "created": 1787677792,
+ "id": "chatcmpl-regression",
+ "model": "us.openai.gpt-5.6-sol",
+ "object": "chat.completion",
+ "usage": {"completion_tokens": 5, "prompt_tokens": 8, "total_tokens": 13},
+ },
+ )
+ assert result.choices[0].message.content == "pong"
+ assert result.usage.prompt_tokens == 8
+ assert result.usage.completion_tokens == 5
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
index 9e05d48a18f..28c6060e5cc 100644
--- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
@@ -8,7 +8,8 @@ gate, the URL construction for both paths, and the shared Bearer auth.
"""
import copy
-
+import json
+from pathlib import Path
import pytest
from botocore.exceptions import (
@@ -1523,7 +1524,7 @@ class TestBedrockMantleResponsesPricing:
assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost)
assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost)
assert info["output_cost_per_token"] == pytest.approx(output_cost)
- assert info["max_input_tokens"] == 1000000
+ assert info["max_input_tokens"] == 1050000
assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2)
assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2)
assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2)
@@ -1565,3 +1566,42 @@ class TestBedrockMantleResponsesPricing:
def test_models_registered(self, local_cost_map):
assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models
+
+
+def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]:
+ repo_root = Path(__file__).resolve().parents[4]
+ paths = {
+ "root": repo_root / "model_prices_and_context_window.json",
+ "bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json",
+ }
+ return json.loads(paths[map_name].read_text())
+
+
+class TestGpt56MantleRegistryEntries:
+ """Locks the gpt-5.6 frontier entries to Bedrock Mantle's live behavior.
+
+ Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna
+ (oversize requests 400 with "prompt tokens (N) exceed model maximum
+ (1050000)", and a 1,030,590-token request completes), matching the OpenAI
+ Bedrock guide. mode must stay "responses": Mantle's native
+ /v1/chat/completions rejects function tools unless reasoning_effort is
+ "none", so chat traffic has to keep bridging to the Responses API
+ (see the responses_api_bridge tests above).
+ """
+
+ @pytest.mark.parametrize("map_name", ("root", "bundled_backup"))
+ @pytest.mark.parametrize(
+ "key",
+ (
+ "bedrock_mantle/openai.gpt-5.6-sol",
+ "bedrock_mantle/openai.gpt-5.6-terra",
+ "bedrock_mantle/openai.gpt-5.6-luna",
+ ),
+ )
+ def test_entry_matches_mantle_enforced_limits(self, map_name, key):
+ entry = _repo_cost_map(map_name)[key]
+ assert entry["max_input_tokens"] == 1050000
+ assert entry["max_output_tokens"] == 128000
+ assert entry["mode"] == "responses"
+ assert entry["use_openai_responses_path"] is True
+ assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"]
diff --git a/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py
new file mode 100644
index 00000000000..3364b1b3872
--- /dev/null
+++ b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py
@@ -0,0 +1,231 @@
+import pytest
+
+import litellm
+from litellm.interactions.utils import get_provider_interactions_api_config
+from litellm.llms.gemini.interactions.transformation import (
+ GoogleAIStudioInteractionsConfig,
+)
+from litellm.llms.vertex_ai.interactions.transformation import (
+ VertexAIInteractionsConfig,
+)
+from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+GLOBAL_BASE = "https://aiplatform.googleapis.com/v1beta1/projects/test-proj/locations/global/interactions"
+
+
+class MinterRecorder:
+ def __init__(self, resolved_project: str = "creds-proj") -> None:
+ self.calls: list[tuple[VERTEX_CREDENTIALS_TYPES | None, str | None]] = []
+ self.resolved_project = resolved_project
+
+ def __call__(
+ self,
+ credentials: VERTEX_CREDENTIALS_TYPES | None,
+ project_id: str | None,
+ ) -> tuple[str, str]:
+ self.calls.append((credentials, project_id))
+ return "test-token", project_id or self.resolved_project
+
+
+@pytest.fixture
+def minter():
+ return MinterRecorder()
+
+
+@pytest.fixture
+def config(minter):
+ return VertexAIInteractionsConfig(mint_access_token=minter)
+
+
+@pytest.fixture
+def litellm_params():
+ return GenericLiteLLMParams(vertex_project="test-proj", vertex_credentials="creds.json")
+
+
+class TestRegistration:
+ def test_vertex_ai_returns_vertex_config(self):
+ assert isinstance(get_provider_interactions_api_config("vertex_ai"), VertexAIInteractionsConfig)
+
+ def test_vertex_ai_beta_returns_vertex_config(self):
+ assert isinstance(get_provider_interactions_api_config("vertex_ai_beta"), VertexAIInteractionsConfig)
+
+ def test_gemini_still_returns_google_ai_studio_config(self):
+ gemini_config = get_provider_interactions_api_config("gemini")
+ assert isinstance(gemini_config, GoogleAIStudioInteractionsConfig)
+ assert not isinstance(gemini_config, VertexAIInteractionsConfig)
+
+ def test_lazy_import_resolves(self):
+ assert litellm.VertexAIInteractionsConfig is VertexAIInteractionsConfig
+
+ def test_custom_llm_provider_is_vertex_ai(self, config):
+ assert config.custom_llm_provider == LlmProviders.VERTEX_AI
+
+
+class TestValidateEnvironment:
+ def test_sets_bearer_auth_without_gemini_headers(self, config, minter, litellm_params):
+ headers = config.validate_environment(
+ headers={},
+ model="gemini-omni-flash-preview",
+ litellm_params=litellm_params,
+ )
+
+ assert headers["Authorization"] == "Bearer test-token"
+ assert headers["Content-Type"] == "application/json"
+ assert "x-goog-api-key" not in headers
+ assert "Api-Revision" not in headers
+ assert minter.calls == [("creds.json", "test-proj")]
+
+ def test_caller_authorization_wins(self, config, litellm_params):
+ headers = config.validate_environment(
+ headers={"Authorization": "Bearer caller-token"},
+ model="gemini-omni-flash-preview",
+ litellm_params=litellm_params,
+ )
+
+ assert headers["Authorization"] == "Bearer caller-token"
+
+
+class TestGetCompleteUrl:
+ def test_defaults_to_global_v1beta1(self, config, litellm_params):
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params=dict(litellm_params),
+ )
+
+ assert url == GLOBAL_BASE
+
+ def test_stream_appends_alt_sse(self, config, litellm_params):
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params=dict(litellm_params),
+ stream=True,
+ )
+
+ assert url == f"{GLOBAL_BASE}?alt=sse"
+
+ def test_multi_region_location_uses_rep_host(self, config):
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={"vertex_project": "test-proj", "vertex_location": "us"},
+ )
+
+ assert url == "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/test-proj/locations/us/interactions"
+
+ def test_regional_location_uses_regional_host(self, config):
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={"vertex_project": "test-proj", "vertex_location": "us-central1"},
+ )
+
+ assert url == (
+ "https://us-central1-aiplatform.googleapis.com"
+ "/v1beta1/projects/test-proj/locations/us-central1/interactions"
+ )
+
+ def test_location_env_fallback_is_ignored(self, config, monkeypatch):
+ monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5")
+
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={"vertex_project": "test-proj"},
+ )
+
+ assert url == GLOBAL_BASE
+
+ def test_api_base_override(self, config, litellm_params):
+ url = config.get_complete_url(
+ api_base="https://proxy.example.test",
+ model="gemini-omni-flash-preview",
+ litellm_params=dict(litellm_params),
+ )
+
+ assert url == "https://proxy.example.test/v1beta1/projects/test-proj/locations/global/interactions"
+
+ def test_project_resolved_from_credentials_when_not_passed(self, config, monkeypatch):
+ monkeypatch.delenv("VERTEXAI_PROJECT", raising=False)
+ monkeypatch.setattr(litellm, "vertex_project", None)
+
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={"vertex_credentials": "creds.json"},
+ )
+
+ assert url == "https://aiplatform.googleapis.com/v1beta1/projects/creds-proj/locations/global/interactions"
+
+ def test_invalid_location_rejected(self, config):
+ with pytest.raises(ValueError, match="Invalid vertex_location"):
+ config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={"vertex_project": "test-proj", "vertex_location": "evil.com#"},
+ )
+
+ def test_missing_project_rejected(self, monkeypatch):
+ monkeypatch.delenv("VERTEXAI_PROJECT", raising=False)
+ monkeypatch.setattr(litellm, "vertex_project", None)
+
+ def unresolved_minter(
+ credentials: VERTEX_CREDENTIALS_TYPES | None,
+ project_id: str | None,
+ ) -> tuple[str, str]:
+ return "test-token", ""
+
+ with pytest.raises(ValueError, match="Vertex AI project is required"):
+ VertexAIInteractionsConfig(mint_access_token=unresolved_minter).get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={},
+ )
+
+
+class TestInteractionByIdRequests:
+ def test_get_url(self, config, litellm_params):
+ url, request_body = config.transform_get_interaction_request(
+ interaction_id="abc123",
+ api_base="",
+ litellm_params=litellm_params,
+ headers={},
+ )
+
+ assert url == f"{GLOBAL_BASE}/abc123"
+ assert request_body == {}
+
+ def test_get_url_encodes_interaction_id(self, config, litellm_params):
+ url, _ = config.transform_get_interaction_request(
+ interaction_id="id/with space",
+ api_base="",
+ litellm_params=litellm_params,
+ headers={},
+ )
+
+ assert url == f"{GLOBAL_BASE}/id%2Fwith%20space"
+
+ def test_delete_url(self, config, litellm_params):
+ url, request_body = config.transform_delete_interaction_request(
+ interaction_id="abc123",
+ api_base="",
+ litellm_params=litellm_params,
+ headers={},
+ )
+
+ assert url == f"{GLOBAL_BASE}/abc123"
+ assert request_body == {}
+
+ def test_cancel_url(self, config, litellm_params):
+ url, request_body = config.transform_cancel_interaction_request(
+ interaction_id="abc123",
+ api_base="",
+ litellm_params=litellm_params,
+ headers={},
+ )
+
+ assert url == f"{GLOBAL_BASE}/abc123:cancel"
+ assert request_body == {}
diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py
index ba2f20e2337..f19e169dc9e 100644
--- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py
+++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py
@@ -538,7 +538,7 @@ def test_messages_thinking_shape_follows_exact_vertex_entry_flag(local_model_cos
)
result = transform()
- assert result.get("thinking") == {"type": "adaptive"}
+ assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert result.get("output_config") == {"effort": "medium"}
monkeypatch.setitem(
diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py
index 04f38b5e2ed..abe73f7d05c 100644
--- a/tests/test_litellm/proxy/auth/test_auth_checks.py
+++ b/tests/test_litellm/proxy/auth/test_auth_checks.py
@@ -47,6 +47,7 @@ from litellm.proxy.auth.auth_checks import (
_virtual_key_soft_budget_check,
get_key_object,
get_user_object,
+ invalidate_team_member_spend_state,
vector_store_access_check,
)
from litellm.caching.in_memory_cache import InMemoryCache
@@ -6939,3 +6940,307 @@ def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is
router = _router_with_a_group_priced_through_model_info()
assert model_has_no_cost_mapping(model="model-info-priced-alias", llm_router=router) is False
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_sets_the_spend_counter_and_clears_both_membership_cache_keys():
+ """A team-member budget reset (new_spend passed) must SET the spend counter to the reset
+ value, clear its DB-floor marker, AND invalidate both independently-keyed membership caches
+ (user_api_key_auth.py's admission check writes one key format, budget_reservation.py and
+ auth_checks.py's own get_team_membership() write the other) or a stale read keeps 429ing
+ after the reset. Asserted against real cache reads, not mock call args, so a change that
+ keeps the call but drops its effect still fails."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ real_cache = UserApiKeyCache()
+ await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership")
+ await real_cache.async_set_cache(key="team_membership:user-1:team-1", value="stale-membership")
+
+ real_spend_counter_cache = DualCache()
+ real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0)
+ real_spend_counter_cache.in_memory_cache.set_cache(
+ key="spend_db_floor:spend:team_member:user-1:team-1", value=999.0
+ )
+
+ with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=real_cache,
+ new_spend=0.0,
+ )
+
+ assert await real_cache.async_get_cache(key="team-1_user-1") is None
+ assert await real_cache.async_get_cache(key="team_membership:user-1:team-1") is None
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 0.0
+ assert (
+ real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1")
+ == 0.0
+ ), "the DB-floor marker kept the pre-reset value; a stale-floor read can raise the counter right back up"
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_leaves_the_live_spend_counter_alone_without_new_spend():
+ """team_member_update only changes the budget cap, not the tracked spend, so it calls
+ invalidate_team_member_spend_state with no new_spend. Deleting the live spend counter in that
+ case would force the next read to reseed from the DB's own spend column, which lags the live
+ counter via periodic batch writes, briefly UNDER-enforcing the raised cap against a spend
+ value lower than what was actually tracked (regression: PR #37971 Bugbot finding)."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ real_cache = UserApiKeyCache()
+ await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership")
+
+ real_spend_counter_cache = DualCache()
+ real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0)
+
+ with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=real_cache,
+ )
+
+ assert await real_cache.async_get_cache(key="team-1_user-1") is None
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 999.0
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_deleting():
+ """/key/{key}/reset_spend SETs its counter to the reset value rather than deleting it, so a
+ worker's next read reflects it directly instead of falling back through a DB reseed. A reset
+ caller passing new_spend must match that precedent, not merely delete the counter."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ real_cache = UserApiKeyCache()
+ real_spend_counter_cache = DualCache()
+ fake_redis_cache = MagicMock()
+ fake_redis_cache.async_set_cache = AsyncMock()
+ real_spend_counter_cache.redis_cache = fake_redis_cache
+
+ with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=real_cache,
+ new_spend=2.5,
+ )
+
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5
+ fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60)
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_deletes_redis_counter_when_set_fails(): # test-quality-ok: only observable effect is the fallback call on the same fake client
+ """Redis reads take priority over the local in-memory copy (get_current_spend reads Redis
+ first), so a failed Redis SET would otherwise leave the OLD pre-reset value authoritative
+ for every worker even though the reset reported success. On a failed SET, the stale Redis
+ entry must be deleted instead, so the next read clean-misses and reseeds from the DB."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ real_cache = UserApiKeyCache()
+ real_spend_counter_cache = DualCache()
+ fake_redis_cache = MagicMock()
+ fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down"))
+ fake_redis_cache.async_delete_cache = AsyncMock()
+ real_spend_counter_cache.redis_cache = fake_redis_cache
+
+ with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=real_cache,
+ new_spend=2.5,
+ )
+
+ fake_redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1")
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_raises_503_when_both_redis_writes_fail():
+ """If the Redis SET fails AND the fallback DELETE fails, the stale pre-reset counter is still
+ authoritative in Redis for every worker. Reporting success would silently keep 429ing the
+ member, so the reset must surface a 503 instead (regression: PR #37971 Greptile finding)."""
+ from fastapi import HTTPException
+
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ real_cache = UserApiKeyCache()
+ real_spend_counter_cache = DualCache()
+ fake_redis_cache = MagicMock()
+ fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down"))
+ fake_redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis still down"))
+ real_spend_counter_cache.redis_cache = fake_redis_cache
+
+ with (
+ patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
+ ),
+ pytest.raises(HTTPException) as exc_info,
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=real_cache,
+ new_spend=2.5,
+ )
+
+ assert exc_info.value.status_code == 503
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_to_remote_workers():
+ """The test above only proves the handling worker's own spend counter is
+ cleared. A remote worker's spend counter is a separate DualCache instance;
+ if the reset never reaches it, that worker keeps enforcing the pre-reset
+ spend the moment its own Redis read for the counter fails and it falls
+ back to its own (now-stale) in-memory copy. Drives the actual message
+ published onto the invalidation channel through a second, independent
+ AuthCacheInvalidationSubscriber standing in for that remote worker, rather
+ than asserting on the publish call args."""
+ from redis.asyncio import Redis
+
+ from litellm.caching.dual_cache import DualCache
+ from litellm.caching.in_memory_cache import InMemoryCache
+ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ published: list[tuple[str, str]] = []
+
+ class _RecordingRedisClient(Redis):
+ def __init__(self) -> None:
+ pass
+
+ async def publish(self, channel: str, message: str) -> int:
+ published.append((channel, message))
+ return 1
+
+ class _FakeRedisCache:
+ def __init__(self) -> None:
+ self.namespace = None
+
+ def init_async_client(self) -> object:
+ return _RecordingRedisClient()
+
+ local_spend_counter_cache = DualCache()
+
+ remote_user_api_key_cache = UserApiKeyCache()
+ remote_spend_counter_in_memory_cache = InMemoryCache()
+ remote_spend_counter_in_memory_cache.set_cache("spend:team_member:user-1:team-1", 999.0)
+ remote_spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:team_member:user-1:team-1", 999.0)
+
+ with (
+ patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache
+ ),
+ patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test
+ "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
+ return_value=_FakeRedisCache(),
+ ),
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=UserApiKeyCache(),
+ new_spend=0.0,
+ )
+
+ def _published_message_for(cache_key: str) -> str:
+ matches = [message for _, message in published if json.loads(message)["cache_key"] == cache_key]
+ assert matches, f"{cache_key} never reached the cross-worker invalidation channel"
+ return matches[-1]
+
+ remote_subscriber = AuthCacheInvalidationSubscriber(
+ redis_cache=_FakeRedisCache(),
+ user_api_key_cache=remote_user_api_key_cache,
+ additional_in_memory_caches=(remote_spend_counter_in_memory_cache,),
+ )
+ for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"):
+ remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
+ {"type": "message", "data": _published_message_for(cache_key)}
+ )
+
+ assert remote_spend_counter_in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0
+ assert (
+ remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0
+ ), "the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor"
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_not_erase_the_reset():
+ """The handling worker subscribes to the same invalidation channel it publishes on, so it
+ receives its own reset message. A delete-style broadcast would erase the post-reset counter
+ and floor marker the handler just wrote, reopening the stale-floor race the reset closed
+ (regression: PR #37971 Greptile finding). The broadcast carries the reset value as a SET, so
+ applying the self-delivered message must leave both keys at the post-reset value."""
+ from redis.asyncio import Redis
+
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ published: list[tuple[str, str]] = []
+
+ class _RecordingRedisClient(Redis):
+ def __init__(self) -> None:
+ pass
+
+ async def publish(self, channel: str, message: str) -> int:
+ published.append((channel, message))
+ return 1
+
+ class _FakeRedisCache:
+ def __init__(self) -> None:
+ self.namespace = None
+
+ def init_async_client(self) -> object:
+ return _RecordingRedisClient()
+
+ local_spend_counter_cache = DualCache()
+ local_user_api_key_cache = UserApiKeyCache()
+
+ with (
+ patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache
+ ),
+ patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test
+ "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
+ return_value=_FakeRedisCache(),
+ ),
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=local_user_api_key_cache,
+ new_spend=0.0,
+ )
+
+ own_subscriber = AuthCacheInvalidationSubscriber(
+ redis_cache=_FakeRedisCache(),
+ user_api_key_cache=local_user_api_key_cache,
+ additional_in_memory_caches=(local_spend_counter_cache.in_memory_cache,),
+ )
+ for _, message in published:
+ own_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
+ {"type": "message", "data": message}
+ )
+
+ assert local_spend_counter_cache.in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0, (
+ "the handler's self-delivered broadcast erased the post-reset spend counter"
+ )
+ assert (
+ local_spend_counter_cache.in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0
+ ), "the handler's self-delivered broadcast erased the post-reset floor marker, reopening the stale-floor race"
diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py
index 468e8aabae8..7d5fc1a3544 100644
--- a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py
+++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py
@@ -6,6 +6,7 @@ from unittest.mock import patch
import pytest
from redis.asyncio import Redis
+from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
AUTH_CACHE_INVALIDATION_CHANNEL,
AuthCacheInvalidationSubscriber,
@@ -144,6 +145,37 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None:
assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL]
+@pytest.mark.asyncio
+async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None:
+ """
+ The spend-counter half of the same cross-worker gap: a remote worker's own
+ spend counter can hold a stale value (its fallback path when that worker's
+ own Redis read for the counter fails), and only clearing user_api_key_cache
+ on message would leave that separate DualCache's in-memory copy untouched.
+ """
+ cache = UserApiKeyCache()
+ spend_counter_in_memory_cache = InMemoryCache()
+ spend_counter_in_memory_cache.set_cache("spend:team_member:u-1:t-1", 999.0)
+ assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is not None
+
+ pubsub = _QueuePubSub(initial_messages=[_invalidation_message("spend:team_member:u-1:t-1")])
+ subscriber = AuthCacheInvalidationSubscriber(
+ redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])),
+ user_api_key_cache=cache,
+ additional_in_memory_caches=(spend_counter_in_memory_cache,),
+ )
+ subscriber.start()
+ try:
+ for _ in range(200):
+ if spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None:
+ break
+ await asyncio.sleep(0.01)
+ finally:
+ await subscriber.stop()
+
+ assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None
+
+
@pytest.mark.asyncio
async def test_subscriber_ignores_malformed_messages() -> None:
cache = UserApiKeyCache()
diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
index f6d74a189bc..7f5d3eb0a14 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
@@ -9,11 +9,13 @@ from unittest.mock import AsyncMock, MagicMock, call, patch
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
+from pydantic import ValidationError
from litellm._uuid import uuid
from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth
from litellm.proxy._types import (
+ LiteLLM_BudgetTable,
LiteLLM_BudgetTableFull,
LiteLLM_ModelTable,
LiteLLM_OrganizationMembershipTable,
@@ -27,7 +29,9 @@ from litellm.proxy._types import (
Member,
ProxyErrorTypes,
ProxyException,
+ ResetSpendRequest,
TeamMemberAddRequest,
+ TeamMemberUpdateRequest,
UpdateTeamRequest,
)
from litellm.proxy.management_endpoints.team_endpoints import (
@@ -42,12 +46,15 @@ from litellm.proxy.management_endpoints.team_endpoints import (
_transform_teams_to_deleted_records,
_update_model_table,
_validate_and_populate_member_user_info,
+ _validate_team_member_reset_spend_value,
_verify_team_access,
delete_team,
list_available_teams,
+ reset_team_member_spend_fn,
router,
team_member_add_duplication_check,
team_member_delete,
+ team_member_update,
update_team,
validate_team_org_change,
)
@@ -12603,3 +12610,376 @@ async def test_invalidate_access_group_cache_deletes_the_cached_object():
"user_api_key_cache": cache,
"proxy_logging_obj": logging_obj,
}
+
+
+def test_validate_team_member_reset_spend_value_rejects_non_numeric():
+ with pytest.raises(HTTPException) as exc:
+ _validate_team_member_reset_spend_value(
+ reset_to="not-a-number",
+ membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
+ )
+ assert exc.value.status_code == 400
+
+
+def test_validate_team_member_reset_spend_value_rejects_negative():
+ with pytest.raises(HTTPException) as exc:
+ _validate_team_member_reset_spend_value(
+ reset_to=-1.0,
+ membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
+ )
+ assert exc.value.status_code == 400
+
+
+@pytest.mark.parametrize("reset_to", [float("nan"), float("inf"), float("-inf")])
+def test_validate_team_member_reset_spend_value_rejects_non_finite(reset_to):
+ """NaN and +/-inf are instances of float and compare False against every bound
+ below (`nan < 0`, `nan > current_spend` are both False), so an isinstance-and-range
+ check alone lets them through to persist as the member's spend and silently
+ disable every later budget comparison against it."""
+ with pytest.raises(HTTPException) as exc:
+ _validate_team_member_reset_spend_value(
+ reset_to=reset_to,
+ membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
+ )
+ assert exc.value.status_code == 400
+
+
+@pytest.mark.parametrize("reset_to", [True, False])
+def test_reset_spend_request_rejects_bool_reset_to(reset_to):
+ """bool is a subclass of int, so pydantic silently coerces True/False into 1.0/0.0 for a
+ ``float`` field: {"reset_to": true} would otherwise reach _validate_team_member_reset_spend_value
+ as an indistinguishable 1.0 and reset the member's spend instead of failing the request."""
+ with pytest.raises(ValidationError):
+ ResetSpendRequest(reset_to=reset_to)
+
+
+def test_validate_team_member_reset_spend_value_rejects_above_current_spend():
+ with pytest.raises(HTTPException) as exc:
+ _validate_team_member_reset_spend_value(
+ reset_to=20.0,
+ membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
+ )
+ assert exc.value.status_code == 400
+
+
+def test_validate_team_member_reset_spend_value_rejects_above_max_budget():
+ with pytest.raises(HTTPException) as exc:
+ _validate_team_member_reset_spend_value(
+ reset_to=10.0,
+ membership=LiteLLM_TeamMembership(
+ user_id="u1",
+ team_id="t1",
+ spend=10.0,
+ litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=5.0),
+ ),
+ )
+ assert exc.value.status_code == 400
+
+
+def test_validate_team_member_reset_spend_value_accepts_valid_reset():
+ result = _validate_team_member_reset_spend_value(
+ reset_to=0.0,
+ membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
+ )
+ assert result == 0.0
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_success(monkeypatch):
+ """A proxy admin resetting a stuck team member's spend must write the DB
+ row to reset_to AND invalidate the cached spend/membership state, or the
+ 429 the endpoint exists to clear keeps firing off the stale cache.
+ Asserted against real cache reads, not mock call args, so a change that
+ keeps the call but drops its effect still fails."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ mock_prisma_client = MagicMock()
+ mock_proxy_logging_obj = MagicMock()
+ real_cache = UserApiKeyCache()
+ await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership")
+ await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership")
+ real_spend_counter_cache = DualCache()
+ real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
+
+ membership_row = LiteLLM_TeamMembership(
+ user_id="member-1",
+ team_id="team-1",
+ spend=10.0,
+ litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0),
+ )
+ mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
+ mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
+
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj)
+ monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
+
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
+ ):
+ response = await reset_team_member_spend_fn(
+ team_id="team-1",
+ user_id="member-1",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+
+ assert response["spend"] == 0.0
+ assert response["previous_spend"] == 10.0
+ assert response["max_budget"] == 50.0
+ mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with(
+ where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}},
+ data={"spend": 0.0},
+ )
+ assert await real_cache.async_get_cache(key="team-1_member-1") is None
+ assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch):
+ mock_prisma_client = MagicMock()
+ mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
+ ):
+ with pytest.raises(HTTPException) as exc:
+ await reset_team_member_spend_fn(
+ team_id="team-1",
+ user_id="ghost-user",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+ assert exc.value.status_code == 404
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_team_not_found(monkeypatch):
+ mock_prisma_client = MagicMock()
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."})),
+ ):
+ with pytest.raises(HTTPException) as exc:
+ await reset_team_member_spend_fn(
+ team_id="ghost-team",
+ user_id="member-1",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+ assert exc.value.status_code == 404
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_forbidden_for_non_admin(monkeypatch):
+ """A caller who is neither proxy admin, org admin, nor this team's admin must be refused,
+ matching every other team-mutating endpoint's authorization."""
+ mock_prisma_client = MagicMock()
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])),
+ ):
+ with pytest.raises(HTTPException) as exc:
+ await reset_team_member_spend_fn(
+ team_id="team-1",
+ user_id="member-1",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="plain-user"
+ ),
+ )
+ assert exc.value.status_code == 403
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_team_admin_cannot_reset_own_spend(monkeypatch):
+ """_verify_team_access authorizes a team admin over their own team with no check that the
+ target differs from the caller. Unchecked, that admin could target their own membership row
+ and repeatedly zero it right before it crosses their per-member cap, consuming the shared
+ team budget without the configured limit ever binding (Veria finding on PR #37971)."""
+ mock_prisma_client = MagicMock()
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+
+ team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-admin", user_id="team-admin-1")
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(
+ return_value=LiteLLM_TeamTable(
+ team_id="team-1",
+ members_with_roles=[Member(user_id="team-admin-1", role="admin")],
+ )
+ ),
+ ):
+ with pytest.raises(HTTPException) as exc:
+ await reset_team_member_spend_fn(
+ team_id="team-1",
+ user_id="team-admin-1",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=team_admin,
+ )
+ assert exc.value.status_code == 403
+ mock_prisma_client.db.litellm_teammembership.update.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkeypatch):
+ """The self-reset guard is scoped to non-proxy-admin roles: a proxy admin resetting their
+ own membership spend is the platform-wide trust boundary, not a team-scoped one."""
+ mock_prisma_client = MagicMock()
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+
+ membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0)
+ mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
+ mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
+
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
+ ):
+ response = await reset_team_member_spend_fn(
+ team_id="team-1",
+ user_id="admin-user",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+ assert response["spend"] == 0.0
+
+
+@pytest.mark.asyncio
+async def test_team_member_update_invalidates_team_member_spend_state_when_budget_patch_applied(monkeypatch):
+ """Raising a stuck member's max_budget_in_team via the documented /team/member_update
+ endpoint must invalidate the cached membership state, or the raised cap never reaches the
+ admission check and the member stays 429ing. The live spend counter itself must be left
+ untouched: only the cap changed, and deleting the counter would force a reseed from the
+ DB's own spend column, which lags the live counter via periodic batch writes, briefly
+ UNDER-enforcing the raised cap against a spend value lower than what was actually tracked.
+ Asserted against real cache reads, not mock call args, so a change that keeps the call but
+ drops its effect still fails."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ mock_prisma_client = MagicMock()
+ real_cache = UserApiKeyCache()
+ await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership")
+ await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership")
+ real_spend_counter_cache = DualCache()
+ real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
+
+ team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[])
+ team_info_response = {
+ "team_info": team_row,
+ "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)],
+ }
+
+ mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
+
+ monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
+ monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
+
+ mock_tx = AsyncMock()
+ mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx)
+ mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None)
+
+ with (
+ patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.team_info",
+ AsyncMock(return_value=team_info_response),
+ ),
+ patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership",
+ AsyncMock(),
+ ),
+ ):
+ await team_member_update(
+ data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1", max_budget_in_team=999999.0),
+ http_request=MagicMock(),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+
+ assert await real_cache.async_get_cache(key="team-1_member-1") is None
+ assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0
+
+
+@pytest.mark.asyncio
+async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(monkeypatch):
+ """A role-only update carries an empty budget_patch and touches no budget state,
+ so the member's cached spend/membership state must be left untouched."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ mock_prisma_client = MagicMock()
+ real_cache = UserApiKeyCache()
+ await real_cache.async_set_cache(key="team-1_member-1", value="still-fresh-membership")
+ real_spend_counter_cache = DualCache()
+ real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=1.5)
+
+ team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[])
+ team_info_response = {
+ "team_info": team_row,
+ "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)],
+ }
+
+ mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
+
+ monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
+ monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
+
+ mock_tx = AsyncMock()
+ mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx)
+ mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None)
+
+ with (
+ patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.team_info",
+ AsyncMock(return_value=team_info_response),
+ ),
+ patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership",
+ AsyncMock(),
+ ),
+ ):
+ await team_member_update(
+ data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1"),
+ http_request=MagicMock(),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+
+ assert await real_cache.async_get_cache(key="team-1_member-1") == "still-fresh-membership"
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5
diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index 31d2a6cef98..3383527e932 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -11347,3 +11347,38 @@ class TestRouterModelNameOnStreamingChunks:
assert len(frames) >= 3
assert '"router_model_name":"deep-model"' in frames[0]
assert all('"router_model_name":"backup-tier"' in frame for frame in frames[1:])
+
+
+@pytest.mark.asyncio
+async def test_authoritative_floor_spend_keeps_a_reset_marker_written_during_the_db_read():
+ """A team-member spend reset writes the post-reset floor to the spend_db_floor marker
+ (auth_checks.invalidate_team_member_spend_state). A floor read already in flight when the
+ reset commits would otherwise cache its stale pre-reset DB value over the fresh marker,
+ letting a budget check raise the counter right back above the just-reset spend
+ (regression: PR #37971 Greptile finding)."""
+ from litellm.proxy.proxy_server import _authoritative_floor_spend
+
+ real_spend_counter_cache = DualCache()
+ counter_key = "spend:team_member:user-1:team-1"
+ marker_key = f"spend_db_floor:{counter_key}"
+
+ async def db_read_racing_with_a_reset(prisma_client, counter_key):
+ real_spend_counter_cache.in_memory_cache.set_cache(key=marker_key, value=0.0)
+ return 999.0
+
+ with (
+ patch.object( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ proxy_server_module, "spend_counter_cache", real_spend_counter_cache
+ ),
+ patch.object( # test-quality-ok: the DB read must race the reset; no injectable seam for module-global prisma reads
+ proxy_server_module.SpendCounterReseed,
+ "from_db",
+ AsyncMock(side_effect=db_read_racing_with_a_reset),
+ ),
+ ):
+ result = await _authoritative_floor_spend(counter_key=counter_key)
+
+ assert result == 0.0
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key=marker_key) == 0.0, (
+ "the in-flight DB read clobbered the post-reset floor marker with the stale pre-reset value"
+ )
diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/test_litellm/rerank_api/test_main.py
index 85777afe81c..587be59c550 100644
--- a/tests/test_litellm/rerank_api/test_main.py
+++ b/tests/test_litellm/rerank_api/test_main.py
@@ -1,6 +1,10 @@
import logging
from unittest.mock import MagicMock, patch
+import httpx
+import pytest
+import respx
+
import litellm
@@ -62,3 +66,66 @@ def test_rerank_does_not_log_request_content_at_info(caplog):
assert all(
r.levelno == logging.DEBUG for r in optional_params_logs
), "optional_rerank_params must be logged at DEBUG, not INFO"
+
+
+TOGETHER_RERANK_BODY = {
+ "id": "rerank-mock-id",
+ "results": [{"index": 0, "relevance_score": 0.95}],
+ "usage": {"prompt_tokens": 10, "total_tokens": 10},
+}
+
+
+def test_together_rerank_defaults_to_together_ai_host(respx_mock: respx.MockRouter, monkeypatch):
+ """Regression for the Together host migration: rerank used to hardcode
+ https://api.together.xyz/v1/rerank. The default must now be api.together.ai."""
+ monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False)
+
+ mock_route = respx_mock.post("https://api.together.ai/v1/rerank")
+ mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
+
+ response = litellm.rerank(
+ model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
+ query=MARKER_QUERY,
+ documents=[MARKER_DOC],
+ api_key="fake-together-key",
+ )
+
+ assert mock_route.called
+ assert response.results[0]["relevance_score"] == 0.95
+
+
+def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter):
+ """Regression: a custom api_base was silently ignored by the Together rerank handler."""
+ mock_route = respx_mock.post("https://custom-together.example/v1/rerank")
+ mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
+
+ litellm.rerank(
+ model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
+ query=MARKER_QUERY,
+ documents=[MARKER_DOC],
+ api_key="fake-together-key",
+ api_base="https://custom-together.example/v1",
+ )
+
+ assert mock_route.called
+ assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key"
+
+
+@pytest.mark.asyncio
+async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch):
+ """Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank."""
+ monkeypatch.setenv("TOGETHER_AI_API_BASE", "https://env-together.example/v1")
+ monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
+
+ mock_route = respx_mock.post("https://env-together.example/v1/rerank")
+ mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
+
+ response = await litellm.arerank(
+ model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
+ query=MARKER_QUERY,
+ documents=[MARKER_DOC],
+ api_key="fake-together-key",
+ )
+
+ assert mock_route.called
+ assert response.results[0]["relevance_score"] == 0.95
diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py
index 99b1cc826aa..8f2b06be4b3 100644
--- a/tests/test_litellm/test_main.py
+++ b/tests/test_litellm/test_main.py
@@ -2944,3 +2944,16 @@ def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map):
assert cost == pytest.approx(
_priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens)
)
+
+
+@pytest.mark.asyncio
+async def test_acompletion_resolves_provider_from_api_base():
+ response = await litellm.acompletion(
+ model="deepseek-chat",
+ api_base="https://api.deepseek.com/v1",
+ api_key="fake-key",
+ messages=[{"role": "user", "content": "hi"}],
+ mock_response="resolved",
+ )
+
+ assert response.choices[0].message.content == "resolved"
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index e8830d05065..0b82ab971bf 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -9162,3 +9162,65 @@ def test_model_group_info_reasoning_efforts_ignore_a_mode_the_operator_declared(
assert result is not None
assert result.supported_reasoning_efforts == expected
+
+
+class TestAddDeploymentApiBaseProviderResolution:
+ def test_bare_model_with_known_api_base_initializes(self):
+ router = litellm.Router(
+ model_list=[
+ {
+ "model_name": "groq-pinned",
+ "litellm_params": {
+ "model": "llama-3.3-70b-versatile",
+ "api_base": "https://api.groq.com/openai/v1",
+ "api_key": "fake-key",
+ },
+ },
+ {
+ "model_name": "deepseek-pinned",
+ "litellm_params": {
+ "model": "deepseek-chat",
+ "api_base": "https://api.deepseek.com/v1",
+ "api_key": "fake-key",
+ },
+ },
+ ]
+ )
+
+ model_list = router.get_model_list()
+ assert model_list is not None
+ assert {m["model_name"] for m in model_list} == {"groq-pinned", "deepseek-pinned"}
+
+ def test_bare_model_with_unknown_api_base_still_raises(self):
+ with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"):
+ litellm.Router(
+ model_list=[
+ {
+ "model_name": "mystery",
+ "litellm_params": {
+ "model": "some-unknown-model",
+ "api_base": "https://llm.internal.example.com/v1",
+ "api_key": "fake-key",
+ },
+ }
+ ]
+ )
+
+ def test_explicit_custom_llm_provider_beats_api_base_endpoint_match(self):
+ router = litellm.Router(
+ model_list=[
+ {
+ "model_name": "openai-via-gateway",
+ "litellm_params": {
+ "model": "gpt-3.5-turbo",
+ "custom_llm_provider": "openai",
+ "api_base": "https://api.groq.com/openai/v1",
+ "api_key": "fake-key",
+ },
+ }
+ ]
+ )
+
+ deployment = router.get_deployment_by_model_group_name("openai-via-gateway")
+ assert deployment is not None
+ assert deployment.litellm_params.custom_llm_provider == "openai"
diff --git a/type-discipline-budget.json b/type-discipline-budget.json
index 627811a7f1d..8799e3e466a 100644
--- a/type-discipline-budget.json
+++ b/type-discipline-budget.json
@@ -1,9 +1,9 @@
{
"LIT001": {
- "limit": 22805
+ "limit": 22804
},
"LIT002": {
- "limit": 26873
+ "limit": 26872
},
"LIT003": {
"limit": 269
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx
index 438caa2f5e6..5aec78ba926 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx
@@ -381,6 +381,44 @@ describe("MCPServerEdit (true passthrough warning)", () => {
});
});
+describe("MCPServerEdit (OAuth authorize temp payload)", () => {
+ beforeEach(() => {
+ vi.clearAllMocks();
+ });
+
+ it("forwards issuer/authorization_url/token_url/registration_url to the temp OAuth session payload", async () => {
+ // Without these fields the ephemeral server the temp OAuth session endpoint builds has no
+ // admin-configured OAuth endpoints on it, discovery falls back to (and fails against) the
+ // plain server url, and Authorize & Fetch Token 400s with "authorization url is not
+ // configured" even though the saved server (and the visible form) has all four fields filled in.
+ render(
+ ,
+ );
+
+ await waitFor(() => {
+ expect(mockOauth.getTemporaryPayload).toBeTruthy();
+ });
+ const payload = mockOauth.getTemporaryPayload!();
+ expect(payload).toBeTruthy();
+ expect(payload?.issuer).toBe("https://github.com/login/oauth");
+ expect(payload?.authorization_url).toBe("https://github.com/login/oauth/authorize");
+ expect(payload?.token_url).toBe("https://github.com/login/oauth/access_token");
+ expect(payload?.registration_url).toBe("https://github.com/login/oauth/register");
+ });
+});
+
describe("MCPServerEdit (auth type switch)", () => {
beforeEach(() => {
vi.clearAllMocks();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx
index 5e79b20825c..8793c45371a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx
@@ -282,6 +282,10 @@ const MCPServerEdit: React.FC = ({
credentials: isClientForwardedTokenMode(values.auth_type)
? preservedAdminCredentials(values.credentials)
: values.credentials,
+ issuer: values.issuer,
+ authorization_url: values.authorization_url,
+ token_url: values.token_url,
+ registration_url: values.registration_url,
mcp_access_groups: values.mcp_access_groups || mcpServer.mcp_access_groups,
static_headers: staticHeaders,
command: values.command,
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 80a8fb041a0..67ec35d36a0 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -14936,6 +14936,33 @@ export interface paths {
patch?: never;
trace?: never;
};
+ "/team/{team_id}/member/{user_id}/reset_spend": {
+ parameters: {
+ query?: never;
+ header?: never;
+ path?: never;
+ cookie?: never;
+ };
+ get?: never;
+ put?: never;
+ /**
+ * Reset Team Member Spend Fn
+ * @description Reset a team member's tracked spend against their per-member budget.
+ *
+ * A member's spend is tracked separately from both their own personal
+ * budget and the team's own budget (LiteLLM_TeamMembership.spend), so
+ * neither /user/update nor /team/update can clear it: this is the only
+ * endpoint that does. The cross-pod spend counter and cached membership
+ * reads are invalidated so the reset takes effect on the member's next
+ * request rather than waiting on the membership cache's TTL.
+ */
+ post: operations["reset_team_member_spend_fn_team__team_id__member__user_id__reset_spend_post"];
+ delete?: never;
+ options?: never;
+ head?: never;
+ patch?: never;
+ trace?: never;
+ };
"/team/{team_id}/members/me": {
parameters: {
query?: never;
@@ -55085,6 +55112,42 @@ export interface operations {
};
};
};
+ reset_team_member_spend_fn_team__team_id__member__user_id__reset_spend_post: {
+ parameters: {
+ query?: never;
+ header?: never;
+ path: {
+ team_id: string;
+ user_id: string;
+ };
+ cookie?: never;
+ };
+ requestBody: {
+ content: {
+ "application/json": components["schemas"]["ResetSpendRequest"];
+ };
+ };
+ responses: {
+ /** @description Successful Response */
+ 200: {
+ headers: {
+ [name: string]: unknown;
+ };
+ content: {
+ "application/json": unknown;
+ };
+ };
+ /** @description Validation Error */
+ 422: {
+ headers: {
+ [name: string]: unknown;
+ };
+ content: {
+ "application/json": components["schemas"]["HTTPValidationError"];
+ };
+ };
+ };
+ };
team_member_me_team__team_id__members_me_get: {
parameters: {
query?: never;