From 8d50956051222b5bb94642506059cee68e2c1ee5 Mon Sep 17 00:00:00 2001 From: Adam Reed Date: Mon, 16 Feb 2026 22:09:07 -0600 Subject: [PATCH 01/16] fix(proxy): preserve and forward OAuth Authorization headers through proxy layer (#19912) PR #21039 fixed OAuth token handling at the LLM layer (Authorization: Bearer instead of x-api-key), but the proxy layer still strips the Authorization header in clean_headers() before it reaches the Anthropic code. This breaks OAuth for proxy users (e.g., Claude Code Max through LiteLLM proxy). Changes: - Add is_anthropic_oauth_key() helper to detect OAuth tokens (sk-ant-oat*) - Preserve OAuth Authorization headers in clean_headers() instead of stripping - Forward OAuth Authorization via ProviderSpecificHeader in add_provider_specific_headers_to_request() so tokens only reach Anthropic-compatible providers (anthropic, bedrock, vertex_ai) Fixes #19618 Co-authored-by: Adam Reed --- litellm/llms/anthropic/common_utils.py | 9 ++ litellm/proxy/litellm_pre_call_utils.py | 19 ++- .../anthropic/test_anthropic_common_utils.py | 139 ++++++++++++++++++ 3 files changed, 166 insertions(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index c665e084261..0cceddd9acf 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -22,6 +22,15 @@ from litellm.types.llms.anthropic import ( from litellm.types.llms.openai import AllMessageValues +def is_anthropic_oauth_key(value: Optional[str]) -> bool: + """Check if a value contains an Anthropic OAuth token (sk-ant-oat*).""" + if value is None: + return False + # Handle both raw token and "Bearer " format + if value.startswith("Bearer "): + value = value[7:] + return value.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX) + def optionally_handle_anthropic_oauth( headers: dict, api_key: Optional[str] ) -> tuple[dict, Optional[str]]: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4d77af513a8..3e8cc46521d 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -239,6 +239,8 @@ def clean_headers( """ Removes litellm api key from headers """ + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + clean_headers = {} litellm_key_lower = ( litellm_key_header_name.lower() if litellm_key_header_name is not None else None @@ -246,8 +248,13 @@ def clean_headers( for header, value in headers.items(): header_lower = header.lower() + # Preserve Authorization header if it contains Anthropic OAuth token (sk-ant-oat*) + # This allows OAuth tokens to be forwarded to Anthropic-compatible providers + # via add_provider_specific_headers_to_request() + if header_lower == "authorization" and is_anthropic_oauth_key(value): + clean_headers[header] = value # Check if header should be excluded: either in special headers cache or matches custom litellm key - if header_lower not in _SPECIAL_HEADERS_CACHE and ( + elif header_lower not in _SPECIAL_HEADERS_CACHE and ( litellm_key_lower is None or header_lower != litellm_key_lower ): clean_headers[header] = value @@ -1717,6 +1724,8 @@ def add_provider_specific_headers_to_request( data: dict, headers: dict, ): + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + anthropic_headers = {} # boolean to indicate if a header was added added_header = False @@ -1726,6 +1735,14 @@ def add_provider_specific_headers_to_request( anthropic_headers[header] = header_value added_header = True + # Check for Authorization header with Anthropic OAuth token (sk-ant-oat*) + # This needs to be handled via provider-specific headers to ensure it only + # goes to Anthropic-compatible providers, not all providers in the router + for header, value in headers.items(): + if header.lower() == "authorization" and is_anthropic_oauth_key(value): + anthropic_headers[header] = value + added_header = True + break if added_header is True: # Anthropic headers work across multiple providers # Store as comma-separated list so retrieval can match any of them diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index a321a24540f..ebffb56446e 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -283,3 +283,142 @@ class TestPassthroughOAuth: assert updated_headers["x-api-key"] == FAKE_REGULAR_KEY assert "authorization" not in updated_headers + + +class TestIsAnthropicOAuthKey: + """Tests for is_anthropic_oauth_key helper function.""" + + def test_oauth_token_raw(self): + """Raw OAuth token should be detected.""" + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + assert is_anthropic_oauth_key("sk-ant-oat01-abc123") is True + assert is_anthropic_oauth_key("sk-ant-oat02-xyz789") is True + + def test_oauth_token_bearer_format(self): + """Bearer-prefixed OAuth token should be detected.""" + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + assert is_anthropic_oauth_key("Bearer sk-ant-oat01-abc123") is True + assert is_anthropic_oauth_key("Bearer sk-ant-oat02-xyz789") is True + + def test_non_oauth_tokens(self): + """Non-OAuth values should return False.""" + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + assert is_anthropic_oauth_key(None) is False + assert is_anthropic_oauth_key("") is False + assert is_anthropic_oauth_key("sk-ant-api01-abc123") is False + assert is_anthropic_oauth_key("Bearer sk-ant-api01-abc123") is False + + def test_case_sensitivity(self): + """OAuth prefix matching should be case-sensitive.""" + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + assert is_anthropic_oauth_key("sk-ant-OAT01-abc123") is False + assert is_anthropic_oauth_key("SK-ANT-OAT01-abc123") is False + + def test_just_prefix(self): + """Just the prefix with no suffix should still match.""" + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + assert is_anthropic_oauth_key("sk-ant-oat") is True + + +class TestProxyOAuthHeaderForwarding: + """Tests for proxy-layer OAuth header preservation and forwarding.""" + + def test_clean_headers_preserves_oauth_authorization(self): + """clean_headers should preserve Authorization header with OAuth tokens.""" + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import clean_headers + + raw_headers = Headers( + raw=[ + (b"authorization", f"Bearer {FAKE_OAUTH_TOKEN}".encode()), + (b"content-type", b"application/json"), + ] + ) + cleaned = clean_headers(raw_headers) + + assert "authorization" in cleaned + assert cleaned["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + assert cleaned["content-type"] == "application/json" + + def test_clean_headers_strips_non_oauth_authorization(self): + """clean_headers should strip Authorization header with regular API keys.""" + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import clean_headers + + raw_headers = Headers( + raw=[ + (b"authorization", b"Bearer sk-regular-key-123"), + (b"content-type", b"application/json"), + ] + ) + cleaned = clean_headers(raw_headers) + + assert "authorization" not in cleaned + assert cleaned["content-type"] == "application/json" + + def test_add_provider_specific_headers_forwards_oauth(self): + """add_provider_specific_headers_to_request should forward OAuth Authorization + as a ProviderSpecificHeader scoped to Anthropic-compatible providers.""" + from litellm.proxy.litellm_pre_call_utils import ( + add_provider_specific_headers_to_request, + ) + + data: dict = {} + headers = { + "authorization": f"Bearer {FAKE_OAUTH_TOKEN}", + "content-type": "application/json", + } + + add_provider_specific_headers_to_request(data=data, headers=headers) + + assert "provider_specific_header" in data + psh = data["provider_specific_header"] + assert "anthropic" in psh["custom_llm_provider"] + assert "bedrock" in psh["custom_llm_provider"] + assert "vertex_ai" in psh["custom_llm_provider"] + assert psh["extra_headers"]["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + + def test_add_provider_specific_headers_ignores_non_oauth(self): + """add_provider_specific_headers_to_request should not create a + ProviderSpecificHeader for non-OAuth Authorization headers.""" + from litellm.proxy.litellm_pre_call_utils import ( + add_provider_specific_headers_to_request, + ) + + data: dict = {} + headers = { + "authorization": "Bearer sk-regular-key-123", + "content-type": "application/json", + } + + add_provider_specific_headers_to_request(data=data, headers=headers) + + assert "provider_specific_header" not in data + + def test_add_provider_specific_headers_combines_anthropic_and_oauth(self): + """When both anthropic-beta and OAuth Authorization are present, both + should be included in the ProviderSpecificHeader.""" + from litellm.proxy.litellm_pre_call_utils import ( + add_provider_specific_headers_to_request, + ) + + data: dict = {} + headers = { + "authorization": f"Bearer {FAKE_OAUTH_TOKEN}", + "anthropic-beta": "oauth-2025-04-20", + "content-type": "application/json", + } + + add_provider_specific_headers_to_request(data=data, headers=headers) + + assert "provider_specific_header" in data + psh = data["provider_specific_header"] + assert psh["extra_headers"]["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + assert psh["extra_headers"]["anthropic-beta"] == "oauth-2025-04-20" From 72af441159e5f77f400fa6c384a5c05f6c724a5c Mon Sep 17 00:00:00 2001 From: Mateusz Szewczyk <139469471+MateuszOssGit@users.noreply.github.com> Date: Tue, 17 Feb 2026 05:12:16 +0100 Subject: [PATCH 02/16] feat: Add IBM watsonx.ai rerank support (#21303) * feat: Add IBM watsonx.ai rerank support * feat: added unit tests * fix docstring * added documentataion * Update litellm/llms/watsonx/rerank/transformation.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update litellm/rerank_api/main.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update litellm/llms/watsonx/rerank/transformation.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * update validate_environment signature * fix ruff check and mypy * fix CR * CR fix * CR fix --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../docs/providers/watsonx/rerank.md | 52 ++++ docs/my-website/docs/rerank.md | 47 ++-- litellm/__init__.py | 1 + litellm/_lazy_imports_registry.py | 2 + litellm/llms/watsonx/__init__.py | 0 litellm/llms/watsonx/chat/__init__.py | 0 litellm/llms/watsonx/completion/__init__.py | 0 litellm/llms/watsonx/embed/__init__.py | 0 litellm/llms/watsonx/rerank/__init__.py | 0 litellm/llms/watsonx/rerank/transformation.py | 204 +++++++++++++++ litellm/rerank_api/main.py | 29 ++- litellm/types/llms/watsonx.py | 1 + litellm/utils.py | 2 + .../llms/watsonx/rerank/__init__.py | 0 .../watsonx/rerank/test_watsonx_rerank.py | 236 ++++++++++++++++++ 15 files changed, 550 insertions(+), 24 deletions(-) create mode 100644 docs/my-website/docs/providers/watsonx/rerank.md create mode 100644 litellm/llms/watsonx/__init__.py create mode 100644 litellm/llms/watsonx/chat/__init__.py create mode 100644 litellm/llms/watsonx/completion/__init__.py create mode 100644 litellm/llms/watsonx/embed/__init__.py create mode 100644 litellm/llms/watsonx/rerank/__init__.py create mode 100644 litellm/llms/watsonx/rerank/transformation.py create mode 100644 tests/test_litellm/llms/watsonx/rerank/__init__.py create mode 100644 tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py diff --git a/docs/my-website/docs/providers/watsonx/rerank.md b/docs/my-website/docs/providers/watsonx/rerank.md new file mode 100644 index 00000000000..0900ce96781 --- /dev/null +++ b/docs/my-website/docs/providers/watsonx/rerank.md @@ -0,0 +1,52 @@ +# watsonx.ai Rerank + +## Overview + +| Property | Details | +|----------|--------------------------------------------------------------------------| +| Description | watsonx.ai rerank integration | +| Provider Route on LiteLLM | `watsonx/` | +| Supported Operations | `/ml/v1/text/rerank` | +| Link to Provider Doc | [IBM WatsonX.ai ↗](https://cloud.ibm.com/apidocs/watsonx-ai#text-rerank) | + +## Quick Start + +### **LiteLLM SDK** + +```python +import os +from litellm import rerank + +os.environ["WATSONX_APIKEY"] = "YOUR_WATSONX_APIKEY" +os.environ["WATSONX_API_BASE"] = "YOUR_WATSONX_API_BASE" +os.environ["WATSONX_PROJECT_ID"] = "YOUR_WATSONX_PROJECT_ID" + +query="Best programming language for beginners?" +documents=[ + "Python is great for beginners due to simple syntax.", + "JavaScript runs in browsers and is versatile.", + "Rust has a steep learning curve but is very safe.", +] + +response = rerank( + model="watsonx/cross-encoder/ms-marco-minilm-l-12-v2", + query=query, + documents=documents, + top_n=2, + return_documents=True, +) + +print(response) +``` + +### **LiteLLM Proxy** + +```yaml +model_list: + - model_name: cross-encoder/ms-marco-minilm-l-12-v2 + litellm_params: + model: watsonx/cross-encoder/ms-marco-minilm-l-12-v2 + api_key: os.environ/WATSONX_APIKEY + api_base: os.environ/WATSONX_API_BASE + project_id: os.environ/WATSONX_PROJECT_ID +``` diff --git a/docs/my-website/docs/rerank.md b/docs/my-website/docs/rerank.md index 90f685d2bbd..9c76883d7fd 100644 --- a/docs/my-website/docs/rerank.md +++ b/docs/my-website/docs/rerank.md @@ -8,15 +8,15 @@ LiteLLM Follows the [cohere api request / response for the rerank api](https://c ## Overview -| Feature | Supported | Notes | -|---------|-----------|-------| -| Cost Tracking | ✅ | Works with all supported models | -| Logging | ✅ | Works across all integrations | -| End-user Tracking | ✅ | | -| Fallbacks | ✅ | Works between supported models | -| Loadbalancing | ✅ | Works between supported models | -| Guardrails | ✅ | Applies to input query only (not documents) | -| Supported Providers | Cohere, Together AI, Azure AI, DeepInfra, Nvidia NIM, Infinity, Fireworks AI, Voyage AI | | +| Feature | Supported | Notes | +|---------|-----------------------------------------------------------------------------------------------------|-------| +| Cost Tracking | ✅ | Works with all supported models | +| Logging | ✅ | Works across all integrations | +| End-user Tracking | ✅ | | +| Fallbacks | ✅ | Works between supported models | +| Loadbalancing | ✅ | Works between supported models | +| Guardrails | ✅ | Applies to input query only (not documents) | +| Supported Providers | Cohere, Together AI, Azure AI, DeepInfra, Nvidia NIM, Infinity, Fireworks AI, Voyage AI, watsonx.ai | | ## **LiteLLM Python SDK Usage** ### Quick Start @@ -123,17 +123,18 @@ curl http://0.0.0.0:4000/rerank \ #### ⚡️See all supported models and providers at [models.litellm.ai](https://models.litellm.ai/) -| Provider | Link to Usage | -|-------------|--------------------| -| Cohere (v1 + v2 clients) | [Usage](#quick-start) | -| Together AI| [Usage](../docs/providers/togetherai) | -| Azure AI| [Usage](../docs/providers/azure_ai#rerank-endpoint) | -| Jina AI| [Usage](../docs/providers/jina_ai) | -| AWS Bedrock| [Usage](../docs/providers/bedrock#rerank-api) | -| HuggingFace| [Usage](../docs/providers/huggingface_rerank) | -| Infinity| [Usage](../docs/providers/infinity) | -| vLLM| [Usage](../docs/providers/vllm#rerank-endpoint) | -| DeepInfra| [Usage](../docs/providers/deepinfra#rerank-endpoint) | -| Vertex AI| [Usage](../docs/providers/vertex#rerank-api) | -| Fireworks AI| [Usage](../docs/providers/fireworks_ai#rerank-endpoint) | -| Voyage AI| [Usage](../docs/providers/voyage#rerank) | \ No newline at end of file +| Provider | Link to Usage | +|--------------------------|------------------------------------------------------| +| Cohere (v1 + v2 clients) | [Usage](#quick-start) | +| Together AI | [Usage](../docs/providers/togetherai) | +| Azure AI | [Usage](../docs/providers/azure_ai#rerank-endpoint) | +| Jina AI | [Usage](../docs/providers/jina_ai) | +| AWS Bedrock | [Usage](../docs/providers/bedrock#rerank-api) | +| HuggingFace | [Usage](../docs/providers/huggingface_rerank) | +| Infinity | [Usage](../docs/providers/infinity) | +| vLLM | [Usage](../docs/providers/vllm#rerank-endpoint) | +| DeepInfra | [Usage](../docs/providers/deepinfra#rerank-endpoint) | +| Vertex AI | [Usage](../docs/providers/vertex#rerank-api) | +| Fireworks AI | [Usage](../docs/providers/fireworks_ai#rerank-endpoint) | +| Voyage AI | [Usage](../docs/providers/voyage#rerank) | +| IBM watsonx.ai | [Usage](../docs/providers/watsonx/rerank) | \ No newline at end of file diff --git a/litellm/__init__.py b/litellm/__init__.py index 4aaddc3da76..c13ae8c2d1c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1333,6 +1333,7 @@ if TYPE_CHECKING: from .llms.vertex_ai.rerank.transformation import VertexAIRerankConfig as VertexAIRerankConfig from .llms.fireworks_ai.rerank.transformation import FireworksAIRerankConfig as FireworksAIRerankConfig from .llms.voyage.rerank.transformation import VoyageRerankConfig as VoyageRerankConfig + from .llms.watsonx.rerank.transformation import IBMWatsonXRerankConfig as IBMWatsonXRerankConfig from .llms.clarifai.chat.transformation import ClarifaiConfig as ClarifaiConfig from .llms.ai21.chat.transformation import AI21ChatConfig as AI21ChatConfig from .llms.meta_llama.chat.transformation import LlamaAPIConfig as LlamaAPIConfig diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 2af6ed8f09e..a3dc12c23a3 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -155,6 +155,7 @@ LLM_CONFIG_NAMES = ( "VertexAIRerankConfig", "FireworksAIRerankConfig", "VoyageRerankConfig", + "IBMWatsonXRerankConfig", "ClarifaiConfig", "AI21ChatConfig", "LlamaAPIConfig", @@ -671,6 +672,7 @@ _LLM_CONFIGS_IMPORT_MAP = { "FireworksAIRerankConfig", ), "VoyageRerankConfig": (".llms.voyage.rerank.transformation", "VoyageRerankConfig"), + "IBMWatsonXRerankConfig": (".llms.watsonx.rerank.transformation", "IBMWatsonXRerankConfig"), "ClarifaiConfig": (".llms.clarifai.chat.transformation", "ClarifaiConfig"), "AI21ChatConfig": (".llms.ai21.chat.transformation", "AI21ChatConfig"), "LlamaAPIConfig": (".llms.meta_llama.chat.transformation", "LlamaAPIConfig"), diff --git a/litellm/llms/watsonx/__init__.py b/litellm/llms/watsonx/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/watsonx/chat/__init__.py b/litellm/llms/watsonx/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/watsonx/completion/__init__.py b/litellm/llms/watsonx/completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/watsonx/embed/__init__.py b/litellm/llms/watsonx/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/watsonx/rerank/__init__.py b/litellm/llms/watsonx/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py new file mode 100644 index 00000000000..7b4c2a07c3c --- /dev/null +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -0,0 +1,204 @@ +""" +Transformation logic for IBM watsonx.ai's /ml/v1/text/rerank endpoint. + +Docs - https://cloud.ibm.com/apidocs/watsonx-ai#text-rerank +""" + +import uuid +from typing import Any, Dict, List, Optional, Union, cast + +import httpx + +from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj +from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.watsonx import ( + WatsonXAIEndpoint, +) +from litellm.types.rerank import ( + RerankResponse, + RerankResponseMeta, + RerankTokens, +) + +from ..common_utils import IBMWatsonXMixin, _generate_watsonx_token, _get_api_params + + +class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): + """ + IBM watsonx.ai Rerank API configuration + """ + + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: + base_url = self._get_base_url(api_base=api_base) + endpoint = WatsonXAIEndpoint.RERANK.value + + url = base_url.rstrip("/") + endpoint + + params = optional_params or {} + + complete_url = self._add_api_version_to_url(url=url, api_version=(params.get("api_version", None))) + return complete_url + + def get_supported_cohere_rerank_params(self, model: str) -> list: + return [ + "query", + "documents", + "top_n", + "return_documents", + "max_tokens_per_doc", + ] + + def validate_environment( # type: ignore[override] + self, + headers: dict, + model: str, + api_key: Optional[str] = None, + optional_params: Optional[dict] = None, + ) -> Dict: + optional_params = optional_params or {} + + default_headers = { + "Content-Type": "application/json", + "Accept": "application/json", + } + + if "Authorization" in headers: + return {**default_headers, **headers} + token = cast( + Optional[str], + optional_params.pop("token", None) or get_secret_str("WATSONX_TOKEN"), + ) + zen_api_key = cast( + Optional[str], + optional_params.pop("zen_api_key", None) or get_secret_str("WATSONX_ZENAPIKEY"), + ) + if token: + headers["Authorization"] = f"Bearer {token}" + elif zen_api_key: + headers["Authorization"] = f"ZenApiKey {zen_api_key}" + else: + token = _generate_watsonx_token(api_key=api_key, token=token) + # build auth headers + headers["Authorization"] = f"Bearer {token}" + return {**default_headers, **headers} + + def map_cohere_rerank_params( + self, + non_default_params: Optional[dict], + model: str, + drop_params: bool, + query: str, + documents: List[Union[str, Dict[str, Any]]], + custom_llm_provider: Optional[str] = None, + top_n: Optional[int] = None, + rank_fields: Optional[List[str]] = None, + return_documents: Optional[bool] = True, + max_chunks_per_doc: Optional[int] = None, + max_tokens_per_doc: Optional[int] = None, + ) -> Dict: + """ + Map Cohere rerank params to IBM watsonx.ai rerank params + """ + optional_rerank_params = {} + if non_default_params is not None: + for k, v in non_default_params.items(): + if k == "query" and v is not None: + optional_rerank_params["query"] = v + elif k == "documents" and v is not None: + optional_rerank_params["inputs"] = [ + {"text": el} if isinstance(el, str) else el for el in v + ] + elif k == "top_n" and v is not None: + optional_rerank_params.setdefault("parameters", {}).setdefault("return_options", {})["top_n"] = v + elif k == "return_documents" and v is not None and isinstance(v, bool): + optional_rerank_params.setdefault("parameters", {}).setdefault("return_options", {})["inputs"] = v + elif k == "max_tokens_per_doc" and v is not None: + optional_rerank_params.setdefault("parameters", {})["truncate_input_tokens"] = v + + # IBM watsonx.ai require one of below parameters + elif k == "project_id" and v is not None: + optional_rerank_params["project_id"] = v + elif k == "space_id" and v is not None: + optional_rerank_params["space_id"] = v + + return dict(optional_rerank_params) + + def transform_rerank_request( + self, + model: str, + optional_rerank_params: Dict, + headers: dict, + ) -> dict: + """ + Transform request to IBM watsonx.ai rerank format + """ + watsonx_api_params = _get_api_params(params=optional_rerank_params, model=model) + watsonx_auth_payload = self._prepare_payload( + model=model, + api_params=watsonx_api_params, + ) + + return optional_rerank_params | watsonx_auth_payload + + def transform_rerank_response( + self, + model: str, + raw_response: httpx.Response, + model_response: RerankResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str] = None, + request_data: dict = {}, + optional_params: dict = {}, + litellm_params: dict = {}, + ) -> RerankResponse: + """ + Transform IBM watsonx.ai rerank response to LiteLLM RerankResponse format + """ + try: + raw_response_json = raw_response.json() + except Exception as e: + raise self.get_error_class( + error_message=f"Failed to parse response: {str(e)}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + _results: Optional[List[dict]] = raw_response_json.get("results") + if _results is None: + raise ValueError(f"No results found in the response={raw_response_json}") + + transformed_results = [] + + for result in _results: + transformed_result: Dict[str, Any] = { + "index": result["index"], + "relevance_score": result["score"], + } + + if "input" in result: + if isinstance(result["input"], str): + transformed_result["document"] = {"text": result["input"]} + else: + transformed_result["document"] = result["input"] + + transformed_results.append(transformed_result) + + response_id = raw_response_json.get("id") or raw_response_json.get("model_id") or str(uuid.uuid4()) + + # Extract usage information + _tokens = RerankTokens( + input_tokens=raw_response_json.get("input_token_count", 0), + ) + rerank_meta = RerankResponseMeta(tokens=_tokens) + + return RerankResponse( + id=response_id, + results=transformed_results, # type: ignore + meta=rerank_meta, + ) diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 8910d37fbe7..f47fd6323f0 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -10,6 +10,7 @@ from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.together_ai.rerank.handler import TogetherAIRerank +from litellm.llms.watsonx.common_utils import IBMWatsonXMixin from litellm.rerank_api.rerank_utils import get_optional_rerank_params from litellm.secret_managers.main import get_secret, get_secret_str from litellm.types.rerank import RerankResponse @@ -29,7 +30,7 @@ async def arerank( model: str, query: str, documents: List[Union[str, Dict[str, Any]]], - custom_llm_provider: Optional[Literal["cohere", "together_ai", "deepinfra", "fireworks_ai", "voyage"]] = None, + custom_llm_provider: Optional[Literal["cohere", "together_ai", "deepinfra", "fireworks_ai", "voyage", "watsonx"]] = None, top_n: Optional[int] = None, rank_fields: Optional[List[str]] = None, return_documents: Optional[bool] = None, @@ -85,6 +86,7 @@ def rerank( # noqa: PLR0915 "deepinfra", "fireworks_ai", "voyage", + "watsonx", ] ] = None, top_n: Optional[int] = None, @@ -478,6 +480,31 @@ def rerank( # noqa: PLR0915 or get_secret_str("VOYAGE_API_BASE") ) + response = base_llm_http_handler.rerank( + model=model, + custom_llm_provider=_custom_llm_provider, + provider_config=rerank_provider_config, + optional_rerank_params=optional_rerank_params, + logging_obj=litellm_logging_obj, + timeout=optional_params.timeout, + api_key=api_key, + api_base=api_base, + _is_async=_is_async, + headers=headers or litellm.headers or {}, + client=client, + model_response=model_response, + ) + elif _custom_llm_provider == litellm.LlmProviders.WATSONX: + credentials = IBMWatsonXMixin.get_watsonx_credentials( + optional_params=dict(optional_params), api_key=dynamic_api_key, api_base=dynamic_api_base + ) + + api_key = credentials["api_key"] + api_base = credentials["api_base"] + + if credentials.get("token") is not None: + optional_rerank_params["token"] = credentials["token"] + response = base_llm_http_handler.rerank( model=model, custom_llm_provider=_custom_llm_provider, diff --git a/litellm/types/llms/watsonx.py b/litellm/types/llms/watsonx.py index 137090b032e..21e58500c6f 100644 --- a/litellm/types/llms/watsonx.py +++ b/litellm/types/llms/watsonx.py @@ -63,6 +63,7 @@ class WatsonXAIEndpoint(str, Enum): EMBEDDINGS = "/ml/v1/text/embeddings" PROMPTS = "/ml/v1/prompts" AVAILABLE_MODELS = "/ml/v1/foundation_model_specs" + RERANK = "/ml/v1/text/rerank" class WatsonXModelPattern(str, Enum): diff --git a/litellm/utils.py b/litellm/utils.py index 0fd21f09919..0e8dada2352 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8145,6 +8145,8 @@ class ProviderConfigManager: return litellm.FireworksAIRerankConfig() elif litellm.LlmProviders.VOYAGE == provider: return litellm.VoyageRerankConfig() + elif litellm.LlmProviders.WATSONX == provider: + return litellm.IBMWatsonXRerankConfig() return litellm.CohereRerankConfig() @staticmethod diff --git a/tests/test_litellm/llms/watsonx/rerank/__init__.py b/tests/test_litellm/llms/watsonx/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py b/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py new file mode 100644 index 00000000000..f50966279b4 --- /dev/null +++ b/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py @@ -0,0 +1,236 @@ +""" +Tests for IBM watsonx.ai rerank transformation functionality. +""" +import json +import re +import uuid +from unittest.mock import MagicMock + +import httpx +import pytest + +from litellm.llms.watsonx.common_utils import ( + WatsonXAIError, +) +from litellm.llms.watsonx.rerank.transformation import IBMWatsonXRerankConfig +from litellm.types.rerank import RerankResponse + + +class TestIBMWatsonXRerankTransform: + def setup_method(self): + self.config = IBMWatsonXRerankConfig() + self.model = "watsonx/cross-encoder/ms-marco-minilm-l-12-v2" + + def test_get_complete_url(self): + """Test URL generation for IBM watsonx.ai rerank API.""" + + api_base = "https://us-south.ml.cloud.ibm.com" + model = "watsonx/cross-encoder/ms-marco-minilm-l-12-v2" + url = self.config.get_complete_url(api_base, model) + assert url == "https://us-south.ml.cloud.ibm.com/ml/v1/text/rerank?version=2024-03-13" + + def test_map_cohere_rerank_params_basic(self): + """Test basic parameter mapping for IBM watsonx.ai rerank.""" + params = self.config.map_cohere_rerank_params( + non_default_params={ + "query": "hello", + "documents": ["hello", "world"], + "top_n": 2, + "return_documents": True, + "max_tokens_per_doc": 100, + }, + model="test", + drop_params=False, + query="hello", + documents=["hello", "world"], + ) + assert params["query"] == "hello" + assert params["inputs"] == [{"text": "hello"}, {"text": "world"}] + assert params["parameters"]["return_options"]["top_n"] == 2 + assert params["parameters"]["return_options"]["inputs"] is True + assert params["parameters"]["truncate_input_tokens"] == 100 + + def test_transform_rerank_request(self): + """Test request transformation for IBM watsonx.ai format.""" + optional_params = { + "query": "What is the capital of France?", + "documents": [ + "Paris is the capital of France.", + "France is a country in Europe.", + ], + "top_n": 2, + "return_documents": True, + "project_id": uuid.uuid4(), + } + + request_body = self.config.transform_rerank_request( + model="cross-encoder/ms-marco-minilm-l-12-v2", optional_rerank_params=optional_params, headers={} + ) + + assert request_body["model_id"] == "cross-encoder/ms-marco-minilm-l-12-v2" + assert request_body["project_id"] is not None + assert request_body["query"] == "What is the capital of France?" + assert request_body["documents"] == optional_params["documents"] + assert request_body["top_n"] == 2 + assert request_body["return_documents"] is True + + def test_transform_rerank_request_missing_scope(self): + """Test that transform_rerank_request raises error for missing scope.""" + optional_params = { + "documents": ["doc1"], + } + expected_error_msg = re.escape( + "Watsonx project_id and space_id not set. Set WX_PROJECT_ID or WX_SPACE_ID in environment variables or pass in as a parameter." + ) + + with pytest.raises(WatsonXAIError, match=expected_error_msg): + self.config.transform_rerank_request(model=self.model, optional_rerank_params=optional_params, headers={}) + + def test_transform_rerank_response_success(self): + """Test successful response transformation.""" + # Mock IBM watsonx.ai response format + response_data = { + "model_id": self.model, + "results": [ + { + "index": 0, + "score": 6.53515625, + "input": {"text": "Python is great for beginners due to simple syntax."}, + }, + {"index": 1, "score": -7.1875, "input": {"text": "JavaScript runs in browsers and is versatile."}}, + ], + "input_token_count": 62, + } + + # Create mock httpx response + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + + # Create mock logging object + mock_logging = MagicMock() + + model_response = RerankResponse() + + result = self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + ) + + # Verify response structure + # IBM watsonx.ai doesn't return "id", so it uses "model" as the id + assert result.id == "watsonx/cross-encoder/ms-marco-minilm-l-12-v2" + assert len(result.results) == 2 + assert result.results[0]["index"] == 0 + assert result.results[0]["relevance_score"] == 6.53515625 + assert result.results[0]["document"]["text"] == "Python is great for beginners due to simple syntax." + assert result.results[1]["index"] == 1 + assert result.results[1]["relevance_score"] == -7.1875 + assert result.results[1]["document"]["text"] == "JavaScript runs in browsers and is versatile." + + # # Verify metadata + assert result.meta["tokens"]["input_tokens"] == 62 + + def test_transform_rerank_response_without_documents(self): + """Test response transformation when return_documents is False.""" + response_data = { + "model_id": self.model, + "results": [ + { + "index": 0, + "score": 6.53515625, + }, + { + "index": 1, + "score": -7.1875, + }, + ], + "input_token_count": 62, + } + + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + + mock_logging = MagicMock() + model_response = RerankResponse() + + result = self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + ) + + # Verify response structure + # IBM watsonx.ai doesn't return "id", so it uses "model" as the id + assert result.id == "watsonx/cross-encoder/ms-marco-minilm-l-12-v2" + assert len(result.results) == 2 + + assert result.results[0]["index"] == 0 + assert result.results[0]["relevance_score"] == 6.53515625 + assert "document" not in result.results[0] + + assert result.results[1]["index"] == 1 + assert result.results[1]["relevance_score"] == -7.1875 + assert "document" not in result.results[1] + + def test_transform_rerank_response_missing_results(self): + """Test that missing results raises ValueError.""" + response_data = { + "model": self.model, + "usage": {"total_tokens": 10}, + } + + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + + mock_logging = MagicMock() + model_response = RerankResponse() + + expected_error_msg = re.escape("No results found") + + with pytest.raises(ValueError, match=expected_error_msg): + self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + ) + + def test_transform_rerank_response_invalid_json(self): + """Test error handling for invalid JSON response.""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "doc", 0) + mock_response.text = "Invalid JSON response" + mock_response.status_code = 500 + mock_response.headers = {} + + mock_logging = MagicMock() + model_response = RerankResponse() + + expected_error_msg = re.escape("Failed to parse response") + + with pytest.raises(Exception, match=expected_error_msg): + self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + ) + + def test_get_supported_cohere_rerank_params(self): + """Test getting supported parameters for IBM watsonx.ai rerank.""" + supported_params = self.config.get_supported_cohere_rerank_params(self.model) + assert "query" in supported_params + assert "documents" in supported_params + assert "top_n" in supported_params + assert "return_documents" in supported_params + assert "max_tokens_per_doc" in supported_params + assert len(supported_params) == 5 From f162371b93df9bd7644ab763025aa6f4b1b9a67e Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Mon, 16 Feb 2026 22:15:46 -0600 Subject: [PATCH 03/16] fix(pod-lock): make release lock compare-and-delete atomic (#21226) --- .../db_transaction_queue/pod_lock_manager.py | 77 +++++++++++-------- .../test_pod_lock_manager.py | 39 ++++++++++ 2 files changed, 85 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index bb5424b0e90..5fee1b28e71 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -24,6 +24,15 @@ class PodLockManager: def __init__(self, redis_cache: Optional[RedisCache] = None): self.pod_id = str(uuid.uuid4()) self.redis_cache = redis_cache + self._release_lock_script: Optional[Any] = None + + _COMPARE_AND_DELETE_LOCK_SCRIPT = """ +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +else + return 0 +end +""" @staticmethod def get_redis_lock_key(cronjob_id: str) -> str: @@ -106,39 +115,20 @@ class PodLockManager: cronjob_id, ) lock_key = PodLockManager.get_redis_lock_key(cronjob_id) - - current_value = await self.redis_cache.async_get_cache(lock_key) - if current_value is not None: - if isinstance(current_value, bytes): - current_value = current_value.decode("utf-8") - if current_value == self.pod_id: - result = await self.redis_cache.async_delete_cache(lock_key) - if result == 1: - verbose_proxy_logger.info( - "Pod %s successfully released Redis lock for cronjob_id=%s", - self.pod_id, - cronjob_id, - ) - self._emit_released_lock_event( - cronjob_id=cronjob_id, - pod_id=self.pod_id, - ) - else: - verbose_proxy_logger.debug( - "Pod %s failed to release Redis lock for cronjob_id=%s", - self.pod_id, - cronjob_id, - ) - else: - verbose_proxy_logger.debug( - "Pod %s cannot release Redis lock for cronjob_id=%s because it is held by pod %s", - self.pod_id, - cronjob_id, - current_value, - ) + result = await self._compare_and_delete_lock(lock_key=lock_key) + if result == 1: + verbose_proxy_logger.info( + "Pod %s successfully released Redis lock for cronjob_id=%s", + self.pod_id, + cronjob_id, + ) + self._emit_released_lock_event( + cronjob_id=cronjob_id, + pod_id=self.pod_id, + ) else: verbose_proxy_logger.debug( - "Pod %s attempted to release Redis lock for cronjob_id=%s, but no lock was found", + "Pod %s failed to release Redis lock for cronjob_id=%s (lock missing or held by another pod)", self.pod_id, cronjob_id, ) @@ -147,6 +137,31 @@ class PodLockManager: f"Error releasing Redis lock for {cronjob_id}: {e}" ) + async def _compare_and_delete_lock(self, lock_key: str) -> int: + """ + Atomically delete lock key only if current pod owns it. + + Falls back to get/delete for non-RedisCache implementations that do not + expose Lua script registration. + """ + script_register = getattr(self.redis_cache, "async_register_script", None) + if callable(script_register): + if self._release_lock_script is None: + self._release_lock_script = script_register( + self._COMPARE_AND_DELETE_LOCK_SCRIPT + ) + script_callable = self._release_lock_script + result = await script_callable(keys=[lock_key], args=[self.pod_id]) + return int(result or 0) + + current_value = await self.redis_cache.async_get_cache(lock_key) # type: ignore + if isinstance(current_value, bytes): + current_value = current_value.decode("utf-8") + if current_value != self.pod_id: + return 0 + result = await self.redis_cache.async_delete_cache(lock_key) # type: ignore + return int(result or 0) + @staticmethod def _emit_acquired_lock_event(cronjob_id: str, pod_id: str): asyncio.create_task( diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py index e83fd75c3a0..7790961eb16 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -307,3 +307,42 @@ async def test_lock_takeover_race_condition(mock_redis): cronjob_id="test_job", ) assert result2 == False + + +@pytest.mark.asyncio +async def test_release_lock_uses_atomic_compare_delete_script_when_available( + pod_lock_manager, mock_redis +): + """ + Test that release_lock prefers atomic compare-and-delete Lua script when + redis cache exposes script registration. + """ + script_callable = AsyncMock(return_value=1) + mock_redis.async_register_script = MagicMock(return_value=script_callable) + + await pod_lock_manager.release_lock(cronjob_id="test_job") + + lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") + mock_redis.async_register_script.assert_called_once_with( + PodLockManager._COMPARE_AND_DELETE_LOCK_SCRIPT + ) + script_callable.assert_called_once_with( + keys=[lock_key], args=[pod_lock_manager.pod_id] + ) + mock_redis.async_get_cache.assert_not_called() + mock_redis.async_delete_cache.assert_not_called() + + +@pytest.mark.asyncio +async def test_release_lock_reuses_registered_script(pod_lock_manager, mock_redis): + """ + Test script registration is cached on manager instance and reused. + """ + script_callable = AsyncMock(return_value=0) + mock_redis.async_register_script = MagicMock(return_value=script_callable) + + await pod_lock_manager.release_lock(cronjob_id="test_job") + await pod_lock_manager.release_lock(cronjob_id="test_job") + + assert mock_redis.async_register_script.call_count == 1 + assert script_callable.call_count == 2 From d184b3cae7ef0fecfc888b560a3179b3a1cee511 Mon Sep 17 00:00:00 2001 From: sahukanishka <34833039+sahukanishka@users.noreply.github.com> Date: Tue, 17 Feb 2026 09:47:58 +0530 Subject: [PATCH 04/16] fix: preserve provider_specific_fields from proxy responses (#21153) (#21220) Co-authored-by: kanishka sahu Co-authored-by: Cursor --- .../convert_dict_to_response.py | 6 +- .../test_convert_dict_to_chat_completion.py | 187 ++++++++++++++++++ 2 files changed, 192 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 25ad0a570cb..a6e502a32b3 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -546,7 +546,11 @@ def convert_to_model_response_object( # noqa: PLR0915 message = litellm.Message(content=json_mode_content_str) finish_reason = "stop" if message is None: - provider_specific_fields = {} + # Preserve provider_specific_fields if already present + # in the response (e.g. from proxy passthrough) + provider_specific_fields = dict( + choice["message"].get("provider_specific_fields", None) or {} + ) message_keys = Message.model_fields.keys() for field in choice["message"].keys(): if field not in message_keys: diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index c151150f634..3b2087d25e9 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -1037,6 +1037,193 @@ def test_convert_to_model_response_object_with_empty_dict_error(): assert result.choices[0].message.content == "Hello!" +def test_convert_to_model_response_object_preserves_provider_specific_fields_from_proxy(): + """ + Test that provider_specific_fields (e.g. Anthropic citations) are preserved + when the response already contains them (e.g. from a proxy passthrough). + + Regression test for https://github.com/BerriAI/litellm/issues/21153 + """ + citations = [ + [ + { + "type": "web_search_result_location", + "cited_text": "The Sony WH-1000XM5 remains one of the best...", + "url": "https://example.com/headphones-review", + "title": "Best Headphones 2025", + "supported_text": "Based on current reviews...", + } + ], + ] + web_search_results = [ + { + "url": "https://example.com/headphones-review", + "title": "Best Headphones 2025", + "snippet": "The Sony WH-1000XM5 remains one of the best...", + } + ] + + response_object = { + "id": "chatcmpl-proxy-123", + "object": "chat.completion", + "created": 1728933352, + "model": "anthropic/claude-opus-4-5-20251101", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Based on current reviews, the Sony WH-1000XM5 remains one of the best headphones.", + "tool_calls": [ + { + "id": "call_ws_123", + "type": "function", + "function": { + "name": "web_search", + "arguments": '{"query": "best headphones 2025"}', + }, + } + ], + "provider_specific_fields": { + "citations": citations, + "web_search_results": web_search_results, + }, + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 50, + "completion_tokens": 20, + "total_tokens": 70, + }, + } + + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + assert result.id == "chatcmpl-proxy-123" + + choice = result.choices[0] + assert choice.message.content == "Based on current reviews, the Sony WH-1000XM5 remains one of the best headphones." + assert choice.message.provider_specific_fields is not None + assert "citations" in choice.message.provider_specific_fields + assert choice.message.provider_specific_fields["citations"] == citations + assert "web_search_results" in choice.message.provider_specific_fields + assert choice.message.provider_specific_fields["web_search_results"] == web_search_results + + +def test_convert_to_model_response_object_provider_specific_fields_merges_extra_keys(): + """ + Test that provider_specific_fields from the response are merged with + any extra non-standard keys present in the message dict. + + Regression test for https://github.com/BerriAI/litellm/issues/21153 + """ + response_object = { + "id": "chatcmpl-merge-123", + "object": "chat.completion", + "created": 1728933352, + "model": "some-model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello!", + "provider_specific_fields": { + "citations": [{"url": "https://example.com"}], + }, + "custom_extra_field": "extra_value", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + } + + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + psf = result.choices[0].message.provider_specific_fields + assert psf is not None + # Both the existing provider_specific_fields and the extra key should be present + assert "citations" in psf + assert psf["citations"] == [{"url": "https://example.com"}] + assert "custom_extra_field" in psf + assert psf["custom_extra_field"] == "extra_value" + + +def test_convert_to_model_response_object_no_provider_specific_fields_still_works(): + """ + Test that responses without provider_specific_fields continue to work as before. + + Ensures the fix for https://github.com/BerriAI/litellm/issues/21153 + doesn't break normal responses. + """ + response_object = { + "id": "chatcmpl-normal-123", + "object": "chat.completion", + "created": 1728933352, + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello!", + "refusal": None, + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + } + + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + psf = result.choices[0].message.provider_specific_fields + # refusal is not a Message model field, so it should be in provider_specific_fields + assert psf is not None + assert "refusal" in psf + + def test_convert_to_model_response_object_with_error_code_only(): """ Test that errors with only a code (no message) are still treated as real errors. From fb2c22eccef55d5924360bbb41e884ed532cb349 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Mon, 16 Feb 2026 22:20:48 -0600 Subject: [PATCH 05/16] perf(router): optimize v2 deployment selection lookup (#21211) --- litellm/router_strategy/lowest_tpm_rpm_v2.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index bf3035fcc9f..70d4c6751db 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -335,13 +335,14 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ): lowest_tpm = float("inf") potential_deployments = [] # if multiple deployments have the same low value + deployment_lookup = { + deployment.get("model_info", {}).get("id"): deployment + for deployment in healthy_deployments + } for item, item_tpm in all_deployments.items(): ## get the item from model list - _deployment = None item = item.split(":")[0] - for m in healthy_deployments: - if item == m["model_info"]["id"]: - _deployment = m + _deployment = deployment_lookup.get(item) if _deployment is None: continue # skip to next one elif item_tpm is None: From ddb48fa1164cc5730336b2b6a12c7331b302e950 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Mon, 16 Feb 2026 22:22:23 -0600 Subject: [PATCH 06/16] perf(router): use set membership in team deployment filter (#21210) --- litellm/router_utils/common_utils.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 10acc343abd..3b0273f4c5d 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -58,7 +58,7 @@ def filter_team_based_models( request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get( "user_api_key_team_id" ) - ids_to_remove = [] + ids_to_remove = set() if isinstance(healthy_deployments, dict): return healthy_deployments for deployment in healthy_deployments: @@ -67,7 +67,7 @@ def filter_team_based_models( if model_team_id is None: continue if model_team_id != request_team_id: - ids_to_remove.append(deployment.get("model_info", {}).get("id")) + ids_to_remove.add(_model_info.get("id")) return [ deployment @@ -125,4 +125,3 @@ def filter_web_search_deployments( if len(healthy_deployments) > 0 and len(final_deployments) == 0: verbose_logger.warning("No deployments support web search for request") return final_deployments - From 4978df8ebdcdce1865b487ae2ae411fefeca518e Mon Sep 17 00:00:00 2001 From: Nick Amabile Date: Mon, 16 Feb 2026 23:28:34 -0500 Subject: [PATCH 07/16] fix: add `store` to OPENAI_CHAT_COMPLETION_PARAMS (#21195) The OpenAI `store` parameter (used for storing completions for distillation/evals) was missing from `OPENAI_CHAT_COMPLETION_PARAMS`. This caused it to be unrecognized by `get_standard_openai_params()` and the `litellm_proxy` provider config. It also meant that code paths using this list (rather than `DEFAULT_CHAT_COMPLETION_PARAM_VALUES`) would treat `store` as a provider-specific parameter and forward it to non-OpenAI providers like Anthropic, resulting in: "store: Extra inputs are not permitted" Fixes #19700 --- litellm/constants.py | 1 + tests/llm_translation/test_optional_params.py | 88 ++++++++++++++++--- 2 files changed, 77 insertions(+), 12 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index a4a0e7882ea..7c21111d313 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -576,6 +576,7 @@ OPENAI_CHAT_COMPLETION_PARAMS = [ "thinking", "web_search_options", "service_tier", + "store", ] OPENAI_TRANSCRIPTION_PARAMS = [ diff --git a/tests/llm_translation/test_optional_params.py b/tests/llm_translation/test_optional_params.py index 4699c31c378..6ecac7b36a2 100644 --- a/tests/llm_translation/test_optional_params.py +++ b/tests/llm_translation/test_optional_params.py @@ -1894,28 +1894,28 @@ def test_validate_openai_optional_params_stop_truncation(): result = validate_openai_optional_params(stop=stop_sequences) assert result == ["stop1", "stop2", "stop3", "stop4"] assert len(result) == 4 - + # Test with exactly 4 stop sequences - should not truncate stop_sequences_4 = ["stop1", "stop2", "stop3", "stop4"] result = validate_openai_optional_params(stop=stop_sequences_4) assert result == ["stop1", "stop2", "stop3", "stop4"] assert len(result) == 4 - + # Test with less than 4 stop sequences - should not truncate stop_sequences_2 = ["stop1", "stop2"] result = validate_openai_optional_params(stop=stop_sequences_2) assert result == ["stop1", "stop2"] assert len(result) == 2 - + # Test with single stop sequence as string - should return as is stop_string = "stop1" result = validate_openai_optional_params(stop=stop_string) assert result == "stop1" - + # Test with None - should return None result = validate_openai_optional_params(stop=None) assert result is None - + # Test with empty list - should return empty list result = validate_openai_optional_params(stop=[]) assert result == [] @@ -1928,7 +1928,7 @@ def test_validate_openai_optional_params_disable_stop_sequence_limit(): """ # Save original value original_value = litellm.disable_stop_sequence_limit - + try: # Test with disable_stop_sequence_limit = True - should NOT truncate litellm.disable_stop_sequence_limit = True @@ -1936,7 +1936,7 @@ def test_validate_openai_optional_params_disable_stop_sequence_limit(): result = validate_openai_optional_params(stop=stop_sequences) assert result == ["stop1", "stop2", "stop3", "stop4", "stop5", "stop6"] assert len(result) == 6 - + # Test with disable_stop_sequence_limit = False - should truncate to 4 litellm.disable_stop_sequence_limit = False stop_sequences = ["stop1", "stop2", "stop3", "stop4", "stop5", "stop6"] @@ -1965,19 +1965,83 @@ def test_validate_openai_optional_params_integration(): mock_response.usage.prompt_tokens = 10 mock_response.usage.completion_tokens = 5 mock_response.usage.total_tokens = 15 - - mock_client.return_value.chat.completions.create.return_value = mock_response - + + mock_client.return_value.chat.completions.create.return_value = ( + mock_response + ) + # Call completion with more than 4 stop sequences response = litellm.completion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello"}], stop=["stop1", "stop2", "stop3", "stop4", "stop5", "stop6"], - mock_response="Test response" # This will use mock + mock_response="Test response", # This will use mock ) - + # Verify the call was made (stop sequences should be truncated internally) assert response is not None except Exception as e: # Should not raise an exception pytest.fail(f"validate_openai_optional_params integration failed: {e}") + + +def test_drop_store_param_for_anthropic(): + """ + Test that the OpenAI-specific `store` parameter is correctly dropped + when calling Anthropic with drop_params=True. + + `store` is an OpenAI Chat Completion parameter (for storing completions + for distillation/evals) that Anthropic does not support. Without proper + handling, it leaks through to the Anthropic API and causes a + "store: Extra inputs are not permitted" error. + + Ref: https://github.com/BerriAI/litellm/issues/19700 + """ + optional_params = get_optional_params( + model="claude-sonnet-4-20250514", + custom_llm_provider="anthropic", + drop_params=True, + store=True, + ) + assert "store" not in optional_params + + +def test_additional_drop_params_store_for_anthropic(): + """ + Test that `additional_drop_params=["store"]` correctly strips the `store` + parameter for non-OpenAI providers like Anthropic. + + Ref: https://github.com/BerriAI/litellm/issues/19700 + """ + optional_params = get_optional_params( + model="claude-sonnet-4-20250514", + custom_llm_provider="anthropic", + additional_drop_params=["store"], + store=True, + ) + assert "store" not in optional_params + + +def test_store_in_openai_chat_completion_params(): + """ + Test that `store` is recognized as a standard OpenAI Chat Completion + parameter. This ensures it is correctly handled by helper functions + like `get_standard_openai_params()` and provider configs that rely on + `OPENAI_CHAT_COMPLETION_PARAMS`. + + Without `store` in this list, functions that filter by known OpenAI + params will silently drop it for OpenAI calls or incorrectly treat + it as a provider-specific param for non-OpenAI providers. + + Ref: https://github.com/BerriAI/litellm/issues/19700 + """ + from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS + + assert "store" in OPENAI_CHAT_COMPLETION_PARAMS + + # Verify get_standard_openai_params recognizes store + from litellm.utils import get_standard_openai_params + + result = get_standard_openai_params({"store": True, "temperature": 0.7}) + assert "store" in result + assert result["store"] is True From b67c1409388b07460ef487751c18285f6fcf9778 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Mon, 16 Feb 2026 22:30:10 -0600 Subject: [PATCH 08/16] Fix Bedrock service_tier cost propagation (#21172) --- litellm/cost_calculator.py | 5 ++- litellm/llms/bedrock/cost_calculation.py | 13 ++++-- tests/test_litellm/test_cost_calculator.py | 50 ++++++++++++++++++++++ 3 files changed, 62 insertions(+), 6 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index dae0bb1c2c0..02df747792d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -448,7 +448,9 @@ def cost_per_token( # noqa: PLR0915 elif custom_llm_provider == "anthropic": return anthropic_cost_per_token(model=model, usage=usage_block) elif custom_llm_provider == "bedrock": - return bedrock_cost_per_token(model=model, usage=usage_block) + return bedrock_cost_per_token( + model=model, usage=usage_block, service_tier=service_tier + ) elif custom_llm_provider == "openai": return openai_cost_per_token( model=model, usage=usage_block, service_tier=service_tier @@ -2146,4 +2148,3 @@ def handle_realtime_stream_cost_calculation( return total_cost - diff --git a/litellm/llms/bedrock/cost_calculation.py b/litellm/llms/bedrock/cost_calculation.py index b20350d7325..ac99d4e36e7 100644 --- a/litellm/llms/bedrock/cost_calculation.py +++ b/litellm/llms/bedrock/cost_calculation.py @@ -3,7 +3,7 @@ Helper util for handling bedrock-specific cost calculation - e.g.: prompt caching """ -from typing import TYPE_CHECKING, Tuple +from typing import TYPE_CHECKING, Optional, Tuple from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token @@ -11,12 +11,17 @@ if TYPE_CHECKING: from litellm.types.utils import Usage -def cost_per_token(model: str, usage: "Usage") -> Tuple[float, float]: +def cost_per_token( + model: str, usage: "Usage", service_tier: Optional[str] = None +) -> Tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. Follows the same logic as Anthropic's cost per token calculation. """ return generic_cost_per_token( - model=model, usage=usage, custom_llm_provider="bedrock" - ) \ No newline at end of file + model=model, + usage=usage, + custom_llm_provider="bedrock", + service_tier=service_tier, + ) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 74f5cf9bdd7..c2c20485b5e 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1600,6 +1600,56 @@ def test_completion_cost_service_tier_priority(): ), "Costs from params and usage should be similar (both flex)" +def test_completion_cost_service_tier_for_bedrock(): + """Test that Bedrock cost calculation applies service_tier-specific pricing.""" + from litellm import completion_cost + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "bedrock/us-east-1/test-bedrock-service-tier-cost-model" + litellm.register_model( + model_cost={ + model: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "input_cost_per_token_priority": 0.01, + "output_cost_per_token_priority": 0.02, + "input_cost_per_token_flex": 0.0005, + "output_cost_per_token_flex": 0.001, + "litellm_provider": "bedrock", + "max_tokens": 8192, + } + } + ) + + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + response = ModelResponse(usage=usage, model=model) + + default_cost = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + ) + + priority_cost = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + optional_params={"service_tier": "priority"}, + ) + + response_with_flex_tier = ModelResponse(usage=usage, model=model) + setattr(response_with_flex_tier, "service_tier", "flex") + flex_cost = completion_cost( + completion_response=response_with_flex_tier, + model=model, + custom_llm_provider="bedrock", + ) + + assert priority_cost > default_cost > flex_cost > 0 + + def test_gemini_cache_tokens_details_no_negative_values(): """ Test for Issue #18750: Negative text_tokens with Gemini caching From b609f5841b030029fa885149731ed43ab26dddbb Mon Sep 17 00:00:00 2001 From: shin-bot-litellm Date: Mon, 16 Feb 2026 20:31:21 -0800 Subject: [PATCH 09/16] fix: add missing OpenAI chat completion params to OPENAI_CHAT_COMPLETION_PARAMS (#21360) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * allow filtering by user in global usage * add server root path test to github actions * Update .github/workflows/test_server_root_path.yml Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * address greptile review feedback (greploop iteration 1) - Fix HTTPException swallowed by broad except block in get_user_daily_activity and get_user_daily_activity_aggregated: re-raise HTTPException before the generic handler so 403 status codes propagate correctly - Add status_code assertions in non-admin access tests Co-Authored-By: Claude Opus 4.6 (1M context) * address greptile review feedback (greploop iteration 2) - Default user_id to caller's own ID for non-admins instead of 403 when omitted, preserving backward compatibility for API consumers - Apply same fix to aggregated endpoint - Update test to verify defaulting behavior instead of expecting 403 - Add useEffect to sync selectedUserId when auth state settles in UsagePageView to handle async auth initialization Co-Authored-By: Claude Opus 4.6 (1M context) * fixing syntax * remove artifacts * feat: guardrail tracing UI - policy, detection method, match details (#21349) * feat: add GuardrailTracingDetail TypedDict and tracing fields to StandardLoggingGuardrailInformation * feat: add policy_template field to Guardrail config TypedDict * feat: accept GuardrailTracingDetail in base guardrail logging method * feat: populate tracing fields in content filter guardrail * test: add tracing fields tests for custom guardrail base class * test: add tracing fields e2e tests for content filter guardrail * feat: add guardrail tracing UI - policy badges, match details, timeline * feat: redesign GuardrailViewer to Guardrails & Policy Compliance layout Two-column layout with request lifecycle timeline on the left and compact evaluation detail cards on the right. Header shows guardrail count, pass/fail status, total overhead, policy info, and an export button. * feat: add clickable guardrail link in metrics + show policy names * feat: add risk_score field to StandardLoggingGuardrailInformation * feat: compute risk_score in content filter guardrail * feat: display backend risk_score badge on evaluation cards * fix: fallback to frontend risk score when backend doesn't provide one * passing in masster key for api calls * Fix: Add blog as incident report * Fix: Add blog as incident report * remove timeline * feat(models): add github_copilot/gpt-5.3-codex and github_copilot/claude-opus-4.6-fast (#21316) Add missing GitHub Copilot model entries for gpt-5.3-codex (GA) and claude-opus-4.6-fast (Public Preview) to both the root and backup model pricing JSON files. * only tests for /ui * bump: version 1.81.12 → 1.81.13 * Fixing mapped tests * fixing no_config test * fixing container tests * fixing test_basic_openai_responses_api * Adding bedrock thinking budget tokens to docs * fixing regen key tests * fix: add missing OpenAI chat completion params to OPENAI_CHAT_COMPLETION_PARAMS Add store, prompt_cache_key, prompt_cache_retention, safety_identifier, and verbosity to OPENAI_CHAT_COMPLETION_PARAMS list. These params were already in DEFAULT_CHAT_COMPLETION_PARAM_VALUES but missing from the OPENAI_CHAT_COMPLETION_PARAMS list, causing them to be dropped when passed to OpenAI-compatible providers. --------- Co-authored-by: yuneng-jiang Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Ishaan Jaff Co-authored-by: Sameer Kankute Co-authored-by: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Co-authored-by: Krish Dholakia --- .github/workflows/test_server_root_path.yml | 96 ++ .../blog/claude_code_beta_headers/index.md | 363 +++---- docs/my-website/docs/proxy/config_settings.md | 1 + litellm/constants.py | 4 + .../litellm_content_filter/content_filter.py | 159 +-- .../internal_user_endpoints.py | 48 +- .../key_management_endpoints.py | 60 +- litellm/types/utils.py | 4 + model_prices_and_context_window.json | 27 + pyproject.toml | 4 +- .../base_responses_api.py | 2 +- .../containers/test_container_integration.py | 11 +- .../test_meta_llama_chat_transformation.py | 82 +- .../test_publicai_chat_transformation.py | 16 +- .../test_vertex_ai_rerank_transformation.py | 12 +- .../test_internal_user_endpoints.py | 134 ++- tests/test_litellm/proxy/test_proxy_cli.py | 36 +- .../(dashboard)/hooks/users/useUsers.test.ts | 339 +++++++ .../app/(dashboard)/hooks/users/useUsers.ts | 41 + .../components/UsagePageView.test.tsx | 467 +++++++++ .../UsagePage/components/UsagePageView.tsx | 157 ++- .../src/components/networking.tsx | 10 +- .../GuardrailViewer/GuardrailViewer.tsx | 932 +++++++++++------- .../LogDetailsDrawer/LogDetailContent.tsx | 23 +- 24 files changed, 2199 insertions(+), 829 deletions(-) create mode 100644 .github/workflows/test_server_root_path.yml create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts diff --git a/.github/workflows/test_server_root_path.yml b/.github/workflows/test_server_root_path.yml new file mode 100644 index 00000000000..bc559817503 --- /dev/null +++ b/.github/workflows/test_server_root_path.yml @@ -0,0 +1,96 @@ +name: Test Proxy SERVER_ROOT_PATH Routing +permissions: + contents: read + +on: + pull_request: + branches: [main] + +jobs: + test-server-root-path: + runs-on: ubuntu-latest + timeout-minutes: 15 + + strategy: + matrix: + root_path: ["/api/v1", "/llmproxy"] + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Build Docker image + uses: docker/build-push-action@v5 + with: + context: . + file: ./docker/Dockerfile.database + tags: litellm-test:${{ github.sha }} + load: true + cache-from: type=gha + cache-to: type=gha,mode=max + + - name: Start LiteLLM container with SERVER_ROOT_PATH + run: | + docker run -d \ + --name litellm-test \ + -p 4000:4000 \ + -e SERVER_ROOT_PATH="${{ matrix.root_path }}" \ + -e LITELLM_MASTER_KEY="sk-1234" \ + litellm-test:${{ github.sha }} \ + --detailed_debug + + - name: Wait for container to be healthy + run: | + echo "Waiting for LiteLLM to start..." + max_attempts=30 + attempt=0 + + while [ $attempt -lt $max_attempts ]; do + if docker logs litellm-test 2>&1 | grep -q "Uvicorn running"; then + echo "LiteLLM started successfully" + break + fi + attempt=$((attempt + 1)) + echo "Attempt $attempt/$max_attempts - waiting for server to start..." + sleep 2 + done + + if [ $attempt -eq $max_attempts ]; then + echo "Server failed to start within timeout" + docker logs litellm-test + exit 1 + fi + + sleep 5 + + - name: Show container logs + if: always() + run: docker logs litellm-test + + - name: Test UI endpoint with root path + run: | + ROOT_PATH="${{ matrix.root_path }}" + echo "Testing UI at: http://localhost:4000${ROOT_PATH}/ui/" + + for i in 1 2 3; do + content=$(curl -sL --max-time 5 -H "Authorization: Bearer sk-1234" "http://localhost:4000${ROOT_PATH}/ui/") + if echo "$content" | grep -q -E "(html|>LP: Request with beta headers Note over CC,LP: anthropic-beta: header1,header2,header3 - + + LP->>Provider: Forward ALL headers (no validation) + Note over LP,Provider: anthropic-beta: header1,header2,header3 + + Provider-->>LP: ❌ Error: invalid beta flag + LP-->>CC: Request fails +``` + +Requests succeeded for Anthropic (native support) but failed for other providers when Claude Code sent headers those providers didn't support. + +--- + +## Root cause + +LiteLLM lacked provider-specific beta header validation. When Claude Code introduced new beta features or sent headers that specific providers didn't support, those headers were blindly forwarded, causing provider API errors. + +--- + +## Remediation + +| # | Action | Status | Code | +|---|---|---|---| +| 1 | Create `anthropic_beta_headers_config.json` with provider-specific mappings | ✅ Done | [`anthropic_beta_headers_config.json`](https://github.com/BerriAI/litellm/blob/main/litellm/anthropic_beta_headers_config.json) | +| 2 | Implement strict validation: headers must be explicitly mapped to be forwarded | ✅ Done | [`litellm_logging.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm_core_utils/litellm_logging.py) | +| 3 | Add `/reload/anthropic_beta_headers` endpoint for dynamic config updates | ✅ Done | Proxy management endpoints | +| 4 | Add `/schedule/anthropic_beta_headers_reload` for automatic periodic updates | ✅ Done | Proxy management endpoints | +| 5 | Support `LITELLM_ANTHROPIC_BETA_HEADERS_URL` for custom config sources | ✅ Done | Environment configuration | +| 6 | Support `LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS` for air-gapped deployments | ✅ Done | Environment configuration | + +Now LiteLLM validates and transforms headers per-provider: + +```mermaid +sequenceDiagram + participant CC as Claude Code + participant LP as LiteLLM (new behavior) + participant Config as Beta Headers Config + participant Provider as Provider (Bedrock/Azure/Vertex) + + CC->>LP: Request with beta headers + Note over CC,LP: anthropic-beta: header1,header2,header3 + LP->>Config: Load header mapping for provider Config-->>LP: Returns mapping (header→value or null) - + Note over LP: Validate & Transform:
1. Check if header exists in mapping
2. Filter out null values
3. Map to provider-specific names - + LP->>Provider: Request with filtered & mapped headers Note over LP,Provider: anthropic-beta: mapped-header2
(header1, header3 filtered out) - - Provider-->>LP: Success response + + Provider-->>LP: ✅ Success response LP-->>CC: Response ``` -### Filtering Rules +--- -1. **Header must exist in mapping**: Unknown headers are filtered out -2. **Header must have non-null value**: Headers with `null` values are filtered out -3. **Header transformation**: Headers are mapped to provider-specific names (e.g., `advanced-tool-use-2025-11-20` → `tool-search-tool-2025-10-19` for Bedrock) +## Dynamic configuration updates -### Example +A key improvement is zero-downtime configuration updates. When Anthropic releases new beta features, users can update their configuration without restarting: -Request with headers: -``` -anthropic-beta: advanced-tool-use-2025-11-20,computer-use-2025-01-24,unknown-header -``` - -For Bedrock Converse: -- ✅ `computer-use-2025-01-24` → `computer-use-2025-01-24` (supported, passed through) -- ❌ `advanced-tool-use-2025-11-20` → filtered out (null value in config) -- ❌ `unknown-header` → filtered out (not in config) - -Result sent to Bedrock: -``` -anthropic-beta: computer-use-2025-01-24 -``` - -## Dynamic Configuration Management (No Restart Required!) - -### Environment Variables - -Control how LiteLLM loads the beta headers configuration: - -| Variable | Description | Default | -|----------|-------------|---------| -| `LITELLM_ANTHROPIC_BETA_HEADERS_URL` | URL to fetch config from | GitHub main branch | -| `LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS` | Set to `True` to use local config only | `False` | - -**Example: Use Custom Config URL** ```bash -export LITELLM_ANTHROPIC_BETA_HEADERS_URL="https://your-company.com/custom-beta-headers.json" +# Manually trigger reload (no restart needed) +curl -X POST "https://your-proxy-url/reload/anthropic_beta_headers" \ + -H "Authorization: Bearer YOUR_ADMIN_TOKEN" + +# Or schedule automatic reloads every 24 hours +curl -X POST "https://your-proxy-url/schedule/anthropic_beta_headers_reload?hours=24" \ + -H "Authorization: Bearer YOUR_ADMIN_TOKEN" ``` -**Example: Use Local Config Only (No Remote Fetching)** -```bash -export LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS=True +This prevents future incidents where Claude Code introduces new headers before LiteLLM configuration is updated. + +--- + +## Configuration format + +The `anthropic_beta_headers_config.json` file maps input headers to provider-specific output headers: + +```json +{ + "description": "Mapping of Anthropic beta headers for each provider.", + "anthropic": { + "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20", + "computer-use-2025-01-24": "computer-use-2025-01-24" + }, + "bedrock_converse": { + "advanced-tool-use-2025-11-20": null, + "computer-use-2025-01-24": "computer-use-2025-01-24" + }, + "azure_ai": { + "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20", + "computer-use-2025-01-24": "computer-use-2025-01-24" + } +} ``` + +**Validation rules:** +1. Headers must exist in the mapping for the target provider +2. Headers with `null` values are filtered out (unsupported) +3. Header names can be transformed per-provider (e.g., Bedrock uses different names for some features) + +--- + +## Resolution steps for users + +For users still experiencing issues, update to the latest LiteLLM version if < v1.81.11-nightly: + +```bash +pip install --upgrade litellm +``` + +Or manually reload the configuration without restarting: + +```bash +curl -X POST "https://your-proxy-url/reload/anthropic_beta_headers" \ + -H "Authorization: Bearer YOUR_ADMIN_TOKEN" +``` + +--- + +## Related documentation + +- [Managing Anthropic Beta Headers](../proxy/sync_anthropic_beta_headers.md) - Complete configuration guide +- [`anthropic_beta_headers_config.json`](https://github.com/BerriAI/litellm/blob/main/litellm/anthropic_beta_headers_config.json) - Current configuration file diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 5e3f56c4206..775cdf6876a 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -450,6 +450,7 @@ router_settings: | BATCH_STATUS_POLL_INTERVAL_SECONDS | Interval in seconds for polling batch status. Default is 3600 (1 hour) | BATCH_STATUS_POLL_MAX_ATTEMPTS | Maximum number of attempts for polling batch status. Default is 24 (for 24 hours) | BEDROCK_MAX_POLICY_SIZE | Maximum size for Bedrock policy. Default is 75 +| BEDROCK_MIN_THINKING_BUDGET_TOKENS | Minimum thinking budget in tokens for Bedrock reasoning models. Bedrock returns a 400 error if budget_tokens is below this value. Requests with lower values are clamped to this minimum. Default is 1024 | BERRISPEND_ACCOUNT_ID | Account ID for BerriSpend service | BRAINTRUST_API_KEY | API key for Braintrust integration | BRAINTRUST_API_BASE | Base URL for Braintrust API. Default is https://api.braintrustdata.com/v1 diff --git a/litellm/constants.py b/litellm/constants.py index 7c21111d313..458f48cb0b6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -577,6 +577,10 @@ OPENAI_CHAT_COMPLETION_PARAMS = [ "web_search_options", "service_tier", "store", + "prompt_cache_key", + "prompt_cache_retention", + "safety_identifier", + "verbosity", ] OPENAI_TRANSCRIPTION_PARAMS = [ diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 9d1c254d1a7..55746e5e527 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -329,10 +329,10 @@ class ContentFilterGuardrail(CustomGuardrail): action if action else category_config_obj.default_action ) - # Handle conditional categories (with identifier_words + inherit_from OR identifier_words + additional_block_words) - if category_config_obj.identifier_words and ( - category_config_obj.inherit_from - or category_config_obj.additional_block_words + # Handle conditional categories (with identifier_words + inherit_from) + if ( + category_config_obj.identifier_words + and category_config_obj.inherit_from ): self._load_conditional_category( category_name, @@ -387,81 +387,64 @@ class ContentFilterGuardrail(CustomGuardrail): categories_dir: str, ) -> None: """ - Load a conditional category that uses identifier_words + block_words. - - Supports two patterns: - 1. Inherit + additional: identifier_words + inherit_from + optional additional_block_words - 2. Standalone: identifier_words + additional_block_words (no inheritance) + Load a conditional category that uses identifier_words + inherited block_words. Args: category_name: Name of the category - category_config_obj: CategoryConfig object with identifier_words and either inherit_from or additional_block_words + category_config_obj: CategoryConfig object with identifier_words and inherit_from category_action: Action to take when match is found severity_threshold: Minimum severity threshold categories_dir: Directory containing category files """ - block_words = [] + # Load the inherited category to get block words inherit_from = category_config_obj.inherit_from + if not inherit_from: + return - # Pattern 1: Load inherited category to get base block words - if inherit_from: - # Remove .json or .yaml extension if included - inherit_base = inherit_from.replace(".json", "").replace(".yaml", "") + # Remove .json or .yaml extension if included + inherit_base = inherit_from.replace(".json", "").replace(".yaml", "") - # Find the inherited category file - inherit_yaml_path = os.path.join(categories_dir, f"{inherit_base}.yaml") - inherit_json_path = os.path.join(categories_dir, f"{inherit_base}.json") + # Find the inherited category file + inherit_yaml_path = os.path.join(categories_dir, f"{inherit_base}.yaml") + inherit_json_path = os.path.join(categories_dir, f"{inherit_base}.json") - if os.path.exists(inherit_yaml_path): - inherit_file_path = inherit_yaml_path - elif os.path.exists(inherit_json_path): - inherit_file_path = inherit_json_path - else: - verbose_proxy_logger.warning( - f"Category {category_name}: inherit_from '{inherit_from}' file not found at {categories_dir}" - ) - verbose_proxy_logger.debug( - f"Tried paths: {inherit_yaml_path}, {inherit_json_path}" - ) - return - - try: - # Load the inherited category - inherited_category = self._load_category_file(inherit_file_path) - - # Extract block words from inherited category that meet severity threshold - for keyword_data in inherited_category.keywords: - keyword = keyword_data["keyword"].lower() - severity = keyword_data["severity"] - if self._should_apply_severity(severity, severity_threshold): - block_words.append(keyword) - except Exception as e: - verbose_proxy_logger.error( - f"Error loading inherited category for {category_name}: {e}" - ) - return - - # Pattern 2 or supplement to Pattern 1: Add additional block words - if category_config_obj.additional_block_words: - block_words.extend(category_config_obj.additional_block_words) - - # Ensure we have block words before storing - if not block_words: + if os.path.exists(inherit_yaml_path): + inherit_file_path = inherit_yaml_path + elif os.path.exists(inherit_json_path): + inherit_file_path = inherit_json_path + else: verbose_proxy_logger.warning( - f"Category {category_name}: no block words found (check inherit_from or additional_block_words)" + f"Category {category_name}: inherit_from '{inherit_from}' file not found at {categories_dir}" + ) + verbose_proxy_logger.debug( + f"Tried paths: {inherit_yaml_path}, {inherit_json_path}" ) return - # Store the conditional category configuration - self.conditional_categories[category_name] = { - "identifier_words": category_config_obj.identifier_words, - "block_words": block_words, - "action": category_action, - "severity": "high", # Combinations are always high severity - } + try: + # Load the inherited category + inherited_category = self._load_category_file(inherit_file_path) + + # Extract block words from inherited category that meet severity threshold + block_words = [] + for keyword_data in inherited_category.keywords: + keyword = keyword_data["keyword"].lower() + severity = keyword_data["severity"] + if self._should_apply_severity(severity, severity_threshold): + block_words.append(keyword) + + # Add additional block words specific to this category + if category_config_obj.additional_block_words: + block_words.extend(category_config_obj.additional_block_words) + + # Store the conditional category configuration + self.conditional_categories[category_name] = { + "identifier_words": category_config_obj.identifier_words, + "block_words": block_words, + "action": category_action, + "severity": "high", # Combinations are always high severity + } - # Log different messages based on pattern - if inherit_from and category_config_obj.additional_block_words: verbose_proxy_logger.info( f"Loaded conditional category {category_name}: " f"{len(category_config_obj.identifier_words)} identifiers + " @@ -469,17 +452,9 @@ class ContentFilterGuardrail(CustomGuardrail): f"({len(category_config_obj.additional_block_words)} additional + " f"{len(block_words) - len(category_config_obj.additional_block_words)} from {inherit_from})" ) - elif inherit_from: - verbose_proxy_logger.info( - f"Loaded conditional category {category_name}: " - f"{len(category_config_obj.identifier_words)} identifiers + " - f"{len(block_words)} block words (from {inherit_from})" - ) - else: - verbose_proxy_logger.info( - f"Loaded conditional category {category_name}: " - f"{len(category_config_obj.identifier_words)} identifiers + " - f"{len(block_words)} block words (standalone)" + except Exception as e: + verbose_proxy_logger.error( + f"Error loading inherited category for {category_name}: {e}" ) def _load_category_file(self, file_path: str) -> CategoryConfig: @@ -1398,6 +1373,41 @@ class ContentFilterGuardrail(CustomGuardrail): names = [cat.description or cat.category_name for cat in self.loaded_categories.values()] return ", ".join(names) if names else None + def _compute_risk_score( + self, + detections: List[ContentFilterDetection], + masked_entity_count: Dict[str, int], + status: "GuardrailStatus", + ) -> float: + """ + Compute a risk score from 0-10 for this guardrail evaluation. + + Factors: + - Match ratio: how many patterns matched vs total checked + - Number of entities masked + - Whether the guardrail blocked the request (max risk) + """ + if status == "guardrail_intervened": + return 10.0 + + total_masked = sum(masked_entity_count.values()) if masked_entity_count else 0 + patterns_checked = self._get_patterns_checked_count() + + # Match ratio contribution (0-7 points) + match_ratio = total_masked / patterns_checked if patterns_checked > 0 else 0.0 + ratio_score = match_ratio * 7.0 + + # Detection count contribution (0-3 points, capped) + detection_score = min(len(detections), 5) * 0.6 + + score = ratio_score + detection_score + + # Floor: if anything matched, minimum risk is 2 + if total_masked > 0 and score < 2.0: + score = 2.0 + + return round(min(10.0, score), 1) + def _log_guardrail_information( self, request_data: dict, @@ -1444,6 +1454,7 @@ class ContentFilterGuardrail(CustomGuardrail): detection_method=self._get_detection_methods(detections) if detections else None, match_details=self._build_match_details(detections) if detections else None, patterns_checked=self._get_patterns_checked_count(), + risk_score=self._compute_risk_score(detections, masked_entity_count, status), ), ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index c0285407855..57b6453ac43 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1911,6 +1911,10 @@ async def get_user_daily_activity( default=None, description="Filter by specific API key", ), + user_id: Optional[str] = fastapi.Query( + default=None, + description="Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id.", + ), page: int = fastapi.Query( default=1, description="Page number for pagination", ge=1 ), @@ -1955,9 +1959,21 @@ async def get_user_daily_activity( ) try: - entity_id: Optional[str] = None - if not _user_has_admin_view(user_api_key_dict): - entity_id = user_api_key_dict.user_id + is_admin = _user_has_admin_view(user_api_key_dict) + + if is_admin: + entity_id = user_id # None means global view, otherwise filter by user + else: + if user_id is None: + user_id = user_api_key_dict.user_id + if user_id != user_api_key_dict.user_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "Non-admin users can only view their own spend data." + }, + ) + entity_id = user_id return await get_daily_activity( prisma_client=prisma_client, @@ -1974,6 +1990,8 @@ async def get_user_daily_activity( timezone_offset_minutes=timezone, ) + except HTTPException: + raise except Exception as e: verbose_proxy_logger.exception( "/spend/daily/analytics: Exception occured - {}".format(str(e)) @@ -2008,6 +2026,10 @@ async def get_user_daily_activity_aggregated( default=None, description="Filter by specific API key", ), + user_id: Optional[str] = fastapi.Query( + default=None, + description="Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id.", + ), timezone: Optional[int] = fastapi.Query( default=None, description="Timezone offset in minutes from UTC (e.g., 480 for PST). " @@ -2034,9 +2056,21 @@ async def get_user_daily_activity_aggregated( ) try: - entity_id: Optional[str] = None - if not _user_has_admin_view(user_api_key_dict): - entity_id = user_api_key_dict.user_id + is_admin = _user_has_admin_view(user_api_key_dict) + + if is_admin: + entity_id = user_id # None means global view, otherwise filter by user + else: + if user_id is None: + user_id = user_api_key_dict.user_id + if user_id != user_api_key_dict.user_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "Non-admin users can only view their own spend data." + }, + ) + entity_id = user_id return await get_daily_activity_aggregated( prisma_client=prisma_client, @@ -2051,6 +2085,8 @@ async def get_user_daily_activity_aggregated( timezone_offset_minutes=timezone, ) + except HTTPException: + raise except Exception as e: verbose_proxy_logger.exception( "/user/daily/activity/aggregated: Exception occured - {}".format(str(e)) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 9dcc25e7a87..21459a1b802 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3277,6 +3277,14 @@ async def _execute_virtual_key_regeneration( update_data.update(non_default_values) update_data = prisma_client.jsonify_object(data=update_data) + # If grace period set, insert deprecated key so old key remains valid + await _insert_deprecated_key( + prisma_client=prisma_client, + old_token_hash=hashed_api_key, + new_token_hash=new_token_hash, + grace_period=data.grace_period if data else None, + ) + updated_token = await prisma_client.db.litellm_verificationtoken.update( where={"token": hashed_api_key}, data=update_data, # type: ignore @@ -3474,58 +3482,6 @@ async def regenerate_key_fn( # noqa: PLR0915 ) verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) - new_token = get_new_token(data=data) - - new_token_hash = hash_token(new_token) - new_token_key_name = f"sk-...{new_token[-4:]}" - - # Prepare the update data - update_data = { - "token": new_token_hash, - "key_name": new_token_key_name, - } - - non_default_values = {} - if data is not None: - # Update with any provided parameters from GenerateKeyRequest - non_default_values = await prepare_key_update_data( - data=data, existing_key_row=_key_in_db - ) - verbose_proxy_logger.debug("non_default_values: %s", non_default_values) - - update_data.update(non_default_values) - update_data = prisma_client.jsonify_object(data=update_data) - - # If grace period set, insert deprecated key so old key remains valid - await _insert_deprecated_key( - prisma_client=prisma_client, - old_token_hash=hashed_api_key, - new_token_hash=new_token_hash, - grace_period=data.grace_period if data else None, - ) - - # Update the token in the database - updated_token = await prisma_client.db.litellm_verificationtoken.update( - where={"token": hashed_api_key}, - data=update_data, # type: ignore - ) - - updated_token_dict = {} - if updated_token is not None: - updated_token_dict = dict(updated_token) - - updated_token_dict["key"] = new_token - updated_token_dict["token_id"] = updated_token_dict.pop("token") - - ### 3. remove existing key entry from cache - ###################################################################### - - if hashed_api_key or key: - await _delete_cache_key_object( - hashed_token=hash_token(key), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) # Normalize litellm_changed_by: if it's a Header object or not a string, convert to None if litellm_changed_by is not None and not isinstance(litellm_changed_by, str): litellm_changed_by = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5fbfd23b2db..5f8798c7712 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2644,6 +2644,9 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False): alert_recipients: Optional[List[str]] """Email addresses that were notified""" + risk_score: Optional[float] + """Risk score 0-10 indicating how risky the request was (higher = riskier). Computed by the guardrail provider.""" + class GuardrailTracingDetail(TypedDict, total=False): """ @@ -2661,6 +2664,7 @@ class GuardrailTracingDetail(TypedDict, total=False): match_details: Optional[List[dict]] patterns_checked: Optional[int] alert_recipients: Optional[List[str]] + risk_score: Optional[float] StandardLoggingPayloadStatus = Literal["success", "failure"] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9ea9f39b1db..41acb5c8101 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -17112,6 +17112,19 @@ "supports_parallel_function_calling": true, "supports_vision": true }, + "github_copilot/claude-opus-4.6-fast": { + "litellm_provider": "github_copilot", + "max_input_tokens": 128000, + "max_output_tokens": 16000, + "max_tokens": 16000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true + }, "github_copilot/claude-opus-41": { "litellm_provider": "github_copilot", "max_input_tokens": 80000, @@ -17363,6 +17376,20 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/gpt-5.3-codex": { + "litellm_provider": "github_copilot", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, "github_copilot/text-embedding-3-small": { "litellm_provider": "github_copilot", "max_input_tokens": 8191, diff --git a/pyproject.toml b/pyproject.toml index 68b38fb5ff8..4deb61836b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.81.12" +version = "1.81.13" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -182,7 +182,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.81.12" +version = "1.81.13" version_files = [ "pyproject.toml:^version" ] diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py index 0850f742231..f38ce67cede 100644 --- a/tests/llm_responses_api_testing/base_responses_api.py +++ b/tests/llm_responses_api_testing/base_responses_api.py @@ -74,7 +74,7 @@ def validate_responses_api_response(response, final_chunk: bool = False): "top_p": (int, float, type(None)), "max_output_tokens": (int, type(None)), "previous_response_id": (str, type(None)), - "reasoning": dict, + "reasoning": (dict, type(None)), "status": str, "text": dict, "truncation": (str, type(None)), diff --git a/tests/test_litellm/containers/test_container_integration.py b/tests/test_litellm/containers/test_container_integration.py index b2f52fcea97..177996abd99 100644 --- a/tests/test_litellm/containers/test_container_integration.py +++ b/tests/test_litellm/containers/test_container_integration.py @@ -385,6 +385,15 @@ class TestContainerIntegration: @pytest.mark.parametrize("provider", ["openai"]) def test_provider_support(self, provider): """Test that the container API works with supported providers.""" + import importlib + import litellm.containers.main as containers_main_module + + # Reload the module to ensure it has a fresh reference to base_llm_http_handler + # after conftest reloads litellm (same pattern as test_error_handling_integration) + importlib.reload(containers_main_module) + + from litellm.containers.main import create_container as create_container_fresh + mock_response = ContainerObject( id="cntr_provider_test", object="container", @@ -398,7 +407,7 @@ class TestContainerIntegration: with patch('litellm.containers.main.base_llm_http_handler') as mock_handler: mock_handler.container_create_handler.return_value = mock_response - response = create_container( + response = create_container_fresh( name="Provider Test Container", custom_llm_provider=provider ) diff --git a/tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py b/tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py index fa605154bb0..7b974aba35c 100644 --- a/tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py +++ b/tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py @@ -1,6 +1,5 @@ import os import sys -from unittest.mock import AsyncMock, patch import pytest @@ -47,67 +46,26 @@ def test_map_openai_params(): assert "response_format" in result -@pytest.mark.asyncio -async def test_llama_api_streaming_no_307_error(): - """Test that streaming works without 307 redirect errors due to follow_redirects=True""" +def test_llama_api_streaming_no_307_error(): + """ + Test that the OpenAI-compatible httpx clients use follow_redirects=True. - # Mock the httpx client to simulate a successful streaming response - with patch( - "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" - ) as mock_get_client: - # Create a mock client - mock_client = AsyncMock() - mock_get_client.return_value = mock_client + meta_llama routes through the OpenAI SDK path (BaseOpenAILLM), so the + follow_redirects setting on that SDK's underlying httpx client is what + actually prevents 307 redirect errors for LLaMA API streaming. + """ + from litellm.llms.openai.common_utils import BaseOpenAILLM - # Mock a successful streaming response (not a 307 redirect) - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "text/plain; charset=utf-8"} + # Verify the async httpx client has follow_redirects enabled + async_client = BaseOpenAILLM._get_async_http_client() + assert async_client is not None + assert ( + async_client.follow_redirects is True + ), "Async httpx client should set follow_redirects=True to prevent 307 errors" - # Mock streaming data that would come from a successful request - async def mock_aiter_lines(): - yield 'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"meta_llama/Llama-4-Maverick-17B-128E-Instruct-FP8","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}' - yield 'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"meta_llama/Llama-4-Maverick-17B-128E-Instruct-FP8","choices":[{"index":0,"delta":{"content":" there"},"finish_reason":null}]}' - yield 'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"meta_llama/Llama-4-Maverick-17B-128E-Instruct-FP8","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' - yield "data: [DONE]" - - mock_response.aiter_lines.return_value = mock_aiter_lines() - mock_client.stream.return_value.__aenter__.return_value = mock_response - - # Test the streaming completion - try: - response = await litellm.acompletion( - model="meta_llama/Llama-4-Maverick-17B-128E-Instruct-FP8", - messages=[{"role": "user", "content": "Tell me about yourself"}], - stream=True, - temperature=0.0, - ) - - # Verify we get a CustomStreamWrapper (streaming response) - from litellm.utils import CustomStreamWrapper - - assert isinstance(response, CustomStreamWrapper) - - # Verify the HTTP client was called with follow_redirects=True - mock_client.stream.assert_called_once() - call_kwargs = mock_client.stream.call_args[1] - assert ( - call_kwargs.get("follow_redirects") is True - ), "follow_redirects should be True to prevent 307 errors" - - # Verify the response status is 200 (not 307) - assert ( - mock_response.status_code == 200 - ), "Should get 200 response, not 307 redirect" - - except Exception as e: - # If there's an exception, make sure it's not a 307 error - error_str = str(e) - assert ( - "307" not in error_str - ), f"Should not get 307 redirect error: {error_str}" - - # Still verify that follow_redirects was set correctly - if mock_client.stream.called: - call_kwargs = mock_client.stream.call_args[1] - assert call_kwargs.get("follow_redirects") is True + # Verify the sync httpx client has follow_redirects enabled + sync_client = BaseOpenAILLM._get_sync_http_client() + assert sync_client is not None + assert ( + sync_client.follow_redirects is True + ), "Sync httpx client should set follow_redirects=True to prevent 307 errors" diff --git a/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py b/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py index f6e5e05fe51..cd530cd3b40 100644 --- a/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py +++ b/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py @@ -7,6 +7,7 @@ PublicAI is an OpenAI-compatible provider with minor customizations. import os import sys +from unittest.mock import patch sys.path.insert( 0, os.path.abspath("../../../../..") @@ -51,9 +52,13 @@ class TestPublicAIConfig: assert result["Authorization"] == f"Bearer {api_key}" assert result["Content-Type"] == "application/json" - def test_get_supported_openai_params(self, config): + @patch("litellm.utils.supports_function_calling", return_value=True) + def test_get_supported_openai_params(self, mock_supports_fc, config): """ - Test that get_supported_openai_params returns correct params + Test that get_supported_openai_params returns correct params. + We mock supports_function_calling because the test model name + 'swiss-ai-apertus' is not in the model registry; this test validates + config behaviour, not registry lookups. """ supported_params = config.get_supported_openai_params(model="swiss-ai-apertus") @@ -66,9 +71,12 @@ class TestPublicAIConfig: # Note: JSON-based configs inherit from OpenAIGPTConfig which includes functions # This is expected behavior for JSON-based providers - def test_map_openai_params_includes_functions(self, config): + @patch("litellm.utils.supports_function_calling", return_value=True) + def test_map_openai_params_includes_functions(self, mock_supports_fc, config): """ - Test that functions parameter is mapped (JSON-based configs don't exclude functions) + Test that functions parameter is mapped (JSON-based configs don't exclude functions). + We mock supports_function_calling because the test model name + 'swiss-ai-apertus' is not in the model registry. """ non_default_params = { "functions": [{"name": "test_function", "description": "Test function"}], diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index fbf5239797f..5e29f927b67 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -22,6 +22,8 @@ class TestVertexAIRerankTransform: "GOOGLE_APPLICATION_CREDENTIALS", "GOOGLE_CLOUD_PROJECT", "VERTEXAI_PROJECT", + "VERTEXAI_CREDENTIALS", + "VERTEX_AI_CREDENTIALS", "VERTEX_PROJECT", "VERTEX_LOCATION", "VERTEX_AI_PROJECT", @@ -471,16 +473,20 @@ class TestVertexAIRerankTransform: } assert headers == expected_headers - @patch('litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig._ensure_access_token') def test_validate_environment_preserves_optional_params_for_get_complete_url( self, - mock_ensure_access_token, ): """ Validate that calling validate_environment does not remove vertex-specific parameters needed later by get_complete_url. + + Uses instance-level mocking to avoid class-reference issues caused by + importlib.reload(litellm) in conftest.py. """ - mock_ensure_access_token.return_value = ("test-access-token", "project-from-token") + mock_ensure_access_token = MagicMock( + return_value=("test-access-token", "project-from-token") + ) + self.config._ensure_access_token = mock_ensure_access_token optional_params = { "vertex_credentials": "path/to/credentials.json", diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 919af96f760..9a417f3566c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1167,4 +1167,136 @@ def test_generate_request_base_validator(): # Test with None req = GenerateRequestBase(max_budget=None) - assert req.max_budget is None \ No newline at end of file + assert req.max_budget is None + + +@pytest.mark.asyncio +async def test_get_user_daily_activity_non_admin_cannot_view_other_users(monkeypatch): + """ + Test that non-admin users cannot view another user's daily activity data. + The endpoint should raise 403 when user_id does not match the caller's own user_id. + Also verifies that omitting user_id defaults to the caller's own user_id. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_daily_activity, + ) + + # Mock the prisma client so the DB-not-connected check passes + mock_prisma_client = MagicMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) + + # Non-admin caller + non_admin_key_dict = UserAPIKeyAuth( + user_id="regular-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + # Case 1: Non-admin tries to view a different user's data — should get 403 + with pytest.raises(HTTPException) as exc_info: + await get_user_daily_activity( + start_date="2025-01-01", + end_date="2025-01-31", + model=None, + api_key=None, + user_id="other-user-456", + page=1, + page_size=50, + timezone=None, + user_api_key_dict=non_admin_key_dict, + ) + + assert exc_info.value.status_code == 403 + assert "Non-admin users can only view their own spend data" in str( + exc_info.value.detail + ) + + # Case 2: Non-admin omits user_id — should default to their own user_id + mock_response = MagicMock() + with patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_get_daily: + result = await get_user_daily_activity( + start_date="2025-01-01", + end_date="2025-01-31", + model=None, + api_key=None, + user_id=None, + page=1, + page_size=50, + timezone=None, + user_api_key_dict=non_admin_key_dict, + ) + + # Verify it called get_daily_activity with the caller's own user_id + mock_get_daily.assert_called_once() + call_kwargs = mock_get_daily.call_args + assert call_kwargs.kwargs["entity_id"] == "regular-user-123" + + +@pytest.mark.asyncio +async def test_get_user_daily_activity_aggregated_admin_global_view(monkeypatch): + """ + Test that admin users can call the aggregated endpoint without a user_id + to get a global view. Also verifies that the correct arguments are forwarded + to the underlying get_daily_activity_aggregated helper. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_daily_activity_aggregated, + ) + + # Mock the prisma client + mock_prisma_client = MagicMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) + + # Mock the downstream helper so we don't need a real DB + mock_response = MagicMock() + mock_get_daily_agg = AsyncMock(return_value=mock_response) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + mock_get_daily_agg, + ) + + # Admin caller + admin_key_dict = UserAPIKeyAuth( + user_id="admin-user-001", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + # Admin calls without user_id → global view (entity_id=None) + result = await get_user_daily_activity_aggregated( + start_date="2025-02-01", + end_date="2025-02-28", + model="gpt-4", + api_key=None, + user_id=None, + timezone=480, + user_api_key_dict=admin_key_dict, + ) + + assert result is mock_response + + # Verify the helper was called with the right parameters + mock_get_daily_agg.assert_called_once_with( + prisma_client=mock_prisma_client, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, # global view: no user_id filter + entity_metadata_field=None, + start_date="2025-02-01", + end_date="2025-02-28", + model="gpt-4", + api_key=None, + timezone_offset_minutes=480, + ) \ No newline at end of file diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index be91800732b..a18c2dba032 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -446,8 +446,24 @@ class TestProxyInitializationHelpers: mock_proxy_config_instance.get_config = mock_get_config mock_proxy_config.return_value = mock_proxy_config_instance - # Ensure DATABASE_URL is not set in the environment - with patch.dict(os.environ, {"DATABASE_URL": ""}, clear=True): + mock_proxy_server_module = MagicMock(app=mock_app) + + # Only remove DATABASE_URL and DIRECT_URL to prevent the database setup + # code path from running. Do NOT use clear=True as it removes PATH, HOME, + # etc., which causes imports inside run_server to break in CI (the real + # litellm.proxy.proxy_server import at line 820 of proxy_cli.py has heavy + # side effects that fail without a proper environment). + env_overrides = { + "DATABASE_URL": "", + "DIRECT_URL": "", + "IAM_TOKEN_DB_AUTH": "", + "USE_AWS_KMS": "", + } + with patch.dict(os.environ, env_overrides): + # Remove DATABASE_URL entirely so the DB setup block is skipped + os.environ.pop("DATABASE_URL", None) + os.environ.pop("DIRECT_URL", None) + with patch.dict( "sys.modules", { @@ -456,7 +472,11 @@ class TestProxyInitializationHelpers: ProxyConfig=mock_proxy_config, KeyManagementSettings=mock_key_mgmt, save_worker_config=mock_save_worker_config, - ) + ), + # Also mock litellm.proxy.proxy_server to prevent the real + # import at line 820 of proxy_cli.py which has heavy side + # effects (FastAPI app init, logging setup, etc.) + "litellm.proxy.proxy_server": mock_proxy_server_module, }, ), patch( "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" @@ -470,7 +490,10 @@ class TestProxyInitializationHelpers: # Test with no config parameter (config=None) result = runner.invoke(run_server, ["--local"]) - assert result.exit_code == 0 + assert result.exit_code == 0, ( + f"run_server failed with exit_code={result.exit_code}, " + f"output={result.output}, exception={result.exception}" + ) # Verify that uvicorn.run was called mock_uvicorn_run.assert_called_once() @@ -481,7 +504,10 @@ class TestProxyInitializationHelpers: # Test with explicit --config None (should behave the same) result = runner.invoke(run_server, ["--local", "--config", "None"]) - assert result.exit_code == 0 + assert result.exit_code == 0, ( + f"run_server failed with exit_code={result.exit_code}, " + f"output={result.output}, exception={result.exception}" + ) # Verify that uvicorn.run was called again mock_uvicorn_run.assert_called_once() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts new file mode 100644 index 00000000000..b0a96eff0e7 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts @@ -0,0 +1,339 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useInfiniteUsers } from "./useUsers"; +import { userListCall } from "@/components/networking"; +import type { UserListResponse } from "@/components/networking"; + +vi.mock("@/components/networking", () => ({ + userListCall: vi.fn(), +})); + +vi.mock("../common/queryKeysFactory", () => ({ + createQueryKeys: vi.fn((resource: string) => ({ + all: [resource], + lists: () => [resource, "list"], + list: (params?: any) => [resource, "list", { params }], + details: () => [resource, "detail"], + detail: (uid: string) => [resource, "detail", uid], + })), +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +const DEFAULT_AUTH = { + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, +}; + +const buildUserListResponse = ( + page: number, + totalPages: number, + userCount = 2, +): UserListResponse => ({ + page, + page_size: 50, + total: totalPages * userCount, + total_pages: totalPages, + users: Array.from({ length: userCount }, (_, i) => ({ + user_id: `user-${page}-${i}`, + user_email: `user-${page}-${i}@example.com`, + user_alias: null, + user_role: "Internal User", + spend: 0, + max_budget: null, + key_count: 0, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + sso_user_id: null, + budget_duration: null, + })), +}); + +describe("useInfiniteUsers", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + vi.clearAllMocks(); + mockUseAuthorized.mockReturnValue(DEFAULT_AUTH); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return paginated user data when query is successful", async () => { + const mockResponse = buildUserListResponse(1, 2); + (userListCall as any).mockResolvedValue(mockResponse); + + const { result } = renderHook(() => useInfiniteUsers(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data?.pages).toHaveLength(1); + expect(result.current.data?.pages[0]).toEqual(mockResponse); + expect(userListCall).toHaveBeenCalledWith( + "test-access-token", + null, + 1, + 50, + null, + ); + }); + + it("should use the default page size of 50", async () => { + const mockResponse = buildUserListResponse(1, 1); + (userListCall as any).mockResolvedValue(mockResponse); + + const { result } = renderHook(() => useInfiniteUsers(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(userListCall).toHaveBeenCalledWith( + "test-access-token", + null, + 1, + 50, + null, + ); + }); + + it("should use a custom page size when provided", async () => { + const customPageSize = 25; + const mockResponse = buildUserListResponse(1, 1, 5); + (userListCall as any).mockResolvedValue(mockResponse); + + const { result } = renderHook(() => useInfiniteUsers(customPageSize), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(userListCall).toHaveBeenCalledWith( + "test-access-token", + null, + 1, + customPageSize, + null, + ); + }); + + it("should pass searchEmail to userListCall when provided", async () => { + const searchEmail = "search@example.com"; + const mockResponse = buildUserListResponse(1, 1, 1); + (userListCall as any).mockResolvedValue(mockResponse); + + const { result } = renderHook(() => useInfiniteUsers(50, searchEmail), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(userListCall).toHaveBeenCalledWith( + "test-access-token", + null, + 1, + 50, + searchEmail, + ); + }); + + it("should pass null for searchEmail when not provided", async () => { + const mockResponse = buildUserListResponse(1, 1); + (userListCall as any).mockResolvedValue(mockResponse); + + const { result } = renderHook(() => useInfiniteUsers(50, undefined), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(userListCall).toHaveBeenCalledWith( + "test-access-token", + null, + 1, + 50, + null, + ); + }); + + it("should fetch the next page when more pages are available", async () => { + const page1 = buildUserListResponse(1, 3); + const page2 = buildUserListResponse(2, 3); + let callCount = 0; + (userListCall as any).mockImplementation(async () => { + callCount++; + return callCount === 1 ? page1 : page2; + }); + + const { result } = renderHook(() => useInfiniteUsers(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.hasNextPage).toBe(true); + + result.current.fetchNextPage(); + + await waitFor(() => { + expect(result.current.isFetchingNextPage).toBe(false); + expect(result.current.data?.pages).toHaveLength(2); + }); + + expect(result.current.data?.pages[1]).toEqual(page2); + expect(userListCall).toHaveBeenCalledTimes(2); + expect(userListCall).toHaveBeenLastCalledWith( + "test-access-token", + null, + 2, + 50, + null, + ); + }); + + it("should not have a next page when on the last page", async () => { + const lastPage = buildUserListResponse(2, 2); + (userListCall as any).mockResolvedValue(lastPage); + + const { result } = renderHook(() => useInfiniteUsers(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.hasNextPage).toBe(false); + }); + + it("should not execute query when accessToken is missing", async () => { + mockUseAuthorized.mockReturnValue({ + ...DEFAULT_AUTH, + accessToken: null, + }); + + const { result } = renderHook(() => useInfiniteUsers(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(userListCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is not an admin role", async () => { + mockUseAuthorized.mockReturnValue({ + ...DEFAULT_AUTH, + userRole: "Internal User", + }); + + const { result } = renderHook(() => useInfiniteUsers(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(userListCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when both accessToken and userRole are invalid", async () => { + mockUseAuthorized.mockReturnValue({ + ...DEFAULT_AUTH, + accessToken: null, + userRole: "App User", + }); + + const { result } = renderHook(() => useInfiniteUsers(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(userListCall).not.toHaveBeenCalled(); + }); + + it("should execute query for each admin role", async () => { + const adminRoles = [ + "Admin", + "Admin Viewer", + "proxy_admin", + "proxy_admin_viewer", + "org_admin", + ]; + + for (const role of adminRoles) { + vi.clearAllMocks(); + queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + const mockResponse = buildUserListResponse(1, 1); + (userListCall as any).mockResolvedValue(mockResponse); + mockUseAuthorized.mockReturnValue({ ...DEFAULT_AUTH, userRole: role }); + + const { result } = renderHook(() => useInfiniteUsers(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(userListCall).toHaveBeenCalledTimes(1); + } + }); + + it("should handle error when userListCall fails", async () => { + const testError = new Error("Failed to fetch users"); + (userListCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useInfiniteUsers(), { wrapper }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + }); + + it("should pass empty string searchEmail as null", async () => { + const mockResponse = buildUserListResponse(1, 1); + (userListCall as any).mockResolvedValue(mockResponse); + + const { result } = renderHook(() => useInfiniteUsers(50, ""), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(userListCall).toHaveBeenCalledWith( + "test-access-token", + null, + 1, + 50, + null, + ); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts new file mode 100644 index 00000000000..cb30299f46f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts @@ -0,0 +1,41 @@ +import { userListCall, UserListResponse } from "@/components/networking"; +import { useInfiniteQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { all_admin_roles } from "@/utils/roles"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +const infiniteUsersKeys = createQueryKeys("infiniteUsers"); + +const DEFAULT_PAGE_SIZE = 50; + +export const useInfiniteUsers = ( + pageSize: number = DEFAULT_PAGE_SIZE, + searchEmail?: string, +) => { + const { accessToken, userRole } = useAuthorized(); + return useInfiniteQuery({ + queryKey: infiniteUsersKeys.list({ + filters: { + pageSize, + ...(searchEmail && { searchEmail }), + }, + }), + queryFn: async ({ pageParam }) => { + return await userListCall( + accessToken!, + null, // userIDs + pageParam as number, // page + pageSize, // page_size + searchEmail || null, // userEmail + ); + }, + initialPageParam: 1, + getNextPageParam: (lastPage) => { + if (lastPage.page < lastPage.total_pages) { + return lastPage.page + 1; + } + return undefined; + }, + enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!), + }); +}; diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx index 1a344d3dd95..5f5ffe83baa 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx @@ -2,6 +2,7 @@ import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; +import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import { act, fireEvent, screen, waitFor } from "@testing-library/react"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../../tests/test-utils"; @@ -116,6 +117,10 @@ vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({ useCurrentUser: vi.fn(), })); +vi.mock("@/app/(dashboard)/hooks/users/useUsers", () => ({ + useInfiniteUsers: vi.fn(), +})); + vi.mock("antd", async (importOriginal) => { const React = await import("react"); const actual = await importOriginal(); @@ -223,6 +228,10 @@ vi.mock("@ant-design/icons", async () => { return React.createElement("span"); } + function LoadingOutlined(props: any) { + return React.createElement("span", { "data-testid": "loading-icon", ...props }); + } + return { GlobalOutlined: Icon, BankOutlined: Icon, @@ -235,6 +244,8 @@ vi.mock("@ant-design/icons", async () => { ClockCircleOutlined: Icon, CalendarOutlined: Icon, InfoCircleOutlined: Icon, + UserOutlined: Icon, + LoadingOutlined, }; }); @@ -320,11 +331,13 @@ vi.mock("@tremor/react", async () => { describe("UsagePage", () => { const mockUserDailyActivityAggregatedCall = vi.mocked(networking.userDailyActivityAggregatedCall); + const mockUserDailyActivityCall = vi.mocked(networking.userDailyActivityCall); const mockTagListCall = vi.mocked(networking.tagListCall); const mockUseCustomers = vi.mocked(useCustomers); const mockUseAgents = vi.mocked(useAgents); const mockUseAuthorized = vi.mocked(useAuthorized); const mockUseCurrentUser = vi.mocked(useCurrentUser); + const mockUseInfiniteUsers = vi.mocked(useInfiniteUsers); const mockSpendData = { results: [ @@ -487,6 +500,8 @@ describe("UsagePage", () => { beforeEach(() => { mockUseAuthorized.mockReturnValue({ + isLoading: false, + isAuthorized: true, token: "mock-token", accessToken: "test-token", userId: "user-123", @@ -505,8 +520,30 @@ describe("UsagePage", () => { error: null, } as any); mockUserDailyActivityAggregatedCall.mockClear(); + mockUserDailyActivityCall.mockClear(); mockTagListCall.mockClear(); mockUserDailyActivityAggregatedCall.mockResolvedValue(mockSpendData); + mockUseInfiniteUsers.mockReturnValue({ + data: { + pages: [ + { + users: [ + { user_id: "user-001", user_alias: "Alice", user_email: "alice@example.com" }, + { user_id: "user-002", user_alias: null, user_email: "bob@example.com" }, + { user_id: "user-003", user_alias: null, user_email: null }, + ], + page: 1, + total_pages: 1, + total_count: 3, + }, + ], + pageParams: [1], + }, + fetchNextPage: vi.fn(), + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + } as any); mockTagListCall.mockResolvedValue({}); mockUseCustomers.mockReturnValue({ data: [], @@ -661,4 +698,434 @@ describe("UsagePage", () => { expect(entityUsageElements.length).toBeGreaterThan(0); }); }); + + describe("admin user selector", () => { + it("should render user selector for admin users in global view", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + // Admin should see the user selector select element with the placeholder attribute + const userSelects = screen.getAllByRole("combobox"); + const userSelect = userSelects.find( + (el) => el.getAttribute("placeholder") === "All Users (Global View)", + ); + expect(userSelect).toBeDefined(); + }); + + it("should format user options with alias when available", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + // User with alias should show "alias (id)" + expect(screen.getByText("Alice (user-001)")).toBeInTheDocument(); + // User without alias but with email should show "email (id)" + expect(screen.getByText("bob@example.com (user-002)")).toBeInTheDocument(); + // User with neither alias nor email should show just the id + expect(screen.getByText("user-003")).toBeInTheDocument(); + }); + + it("should call useInfiniteUsers with debounced search", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + // useInfiniteUsers should be called with default page size + expect(mockUseInfiniteUsers).toHaveBeenCalledWith(50, undefined); + }); + + it("should deduplicate users across pages", async () => { + mockUseInfiniteUsers.mockReturnValue({ + data: { + pages: [ + { + users: [ + { user_id: "user-dup", user_alias: "DupUser", user_email: null }, + ], + page: 1, + total_pages: 2, + total_count: 2, + }, + { + users: [ + { user_id: "user-dup", user_alias: "DupUser", user_email: null }, + { user_id: "user-unique", user_alias: "UniqueUser", user_email: null }, + ], + page: 2, + total_pages: 2, + total_count: 2, + }, + ], + pageParams: [1, 2], + }, + fetchNextPage: vi.fn(), + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + // Duplicate user should appear only once + const dupElements = screen.getAllByText("DupUser (user-dup)"); + expect(dupElements).toHaveLength(1); + // Unique user should also appear + expect(screen.getByText("UniqueUser (user-unique)")).toBeInTheDocument(); + }); + + it("should pass selected userId to aggregated call", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + // Initially called with null (global view for admin) + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledWith( + "test-token", + expect.any(Date), + expect.any(Date), + null, + ); + }); + }); + + describe("non-admin user behavior", () => { + it("should not render user selector for non-admin users", async () => { + mockUseAuthorized.mockReturnValue({ + isLoading: false, + isAuthorized: true, + token: "mock-token", + accessToken: "test-token", + userId: "user-123", + userEmail: "test@example.com", + userRole: "Internal User", + premiumUser: false, + disabledPersonalKeyCreation: false, + showSSOBanner: false, + }); + + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + // Non-admin should not see the user selector + const userSelects = screen.getAllByRole("combobox"); + const userSelect = userSelects.find( + (el) => el.getAttribute("placeholder") === "All Users (Global View)", + ); + expect(userSelect).toBeUndefined(); + }); + + it("should always pass own userId for non-admin users", async () => { + mockUseAuthorized.mockReturnValue({ + isLoading: false, + isAuthorized: true, + token: "mock-token", + accessToken: "test-token", + userId: "user-123", + userEmail: "test@example.com", + userRole: "Internal User", + premiumUser: false, + disabledPersonalKeyCreation: false, + showSSOBanner: false, + }); + + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledWith( + "test-token", + expect.any(Date), + expect.any(Date), + "user-123", + ); + }); + }); + }); + + describe("aggregated endpoint fallback", () => { + it("should fall back to paginated calls when aggregated endpoint fails", async () => { + mockUserDailyActivityAggregatedCall.mockRejectedValue(new Error("Aggregated endpoint not available")); + mockUserDailyActivityCall.mockResolvedValue({ + ...mockSpendData, + metadata: { + ...mockSpendData.metadata, + total_pages: 1, + page: 1, + }, + }); + + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + expect(mockUserDailyActivityCall).toHaveBeenCalled(); + }); + + // Should still render the data from the paginated fallback + expect(screen.getByText("1,500")).toBeInTheDocument(); + }); + + it("should aggregate multiple pages when paginated endpoint has more than 1 page", async () => { + mockUserDailyActivityAggregatedCall.mockRejectedValue(new Error("Not available")); + + const page1Data = { + results: [mockSpendData.results[0]], + metadata: { + total_spend: 60, + total_api_requests: 700, + total_successful_requests: 680, + total_failed_requests: 20, + total_tokens: 35000, + total_pages: 2, + page: 1, + }, + }; + + const page2Data = { + results: [ + { + ...mockSpendData.results[0], + date: "2025-01-02", + }, + ], + metadata: { + total_spend: 65.75, + total_api_requests: 800, + total_successful_requests: 770, + total_failed_requests: 30, + total_tokens: 40000, + total_pages: 2, + page: 2, + }, + }; + + mockUserDailyActivityCall + .mockResolvedValueOnce(page1Data) + .mockResolvedValueOnce(page2Data); + + renderWithProviders(); + + await waitFor(() => { + // Both pages should have been fetched + expect(mockUserDailyActivityCall).toHaveBeenCalledTimes(2); + }); + + // Verify first page call + expect(mockUserDailyActivityCall).toHaveBeenCalledWith( + "test-token", + expect.any(Date), + expect.any(Date), + 1, + null, + ); + + // Verify second page call + expect(mockUserDailyActivityCall).toHaveBeenCalledWith( + "test-token", + expect.any(Date), + expect.any(Date), + 2, + null, + ); + }); + }); + + describe("MCP Server Activity tab", () => { + it("should render MCP Server Activity tab", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + // The tab list should contain MCP Server Activity + expect(screen.getByText("MCP Server Activity")).toBeInTheDocument(); + }); + }); + + describe("User Agent Activity view", () => { + it("should render User Agent Activity component when view is selected", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + const usageSelect = screen.getByTestId("usage-view-select"); + act(() => { + fireEvent.change(usageSelect, { target: { value: "user-agent-activity" } }); + }); + + await waitFor(() => { + // "User Agent Activity" appears both in the select option and in the rendered component + const elements = screen.getAllByText("User Agent Activity"); + expect(elements.length).toBeGreaterThanOrEqual(2); + }); + }); + }); + + describe("Export Data button", () => { + it("should render Export Data button in global view for admin", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + expect(screen.getByText("Export Data")).toBeInTheDocument(); + }); + }); + + describe("model view toggle", () => { + it("should show Public Model Name view by default", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + // Default should be "groups" view showing "Top Public Model Names" + expect(screen.getByText("Top Public Model Names")).toBeInTheDocument(); + expect(screen.getByText("Public Model Name")).toBeInTheDocument(); + expect(screen.getByText("Litellm Model Name")).toBeInTheDocument(); + }); + + it("should switch to Litellm Model Name view on toggle click", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + // Click the "Litellm Model Name" toggle + const litellmToggle = screen.getByText("Litellm Model Name"); + act(() => { + fireEvent.click(litellmToggle); + }); + + // Title should change to "Top Litellm Models" + await waitFor(() => { + expect(screen.getByText("Top Litellm Models")).toBeInTheDocument(); + }); + }); + + it("should switch back to Public Model Name view", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + // Switch to individual first + const litellmToggle = screen.getByText("Litellm Model Name"); + act(() => { + fireEvent.click(litellmToggle); + }); + + await waitFor(() => { + expect(screen.getByText("Top Litellm Models")).toBeInTheDocument(); + }); + + // Switch back to groups + const publicToggle = screen.getByText("Public Model Name"); + act(() => { + fireEvent.click(publicToggle); + }); + + await waitFor(() => { + expect(screen.getByText("Top Public Model Names")).toBeInTheDocument(); + }); + }); + }); + + describe("customer usage banner", () => { + it("should show and be dismissible in customer view", async () => { + mockUseCustomers.mockReturnValue({ + data: mockCustomers, + isLoading: false, + error: null, + } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + const usageSelect = screen.getByTestId("usage-view-select"); + act(() => { + fireEvent.change(usageSelect, { target: { value: "customer" } }); + }); + + await waitFor(() => { + expect(screen.getByText("Customer usage is a new feature.")).toBeInTheDocument(); + }); + + // Click the close button + const closeButton = screen.getByLabelText("Close"); + act(() => { + fireEvent.click(closeButton); + }); + + await waitFor(() => { + expect(screen.queryByText("Customer usage is a new feature.")).not.toBeInTheDocument(); + }); + }); + }); + + describe("agent usage banner", () => { + it("should show agent usage banner with A2A info", async () => { + mockUseAgents.mockReturnValue({ + data: { agents: mockAgents }, + isLoading: false, + error: null, + } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + const usageSelect = screen.getByTestId("usage-view-select"); + act(() => { + fireEvent.change(usageSelect, { target: { value: "agent" } }); + }); + + await waitFor(() => { + expect(screen.getByText("Agent usage (A2A) is a new feature.")).toBeInTheDocument(); + }); + }); + }); + + describe("tab navigation in global view", () => { + it("should render all expected tabs", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + expect(screen.getByText("Cost")).toBeInTheDocument(); + expect(screen.getByText("Model Activity")).toBeInTheDocument(); + expect(screen.getByText("Key Activity")).toBeInTheDocument(); + expect(screen.getByText("MCP Server Activity")).toBeInTheDocument(); + expect(screen.getByText("Endpoint Activity")).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 688ee73767f..f81da6e2455 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -6,7 +6,7 @@ * Works at 1m+ spend logs, by querying an aggregate table instead. */ -import { InfoCircleOutlined } from "@ant-design/icons"; +import { InfoCircleOutlined, LoadingOutlined, UserOutlined } from "@ant-design/icons"; import { BarChart, Card, @@ -21,13 +21,15 @@ import { Text, Title } from "@tremor/react"; -import { Alert, Segmented, Tooltip } from "antd"; -import React, { useCallback, useEffect, useMemo, useState } from "react"; +import { Alert, Segmented, Select, Tooltip } from "antd"; +import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; +import React, { useCallback, useEffect, useMemo, useState, type UIEvent } from "react"; import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; +import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { Button } from "@tremor/react"; import { all_admin_roles } from "../../../utils/roles"; @@ -81,6 +83,62 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const { data: currentUser } = useCurrentUser(); console.log(`currentUser: ${JSON.stringify(currentUser)}`); console.log(`currentUser max budget: ${currentUser?.max_budget}`); + const isAdmin = all_admin_roles.includes(userRole || ""); + + // Debounced search for user selector + const [userSearchInput, setUserSearchInput] = useState(""); + const [debouncedUserSearch, setDebouncedUserSearch] = useDebouncedState("", { + wait: 300, + }); + + const { + data: usersInfiniteData, + fetchNextPage: fetchNextUsersPage, + hasNextPage: hasNextUsersPage, + isFetchingNextPage: isFetchingNextUsersPage, + isLoading: isLoadingUsers, + } = useInfiniteUsers(50, debouncedUserSearch || undefined); + + const userOptions = useMemo(() => { + if (!usersInfiniteData?.pages) return []; + const seen = new Set(); + const result: { value: string; label: string }[] = []; + for (const page of usersInfiniteData.pages) { + for (const user of page.users) { + if (seen.has(user.user_id)) continue; + seen.add(user.user_id); + result.push({ + value: user.user_id, + label: user.user_alias + ? `${user.user_alias} (${user.user_id})` + : user.user_email + ? `${user.user_email} (${user.user_id})` + : user.user_id, + }); + } + } + return result; + }, [usersInfiniteData]); + + const handleUserSearchChange = (value: string) => { + setUserSearchInput(value); + setDebouncedUserSearch(value); + }; + + const handleUserPopupScroll = (e: UIEvent) => { + const target = e.currentTarget; + const scrollRatio = + (target.scrollTop + target.clientHeight) / target.scrollHeight; + if (scrollRatio >= 0.8 && hasNextUsersPage && !isFetchingNextUsersPage) { + fetchNextUsersPage(); + } + }; + + // For admins: null means global view (all users), a string means filter by that user + // For non-admins: always set to their own user ID + const [selectedUserId, setSelectedUserId] = useState( + isAdmin ? null : (userID || null) + ); const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups"); const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false); const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false); @@ -107,6 +165,13 @@ const UsagePage: React.FC = ({ teams, organizations }) => { getAllTags(); }, [accessToken]); + // Sync selectedUserId when auth state settles (isAdmin/userID may be null on initial render) + useEffect(() => { + if (!isAdmin && userID) { + setSelectedUserId(userID); + } + }, [isAdmin, userID]); + // Derived states from userSpendData const totalSpend = userSpendData.metadata?.total_spend || 0; @@ -301,6 +366,9 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const fetchUserSpendData = useCallback(async () => { if (!accessToken || !dateValue.from || !dateValue.to) return; + // For non-admins, always pass their own user_id + const effectiveUserId = isAdmin ? selectedUserId : (userID || null); + setLoading(true); // Create new Date objects to avoid mutating the original dates @@ -310,14 +378,14 @@ const UsagePage: React.FC = ({ teams, organizations }) => { try { // Prefer aggregated endpoint to avoid many page requests try { - const aggregated = await userDailyActivityAggregatedCall(accessToken, startTime, endTime); + const aggregated = await userDailyActivityAggregatedCall(accessToken, startTime, endTime, effectiveUserId); setUserSpendData(aggregated); return; } catch (e) { // Fallback to paginated calls if aggregated endpoint is unavailable } - const firstPageData = await userDailyActivityCall(accessToken, startTime, endTime); + const firstPageData = await userDailyActivityCall(accessToken, startTime, endTime, 1, effectiveUserId); if (firstPageData.metadata.total_pages <= 1) { setUserSpendData(firstPageData); @@ -328,7 +396,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const aggregatedMetadata = { ...firstPageData.metadata }; for (let page = 2; page <= firstPageData.metadata.total_pages; page++) { - const pageData = await userDailyActivityCall(accessToken, startTime, endTime, page); + const pageData = await userDailyActivityCall(accessToken, startTime, endTime, page, effectiveUserId); allResults.push(...pageData.results); if (pageData.metadata) { aggregatedMetadata.total_spend += pageData.metadata.total_spend || 0; @@ -349,7 +417,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { setLoading(false); setIsDateChanging(false); } - }, [accessToken, dateValue.from, dateValue.to]); + }, [accessToken, dateValue.from, dateValue.to, selectedUserId, isAdmin, userID]); // Super responsive date change handler const handleDateChange = useCallback((newValue: DateRangePickerValue) => { @@ -423,12 +491,13 @@ const UsagePage: React.FC = ({ teams, organizations }) => { setUsageView(value)} - isAdmin={all_admin_roles.includes(userRole || "")} + isAdmin={isAdmin} /> {/* Your Usage Panel */} {usageView === "global" && ( + <>
@@ -460,24 +529,61 @@ const UsagePage: React.FC = ({ teams, organizations }) => { {/* Total Spend Card */} - - Project Spend{" "} - {dateValue.from && dateValue.to && ( - <> - {dateValue.from.toLocaleDateString("en-US", { - month: "short", - day: "numeric", - year: dateValue.from.getFullYear() !== dateValue.to.getFullYear() ? "numeric" : undefined, - })} - {" - "} - {dateValue.to.toLocaleDateString("en-US", { - month: "short", - day: "numeric", - year: "numeric", - })} - +
+ + Project Spend{" "} + {dateValue.from && dateValue.to && ( + <> + {dateValue.from.toLocaleDateString("en-US", { + month: "short", + day: "numeric", + year: dateValue.from.getFullYear() !== dateValue.to.getFullYear() ? "numeric" : undefined, + })} + {" - "} + {dateValue.to.toLocaleDateString("en-US", { + month: "short", + day: "numeric", + year: "numeric", + })} + + )} + + {isAdmin && ( +
+ +