From 432d99f3ee31bd6a19fa014d75cd9bfdda0283a8 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 25 Jun 2026 23:48:21 -0700 Subject: [PATCH 1/9] test(pass-through): grant allowed_passthrough_routes so langfuse auth=true test reaches rpm path #29256 made auth=true pass-through routes deny-by-default unless the key/team has allowed_passthrough_routes configured, but this integration test was not updated. The test key had no allowlist, so the auth=true parametrizations (rpm_limit=0 -> expect 429, rpm_limit=2 -> expect 207) now hit the 403 gate in auth before reaching the rpm/forwarding logic they mean to exercise. Grant the test key allowed_passthrough_routes for /api/public/ingestion so it clears the gate. Also removes a latent order-dependency: the case only passed locally when an earlier (auth=false) parametrization registered the route first; under worker isolation (CI xdist) it failed with 403. --- tests/local_testing/test_pass_through_endpoints.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/local_testing/test_pass_through_endpoints.py b/tests/local_testing/test_pass_through_endpoints.py index 68ba62bcbab..eeb29dea531 100644 --- a/tests/local_testing/test_pass_through_endpoints.py +++ b/tests/local_testing/test_pass_through_endpoints.py @@ -402,7 +402,9 @@ async def test_aaapass_through_endpoint_pass_through_keys_langfuse( mock_api_key = "sk-my-test-key" cache_value = UserAPIKeyAuth( - token=hash_token(mock_api_key), rpm_limit=rpm_limit + token=hash_token(mock_api_key), + rpm_limit=rpm_limit, + metadata={"allowed_passthrough_routes": ["/api/public/ingestion"]}, ) _cohere_api_key = os.environ.get("COHERE_API_KEY") From 5a1c7839be05e84aa57d463b71478141387a0f0a Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 26 Jun 2026 10:29:07 -0700 Subject: [PATCH 2/9] feat(mistral): add mistral/mistral-ocr-2512 (OCR 3) to cost map (#31463) Adds the OCR 3 model (mistral-ocr-2512) released 2025-12-18 to both the root and bundled backup cost maps at $2 / 1000 pages and $3 / 1000 annotated pages, mirroring the existing Mistral OCR entries. Regresses the pricing in both maps and verifies completion_cost scales per page. --- ...odel_prices_and_context_window_backup.json | 10 ++++ model_prices_and_context_window.json | 10 ++++ .../llms/mistral/ocr/test_mistral_ocr_cost.py | 51 +++++++++++++++++++ 3 files changed, 71 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5cf5af8a22a..6a65977e909 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25610,6 +25610,16 @@ ], "source": "https://mistral.ai/pricing#api-pricing" }, + "mistral/mistral-ocr-2512": { + "litellm_provider": "mistral", + "ocr_cost_per_page": 0.002, + "annotation_cost_per_page": 0.003, + "mode": "ocr", + "supported_endpoints": [ + "/v1/ocr" + ], + "source": "https://mistral.ai/pricing#api-pricing" + }, "mistral/magistral-medium-latest": { "input_cost_per_token": 2e-06, "litellm_provider": "mistral", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ed1b2fb7db9..d5af7391439 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25774,6 +25774,16 @@ ], "source": "https://mistral.ai/pricing#api-pricing" }, + "mistral/mistral-ocr-2512": { + "litellm_provider": "mistral", + "ocr_cost_per_page": 0.002, + "annotation_cost_per_page": 0.003, + "mode": "ocr", + "supported_endpoints": [ + "/v1/ocr" + ], + "source": "https://mistral.ai/pricing#api-pricing" + }, "mistral/magistral-medium-latest": { "input_cost_per_token": 2e-06, "litellm_provider": "mistral", diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index b98b0c22e50..c7f959826fe 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -5,6 +5,9 @@ for mistral-ocr-4-0 and mistral-ocr-latest, which now both resolve to OCR 4 at $4 / 1000 pages. """ +import json +from pathlib import Path + import pytest import litellm @@ -13,6 +16,14 @@ from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUs OCR4_COST_PER_PAGE = 0.004 +REPO_ROOT = Path(__file__).parents[5] +MAIN_COST_MAP = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_COST_MAP = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +OCR3_MODEL = "mistral/mistral-ocr-2512" +OCR3_COST_PER_PAGE = 0.002 +OCR3_ANNOTATION_COST_PER_PAGE = 0.003 + def _ocr_response(model: str, pages_processed: int) -> OCRResponse: return OCRResponse( @@ -38,3 +49,43 @@ def test_ocr4_cost_scales_with_pages(model: str, pages_processed: int) -> None: call_type="ocr", ) assert cost == pytest.approx(OCR4_COST_PER_PAGE * pages_processed) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force get_model_info to resolve against the in-repo cost map instead of the + remote one fetched at import time, which does not yet carry OCR 3 pricing.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("cost_map_path", [MAIN_COST_MAP, BACKUP_COST_MAP]) +def test_ocr3_pricing_entry(cost_map_path: Path) -> None: + with open(cost_map_path) as f: + info = json.load(f).get(OCR3_MODEL) + + assert info is not None, f"{OCR3_MODEL} missing from {cost_map_path.name}" + assert info["litellm_provider"] == "mistral" + assert info["mode"] == "ocr" + assert info["supported_endpoints"] == ["/v1/ocr"] + assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE + assert info["annotation_cost_per_page"] == OCR3_ANNOTATION_COST_PER_PAGE + + +def test_ocr3_model_info_price(local_model_cost_map) -> None: + info = litellm.get_model_info(model=OCR3_MODEL, custom_llm_provider="mistral") + assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE + + +@pytest.mark.parametrize("pages_processed", [1, 3, 10]) +def test_ocr3_cost_scales_with_pages(local_model_cost_map, pages_processed: int) -> None: + cost = completion_cost( + completion_response=_ocr_response("mistral-ocr-2512", pages_processed), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(OCR3_COST_PER_PAGE * pages_processed) From e99151bb95eaf1de26009940aa780b7f6f8b8cda Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 26 Jun 2026 11:25:56 -0700 Subject: [PATCH 3/9] feat(guardrails): make the Generic Guardrail resilient to built-in tools and errors (adopted from #31286) (#31461) * fix(guardrails): stop Generic Guardrail API 500 on built-in tools Requests carrying built-in tools (code_interpreter, file_search, ...) crashed the Generic Guardrail with a 500. GenericGuardrailAPIRequest.tools validated each tool against ChatCompletionToolParam, whose base TypedDict requires a function block, so a tool like {"type": "code_interpreter"} raised a Pydantic ValidationError before the request was ever sent. Type the field with a permissive GuardrailToolParam model (type required, extra=allow) so built-in tools validate and their config is forwarded to the guardrail intact instead of being stripped. * feat(guardrails): add complete fail-open (fail_on_error) to Generic Guardrail The Generic Guardrail already honored unreachable_fallback, which fails open only on network-unreachable errors. This wires up the existing generic fail_on_error config (so far implemented only by Model Armor) so that fail_on_error=false degrades any guardrail error to a critical-log warning and lets the request proceed as if the guardrail were absent. Only a valid guardrail response can act: a parsed BLOCKED decision still raises, while endpoint errors, malformed responses, and internal serialization or validation errors all fall through when fail_on_error=false. To cover that last class, the request construction now runs inside the protected block, so the kind of validation error that previously surfaced as a 500 is caught here too. Defaults to true (fail closed), matching today's behavior; turning it off is an explicit availability-over-security choice and is logged at critical level on every bypass. * test(guardrails): cover fail_on_error on the response path The existing fail_on_error tests all drive the request path. Add response-path (input_type=response) coverage: an endpoint error proceeds unchanged under fail_on_error=false, and a valid BLOCKED decision still raises. Guards against a future regression that special-cases input_type in the error handling. * style(guardrails): black-format the fail-open guard expression CI runs black (line-length 88) over litellm/; the unreachable_fail_open assignment exceeded it. Wrap it to satisfy the formatter. * fix(guardrails): validate tools into GuardrailToolParam at the call site Changing the request field to List[GuardrailToolParam] left the construction passing List[ChatCompletionToolParam] (list is invariant), which tripped the basedpyright reportArgumentType budget gate. Validate each tool explicitly, which is what Pydantic did implicitly, so the types line up with no Any or suppression and the serialized payload is unchanged. * fix(guardrails): make fail-open log message accurate for non-network errors The fail-open path is now shared by fail_on_error, so it fires for any guardrail error, not just unreachability. The log said 'unreachable' even for an HTTP 400 or a malformed response; reword to 'error' (the status code and exception are already logged). Addresses the Greptile review's only finding. * fix(guardrails): align GenericGuardrailAPIResponse.tools with GuardrailToolParam Greptile flagged that the request side moved to GuardrailToolParam but the response side still annotated tools as List[ChatCompletionToolParam], which mandates a function block and contradicts the new built-in-tools support. Update the response annotation (and the now-unused import) so the two sides agree. Runtime is unchanged; from_dict stores the raw dicts and the only consumer assigns through to GenericGuardrailAPIInputs without inspecting the elements. --------- Co-authored-by: Itay Ovadia --- .../generic_guardrail_api/__init__.py | 1 + .../generic_guardrail_api.py | 53 +++--- litellm/types/guardrails.py | 7 +- .../guardrail_hooks/generic_guardrail_api.py | 31 +++- .../test_generic_guardrail_api.py | 168 ++++++++++++++++++ 5 files changed, 233 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index bb0d0a99b31..777c45bcd83 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -21,6 +21,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" unreachable_fallback=getattr( litellm_params, "unreachable_fallback", "fail_closed" ), + fail_on_error=getattr(litellm_params, "fail_on_error", True), extra_headers=getattr(litellm_params, "extra_headers", None), guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 790ee31f2e0..8d3abe2bdc3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -27,6 +27,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest, GenericGuardrailAPIResponse, + GuardrailToolParam, ) from litellm.types.utils import GenericGuardrailAPIInputs @@ -187,6 +188,7 @@ class GenericGuardrailAPI(CustomGuardrail): api_key: Optional[str] = None, additional_provider_specific_params: Optional[Dict[str, Any]] = None, unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", + fail_on_error: Optional[bool] = True, extra_headers: Optional[list] = None, **kwargs, ): @@ -223,6 +225,8 @@ class GenericGuardrailAPI(CustomGuardrail): unreachable_fallback ) + self.fail_on_error: bool = True if fail_on_error is None else fail_on_error + # Set supported event hooks if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [ @@ -299,7 +303,7 @@ class GenericGuardrailAPI(CustomGuardrail): f" http_status_code={http_status_code}" if http_status_code else "" ) verbose_proxy_logger.critical( - "Generic Guardrail API unreachable (fail-open). Proceeding without guardrail.%s " + "Generic Guardrail API error (fail-open). Proceeding without guardrail.%s " "guardrail_name=%s api_base=%s input_type=%s litellm_call_id=%s litellm_trace_id=%s", status_suffix, getattr(self, "guardrail_name", None), @@ -351,7 +355,10 @@ class GenericGuardrailAPI(CustomGuardrail): logging_obj: Optional["LiteLLMLoggingObj"], is_unreachable: bool = True, ) -> GenericGuardrailAPIInputs: - if is_unreachable and self.unreachable_fallback == "fail_open": + unreachable_fail_open = ( + is_unreachable and self.unreachable_fallback == "fail_open" + ) + if unreachable_fail_open or not self.fail_on_error: http_status_code = getattr( getattr(error, "response", None), "status_code", None ) @@ -432,26 +439,30 @@ class GenericGuardrailAPI(CustomGuardrail): extra_allowlist=extra_allowlist, ) - # Create request payload - guardrail_request = GenericGuardrailAPIRequest( - litellm_call_id=logging_obj.litellm_call_id if logging_obj else None, - litellm_trace_id=logging_obj.litellm_trace_id if logging_obj else None, - texts=texts, - request_data=user_metadata, - request_headers=inbound_headers, - litellm_version=litellm_version, - images=images, - tools=tools, - structured_messages=structured_messages, - tool_calls=tool_calls, - additional_provider_specific_params=additional_params, - input_type=input_type, - model=model, - ) - - headers = self._build_request_headers() - try: + # Create request payload + guardrail_request = GenericGuardrailAPIRequest( + litellm_call_id=logging_obj.litellm_call_id if logging_obj else None, + litellm_trace_id=logging_obj.litellm_trace_id if logging_obj else None, + texts=texts, + request_data=user_metadata, + request_headers=inbound_headers, + litellm_version=litellm_version, + images=images, + tools=( + [GuardrailToolParam.model_validate(t) for t in tools] + if tools + else None + ), + structured_messages=structured_messages, + tool_calls=tool_calls, + additional_provider_specific_params=additional_params, + input_type=input_type, + model=model, + ) + + headers = self._build_request_headers() + # Make the API request # Use mode="json" to ensure all iterables are converted to lists response = await self.async_handler.post( diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c9623d8595a..6a76dd337fa 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -750,7 +750,12 @@ class BaseLitellmParams( ) fail_on_error: Optional[bool] = Field( default=True, - description="Whether to fail the request if Model Armor encounters an error", + description=( + "Whether to fail the request if the guardrail encounters an error. " + "Implemented by guardrail='model_armor' and 'generic_guardrail_api'. " + "True (default) raises the error. False logs a critical error and lets the request proceed, " + "so only a valid guardrail response can block or modify it." + ), ) additional_provider_specific_params: Optional[Dict[str, Any]] = Field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 94f219a5fc6..2fe976463c4 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -1,17 +1,28 @@ from typing import Any, Dict, List, Literal, Optional, Union -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from typing_extensions import TYPE_CHECKING, TypedDict from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolCallChunk, - ChatCompletionToolParam, ) from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel from litellm.types.utils import ChatCompletionMessageToolCall +class GuardrailToolParam(BaseModel): + """A tool forwarded verbatim to the guardrail for inspection. + + Built-in tools (code_interpreter, file_search, ...) have no ``function`` block + and stash their config in tool-specific keys, so only ``type`` is required and + ``extra="allow"`` preserves the rest instead of stripping it. + """ + + model_config = ConfigDict(extra="allow") + type: str + + class GenericGuardrailAPIMetadata(TypedDict, total=False): user_api_key_hash: Optional[str] user_api_key_alias: Optional[str] @@ -39,6 +50,16 @@ class GenericGuardrailAPIOptionalParams(BaseModel): ), ) + fail_on_error: Optional[bool] = Field( + default=True, + description=( + "Behavior on any guardrail error, not just unreachability. " + "True (default) raises and blocks the request on error. " + "False logs a critical error and allows the request to proceed, so only a valid " + "guardrail response can block or modify it; broader than unreachable_fallback." + ), + ) + class GenericGuardrailAPIConfigModel( GuardrailConfigModel[GenericGuardrailAPIOptionalParams], @@ -65,7 +86,7 @@ class GenericGuardrailAPIRequest(BaseModel): ) structured_messages: Optional[List[AllMessageValues]] = None images: Optional[List[str]] = None - tools: Optional[List[ChatCompletionToolParam]] = None + tools: Optional[List[GuardrailToolParam]] = None texts: Optional[List[str]] = None request_data: GenericGuardrailAPIMetadata request_headers: Optional[Dict[str, str]] = Field( @@ -88,7 +109,7 @@ class GenericGuardrailAPIResponse: texts: Optional[List[str]] images: Optional[List[str]] - tools: Optional[List[ChatCompletionToolParam]] + tools: Optional[List[GuardrailToolParam]] action: str blocked_reason: Optional[str] @@ -98,7 +119,7 @@ class GenericGuardrailAPIResponse: texts: Optional[List[str]] = None, blocked_reason: Optional[str] = None, images: Optional[List[str]] = None, - tools: Optional[List[ChatCompletionToolParam]] = None, + tools: Optional[List[GuardrailToolParam]] = None, ): self.action = action self.blocked_reason = blocked_reason diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 6ec793a1bb0..399442a5f71 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -1021,3 +1021,171 @@ class TestMultimodalSupport: call_args = mock_post.call_args json_payload = call_args.kwargs["json"] assert isinstance(json_payload["structured_messages"], list) + + +class TestToolSupport: + """Test tool handling in guardrail requests""" + + @pytest.mark.asyncio + async def test_builtin_tools_without_function_block_do_not_crash( + self, generic_guardrail + ): + """Built-in tools (code_interpreter, file_search) have no `function` block. + + Regression for a 500 where serializing them raised a Pydantic + ValidationError because the tool schema required `function`. The full + tool, including built-in tool config, must reach the guardrail intact. + """ + tools = [ + {"type": "function", "function": {"name": "get_weather", "parameters": {}}}, + {"type": "code_interpreter"}, + { + "type": "file_search", + "vector_store_ids": ["vs_1"], + "max_num_results": 5, + }, + ] + + mock_response = MagicMock() + mock_response.json.return_value = {"action": "NONE", "texts": ["hi"]} + mock_response.raise_for_status = MagicMock() + + with patch.object( + generic_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + await generic_guardrail.apply_guardrail( + inputs={"texts": ["hi"], "tools": tools}, + request_data={}, + input_type="request", + ) + + forwarded_tools = mock_post.call_args.kwargs["json"]["tools"] + + assert forwarded_tools == tools + + +class TestFailOnError: + """Test fail_on_error: complete fail-open on any guardrail error""" + + @pytest.fixture + def fail_open_guardrail(self): + return GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-fail-open-guardrail", + event_hook="pre_call", + default_on=True, + fail_on_error=False, + ) + + @pytest.mark.asyncio + async def test_endpoint_error_continues_when_fail_on_error_false( + self, fail_open_guardrail + ): + """A non-unreachable endpoint error (HTTP 400) is swallowed and the request proceeds unchanged.""" + error = httpx.HTTPStatusError( + "bad request", request=MagicMock(), response=MagicMock(status_code=400) + ) + with patch.object( + fail_open_guardrail.async_handler, "post", side_effect=error + ): + result = await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="request", + ) + + assert result == {"texts": ["hi"]} + + @pytest.mark.asyncio + async def test_internal_error_continues_without_calling_endpoint( + self, fail_open_guardrail + ): + """An error while building the request (here: invalid input_type) fails open too. + + Proves the request construction runs inside the protected block: the + endpoint is never called, yet the request still proceeds unchanged. + """ + with patch.object(fail_open_guardrail.async_handler, "post") as mock_post: + result = await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="bogus", # type: ignore[arg-type] + ) + + mock_post.assert_not_called() + assert result == {"texts": ["hi"]} + + @pytest.mark.asyncio + async def test_valid_block_still_blocks_when_fail_on_error_false( + self, fail_open_guardrail + ): + """Only a valid response acts: a BLOCKED decision still raises even with fail_on_error=False.""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "BLOCKED", + "blocked_reason": "policy violation", + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + fail_open_guardrail.async_handler, "post", return_value=mock_response + ): + with pytest.raises(GuardrailRaisedException): + await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="request", + ) + + @pytest.mark.asyncio + async def test_endpoint_error_raises_by_default(self, generic_guardrail): + """Default fail_on_error=True keeps blocking on a non-unreachable endpoint error.""" + error = httpx.HTTPStatusError( + "bad request", request=MagicMock(), response=MagicMock(status_code=400) + ) + with patch.object(generic_guardrail.async_handler, "post", side_effect=error): + with pytest.raises(Exception, match="Generic Guardrail API failed"): + await generic_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="request", + ) + + @pytest.mark.asyncio + async def test_response_path_continues_when_fail_on_error_false( + self, fail_open_guardrail + ): + """fail_on_error governs the response path identically to the request path.""" + error = httpx.HTTPStatusError( + "bad request", request=MagicMock(), response=MagicMock(status_code=400) + ) + with patch.object( + fail_open_guardrail.async_handler, "post", side_effect=error + ): + result = await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["model output"]}, + request_data={}, + input_type="response", + ) + + assert result == {"texts": ["model output"]} + + @pytest.mark.asyncio + async def test_response_path_valid_block_still_blocks(self, fail_open_guardrail): + """On the response path too, a valid BLOCKED decision raises despite fail_on_error=False.""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "BLOCKED", + "blocked_reason": "policy violation", + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + fail_open_guardrail.async_handler, "post", return_value=mock_response + ): + with pytest.raises(GuardrailRaisedException): + await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["model output"]}, + request_data={}, + input_type="response", + ) From 8622df5a6335aaeadb02ff46d0a926837e641acc Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 26 Jun 2026 12:44:07 -0700 Subject: [PATCH 4/9] fix(build): restore pure-Python uv_build backend to unblock PyPI publish #31267 switched the build backend to maturin to bundle the Rust OCR bridge into the litellm wheel. On the release runner, maturin builds with `--compatibility off` and emits a `cp312-cp312-linux_x86_64` wheel. PyPI only accepts `manylinux*` Linux wheels, so the upload is rejected with `400 Bad Request - unsupported platform tag 'linux_x86_64'`, which blocked the 1.91.0.dev2 publish. Revert `[build-system]` to the pure-Python `uv_build` backend and restore `[tool.uv.build-backend]` so the release once again produces a universal `litellm--py3-none-any.whl` that installs on every OS / arch / Python. The Rust bridge stays optional: `litellm/rust_bridge/loader.py` already falls back gracefully when `_native` is absent (`try/except ImportError`), so runtime behavior is unchanged - the native module is simply not bundled. Rust acceleration can ship as a separate platform-specific package later. --- pyproject.toml | 32 ++++++++++++++------------------ 1 file changed, 14 insertions(+), 18 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6e99d81f8f3..23f02cb5762 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -234,24 +234,8 @@ healthcheck = [ ] [build-system] -requires = ["maturin==1.9.4"] -build-backend = "maturin" - -[tool.maturin] -manifest-path = "litellm-rust/crates/python-bridge/Cargo.toml" -module-name = "litellm.rust_bridge._native" -python-source = "." -bindings = "pyo3" -exclude = [ - "litellm/proxy/enterprise", - "litellm/proxy/enterprise/**", - "**/__pycache__", - "**/__pycache__/**", - "**/.pytest_cache", - "**/.pytest_cache/**", - "**/.ruff_cache", - "**/.ruff_cache/**", -] +requires = ["uv_build==0.11.8"] +build-backend = "uv_build" [tool.uv] constraint-dependencies = [ @@ -269,6 +253,18 @@ litellm-enterprise = { workspace = true } [tool.uv.workspace] members = ["enterprise", "litellm-proxy-extras"] +[tool.uv.build-backend] +module-root = "" +source-exclude = [ + "litellm/proxy/enterprise", + "**/__pycache__", + "**/__pycache__/**", + "**/.pytest_cache", + "**/.pytest_cache/**", + "**/.ruff_cache", + "**/.ruff_cache/**", +] + [tool.isort] profile = "black" From f55d13ebba6981899f0ba6cebc0d6c1d30aaebd6 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 22:59:30 +0300 Subject: [PATCH 5/9] fix(team): persist budget_duration on /team/member_add member budgets (#31443) /team/member_add could not set budget_duration on an individual member budget. add_new_member created the budget row with only max_budget and allowed_models, and TeamMemberAddRequest had no budget_duration field, so a member added with an explicit per-member budget while the team ran a recurring member budget got a lifetime cap instead of a recurring allowance. Thread budget_duration from TeamMemberAddRequest through _process_team_members into add_new_member, and pull the member-budget resolution into a helper that writes budget_duration plus a computed budget_reset_at. When only a budget_duration is supplied and the team has a default member budget, the default is cloned and its reset window overridden so the member keeps the default's max_budget rather than becoming uncapped; a duration with no team default creates a window-only budget. Invalid durations are rejected with a 400 before any DB write, symmetric with /team/member_update. The available-team self-join bypass only grants the ability to join, so reject per-member budget and model controls (max_budget_in_team, budget_duration, allowed_models) for non-admin self-join callers in _validate_team_member_add_permissions, before any DB write. Otherwise a self-joining non-admin could set their own cap, reset window, or model scope past the team default; admins, team admins, and org admins are unaffected and a clean self-join still inherits the team default budget. Resolves LIT-4052 --- litellm/proxy/_types.py | 4 + .../management_endpoints/team_endpoints.py | 24 +++ litellm/proxy/management_helpers/utils.py | 95 ++++++--- .../test_team_endpoints.py | 90 ++++++++ .../test_management_helpers_utils.py | 200 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 6 files changed, 390 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ed030c31b19..559720eb64e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3658,6 +3658,10 @@ class TeamMemberAddRequest(MemberAddRequest): default=None, description="Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits", ) + budget_duration: Optional[str] = Field( + default=None, + description="Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets.", + ) allowed_models: Optional[List[str]] = Field( default=None, description="List of models this team member can access. If not set, inherits the team's default_team_member_models or all team models.", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 249cfdae4eb..f533774ce18 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2225,6 +2225,26 @@ async def _validate_team_member_add_permissions( }, ) + # Available-team self-join grants only the ability to join; per-member + # budget and model controls stay admin-only. Reject them here so a + # self-joining non-admin cannot set their own cap, reset window, or model + # scope via the bypass. + if ( + data.max_budget_in_team is not None + or data.budget_duration is not None + or data.allowed_models is not None + ): + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Available-team self-join cannot set per-member budget or " + "model controls (max_budget_in_team, budget_duration, " + "allowed_models); these are admin-only." + ) + }, + ) + # Available-team self-join: caller may add only themselves, only as a # standard user. Enforce that here so the bypass cannot be used as a # privilege-escalation or cross-user-injection primitive. @@ -2290,6 +2310,7 @@ async def _process_team_members( team_id=data.team_id, default_team_budget_id=default_team_budget_id, allowed_models=member_allowed_models, + budget_duration=data.budget_duration, ) except Exception as e: raise HTTPException( @@ -2315,6 +2336,7 @@ async def _process_team_members( team_id=data.team_id, default_team_budget_id=default_team_budget_id, allowed_models=member_allowed_models, + budget_duration=data.budget_duration, ) except Exception as e: raise HTTPException( @@ -2578,6 +2600,8 @@ async def team_member_add( except HTTPException as e: raise e + _validate_budget_duration(data.budget_duration) + prisma_client = cast(PrismaClient, prisma_client) existing_team_row = await get_team_object( diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 4ce8633ed57..0c5439e26e2 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -167,6 +167,7 @@ async def _clone_team_default_budget_for_member( default_team_budget_id: str, user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, + budget_duration_override: Optional[str] = None, ) -> Optional[str]: """ Create a new budget row that copies the values from the team's default @@ -176,6 +177,10 @@ async def _clone_team_default_budget_for_member( Used when adding a new team member without an explicit per-member budget, so the member starts with the team default's values but gets their own private budget row (which can be edited independently). + + ``budget_duration_override`` replaces the default's reset window for this + member while keeping the default's other limits, so an admin can set a + member's reset cadence without discarding the team default's max_budget. """ default_budget = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": default_team_budget_id} @@ -198,6 +203,9 @@ async def _clone_team_default_budget_for_member( continue cloned_data[field] = value + if budget_duration_override is not None: + cloned_data["budget_duration"] = budget_duration_override + # Start the member's budget window at clone time, not the pool's reset # timestamp — otherwise a member joining mid-cycle inherits a stale reset. if cloned_data.get("budget_duration"): @@ -209,6 +217,55 @@ async def _clone_team_default_budget_for_member( return new_budget.budget_id +async def _resolve_member_budget_id( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str, + max_budget_in_team: Optional[float], + allowed_models: Optional[list[str]], + budget_duration: Optional[str], + default_team_budget_id: Optional[str], +) -> Optional[str]: + """ + Resolve the budget a new team member should be linked to. + + Explicit per-member limits create a fresh budget. Otherwise the team's + default member budget is cloned (with ``budget_duration`` overriding its + reset window while keeping its other limits). A lone ``budget_duration`` + with no team default creates a window-only budget. With nothing set the + member gets no budget. + """ + has_explicit_limit = max_budget_in_team is not None or allowed_models is not None + + if not has_explicit_limit and default_team_budget_id is not None: + return await _clone_team_default_budget_for_member( + prisma_client=prisma_client, + default_team_budget_id=default_team_budget_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + budget_duration_override=budget_duration, + ) + + if not has_explicit_limit and budget_duration is None: + return None + + budget_data: dict = { + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + if max_budget_in_team is not None: + budget_data["max_budget"] = max_budget_in_team + if allowed_models is not None: + budget_data["allowed_models"] = allowed_models + if budget_duration is not None: + budget_data["budget_duration"] = budget_duration + budget_data["budget_reset_at"] = get_budget_reset_time( + budget_duration=budget_duration + ) + response = await BudgetRepository(prisma_client).table.create(data=budget_data) + return response.budget_id + + async def add_new_member( new_member: Member, max_budget_in_team: Optional[float], @@ -218,6 +275,7 @@ async def add_new_member( litellm_proxy_admin_name: str, default_team_budget_id: Optional[str] = None, allowed_models: Optional[List[str]] = None, + budget_duration: Optional[str] = None, ) -> Tuple[LiteLLM_UserTable, Optional[LiteLLM_TeamMembership]]: """ Add a new member to a team @@ -278,34 +336,15 @@ async def add_new_member( }, ) - # Check if trying to set a budget or model scope for team member - if max_budget_in_team is not None or allowed_models is not None: - # create a new budget item for this member - budget_data: dict = { - "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - } - if max_budget_in_team is not None: - budget_data["max_budget"] = max_budget_in_team - if allowed_models is not None: - budget_data["allowed_models"] = allowed_models - response = await BudgetRepository(prisma_client).table.create(data=budget_data) - - _budget_id = response.budget_id - elif default_team_budget_id is not None: - # No per-member budget was provided, but the team has a default member - # budget. Clone the default budget into a new row for this user so that - # later edits to one member's budget do not bleed into other members. - # If the default no longer exists in the DB, fall back to no budget. - _budget_id = await _clone_team_default_budget_for_member( - prisma_client=prisma_client, - default_team_budget_id=default_team_budget_id, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ) - else: - # No per-member budget and no team default → member gets no budget. - _budget_id = None + _budget_id = await _resolve_member_budget_id( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + max_budget_in_team=max_budget_in_team, + allowed_models=allowed_models, + budget_duration=budget_duration, + default_team_budget_id=default_team_budget_id, + ) if _budget_id and returned_user is not None and returned_user.user_id is not None: _returned_team_membership = await TeamMembershipRepository( diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a649bc7225e..ecce92e7324 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1270,6 +1270,95 @@ async def test_available_team_self_join_blocks_admin_role_in_member_list(): assert exc_info.value.status_code == 403 +@pytest.mark.asyncio +@pytest.mark.parametrize( + "budget_control", + [ + {"max_budget_in_team": 1000.0}, + {"budget_duration": "1h"}, + {"allowed_models": ["gpt-4o"]}, + ], +) +async def test_available_team_self_join_blocks_member_budget_controls(budget_control): + """A self-joining non-admin must not be able to set their own per-member + budget or model controls via the available-team bypass; only proxy/team/org + admins may. Without this guard a self-joiner could shorten their budget + reset window or widen their cap/model scope past the team default.""" + from litellm.proxy._types import Member, TeamMemberAddRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + data = TeamMemberAddRequest( + team_id="public-team", + member=Member(role="user", user_id="alice"), + **budget_control, + ) + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + pytest.raises(HTTPException) as exc_info, + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=data, + ) + + assert exc_info.value.status_code == 403 + assert "admin-only" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_available_team_self_join_allows_no_budget_controls(): + """The clean self-join (no per-member budget/model controls) must still be + permitted, so the new guard does not break the legitimate join path.""" + from litellm.proxy._types import Member, TeamMemberAddRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + data = TeamMemberAddRequest( + team_id="public-team", + member=Member(role="user", user_id="alice"), + ) + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=data, + ) + + @pytest.mark.asyncio async def test_update_team_member_permissions_blocks_non_admin_via_available_team( mock_db_client, @@ -1385,6 +1474,7 @@ async def test_process_team_members_single_member(): team_id="test-team-123", default_team_budget_id="budget-123", allowed_models=None, + budget_duration=None, ) diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index 463aba6f744..dbbdca65cc6 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -1,6 +1,7 @@ import json import os import sys +from datetime import datetime, timezone from litellm._uuid import uuid from unittest.mock import AsyncMock, MagicMock @@ -283,6 +284,81 @@ async def test_add_new_member_clones_default_team_budget_id(): assert create_data["budget_id"] == test_cloned_budget_id +@pytest.mark.asyncio +async def test_add_new_member_budget_duration_only_clones_default_max_budget(): + """When only a budget_duration is given and the team has a default member + budget, the member must clone the default (keeping its max_budget) and just + override the reset window. Creating a fresh duration-only row instead would + silently drop the team default's cap, leaving the member uncapped.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="dur-clone-user", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": "dur-clone-user", + "user_email": None, + "teams": ["team-dc"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + mock_default_budget_row = MagicMock() + mock_default_budget_row.model_dump.return_value = { + "budget_id": "default-dc", + "max_budget": 100.0, + "soft_budget": None, + "max_parallel_requests": None, + "tpm_limit": 1000, + "rpm_limit": None, + "model_max_budget": None, + "budget_duration": "1d", + "allowed_models": [], + } + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_default_budget_row + ) + mock_cloned_budget_row = MagicMock() + mock_cloned_budget_row.budget_id = "cloned-dc" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_cloned_budget_row + ) + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": "team-dc", + "user_id": "dur-clone-user", + "budget_id": "cloned-dc", + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-dc", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id="default-dc", + budget_duration="7d", + ) + + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + cloned_create_data = ( + mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs["data"] + ) + assert cloned_create_data["max_budget"] == 100.0 # kept from the team default + assert cloned_create_data["budget_duration"] == "7d" # overridden by the caller + assert cloned_create_data["budget_reset_at"] > datetime.now(timezone.utc) + + @pytest.mark.asyncio async def test_add_new_member_no_budget_when_no_default_and_no_max_budget(): """ @@ -434,6 +510,130 @@ async def test_add_new_member_creates_new_budget_when_max_budget_provided(): assert create_data["budget_id"] == test_new_budget_id +@pytest.mark.asyncio +async def test_add_new_member_persists_budget_duration(): + """Regression for the member_add half of the recurring-member-budget gap: + a budget_duration passed to add_new_member must be written to the new + member budget along with a future budget_reset_at, so the per-member budget + recurs instead of acting as a lifetime cap.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="user-dur", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": "user-dur", + "user_email": None, + "teams": ["team-dur"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + mock_budget_response = MagicMock() + mock_budget_response.budget_id = "budget-dur" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_budget_response + ) + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": "team-dur", + "user_id": "user-dur", + "budget_id": "budget-dur", + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + await add_new_member( + new_member=new_member, + max_budget_in_team=10.0, + prisma_client=mock_prisma_client, + team_id="team-dur", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id=None, + allowed_models=["gpt-4o-mini"], + budget_duration="30d", + ) + + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + budget_data = mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs[ + "data" + ] + assert budget_data["max_budget"] == 10.0 + assert budget_data["allowed_models"] == ["gpt-4o-mini"] + assert budget_data["budget_duration"] == "30d" + reset_at = budget_data["budget_reset_at"] + assert isinstance(reset_at, datetime) + assert reset_at.tzinfo is not None + assert reset_at > datetime.now(timezone.utc) + + +@pytest.mark.asyncio +async def test_add_new_member_persists_budget_duration_without_max_budget(): + """budget_duration alone must still create a member budget; otherwise an + explicit recurring window passed without a cap would be silently dropped.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="user-dur2", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": "user-dur2", + "user_email": None, + "teams": ["team-dur2"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + mock_budget_response = MagicMock() + mock_budget_response.budget_id = "budget-dur2" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_budget_response + ) + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": "team-dur2", + "user_id": "user-dur2", + "budget_id": "budget-dur2", + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + _, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-dur2", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id=None, + budget_duration="7d", + ) + + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + budget_data = mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs[ + "data" + ] + assert budget_data["budget_duration"] == "7d" + assert budget_data["budget_reset_at"] > datetime.now(timezone.utc) + assert result_team_membership is not None + assert result_team_membership.budget_id == "budget-dur2" + + @pytest.mark.asyncio async def test_add_new_member_with_user_email_clones_default_budget(): """ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cdc91371dfe..43be9ad6e39 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30842,6 +30842,11 @@ export interface components { * @description List of models this team member can access. If not set, inherits the team's default_team_member_models or all team models. */ allowed_models?: string[] | null; + /** + * Budget Duration + * @description Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets. + */ + budget_duration?: string | null; /** * Max Budget In Team * @description Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits From 7209e139d64d9ee29ed2e0d51b561e4cb8c59b48 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 23:01:05 +0300 Subject: [PATCH 6/9] fix(spend): fold logs-tab total into the page query to avoid a separate COUNT(*) (#31423) The spend-logs UI list endpoint (/spend/logs/ui, /spend/logs/v2) ran a standalone SELECT COUNT(*) before the page query to compute total_pages. On sharded engines like YugabyteDB a COUNT(*) is a distributed RPC that contacts every tablet leader and aggregates partial results regardless of row count, so it hits the distributed RPC timeout and the logs tab 500s even on a one-minute window with a couple of rows. The startTime range cannot prune tablets because rows hash to tablets on request_id, not startTime. Fold the count into the same scan as the page data with COUNT(*) OVER () and read total off the returned rows, dropping the helper column before serialisation. One distributed scan per page load instead of two; the response shape is unchanged. An empty page carries no count row, in which case the total is zero. Resolves LIT-4027 --- .../spend_management_endpoints.py | 31 +++- .../test_spend_management_endpoints.py | 124 ++++++++++++-- .../test_spend_query_optimization.py | 151 ++++++++++++++++++ 3 files changed, 284 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 8ca048a0caa..3374877572b 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2018,11 +2018,6 @@ async def ui_view_spend_logs( order_column = sort_by order_direction = (sort_order or "desc").lower() - # Get total count of records - total_records = await SpendLogsRepository(prisma_client).table.count( - where=where_conditions, - ) - # Build raw SQL to fetch paginated data WITHOUT heavy columns # (messages, response, proxy_server_request can be hundreds of KB per row). # These are only needed in the detail endpoint /spend/logs/ui/{request_id}. @@ -2136,7 +2131,8 @@ async def ui_view_spend_logs( cache_hit, cache_key, request_tags, team_id, organization_id, end_user, requester_ip_address, session_id, status, mcp_namespaced_tool_name, agent_id, - COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms + COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms, + COUNT(*) OVER () AS total_count FROM "LiteLLM_SpendLogs" WHERE {" AND ".join(sql_conditions)} ORDER BY {_order_expr} {_sql_dir}{_nulls_clause} @@ -2146,13 +2142,34 @@ async def ui_view_spend_logs( data = await prisma_client.db.query_raw(sql_query, *sql_params) + # `COUNT(*) OVER ()` folds the total-match count into the same scan as the + # page data; a standalone `COUNT(*)` is a distributed RPC on sharded + # engines like YugabyteDB that contacts every tablet and times out + # regardless of row count (LIT-4027). The hot path (page 1 and in-range + # pages) always carries the count on its rows, so the count round trip is + # gone there. Only an out-of-range page overshoots the last row and comes + # back empty; fall back to a direct count there so total/total_pages stay + # accurate rather than collapsing to zero. + if data: + total_records = int(data[0]["total_count"]) + elif page > 1: + total_records = int( + await SpendLogsRepository(prisma_client).table.count( + where=where_conditions, + ) + ) + else: + total_records = 0 + # query_raw returns the JSONB `metadata` column as a string (the Prisma # serialiser bypasses the model-layer JSON hydration we get on the ORM # path). The UI reads `metadata.status` / `metadata.error_information` # as object fields, so failure rows looked like successes (#29674). - # Re-hydrate to dict here. + # Re-hydrate to dict here. Also drop the window-function `total_count` + # helper column so it does not leak into the serialised rows. for row in data: if isinstance(row, dict): + row.pop("total_count", None) md = row.get("metadata") if isinstance(md, str): try: diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index de5250fc5f3..b37bb18744d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -2,6 +2,7 @@ import asyncio import datetime import json import os +import re import sys from datetime import timezone @@ -60,31 +61,115 @@ def _filter_logs_by_date_range(logs, where): return filtered +def _reconstruct_ui_where_from_sql(sql_query, params): + """ + Rebuild the Prisma-style ``where`` dict the filter_fns below expect from the + raw SQL + params the endpoint emits. + + ``ui_view_spend_logs`` folds the total into the page query via + ``COUNT(*) OVER ()`` and no longer issues a separate ``count(where=...)`` + call, so the mock derives the active filter from the one query it sees + instead of from the (now absent) count call. + """ + where: dict = {} + clause = re.search(r"WHERE (.*) ORDER BY", sql_query, re.DOTALL) + if clause is None: + return where + + def _iso(value): + return value.isoformat() if hasattr(value, "isoformat") else str(value) + + eq_cols = { + "team_id": "team_id", + '"user"': "user", + "api_key": "api_key", + "request_id": "request_id", + "model": "model", + "model_id": "model_id", + "model_group": "model_group", + "end_user": "end_user", + } + date_bounds: dict = {} + metadata_conds: list = [] + for cond in (c.strip() for c in clause.group(1).split(" AND ")): + gte = re.search(r'"startTime" >= \(\$(\d+)', cond) + lte = re.search(r'"startTime" <= \(\$(\d+)', cond) + alias = re.search(r"user_api_key_alias' LIKE \$(\d+)", cond) + code = re.search(r"error_code' = \$(\d+)", cond) + msg = re.search(r"error_message' LIKE \$(\d+)", cond) + status = re.fullmatch(r"status = \$(\d+)", cond) + if gte: + date_bounds["gte"] = _iso(params[int(gte.group(1)) - 1]) + elif lte: + date_bounds["lte"] = _iso(params[int(lte.group(1)) - 1]) + elif "OR team_id = ANY" in cond: + where["OR"] = where.get("OR", []) + [{"multi_team": True}] + elif "status = 'success'" in cond: + where["OR"] = where.get("OR", []) + [{"status": "success"}] + elif status: + where["status"] = {"equals": params[int(status.group(1)) - 1]} + elif alias: + metadata_conds.append( + { + "path": ["user_api_key_alias"], + "string_contains": str(params[int(alias.group(1)) - 1]).strip("%"), + } + ) + elif code: + metadata_conds.append( + { + "path": ["error_information", "error_code"], + "equals": params[int(code.group(1)) - 1], + } + ) + elif msg: + metadata_conds.append( + { + "path": ["error_information", "error_message"], + "string_contains": str(params[int(msg.group(1)) - 1]).strip("%"), + } + ) + else: + for sql_col, key in eq_cols.items(): + eq = re.fullmatch(rf"{re.escape(sql_col)} = \$(\d+)", cond) + if eq: + where[key] = params[int(eq.group(1)) - 1] + break + + if date_bounds: + where["startTime"] = date_bounds + if len(metadata_conds) == 1: + where["metadata"] = metadata_conds[0] + elif len(metadata_conds) > 1: + where["AND"] = [{"metadata": cond} for cond in metadata_conds] + return where + + def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None): """ Create a MockPrismaClient for /spend/logs/ui endpoint tests. Args: mock_spend_logs: List of mock spend log dicts. - filter_fn: Callable[[dict], list] - receives where_conditions from count(), - returns the filtered list of logs for that query. + filter_fn: Callable[[dict], list] - receives the reconstructed + where_conditions, returns the filtered list of logs. team_lookup_fn: Optional async callable for team RBAC (find_unique). If provided, adds litellm_teamtable to db. """ - filtered_holder = [] class MockDB: async def count(self, *args, **kwargs): - where = kwargs.get("where", {}) - filtered = filter_fn(where) - filtered_holder.clear() - filtered_holder.extend(filtered) - return len(filtered) + return len(filter_fn(kwargs.get("where", {}))) async def query_raw(self, sql_query, *params): + filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params)) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return filtered_holder[skip : skip + page_size] + total = len(filtered) + return [ + {**row, "total_count": total} + for row in filtered[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -608,7 +693,10 @@ async def test_ui_view_spend_logs_sort_by_and_sort_order( sorted_logs = _sort_logs(base_logs, order) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return sorted_logs[skip : skip + page_size] + return [ + {**row, "total_count": len(base_logs)} + for row in sorted_logs[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -748,7 +836,10 @@ async def test_ui_view_spend_logs_sort_by_request_duration_ms(client, monkeypatc ) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return sorted_logs[skip : skip + page_size] + return [ + {**row, "total_count": len(base_logs)} + for row in sorted_logs[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -846,7 +937,10 @@ async def test_ui_view_spend_logs_sort_by_model( ) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return sorted_logs[skip : skip + page_size] + return [ + {**row, "total_count": len(base_logs)} + for row in sorted_logs[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -957,7 +1051,7 @@ async def test_ui_view_spend_logs_sort_by_ttft_ms(client, monkeypatch): page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 return [ - {k: v for k, v in row.items() if k != "_ttft_ms"} + {**{k: v for k, v in row.items() if k != "_ttft_ms"}, "total_count": len(base_logs)} for row in sorted_logs[skip : skip + page_size] ] @@ -3668,7 +3762,7 @@ async def test_ui_view_spend_logs_rehydrates_metadata_jsonb_text(client, monkeyp return 1 async def mock_query_raw(sql_query, *params): - return [raw_row] + return [{**raw_row, "total_count": 1}] class MockPrismaClient: def __init__(self): @@ -3754,7 +3848,7 @@ async def test_ui_view_spend_logs_metadata_invalid_json_falls_back_to_empty_dict return 1 async def mock_query_raw(sql_query, *params): - return [raw_row] + return [{**raw_row, "total_count": 1}] class MockPrismaClient: def __init__(self): diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index fbac71e6372..5b793adbbb4 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -222,3 +222,154 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): "/spend/logs/ui must wrap both `startTime` bounds with " f"`AT TIME ZONE 'UTC'`. SQL was:\n{sql}" ) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_folds_count_into_window_function(monkeypatch): + """ + /spend/logs/ui must not issue a separate `COUNT(*)` round trip to compute + the total. On sharded engines like YugabyteDB a standalone `COUNT(*)` is a + distributed RPC that contacts every tablet and times out regardless of row + count, so the logs tab 500s (LIT-4027). The total is folded into the page + query via `COUNT(*) OVER ()` and read off the returned rows instead. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + ui_view_spend_logs, + ) + + rows = [ + {"request_id": "req-1", "metadata": "{}", "session_id": None, "total_count": 137}, + {"request_id": "req-2", "metadata": "{}", "session_id": None, "total_count": 137}, + ] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=rows) + mock_prisma.db.litellm_spendlogs = MagicMock() + mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + ) + + mock_prisma.db.litellm_spendlogs.count.assert_not_called() + + sql = mock_prisma.db.query_raw.call_args[0][0] + assert "COUNT(*) OVER ()" in sql, ( + "the page query must carry a window-function count so a separate " + f"distributed COUNT(*) is avoided. SQL was:\n{sql}" + ) + + assert response["total"] == 137 + assert response["total_pages"] == (137 + 50 - 1) // 50 + + for row in response["data"]: + assert "total_count" not in row, ( + "the window-function helper column must be stripped before " + "serialising rows" + ) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_empty_page_reports_zero_total(monkeypatch): + """ + When a page matches no rows the window-function count row is absent, so the + total must fall back to zero without issuing a separate `COUNT(*)`. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + ui_view_spend_logs, + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + mock_prisma.db.litellm_spendlogs = MagicMock() + mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + ) + + mock_prisma.db.litellm_spendlogs.count.assert_not_called() + assert response["total"] == 0 + assert response["total_pages"] == 0 + assert response["data"] == [] + + +@pytest.mark.asyncio +async def test_spend_logs_ui_out_of_range_page_falls_back_to_count(monkeypatch): + """ + An out-of-range page (offset past the last matching row) returns no rows, so + the window-function count is unavailable. The total must not collapse to zero + there; it falls back to a direct count so total/total_pages stay accurate. + This fallback only fires off the hot path (page > 1 with an empty result), so + the YugabyteDB timeout the fix removes from page 1 stays removed. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + ui_view_spend_logs, + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + mock_prisma.db.litellm_spendlogs = MagicMock() + mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=7) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=99, + page_size=2, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + ) + + mock_prisma.db.litellm_spendlogs.count.assert_called_once() + assert response["total"] == 7 + assert response["total_pages"] == (7 + 2 - 1) // 2 From bd5e0464643666e20a5bc2a2be5964f2e0430942 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 23:03:09 +0300 Subject: [PATCH 7/9] fix(bedrock): surface web identity token aud/iss on InvalidIdentityToken (#31412) When STS rejects a web identity token with InvalidIdentityToken (the "Incorrect token audience" case), litellm propagated the raw botocore error, which never names the aud LiteLLM actually sent. Diagnosing an audience mismatch then required enabling LITELLM_LOG=DEBUG on the prod instance, which degrades performance. _auth_with_web_identity_token now catches InvalidIdentityTokenException, decodes the public aud/iss claims of the resolved JWT without verifying its signature (no secret is read), and raises an AwsAuthError that preserves the STS reason and names the token audience and issuer, so the mismatch is visible from the error alone. Resolves LIT-4026 --- litellm/llms/bedrock/base_aws_llm.py | 43 ++++++++++- .../test_web_identity_session_policy.py | 73 +++++++++++++++++++ 2 files changed, 114 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index c31462a735b..b71f37023e8 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1,3 +1,4 @@ +import base64 import hashlib import json import os @@ -19,7 +20,7 @@ from typing import ( ) import httpx -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from litellm._logging import verbose_logger from litellm.caching.caching import DualCache @@ -56,6 +57,11 @@ class Boto3CredentialsInfo(BaseModel): aws_bedrock_runtime_endpoint: Optional[str] +class _WebIdentityTokenClaims(BaseModel): + aud: Optional[Union[str, list[str]]] = None + iss: Optional[str] = None + + class AwsAuthError(Exception): def __init__(self, status_code, message): self.status_code = status_code @@ -817,6 +823,25 @@ class BaseAWSLLM: return False + @staticmethod + def _unverified_web_identity_audience(oidc_token: str) -> Optional[str]: + """Return the public ``aud``/``iss`` claims of a web identity JWT + without verifying its signature, so a rejected-token error can name + the audience LiteLLM actually sent. The signature is never read, so no + secret is exposed.""" + segments = oidc_token.split(".") + if len(segments) != 3: + return None + payload = segments[1] + try: + decoded = base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4)) + claims = _WebIdentityTokenClaims.model_validate_json(decoded) + except (ValueError, ValidationError): + return None + if claims.aud is None and claims.iss is None: + return None + return f"aud={claims.aud!r}, iss={claims.iss!r}" + @tracer.wrap() def _auth_with_web_identity_token( self, @@ -925,7 +950,21 @@ class BaseAWSLLM: if aws_external_id is not None: assume_role_params["ExternalId"] = aws_external_id - sts_response = sts_client.assume_role_with_web_identity(**assume_role_params) + try: + sts_response = sts_client.assume_role_with_web_identity( + **assume_role_params + ) + except sts_client.exceptions.InvalidIdentityTokenException as e: + audience = ( + self._unverified_web_identity_audience(oidc_token) + if isinstance(oidc_token, str) + else None + ) + detail = f" Token {audience}" if audience else "" + raise AwsAuthError( + status_code=401, + message=f"AWS STS rejected the web identity token: {e}.{detail}", + ) from e iam_creds_dict = { "aws_access_key_id": sts_response["Credentials"]["AccessKeyId"], diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index 2cf1fa16e91..7e9c8a273ae 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -29,6 +29,7 @@ claude_platform statement are present and cover every documented action. """ +import base64 import json from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch @@ -157,6 +158,78 @@ class TestClaudePlatformActionsCovered: ) +def _make_jwt(payload: dict) -> str: + def _segment(data: dict) -> str: + return base64.urlsafe_b64encode(json.dumps(data).encode()).rstrip(b"=").decode() + + return f"{_segment({'alg': 'RS256', 'typ': 'JWT'})}.{_segment(payload)}.signature" + + +class TestInvalidIdentityTokenSurfacesAudience: + """LIT-4026: when STS rejects the web identity token with + ``InvalidIdentityToken`` (the "Incorrect token audience" case), the raised + error must name the ``aud``/``iss`` the token actually carries so an + operator can diagnose the mismatch without enabling LITELLM_LOG=DEBUG on a + prod instance.""" + + _AUD = "https://guidepoint.litellm-prod.ai" + _ISS = "https://accounts.google.com" + _STS_MESSAGE = ( + "An error occurred (InvalidIdentityToken) when calling the " + "AssumeRoleWithWebIdentity operation: Incorrect token audience" + ) + + def _raise_invalid_identity_token(self) -> Exception: + from litellm.llms.bedrock.base_aws_llm import AwsAuthError, BaseAWSLLM + + token = _make_jwt({"aud": self._AUD, "iss": self._ISS, "sub": "svc-account"}) + + mock_sts = MagicMock() + + class _InvalidIdentityTokenException(Exception): + pass + + mock_sts.exceptions.InvalidIdentityTokenException = ( + _InvalidIdentityTokenException + ) + mock_sts.assume_role_with_web_identity.side_effect = ( + _InvalidIdentityTokenException(self._STS_MESSAGE) + ) + + with ( + patch("boto3.client", return_value=mock_sts), + patch( + "litellm.llms.bedrock.base_aws_llm.get_secret", + return_value=token, + ), + pytest.raises(AwsAuthError) as exc_info, + ): + BaseAWSLLM()._auth_with_web_identity_token( + aws_web_identity_token="oidc/google/" + self._AUD, + aws_role_name="arn:aws:iam::123456789012:role/litellm-bedrock-role", + aws_session_name="test-session", + aws_region_name="us-east-1", + aws_sts_endpoint=None, + ) + return exc_info.value + + def test_error_names_token_audience(self): + err = self._raise_invalid_identity_token() + assert self._AUD in str(err) + + def test_error_names_token_issuer(self): + err = self._raise_invalid_identity_token() + assert self._ISS in str(err) + + def test_error_preserves_original_sts_reason(self): + err = self._raise_invalid_identity_token() + assert "Incorrect token audience" in str(err) + + def test_error_is_401(self): + err = self._raise_invalid_identity_token() + assert err.status_code == 401 + + class TestPolicyTransportConditions: def test_bedrock_statement_keeps_secure_transport_condition(self): policy = _captured_policy() From f2fa23b0ec9f59987cc73e2783c09360dcbff578 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 23:07:50 +0300 Subject: [PATCH 8/9] fix(guardrails): instrument during-call and post-call guardrail latency (#31414) litellm_guardrail_latency_seconds was only emitted for pre-call guardrails. during_call_hook and post_call_success_hook ran guardrails without recording any latency, so during-call and post-call guardrail time was invisible in the metric and leaked into litellm_overhead_latency_metric, making the documented "subtract guardrail latency from overhead" workaround under-report total guardrail time. Extract the find-the-PrometheusLogger-and-record step into _emit_guardrail_metrics and add _run_guardrail_with_metrics, a single wrapper that times a guardrail coroutine, classifies its outcome (success / intervened / error), enriches any raised HTTPException, and records the latency under the given hook_type. Route the pre-call emit, during_call_hook, and post_call_success_hook through it so every guardrail phase contributes to the metric the same way. Resolves LIT-3999 --- litellm/proxy/utils.py | 115 +++++++++++------- .../proxy_logging/test_guardrail_pipeline.py | 104 ++++++++++++++-- 2 files changed, 171 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 781f0e9f301..f6d9d8975be 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -100,6 +100,7 @@ from litellm.proxy.hooks.sensitive_data_routing import ( _PROXY_SensitiveDataRoutingHandler, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.litellm_logging import Logging @@ -1124,7 +1125,6 @@ class ProxyLogging: Returns: Updated data dictionary if guardrail passes, None if guardrail should be skipped """ - from litellm.integrations.prometheus import PrometheusLogger from litellm.types.guardrails import GuardrailEventHooks # Determine the event type based on call type @@ -1197,17 +1197,13 @@ class ProxyLogging: or "unknown" ) - # Find PrometheusLogger in callbacks and record metrics - for prom_callback in litellm.callbacks: - if isinstance(prom_callback, PrometheusLogger): - prom_callback._record_guardrail_metrics( - guardrail_name=metrics_guardrail_name, - latency_seconds=latency_seconds, - status=status, - error_type=error_type, - hook_type="pre_call", - ) - break + self._emit_guardrail_metrics( + guardrail_name=metrics_guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type="pre_call", + ) return data @@ -1625,19 +1621,58 @@ class ProxyLogging: return data @staticmethod - async def _run_guardrail_task_with_enrichment( - callback: Any, coro: Awaitable[Any] + def _emit_guardrail_metrics( + guardrail_name: str, + latency_seconds: float, + status: str, + error_type: Optional[str], + hook_type: str, + ) -> None: + for prom_callback in litellm.callbacks: + if isinstance(prom_callback, PrometheusLogger): + prom_callback._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + break + + @staticmethod + async def _run_guardrail_with_metrics( + callback: Any, coro: Awaitable[Any], hook_type: str ) -> Any: """ - Await `coro`; if it raises an HTTPException with dict detail, - enrich the detail with the originating callback's `guardrail_name` - and `guardrail_mode` before re-raising. + Await `coro`, recording its latency and status to the + `litellm_guardrail_latency_seconds` metric under `hook_type`, and + enriching any raised HTTPException with the originating callback's + `guardrail_name`/`guardrail_mode` before re-raising. """ + guardrail_name = ( + getattr(callback, "guardrail_name", None) or type(callback).__name__ + ) + start_time = time.perf_counter() + status = "success" + error_type: Optional[str] = None try: return await coro + except SensitiveDataRouteException: + status = "intervened" + raise except Exception as e: + status = "error" + error_type = type(e).__name__ _enrich_http_exception_with_guardrail_context(e, callback) raise + finally: + ProxyLogging._emit_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=time.perf_counter() - start_time, + status=status, + error_type=error_type, + hook_type=hook_type, + ) @staticmethod async def _wrap_streaming_iterator_with_enrichment( @@ -1853,22 +1888,24 @@ class ProxyLogging: and not getattr(callback, "use_native_during_call_hook", False) ): data["guardrail_to_apply"] = callback - guardrail_task = self._run_guardrail_task_with_enrichment( + guardrail_task = self._run_guardrail_with_metrics( callback, unified_guardrail.async_moderation_hook( user_api_key_dict=user_api_key_dict, data=data, call_type=call_type, ), + "during_call", ) else: - guardrail_task = self._run_guardrail_task_with_enrichment( + guardrail_task = self._run_guardrail_with_metrics( callback, callback.async_moderation_hook( data=data, user_api_key_dict=user_api_key_auth_dict, # type: ignore call_type=call_type, # type: ignore ), + "during_call", ) guardrail_tasks.append(guardrail_task) @@ -2394,29 +2431,25 @@ class ProxyLogging: if "apply_guardrail" in type(callback).__dict__: data["guardrail_to_apply"] = callback - try: - guardrail_response = ( - await unified_guardrail.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ) - ) - except Exception as e: - _enrich_http_exception_with_guardrail_context(e, callback) - raise + guardrail_response = await self._run_guardrail_with_metrics( + callback, + unified_guardrail.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ), + "post_call", + ) else: - try: - guardrail_response = ( - await callback.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ) - ) - except Exception as e: - _enrich_http_exception_with_guardrail_context(e, callback) - raise + guardrail_response = await self._run_guardrail_with_metrics( + callback, + callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ), + "post_call", + ) if guardrail_response is not None: response = guardrail_response diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 1ff9fbf8d83..b9f209e5251 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -4,7 +4,7 @@ Covers ``_should_use_guardrail_load_balancing``, ``_execute_guardrail_hook``, ``_execute_guardrail_with_load_balancing``, ``_process_guardrail_callback``, ``_process_prompt_template``, ``_process_guardrail_metadata``, ``_maybe_execute_pipelines``, ``_handle_pipeline_result``, -``_run_guardrail_task_with_enrichment``. +``_run_guardrail_with_metrics``, ``_emit_guardrail_metrics``. """ from __future__ import annotations @@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, ) +from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks @@ -425,23 +426,42 @@ def test_handle_pipeline_result_unknown_action_returns_data(): # --------------------------------------------------------------------------- -# _run_guardrail_task_with_enrichment +# _run_guardrail_with_metrics # --------------------------------------------------------------------------- +def _prometheus_callback() -> MagicMock: + """Stand-in PrometheusLogger that records ``_record_guardrail_metrics`` calls. + + ``MagicMock(spec=PrometheusLogger)`` passes the ``isinstance`` check inside + ``_emit_guardrail_metrics`` while letting us capture the recorded labels. + """ + return MagicMock(spec=PrometheusLogger) + + @pytest.mark.asyncio -async def test_run_guardrail_task_with_enrichment_passes_result(): +async def test_run_guardrail_with_metrics_passes_result_and_records_success(monkeypatch): async def task(): return {"a": 1, "b": 2, "c": 3} - out = await ProxyLogging._run_guardrail_task_with_enrichment( - callback=MagicMock(guardrail_name="g"), coro=task() + prom = _prometheus_callback() + monkeypatch.setattr(litellm, "callbacks", [prom]) + + out = await ProxyLogging._run_guardrail_with_metrics( + callback=MagicMock(guardrail_name="g"), coro=task(), hook_type="during_call" ) + assert out == {"a": 1, "b": 2, "c": 3} + recorded = prom._record_guardrail_metrics.call_args.kwargs + assert recorded["guardrail_name"] == "g" + assert recorded["status"] == "success" + assert recorded["error_type"] is None + assert recorded["hook_type"] == "during_call" + assert recorded["latency_seconds"] >= 0 @pytest.mark.asyncio -async def test_run_guardrail_task_with_enrichment_enriches_http_exception_raises(): +async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch): detail = {"error": "blocked"} async def task(): @@ -450,9 +470,79 @@ async def test_run_guardrail_task_with_enrichment_enriches_http_exception_raises cb = MagicMock() cb.guardrail_name = "presidio" cb.event_hook = "pre_call" + prom = _prometheus_callback() + monkeypatch.setattr(litellm, "callbacks", [prom]) + with pytest.raises(HTTPException): - await ProxyLogging._run_guardrail_task_with_enrichment(callback=cb, coro=task()) + await ProxyLogging._run_guardrail_with_metrics( + callback=cb, coro=task(), hook_type="post_call" + ) + assert detail["guardrail_name"] == "presidio" + recorded = prom._record_guardrail_metrics.call_args.kwargs + assert recorded["status"] == "error" + assert recorded["error_type"] == "HTTPException" + assert recorded["hook_type"] == "post_call" + + +# --------------------------------------------------------------------------- +# during_call / post_call phases emit the latency metric (LIT-3999 regression) +# --------------------------------------------------------------------------- + + +def _moderation_guardrail() -> MagicMock: + cb = MagicMock(spec=CustomGuardrail) + cb.__class__ = CustomGuardrail + cb.guardrail_name = "g" + cb.event_hook = GuardrailEventHooks.during_call + cb.use_native_during_call_hook = False + cb.should_run_guardrail = MagicMock(return_value=True) + cb.async_moderation_hook = AsyncMock(return_value=None) + cb.async_post_call_success_hook = AsyncMock(return_value=None) + return cb + + +@pytest.mark.asyncio +async def test_during_call_hook_records_latency_metric( + proxy_logging, make_user_api_key_auth, monkeypatch +): + cb = _moderation_guardrail() + prom = _prometheus_callback() + monkeypatch.setattr(litellm, "callbacks", [prom, cb]) + + await proxy_logging.during_call_hook( + data={"model": "m"}, + user_api_key_dict=make_user_api_key_auth(), + call_type="completion", + ) + + cb.async_moderation_hook.assert_awaited_once() + recorded = prom._record_guardrail_metrics.call_args.kwargs + assert recorded["hook_type"] == "during_call" + assert recorded["guardrail_name"] == "g" + assert recorded["status"] == "success" + + +@pytest.mark.asyncio +async def test_post_call_success_hook_records_latency_metric( + proxy_logging, make_user_api_key_auth, monkeypatch +): + cb = _moderation_guardrail() + prom = _prometheus_callback() + monkeypatch.setattr(litellm, "callbacks", [prom, cb]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + await proxy_logging.post_call_success_hook( + data={"model": "m"}, + response=litellm.ModelResponse(), + user_api_key_dict=make_user_api_key_auth(), + ) + + cb.async_post_call_success_hook.assert_awaited_once() + recorded = prom._record_guardrail_metrics.call_args.kwargs + assert recorded["hook_type"] == "post_call" + assert recorded["guardrail_name"] == "g" + assert recorded["status"] == "success" # --------------------------------------------------------------------------- From ce658367a4171c078db03b71ef4183ad2de48709 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 23:36:50 +0300 Subject: [PATCH 9/9] fix(auth): cache auth-path team object under canonical team_id key (#31418) The auth builder cached the team object under the raw `valid_token.team_id`, while `get_team_object`, `_cache_team_object`, and `_update_team_cache` all read and write under `team_id:{id}`. The raw-key write was therefore never served back, and on a non-team (personal) key, whose team_id is None, the original unguarded version passed a None key straight to the cache layer; the in-memory cache tolerates None keys but Redis rejects them with a NoneType key error, so with `enable_redis_auth_cache: true` the team object never reached the L2 cache and every request fell back to Postgres. Write under the canonical `team_id:{id}` key, keeping the existing guard that skips the write when team_id is None. Add a regression test that drives the real auth builder for a team-scoped key against an in-memory cache and asserts the team object is served back under `team_id:{id}` and never under the raw team_id or a None key. Resolves LIT-4000 --- litellm/proxy/auth/user_api_key_auth.py | 9 +- .../proxy/auth/test_user_api_key_auth.py | 95 +++++++++++++++++++ 2 files changed, 101 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2dcc7bdbd19..184d91fb103 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1991,13 +1991,16 @@ async def _user_api_key_auth_builder( else: valid_token.team_object_permission = None - # Only cache when the key is a real team_id (non-team keys must not use key=None). + # Cache under the canonical "team_id:{id}" key so get_team_object and + # _update_team_cache serve this write from the L2 cache. The guard keeps a + # non-team (personal) key, whose team_id is None, from reaching the cache + # layer, which Redis rejects with a NoneType key error. if valid_token.team_id is not None and _team_obj is not None: await user_api_key_cache.async_set_cache( - key=valid_token.team_id, + key=f"team_id:{valid_token.team_id}", value=_team_obj, model_type=LiteLLM_TeamTableCachedObj, - ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py + ) # Fetch project object if key belongs to a project _project_obj = None diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index f80a8b28cc3..4141b1f20c3 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3989,3 +3989,98 @@ async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypa assert call_kwargs["valid_token_dict"]["is_session_token"] is True assert call_kwargs["valid_token_dict"]["user_role"] == LitellmUserRoles.INTERNAL_USER assert result.is_session_token is True + + +@pytest.mark.asyncio +async def test_auth_path_caches_team_object_under_canonical_team_id_key(): + """Regression for LIT-4000: the auth builder must cache the team object under + the canonical ``team_id:{id}`` key that ``get_team_object`` and + ``_update_team_cache`` read, never under the raw ``team_id`` (and never under + a ``None`` key, which Redis rejects with a NoneType key error). A raw or None + key is silently dropped by Redis / never served back, so every request + re-hits Postgres for the team object instead of the L2 cache. + + Drives the real builder for a team-scoped key against a real in-memory + ``UserApiKeyCache`` and reads the team object back. Mutating the cache key at + the write site to the raw ``valid_token.team_id`` (or ``None``) makes the + canonical-key read miss and fails this test. + """ + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import hash_token + + team_id = "team-lit-4000" + api_key = "sk-lit-4000-team-key" + cache = UserApiKeyCache() + + team_token = UserAPIKeyAuth(token=hash_token(api_key), team_id=team_id) + team_obj = LiteLLM_TeamTableCachedObj(team_id=team_id) + + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + attrs = { + "prisma_client": MagicMock(), + "user_api_key_cache": cache, + "proxy_logging_obj": proxy_logging_obj, + "master_key": "sk-test-master", + "general_settings": {"allow_requests_on_db_unavailable": False}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + with ( + patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + AsyncMock(return_value=team_token), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + AsyncMock(return_value=team_obj), + ), + patch( + "litellm.proxy.auth.user_api_key_auth._enforce_key_and_fallback_model_access", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj", + new_callable=AsyncMock, + return_value=team_token, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + ): + await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + served = cache.get_cache( + key=f"team_id:{team_id}", model_type=LiteLLM_TeamTableCachedObj + ) + assert served is not None and served.team_id == team_id + assert cache.get_cache(key=team_id) is None + assert cache.get_cache(key=None) is None