From 9d4bab3b700c1170b7451bedfbda67cf830f58f3 Mon Sep 17 00:00:00 2001 From: zoroyihan7 Date: Tue, 8 Sep 2026 10:46:16 +0000 Subject: [PATCH 01/61] fix(responses): emit typed streaming failure events --- litellm/exceptions.py | 3 +- .../common_utils/responses_stream_errors.py | 119 ++++++++++++++++++ litellm/proxy/proxy_server.py | 21 +++- .../proxy/response_api_endpoints/endpoints.py | 6 +- litellm/responses/streaming_iterator.py | 5 + .../proxy_server/test_streaming_helpers.py | 85 ++++++++++++- .../response_api_endpoints/test_endpoints.py | 95 +++++++++++++- 7 files changed, 328 insertions(+), 6 deletions(-) create mode 100644 litellm/proxy/common_utils/responses_stream_errors.py diff --git a/litellm/exceptions.py b/litellm/exceptions.py index f9215267bf3..318b227ffa8 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -789,6 +789,7 @@ class APIError(openai.APIError): litellm_debug_info: str | None = None, max_retries: int | None = None, num_retries: int | None = None, + body: object | None = None, ): self.status_code = status_code self.message = f"litellm.APIError: {message}" @@ -799,7 +800,7 @@ class APIError(openai.APIError): self.num_retries = num_retries if request is None: request = httpx.Request(method="POST", url="https://api.openai.com/v1") - super().__init__(self.message, request=request, body=None) + super().__init__(self.message, request=request, body=body) def __str__(self): _message = self.message diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py new file mode 100644 index 00000000000..a3bee06e912 --- /dev/null +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -0,0 +1,119 @@ +import time +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel, ConfigDict + +from litellm._logging import redact_internal_details_from_client_message +from litellm._uuid import uuid +from litellm.exceptions import MidStreamFallbackError +from litellm.types.llms.openai import ResponseFailedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents + + +class _ResponseIdentity(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + id: str | None = None + model: str | None = None + created_at: int | None = None + + +class _StreamEvent(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + type: str | None = None + sequence_number: int | None = None + response: _ResponseIdentity | None = None + + +class _FailureDetails(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + message: str | None = None + code: str | int | None = None + type: str | None = None + status_code: int | None = None + + +def _original_failure(exception: Exception) -> Exception: + if isinstance(exception, MidStreamFallbackError) and exception.original_exception is not None: + return _original_failure(exception.original_exception) + return exception + + +def _response_error_code(details: _FailureDetails) -> str: + for value in (details.code, details.type): + if value == "insufficient_quota": + return "insufficient_quota" + if value in (429, "429") or isinstance(value, str) and value.startswith("rate_limit"): + return "rate_limit_exceeded" + if isinstance(details.code, str) and details.code and not details.code.isdecimal(): + return details.code + if details.status_code == 429: + return "rate_limit_exceeded" + return "server_error" + + +class ResponsesStreamErrorState: + def __init__(self) -> None: + self.response_id: str | None = None + self.model: str | None = None + self.created_at: int | None = None + self.sequence_number = -1 + self.terminal_emitted = False + + @staticmethod + def observe_chunk(chunk: object) -> _StreamEvent | None: + if not isinstance(chunk, (BaseModel, Mapping)): + return None + return _StreamEvent.model_validate(chunk) + + def mark_emitted(self, event: _StreamEvent | None) -> None: + if event is None: + return + if event.sequence_number is not None: + self.sequence_number = max(self.sequence_number, event.sequence_number) + if event.response is not None: + self.response_id = event.response.id or self.response_id + self.model = event.response.model or self.model + if event.response.created_at is not None: + self.created_at = event.response.created_at + if event.type in ("response.completed", "response.failed", "response.incomplete"): + self.terminal_emitted = True + + def format_failure(self, exception: Exception) -> str | None: + if self.terminal_emitted: + return None + original: Final = _original_failure(exception) + details: Final = _FailureDetails.model_validate(original) + response: Final = ResponsesAPIResponse.model_validate( + MappingProxyType( + { + "id": self.response_id or f"resp_{uuid.uuid4().hex}", + "object": "response", + "created_at": self.created_at if self.created_at is not None else int(time.time()), + "model": self.model, + "status": "failed", + "output": (), + "error": MappingProxyType( + { + "code": _response_error_code(details), + "message": redact_internal_details_from_client_message(details.message or str(original)), + } + ), + } + ) + ) + event: Final = ResponseFailedEvent.model_validate( + MappingProxyType( + { + "type": ResponsesAPIStreamEvents.RESPONSE_FAILED, + "response": response, + "sequence_number": self.sequence_number + 1, + } + ) + ) + payload: Final = event.model_dump_json(exclude_none=True) + self.terminal_emitted = True + return f"event: response.failed\ndata: {payload}\n\n" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 32b6b841af7..9f1d69acefa 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -385,6 +385,7 @@ from litellm.proxy.common_utils.periodic_reload_schedule import ( ) from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob +from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState from litellm.proxy.common_utils.scheduled_job_stagger import ( apply_scheduled_job_stagger, attach_job_timing_logger, @@ -8723,10 +8724,13 @@ async def async_data_generator( user_api_key_dict: UserAPIKeyAuth, request_data: dict, request: Request | None = None, + *, + responses_stream_errors: bool = False, ): verbose_proxy_logger.debug("inside generator") stream_completed = False client_disconnected = False + error_state: Final = ResponsesStreamErrorState() if responses_stream_errors else None try: error_message: str | None = None requested_model_from_client: Final = _get_client_requested_model_for_streaming(request_data=request_data) @@ -8837,6 +8841,7 @@ async def async_data_generator( fallback_metadata_event_sent = True continue + responses_event: Final = error_state.observe_chunk(chunk) if error_state is not None else None raw_passthrough = False if isinstance(chunk, BaseModel): chunk = _serialize_streaming_chunk(chunk) @@ -8871,8 +8876,13 @@ async def async_data_generator( if not raw_passthrough: try: - yield _format_streaming_sse_chunk(chunk=chunk) + formatted_chunk: Final = _format_streaming_sse_chunk(chunk=chunk) + if error_state is not None: + error_state.mark_emitted(responses_event) + yield formatted_chunk except Exception as e: + if error_state is not None: + raise yield f"data: {e}\n\n" if pending_fallback_event: @@ -8922,6 +8932,12 @@ async def async_data_generator( e, ) + if error_state is not None: + stream_completed = True + error_frame: Final = error_state.format_failure(e) + if error_frame is not None: + yield error_frame + return if isinstance(e, HTTPException): raise e elif isinstance(e, StreamingCallbackError): @@ -8958,12 +8974,15 @@ def select_data_generator( user_api_key_dict: UserAPIKeyAuth, request_data: dict, request: Request | None = None, + *, + responses_stream_errors: bool = False, ): return async_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, request=request, + responses_stream_errors=responses_stream_errors, ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 5907ffc64eb..3202fabe74e 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -3,6 +3,7 @@ import json import time from collections.abc import AsyncIterator, Awaitable, Mapping from enum import Enum +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, cast, get_args from uuid import uuid4 @@ -243,6 +244,7 @@ async def responses_api( version, ) + native_data_generator: Final = partial(select_data_generator, responses_stream_errors=True) data = await _read_request_body(request=request) # Check if polling via cache should be used for this request @@ -329,7 +331,7 @@ async def responses_api( llm_router=llm_router, proxy_config=proxy_config, proxy_logging_obj=proxy_logging_obj, - select_data_generator=select_data_generator, + select_data_generator=native_data_generator, user_model=user_model, user_temperature=user_temperature, user_request_timeout=user_request_timeout, @@ -355,7 +357,7 @@ async def responses_api( llm_router=llm_router, general_settings=general_settings, proxy_config=proxy_config, - select_data_generator=select_data_generator, + select_data_generator=native_data_generator, model=None, user_model=user_model, user_temperature=user_temperature, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 9f9016c5a7f..1134f6b07e3 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -579,6 +579,11 @@ class BaseResponsesAPIStreamingIterator: message=error_message, llm_provider=self.custom_llm_provider or "", model=self.model or "", + body={ # mutable-ok: OpenAI APIError reads code/type only from a dict body + "code": error_code, + "type": error_type, + "message": error_message, + }, ) if 400 <= status_code < 500 and status_code != 429: raise mapped_exception diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 87e10ce7e8d..6f3a47a4b79 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -17,15 +17,18 @@ from __future__ import annotations import asyncio import json +from collections.abc import AsyncIterator +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import Response from fastapi.responses import StreamingResponse +from pydantic import BaseModel import litellm -from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY import litellm.proxy.proxy_server as ps +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.proxy_server import ( _apply_streaming_chunk_hooks, @@ -42,6 +45,12 @@ from litellm.proxy.proxy_server import ( data_generator, select_data_generator, ) +from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponseCreatedEvent, + ResponseFailedEvent, + ResponsesAPIResponse, +) from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage from .conftest import normalize @@ -872,6 +881,80 @@ async def test_async_data_generator_mid_stream_exception_yields_error_payload( assert any(isinstance(item, str) and item.startswith('data: {"error":') for item in out) +@pytest.mark.asyncio +@pytest.mark.parametrize("terminal", ["completed", "serialization_failure", "failure_after_completed"]) +async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_terminal( + terminal: Literal["completed", "serialization_failure", "failure_after_completed"], +) -> None: + class ToolDelta(BaseModel): + type: Literal["response.function_call_arguments.delta"] + sequence_number: int + item_id: str + output_index: int + delta: str + + class UnserializableTerminal(BaseModel): + type: Literal["response.completed"] + sequence_number: int + response: ResponsesAPIResponse + invalid: object + + response: Final = ResponsesAPIResponse(id="resp_visible", created_at=1, model="gpt-6-astra", output=[]) + created: Final = ResponseCreatedEvent.model_validate( + {"type": "response.created", "sequence_number": 0, "response": response} + ) + completed: Final = ResponseCompletedEvent.model_validate( + {"type": "response.completed", "sequence_number": 2, "response": response} + ) + tool_delta: Final = ToolDelta( + type="response.function_call_arguments.delta", sequence_number=1, item_id="fc_stream_error", + output_index=0, delta='{"path":"partial', + ) + + async def upstream() -> AsyncIterator[BaseModel]: + yield created + yield tool_delta + yield ( + UnserializableTerminal(type="response.completed", sequence_number=2, response=response, invalid=object()) + if terminal == "serialization_failure" else completed + ) + if terminal == "failure_after_completed": + raise litellm.APIError(status_code=500, message="Stream close failed", llm_provider="openai", model="gpt-6-astra") + + frames: Final = [ + frame + async for frame in select_data_generator( + response=upstream(), + user_api_key_dict=_user_auth(), + request_data={}, + responses_stream_errors=True, + ) + ] + decoded: Final = tuple(frame.decode() if isinstance(frame, bytes) else frame for frame in frames) + event_frames: Final = tuple(frame for frame in decoded if frame != "data: [DONE]\n\n") + payloads: Final = tuple( + json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) + for frame in event_frames + ) + + assert payloads[0]["response"]["id"] == "resp_visible" + assert payloads[1] == tool_delta.model_dump() + assert len(payloads) == 3 + if terminal == "serialization_failure": + failure: Final = ResponseFailedEvent.model_validate(payloads[-1]) + assert event_frames[-1].startswith("event: response.failed\n") + assert failure.response.id == "resp_visible" + assert failure.response.status == "failed" + assert failure.response.error is not None + assert failure.response.error["code"] == "server_error" + assert "serialize" in failure.response.error["message"].lower() + assert payloads[-1]["sequence_number"] > payloads[1]["sequence_number"] + else: + assert payloads[-1]["type"] == "response.completed" + assert payloads[-1]["sequence_number"] == 2 + assert "error" not in payloads[-1] + + # --------------------------------------------------------------------------- # select_data_generator # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index d7010de6405..9e2f70a3d64 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -3,10 +3,12 @@ Test for response_api_endpoints/endpoints.py """ import unittest -from typing import Any +from typing import Any, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx from fastapi.testclient import TestClient from httpx import Response @@ -14,6 +16,97 @@ import litellm from litellm.proxy.proxy_server import app +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path,error_kind", + [ + ("/v1/responses", "rate_limit"), + ("/v1/responses", "numeric_rate_limit"), + ("/v1/responses", "server_error"), + ("/v1/responses", "response_failed"), + ("/cursor/chat/completions", "server_error"), + ("/v1/chat/completions", "server_error"), + ], +) +async def test_streaming_upstream_errors_keep_the_client_protocol( + monkeypatch: pytest.MonkeyPatch, + path: str, + error_kind: Literal["rate_limit", "numeric_rate_limit", "server_error", "response_failed"], +) -> None: + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + model: Final = "gpt-6-astra" + message: Final = "Upstream cannot complete this response" + code: Final = { + "rate_limit": "rate_limit_exceeded", "numeric_rate_limit": "429", + "server_error": "server_error", "response_failed": "server_error", + }[error_kind] + error: Final = {"message": message, "code": code, "type": None, "param": "input"} + response: Final = {"id": "resp_upstream", "object": "response", "created_at": 1, + "status": "in_progress", "model": model, "output": [], + "parallel_tool_calls": True, "tool_choice": "auto", "tools": []} + created: Final = {"type": "response.created", "sequence_number": 0, "response": response} + tool_added: Final = {"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, + "item": {"type": "function_call", "id": "fc_partial", "call_id": "call_partial", + "name": "read_file", "arguments": "", "status": "in_progress"}} + tool_delta: Final = {"type": "response.function_call_arguments.delta", "sequence_number": 2, + "item_id": "fc_partial", "output_index": 0, "delta": '{"path":"partial'} + failed: Final = ( + {"type": "response.failed", "sequence_number": 9, + "response": {**response, "status": "failed", "error": error}} + if error_kind == "response_failed" else {"type": "error", "error": error} + ) + chat: Final = {"id": "chatcmpl_partial", "object": "chat.completion.chunk", "created": 1, + "model": model, "choices": [{"index": 0, "delta": {"content": "partial"}, + "finish_reason": None}]} + is_chat: Final = path == "/v1/chat/completions" + upstream_events: Final = (chat, {"error": error}) if is_chat else (created, tool_added, tool_delta, failed) + wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in upstream_events) + upstream_url: Final = "https://streaming.example/v1" + router: Final = litellm.Router( + model_list=[{"model_name": model, "litellm_params": { + "model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"}}], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock as transport: + transport.post(upstream_url + ("/chat/completions" if is_chat else "/responses")).respond( + 200, content=wire, headers={"Content-Type": "text/event-stream"} + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.post( + path, json={"model": model, "stream": True, + **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"})}, + ) + frames: Final = tuple(frame for frame in result.text.split("\n\n") if "data: " in frame) + events: Final = tuple( + json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) + for frame in frames if "data: [DONE]" not in frame + ) + + assert result.status_code == 200, result.text + assert message in result.text + if path == "/v1/responses": + assert frames[-1].startswith("event: response.failed\n"), result.text + assert [event["type"] for event in events] == [ + "response.created", "response.output_item.added", "response.function_call_arguments.delta", "response.failed" + ] + assert events[2]["delta"] == tool_delta["delta"] + assert events[-1]["sequence_number"] == events[-2]["sequence_number"] + 1 + assert events[-1]["response"]["id"] == events[0]["response"]["id"] + assert events[-1]["response"]["status"] == "failed" + assert events[-1]["response"]["error"]["code"] == ( + "rate_limit_exceeded" if error_kind in ("rate_limit", "numeric_rate_limit") else "server_error" + ) + else: + assert events[0]["object"] == "chat.completion.chunk", result.text + assert "response.failed" not in result.text + assert "error" in events[-1] + + class TestResponsesAPIEndpoints(unittest.TestCase): @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.llm_router") From 089703d10b2eebbec1a5b780494686451ab72161 Mon Sep 17 00:00:00 2001 From: zoroyihan7 Date: Tue, 8 Sep 2026 11:02:38 +0000 Subject: [PATCH 02/61] fix(responses): preserve failure metadata at streaming boundaries --- .../common_utils/responses_stream_errors.py | 45 +++++++++---- litellm/proxy/proxy_server.py | 9 +-- .../proxy_server/test_streaming_helpers.py | 63 ++++++++++++++++--- .../response_api_endpoints/test_endpoints.py | 28 ++++++--- 4 files changed, 114 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index a3bee06e912..356e948a2df 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -1,9 +1,10 @@ import time from collections.abc import Mapping +from http import HTTPStatus from types import MappingProxyType from typing import Final -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, field_validator from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid @@ -35,6 +36,16 @@ class _FailureDetails(BaseModel): type: str | None = None status_code: int | None = None + @field_validator("code", mode="before") + @classmethod + def normalize_code(cls, value: object) -> str | int | None: + return value if isinstance(value, (str, int)) and not isinstance(value, bool) else None + + @field_validator("type", mode="before") + @classmethod + def normalize_type(cls, value: object) -> str | None: + return value if isinstance(value, str) else None + def _original_failure(exception: Exception) -> Exception: if isinstance(exception, MidStreamFallbackError) and exception.original_exception is not None: @@ -50,9 +61,21 @@ def _response_error_code(details: _FailureDetails) -> str: return "rate_limit_exceeded" if isinstance(details.code, str) and details.code and not details.code.isdecimal(): return details.code - if details.status_code == 429: - return "rate_limit_exceeded" - return "server_error" + match details.status_code: + case HTTPStatus.UNAUTHORIZED: + return "authentication_error" + case HTTPStatus.FORBIDDEN: + return "permission_error" + case HTTPStatus.NOT_FOUND: + return "not_found_error" + case HTTPStatus.REQUEST_TIMEOUT: + return "request_timeout" + case HTTPStatus.TOO_MANY_REQUESTS: + return "rate_limit_exceeded" + case int(status) if HTTPStatus.BAD_REQUEST <= status < HTTPStatus.INTERNAL_SERVER_ERROR: + return "invalid_request_error" + case _: + return "server_error" class ResponsesStreamErrorState: @@ -62,16 +85,15 @@ class ResponsesStreamErrorState: self.created_at: int | None = None self.sequence_number = -1 self.terminal_emitted = False + self._pending_event: _StreamEvent | None = None - @staticmethod - def observe_chunk(chunk: object) -> _StreamEvent | None: - if not isinstance(chunk, (BaseModel, Mapping)): - return None - return _StreamEvent.model_validate(chunk) + def observe_chunk(self, chunk: object) -> None: + self._pending_event = _StreamEvent.model_validate(chunk) if isinstance(chunk, (BaseModel, Mapping)) else None - def mark_emitted(self, event: _StreamEvent | None) -> None: + def mark_emitted(self, frame: str | bytes) -> str | bytes: + event: Final = self._pending_event if event is None: - return + return frame if event.sequence_number is not None: self.sequence_number = max(self.sequence_number, event.sequence_number) if event.response is not None: @@ -81,6 +103,7 @@ class ResponsesStreamErrorState: self.created_at = event.response.created_at if event.type in ("response.completed", "response.failed", "response.incomplete"): self.terminal_emitted = True + return frame def format_failure(self, exception: Exception) -> str | None: if self.terminal_emitted: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9f1d69acefa..870cd78aa85 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8841,7 +8841,8 @@ async def async_data_generator( fallback_metadata_event_sent = True continue - responses_event: Final = error_state.observe_chunk(chunk) if error_state is not None else None + if error_state is not None: + error_state.observe_chunk(cast(object, chunk)) # cast-ok: the helper validates legacy untyped chunks raw_passthrough = False if isinstance(chunk, BaseModel): chunk = _serialize_streaming_chunk(chunk) @@ -8876,10 +8877,10 @@ async def async_data_generator( if not raw_passthrough: try: - formatted_chunk: Final = _format_streaming_sse_chunk(chunk=chunk) if error_state is not None: - error_state.mark_emitted(responses_event) - yield formatted_chunk + yield error_state.mark_emitted(_format_streaming_sse_chunk(chunk=chunk)) + else: + yield _format_streaming_sse_chunk(chunk=chunk) except Exception as e: if error_state is not None: raise diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 6f3a47a4b79..53e055882e8 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -21,9 +21,11 @@ from collections.abc import AsyncIterator from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock +import httpx import pytest -from fastapi import Response +from fastapi import HTTPException, Response from fastapi.responses import StreamingResponse +from openai import APIError as OpenAIAPIError from pydantic import BaseModel import litellm @@ -882,9 +884,44 @@ async def test_async_data_generator_mid_stream_exception_yields_error_payload( @pytest.mark.asyncio -@pytest.mark.parametrize("terminal", ["completed", "serialization_failure", "failure_after_completed"]) +@pytest.mark.parametrize( + "terminal,upstream_error,expected_code", + [ + ("completed", None, None), + ("serialization_failure", None, "server_error"), + ("failure_after_completed", None, None), + pytest.param( + "upstream_failure", + litellm.AuthenticationError( + message="Upstream rejected request", llm_provider="openai", model="gpt-6-astra" + ), + "authentication_error", id="authentication_error", + ), + pytest.param( + "upstream_failure", + OpenAIAPIError( + message="Upstream rejected request", + request=httpx.Request("POST", "https://streaming.example/v1/responses"), + body={"code": {"reason": "overloaded"}, "type": {"unexpected": "object"}}, + ), + "server_error", id="structured_provider_error_fields", + ), + *( + pytest.param( + "upstream_failure", HTTPException(status_code=status, detail="Upstream rejected request"), + code, id=f"http_{status}", + ) + for status, code in ( + (400, "invalid_request_error"), (403, "permission_error"), (404, "not_found_error"), + (408, "request_timeout"), (422, "invalid_request_error"), (500, "server_error"), (503, "server_error"), + ) + ), + ], +) async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_terminal( - terminal: Literal["completed", "serialization_failure", "failure_after_completed"], + terminal: Literal["completed", "serialization_failure", "failure_after_completed", "upstream_failure"], + upstream_error: HTTPException | OpenAIAPIError | None, + expected_code: str | None, ) -> None: class ToolDelta(BaseModel): type: Literal["response.function_call_arguments.delta"] @@ -910,16 +947,23 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina type="response.function_call_arguments.delta", sequence_number=1, item_id="fc_stream_error", output_index=0, delta='{"path":"partial', ) + original_status: Final = ( + upstream_error.status_code if isinstance(upstream_error, (HTTPException, litellm.AuthenticationError)) else None + ) async def upstream() -> AsyncIterator[BaseModel]: yield created yield tool_delta + if upstream_error is not None: + raise upstream_error yield ( UnserializableTerminal(type="response.completed", sequence_number=2, response=response, invalid=object()) if terminal == "serialization_failure" else completed ) if terminal == "failure_after_completed": - raise litellm.APIError(status_code=500, message="Stream close failed", llm_provider="openai", model="gpt-6-astra") + raise litellm.APIError( + status_code=500, message="Stream close failed", llm_provider="openai", model="gpt-6-astra" + ) frames: Final = [ frame @@ -940,14 +984,19 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina assert payloads[0]["response"]["id"] == "resp_visible" assert payloads[1] == tool_delta.model_dump() assert len(payloads) == 3 - if terminal == "serialization_failure": + if terminal in ("serialization_failure", "upstream_failure"): failure: Final = ResponseFailedEvent.model_validate(payloads[-1]) assert event_frames[-1].startswith("event: response.failed\n") assert failure.response.id == "resp_visible" assert failure.response.status == "failed" assert failure.response.error is not None - assert failure.response.error["code"] == "server_error" - assert "serialize" in failure.response.error["message"].lower() + assert failure.response.error["code"] == expected_code + if upstream_error is None: + assert "serialize" in failure.response.error["message"].lower() + else: + assert "Upstream rejected request" in failure.response.error["message"] + if isinstance(upstream_error, (HTTPException, litellm.AuthenticationError)): + assert upstream_error.status_code == original_status assert payloads[-1]["sequence_number"] > payloads[1]["sequence_number"] else: assert payloads[-1]["type"] == "response.completed" diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 9e2f70a3d64..7b79e1b8613 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -61,7 +61,9 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( "model": model, "choices": [{"index": 0, "delta": {"content": "partial"}, "finish_reason": None}]} is_chat: Final = path == "/v1/chat/completions" - upstream_events: Final = (chat, {"error": error}) if is_chat else (created, tool_added, tool_delta, failed) + partial: Final = path != "/v1/responses" or error_kind in ("numeric_rate_limit", "response_failed") + response_events: Final = (created, tool_added, tool_delta, failed) if partial else (failed,) + upstream_events: Final = (chat, {"error": error}) if is_chat else response_events wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in upstream_events) upstream_url: Final = "https://streaming.example/v1" router: Final = litellm.Router( @@ -78,8 +80,10 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( ) async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: result: Final = await client.post( - path, json={"model": model, "stream": True, - **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"})}, + path, json={ + "model": model, "stream": True, + **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"}), + }, ) frames: Final = tuple(frame for frame in result.text.split("\n\n") if "data: " in frame) events: Final = tuple( @@ -91,12 +95,18 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert message in result.text if path == "/v1/responses": assert frames[-1].startswith("event: response.failed\n"), result.text - assert [event["type"] for event in events] == [ - "response.created", "response.output_item.added", "response.function_call_arguments.delta", "response.failed" - ] - assert events[2]["delta"] == tool_delta["delta"] - assert events[-1]["sequence_number"] == events[-2]["sequence_number"] + 1 - assert events[-1]["response"]["id"] == events[0]["response"]["id"] + if partial: + assert [event["type"] for event in events] == [ + "response.created", "response.output_item.added", + "response.function_call_arguments.delta", "response.failed", + ] + assert events[2]["delta"] == tool_delta["delta"] + assert events[-1]["sequence_number"] == events[-2]["sequence_number"] + 1 + assert events[-1]["response"]["id"] == events[0]["response"]["id"] + else: + assert [event["type"] for event in events] == ["response.failed"] + assert events[0]["sequence_number"] == 0 + assert events[0]["response"]["id"].startswith("resp_") assert events[-1]["response"]["status"] == "failed" assert events[-1]["response"]["error"]["code"] == ( "rate_limit_exceeded" if error_kind in ("rate_limit", "numeric_rate_limit") else "server_error" From dc7895c1eacf55733953c3f802f46ea9103a76b4 Mon Sep 17 00:00:00 2001 From: zoroyihan7 Date: Tue, 8 Sep 2026 11:44:41 +0000 Subject: [PATCH 03/61] fix(responses): satisfy streaming regression checks --- litellm/proxy/common_utils/responses_stream_errors.py | 7 ++++--- .../test_response_polling_pre_call_checks.py | 1 - 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 356e948a2df..706b4c298d7 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -48,9 +48,10 @@ class _FailureDetails(BaseModel): def _original_failure(exception: Exception) -> Exception: - if isinstance(exception, MidStreamFallbackError) and exception.original_exception is not None: - return _original_failure(exception.original_exception) - return exception + current = exception # rebind-ok: the recursion gate requires iterative wrapper traversal + while isinstance(current, MidStreamFallbackError) and current.original_exception is not None: + current = current.original_exception + return current def _response_error_code(details: _FailureDetails) -> str: diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py index 459834d0fd2..3fabdcefe5a 100644 --- a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py +++ b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py @@ -130,7 +130,6 @@ class TestPollingEndpointPreCallGuard: "litellm.proxy.proxy_server.proxy_config": MagicMock(), "litellm.proxy.proxy_server.proxy_logging_obj": AsyncMock(), "litellm.proxy.proxy_server.redis_usage_cache": AsyncMock(), - "litellm.proxy.proxy_server.select_data_generator": None, "litellm.proxy.proxy_server.user_api_base": None, "litellm.proxy.proxy_server.user_max_tokens": None, "litellm.proxy.proxy_server.user_model": None, From 97dbd2dfbf3de0386efbe77057bbf75bc1aa1336 Mon Sep 17 00:00:00 2001 From: tusharjamunkar Date: Sat, 12 Sep 2026 22:16:19 +0530 Subject: [PATCH 04/61] fix(gemini): preserve candidates with finishReason and no content (#40477) --- litellm/litellm_core_utils/core_helpers.py | 1 + .../adapters/transformation.py | 2 + .../vertex_and_google_ai_studio_gemini.py | 30 ++-- .../transformation.py | 20 ++- ...test_vertex_and_google_ai_studio_gemini.py | 146 ++++++++++++++++++ 5 files changed, 184 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index aa7d6ca1699..66180b165f8 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -224,6 +224,7 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = { "IMAGE_PROHIBITED_CONTENT": "content_filter", "TOO_MANY_TOOL_CALLS": "stop", "MALFORMED_RESPONSE": "stop", + "NO_IMAGE": "content_filter", # Zhipu GLM "network_error": "stop", "sensitive": "content_filter", diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 8ff9f2e0679..f4cc569bcef 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1367,6 +1367,8 @@ class LiteLLMAnthropicMessagesAdapter: return "max_tokens" elif openai_finish_reason == "tool_calls": return "tool_use" + elif openai_finish_reason in ["content_filter", "refusal"]: + return "refusal" return "end_turn" @staticmethod diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d113b2b4f6b..01d1f063b1b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1340,6 +1340,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "IMAGE_PROHIBITED_CONTENT", "TOO_MANY_TOOL_CALLS", "MALFORMED_RESPONSE", + "NO_IMAGE", } ) @@ -2224,22 +2225,23 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): grounding_metadata: Final[list[dict]] = [] url_context_metadata: Final[list[dict]] = [] - image_response: list[ImageURLListItem] | None = None safety_ratings: Final[list] = [] citation_metadata: Final[list] = [] - chat_completion_message: Final[ChatCompletionResponseMessage] = {"role": "assistant"} - chat_completion_logprobs: ChoiceLogprobs | None = None - tools: list[ChatCompletionToolCallChunk] | None = [] - functions: ChatCompletionToolCallFunctionChunk | None = None - thinking_blocks: list[ChatCompletionThinkingBlock] | None = None - reasoning_content: str | None = None - thought_signatures: Sequence[str] | None = None - server_side_tool_invocations: list[dict[str, object]] | None = None for idx, candidate in enumerate(_candidates): - if "content" not in candidate: + if "content" not in candidate and "finishReason" not in candidate: continue + image_response: list[ImageURLListItem] | None = None + chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} + chat_completion_logprobs: ChoiceLogprobs | None = None + tools: list[ChatCompletionToolCallChunk] | None = [] + functions: ChatCompletionToolCallFunctionChunk | None = None + thinking_blocks: list[ChatCompletionThinkingBlock] | None = None + reasoning_content: str | None = None + thought_signatures: Sequence[str] | None = None + server_side_tool_invocations: list[dict[str, object]] | None = None + # Extract metadata using helper function ( candidate_grounding_metadata, @@ -2253,7 +2255,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): safety_ratings.extend(candidate_safety_ratings) citation_metadata.extend(candidate_citation_metadata) - if "parts" in candidate["content"]: + if "content" in candidate and candidate["content"] and "parts" in candidate["content"]: ( content, reasoning_content, @@ -2348,6 +2350,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tool_invocation_fields["server_side_tool_invocations"] = server_side_tool_invocations chat_completion_message["provider_specific_fields"] = tool_invocation_fields + if candidate.get("finishReason"): + finish_reason_fields = chat_completion_message.get("provider_specific_fields") or {} + finish_reason_fields["native_finish_reason"] = candidate.get("finishReason") + chat_completion_message["provider_specific_fields"] = finish_reason_fields + if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( chat_completion_message=chat_completion_message, @@ -2368,6 +2375,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): message=chat_completion_message, logprobs=chat_completion_logprobs, enhancements=None, + provider_specific_fields=chat_completion_message.get("provider_specific_fields"), ) model_response.choices.append(choice) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index fca5b0d11cf..13119085e46 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2272,13 +2272,27 @@ class LiteLLMCompletionResponsesConfig: if choices and len(choices) > 0: finish_reason = choices[0].finish_reason + status: Final[ResponsesAPIStatus] = ( + LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( + finish_reason + ) + ) + incomplete_details = getattr(chat_completion_response, "incomplete_details", None) + if incomplete_details is None and status == "incomplete": + from openai.types.responses.response import IncompleteDetails + + if finish_reason == "length": + incomplete_details = IncompleteDetails(reason="max_output_tokens") + elif finish_reason in ["content_filter", "refusal"]: + incomplete_details = IncompleteDetails(reason="content_filter") + responses_api_response: Final[ResponsesAPIResponse] = ResponsesAPIResponse( id=chat_completion_response.id, created_at=chat_completion_response.created, model=chat_completion_response.model, object="response", error=getattr(chat_completion_response, "error", None), - incomplete_details=getattr(chat_completion_response, "incomplete_details", None), + incomplete_details=incomplete_details, instructions=getattr(chat_completion_response, "instructions", None), metadata=getattr(chat_completion_response, "metadata", {}), output=LiteLLMCompletionResponsesConfig._transform_chat_completion_choices_to_responses_output( @@ -2296,9 +2310,7 @@ class LiteLLMCompletionResponsesConfig: max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None), previous_response_id=getattr(chat_completion_response, "previous_response_id", None), reasoning=None, - status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( - finish_reason - ), + status=status, text={}, truncation=getattr(chat_completion_response, "truncation", None), usage=LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 101f6e6fa5d..a048ad4f171 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5836,3 +5836,149 @@ def test_supported_reasoning_efforts_still_map(model): drop_params=False, ) assert "thinkingConfig" in result + + +def test_gemini_candidate_with_finish_reason_no_content_chat_completion(): + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + model_response = ModelResponse() + logging_obj = MagicMock() + raw_response = MagicMock() + raw_response.headers = {} + + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=model_response, + model="gemini-2.5-flash-image", + logging_obj=logging_obj, + raw_response=raw_response, + ) + assert len(resp.choices) == 1 + assert resp.choices[0].finish_reason == "content_filter" + assert resp.choices[0].message.content is None + assert resp.choices[0].provider_specific_fields["native_finish_reason"] == "NO_IMAGE" + + +def test_gemini_candidate_with_finish_reason_no_content_anthropic_messages(): + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) + + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash-image", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_resp = adapter.translate_openai_response_to_anthropic( + response=resp, + tool_name_mapping={}, + ) + assert anthropic_resp["stop_reason"] == "refusal" + assert anthropic_resp["content"] == [] + + +def test_gemini_candidate_with_finish_reason_no_content_responses_api(): + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash-image", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + responses_resp = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Generate picture", + responses_api_request={}, + chat_completion_response=resp, + ) + assert responses_resp.status == "incomplete" + assert responses_resp.incomplete_details is not None + assert responses_resp.incomplete_details.reason == "content_filter" + + +def test_gemini_candidate_other_finish_reasons_no_content(): + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + config = VertexGeminiConfig() + max_tokens_response = { + "candidates": [{"finishReason": "MAX_TOKENS", "index": 0}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 50, "totalTokenCount": 60}, + } + resp_length = config._transform_google_generate_content_to_openai_model_response( + completion_response=max_tokens_response, + model_response=ModelResponse(), + model="gemini-2.5-flash", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + assert len(resp_length.choices) == 1 + assert resp_length.choices[0].finish_reason == "length" + assert resp_length.choices[0].provider_specific_fields["native_finish_reason"] == "MAX_TOKENS" + + anthropic_length = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( + response=resp_length, + tool_name_mapping={}, + ) + assert anthropic_length["stop_reason"] == "max_tokens" + + responses_length = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="thinking request", + responses_api_request={}, + chat_completion_response=resp_length, + ) + assert responses_length.status == "incomplete" + assert responses_length.incomplete_details.reason == "max_output_tokens" + From cd66b34b45dace036126f4ea809df132407d907c Mon Sep 17 00:00:00 2001 From: tusharjamunkar Date: Sat, 12 Sep 2026 22:48:08 +0530 Subject: [PATCH 05/61] style(responses): apply ruff formatting to transformation.py --- .../litellm_completion_transformation/transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 13119085e46..10e56e85bff 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2273,9 +2273,7 @@ class LiteLLMCompletionResponsesConfig: finish_reason = choices[0].finish_reason status: Final[ResponsesAPIStatus] = ( - LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( - finish_reason - ) + LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status(finish_reason) ) incomplete_details = getattr(chat_completion_response, "incomplete_details", None) if incomplete_details is None and status == "incomplete": From e54399eff7a13e649b6a353486922166b1bf5688 Mon Sep 17 00:00:00 2001 From: tusharjamunkar Date: Sat, 12 Sep 2026 23:07:31 +0530 Subject: [PATCH 06/61] test: add direct coverage for content_filter and refusal in anthropic and responses adapters --- ...al_pass_through_adapters_transformation.py | 40 ++++++++++++ .../test_litellm_completion_responses.py | 62 +++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 03b9840b1c3..0e09e27f4db 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -102,6 +102,46 @@ def test_translate_chat_length_takes_precedence_over_refusal(): assert result.get("stop_details") is None +def test_translate_chat_content_filter_to_anthropic_response(): + response = ModelResponse( + id="chatcmpl-content-filter", + model="openai-model", + choices=[ + Choices( + index=0, + finish_reason="content_filter", + message=Message(content=None, role="assistant"), + ) + ], + usage=Usage(prompt_tokens=1, completion_tokens=0, total_tokens=1), + ) + + result = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "refusal" + + +def test_translate_chat_refusal_finish_reason_to_anthropic_response(): + response = ModelResponse( + id="chatcmpl-refusal-reason", + model="openai-model", + choices=[ + Choices( + index=0, + finish_reason="refusal", + message=Message(content=None, role="assistant"), + ) + ], + usage=Usage(prompt_tokens=1, completion_tokens=0, total_tokens=1), + ) + + result = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "refusal" + + def test_translate_streaming_openai_chunk_to_anthropic_content_block(): choices = [ StreamingChoices( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 46249e50572..f2ebbf316c3 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -4246,3 +4246,65 @@ class TestStreamingSnapshotItemIds: reasoning_items = _bridged_output_items(completed_event.response, "reasoning") assert len(reasoning_items) == 1 assert reasoning_items[0].id == streamed_event.item_id + + +def test_transform_chat_completion_response_incomplete_details(): + from openai.types.responses.response import IncompleteDetails + + resp_length = ModelResponse( + id="resp-length", + choices=[Choices(index=0, finish_reason="length", message=Message(content="cutoff", role="assistant"))], + model="gpt-4o", + ) + result_length = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_length, + ) + assert result_length.status == "incomplete" + assert result_length.incomplete_details is not None + assert result_length.incomplete_details.reason == "max_output_tokens" + + resp_filter = ModelResponse( + id="resp-filter", + choices=[Choices(index=0, finish_reason="content_filter", message=Message(content=None, role="assistant"))], + model="gpt-4o", + ) + result_filter = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_filter, + ) + assert result_filter.status == "incomplete" + assert result_filter.incomplete_details is not None + assert result_filter.incomplete_details.reason == "content_filter" + + resp_refusal = ModelResponse( + id="resp-refusal", + choices=[Choices(index=0, finish_reason="refusal", message=Message(content=None, role="assistant"))], + model="gpt-4o", + ) + result_refusal = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_refusal, + ) + assert result_refusal.status == "incomplete" + assert result_refusal.incomplete_details is not None + assert result_refusal.incomplete_details.reason == "content_filter" + + existing_details = IncompleteDetails(reason="content_filter") + resp_existing = ModelResponse( + id="resp-existing", + choices=[Choices(index=0, finish_reason="length", message=Message(content="cutoff", role="assistant"))], + model="gpt-4o", + ) + resp_existing.incomplete_details = existing_details + result_existing = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_existing, + ) + assert result_existing.status == "incomplete" + assert result_existing.incomplete_details == existing_details + From 5b5bbac769e548199393e54aacf346953c5c5528 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:16:21 +0000 Subject: [PATCH 07/61] fix(team): link new members to the shared team member budget so /team/update applies to them Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_helpers/utils.py | 23 +-- .../test_management_helpers_utils.py | 135 +++++++++--------- 2 files changed, 82 insertions(+), 76 deletions(-) diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index f3bd4b0f6dd..2e7458232cc 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -291,9 +291,9 @@ async def _clone_team_default_budget_for_member( member budget. Returns the new budget_id, or None if the default budget no longer exists in the DB. - 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). + Used when adding a new team member with a per-member ``budget_duration`` + but no other per-member limit, so the member keeps the team default's + values in their own private budget row while the reset window differs. ``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 @@ -344,14 +344,21 @@ async def _resolve_member_budget_id( """ 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. + Explicit per-member limits create a fresh budget. Otherwise the member is + linked to the team's shared default member budget, so later ``/team/update`` + changes reach them; ``/team/member_update`` clones that row on first write. + A lone ``budget_duration`` clones the default with the reset window + overridden, or creates a window-only budget when there is no team default. + With nothing set the member gets no budget. """ has_explicit_limit: Final = 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 and budget_duration is None: + default_budget: Final = await _budget_table(prisma_client, tx).find_unique( + where={"budget_id": default_team_budget_id} + ) + return default_team_budget_id if default_budget is not None else 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, 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 a6b1fc32eda..00de5171aa7 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 @@ -164,13 +164,16 @@ async def test_management_otel_span_redacts_nested_submission_env_var_secrets( @pytest.mark.asyncio -async def test_add_new_member_clones_default_team_budget_id(): +async def test_add_new_member_links_default_team_budget_id(): """ - Test that add_new_member CLONES the team's default member budget when - max_budget_in_team is None and a default_team_budget_id is provided. + A member added without any per-member limit must be LINKED to the team's + shared default member budget, not given a private copy of it. - Cloning (rather than sharing the same budget row) is what lets admins later - edit one member's budget without mutating every other member's budget. + Linking is what makes a later ``/team/update team_member_budget=...`` + reach existing members: the auth check reads the budget row behind the + membership, so a private clone would freeze the member at the old cap. + Per-member isolation is handled by ``/team/member_update`` cloning the + shared row on first write. """ from litellm.proxy._types import LitellmUserRoles @@ -178,7 +181,6 @@ async def test_add_new_member_clones_default_team_budget_id(): test_user_id = "test_user_123" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" - test_cloned_budget_id = "cloned_budget_xyz" test_admin_name = "test_admin" new_member = Member(user_id=test_user_id, role="user") @@ -202,36 +204,19 @@ async def test_add_new_member_clones_default_team_budget_id(): return_value=mock_user_response ) - # Mock the default budget row fetched for cloning. mock_default_budget_row = MagicMock() - mock_default_budget_row.model_dump.return_value = { - "budget_id": test_default_budget_id, - "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_default_budget_row.budget_id = test_default_budget_id mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( return_value=mock_default_budget_row ) - - # Mock the cloned budget row that .create() returns. - mock_cloned_budget_row = MagicMock() - mock_cloned_budget_row.budget_id = test_cloned_budget_id - mock_prisma_client.db.litellm_budgettable.create = AsyncMock( - return_value=mock_cloned_budget_row - ) + mock_prisma_client.db.litellm_budgettable.create = AsyncMock() # Mock the team membership creation mock_team_membership_response = MagicMock() mock_team_membership_response.model_dump.return_value = { "team_id": test_team_id, "user_id": test_user_id, - "budget_id": test_cloned_budget_id, + "budget_id": test_default_budget_id, "litellm_budget_table": None, } mock_prisma_client.db.litellm_teammembership.create = AsyncMock( @@ -251,33 +236,67 @@ async def test_add_new_member_clones_default_team_budget_id(): assert result_user is not None assert result_user.user_id == test_user_id - # Membership should be linked to the new cloned budget, not the shared default. + # Membership points at the shared default row itself. assert result_team_membership is not None - assert result_team_membership.budget_id == test_cloned_budget_id - assert result_team_membership.budget_id != test_default_budget_id + assert result_team_membership.budget_id == test_default_budget_id mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() mock_prisma_client.db.litellm_teammembership.create.assert_called_once() - # The clone must have happened: find_unique on the default, create for the clone. + # The default is only checked for existence; no private budget row is created. mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( where={"budget_id": test_default_budget_id} ) - mock_prisma_client.db.litellm_budgettable.create.assert_called_once() - cloned_create_data = ( - mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs["data"] - ) - # Cloned values from the default budget row - assert cloned_create_data["max_budget"] == 100.0 - assert cloned_create_data["tpm_limit"] == 1000 - assert cloned_create_data["budget_duration"] == "1d" - assert cloned_create_data["created_by"] == user_api_key_dict.user_id + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() team_membership_call_args = ( mock_prisma_client.db.litellm_teammembership.create.call_args ) create_data = team_membership_call_args.kwargs["data"] - assert create_data["budget_id"] == test_cloned_budget_id + assert create_data["budget_id"] == test_default_budget_id + + +@pytest.mark.asyncio +async def test_add_new_member_no_budget_when_default_budget_row_is_missing(): + """If team metadata still names a default member budget whose row was + deleted, the member must get no budget rather than a dangling link that + the membership foreign key would reject.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="missing-default-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": "missing-default-user", + "user_email": None, + "teams": ["team-md"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mock_user_response + ) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_budgettable.create = AsyncMock() + mock_prisma_client.db.litellm_teammembership.create = AsyncMock() + + _, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-md", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id="deleted-default", + ) + + assert result_team_membership is None + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() + mock_prisma_client.db.litellm_teammembership.create.assert_not_called() @pytest.mark.asyncio @@ -636,18 +655,17 @@ async def test_add_new_member_persists_budget_duration_without_max_budget(): @pytest.mark.asyncio -async def test_add_new_member_with_user_email_clones_default_budget(): +async def test_add_new_member_with_user_email_links_default_budget(): """ Test add_new_member with user_email instead of user_id and a team default - budget. The default budget should be CLONED into a new private row for - this user, not shared with other members of the team. + budget. The membership must link the shared default row so team-level + budget updates keep applying to this member. """ from litellm.proxy._types import LitellmUserRoles test_user_email = "test@example.com" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" - test_cloned_budget_id = "cloned_budget_for_email_user" test_admin_name = "test_admin" new_member = Member(user_email=test_user_email, role="user") @@ -669,35 +687,18 @@ async def test_add_new_member_with_user_email_clones_default_budget(): } mock_prisma_client.insert_data = AsyncMock(return_value=mock_user_response) - # Default budget that will be cloned mock_default_budget_row = MagicMock() - mock_default_budget_row.model_dump.return_value = { - "budget_id": test_default_budget_id, - "max_budget": 25.0, - "soft_budget": None, - "max_parallel_requests": None, - "tpm_limit": None, - "rpm_limit": None, - "model_max_budget": None, - "budget_duration": None, - "allowed_models": [], - } + mock_default_budget_row.budget_id = test_default_budget_id mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( return_value=mock_default_budget_row ) - - # Cloned budget result - mock_cloned_budget_row = MagicMock() - mock_cloned_budget_row.budget_id = test_cloned_budget_id - mock_prisma_client.db.litellm_budgettable.create = AsyncMock( - return_value=mock_cloned_budget_row - ) + mock_prisma_client.db.litellm_budgettable.create = AsyncMock() mock_team_membership_response = MagicMock() mock_team_membership_response.model_dump.return_value = { "team_id": test_team_id, "user_id": "generated_user_id", - "budget_id": test_cloned_budget_id, + "budget_id": test_default_budget_id, "litellm_budget_table": None, } mock_prisma_client.db.litellm_teammembership.create = AsyncMock( @@ -717,9 +718,8 @@ async def test_add_new_member_with_user_email_clones_default_budget(): assert result_user is not None assert result_user.user_email == test_user_email - # Membership should point at the cloned (private) budget, not the shared default. assert result_team_membership is not None - assert result_team_membership.budget_id == test_cloned_budget_id + assert result_team_membership.budget_id == test_default_budget_id mock_prisma_client.get_data.assert_called_once_with( key_val={"user_email": test_user_email}, @@ -733,11 +733,10 @@ async def test_add_new_member_with_user_email_clones_default_budget(): assert insert_data["user_email"] == test_user_email assert insert_data["teams"] == [test_team_id] - # Confirm the clone path ran mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( where={"budget_id": test_default_budget_id} ) - mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() @pytest.mark.asyncio From ab99be9dad023d7d5e25d3e9352f7954357d4963 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:43:21 +0000 Subject: [PATCH 08/61] test(team): cover team_member_budget propagation and per-member isolation end to end Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_management_helpers_utils.py | 172 +++++++++++++++--- 1 file changed, 148 insertions(+), 24 deletions(-) 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 00de5171aa7..7512a3dfae8 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,13 +1,17 @@ import json +from collections.abc import Mapping from datetime import datetime, timezone -from litellm._uuid import uuid -from unittest.mock import AsyncMock, MagicMock +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest - +import litellm +from litellm._uuid import uuid from litellm.proxy._types import ( + LiteLLM_BudgetTable, LiteLLM_TeamMembership, + LiteLLM_TeamTable, LiteLLM_UserTable, Member, UserAPIKeyAuth, @@ -165,19 +169,8 @@ async def test_management_otel_span_redacts_nested_submission_env_var_secrets( @pytest.mark.asyncio async def test_add_new_member_links_default_team_budget_id(): - """ - A member added without any per-member limit must be LINKED to the team's - shared default member budget, not given a private copy of it. - - Linking is what makes a later ``/team/update team_member_budget=...`` - reach existing members: the auth check reads the budget row behind the - membership, so a private clone would freeze the member at the old cap. - Per-member isolation is handled by ``/team/member_update`` cloning the - shared row on first write. - """ from litellm.proxy._types import LitellmUserRoles - # Setup test data test_user_id = "test_user_123" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" @@ -236,14 +229,12 @@ async def test_add_new_member_links_default_team_budget_id(): assert result_user is not None assert result_user.user_id == test_user_id - # Membership points at the shared default row itself. assert result_team_membership is not None assert result_team_membership.budget_id == test_default_budget_id mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() mock_prisma_client.db.litellm_teammembership.create.assert_called_once() - # The default is only checked for existence; no private budget row is created. mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( where={"budget_id": test_default_budget_id} ) @@ -258,9 +249,6 @@ async def test_add_new_member_links_default_team_budget_id(): @pytest.mark.asyncio async def test_add_new_member_no_budget_when_default_budget_row_is_missing(): - """If team metadata still names a default member budget whose row was - deleted, the member must get no budget rather than a dangling link that - the membership foreign key would reject.""" from litellm.proxy._types import LitellmUserRoles new_member = Member(user_id="missing-default-user", role="user") @@ -656,11 +644,6 @@ async def test_add_new_member_persists_budget_duration_without_max_budget(): @pytest.mark.asyncio async def test_add_new_member_with_user_email_links_default_budget(): - """ - Test add_new_member with user_email instead of user_id and a team default - budget. The membership must link the shared default row so team-level - budget updates keep applying to this member. - """ from litellm.proxy._types import LitellmUserRoles test_user_email = "test@example.com" @@ -739,6 +722,147 @@ async def test_add_new_member_with_user_email_links_default_budget(): mock_prisma_client.db.litellm_budgettable.create.assert_not_called() +class _FakeBudgetTable: + def __init__(self) -> None: + self.rows: dict[str, dict[str, object]] = {} + + def _record(self, budget_id: str) -> LiteLLM_BudgetTable: + row: Final = self.rows[budget_id] + return LiteLLM_BudgetTable(**{k: v for k, v in row.items() if k in LiteLLM_BudgetTable.model_fields}) + + async def create( + self, *, data: Mapping[str, object], include: Mapping[str, bool] | None = None + ) -> LiteLLM_BudgetTable: + budget_id: Final = str(data.get("budget_id") or uuid.uuid4()) + self.rows[budget_id] = {**data, "budget_id": budget_id} + return self._record(budget_id) + + async def find_unique(self, *, where: Mapping[str, str]) -> LiteLLM_BudgetTable | None: + return self._record(where["budget_id"]) if where["budget_id"] in self.rows else None + + async def update(self, *, where: Mapping[str, str], data: Mapping[str, object]) -> LiteLLM_BudgetTable: + self.rows[where["budget_id"]] = {**self.rows[where["budget_id"]], **data} + return self._record(where["budget_id"]) + + +class _FakeMembershipTable: + def __init__(self, budgets: _FakeBudgetTable) -> None: + self.budgets: Final = budgets + self.budget_ids: dict[tuple[str, str], str | None] = {} + + def membership(self, team_id: str, user_id: str) -> LiteLLM_TeamMembership: + budget_id: Final = self.budget_ids[(team_id, user_id)] + return LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + budget_id=budget_id, + litellm_budget_table=self.budgets._record(budget_id) if budget_id is not None else None, + ) + + async def create(self, *, data: Mapping[str, str], include: Mapping[str, bool]) -> LiteLLM_TeamMembership: + self.budget_ids[(data["team_id"], data["user_id"])] = data["budget_id"] + return self.membership(data["team_id"], data["user_id"]) + + async def upsert(self, *, where: Mapping[str, Mapping[str, str]], data: Mapping[str, Mapping[str, object]]) -> None: + key: Final = where["user_id_team_id"] + connect: Final = data["update"]["litellm_budget_table"] + assert isinstance(connect, dict) + self.budget_ids[(key["team_id"], key["user_id"])] = connect["connect"]["budget_id"] + + +class _FakeUserTable: + async def upsert(self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id=where["user_id"], teams=list(data["create"].get("teams", []))) + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: + return 1 + + +class _FakeDb: + def __init__(self) -> None: + self.litellm_budgettable: Final = _FakeBudgetTable() + self.litellm_teammembership: Final = _FakeMembershipTable(self.litellm_budgettable) + self.litellm_usertable: Final = _FakeUserTable() + + +@pytest.mark.asyncio +async def test_team_update_reaches_inherited_members_but_not_overridden_ones(): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.auth_checks import _check_team_member_budget + from litellm.proxy.management_endpoints.common_utils import _upsert_budget_and_membership + from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + from litellm.proxy.utils import ProxyLogging + + db: Final = _FakeDb() + prisma_client: Final = MagicMock() + prisma_client.db = db + admin: Final = UserAPIKeyAuth(user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN) + team_id: Final = "team-shared-default" + default_budget: Final = await db.litellm_budgettable.create(data={"budget_id": "team-default", "max_budget": 100.0}) + team: Final = LiteLLM_TeamTable(team_id=team_id, metadata={"team_member_budget_id": default_budget.budget_id}) + + for user_id in ("inherits", "overridden"): + await add_new_member( + new_member=Member(user_id=user_id, role="user"), + max_budget_in_team=None, + prisma_client=prisma_client, + team_id=team_id, + user_api_key_dict=admin, + litellm_proxy_admin_name="admin", + default_team_budget_id=default_budget.budget_id, + ) + + await _upsert_budget_and_membership( + db, + team_id=team_id, + user_id="overridden", + existing_budget_id=default_budget.budget_id, + user_api_key_dict=admin, + budget_patch={"max_budget": 50.0}, + team_default_budget_id=default_budget.budget_id, + ) + assert db.litellm_teammembership.membership(team_id, "inherits").budget_id == default_budget.budget_id + assert db.litellm_teammembership.membership(team_id, "overridden").budget_id != default_budget.budget_id + assert db.litellm_budgettable.rows[default_budget.budget_id]["max_budget"] == 100.0 + + with patch( # test-quality-ok: update_budget reads this module global; no dependency injection seam exists + "litellm.proxy.proxy_server.prisma_client", prisma_client + ): + await TeamMemberBudgetHandler.upsert_team_member_budget_table( + team_table=team, + user_api_key_dict=admin, + updated_kv={}, + team_member_budget=1.0, + ) + + async def spend_from_membership(counter_key: str, fallback_spend: float, max_budget: float | None = None) -> float: + return fallback_spend + + async def check(user_id: str, spend: float) -> None: + membership: Final = db.litellm_teammembership.membership(team_id, user_id).model_copy(update={"spend": spend}) + with patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists + "litellm.proxy.proxy_server.get_current_spend", spend_from_membership + ): + await _check_team_member_budget( + team_object=team, + user_object=LiteLLM_UserTable(user_id=user_id), + valid_token=UserAPIKeyAuth(token="tok", user_id=user_id, team_id=team_id), + prisma_client=prisma_client, + user_api_key_cache=MagicMock(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + team_membership=membership, + team_membership_loaded=True, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await check("inherits", spend=2.0) + assert exc_info.value.max_budget == 1.0 + await check("overridden", spend=2.0) + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await check("overridden", spend=60.0) + assert exc_info.value.max_budget == 50.0 + + @pytest.mark.asyncio async def test_attach_object_permission_to_dict_with_object_permission_id(): """ From 1a749d84bdd66706bb41cafdba28e7a8b6a20fa9 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:56:00 +0000 Subject: [PATCH 09/61] fix(proxy): track project spend and enforce project budgets additively Project-scoped keys never wrote spend to LiteLLM_ProjectTable, so /project/info stayed at 0 and project budgets could not block. Wire the PROJECT entity through the spend queue, redis buffer, and db writer, reserve and increment a spend:project counter, reseed it from the project row, reset project spend in the budget cascade, and read the live counter in the project max budget check. Team member budgets keep gating project-scoped keys alongside the project budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_checks.py | 34 ++-- .../proxy/common_utils/reset_budget_job.py | 24 +++ .../proxy/common_utils/user_api_key_cache.py | 10 ++ litellm/proxy/db/db_spend_update_writer.py | 79 ++++++++- .../redis_update_buffer.py | 7 + .../spend_update_queue.py | 4 + litellm/proxy/db/spend_counter_reseed.py | 5 + .../proxy/hooks/proxy_track_cost_callback.py | 6 + litellm/proxy/proxy_server.py | 30 ++++ .../spend_tracking/budget_reservation.py | 41 +++++ .../spend_tracking/spend_counter_batch.py | 11 +- litellm/repositories/prisma_protocols.py | 3 + litellm/repositories/unit_of_work.py | 2 + .../proxy/auth/test_auth_checks.py | 61 +++++++ .../common_utils/test_reset_budget_job.py | 30 +++- .../proxy/db/test_db_spend_update_writer.py | 76 +++++++++ .../proxy/db/test_spend_counter_reseed.py | 19 +++ .../hooks/test_proxy_track_cost_callback.py | 1 + .../test_spend_tracking_utils.py | 1 + .../proxy/test_budget_reservation.py | 156 ++++++++++++++++++ .../repositories/test_unit_of_work.py | 3 + 22 files changed, 583 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 321f8190f13..228a91ad446 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5261,6 +5261,7 @@ class DBSpendUpdateTransactions(TypedDict): team_member_list_transactions: dict[str, float] | None org_list_transactions: dict[str, float] | None org_member_list_transactions: ReadOnly[dict[str, float] | None] + project_list_transactions: ReadOnly[dict[str, float] | None] tag_list_transactions: dict[str, float] | None agent_list_transactions: dict[str, float] | None model_access_group_list_transactions: ReadOnly[dict[str, float] | None] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e3783c94dc7..ef17913d9ec 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -92,6 +92,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( model_access_group_registry_cache_key, model_access_group_spend_counter_key, object_permission_cache_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, tag_registry_cache_key, team_membership_auth_cache_key, @@ -5586,16 +5588,22 @@ async def _project_max_budget_check( if project_object.litellm_budget_table is not None: max_budget = project_object.litellm_budget_table.max_budget - if ( - max_budget is not None - and project_object.spend is not None - and math.isfinite(max_budget) - and project_object.spend > max_budget - ): + if max_budget is None or not math.isfinite(max_budget): + return + + from litellm.proxy.proxy_server import get_current_spend + + project_spend: Final = await get_current_spend( + counter_key=project_spend_counter_key(project_object.project_id), + fallback_spend=project_object.spend or 0.0, + max_budget=max_budget, + ) + + if project_spend >= max_budget: if valid_token: call_info: Final = CallInfo( token=valid_token.token, - spend=project_object.spend, + spend=project_spend, max_budget=max_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, @@ -5611,9 +5619,9 @@ async def _project_max_budget_check( ) raise litellm.BudgetExceededError( - current_cost=project_object.spend, + current_cost=project_spend, max_budget=max_budget, - message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_object.spend}, Max budget: {max_budget}", + message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}", entity_type=Litellm_EntityType.PROJECT.value, entity_id=project_object.project_id, ) @@ -5663,10 +5671,6 @@ async def _project_soft_budget_check( ) -def _project_cache_key(project_id: str) -> str: - return f"project_id:{project_id}" - - async def get_project_object( project_id: str, prisma_client: PrismaClient | None, @@ -5684,7 +5688,7 @@ async def get_project_object( return None # Check cache first - cache_key: Final = _project_cache_key(project_id) + cache_key: Final = project_cache_key(project_id) deserialized_project: Final = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_ProjectTableCachedObj, @@ -5726,7 +5730,7 @@ async def delete_cached_project_object( from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast await evict_and_broadcast( - cache_keys=(_project_cache_key(project_id),), + cache_keys=(project_cache_key(project_id),), user_api_key_cache=user_api_key_cache, ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index acb51e73daf..2baefa89943 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -41,6 +41,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, model_access_group_cache_key, model_access_group_spend_counter_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, ) from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row @@ -49,6 +51,7 @@ from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import SpendLinkedTable +from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( EndUserRepository, ModelAccessGroupBudgetRepository, @@ -115,6 +118,11 @@ class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol): def access_group_name(self) -> str: ... +class _ProjectRow(_BudgetLinkedRow, Protocol): + @property + def project_id(self) -> str: ... + + class _EndUserRow(_BudgetLinkedRow, Protocol): @property def user_id(self) -> str: ... @@ -185,6 +193,14 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...] return (model_access_group_cache_key(row.access_group_name),) +def _project_counter_key(row: _ProjectRow) -> str: + return project_spend_counter_key(row.project_id) + + +def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]: + return (project_cache_key(row.project_id),) + + def _enduser_counter_key(row: _EndUserRow) -> str: return f"spend:end_user:{row.user_id}" @@ -661,6 +677,11 @@ class ResetBudgetJob: where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), log_subject="model access groups", ) + projects: Final[tuple[_ProjectRow, ...]] = await self._fetch_linked_rows( + table=ProjectRepository(self.prisma_client).table, + where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), + log_subject="projects", + ) rollover_caps: Final[Mapping[str, float]] = MappingProxyType( { # mutable-ok: MappingProxyType wraps a one-shot dict comprehension b.budget_id: cap @@ -695,6 +716,7 @@ class ResetBudgetJob: (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in model_access_groups ), + *((_project_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in projects), *((_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) for row in endusers), ), rollover_caps=rollover_caps, @@ -704,6 +726,7 @@ class ResetBudgetJob: *(key for row in orgs for key in _org_cache_keys(row)), *(key for row in tags for key in _tag_cache_keys(row)), *(key for row in model_access_groups for key in _model_access_group_cache_keys(row)), + *(key for row in projects for key in _project_cache_keys(row)), *(key for row in endusers for key in _enduser_cache_keys(row)), ), ) @@ -731,6 +754,7 @@ class ResetBudgetJob: _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) + _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) _queue_enduser_resets(uow.endusers, cascade) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 1c7a379897f..2187ed63ea5 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -306,6 +306,16 @@ def model_access_group_spend_counter_key(access_group_name: str) -> str: return f"spend:model_access_group:{access_group_name}" +def project_cache_key(project_id: str) -> str: + """Cache key one project row is stored under; shared by auth, spend tracking and the spend writer.""" + return f"project_id:{project_id}" + + +def project_spend_counter_key(project_id: str) -> str: + """Spend counter key for one project; the reservation, cost callback, auth and reseed paths all read it.""" + return f"spend:project:{project_id}" + + #: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds #: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch. END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__" diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index a90d1351fd7..51d00b9789c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -44,6 +44,7 @@ from litellm.proxy._types import ( SpendUpdateQueueItem, ToolDiscoveryQueueItem, ) +from litellm.proxy.common_utils.user_api_key_cache import project_cache_key from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, build_bulk_upsert, @@ -116,6 +117,7 @@ class _SpendBatch(Protocol): litellm_teammembership: BatchTable litellm_organizationtable: BatchTable litellm_organizationmembership: BatchTable + litellm_projecttable: BatchTable litellm_tagtable: BatchTable litellm_agentstable: BatchTable litellm_modelaccessgroupbudgettable: BatchTable @@ -254,6 +256,7 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, + project_id: str | None = None, ) -> bool: """Record the request's spend, answering whether its cost still needs charging. @@ -335,6 +338,7 @@ class DBSpendUpdateWriter: hashed_token=hashed_token, team_id=team_id, org_id=org_id, + project_id=project_id, end_user_id=end_user_id, prisma_client=prisma_client, litellm_proxy_budget_name=litellm_proxy_budget_name, @@ -631,6 +635,7 @@ class DBSpendUpdateWriter: litellm_proxy_budget_name: str | None, payload: SpendLogsPayload, request_model_access_groups: Sequence[str] = (), + project_id: str | None = None, ): """ Runs all 13 spend-update helpers sequentially inside a single asyncio task. @@ -694,6 +699,18 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + try: + await self._update_project_db( + response_cost=response_cost, + project_id=project_id, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_project_db failed: %s", + traceback.format_exc(), + ) + try: await self._update_tag_db( response_cost=response_cost, @@ -956,6 +973,33 @@ class DBSpendUpdateWriter: ) raise e + async def _update_project_db( + self, + response_cost: float | None, + project_id: str | None, + prisma_client: PrismaClient | None, + ): + try: + if project_id is None or prisma_client is None: + return + + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.PROJECT, + entity_id=project_id, + response_cost=response_cost, + ) + ) + except Exception as e: + spend_log_error( + "Spend tracking - failed to enqueue project spend update. project_id=%s, response_cost=%s - %s", + project_id, + response_cost, + str(e), + exc=e, + ) + raise e + async def _update_agent_db( self, response_cost: float | None, @@ -1193,8 +1237,8 @@ class DBSpendUpdateWriter: if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " - "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, tags=%d, " - "agents=%d, model_access_groups=%d", + "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, " + "projects=%d, tags=%d, agents=%d, model_access_groups=%d", len(db_spend_update_transactions.get("key_list_transactions") or {}), len(db_spend_update_transactions.get("user_list_transactions") or {}), len(db_spend_update_transactions.get("team_list_transactions") or {}), @@ -1202,6 +1246,7 @@ class DBSpendUpdateWriter: len(db_spend_update_transactions.get("end_user_list_transactions") or {}), len(db_spend_update_transactions.get("team_member_list_transactions") or {}), len(db_spend_update_transactions.get("org_member_list_transactions") or {}), + len(db_spend_update_transactions.get("project_list_transactions") or {}), len(db_spend_update_transactions.get("tag_list_transactions") or {}), len(db_spend_update_transactions.get("agent_list_transactions") or {}), len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}), @@ -1762,6 +1807,22 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + ### UPDATE PROJECT TABLE ### + project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions") + await DBSpendUpdateWriter._update_entity_spend_in_db( + entity_name="Project", + transactions=project_list_transactions, + table_accessor="litellm_projecttable", + where_field="project_id", + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + await DBSpendUpdateWriter._invalidate_project_caches( + project_ids=tuple(project_list_transactions or ()), + proxy_logging_obj=proxy_logging_obj, + ) + ### UPDATE TAG TABLE ### tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( @@ -1800,11 +1861,23 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + @staticmethod + async def _invalidate_project_caches(project_ids: Sequence[str], proxy_logging_obj: ProxyLogging | None) -> None: + if not project_ids or proxy_logging_obj is None: + return + user_api_key_cache: Final = proxy_logging_obj.call_details.get("user_api_key_cache") + if user_api_key_cache is None: + return + for project_id in project_ids: + await user_api_key_cache.async_delete_cache(key=project_cache_key(project_id)) + @staticmethod async def _update_entity_spend_in_db( entity_name: str, transactions: dict[str, float] | None, - table_accessor: Literal["litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable"], + table_accessor: Literal[ + "litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable", "litellm_projecttable" + ], where_field: str, n_retry_times: int, prisma_client: PrismaClient, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 6f49a00b763..cead63795a2 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -70,6 +70,7 @@ _SpendTransactionField: TypeAlias = Literal[ "team_member_list_transactions", "org_list_transactions", "org_member_list_transactions", + "project_list_transactions", "tag_list_transactions", "agent_list_transactions", "model_access_group_list_transactions", @@ -83,6 +84,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( "team_member_list_transactions", "org_list_transactions", "org_member_list_transactions", + "project_list_transactions", "tag_list_transactions", "agent_list_transactions", "model_access_group_list_transactions", @@ -418,6 +420,10 @@ class RedisUpdateBuffer: Litellm_EntityType.ORGANIZATION_MEMBER, db_spend_update_transactions.get("org_member_list_transactions"), ), + ( + Litellm_EntityType.PROJECT, + db_spend_update_transactions.get("project_list_transactions"), + ), ( Litellm_EntityType.TAG, db_spend_update_transactions.get("tag_list_transactions"), @@ -885,6 +891,7 @@ class RedisUpdateBuffer: org_member_list_transactions=_merged_entity_transactions( list_of_transactions, "org_member_list_transactions" ), + project_list_transactions=_merged_entity_transactions(list_of_transactions, "project_list_transactions"), tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"), agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"), model_access_group_list_transactions=_merged_entity_transactions( diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index bc068d10daf..2b8535cb113 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -138,6 +138,7 @@ class SpendUpdateQueue(BaseUpdateQueue): team_member_list_transactions={}, org_list_transactions={}, org_member_list_transactions={}, + project_list_transactions={}, tag_list_transactions={}, agent_list_transactions={}, model_access_group_list_transactions={}, @@ -152,6 +153,7 @@ class SpendUpdateQueue(BaseUpdateQueue): Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions", Litellm_EntityType.ORGANIZATION: "org_list_transactions", Litellm_EntityType.ORGANIZATION_MEMBER: "org_member_list_transactions", + Litellm_EntityType.PROJECT: "project_list_transactions", Litellm_EntityType.TAG: "tag_list_transactions", Litellm_EntityType.AGENT: "agent_list_transactions", Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions", @@ -192,6 +194,8 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict = db_spend_update_transactions["org_list_transactions"] elif dict_key == "org_member_list_transactions": transactions_dict = db_spend_update_transactions["org_member_list_transactions"] + elif dict_key == "project_list_transactions": + transactions_dict = db_spend_update_transactions["project_list_transactions"] elif dict_key == "tag_list_transactions": transactions_dict = db_spend_update_transactions["tag_list_transactions"] elif dict_key == "agent_list_transactions": diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 89a07234c6c..2dd028454d6 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -26,6 +26,7 @@ from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_lookup_gate import db_lookup_gate from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( BudgetWindowSpendRepository, EndUserRepository, @@ -77,6 +78,7 @@ class SpendCounterReseed: spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend spend:user:{user_id} -> LiteLLM_UserTable.spend spend:org:{org_id} -> LiteLLM_OrganizationTable.spend + spend:project:{project_id} -> LiteLLM_ProjectTable.spend End-user and tag spend counters intentionally do not reseed here. Their auth paths already load the corresponding objects via get_end_user_object() @@ -157,6 +159,9 @@ class SpendCounterReseed: row = await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": org_id} ) + elif counter_key.startswith("spend:project:"): + project_id: Final = counter_key[len("spend:project:") :] + row = await ProjectRepository(prisma_client).table.find_unique(where={"project_id": project_id}) else: return None except Exception: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 1ae106be390..0c562cf37ef 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -267,6 +267,7 @@ class _ProxyDBLogger(CustomLogger): start_time=actual_start_time, end_time=datetime.now(), org_id=user_api_key_dict.org_id, + project_id=user_api_key_dict.project_id, ) @log_db_metrics @@ -318,6 +319,7 @@ class _ProxyDBLogger(CustomLogger): user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) + project_id: Final = cast(str | None, metadata.get("user_api_key_project_id", None)) key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None)) end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None) sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) @@ -368,6 +370,7 @@ class _ProxyDBLogger(CustomLogger): budget_reservation=budget_reservation, request_tags=tags, model_access_groups=model_access_groups, + project_id=project_id, ) if not charged: return @@ -651,6 +654,7 @@ async def _update_database_and_spend_counters( budget_reservation: dict | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, ) -> bool: if budget_reservation is not None: await _reconcile_budget_reservation_before_db_update( @@ -668,6 +672,7 @@ async def _update_database_and_spend_counters( start_time=start_time, end_time=end_time, org_id=org_id, + project_id=project_id, ) except Exception: if budget_reservation is not None: @@ -698,6 +703,7 @@ async def _update_database_and_spend_counters( tags=request_tags, request_started_at=start_time, model_access_groups=model_access_groups, + project_id=project_id, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7964556531..1e84e5f56e2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -417,6 +417,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( get_management_object_ttl, model_access_group_cache_key, model_access_group_spend_counter_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, ) from litellm.proxy.config_resolvers import resolve_fields @@ -2780,6 +2782,7 @@ async def increment_spend_counters( tags: list[str] | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2801,6 +2804,7 @@ async def increment_spend_counters( end_user_id=end_user_id, tags=tags, model_access_groups=model_access_groups, + project_id=project_id, ), ): await _increment_spend_counters_batched( @@ -2814,6 +2818,7 @@ async def increment_spend_counters( tags=tags, request_started_at=request_started_at, model_access_groups=model_access_groups, + project_id=project_id, ) @@ -2828,6 +2833,7 @@ async def _increment_spend_counters_batched( tags: list[str] | None, request_started_at: datetime | None, model_access_groups: Sequence[str] | None, + project_id: str | None = None, ): """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET.""" reserved_counter_keys: Final = await _reconcile_budget_reservation_for_counter_update( @@ -3028,6 +3034,13 @@ async def _increment_spend_counters_batched( ) if org_id is not None else None, + _prepare_project_spend_increment( + project_id=project_id, + response_cost=cost, + reserved_counter_keys=reserved_counter_keys, + ) + if project_id is not None + else None, ) if coro is not None ) @@ -3180,6 +3193,23 @@ async def _prepare_org_spend_increment( return (pending,) if pending is not None else () +async def _prepare_project_spend_increment( + project_id: str | None, + response_cost: float, + reserved_counter_keys: set[str], +) -> tuple[PendingSpendIncrement, ...]: + if project_id is None: + return () + + pending: Final = await _prepare_unreserved_spend_counter_increment( + counter_key=project_spend_counter_key(project_id), + source_cache_key=project_cache_key(project_id), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + return (pending,) if pending is not None else () + + async def _prepare_unreserved_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 373f2d0fe36..24b8470145c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import math from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -29,6 +30,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, model_access_group_cache_key, model_access_group_spend_counter_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, team_membership_reservation_cache_key, ) @@ -62,6 +65,7 @@ _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { "Tag": Litellm_EntityType.TAG.value, "Model access group": Litellm_EntityType.MODEL_ACCESS_GROUP.value, "Organization": Litellm_EntityType.ORGANIZATION.value, + "Project": Litellm_EntityType.PROJECT.value, } @@ -542,6 +546,13 @@ async def _get_budget_counters( if org_counter is not None: counters.append(org_counter) + project_counter: Final = await _get_project_budget_counter( + valid_token=valid_token, + user_api_key_cache=user_api_key_cache, + ) + if project_counter is not None: + counters.append(project_counter) + return counters @@ -751,6 +762,36 @@ async def _get_org_budget_counter( ) +async def _get_project_budget_counter( + valid_token: UserAPIKeyAuth, + user_api_key_cache: UserApiKeyCache, +) -> _BudgetCounter | None: + if valid_token.project_id is None: + return None + + source_cache_key: Final = project_cache_key(valid_token.project_id) + project_object: Final = await user_api_key_cache.async_get_cache(key=source_cache_key) + if project_object is None: + return None + + project_budget_table: Final = _get_value(project_object, "litellm_budget_table") + if project_budget_table is None: + return None + + project_max_budget: Final = _to_float(_get_value(project_budget_table, "max_budget")) + if project_max_budget is None or project_max_budget <= 0 or not math.isfinite(project_max_budget): + return None + + return _BudgetCounter( + counter_key=project_spend_counter_key(valid_token.project_id), + source_cache_key=source_cache_key, + max_budget=project_max_budget, + fallback_spend=_to_float(_get_value(project_object, "spend")) or 0.0, + entity_type="Project", + entity_id=valid_token.project_id, + ) + + def _get_budget_limit_counters( entity_prefix: str, entity_type: str, diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 7106d88c655..ddb074ae023 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -12,7 +12,10 @@ from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_utils.user_api_key_cache import model_access_group_spend_counter_key +from litellm.proxy.common_utils.user_api_key_cache import ( + model_access_group_spend_counter_key, + project_spend_counter_key, +) _CounterValues: Final = TypeAdapter(dict[str, float | None]) _NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({}) @@ -154,6 +157,8 @@ def _iter_admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) yield f"spend:end_user:{end_user_id}" if token.org_id is not None: yield f"spend:org:{token.org_id}" + if token.project_id is not None: + yield project_spend_counter_key(token.project_id) def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> frozenset[str]: @@ -168,10 +173,12 @@ def post_call_counter_keys( end_user_id: str | None, tags: Sequence[object] | None, model_access_groups: Sequence[object] | None, + project_id: str | None = None, ) -> frozenset[str]: """Every counter ``increment_spend_counters`` warm-checks, except budget windows which bind on read.""" entity_keys: Final = admission_counter_keys( - UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id), end_user_id + UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id, project_id=project_id), + end_user_id, ) tag_keys: Final = frozenset(f"spend:tag:{tag}" for tag in tags or () if tag and isinstance(tag, str)) group_keys: Final = frozenset( diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 93b8c5c7cd7..60c16fbd746 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -152,4 +152,7 @@ class PrismaBatch(Protocol): @property def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ... + @property + def litellm_projecttable(self) -> BatchTable: ... + async def commit(self) -> None: ... diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 0cdce307f9b..c09e5eb75d4 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -109,6 +109,7 @@ class BudgetCascadeUnitOfWork: organizations: LinkedSpendResetWrites tags: LinkedSpendResetWrites model_access_groups: LinkedSpendResetWrites + projects: LinkedSpendResetWrites endusers: LinkedSpendResetWrites budgets: BudgetWindowWrites @@ -135,6 +136,7 @@ async def budget_cascade_unit_of_work( organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable), tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable), + projects=LinkedSpendResetWrites(table=batch.litellm_projecttable), endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable), budgets=BudgetWindowWrites(table=batch.litellm_budgettable), ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 8c8b755195f..0d5c0dd5d72 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7138,6 +7138,67 @@ async def test_project_allowlist_enforced_when_key_models_empty(): assert exc_info.value.code == "403" +def _project_with_budget(spend: float, max_budget: float): + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_ProjectTableCachedObj + + return LiteLLM_ProjectTableCachedObj( + project_id="p-budget", + team_id="t-1", + budget_id="b-1", + spend=spend, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b-1", max_budget=max_budget), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "counter_spend, db_spend, blocks", + [ + pytest.param(5.0, 0.0, True, id="counter-at-budget-blocks-despite-stale-db-row"), + pytest.param(4.99, 0.0, False, id="counter-under-budget-admits"), + pytest.param(None, 5.0, True, id="no-counter-falls-back-to-persisted-spend"), + pytest.param(None, 0.0, False, id="no-counter-and-no-persisted-spend-admits"), + ], +) +async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks): + """LIT-3269: project budget enforcement must read the cross-pod + ``spend:project:{id}`` counter first and only fall back to the cached row's + spend, matching key/team/org checks. The boundary is inclusive (>=).""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.auth_checks import _project_max_budget_check + + real_spend_counter_cache = DualCache() + if counter_spend is not None: + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=counter_spend) + valid_token = UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget", team_id="t-1", user_id="u-1") + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + if not blocks: + await _project_max_budget_check( + project_object=_project_with_budget(spend=db_spend, max_budget=5.0), + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + return + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _project_max_budget_check( + project_object=_project_with_budget(spend=db_spend, max_budget=5.0), + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + await asyncio.sleep(0) + + assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value + assert exc_info.value.entity_id == "p-budget" + assert exc_info.value.current_cost == 5.0 + proxy_logging_obj.budget_alerts.assert_awaited_once() + assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget" + + def test_is_user_proxy_admin_rejects_view_only_admin(): """This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an Admin Viewer answering True here would gain every write route. Read parity for diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 943a6c905c0..8c254686385 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -78,6 +78,7 @@ class MockBatcher: self.litellm_organizationtable = _Table("org", self) self.litellm_tagtable = _Table("tag", self) self.litellm_modelaccessgroupbudgettable = _Table("model_access_group", self) + self.litellm_projecttable = _Table("project", self) self.litellm_endusertable = _Table("enduser", self) async def commit(self): @@ -93,6 +94,7 @@ class MockDB: self.litellm_organizationtable = MockTable() self.litellm_tagtable = MockTable() self.litellm_modelaccessgroupbudgettable = MockTable() + self.litellm_projecttable = MockTable() self.batch_calls: List[Dict[str, Any]] = [] self.batchers: List[MockBatcher] = [] @@ -1507,13 +1509,19 @@ _INVALIDATION_CASES = [ "spend:model_access_group:gpt-4-group", {"model_access_group:gpt-4-group"}, ), + ( + "litellm_projecttable", + type("Project", (), {"project_id": "proj-1"}), + "spend:project:proj-1", + {"project_id:proj-1"}, + ), ] @pytest.mark.parametrize( "table_attr, linked_row, counter_key, cache_keys", _INVALIDATION_CASES, - ids=["team_membership", "key", "org", "tag", "model_access_group"], + ids=["team_membership", "key", "org", "tag", "model_access_group", "project"], ) def test_budget_table_reset_invalidates_counters_and_management_cache( reset_budget_job, mock_prisma_client, monkeypatch, table_attr, linked_row, counter_key, cache_keys @@ -1657,6 +1665,25 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( counter_cache.in_memory_cache.delete_cache.assert_any_call(key=f"spend:model_access_group:{name}") +def test_project_reset_zeroes_spend_on_due_tiers(reset_budget_job, mock_prisma_client, monkeypatch): + """A project linked to an expiring budget tier has its spend zeroed in the same cascade transaction.""" + _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")] + mock_prisma_client.db.litellm_projecttable.set_find_many_results( + [type("Project", (), {"project_id": "proj-1", "spend": 12.0, "budget_id": "budget-due"})] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + expected_where = {"budget_id": {"in": ["budget-due"]}, "spend": {"gt": 0}} + assert mock_prisma_client.db.litellm_projecttable.find_many_calls == [{"where": expected_where}] + writes = _batch_writes(mock_prisma_client, "project", op="update_many") + assert len(writes) == 1 + assert writes[0]["where"] == expected_where + assert writes[0]["data"] == {"spend": 0} + assert mock_prisma_client.db.batchers[0].committed is True + + def test_budget_cascade_carries_access_group_overage_when_rollover_enabled( rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch ): @@ -1802,6 +1829,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo ("org", "update_many"), ("tag", "update_many"), ("model_access_group", "update_many"), + ("project", "update_many"), ("enduser", "update_many"), ("budget", "update_many"), } diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index c547d06904b..da5879a375a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1060,6 +1060,82 @@ async def test_batch_database_updates_queues_org_member_spend_for_the_request_us assert transactions["org_member_list_transactions"] == {"organization_id::org1::user_id::u1": 0.1} +@pytest.mark.asyncio +async def test_project_spend_is_persisted_to_project_table_and_project_cache_is_evicted(): + """Regression for LIT-3269: a request made with a project-scoped key must + increment LiteLLM_ProjectTable.spend, otherwise /project/info stays at 0 + and the project budget never blocks. The cached project row is evicted so + the next auth check reads the fresh spend.""" + db_writer: Final = DBSpendUpdateWriter() + await db_writer._batch_database_updates( + response_cost=0.25, + user_id="u1", + hashed_token="t1", + team_id="team-1", + org_id=None, + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25}, + project_id="proj-1", + ) + await db_writer._batch_database_updates( + response_cost=0.5, + user_id="u1", + hashed_token="t1", + team_id="team-1", + org_id=None, + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-2", "model": "gpt-4o-mini", "spend": 0.5}, + project_id="proj-1", + ) + transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert transactions["project_list_transactions"] == {"proj-1": 0.75} + assert transactions["team_member_list_transactions"] == {"team_id::team-1::user_id::u1": 0.75} + + mock_batcher: Final = MagicMock() + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher)) + user_api_key_cache: Final = MagicMock() + user_api_key_cache.async_delete_cache = AsyncMock() + proxy_logging: Final = MagicMock() + proxy_logging.call_details = {"user_api_key_cache": user_api_key_cache} + + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=transactions, + ) + + mock_batcher.litellm_projecttable.update_many.assert_called_once_with( + where={"project_id": "proj-1"}, + data={"spend": {"increment": 0.75}}, + ) + user_api_key_cache.async_delete_cache.assert_any_await(key="project_id:proj-1") + + +@pytest.mark.asyncio +async def test_batch_database_updates_without_project_id_touches_no_project_row(): + db_writer: Final = DBSpendUpdateWriter() + await db_writer._batch_database_updates( + response_cost=0.1, + user_id="u1", + hashed_token="t1", + team_id=None, + org_id=None, + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.1}, + ) + transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + + assert transactions["project_list_transactions"] == {} + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id(): """ diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index bca6344b3f7..53e91b8792e 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -67,12 +67,14 @@ class _FakePrismaClient: error: Exception | None = None, end_user_row: SimpleNamespace | None = None, end_user_error: Exception | None = None, + project_row: SimpleNamespace | None = None, ) -> None: self.db = SimpleNamespace( litellm_budgetwindowspend=_FakeFindUniqueTable(row=row, error=error), litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), litellm_endusertable=_FakeFindUniqueTable(row=end_user_row, error=end_user_error), litellm_verificationtoken=_InFlightCountingTable(), + litellm_projecttable=_FakeFindUniqueTable(row=project_row), ) @@ -428,6 +430,23 @@ async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys(): assert prisma.db.litellm_verificationtoken.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY +@pytest.mark.asyncio +async def test_from_db_reseeds_project_counter_from_the_project_row(): + """LIT-3269: a cold ``spend:project:{id}`` counter seeds from LiteLLM_ProjectTable.spend, + so a fresh pod enforces the project budget against persisted spend rather than 0.""" + prisma: Final = _FakePrismaClient(project_row=SimpleNamespace(project_id="proj-1", spend=7.25)) + + assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") == 7.25 + assert prisma.db.litellm_projecttable.where_clauses == [{"project_id": "proj-1"}] + + +@pytest.mark.asyncio +async def test_from_db_returns_none_for_a_missing_project_row(): + prisma: Final = _FakePrismaClient(project_row=None) + + assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") is None + + @pytest.mark.asyncio async def test_from_db_still_never_reads_the_end_user_row(): """A cold end-user counter keeps seeding from the cached end-user object the auth diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index dfc95db3e14..965e134772d 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -586,6 +586,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda tags=["tag-a"], request_started_at=start_time, model_access_groups=("premium",), + project_id=None, ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 8b105e94d19..834cf8d100d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1187,6 +1187,7 @@ async def test_api_key_preserved_through_failure_hook_to_database(): start_time, end_time, org_id, + project_id=None, ): """Mock update_database and capture the payload it creates""" from litellm.proxy.spend_tracking.spend_tracking_utils import ( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 032722d3259..014f240d9cc 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ( LiteLLM_EndUserTable, Litellm_EntityType, LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, LiteLLM_TagTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, @@ -631,6 +632,161 @@ async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_ await release_budget_reservation(reservation) +def _project_scoped_token() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + token="key-project-scoped", + spend=0.0, + user_id="user-proj", + team_id="team-proj", + project_id="proj-1", + ) + + +async def _seed_project_scoped_budgets( + key_cache: DualCache, + team_member_spend: float, + team_member_max_budget: float, + project_spend: float, + project_max_budget: float, +) -> None: + await key_cache.async_set_cache( + key="team_membership:user-proj:team-proj", + value=LiteLLM_TeamMembership( + user_id="user-proj", + team_id="team-proj", + spend=team_member_spend, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=team_member_max_budget), + ).model_dump(), + ) + await key_cache.async_set_cache( + key="project_id:proj-1", + value=LiteLLM_ProjectTableCachedObj( + project_id="proj-1", + team_id="team-proj", + budget_id="project-budget-id", + spend=project_spend, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=project_max_budget), + ).model_dump(), + ) + + +@pytest.mark.asyncio +async def test_should_reserve_project_and_team_member_counters_for_project_scoped_key(spend_counter_state): + """LIT-3269: a key carrying user_id, team_id and project_id reserves against + both the team member counter and the project counter; neither replaces the + other. After the call the project counter reflects the real cost once, not + the reservation plus the post-call increment.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + await _seed_project_scoped_budgets( + key_cache, + team_member_spend=0.1, + team_member_max_budget=1.0, + project_spend=0.2, + project_max_budget=1.0, + ) + + estimated = estimate_request_max_cost(request_body=_request_body(), route="/chat/completions", llm_router=None) + assert estimated is not None and estimated > 0 + + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=_project_scoped_token(), + team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None), + user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0), + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx( + 0.1 + estimated + ) + assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.2 + estimated) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token="key-project-scoped", + team_id="team-proj", + user_id="user-proj", + response_cost=0.05, + budget_reservation=reservation, + project_id="proj-1", + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.25) + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx(0.15) + + +@pytest.mark.asyncio +async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spend_counter_state): + """LIT-3269: the project budget is additive. A project with plenty of + headroom must not let a key through once its team member budget is spent.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + await _seed_project_scoped_budgets( + key_cache, + team_member_spend=1.0, + team_member_max_budget=1.0, + project_spend=0.0, + project_max_budget=100.0, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=_project_scoped_token(), + team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None), + user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0), + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert "TeamMember=user-proj:team-proj" in str(exc_info.value) + assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") in (None, pytest.approx(0.0)) + + +@pytest.mark.asyncio +async def test_exhausted_project_budget_blocks_project_scoped_key(spend_counter_state): + """LIT-3269: with team member headroom left, the project budget alone blocks the key.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + await _seed_project_scoped_budgets( + key_cache, + team_member_spend=0.0, + team_member_max_budget=100.0, + project_spend=5.0, + project_max_budget=5.0, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=_project_scoped_token(), + team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None), + user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0), + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert "Project=proj-1" in str(exc_info.value) + assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") in ( + None, + pytest.approx(0.0), + ) + + @pytest.mark.asyncio async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter_state): """The reservation path mirrors the read path: no personal user counter for a team key. diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index b52b8ced31e..9eff248b917 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -35,6 +35,7 @@ class FakeBatch: self.litellm_organizationtable = FakeBatchTable("litellm_organizationtable", self.calls) self.litellm_tagtable = FakeBatchTable("litellm_tagtable", self.calls) self.litellm_modelaccessgroupbudgettable = FakeBatchTable("litellm_modelaccessgroupbudgettable", self.calls) + self.litellm_projecttable = FakeBatchTable("litellm_projecttable", self.calls) self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls) async def commit(self) -> None: @@ -94,6 +95,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): uow.organizations.queue_spend_zero(where=linked) uow.tags.queue_spend_zero(where=linked) uow.model_access_groups.queue_spend_zero(where=linked) + uow.projects.queue_spend_zero(where=linked) uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}}) uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) assert batch.commit_count == 0 @@ -105,6 +107,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): ("litellm_organizationtable.update_many", linked, {"spend": 0}), ("litellm_tagtable.update_many", linked, {"spend": 0}), ("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}), + ("litellm_projecttable.update_many", linked, {"spend": 0}), ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}), ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}), ] From b00cd15bd73a380c77aad0d04672d27ee4bc13cf Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 02:07:11 +0000 Subject: [PATCH 10/61] test(proxy): import project_cache_key from user_api_key_cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/test_key_management_endpoints.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 4e70063015d..b57f8b7857d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -38,11 +38,10 @@ from litellm.proxy._types import ( from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.proxy.auth.auth_checks import ( _delete_cache_key_object, - _project_cache_key, jwt_key_mapping_cache_key, ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_org_key_limits, @@ -18951,7 +18950,7 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( async def _cache_with_project(project_id: str, project_models: list[str]) -> UserApiKeyCache: user_api_key_cache = UserApiKeyCache() await user_api_key_cache.async_set_cache( - key=_project_cache_key(project_id), + key=project_cache_key(project_id), value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id="team-lit-5823", models=project_models), model_type=LiteLLM_ProjectTableCachedObj, ) From b20f1422eb204c2cb2fba26912a548454cd18b02 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 02:28:38 +0000 Subject: [PATCH 11/61] fix(proxy): carry project_id through key metadata enrichment and drop docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/user_api_key_cache.py | 2 -- litellm/proxy/hooks/proxy_track_cost_callback.py | 2 ++ tests/test_litellm/proxy/auth/test_auth_checks.py | 3 --- .../proxy/common_utils/test_reset_budget_job.py | 1 - tests/test_litellm/proxy/db/test_db_spend_update_writer.py | 4 ---- tests/test_litellm/proxy/db/test_spend_counter_reseed.py | 2 -- .../proxy/hooks/test_proxy_track_cost_callback.py | 3 +++ tests/test_litellm/proxy/test_budget_reservation.py | 7 ------- 8 files changed, 5 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 2187ed63ea5..0386b58070d 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -307,12 +307,10 @@ def model_access_group_spend_counter_key(access_group_name: str) -> str: def project_cache_key(project_id: str) -> str: - """Cache key one project row is stored under; shared by auth, spend tracking and the spend writer.""" return f"project_id:{project_id}" def project_spend_counter_key(project_id: str) -> str: - """Spend counter key for one project; the reservation, cost callback, auth and reseed paths all read it.""" return f"spend:project:{project_id}" diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0c562cf37ef..5e525108ade 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -504,6 +504,8 @@ class _ProxyDBLogger(CustomLogger): metadata["user_api_key_team_id"] = key_obj.team_id if metadata.get("user_api_key_org_id") is None: metadata["user_api_key_org_id"] = key_obj.org_id + if metadata.get("user_api_key_project_id") is None: + metadata["user_api_key_project_id"] = key_obj.project_id except Exception: verbose_proxy_logger.debug( "Failed to enrich failure metadata with key info for api_key=%s", diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 0d5c0dd5d72..fe469fc574e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7161,9 +7161,6 @@ def _project_with_budget(spend: float, max_budget: float): ], ) async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks): - """LIT-3269: project budget enforcement must read the cross-pod - ``spend:project:{id}`` counter first and only fall back to the cached row's - spend, matching key/team/org checks. The boundary is inclusive (>=).""" from litellm.caching.dual_cache import DualCache from litellm.proxy.auth.auth_checks import _project_max_budget_check diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 8c254686385..48b06649237 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1666,7 +1666,6 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( def test_project_reset_zeroes_spend_on_due_tiers(reset_budget_job, mock_prisma_client, monkeypatch): - """A project linked to an expiring budget tier has its spend zeroed in the same cascade transaction.""" _make_counter_invalidation_job(monkeypatch) mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")] mock_prisma_client.db.litellm_projecttable.set_find_many_results( diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index da5879a375a..60cc742577b 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1062,10 +1062,6 @@ async def test_batch_database_updates_queues_org_member_spend_for_the_request_us @pytest.mark.asyncio async def test_project_spend_is_persisted_to_project_table_and_project_cache_is_evicted(): - """Regression for LIT-3269: a request made with a project-scoped key must - increment LiteLLM_ProjectTable.spend, otherwise /project/info stays at 0 - and the project budget never blocks. The cached project row is evicted so - the next auth check reads the fresh spend.""" db_writer: Final = DBSpendUpdateWriter() await db_writer._batch_database_updates( response_cost=0.25, diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 53e91b8792e..ff0b67d426b 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -432,8 +432,6 @@ async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys(): @pytest.mark.asyncio async def test_from_db_reseeds_project_counter_from_the_project_row(): - """LIT-3269: a cold ``spend:project:{id}`` counter seeds from LiteLLM_ProjectTable.spend, - so a fresh pod enforces the project budget against persisted spend rather than 0.""" prisma: Final = _FakePrismaClient(project_row=SimpleNamespace(project_id="proj-1", spend=7.25)) assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") == 7.25 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 965e134772d..202495517ad 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1372,6 +1372,7 @@ async def test_enrich_failure_metadata_with_full_key_lookup(): mock_key_obj.user_id = "fetched-user-id" mock_key_obj.team_id = "fetched-team-id" mock_key_obj.org_id = "fetched-org-id" + mock_key_obj.project_id = "fetched-project-id" mock_team_obj = MagicMock() mock_team_obj.team_alias = "fetched-team-alias" @@ -1395,12 +1396,14 @@ async def test_enrich_failure_metadata_with_full_key_lookup(): "user_api_key_team_id": None, "user_api_key_team_alias": None, "user_api_key_org_id": None, + "user_api_key_project_id": None, } result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) assert result["user_api_key_alias"] == "fetched-key-alias" assert result["user_api_key_user_id"] == "fetched-user-id" assert result["user_api_key_team_id"] == "fetched-team-id" assert result["user_api_key_org_id"] == "fetched-org-id" + assert result["user_api_key_project_id"] == "fetched-project-id" assert result["user_api_key_team_alias"] == "fetched-team-alias" diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 014f240d9cc..c834ac05f0a 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -672,10 +672,6 @@ async def _seed_project_scoped_budgets( @pytest.mark.asyncio async def test_should_reserve_project_and_team_member_counters_for_project_scoped_key(spend_counter_state): - """LIT-3269: a key carrying user_id, team_id and project_id reserves against - both the team member counter and the project counter; neither replaces the - other. After the call the project counter reflects the real cost once, not - the reservation plus the post-call increment.""" counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) await _seed_project_scoped_budgets( @@ -724,8 +720,6 @@ async def test_should_reserve_project_and_team_member_counters_for_project_scope @pytest.mark.asyncio async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spend_counter_state): - """LIT-3269: the project budget is additive. A project with plenty of - headroom must not let a key through once its team member budget is spent.""" counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) await _seed_project_scoped_budgets( @@ -755,7 +749,6 @@ async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spen @pytest.mark.asyncio async def test_exhausted_project_budget_blocks_project_scoped_key(spend_counter_state): - """LIT-3269: with team member headroom left, the project budget alone blocks the key.""" counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) await _seed_project_scoped_budgets( From 47be6c8aeb578c71c13f56eea8d1db6a5b81ba3f Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 18:34:36 +0000 Subject: [PATCH 12/61] feat(mcp): allowlist client applications for MCP gateway access Adds the mcp_allowed_clients general setting, enforced against the clientInfo.name each MCP client sends in its initialize request. A client not on the list, or one that does not identify itself, is rejected with 403 before any stateful session is created. The setting is configurable from config.yaml and from the Admin UI MCP network settings page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/client_allowlist.py | 67 +++++ .../proxy/_experimental/mcp_server/server.py | 66 ++++- litellm/proxy/_types.py | 4 + litellm/proxy/proxy_server.py | 4 + .../mcp_server/test_client_allowlist.py | 105 ++++++++ .../mcp_server/test_mcp_server.py | 244 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 165 +++++++++--- .../_components/MCPNetworkSettings.test.tsx | 66 ++++- .../_components/MCPNetworkSettings.tsx | 77 +++++- 9 files changed, 745 insertions(+), 53 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/client_allowlist.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py diff --git a/litellm/proxy/_experimental/mcp_server/client_allowlist.py b/litellm/proxy/_experimental/mcp_server/client_allowlist.py new file mode 100644 index 00000000000..57343c39565 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/client_allowlist.py @@ -0,0 +1,67 @@ +""" +Gateway-level allowlist of MCP client applications, matched against the +``clientInfo.name`` a client sends in its JSON-RPC ``initialize`` request. The +name is client-supplied, so this is a policy control and not a security boundary. +""" + +import json +from dataclasses import dataclass +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger + +MCP_ALLOWED_CLIENTS_SETTING: Final = "mcp_allowed_clients" + +_ALLOWED_CLIENTS_ADAPTER: Final = TypeAdapter(list[str]) + + +@dataclass(frozen=True, slots=True) +class MCPClientRejection: + client_name: str | None + + @property + def details(self) -> str: + if self.client_name is None: + return ( + "MCP initialize request did not identify the client application (clientInfo.name). " + f"This gateway only admits clients listed in {MCP_ALLOWED_CLIENTS_SETTING}." + ) + return f"MCP client '{self.client_name}' is not listed in this gateway's {MCP_ALLOWED_CLIENTS_SETTING}." + + +def parse_allowed_mcp_clients(raw_setting: object) -> frozenset[str] | None: + """None when the setting is absent (not enforced). A malformed setting admits nobody.""" + if raw_setting is None: + return None + try: + return frozenset(_ALLOWED_CLIENTS_ADAPTER.validate_python(raw_setting)) + except ValidationError: + verbose_logger.warning( + "%s is not a list of client names (%r); rejecting every MCP client until it is fixed", + MCP_ALLOWED_CLIENTS_SETTING, + raw_setting, + ) + return frozenset() + + +def extract_mcp_client_name(body: bytes) -> str | None: + try: + data: Final = json.loads(body) + except (json.JSONDecodeError, UnicodeDecodeError): + return None + params: Final = data.get("params") if isinstance(data, dict) else None + client_info: Final = params.get("clientInfo") if isinstance(params, dict) else None + name: Final = client_info.get("name") if isinstance(client_info, dict) else None + return name if isinstance(name, str) and name else None + + +def check_mcp_client_allowed(body: bytes, allowed_clients: frozenset[str] | None) -> MCPClientRejection | None: + """None when the initialize is admitted, otherwise the rejection to send back as a 403.""" + if allowed_clients is None: + return None + client_name: Final = extract_mcp_client_name(body) + if client_name is not None and client_name in allowed_clients: + return None + return MCPClientRejection(client_name=client_name) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7feb1fd468d..cca867d4d2a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -463,6 +463,11 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( MCPAuthenticatedUser, ) + from litellm.proxy._experimental.mcp_server.client_allowlist import ( + MCP_ALLOWED_CLIENTS_SETTING, + check_mcp_client_allowed, + parse_allowed_mcp_clients, + ) from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( SERVER_OUTCOMES_META_KEY, AggregateToolListing, @@ -3810,6 +3815,43 @@ if MCP_AVAILABLE: except (json.JSONDecodeError, TypeError): return False + def _load_allowed_mcp_clients() -> frozenset[str] | None: + from litellm.proxy.proxy_server import general_settings + + return parse_allowed_mcp_clients(general_settings.get(MCP_ALLOWED_CLIENTS_SETTING)) + + async def _reject_initialize_from_disallowed_client( + scope: Scope, + receive: Receive, + send: Send, + body: bytes, + client_ip: str | None, + ) -> bool: + """Send a 403 and return True when the initialize body names a client the gateway does not admit.""" + rejection: Final = check_mcp_client_allowed(body, _load_allowed_mcp_clients()) + if rejection is None: + return False + verbose_logger.warning( + "Rejecting MCP initialize from client %r (ip=%s): not listed in %s", + rejection.client_name, + client_ip, + MCP_ALLOWED_CLIENTS_SETTING, + ) + forbidden: Final = JSONResponse( + status_code=403, + content={"error": "Forbidden", "details": rejection.details}, + ) + await forbidden(scope, receive, send) + return True + + def _replay_consumed_messages(consumed_messages: list[Message], receive: Receive) -> Receive: + async def wrapped_receive() -> Message: + if consumed_messages: + return consumed_messages.pop(0) + return await receive() + + return wrapped_receive + async def _read_request_body_for_routing( receive: Receive, ) -> tuple[list[Message], bytes]: @@ -4510,6 +4552,10 @@ if MCP_AVAILABLE: if scope.get("method") == "POST": consumed_messages, body = await _read_request_body_for_routing(receive) is_initialize = _is_initialize_request(body) + if is_initialize and await _reject_initialize_from_disallowed_client( + scope, receive, send, body, _client_ip + ): + return use_stateful: Final = bool(session_id or is_initialize) target_manager: Final = session_manager_stateful if use_stateful else session_manager_stateless @@ -4540,15 +4586,8 @@ if MCP_AVAILABLE: return # Replay body messages if we consumed them for peeking - original_receive: Final = receive if consumed_messages: - - async def wrapped_receive(): - if consumed_messages: - return consumed_messages.pop(0) - return await original_receive() - - receive = wrapped_receive + receive = _replay_consumed_messages(consumed_messages, receive) # Serialize requests on the same stateful session so concurrent # callers don't clobber each other's auth context mid-flight. @@ -4785,6 +4824,15 @@ if MCP_AVAILABLE: await initialize_session_managers() await asyncio.sleep(0.1) + sse_consumed_messages, sse_body = ( + await _read_request_body_for_routing(receive) if scope.get("method") == "POST" else ([], b"") + ) + if _is_initialize_request(sse_body) and await _reject_initialize_from_disallowed_client( + scope, receive, send, sse_body, _sse_client_ip + ): + return + sse_receive: Final = _replay_consumed_messages(sse_consumed_messages, receive) + async with _gateway_initialize_instructions_request_scope( user_api_key_auth, mcp_servers, @@ -4792,7 +4840,7 @@ if MCP_AVAILABLE: scoped_server_endpoint=scoped_server_endpoint, is_initialize=scope.get("method") == "GET", ): - await sse_session_manager.handle_request(scope, receive, send) + await sse_session_manager.handle_request(scope, sse_receive, send) except MCPUpstreamAuthError as e: # Upstream delegated auth returned 401; surface it to the client so # standards-compliant MCP clients trigger the upstream OAuth flow. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 680c63393e8..bbebb99055f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2853,6 +2853,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).", ) + mcp_allowed_clients: list[str] | None = Field( + None, + description="MCP client applications admitted by the gateway, matched exactly against the clientInfo.name the client sends in its initialize request (for example 'claude-code'). When set, an initialize from any other client, or one that does not identify itself, is rejected with 403. Unset means every client is admitted. The name is client-supplied, so this is a policy control rather than a security boundary.", + ) mcp_trusted_proxy_ranges: list[str] | None = Field( None, description="CIDR ranges of trusted reverse proxies. When set, X-Forwarded-For and X-Forwarded-* origin headers are only trusted from these IPs.", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7d8413d2ce..aa98491cf3b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7179,6 +7179,9 @@ class ProxyConfig: "enable_openai_websocket_passthrough" ) + if "mcp_allowed_clients" not in self._yaml_general_settings_keys: + general_settings["mcp_allowed_clients"] = _general_settings.get("mcp_allowed_clients") + if "user_api_key_cache_max_size" not in self._yaml_general_settings_keys: db_cache_max_size: Final = _general_settings.get("user_api_key_cache_max_size") try: @@ -17137,6 +17140,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "maximum_spend_logs_cleanup_run_budget": "String", "maximum_spend_logs_cleanup_batch_timeout": "String", "mcp_internal_ip_ranges": "List", + "mcp_allowed_clients": "List", "mcp_trusted_proxy_ranges": "List", "mcp_xff_num_trusted_hops": "Integer", "always_include_stream_usage": "Boolean", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py new file mode 100644 index 00000000000..fce8a0a6c3b --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py @@ -0,0 +1,105 @@ +import json +from typing import Final + +import pytest + +from litellm.proxy._experimental.mcp_server.client_allowlist import ( + MCP_ALLOWED_CLIENTS_SETTING, + MCPClientRejection, + check_mcp_client_allowed, + extract_mcp_client_name, + parse_allowed_mcp_clients, +) + + +def _initialize_body(client_info: object) -> bytes: + return json.dumps( + { + "jsonrpc": "2.0", + "id": 0, + "method": "initialize", + "params": {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": client_info}, + } + ).encode() + + +CLAUDE_CODE: Final = _initialize_body({"name": "claude-code", "version": "2.1.274"}) +ANTIGRAVITY: Final = _initialize_body({"name": "antigravity-cli", "version": "1.0.0"}) + + +@pytest.mark.parametrize( + ("raw_setting", "expected"), + ( + (None, None), + ([], frozenset()), + (["antigravity-cli"], frozenset({"antigravity-cli"})), + (["antigravity-cli", "codex-mcp-client"], frozenset({"antigravity-cli", "codex-mcp-client"})), + ("antigravity-cli", frozenset()), + ([1, "antigravity-cli"], frozenset()), + ({"name": "antigravity-cli"}, frozenset()), + ), +) +def test_parse_allowed_mcp_clients(raw_setting: object, expected: frozenset[str] | None) -> None: + assert parse_allowed_mcp_clients(raw_setting) == expected + + +@pytest.mark.parametrize( + ("body", "expected"), + ( + (CLAUDE_CODE, "claude-code"), + (_initialize_body({"name": "", "version": "1"}), None), + (_initialize_body({"version": "1"}), None), + (_initialize_body({"name": 7}), None), + (_initialize_body("claude-code"), None), + (b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', None), + (b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":[]}', None), + (b'["not", "an", "object"]', None), + (b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"clientInfo":{"name":"clau', None), + (b"\xff\xfe", None), + (b"", None), + ), +) +def test_extract_mcp_client_name(body: bytes, expected: str | None) -> None: + assert extract_mcp_client_name(body) == expected + + +def test_unconfigured_allowlist_admits_every_client_including_unidentified_ones() -> None: + assert check_mcp_client_allowed(CLAUDE_CODE, None) is None + assert check_mcp_client_allowed(b'{"method":"initialize","params":{}}', None) is None + assert check_mcp_client_allowed(b"garbage", None) is None + + +def test_listed_client_is_admitted_and_unlisted_client_is_rejected_by_name() -> None: + allowed: Final = frozenset({"antigravity-cli"}) + assert check_mcp_client_allowed(ANTIGRAVITY, allowed) is None + assert check_mcp_client_allowed(CLAUDE_CODE, allowed) == MCPClientRejection(client_name="claude-code") + + +def test_matching_is_exact_not_prefix_or_case_insensitive() -> None: + allowed: Final = frozenset({"claude-code"}) + assert check_mcp_client_allowed(_initialize_body({"name": "Claude-Code"}), allowed) is not None + assert check_mcp_client_allowed(_initialize_body({"name": "claude-code-sdk"}), allowed) is not None + assert check_mcp_client_allowed(_initialize_body({"name": " claude-code"}), allowed) is not None + + +def test_empty_allowlist_rejects_every_client() -> None: + assert check_mcp_client_allowed(ANTIGRAVITY, frozenset()) == MCPClientRejection(client_name="antigravity-cli") + assert check_mcp_client_allowed(CLAUDE_CODE, frozenset()) == MCPClientRejection(client_name="claude-code") + + +def test_missing_or_malformed_client_metadata_is_rejected_when_allowlist_is_set() -> None: + allowed: Final = frozenset({"antigravity-cli"}) + assert check_mcp_client_allowed(_initialize_body({"version": "1"}), allowed) == MCPClientRejection(None) + assert check_mcp_client_allowed(b'{"method":"initialize","params":{}}', allowed) == MCPClientRejection(None) + assert check_mcp_client_allowed(b"{not json", allowed) == MCPClientRejection(None) + + +def test_rejection_details_name_the_setting_and_the_offending_client() -> None: + named: Final = MCPClientRejection(client_name="claude-code").details + assert "claude-code" in named + assert MCP_ALLOWED_CLIENTS_SETTING in named + + anonymous: Final = MCPClientRejection(client_name=None).details + assert "clientInfo.name" in anonymous + assert MCP_ALLOWED_CLIENTS_SETTING in anonymous + assert "None" not in anonymous diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index f5e4a420496..8550226e19d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import contextvars import os from datetime import datetime, timedelta @@ -2035,6 +2036,249 @@ async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless( assert not any(name.startswith(b"x-mcp-debug") for name in headers) +_CLAUDE_CODE_INITIALIZE: Final = ( + b'{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-06-18",' + b'"capabilities":{},"clientInfo":{"name":"claude-code","version":"2.1.274"}}}' +) +_ANTIGRAVITY_INITIALIZE: Final = ( + b'{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-06-18",' + b'"capabilities":{},"clientInfo":{"name":"antigravity-cli","version":"1.0.0"}}}' +) +_ANONYMOUS_INITIALIZE: Final = b'{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"capabilities":{}}}' +_TOOLS_LIST: Final = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}' + + +async def _drain_body(receive) -> bytes: + chunks: list[bytes] = [] + while True: + message = await receive() + chunks.append(message.get("body", b"")) + if not message.get("more_body", False): + return b"".join(chunks) + + +def _forbidden_client_response(send: AsyncMock) -> tuple[int, dict[str, str]]: + import json as _json + + start: Final = send.call_args_list[0].args[0] + body: Final = b"".join(call.args[0].get("body", b"") for call in send.call_args_list[1:]) + return start["status"], _json.loads(body) + + +@contextlib.contextmanager +def _client_allowlist_patches(allowed_clients: object): + settings: Final = {} if allowed_clients is None else {"mcp_allowed_clients": allowed_clients} + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(UserAPIKeyAuth(user_id="allowlist-user"), None, None, None, None, {}), + ), + patch("litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True), + patch("litellm.proxy.proxy_server.general_settings", settings), + ): + yield + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_body", "expected_details"), + ( + ( + _CLAUDE_CODE_INITIALIZE, + "MCP client 'claude-code' is not listed in this gateway's mcp_allowed_clients.", + ), + ( + _ANONYMOUS_INITIALIZE, + "MCP initialize request did not identify the client application (clientInfo.name). " + "This gateway only admits clients listed in mcp_allowed_clients.", + ), + ), +) +async def test_streamable_http_rejects_initialize_from_unlisted_client_before_session_creation( + request_body: bytes, expected_details: str +) -> None: + from starlette.types import Scope + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []} + receive: Final = AsyncMock(return_value={"type": "http.request", "body": request_body, "more_body": False}) + send: Final = AsyncMock() + stateful_handle: Final = AsyncMock() + stateless_handle: Final = AsyncMock() + session_cap: Final = AsyncMock(return_value=True) + + with ( + _client_allowlist_patches(["antigravity-cli"]), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), + ), + patch("litellm.proxy._experimental.mcp_server.server._enforce_stateful_session_cap_for_owner", session_cap), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + assert _forbidden_client_response(send) == (403, {"error": "Forbidden", "details": expected_details}) + stateful_handle.assert_not_awaited() + stateless_handle.assert_not_awaited() + session_cap.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("allowed_clients", "request_body"), + ( + (["antigravity-cli"], _ANTIGRAVITY_INITIALIZE), + (["claude-code", "antigravity-cli"], _CLAUDE_CODE_INITIALIZE), + (None, _CLAUDE_CODE_INITIALIZE), + (None, _ANONYMOUS_INITIALIZE), + ), +) +async def test_streamable_http_admits_listed_or_unrestricted_initialize_and_replays_body( + allowed_clients: list[str] | None, request_body: bytes +) -> None: + from starlette.types import Receive, Scope, Send + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []} + receive: Final = AsyncMock( + side_effect=[ + {"type": "http.request", "body": request_body[:20], "more_body": True}, + {"type": "http.request", "body": request_body[20:], "more_body": False}, + ] + ) + send: Final = AsyncMock() + downstream_bodies: Final[list[bytes]] = [] + + async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None: + downstream_bodies.append(await _drain_body(downstream_receive)) + + stateful_handle: Final = AsyncMock(side_effect=handle_request) + stateless_handle: Final = AsyncMock() + + with ( + _client_allowlist_patches(allowed_clients), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), + ), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + assert downstream_bodies == [request_body] + stateless_handle.assert_not_awaited() + send.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("allowed_clients", ([], "claude-code", [{"name": "claude-code"}])) +async def test_streamable_http_empty_or_malformed_allowlist_admits_nobody(allowed_clients: object) -> None: + from starlette.types import Scope + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []} + receive: Final = AsyncMock( + return_value={"type": "http.request", "body": _CLAUDE_CODE_INITIALIZE, "more_body": False} + ) + send: Final = AsyncMock() + stateful_handle: Final = AsyncMock() + + with ( + _client_allowlist_patches(allowed_clients), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + status, body = _forbidden_client_response(send) + assert status == 403 + assert body["details"] == "MCP client 'claude-code' is not listed in this gateway's mcp_allowed_clients." + stateful_handle.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_streamable_http_allowlist_only_inspects_initialize_requests() -> None: + from starlette.types import Receive, Scope, Send + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []} + receive: Final = AsyncMock(side_effect=[{"type": "http.request", "body": _TOOLS_LIST, "more_body": False}]) + send: Final = AsyncMock() + downstream_bodies: Final[list[bytes]] = [] + + async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None: + downstream_bodies.append(await _drain_body(downstream_receive)) + + with ( + _client_allowlist_patches(["antigravity-cli"]), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=AsyncMock(side_effect=handle_request)), + ), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + assert downstream_bodies == [_TOOLS_LIST] + send.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_body", "admitted"), + ((_ANTIGRAVITY_INITIALIZE, True), (_CLAUDE_CODE_INITIALIZE, False), (_ANONYMOUS_INITIALIZE, False)), +) +async def test_sse_endpoint_applies_the_same_client_allowlist(request_body: bytes, admitted: bool) -> None: + from starlette.types import Receive, Scope, Send + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp/sse", "headers": []} + receive: Final = AsyncMock(side_effect=[{"type": "http.request", "body": request_body, "more_body": False}]) + send: Final = AsyncMock() + downstream_bodies: Final[list[bytes]] = [] + + async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None: + downstream_bodies.append(await _drain_body(downstream_receive)) + + with ( + _client_allowlist_patches(["antigravity-cli"]), + patch( + "litellm.proxy._experimental.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._check_passthrough_upstream_auth", + new_callable=AsyncMock, + ), + patch.object(mcp_module.sse_session_manager, "handle_request", side_effect=handle_request), + ): + await mcp_module.handle_sse_mcp(scope, receive, send) + + if admitted: + assert downstream_bodies == [request_body] + send.assert_not_awaited() + return + assert downstream_bodies == [] + status, body = _forbidden_client_response(send) + assert status == 403 + assert body["error"] == "Forbidden" + assert "mcp_allowed_clients" in body["details"] + + @pytest.mark.asyncio async def test_mcp_routing_chunked_initialize_to_stateful(): """ diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 41c4956dba6..4a0f543dc38 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7428,10 +7428,18 @@ async def test_update_general_settings_keeps_yaml_pass_through_endpoints_next_to request.query_params = {} return request - settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam - yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in - initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here - master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401 + settings: Final = patch( + "litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]} + ) # test-quality-ok: the method reads this module global; no injection seam + yaml_endpoints: Final = patch( + "litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint] + ) # test-quality-ok: module global holding the YAML endpoints the fix merges in + initialize: Final = patch( + "litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock() + ) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here + master_key: Final = patch( + "litellm.proxy.proxy_server.master_key", "sk-master" + ) # test-quality-ok: a set master key is what makes a missing Authorization header a 401 with settings, yaml_endpoints, initialize, master_key: await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]}) @@ -7479,10 +7487,18 @@ async def test_update_general_settings_db_pass_through_endpoint_overrides_yaml_e request.headers = {} request.query_params = {} - settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam - yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in - initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here - master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401 + settings: Final = patch( + "litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]} + ) # test-quality-ok: the method reads this module global; no injection seam + yaml_endpoints: Final = patch( + "litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint] + ) # test-quality-ok: module global holding the YAML endpoints the fix merges in + initialize: Final = patch( + "litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock() + ) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here + master_key: Final = patch( + "litellm.proxy.proxy_server.master_key", "sk-master" + ) # test-quality-ok: a set master key is what makes a missing Authorization header a 401 with settings, yaml_endpoints, initialize, master_key: await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]}) @@ -8597,9 +8613,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): async def assert_reservation_not_finalized_yet(**kwargs): assert budget_reservation["finalized"] is False incremented_counters.append(kwargs["counter_key"]) - return ps.PendingSpendIncrement( - counter_key=kwargs["counter_key"], increment=kwargs["increment"] - ) + return ps.PendingSpendIncrement(counter_key=kwargs["counter_key"], increment=kwargs["increment"]) import litellm.proxy.proxy_server as ps @@ -10144,9 +10158,15 @@ async def _lit6973_drive_realtime_session( side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj) ) ws: Final = websocket if websocket is not None else _lit6973_fake_realtime_ws() - can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test - pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state - route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object + can_call = patch.object( + ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error) + ) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test + pre = patch.object( + ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call + ) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state + route = patch.object( + ps, "route_request", new=AsyncMock(return_value=fake_llm_call()) + ) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object with can_call, pre, route: await ps.realtime_websocket_endpoint( websocket=ws, @@ -10278,13 +10298,9 @@ async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot( from litellm.proxy.utils import InternalUsageCache dual_cache: Final = DualCache() - await dual_cache.async_set_cache( - key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True - ) + await dual_cache.async_set_cache(key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True) limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache)) - stash: Final = RequestRateLimiterStash( - parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} - ) + stash: Final = RequestRateLimiterStash(parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]}) reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} stash_token: Final = _request_stash.set(stash) @@ -10336,9 +10352,7 @@ async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_ limiter's integer in-memory fallback, double-decrement the counter so the key admits more sessions than max_parallel_requests allows. With the success stamp present the route leaves the slot and the stash alone.""" - dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( - backend_logged_success=True - ) + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(backend_logged_success=True) assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == { "slot-1": 1.0, @@ -10384,8 +10398,12 @@ async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): async def _record(counter_key: str) -> None: invalidated.append(counter_key) - failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated - sink = patch.object(ps, "_invalidate_spend_counter", new=_record) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable + failing_release = patch.object( + br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down")) + ) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated + sink = patch.object( + ps, "_invalidate_spend_counter", new=_record + ) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable with failing_release, sink: await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) @@ -10401,8 +10419,12 @@ async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback from litellm.proxy.spend_tracking import budget_reservation as br reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} - failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch - failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail + failing_release = patch.object( + br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down")) + ) # test-quality-ok: forces the fallback branch + failing_invalidate = patch.object( + br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down")) + ) # test-quality-ok: forces the fallback itself to fail with failing_release, failing_invalidate: await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) @@ -12987,9 +13009,15 @@ async def test_moderations_response_carries_litellm_call_id_header(): user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", spend=0.0) with ( - patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point - patch.object(proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call())), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable - patch.object(proxy_server_module, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global, no injection point + patch.object( + proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data + ), # test-quality-ok: the route reads this module global, no injection point + patch.object( + proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call()) + ), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable + patch.object( + proxy_server_module, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global, no injection point ): mock_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) mock_logging.update_request_status = AsyncMock() @@ -13026,9 +13054,15 @@ async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplo verbose_proxy_logger.propagate = True try: with ( - patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point - patch.object(proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))), # test-quality-ok: fakes the provider failure so the real route's error log is observable - patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + patch.object( + proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data + ), # test-quality-ok: the route reads this module global, no injection point + patch.object( + proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key")) + ), # test-quality-ok: fakes the provider failure so the real route's error log is observable + patch.object( + proxy_server_module, "proxy_logging_obj", new=fake_logging + ), # test-quality-ok: module global, no injection point caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised, ): @@ -13061,7 +13095,9 @@ async def test_moderations_unparseable_body_bills_the_callers_litellm_call_id(): fake_logging.post_call_failure_hook = AsyncMock() with ( - patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + patch.object( + proxy_server_module, "proxy_logging_obj", new=fake_logging + ), # test-quality-ok: module global, no injection point pytest.raises(ProxyException) as raised, ): await proxy_server_module.moderations( @@ -13089,8 +13125,12 @@ async def test_moderations_already_shaped_failure_answers_with_the_callers_litel fake_logging.post_call_failure_hook = AsyncMock() with ( - patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point - patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + patch.object( + proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc) + ), # test-quality-ok: the route reads this module global, no injection point + patch.object( + proxy_server_module, "proxy_logging_obj", new=fake_logging + ), # test-quality-ok: module global, no injection point pytest.raises(ProxyException) as raised, ): await proxy_server_module.moderations( @@ -13125,8 +13165,12 @@ async def test_audio_speech_already_shaped_failure_answers_with_the_callers_lite fake_logging.post_call_failure_hook = AsyncMock() with ( - patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point - patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + patch.object( + proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc) + ), # test-quality-ok: the route reads this module global, no injection point + patch.object( + proxy_server_module, "proxy_logging_obj", new=fake_logging + ), # test-quality-ok: module global, no injection point pytest.raises(type(exc)) as raised, ): await proxy_server_module.audio_speech( @@ -13814,6 +13858,43 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough() assert ps.general_settings["enable_openai_websocket_passthrough"] is False +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_general_settings, expected", + [ + ({"mcp_allowed_clients": ["antigravity-cli"]}, ["antigravity-cli"]), + ({"mcp_allowed_clients": []}, []), + ({}, None), + ], +) +async def test_update_general_settings_propagates_mcp_allowed_clients(db_general_settings, expected): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + with patch("litellm.proxy.proxy_server.general_settings", {"mcp_allowed_clients": ["claude-code"]}): + await proxy_config._update_general_settings(db_general_settings=db_general_settings) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["mcp_allowed_clients"] == expected + + +@pytest.mark.asyncio +async def test_update_general_settings_keeps_yaml_mcp_allowed_clients(): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._yaml_general_settings_keys = {"mcp_allowed_clients"} + + with patch("litellm.proxy.proxy_server.general_settings", {"mcp_allowed_clients": ["claude-code"]}): + await proxy_config._update_general_settings(db_general_settings={"mcp_allowed_clients": ["codex-mcp-client"]}) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["mcp_allowed_clients"] == ["claude-code"] + + async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch): from tests.large_text import text from tests.test_litellm.litellm_core_utils.event_loop_lag import ( @@ -13855,14 +13936,18 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp { "model_name": "self-hosted", "litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"}, - "model_info": {"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}}, + "model_info": { + "custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None} + }, } ] ), ) response, took, lags = await timed_with_loop_lags( - lambda: proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me off the loop")) + lambda: proxy_server_module.token_counter( + TokenCountRequest(model="self-hosted", prompt="count me off the loop") + ) ) assert response.tokenizer_type == "huggingface_tokenizer" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx index 9526f5de074..b6521acd7e2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx @@ -93,7 +93,7 @@ describe("MCPNetworkSettings", () => { await waitFor(() => expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges", ["10.0.0.0/8"]), ); - expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges"); }); it("clears the setting instead of saving an empty list", async () => { @@ -103,4 +103,68 @@ describe("MCPNetworkSettings", () => { await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges")); expect(updateConfigFieldSetting).not.toHaveBeenCalled(); }); + + it("renders the stored allowed client names once settings load", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: ["antigravity-cli", "codex-mcp-client"] }, + ]); + + renderSettings(); + + expect(await screen.findByText("antigravity-cli")).toBeInTheDocument(); + expect(screen.getByText("codex-mcp-client")).toBeInTheDocument(); + }); + + it("adds typed client names on Enter and saves them under mcp_allowed_clients", async () => { + renderSettings(); + const input = await screen.findByRole("textbox", { name: "Allowed client names" }); + + await userEvent.type(input, "antigravity-cli, codex-mcp-client{Enter}"); + + expect(screen.getByText("antigravity-cli")).toBeInTheDocument(); + expect(screen.getByText("codex-mcp-client")).toBeInTheDocument(); + expect(input).toHaveValue(""); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ + "antigravity-cli", + "codex-mcp-client", + ]), + ); + expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients"); + }); + + it("removes a client name and clears the setting when the list becomes empty", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: ["claude-code"] }, + ]); + + renderSettings(); + await userEvent.click(await screen.findByRole("button", { name: "Remove claude-code" })); + + expect(screen.queryByText("claude-code")).not.toBeInTheDocument(); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients")); + expect(updateConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients", expect.anything()); + }); + + it("keeps the private ranges and the allowed clients as independent settings on save", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_internal_ip_ranges", field_value: ["10.0.0.0/8"] }, + { field_name: "mcp_allowed_clients", field_value: ["antigravity-cli"] }, + ]); + + renderSettings(); + await userEvent.click(await screen.findByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", ["antigravity-cli"]), + ); + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges", ["10.0.0.0/8"]); + expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index 8b4d2a58652..18377a4bb82 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -30,8 +30,10 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) const [loading, setLoading] = useState(true); const [saving, setSaving] = useState(false); const [privateRanges, setPrivateRanges] = useState([]); + const [allowedClients, setAllowedClients] = useState([]); const [currentIp, setCurrentIp] = useState(null); const [rangeDraft, setRangeDraft] = useState(""); + const [clientDraft, setClientDraft] = useState(""); useEffect(() => { loadSettings(); @@ -47,6 +49,9 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) if (field.field_name === "mcp_internal_ip_ranges" && field.field_value) { setPrivateRanges(field.field_value); } + if (field.field_name === "mcp_allowed_clients" && field.field_value) { + setAllowedClients(field.field_value); + } } } catch (error) { console.error("Failed to load MCP network settings:", error); @@ -72,6 +77,11 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) } else { await deleteConfigFieldSetting(accessToken, "mcp_internal_ip_ranges"); } + if (allowedClients.length > 0) { + await updateConfigFieldSetting(accessToken, "mcp_allowed_clients", allowedClients); + } else { + await deleteConfigFieldSetting(accessToken, "mcp_allowed_clients"); + } } catch (error) { console.error("Failed to save MCP network settings:", error); } finally { @@ -86,17 +96,28 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) }; // Commas separate entries, matching the old tokenised input. - const commitDraft = () => { - const added = rangeDraft + const splitDraft = (draft: string, existing: string[]) => + draft .split(",") .map((r) => r.trim()) - .filter((r) => r !== "" && !privateRanges.includes(r)); + .filter((r) => r !== "" && !existing.includes(r)); + + const commitDraft = () => { + const added = splitDraft(rangeDraft, privateRanges); if (added.length > 0) { setPrivateRanges([...privateRanges, ...added]); } setRangeDraft(""); }; + const commitClientDraft = () => { + const added = splitDraft(clientDraft, allowedClients); + if (added.length > 0) { + setAllowedClients([...allowedClients, ...added]); + } + setClientDraft(""); + }; + if (loading) { return (
@@ -178,6 +199,56 @@ const MCPNetworkSettings: React.FC = ({ accessToken })

+
+

Allowed Client Applications

+

+ Only the MCP client applications listed here can connect to the gateway. Names are matched exactly against + the clientInfo.name each client sends in its MCP initialize request (for example claude-code or + codex-mcp-client). Leave empty to allow every client. Clients choose the name they send, so treat this as a + policy control rather than a security boundary. +

+
+ + +
+

Allowed Client Names

+
+ {allowedClients.length > 0 && ( +
+ {allowedClients.map((client) => ( + + {client} + + + ))} +
+ )} + setClientDraft(e.target.value)} + onBlur={commitClientDraft} + onKeyDown={(e) => { + if (e.key === "Enter" || e.key === ",") { + e.preventDefault(); + commitClientDraft(); + } + }} + /> +

+ Enter the clientInfo.name values to admit. Any other client, or one that does not identify itself, gets a + 403 on its MCP initialize request. +

+ +
- + + + ))}
)} - setClientDraft(e.target.value)} - onBlur={commitClientDraft} - onKeyDown={(e) => { - if (e.key === "Enter" || e.key === ",") { - e.preventDefault(); - commitClientDraft(); - } - }} - /> +

- Enter the exact JWT claim or header values to admit. Every MCP request from any other client, or from one with - no resolvable identity, gets a 403. + The alias is the name shown here and in gateway logs. The value is the exact JWT claim or header value that + identifies the client, such as the OAuth client ID your identity provider issues. Leave the list empty to + allow every client. Every MCP request from an unlisted client, or from one with no resolvable identity, gets a + 403.

diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 19c0e29f17c..45a11aff0bc 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26849,9 +26849,9 @@ export interface components { maximum_spend_logs_retention_period?: string | null; /** * Mcp Allowed Clients - * @description MCP client applications admitted by the gateway. When set, every MCP request must carry a client identity that matches one of these values exactly: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted. + * @description MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted. */ - mcp_allowed_clients?: string[] | null; + mcp_allowed_clients?: components["schemas"]["MCPAllowedClient"][] | null; /** * Mcp Client Id Header * @description Request header whose value names the calling MCP client application (for example 'x-mcp-client') for callers that did not authenticate with a JWT, used only while mcp_allowed_clients is set. The client picks this value itself, so it is a policy control rather than a security boundary; prefer litellm_jwtauth.mcp_client_id_jwt_field where callers use JWTs. @@ -32491,6 +32491,22 @@ export interface components { */ status?: "healthy" | "unhealthy"; }; + /** + * MCPAllowedClient + * @description One entry of `general_settings.mcp_allowed_clients`. + */ + MCPAllowedClient: { + /** + * Alias + * @description Human-readable name for this client application, shown in the dashboard and in gateway logs. + */ + alias: string; + /** + * Value + * @description Exact value of the JWT claim named in litellm_jwtauth.mcp_client_id_jwt_field, or of the mcp_client_id_header header, that identifies this client application. Matched case-sensitively. + */ + value: string; + }; /** MCPConnectorEntry */ MCPConnectorEntry: { /** Args */ From da603c629ba465b8e81709847506580623450cb7 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 22:37:30 +0000 Subject: [PATCH 48/61] fix(ui): surface a malformed stored MCP allowlist as deny-all and let Save replace or remove it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/MCPNetworkSettings.test.tsx | 29 +++++++++- .../_components/MCPNetworkSettings.tsx | 55 +++++++++++++------ 2 files changed, 65 insertions(+), 19 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx index 92f6f7554d4..d27c18c5ae3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx @@ -154,16 +154,39 @@ describe("MCPNetworkSettings", () => { expect(screen.getByRole("textbox", { name: "Client 2 value" })).toHaveValue("codex-mcp-client"); }); - it("ignores a stored allowlist in the old plain-string shape instead of rendering it", async () => { + it("warns that a stored allowlist in the old plain-string shape denies every client and lets Save remove it", async () => { vi.mocked(getGeneralSettingsCall).mockResolvedValue([ { field_name: "mcp_allowed_clients", field_value: ["antigravity-cli"] }, ]); renderSettings(); - await screen.findByText("Allowed Clients"); + expect(await screen.findByText(/stored allowlist is not a list of alias and value pairs/)).toBeVisible(); expect(screen.queryByRole("textbox", { name: "Client 1 value" })).not.toBeInTheDocument(); - expect(screen.queryByText(/every client is denied/)).not.toBeInTheDocument(); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients")); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + await waitFor(() => expect(screen.queryByText(/stored allowlist is not a list/)).not.toBeInTheDocument()); + }); + + it("replaces a stored allowlist in the old plain-string shape with the clients the admin adds", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: ["antigravity-cli"] }, + ]); + + renderSettings(); + + await screen.findByText(/stored allowlist is not a list of alias and value pairs/); + await addClient(ANTIGRAVITY.alias, ANTIGRAVITY.value); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ANTIGRAVITY]), + ); + expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients"); + await waitFor(() => expect(screen.queryByText(/stored allowlist is not a list/)).not.toBeInTheDocument()); }); it("adds clients as alias and value pairs and saves them under mcp_allowed_clients", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index 2ef3ee8707d..ae1fad36599 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -42,10 +42,20 @@ const isAllowedClient = (entry: unknown): entry is AllowedClient => { return typeof alias === "string" && typeof value === "string"; }; -const parseStoredClients = (fieldValue: unknown): AllowedClient[] | null => - Array.isArray(fieldValue) && fieldValue.every(isAllowedClient) - ? fieldValue.map(({ alias, value }) => ({ alias, value })) - : null; +type StoredAllowlist = + | { readonly kind: "absent" } + | { readonly kind: "clients"; readonly clients: AllowedClient[] } + | { readonly kind: "malformed" }; + +const ABSENT: StoredAllowlist = { kind: "absent" }; + +const parseStoredClients = (fieldValue: unknown): StoredAllowlist => { + if (fieldValue === null || fieldValue === undefined) return ABSENT; + if (Array.isArray(fieldValue) && fieldValue.every(isAllowedClient)) { + return { kind: "clients", clients: fieldValue.map(({ alias, value }) => ({ alias, value })) }; + } + return { kind: "malformed" }; +}; let nextRowKey = 0; const newRow = (client: AllowedClient = { alias: "", value: "" }): AllowedClientRow => ({ @@ -66,8 +76,16 @@ const sameClients = (a: AllowedClient[], b: AllowedClient[]) => const unchangedSinceLoad = (value: string[], stored: string[] | null) => stored === null ? value.length === 0 : value.length > 0 && sameList(value, stored); -const clientsUnchangedSinceLoad = (value: AllowedClient[], stored: AllowedClient[] | null) => - stored === null ? value.length === 0 : value.length > 0 && sameClients(value, stored); +const clientsUnchangedSinceLoad = (value: AllowedClient[], stored: StoredAllowlist) => { + switch (stored.kind) { + case "absent": + return value.length === 0; + case "clients": + return value.length > 0 && sameClients(value, stored.clients); + case "malformed": + return false; + } +}; const headerUnchangedSinceLoad = (value: string, stored: string | null) => stored === null ? value === "" : value !== "" && value === stored; @@ -79,7 +97,7 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) const [allowedClients, setAllowedClients] = useState([]); const [clientIdHeader, setClientIdHeader] = useState(""); const [storedRanges, setStoredRanges] = useState(null); - const [storedClients, setStoredClients] = useState(null); + const [storedClients, setStoredClients] = useState(ABSENT); const [storedClientIdHeader, setStoredClientIdHeader] = useState(null); const [currentIp, setCurrentIp] = useState(null); const [rangeDraft, setRangeDraft] = useState(""); @@ -100,11 +118,9 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) setStoredRanges(field.field_value); } if (field.field_name === "mcp_allowed_clients") { - const clients = parseStoredClients(field.field_value); - if (clients !== null) { - setAllowedClients(clients.map(newRow)); - setStoredClients(clients); - } + const stored = parseStoredClients(field.field_value); + setAllowedClients(stored.kind === "clients" ? stored.clients.map(newRow) : []); + setStoredClients(stored); } if (field.field_name === "mcp_client_id_header" && typeof field.field_value === "string") { setClientIdHeader(field.field_value); @@ -145,11 +161,11 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) if (clientsUnchangedSinceLoad(clients, storedClients)) return; if (clients.length > 0) { await updateConfigFieldSetting(token, "mcp_allowed_clients", clients); - setStoredClients(clients); + setStoredClients({ kind: "clients", clients }); return; } await deleteConfigFieldSetting(token, "mcp_allowed_clients"); - setStoredClients(null); + setStoredClients(ABSENT); }; const persistClientIdHeader = async (token: string) => { @@ -216,7 +232,8 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) } const suggestedRange = currentIp ? ipToSlash24(currentIp) : null; - const storedAllowlistDeniesEveryone = storedClients !== null && storedClients.length === 0; + const storedAllowlistIsMalformed = storedClients.kind === "malformed"; + const storedAllowlistIsEmpty = storedClients.kind === "clients" && storedClients.clients.length === 0; return (
@@ -303,7 +320,13 @@ const MCPNetworkSettings: React.FC = ({ accessToken })

Allowed Clients

- {storedAllowlistDeniesEveryone && ( + {storedAllowlistIsMalformed && ( +

+ The stored allowlist is not a list of alias and value pairs, so every client is denied. Add the clients you + want and save to replace it, or save with the list empty to remove it and allow every client again. +

+ )} + {storedAllowlistIsEmpty && (

An empty allowlist is currently stored, so every client is denied. Save with the list empty to remove it and allow every client again. From 66c01cf35c14d06877f7560af7604a5c8e36154c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:39:02 -0700 Subject: [PATCH 49/61] refactor(responses): map finish reasons to incomplete_details through a lookup table The match statement in _incomplete_details_for_finish_reason tripped CodeQL's mixed explicit and implicit returns alert (code-scanning 12640). A module-level MappingProxyType keyed by finish reason gives the same three mappings with one explicit return path --- .../transformation.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 961fffd3d42..7337beba45c 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -112,6 +112,9 @@ ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"}) +_INCOMPLETE_REASON_BY_FINISH_REASON: Final[Mapping[str, Literal["max_output_tokens", "content_filter"]]] = ( + MappingProxyType({"length": "max_output_tokens", "content_filter": "content_filter", "refusal": "content_filter"}) +) @dataclass(frozen=True, slots=True) @@ -2303,13 +2306,10 @@ class LiteLLMCompletionResponsesConfig: ) -> IncompleteDetails | None: if existing is not None: return existing - match finish_reason: - case "length": - return IncompleteDetails(reason="max_output_tokens") - case "content_filter" | "refusal": - return IncompleteDetails(reason="content_filter") - case _: - return None + if finish_reason is None: + return None + reason: Final = _INCOMPLETE_REASON_BY_FINISH_REASON.get(finish_reason) + return IncompleteDetails(reason=reason) if reason is not None else None @staticmethod def _tool_call_id_from_responses_item(item_id: str | None, call_id: str | None) -> str: From c181c927d0b7a4ad214a4dd160c2f16f1373385e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:58:15 -0700 Subject: [PATCH 50/61] fix(proxy): record response.failed frames in background polling --- .../response_polling/background_streaming.py | 7 +++- .../test_response_polling_handler.py | 38 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index fac45d4391c..b13042dfb6c 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -75,6 +75,10 @@ class _StreamEventParser: parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) +def _sse_frame_data(frame: str) -> str | None: + return next((line[6:].strip() for line in frame.splitlines() if line.startswith("data: ")), None) + + async def _never_receive() -> Message: await asyncio.Event().wait() raise AssertionError("unreachable") @@ -224,8 +228,7 @@ async def background_streaming_task( if isinstance(chunk, bytes): chunk = chunk.decode("utf-8") - if isinstance(chunk, str) and chunk.startswith("data: "): - chunk_data = chunk[6:].strip() + if isinstance(chunk, str) and (chunk_data := _sse_frame_data(chunk)) is not None: if chunk_data == "[DONE]": break diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index 467c1332325..81ca0114a8d 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -1482,6 +1482,44 @@ class TestBackgroundStreamingTerminalEvents: assert final_call.kwargs["status"] == "failed" assert final_call.kwargs["error"] == error_payload + @pytest.mark.asyncio + async def test_named_event_failed_frame_sets_failed_status_and_error(self): + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + error_payload = { + "code": "cyber_policy", + "message": "Your request was flagged for possible cybersecurity risk and was not completed", + } + failed_event = { + "type": "response.failed", + "sequence_number": 5, + "response": {"id": "resp_123", "status": "failed", "error": error_payload, "output": []}, + } + + async def _body_iterator(): + yield b'data: {"type": "response.in_progress"}\n\n' + yield f"event: response.failed\ndata: {json.dumps(failed_event)}\n\n".encode() + yield b"data: [DONE]\n\n" + + mock_response = Mock() + mock_response.body_iterator = _body_iterator() + handler = AsyncMock(spec=ResponsePollingHandler) + kwargs = _make_background_streaming_kwargs("poll_named_event", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock( + return_value=mock_response + ) + await background_streaming_task(**kwargs) + + final_call = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "failed" + assert final_call.kwargs["error"] == error_payload + @pytest.mark.asyncio async def test_response_incomplete_sets_incomplete_status_and_details(self): """Test that a response.incomplete stream event results in incomplete status""" From b3cf45e9f232c094e2f2b2a6bf59f609464bf742 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:58:24 -0700 Subject: [PATCH 51/61] fix(proxy): drop daily spend batches that cannot be re-sent safely instead of requeueing them --- litellm/proxy/db/db_spend_update_writer.py | 24 ++++++++- litellm/proxy/db/exception_handler.py | 18 +++++++ .../proxy/db/test_db_spend_update_writer.py | 51 ++++++++++++++++++- .../proxy/db/test_exception_handler.py | 22 ++++++++ 4 files changed, 112 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index b2e6f9dc54d..e9967fb0d67 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -31,6 +31,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( + DB_CONNECTION_ERROR_TYPES, DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, @@ -64,6 +65,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, ) +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, @@ -157,6 +159,16 @@ class _DailySpendCommit(Protocol[_DailySpendTransactionT]): ) -> None: ... +_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) + + +def _daily_spend_commit_failure_is_requeue_safe(e: Exception) -> bool: + if isinstance(e, DB_CONNECTION_ERROR_TYPES): + return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) + sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) + return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES + + def _timed_request_duration_ms( payload: dict | SpendLogsPayload, request_status: Literal["success", "failure"], @@ -1319,7 +1331,17 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions), ) - except Exception as e: # noqa: BLE001 # the uncommitted rows go back on the queue; the other tables must still flush + except Exception as e: # noqa: BLE001 # whatever failed here, the other tables must still flush + if not _daily_spend_commit_failure_is_requeue_safe(e): + spend_log_error( + "Spend tracking - dropped %d daily %s spend rows: the failed commit may have applied " + "or the database refused the data, so re-sending it is not safe. Error: %s", + len(transactions), + entity_type, + str(e), + exc=e, + ) + return spend_log_error( "Spend tracking - failed to commit daily %s spend updates. " "Re-queued %d rows for retry on next tick. Error: %s", diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 2cee5128c66..460bf5db3b1 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,6 +1,8 @@ from collections.abc import Awaitable, Callable, Iterator from typing import Any, Final, TypeVar +from pydantic import TypeAdapter, ValidationError + from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, @@ -17,6 +19,8 @@ _TRANSIENT_DB_UNAVAILABLE_MESSAGE: Final = ( "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly." ) +_DATABASE_ERROR_META: Final = TypeAdapter(dict[str, object]) + def _exception_chain(e: BaseException) -> Iterator[BaseException]: current = e # rebind-ok: advances one link per iteration of the bounded walk @@ -221,6 +225,20 @@ class PrismaDBExceptionHandler: or "write conflict or a deadlock" in error_message ) + @staticmethod + def postgres_sqlstate(e: Exception) -> str | None: + """The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None.""" + import prisma + + if not isinstance(e, _exception_types(prisma.errors.DataError)): + return None + try: + meta: Final = _DATABASE_ERROR_META.validate_python(getattr(e, "meta", None)) + except ValidationError: + return None + code: Final = meta.get("code") + return code if isinstance(code, str) else None + @staticmethod def is_read_only_transaction_error(e: Exception) -> bool: """True iff ``e`` is Postgres SQLSTATE 25006 surfaced through prisma: the diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index b27b838133b..155bca656d5 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -11,7 +11,9 @@ from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, call, patch +import httpx import pytest +from prisma.errors import RawQueryError from redis.exceptions import DataError import litellm @@ -2812,14 +2814,15 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ class _DailySpendFakeDB(_WindowSpendFakeDB): """Records the daily rollup upserts it is handed and fails the ones aimed at one table.""" - def __init__(self, failing_table: str | None) -> None: + def __init__(self, failing_table: str | None, failure: Exception | None = None) -> None: super().__init__() self.failing_table = failing_table + self.failure = failure self.execute_raw_calls: list[Statement] = [] async def execute_raw(self, query: str, *args: object) -> int: if self.failing_table is not None and self.failing_table in query: - raise Exception("connection reset") + raise self.failure if self.failure is not None else Exception("connection reset") self.execute_raw_calls.append((query, args)) return len(args) @@ -2828,6 +2831,50 @@ def _daily_upserts(db: _DailySpendFakeDB, table: str) -> list[Statement]: return [statement for statement in db.execute_raw_calls if table in statement[0]] +def _postgres_rejection(sqlstate: str) -> RawQueryError: + return RawQueryError( + data={"user_facing_error": {"error_code": "P2010", "meta": {"code": sqlstate, "message": "db error"}}} + ) + + +@pytest.mark.parametrize( + ("failure", "lands_on_the_next_tick"), + [ + pytest.param(httpx.ReadTimeout("no reply"), False, id="reply lost after the statement was sent"), + pytest.param(httpx.ConnectError("refused"), True, id="statement never reached the database"), + pytest.param(_postgres_rejection("22021"), False, id="postgres refused the data itself"), + pytest.param(_postgres_rejection("23502"), False, id="postgres refused a constraint violation"), + pytest.param(_postgres_rejection("42P01"), True, id="table missing"), + pytest.param(_postgres_rejection("57014"), True, id="statement cancelled"), + ], +) +@pytest.mark.asyncio +async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_provably_uncommitted( + failure: Exception, lands_on_the_next_tick: bool +): + """A lost reply means the statement may already have applied, and re-sending it stacks a + second increment into the same transaction (LIT-4823); a row Postgres refuses would fail + every tick forever. Both are dropped loudly. Every other failure left nothing committed, + so its rows go back on the queue and land on the next tick.""" + db_writer = DBSpendUpdateWriter() + await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")}) + db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=failure) + db_writer._flush_tool_discovery_queue = AsyncMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + db.failing_table = None + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + + assert len(_daily_upserts(db, "LiteLLM_DailyUserSpend")) == (1 if lands_on_the_next_tick else 0) + assert db_writer.daily_spend_update_queue.update_queue.empty() + + @pytest.mark.asyncio async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other_tables(): """With the Redis buffer off, a daily batch that failed to commit was discarded along diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 3f009137a1c..26ac1ea65ad 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -665,6 +665,28 @@ def test_is_deadlock_error_excludes_non_deadlocks(error): assert PrismaDBExceptionHandler.is_deadlock_error(error) is False +@pytest.mark.parametrize( + ("error", "sqlstate"), + [ + ( + RawQueryError( + data={"user_facing_error": {"error_code": "P2010", "meta": {"code": "22021", "message": "m"}}} + ), + "22021", + ), + (RawQueryError(data={"user_facing_error": {"error_code": "P2010", "meta": {"message": "m"}}}), None), + (RawQueryError(data={"user_facing_error": {"error_code": "P2010", "meta": {"code": 42, "message": "m"}}}), None), + (prisma_errors.DataError(data={"user_facing_error": {"meta": None}}), None), + (PrismaError("db error"), None), + (httpx.ReadTimeout("no reply"), None), + ], +) +def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error, sqlstate): + """Only a prisma data error carrying Postgres's own error code yields a SQLSTATE; a + codeless or malformed payload, an engine-level error, and a transport error yield None.""" + assert PrismaDBExceptionHandler.postgres_sqlstate(error) == sqlstate + + READ_ONLY_CONNECTOR_ERROR: Final = ( "Error occurred during query execution:\nConnectorError(ConnectorError { user_facing_error: None, " 'kind: QueryError(PostgresError { code: "25006", message: "cannot execute UPDATE in a read-only transaction", ' From 8005856411ae43091c86ba2632d389916b8fa0ec Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:03:27 -0700 Subject: [PATCH 52/61] ci(build_and_test): seed the routing strategy through /config/update --- .circleci/config.yml | 6 ++++++ proxy_server_config.yaml | 1 - 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index df17a9e4402..87f1ee604cf 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1785,6 +1785,12 @@ jobs: - wait_for_service: url: http://localhost:4000 timeout: "300" + - run: + name: Seed the routing strategy through /config/update + command: | + curl --noproxy '*' -sSf -X POST http://localhost:4000/config/update \ + -H 'Authorization: Bearer sk-1234' -H 'Content-Type: application/json' \ + -d '{"router_settings": {"routing_strategy": "usage-based-routing-v2"}}' - run: name: Run tests command: | diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml index 73990153227..703d56bc0cd 100644 --- a/proxy_server_config.yaml +++ b/proxy_server_config.yaml @@ -213,7 +213,6 @@ files_settings: api_key: os.environ/OPENAI_API_KEY router_settings: - routing_strategy: usage-based-routing-v2 redis_host: os.environ/REDIS_HOST redis_password: os.environ/REDIS_PASSWORD redis_port: os.environ/REDIS_PORT From a4624b6c6c4bb1ddf753bd57bbfd37b711598c90 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:03:45 -0700 Subject: [PATCH 53/61] fix(gemini): derive the finish reason key set from the Candidates type Candidates.finishReason listed eleven values while the mapping key set carried twenty-one, so typed fixtures could not spell the reasons this PR handles. GeminiFinishReason is now the one list, the key set derives from it, and a test checks every documented reason has an explicit mapping instead of falling through to "stop" --- .../vertex_and_google_ai_studio_gemini.py | 29 ++------------ litellm/types/llms/vertex_ai.py | 39 ++++++++++++------- ...test_vertex_and_google_ai_studio_gemini.py | 9 ++++- 3 files changed, 36 insertions(+), 41 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a95b845718a..46f1b948026 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -6,7 +6,7 @@ import time from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args import httpx @@ -57,6 +57,7 @@ from litellm.types.llms.vertex_ai import ( ContentType, FunctionCallingConfig, FunctionDeclaration, + GeminiFinishReason, GeminiThinkingConfig, GenerateContentResponseBody, HttpxPartType, @@ -1330,31 +1331,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "IMAGE_PROHIBITED_CONTENT": "The token generation was stopped as the response was flagged for prohibited image content.", } - _GEMINI_FINISH_REASON_KEYS = frozenset( - { - "STOP", - "MAX_TOKENS", - "SAFETY", - "RECITATION", - "FINISH_REASON_UNSPECIFIED", - "MALFORMED_FUNCTION_CALL", - "LANGUAGE", - "OTHER", - "BLOCKLIST", - "PROHIBITED_CONTENT", - "SPII", - "IMAGE_SAFETY", - "IMAGE_PROHIBITED_CONTENT", - "TOO_MANY_TOOL_CALLS", - "MALFORMED_RESPONSE", - "NO_IMAGE", - "IMAGE_RECITATION", - "IMAGE_OTHER", - "ESCALATION", - "UNEXPECTED_TOOL_CALL", - "MISSING_THOUGHT_SIGNATURE", - } - ) + _GEMINI_FINISH_REASON_KEYS: Final[frozenset[str]] = frozenset(get_args(GeminiFinishReason)) @staticmethod def get_finish_reason_mapping() -> dict[str, OpenAIChatCompletionFinishReason]: diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 3b95b786631..ce51e46ef15 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -425,22 +425,35 @@ class UrlContextMetadata(TypedDict, total=False): urlMetadata: list[UrlMetadata] +GeminiFinishReason = Literal[ + "FINISH_REASON_UNSPECIFIED", + "STOP", + "MAX_TOKENS", + "SAFETY", + "RECITATION", + "LANGUAGE", + "OTHER", + "BLOCKLIST", + "PROHIBITED_CONTENT", + "SPII", + "MALFORMED_FUNCTION_CALL", + "IMAGE_SAFETY", + "IMAGE_PROHIBITED_CONTENT", + "TOO_MANY_TOOL_CALLS", + "MALFORMED_RESPONSE", + "NO_IMAGE", + "IMAGE_RECITATION", + "IMAGE_OTHER", + "ESCALATION", + "UNEXPECTED_TOOL_CALL", + "MISSING_THOUGHT_SIGNATURE", +] + + class Candidates(TypedDict, total=False): index: int content: HttpxContentType - finishReason: Literal[ - "FINISH_REASON_UNSPECIFIED", - "STOP", - "MAX_TOKENS", - "SAFETY", - "RECITATION", - "OTHER", - "BLOCKLIST", - "PROHIBITED_CONTENT", - "SPII", - "MALFORMED_FUNCTION_CALL", - "IMAGE_SAFETY", - ] + finishReason: GeminiFinishReason safetyRatings: list[SafetyRatings] citationMetadata: CitationMetadata groundingMetadata: GroundingMetadata diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 4ee199bd9af..6c818016c87 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2,7 +2,7 @@ import asyncio import json import re from copy import deepcopy -from typing import Final, List, cast +from typing import Final, List, cast, get_args from unittest.mock import MagicMock, patch import httpx @@ -18,7 +18,7 @@ from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) -from litellm.types.llms.vertex_ai import UsageMetadata +from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata from litellm.types.utils import ChoiceLogprobs, Usage from litellm.utils import CustomStreamWrapper @@ -940,6 +940,11 @@ def test_check_finish_reason(): ) +def test_every_documented_gemini_finish_reason_has_an_explicit_mapping(): + documented: Final = frozenset(get_args(GeminiFinishReason)) + assert set(VertexGeminiConfig.get_finish_reason_mapping()) == documented + + def test_finish_reason_unspecified_and_malformed_function_call(): """ Test that FINISH_REASON_UNSPECIFIED and MALFORMED_FUNCTION_CALL From 5db2a0c8850385b9e56b17bcd8053bab4fac0b59 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:05:42 -0700 Subject: [PATCH 54/61] test(proxy): type the sqlstate test parameters --- tests/test_litellm/proxy/db/test_exception_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 26ac1ea65ad..f7cc5e3ed83 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -681,7 +681,7 @@ def test_is_deadlock_error_excludes_non_deadlocks(error): (httpx.ReadTimeout("no reply"), None), ], ) -def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error, sqlstate): +def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error: Exception, sqlstate: str | None): """Only a prisma data error carrying Postgres's own error code yields a SQLSTATE; a codeless or malformed payload, an engine-level error, and a transport error yield None.""" assert PrismaDBExceptionHandler.postgres_sqlstate(error) == sqlstate From 1704aeebb49eea499a863a2dfbd9c990b0e2b501 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 23:06:47 +0000 Subject: [PATCH 55/61] fix(enterprise): resolve openai_moderations model at call time and default to omni-moderation-latest Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../enterprise_hooks/openai_moderation.py | 9 +++-- litellm/constants.py | 2 ++ .../guardrails/test_guardrail_coverage.py | 36 +++++++++++++++++++ 3 files changed, 42 insertions(+), 5 deletions(-) diff --git a/enterprise/enterprise_hooks/openai_moderation.py b/enterprise/enterprise_hooks/openai_moderation.py index 2162370804a..017f51bfabd 100644 --- a/enterprise/enterprise_hooks/openai_moderation.py +++ b/enterprise/enterprise_hooks/openai_moderation.py @@ -17,6 +17,7 @@ from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_OPENAI_MODERATIONS_MODEL from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import iter_message_text @@ -24,11 +25,9 @@ from litellm.types.utils import CallTypesLiteral class _ENTERPRISE_OpenAI_Moderation(CustomLogger): - def __init__(self): - self.model_name = ( - litellm.openai_moderations_model_name or "text-moderation-latest" - ) # pass the model_name you initialized on litellm.Router() - pass + @property + def model_name(self) -> str: + return litellm.openai_moderations_model_name or DEFAULT_OPENAI_MODERATIONS_MODEL #### CALL HOOKS - proxy only #### diff --git a/litellm/constants.py b/litellm/constants.py index e7cb21a3a7d..a7d4eba0f15 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -158,6 +158,8 @@ DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL: Final = str( ) DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD", 0.75)) +DEFAULT_OPENAI_MODERATIONS_MODEL: Final = "omni-moderation-latest" + # MCP OAuth2 Client Credentials Defaults MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60")) MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200")) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py index f25e83b1672..548677c70bc 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py @@ -18,7 +18,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from httpx import Request, Response +import litellm from litellm import DualCache +from litellm.constants import DEFAULT_OPENAI_MODERATIONS_MODEL from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import Choices, Message, ModelResponse @@ -764,6 +766,40 @@ async def test_openai_moderation_inspects_multimodal_content(monkeypatch, user_a assert seen_inputs == ["alpha beta"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("configured_after_init", "expected_model"), + [("omni-moderation-2024-09-26", "omni-moderation-2024-09-26"), (None, DEFAULT_OPENAI_MODERATIONS_MODEL)], +) +async def test_openai_moderation_reads_model_name_at_call_time( + monkeypatch, user_api_key, configured_after_init, expected_model +): + """``litellm_settings`` applies ``callbacks`` and ``openai_moderations_model_name`` in YAML + order, so the hook must resolve the model when it runs, not when it is constructed.""" + from enterprise.enterprise_hooks.openai_moderation import ( + _ENTERPRISE_OpenAI_Moderation, + ) + + monkeypatch.setattr(litellm, "openai_moderations_model_name", None) + guard = _ENTERPRISE_OpenAI_Moderation() + monkeypatch.setattr(litellm, "openai_moderations_model_name", configured_after_init) + + class FakeModeration: + results = [type("R", (), {"flagged": False})()] + + fake_router = MagicMock() + fake_router.amoderation = AsyncMock(return_value=FakeModeration()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router, raising=False) + + await guard.async_moderation_hook( + data={"messages": [{"role": "user", "content": "hello"}]}, + user_api_key_dict=user_api_key, + call_type="acompletion", + ) + + fake_router.amoderation.assert_awaited_once_with(model=expected_model, input="hello") + + # ── Google Text Moderation ──────────────────────────────────────────────────── From c6c8aed3f8594f89a86b91df553b84f6aab2fb20 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:22:04 -0700 Subject: [PATCH 56/61] fix(proxy): drop only the daily spend batch whose failure cannot be re-sent, requeue the unsent ones --- litellm/proxy/db/db_spend_update_writer.py | 30 ++++++----- .../proxy/db/test_db_spend_update_writer.py | 54 +++++++++++++++++++ 2 files changed, 71 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e9967fb0d67..37eac8604bd 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1332,15 +1332,7 @@ class DBSpendUpdateWriter: daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions), ) except Exception as e: # noqa: BLE001 # whatever failed here, the other tables must still flush - if not _daily_spend_commit_failure_is_requeue_safe(e): - spend_log_error( - "Spend tracking - dropped %d daily %s spend rows: the failed commit may have applied " - "or the database refused the data, so re-sending it is not safe. Error: %s", - len(transactions), - entity_type, - str(e), - exc=e, - ) + if not transactions: return spend_log_error( "Spend tracking - failed to commit daily %s spend updates. " @@ -2050,13 +2042,25 @@ class DBSpendUpdateWriter: sql, params = build_bulk_upsert(table=table, batch=merged_batch) await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: - # Log detailed error information for debugging batch upsert failures - # This helps diagnose issues like unique constraint violations + if _daily_spend_commit_failure_is_requeue_safe(batch_error): + spend_log_error( + "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", + entity_type, + table.name, + len(transactions_to_process), + str(batch_error), + exc=batch_error, + ) + raise + for key in transactions_to_process: + daily_spend_transactions.pop(key, None) spend_log_error( - "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", + "Spend tracking - dropped %d daily %s spend rows: the failed statement may have " + "applied or the database refused the data, so re-sending it is not safe. " + "Table: %s, Error: %s", + len(transactions_to_process), entity_type, table.name, - len(transactions_to_process), str(batch_error), exc=batch_error, ) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 155bca656d5..20f1fa9d363 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1624,6 +1624,33 @@ async def test_update_daily_spend_keeps_failed_transactions_for_retry(): assert daily_spend_transactions == expected +@pytest.mark.asyncio +async def test_update_daily_spend_drops_the_batch_whose_failure_cannot_be_resent(): + """A reply lost after the statement was sent may already have applied, so the batch is + taken out of the caller's dict before the error propagates: whichever requeue the caller + runs afterwards, the Redis restore included, cannot send it a second time.""" + + def lose_the_reply() -> int: + raise httpx.ReadTimeout("no reply") + + prisma_client = _RecordingPrisma(execute_raw=lose_the_reply) + daily_spend_transactions = {"user-key": _daily_txn(user_id="user-1")} + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + with pytest.raises(httpx.ReadTimeout): + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_spend_transactions, + entity_type="user", + entity_id_field="user_id", + ) + + assert daily_spend_transactions == {} + + @pytest.mark.asyncio async def test_commit_key_spend_updates_includes_last_active(): """ @@ -2875,6 +2902,33 @@ async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_prov assert db_writer.daily_spend_update_queue.update_queue.empty() +@pytest.mark.asyncio +async def test_failed_daily_spend_commit_drops_only_the_batch_that_was_sent(): + """A tick holding more than one batch of 100 rows sends them one statement at a time, and + a reply lost on one statement says nothing about the batches after it: only the batch that + was on the wire is dropped, the ones never sent go back on the queue and land next tick.""" + db_writer = DBSpendUpdateWriter() + await db_writer.daily_spend_update_queue.add_update( + {f"user-{i:03d}": _daily_txn(user_id=f"user-{i:03d}") for i in range(150)} + ) + db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=httpx.ReadTimeout("no reply")) + db_writer._flush_tool_discovery_queue = AsyncMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + db.failing_table = None + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + + (upsert,) = _daily_upserts(db, "LiteLLM_DailyUserSpend") + assert _row_values(upsert, "user_id") == [f"user-{i:03d}" for i in range(100, 150)] + assert db_writer.daily_spend_update_queue.update_queue.empty() + + @pytest.mark.asyncio async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other_tables(): """With the Redis buffer off, a daily batch that failed to commit was discarded along From 83d89aa134bbeac391ce4a49dd63223f52b97019 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:08:54 +0000 Subject: [PATCH 57/61] test(integration): cover off-peak pricing on a live proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/client.py | 4 +- tests/integration/contracts.json | 6 ++ .../pricing/test_off_peak_pricing.py | 82 +++++++++++++++++++ 3 files changed, 90 insertions(+), 2 deletions(-) create mode 100644 tests/integration/pricing/test_off_peak_pricing.py diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 97522e5728c..9f1118ab1e3 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -161,7 +161,7 @@ class Scenario: assert all(object_value(object_value(entry)["model_info"])["id"] != identity for entry in entries) assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [] - def model(self, **parameters: JsonValue) -> str: + def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str: name: Final = f"integration-{uuid.uuid4().hex}" created: Final = self.gateway.post( "/model/new", @@ -173,7 +173,7 @@ class Scenario: "api_base": f"{self.gateway.upstream_url}/v1", **parameters, }, - "model_info": {}, + "model_info": dict(model_info) if model_info is not None else {}, }, ) identity: Final = string_value(object_value(created["model_info"])["id"]) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 91b1bd86954..6958ade50f7 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -92,6 +92,12 @@ "tests/integration/pricing/test_price_precedence.py::test_same_upstream_aliases_keep_distinct_prices_after_reload": [ "quota_management.spend_tracking.alias_prices.remain_independent_on_reload" ], + "tests/integration/pricing/test_off_peak_pricing.py::test_open_off_peak_window_bills_off_peak_rates": [ + "quota_management.spend_tracking.off_peak_pricing.open_window_bills_off_peak_rates" + ], + "tests/integration/pricing/test_off_peak_pricing.py::test_closed_off_peak_window_bills_standard_rates": [ + "quota_management.spend_tracking.off_peak_pricing.closed_window_bills_standard_rates" + ], "tests/integration/spend/test_cache_and_quota.py::test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost": [ "quota_management.response_cache.generated_sequences_preserve_content_and_accounting" ], diff --git a/tests/integration/pricing/test_off_peak_pricing.py b/tests/integration/pricing/test_off_peak_pricing.py new file mode 100644 index 00000000000..5623356c078 --- /dev/null +++ b/tests/integration/pricing/test_off_peak_pricing.py @@ -0,0 +1,82 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest +from pydantic import JsonValue + +from tests.integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from tests.integration._support.database import read_rows + +STANDARD_INPUT_RATE: Final = 0.001 +STANDARD_OUTPUT_RATE: Final = 0.002 +OFF_PEAK_INPUT_RATE: Final = 0.0001 +OFF_PEAK_OUTPUT_RATE: Final = 0.0002 + + +def off_peak_window(start_offset_hours: int, end_offset_hours: int) -> Mapping[str, JsonValue]: + now: Final = datetime.now(timezone.utc) + start: Final = now + timedelta(hours=start_offset_hours) + end: Final = now + timedelta(hours=end_offset_hours) + return { + "hours_utc": f"{start:%H:%M}-{end:%H:%M}", + "input_cost_per_token": OFF_PEAK_INPUT_RATE, + "output_cost_per_token": OFF_PEAK_OUTPUT_RATE, + } + + +def billed_model(scenario: Scenario, off_peak: Mapping[str, JsonValue]) -> str: + return scenario.model( + input_cost_per_token=STANDARD_INPUT_RATE, + output_cost_per_token=STANDARD_OUTPUT_RATE, + model_info={"off_peak_pricing": dict(off_peak)}, + ) + + +def assert_chat_bills_rates(gateway: Gateway, model: str, input_rate: float, output_rate: float) -> None: + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "off peak control"}]} + ) + assert response.status_code == 200, response.text + expected: Final = 20 * input_rate + 20 * output_rate + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6) + request_id: Final = string_value(object_value(response.json())["id"]) + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id = %s', + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["prompt_tokens"] == 20 + assert rows[0]["completion_tokens"] == 20 + assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6) + metadata: Final = rows[0]["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + breakdown: Final = object_value(parsed["cost_breakdown"]) + assert float(breakdown["input_cost"]) == pytest.approx(20 * input_rate, rel=1e-6) + assert float(breakdown["output_cost"]) == pytest.approx(20 * output_rate, rel=1e-6) + + +@pytest.mark.covers("quota_management.spend_tracking.off_peak_pricing.open_window_bills_off_peak_rates") +def test_open_off_peak_window_bills_off_peak_rates(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = billed_model(scenario, off_peak_window(-1, 1)) + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + matching: Final = tuple(object_value(entry) for entry in entries if object_value(entry)["model_name"] == model) + assert len(matching) == 1 + info: Final = object_value(matching[0]["model_info"]) + off_peak: Final = object_value(info["off_peak_pricing"]) + assert off_peak["input_cost_per_token"] == OFF_PEAK_INPUT_RATE + assert off_peak["output_cost_per_token"] == OFF_PEAK_OUTPUT_RATE + assert_chat_bills_rates(gateway, model, OFF_PEAK_INPUT_RATE, OFF_PEAK_OUTPUT_RATE) + + +@pytest.mark.covers("quota_management.spend_tracking.off_peak_pricing.closed_window_bills_standard_rates") +def test_closed_off_peak_window_bills_standard_rates(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = billed_model(scenario, off_peak_window(2, 3)) + assert_chat_bills_rates(gateway, model, STANDARD_INPUT_RATE, STANDARD_OUTPUT_RATE) From 12120fe59bd9dd36486fa683f84b06fb91bd9c9f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:11:29 -0700 Subject: [PATCH 58/61] refactor(bedrock): inline maxTokens clamp and cover inference-profile ARNs in tests --- litellm/llms/bedrock/chat/converse_transformation.py | 10 ++-------- .../llms/bedrock/chat/test_converse_transformation.py | 3 +++ 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 801ec571376..176819c0dab 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -382,12 +382,6 @@ class AmazonConverseConfig(BaseConfig): def _requires_min_max_tokens(model: str) -> bool: return re.search(r"openai\.gpt-\d|xai\.grok-", model) is not None - @staticmethod - def _enforce_min_max_tokens(max_tokens: object) -> object: - if isinstance(max_tokens, int) and max_tokens < BEDROCK_OPENAI_COMPAT_MIN_MAX_TOKENS: - return BEDROCK_OPENAI_COMPAT_MIN_MAX_TOKENS - return max_tokens - def _is_nova_2_model(self, model: str) -> bool: """ Check if the model is a Nova 2 model that supports reasoningConfig. @@ -1011,8 +1005,8 @@ class AmazonConverseConfig(BaseConfig): ) if param == "max_tokens" or param == "max_completion_tokens": optional_params["maxTokens"] = ( - self._enforce_min_max_tokens(value) - if self._requires_min_max_tokens(model) and isinstance(value, int) + max(value, BEDROCK_OPENAI_COMPAT_MIN_MAX_TOKENS) + if isinstance(value, int) and self._requires_min_max_tokens(model) else value ) if param == "stream": diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 9d8bf786829..086a7e59f56 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -467,6 +467,9 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto(): ("global.xai.grok-4.6", "max_completion_tokens", 1, 16), ("us.xai.grok-4.6", "max_tokens", 32, 32), ("anthropic.claude-sonnet-4-5-20250929-v1:0", "max_tokens", 1, 1), + ("arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.openai.gpt-6-astra", "max_tokens", 1, 16), + ("arn:aws:bedrock:us-east-1:123456789012:inference-profile/global.xai.grok-4.6", "max_tokens", 1, 16), + ("arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz", "max_tokens", 1, 1), ], ) def test_map_openai_params_enforces_minimum_max_tokens_for_openai_compat_models( From c32309fb2de108768fa8704ee0696b29d383733c Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 19 Sep 2026 00:13:44 +0000 Subject: [PATCH 59/61] feat(ui): show MCP allowed clients as cards edited in a dialog Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/MCPNetworkSettings.test.tsx | 114 +++++++++--- .../_components/MCPNetworkSettings.tsx | 167 ++++++++++++------ 2 files changed, 201 insertions(+), 80 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx index d27c18c5ae3..4cc87f1455c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import MCPNetworkSettings from "./MCPNetworkSettings"; @@ -26,12 +26,20 @@ const renderSettings = () => render(); const ANTIGRAVITY = { alias: "Antigravity CLI", value: "antigravity-cli" }; const CODEX = { alias: "Codex", value: "codex-mcp-client" }; +const clientCard = (alias: string) => screen.getByRole("button", { name: new RegExp(`^${alias}`) }); + +const fillClientDialog = async (alias: string, value: string) => { + const dialog = await screen.findByRole("dialog"); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Alias" }), { target: { value: alias } }); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Value" }), { target: { value } }); + return dialog; +}; + const addClient = async (alias: string, value: string) => { await userEvent.click(screen.getByRole("button", { name: "Add client" })); - const aliases = screen.getAllByRole("textbox", { name: /^Client \d+ alias$/ }); - const values = screen.getAllByRole("textbox", { name: /^Client \d+ value$/ }); - fireEvent.change(aliases[aliases.length - 1], { target: { value: alias } }); - fireEvent.change(values[values.length - 1], { target: { value } }); + const dialog = await fillClientDialog(alias, value); + await userEvent.click(within(dialog).getByRole("button", { name: "Add" })); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); }; describe("MCPNetworkSettings", () => { @@ -139,7 +147,7 @@ describe("MCPNetworkSettings", () => { expect(updateConfigFieldSetting).not.toHaveBeenCalled(); }); - it("labels the section Allowed Clients and renders each stored client as an alias and value row", async () => { + it("labels the section Allowed Clients and renders each stored client as a card showing alias and value", async () => { vi.mocked(getGeneralSettingsCall).mockResolvedValue([ { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY, CODEX] }, ]); @@ -148,10 +156,24 @@ describe("MCPNetworkSettings", () => { expect(await screen.findByText("Allowed Clients")).toBeVisible(); expect(screen.queryByText(/Allowed Client IDs/)).not.toBeInTheDocument(); - expect(screen.getByRole("textbox", { name: "Client 1 alias" })).toHaveValue("Antigravity CLI"); - expect(screen.getByRole("textbox", { name: "Client 1 value" })).toHaveValue("antigravity-cli"); - expect(screen.getByRole("textbox", { name: "Client 2 alias" })).toHaveValue("Codex"); - expect(screen.getByRole("textbox", { name: "Client 2 value" })).toHaveValue("codex-mcp-client"); + expect(screen.queryByText(/Allowed Client Applications/)).not.toBeInTheDocument(); + expect(clientCard("Antigravity CLI")).toHaveTextContent("antigravity-cli"); + expect(clientCard("Codex")).toHaveTextContent("codex-mcp-client"); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); + + it("opens an edit dialog when a client card is clicked, prefilled with that client's alias and value", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY, CODEX] }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Codex")); + + const dialog = await screen.findByRole("dialog", { name: "Edit client" }); + expect(within(dialog).getByRole("textbox", { name: "Alias" })).toHaveValue("Codex"); + expect(within(dialog).getByRole("textbox", { name: "Value" })).toHaveValue("codex-mcp-client"); }); it("warns that a stored allowlist in the old plain-string shape denies every client and lets Save remove it", async () => { @@ -162,7 +184,7 @@ describe("MCPNetworkSettings", () => { renderSettings(); expect(await screen.findByText(/stored allowlist is not a list of alias and value pairs/)).toBeVisible(); - expect(screen.queryByRole("textbox", { name: "Client 1 value" })).not.toBeInTheDocument(); + expect(screen.queryByText("antigravity-cli")).not.toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: /Save/ })); @@ -203,15 +225,22 @@ describe("MCPNetworkSettings", () => { expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients"); }); - it("edits a stored client's value in place and saves the new value", async () => { + it("edits a stored client's value through its dialog and saves the new value", async () => { vi.mocked(getGeneralSettingsCall).mockResolvedValue([ { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY] }, ]); renderSettings(); - fireEvent.change(await screen.findByRole("textbox", { name: "Client 1 value" }), { - target: { value: "0oa1b2c3d4e5f6g7h8i9" }, + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Antigravity CLI")); + const dialog = await screen.findByRole("dialog"); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Value" }), { + target: { value: " 0oa1b2c3d4e5f6g7h8i9 " }, }); + await userEvent.click(within(dialog).getByRole("button", { name: "Done" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(clientCard("Antigravity CLI")).toHaveTextContent("0oa1b2c3d4e5f6g7h8i9"); await userEvent.click(screen.getByRole("button", { name: /Save/ })); await waitFor(() => @@ -221,22 +250,37 @@ describe("MCPNetworkSettings", () => { ); }); - it("refuses to save a client that has an alias but no value, and reports why", async () => { + it("keeps a stored client untouched when its dialog is cancelled", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY] }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Antigravity CLI")); + const dialog = await screen.findByRole("dialog"); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Value" }), { target: { value: "changed" } }); + await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(clientCard("Antigravity CLI")).toHaveTextContent("antigravity-cli"); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("MCP network settings saved")); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("will not add a client that has an alias but no value", async () => { renderSettings(); await screen.findByText("Allowed Clients"); - await addClient("Antigravity CLI", ""); - await userEvent.click(screen.getByRole("button", { name: /Save/ })); + await userEvent.click(screen.getByRole("button", { name: "Add client" })); + const dialog = await fillClientDialog("Antigravity CLI", " "); - await waitFor(() => - expect(toast.fromError).toHaveBeenCalledWith(new Error("Every allowed client needs both an alias and a value")), - ); - expect(updateConfigFieldSetting).not.toHaveBeenCalled(); - expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); - expect(toast.success).not.toHaveBeenCalled(); + expect(within(dialog).getByRole("button", { name: "Add" })).toBeDisabled(); }); - it("drops rows left completely blank instead of saving or failing on them", async () => { + it("adds nothing when the add dialog is cancelled", async () => { vi.mocked(getGeneralSettingsCall).mockResolvedValue([ { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY] }, ]); @@ -244,6 +288,11 @@ describe("MCPNetworkSettings", () => { renderSettings(); await screen.findByText("Allowed Clients"); await userEvent.click(screen.getByRole("button", { name: "Add client" })); + const dialog = await fillClientDialog("Codex", "codex-mcp-client"); + await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.queryByText("Codex")).not.toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: /Save/ })); await waitFor(() => expect(toast.success).toHaveBeenCalledWith("MCP network settings saved")); @@ -260,13 +309,17 @@ describe("MCPNetworkSettings", () => { ]); renderSettings(); - await userEvent.click(await screen.findByRole("button", { name: "Remove client Claude Code" })); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Claude Code")); + await userEvent.click(within(await screen.findByRole("dialog")).getByRole("button", { name: "Remove client" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.queryByText("claude-code")).not.toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: /Save/ })); await waitFor(() => expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ANTIGRAVITY, CODEX]), ); - expect(screen.queryByDisplayValue("claude-code")).not.toBeInTheDocument(); }); it("removes a client and clears the setting when the list becomes empty", async () => { @@ -275,9 +328,12 @@ describe("MCPNetworkSettings", () => { ]); renderSettings(); - await userEvent.click(await screen.findByRole("button", { name: "Remove client Claude Code" })); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Claude Code")); + await userEvent.click(within(await screen.findByRole("dialog")).getByRole("button", { name: "Remove client" })); - expect(screen.queryByDisplayValue("claude-code")).not.toBeInTheDocument(); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.queryByText("claude-code")).not.toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: /Save/ })); @@ -304,7 +360,7 @@ describe("MCPNetworkSettings", () => { renderSettings(); - await screen.findByText("Allowed Client Applications"); + await screen.findByText("Allowed Clients"); expect(screen.queryByText(/every client is denied/)).not.toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index ae1fad36599..db45cfdedd1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -1,9 +1,18 @@ -import React, { useState, useEffect } from "react"; +import React, { useState, useEffect, useId } from "react"; import { Save, Plus, X } from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { DeprecationBanner } from "@/components/DeprecationBanner"; import { toast } from "@/lib/toast"; @@ -36,6 +45,10 @@ interface AllowedClientRow extends AllowedClient { readonly key: string; } +interface ClientDraft extends AllowedClient { + readonly key: string | null; +} + const isAllowedClient = (entry: unknown): entry is AllowedClient => { if (typeof entry !== "object" || entry === null) return false; const { alias, value } = entry as Partial>; @@ -58,14 +71,10 @@ const parseStoredClients = (fieldValue: unknown): StoredAllowlist => { }; let nextRowKey = 0; -const newRow = (client: AllowedClient = { alias: "", value: "" }): AllowedClientRow => ({ - ...client, - key: `client-${nextRowKey++}`, -}); +const newRow = (client: AllowedClient): AllowedClientRow => ({ ...client, key: `client-${nextRowKey++}` }); const trimClient = ({ alias, value }: AllowedClient): AllowedClient => ({ alias: alias.trim(), value: value.trim() }); -const isBlank = ({ alias, value }: AllowedClient) => alias === "" && value === ""; const isIncomplete = ({ alias, value }: AllowedClient) => alias === "" || value === ""; const sameList = (a: string[], b: string[]) => a.length === b.length && a.every((value, i) => value === b[i]); @@ -90,6 +99,67 @@ const clientsUnchangedSinceLoad = (value: AllowedClient[], stored: StoredAllowli const headerUnchangedSinceLoad = (value: string, stored: string | null) => stored === null ? value === "" : value !== "" && value === stored; +interface AllowedClientDialogProps { + readonly draft: ClientDraft | null; + readonly onChange: (draft: ClientDraft) => void; + readonly onCommit: () => void; + readonly onRemove: () => void; + readonly onClose: () => void; +} + +const AllowedClientDialog: React.FC = ({ draft, onChange, onCommit, onRemove, onClose }) => { + const aliasId = useId(); + const valueId = useId(); + if (draft === null) return null; + return ( +

!open && onClose()}> + + + {draft.key === null ? "Add client" : "Edit client"} + + The alias is the name shown in the dashboard and gateway logs. The value is the exact JWT claim or header + value that identifies the client, such as the OAuth client ID your identity provider issues. + + +
+
+ + onChange({ ...draft, alias: e.target.value })} + /> +
+
+ + onChange({ ...draft, value: e.target.value })} + /> +
+
+ + {draft.key !== null && ( + + )} + + + + + + ); +}; + const MCPNetworkSettings: React.FC = ({ accessToken }) => { const [loading, setLoading] = useState(true); const [saving, setSaving] = useState(false); @@ -101,6 +171,7 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) const [storedClientIdHeader, setStoredClientIdHeader] = useState(null); const [currentIp, setCurrentIp] = useState(null); const [rangeDraft, setRangeDraft] = useState(""); + const [clientDraft, setClientDraft] = useState(null); useEffect(() => { loadSettings(); @@ -154,10 +225,7 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) }; const persistAllowedClients = async (token: string) => { - const clients = allowedClients.map(trimClient).filter((client) => !isBlank(client)); - if (clients.some(isIncomplete)) { - throw new Error("Every allowed client needs both an alias and a value"); - } + const clients = allowedClients.map(({ alias, value }) => ({ alias, value })); if (clientsUnchangedSinceLoad(clients, storedClients)) return; if (clients.length > 0) { await updateConfigFieldSetting(token, "mcp_allowed_clients", clients); @@ -218,10 +286,22 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) setRangeDraft(""); }; - const updateClient = (key: string, patch: Partial) => - setAllowedClients(allowedClients.map((row) => (row.key === key ? { ...row, ...patch } : row))); + const commitClientDraft = () => { + if (clientDraft === null) return; + const client = trimClient(clientDraft); + setAllowedClients( + clientDraft.key === null + ? [...allowedClients, newRow(client)] + : allowedClients.map((row) => (row.key === clientDraft.key ? { ...row, ...client } : row)), + ); + setClientDraft(null); + }; - const removeClient = (key: string) => setAllowedClients(allowedClients.filter((row) => row.key !== key)); + const removeDraftedClient = () => { + if (clientDraft === null) return; + setAllowedClients(allowedClients.filter((row) => row.key !== clientDraft.key)); + setClientDraft(null); + }; if (loading) { return ( @@ -307,7 +387,7 @@ const MCPNetworkSettings: React.FC = ({ accessToken })
-

Allowed Client Applications

+

Allowed Clients

Only the MCP client applications listed here can use the gateway. Leave empty to allow every client. A client that authenticates with a JWT is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field in @@ -317,9 +397,6 @@ const MCPNetworkSettings: React.FC = ({ accessToken })

-
-

Allowed Clients

-
{storedAllowlistIsMalformed && (

The stored allowlist is not a list of alias and value pairs, so every client is denied. Add the clients you @@ -333,35 +410,17 @@ const MCPNetworkSettings: React.FC = ({ accessToken })

)} {allowedClients.length > 0 && ( -
-

Alias

-

Value

- - {allowedClients.map((row, index) => ( - - updateClient(row.key, { alias: e.target.value })} - /> - updateClient(row.key, { value: e.target.value })} - /> - - +
+ {allowedClients.map((row) => ( + ))}
)} @@ -369,16 +428,14 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) type="button" variant="outline" size="sm" - onClick={() => setAllowedClients([...allowedClients, newRow()])} + onClick={() => setClientDraft({ key: null, alias: "", value: "" })} > Add client

- The alias is the name shown here and in gateway logs. The value is the exact JWT claim or header value that - identifies the client, such as the OAuth client ID your identity provider issues. Leave the list empty to - allow every client. Every MCP request from an unlisted client, or from one with no resolvable identity, gets a - 403. + Click a client to edit or remove it. Leave the list empty to allow every client. Every MCP request from an + unlisted client, or from one with no resolvable identity, gets a 403.

@@ -403,6 +460,14 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) Save
+ + setClientDraft(null)} + />
); }; From 92d82841fd6d00f309e8afcc7044938e598f25bf Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:14:40 +0000 Subject: [PATCH 60/61] chore(model_info): backfill reseller Gemini entries from provider catalogs and prune retired ids Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 268 ++++++++++++------ model_prices_and_context_window.json | 268 ++++++++++++------ 2 files changed, 370 insertions(+), 166 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 30b08e54410..0e63653e2f2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19244,7 +19244,20 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-2-5-pro": { "cache_creation_input_token_cost": 1.24999e-06, @@ -19265,7 +19278,21 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2026-10-02", + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, @@ -19285,7 +19312,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-image": { "litellm_provider": "databricks", @@ -19347,7 +19386,20 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, @@ -19367,7 +19419,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, @@ -21433,7 +21497,11 @@ "mode": "chat", "supports_tool_choice": true, "supports_function_calling": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemini-2.5-pro": { "max_tokens": 1000000, @@ -21444,7 +21512,11 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_function_calling": true + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemma-3-12b-it": { "max_tokens": 131072, @@ -29340,26 +29412,6 @@ "supports_parallel_function_calling": true, "supports_vision": true }, - "github_copilot/gemini-2.5-pro": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, - "github_copilot/gemini-3-pro-preview": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, "github_copilot/gpt-3.5-turbo": { "litellm_provider": "github_copilot", "max_input_tokens": 16384, @@ -30014,17 +30066,6 @@ "output_cost_per_token": 8.8e-07, "supports_function_calling": true }, - "gmi/google/gemini-3-pro-preview": { - "input_cost_per_token": 2e-06, - "litellm_provider": "gmi", - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 1.2e-05, - "supports_function_calling": true, - "supports_vision": true - }, "gmi/google/gemini-3-flash-preview": { "input_cost_per_token": 5e-07, "litellm_provider": "gmi", @@ -30034,7 +30075,8 @@ "mode": "chat", "output_cost_per_token": 3e-06, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "supports_system_messages": true }, "gmi/moonshotai/Kimi-K2-Thinking": { "input_cost_per_token": 8e-07, @@ -40163,7 +40205,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-pro": { "input_cost_per_token": 1.25e-06, @@ -40177,7 +40224,12 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_vision": true, - "supports_native_streaming": true + "supports_native_streaming": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-flash-lite": { "input_cost_per_token": 7.5e-08, @@ -40192,7 +40244,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": false, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/cohere.command-a-vision": { "input_cost_per_token": 1.56e-06, @@ -41435,7 +41492,7 @@ "max_tokens": 65535, "mode": "chat", "output_cost_per_token": 2.5e-06, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41449,7 +41506,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro": { "cache_creation_input_token_cost": 3.75e-07, @@ -41462,7 +41520,7 @@ "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41478,7 +41536,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -41563,7 +41622,8 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": false, - "tpm": 800000 + "tpm": 800000, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite-preview": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -41690,7 +41750,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/gryphe/mythomax-l2-13b": { "input_cost_per_token": 8e-08, @@ -44413,12 +44474,16 @@ "output_cost_per_token": 1.2e-05, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true + "supports_tool_choice": false, + "supports_response_schema": false, + "input_cost_per_token_above_200k_tokens": 4e-06, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "supports_audio_input": true, + "supports_video_input": true }, "replicate/anthropic/claude-4.5-sonnet": { "input_cost_per_token": 3e-06, @@ -44487,17 +44552,19 @@ "supports_response_schema": true }, "replicate/google/gemini-2.5-flash": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 3e-07, "output_cost_per_token": 2.5e-06, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_image_size": false + "supports_tool_choice": false, + "supports_response_schema": false, + "supports_image_size": false, + "supports_reasoning": true, + "supports_video_input": true }, "replicate/openai/gpt-oss-120b": { "input_cost_per_token": 1.8e-07, @@ -48076,10 +48143,15 @@ "supports_function_calling": true, "supports_tool_choice": true, "supports_response_schema": true, - "supports_image_size": false + "supports_image_size": false, + "cache_read_input_token_cost": 3e-08, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-2.5-pro": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -48089,7 +48161,15 @@ "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, - "supports_response_schema": true + "supports_response_schema": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -62224,7 +62304,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5": { "max_tokens": 262144, @@ -62524,7 +62605,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/google/gemini-3.7-flash": { "max_tokens": 1000000, @@ -62538,7 +62620,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/inclusionAI/Ling-3.0-flash": { "max_tokens": 131072, @@ -62970,7 +63053,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5-Pro": { "max_tokens": 1048576, @@ -65365,7 +65449,8 @@ "deprecation_date": "2026-10-20", "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65388,7 +65473,8 @@ "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 3e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65411,7 +65497,8 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65434,7 +65521,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65457,7 +65545,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65480,7 +65569,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/openai/gpt-4o-mini": { "input_cost_per_token": 1.5e-07, @@ -67108,7 +67198,8 @@ "supports_pdf_input": true, "supports_audio_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/qwen/qwen3-max-thinking": { "input_cost_per_token": 7.8e-07, @@ -71974,7 +72065,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-flash:batch": { "cache_read_input_audio_token_cost": 1e-07, @@ -71997,7 +72089,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro:batch": { "cache_read_input_audio_token_cost": 1.25e-07, @@ -72023,7 +72116,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-flash-preview:batch": { "input_cost_per_audio_token": 5e-07, @@ -72043,7 +72137,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite:batch": { "cache_read_input_audio_token_cost": 2.5e-08, @@ -72065,7 +72160,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-pro-preview:batch": { "input_cost_per_audio_token": 1e-06, @@ -72087,7 +72183,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite:batch": { "cache_read_input_audio_token_cost": 1.5e-08, @@ -72109,7 +72206,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash:batch": { "cache_read_input_audio_token_cost": 1.5e-07, @@ -72131,7 +72229,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72154,7 +72253,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72177,7 +72277,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72200,7 +72301,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/ibm-granite/granite-4.0-h-micro": { "input_cost_per_token": 1.7e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 30b08e54410..0e63653e2f2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19244,7 +19244,20 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-2-5-pro": { "cache_creation_input_token_cost": 1.24999e-06, @@ -19265,7 +19278,21 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2026-10-02", + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, @@ -19285,7 +19312,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-image": { "litellm_provider": "databricks", @@ -19347,7 +19386,20 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, @@ -19367,7 +19419,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, @@ -21433,7 +21497,11 @@ "mode": "chat", "supports_tool_choice": true, "supports_function_calling": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemini-2.5-pro": { "max_tokens": 1000000, @@ -21444,7 +21512,11 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_function_calling": true + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemma-3-12b-it": { "max_tokens": 131072, @@ -29340,26 +29412,6 @@ "supports_parallel_function_calling": true, "supports_vision": true }, - "github_copilot/gemini-2.5-pro": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, - "github_copilot/gemini-3-pro-preview": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, "github_copilot/gpt-3.5-turbo": { "litellm_provider": "github_copilot", "max_input_tokens": 16384, @@ -30014,17 +30066,6 @@ "output_cost_per_token": 8.8e-07, "supports_function_calling": true }, - "gmi/google/gemini-3-pro-preview": { - "input_cost_per_token": 2e-06, - "litellm_provider": "gmi", - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 1.2e-05, - "supports_function_calling": true, - "supports_vision": true - }, "gmi/google/gemini-3-flash-preview": { "input_cost_per_token": 5e-07, "litellm_provider": "gmi", @@ -30034,7 +30075,8 @@ "mode": "chat", "output_cost_per_token": 3e-06, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "supports_system_messages": true }, "gmi/moonshotai/Kimi-K2-Thinking": { "input_cost_per_token": 8e-07, @@ -40163,7 +40205,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-pro": { "input_cost_per_token": 1.25e-06, @@ -40177,7 +40224,12 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_vision": true, - "supports_native_streaming": true + "supports_native_streaming": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-flash-lite": { "input_cost_per_token": 7.5e-08, @@ -40192,7 +40244,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": false, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/cohere.command-a-vision": { "input_cost_per_token": 1.56e-06, @@ -41435,7 +41492,7 @@ "max_tokens": 65535, "mode": "chat", "output_cost_per_token": 2.5e-06, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41449,7 +41506,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro": { "cache_creation_input_token_cost": 3.75e-07, @@ -41462,7 +41520,7 @@ "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41478,7 +41536,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -41563,7 +41622,8 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": false, - "tpm": 800000 + "tpm": 800000, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite-preview": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -41690,7 +41750,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/gryphe/mythomax-l2-13b": { "input_cost_per_token": 8e-08, @@ -44413,12 +44474,16 @@ "output_cost_per_token": 1.2e-05, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true + "supports_tool_choice": false, + "supports_response_schema": false, + "input_cost_per_token_above_200k_tokens": 4e-06, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "supports_audio_input": true, + "supports_video_input": true }, "replicate/anthropic/claude-4.5-sonnet": { "input_cost_per_token": 3e-06, @@ -44487,17 +44552,19 @@ "supports_response_schema": true }, "replicate/google/gemini-2.5-flash": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 3e-07, "output_cost_per_token": 2.5e-06, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_image_size": false + "supports_tool_choice": false, + "supports_response_schema": false, + "supports_image_size": false, + "supports_reasoning": true, + "supports_video_input": true }, "replicate/openai/gpt-oss-120b": { "input_cost_per_token": 1.8e-07, @@ -48076,10 +48143,15 @@ "supports_function_calling": true, "supports_tool_choice": true, "supports_response_schema": true, - "supports_image_size": false + "supports_image_size": false, + "cache_read_input_token_cost": 3e-08, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-2.5-pro": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -48089,7 +48161,15 @@ "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, - "supports_response_schema": true + "supports_response_schema": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -62224,7 +62304,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5": { "max_tokens": 262144, @@ -62524,7 +62605,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/google/gemini-3.7-flash": { "max_tokens": 1000000, @@ -62538,7 +62620,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/inclusionAI/Ling-3.0-flash": { "max_tokens": 131072, @@ -62970,7 +63053,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5-Pro": { "max_tokens": 1048576, @@ -65365,7 +65449,8 @@ "deprecation_date": "2026-10-20", "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65388,7 +65473,8 @@ "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 3e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65411,7 +65497,8 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65434,7 +65521,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65457,7 +65545,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65480,7 +65569,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/openai/gpt-4o-mini": { "input_cost_per_token": 1.5e-07, @@ -67108,7 +67198,8 @@ "supports_pdf_input": true, "supports_audio_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/qwen/qwen3-max-thinking": { "input_cost_per_token": 7.8e-07, @@ -71974,7 +72065,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-flash:batch": { "cache_read_input_audio_token_cost": 1e-07, @@ -71997,7 +72089,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro:batch": { "cache_read_input_audio_token_cost": 1.25e-07, @@ -72023,7 +72116,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-flash-preview:batch": { "input_cost_per_audio_token": 5e-07, @@ -72043,7 +72137,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite:batch": { "cache_read_input_audio_token_cost": 2.5e-08, @@ -72065,7 +72160,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-pro-preview:batch": { "input_cost_per_audio_token": 1e-06, @@ -72087,7 +72183,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite:batch": { "cache_read_input_audio_token_cost": 1.5e-08, @@ -72109,7 +72206,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash:batch": { "cache_read_input_audio_token_cost": 1.5e-07, @@ -72131,7 +72229,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72154,7 +72253,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72177,7 +72277,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72200,7 +72301,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/ibm-granite/granite-4.0-h-micro": { "input_cost_per_token": 1.7e-08, From e2141da81ea059c6946c7c7c674babd7ef62443a Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 19 Sep 2026 00:23:11 +0000 Subject: [PATCH 61/61] fix(ui): treat MCP allowed clients with an empty alias or value as malformed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/MCPNetworkSettings.test.tsx | 11 +++++++++++ .../mcp-servers/_components/MCPNetworkSettings.tsx | 2 +- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx index 4cc87f1455c..4a486bfc648 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx @@ -211,6 +211,17 @@ describe("MCPNetworkSettings", () => { await waitFor(() => expect(screen.queryByText(/stored allowlist is not a list/)).not.toBeInTheDocument()); }); + it("treats a stored entry with an empty alias or value as denying every client, like the gateway does", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY, { alias: "", value: "claude-code" }] }, + ]); + + renderSettings(); + + expect(await screen.findByText(/stored allowlist is not a list of alias and value pairs/)).toBeVisible(); + expect(screen.queryByRole("button", { name: /^Antigravity CLI/ })).not.toBeInTheDocument(); + }); + it("adds clients as alias and value pairs and saves them under mcp_allowed_clients", async () => { renderSettings(); await screen.findByText("Allowed Clients"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index db45cfdedd1..2fd62c7f1ef 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -52,7 +52,7 @@ interface ClientDraft extends AllowedClient { const isAllowedClient = (entry: unknown): entry is AllowedClient => { if (typeof entry !== "object" || entry === null) return false; const { alias, value } = entry as Partial>; - return typeof alias === "string" && typeof value === "string"; + return typeof alias === "string" && typeof value === "string" && !isIncomplete({ alias, value }); }; type StoredAllowlist =