From 0135e1d634c4f9068fcd345e305767e76dbf5a4e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:01:44 -0700 Subject: [PATCH 001/155] fix(oci): pin one response id per streamed completion, skip the [DONE] sentinel OCIStreamWrapper.chunk_creator built every chunk straight from the apiFormat handlers, so it never reached model_response_creator and OCI streams came back with a fresh chatcmpl id, a drifting created value and no model on every chunk. Both exits now go through the shared creator. The GENERIC apiFormat also closes its stream with a literal `data: [DONE]` line, which chunk_creator json-parsed and turned into a 500 on every OCI streaming completion. It is skipped now. --- litellm/llms/oci/chat/transformation.py | 16 +- .../oci/chat/test_oci_chat_transformation.py | 145 ++++++++++++++++++ 2 files changed, 157 insertions(+), 4 deletions(-) diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 98e23a59eea..17aed74b9d3 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -745,13 +745,21 @@ class OCIStreamWrapper(CustomStreamWrapper): # single-event case (terminal chunk carries the only copy of the text). self._cohere_text_emitted = False - def chunk_creator(self, chunk: Any) -> ModelResponseStream: + def _with_stream_identity(self, parsed: ModelResponseStream) -> ModelResponseStream: + model_response: Final = self.model_response_creator() + model_response.choices = parsed.choices + return model_response + + def chunk_creator(self, chunk: Any) -> ModelResponseStream | None: if not isinstance(chunk, str): raise ValueError(f"Chunk is not a string: {chunk}") if not chunk.startswith("data:"): raise ValueError(f"Chunk does not start with 'data:': {chunk}") + payload: Final = chunk[5:].strip() + if payload == "[DONE]": + return None try: - dict_chunk: Final = json.loads(chunk[5:]) + dict_chunk: Final = json.loads(payload) except json.JSONDecodeError as e: raise OCIError( status_code=500, @@ -774,8 +782,8 @@ class OCIStreamWrapper(CustomStreamWrapper): if getattr(choice.delta, "content", None): self._cohere_text_emitted = True break - return result - return handle_generic_stream_chunk(dict_chunk) + return self._with_stream_identity(result) + return self._with_stream_identity(handle_generic_stream_chunk(dict_chunk)) __all__ = [ diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py index 86c534c73c2..6601de2853d 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py @@ -1900,3 +1900,148 @@ class TestOCIImageUrlTransformation: adapt_messages_to_generic_oci_standard(messages) assert "image_url" in str(exc_info.value) + + +# --------------------------------------------------------------------------- +# OCIStreamWrapper: per-stream identity and the GENERIC `[DONE]` sentinel +# --------------------------------------------------------------------------- + +import itertools +from unittest.mock import patch + +from litellm.llms.oci.chat.transformation import OCIStreamWrapper, _iter_sse_events + +_STREAM_GENERIC_MODEL = "xai.grok-4" +_STREAM_COHERE_MODEL = "cohere.command-latest" + +_GENERIC_TEXT_EVENT = ( + 'data: {{"index":0,"message":{{"role":"ASSISTANT","content":[{{"type":"TEXT","text":"{text}"}}]}},"pad":"aaa"}}' +) +_GENERIC_TERMINAL_EVENT = ( + 'data: {"message":{"role":"ASSISTANT","content":[{"type":"TEXT","text":""}]},"finishReason":"stop","pad":"a"}' +) +_COHERE_TEXT_EVENT = 'data: {{"apiFormat":"COHERE","text":"{text}","pad":"aaaaaa"}}' +_COHERE_TERMINAL_EVENT = ( + 'data: {"apiFormat":"COHERE","text":"123","finishReason":"COMPLETE",' + '"chatHistory":[{"role":"USER","message":"count"},{"role":"CHATBOT","message":"123"}]}' +) + + +def _make_stream_wrapper(model: str) -> OCIStreamWrapper: + logging_obj = MagicMock() + logging_obj.model_call_details = {"custom_llm_provider": "oci", "litellm_params": {}} + return OCIStreamWrapper( + completion_stream=iter([]), + model=model, + custom_llm_provider="oci", + logging_obj=logging_obj, + ) + + +def _ticking_clock(): + """A ``time.time`` stand-in that advances a full second on every call. + + Without it the whole test runs inside one wall-clock second, so a per-chunk + ``created`` would coincidentally match and the drift would go unnoticed. + """ + return itertools.count(1_700_000_000.0) + + +class TestOCIStreamWrapperIdentityPinning: + """One OCI streaming completion must present one id, one created and the + wrapper's model on every chunk, the way every other provider does.""" + + def test_generic_stream_shares_one_id_created_and_model(self): + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + events = [ + _GENERIC_TEXT_EVENT.format(text="1"), + _GENERIC_TEXT_EVENT.format(text="2"), + _GENERIC_TEXT_EVENT.format(text="3"), + _GENERIC_TERMINAL_EVENT, + ] + + with patch("time.time", side_effect=_ticking_clock()): + chunks = [wrapper.chunk_creator(event) for event in events] + + assert len(chunks) == 4 + assert len({chunk.id for chunk in chunks}) == 1 + assert chunks[0].id.startswith("chatcmpl-") + assert len({chunk.created for chunk in chunks}) == 1 + assert {chunk.model for chunk in chunks} == {_STREAM_GENERIC_MODEL} + assert [chunk.choices[0].delta.content for chunk in chunks[:3]] == ["1", "2", "3"] + assert chunks[-1].choices[0].finish_reason == "stop" + assert all(chunk._hidden_params["custom_llm_provider"] == "oci" for chunk in chunks) + + def test_cohere_stream_shares_one_id_created_and_model(self): + """Rebuilding each chunk through the shared creator must not disturb the + Cohere bookkeeping that suppresses the terminal event's repeated text.""" + wrapper = _make_stream_wrapper(_STREAM_COHERE_MODEL) + events = [ + _COHERE_TEXT_EVENT.format(text="1"), + _COHERE_TEXT_EVENT.format(text="2"), + _COHERE_TEXT_EVENT.format(text="3"), + _COHERE_TERMINAL_EVENT, + ] + + with patch("time.time", side_effect=_ticking_clock()): + chunks = [wrapper.chunk_creator(event) for event in events] + + assert len(chunks) == 4 + assert len({chunk.id for chunk in chunks}) == 1 + assert len({chunk.created for chunk in chunks}) == 1 + assert {chunk.model for chunk in chunks} == {_STREAM_COHERE_MODEL} + assert [chunk.choices[0].delta.content for chunk in chunks[:3]] == ["1", "2", "3"] + assert chunks[-1].choices[0].finish_reason == "stop" + assert chunks[-1].choices[0].delta.content is None + assert wrapper._cohere_text_emitted is True + + def test_id_is_pinned_to_the_wrapper_response_id(self): + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + + first = wrapper.chunk_creator(_GENERIC_TEXT_EVENT.format(text="1")) + + assert wrapper.response_id == first.id + assert wrapper.created == first.created + + +class TestOCIStreamWrapperDoneSentinel: + """OCI's GENERIC apiFormat closes the stream with a literal `[DONE]` line; + parsing it as JSON turned every streaming completion into a 500.""" + + @pytest.mark.parametrize("done_event", ["data: [DONE]", "data:[DONE]", "data: [DONE] "]) + def test_done_sentinel_returns_none(self, done_event): + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + assert wrapper.chunk_creator(done_event) is None + + def test_done_sentinel_off_the_sse_splitter_is_skipped(self): + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + wire = ( + f"{_GENERIC_TEXT_EVENT.format(text='1')}\n\n" + f"{_GENERIC_TEXT_EVENT.format(text='2')}\n\n" + f"{_GENERIC_TERMINAL_EVENT}\n\n" + "data: [DONE]\n\n" + ) + + events = list(_iter_sse_events(iter([wire]))) + assert events[-1] == "data: [DONE]" + + chunks = [wrapper.chunk_creator(event) for event in events] + assert chunks[-1] is None + + emitted = [chunk for chunk in chunks if chunk is not None] + assert len(emitted) == 3 + assert len({chunk.id for chunk in emitted}) == 1 + + def test_unparseable_payload_still_raises_oci_error(self): + from litellm.llms.oci.common_utils import OCIError + + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + with pytest.raises(OCIError, match="Chunk cannot be parsed as JSON"): + wrapper.chunk_creator("data: not-json-at-all") + + def test_done_lookalike_payload_still_raises_oci_error(self): + from litellm.llms.oci.common_utils import OCIError + + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + with pytest.raises(OCIError, match="Chunk cannot be parsed as JSON"): + wrapper.chunk_creator("data: [DONE] trailing garbage") From fd2104a61efe91eb1428d40c3261527c363ea5d6 Mon Sep 17 00:00:00 2001 From: mateo-berri Date: Thu, 3 Sep 2026 00:27:54 -0700 Subject: [PATCH 002/155] refactor(oci): build the identity-pinned stream chunk in one shot Address the two Greptile P2 notes: construct the chunk through the shared creator instead of mutating its choices afterwards, and drop the decorative divider comment from the new tests. --- litellm/llms/oci/chat/transformation.py | 4 +--- .../llms/oci/chat/test_oci_chat_transformation.py | 4 ---- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 17aed74b9d3..b1167d447fc 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -746,9 +746,7 @@ class OCIStreamWrapper(CustomStreamWrapper): self._cohere_text_emitted = False def _with_stream_identity(self, parsed: ModelResponseStream) -> ModelResponseStream: - model_response: Final = self.model_response_creator() - model_response.choices = parsed.choices - return model_response + return self.model_response_creator(chunk={"choices": parsed.choices}) def chunk_creator(self, chunk: Any) -> ModelResponseStream | None: if not isinstance(chunk, str): diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py index 6601de2853d..e77b2c24d01 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py @@ -1902,10 +1902,6 @@ class TestOCIImageUrlTransformation: assert "image_url" in str(exc_info.value) -# --------------------------------------------------------------------------- -# OCIStreamWrapper: per-stream identity and the GENERIC `[DONE]` sentinel -# --------------------------------------------------------------------------- - import itertools from unittest.mock import patch From 445ccc14b29bb4a93e3b21bf907a7b2a26da507a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:53:55 -0700 Subject: [PATCH 003/155] fix(vector-stores): surface retrieval failures to the API caller A vector store search that fails is swallowed by the pre-call hook, so the request goes to the model with an un-augmented prompt and the caller gets a 200 answering from the model's own knowledge with no way to tell the knowledge base was skipped. Failed searches now ride the same channel their successes already use: a vector_store_search_failures entry on provider_specific_fields naming the store id, provider, and error. That is additive and always on. For callers who would rather fail than answer ungrounded, litellm_settings vector_store_search_failure_mode: error raises VectorStoreSearchError (400) instead; the default stays annotate, today's permissive behavior. The hook's outer catch-all also now names the requested vector store ids in its log line, and only wraps the augmentation itself, so the fail-closed raise is not swallowed by it. --- litellm/__init__.py | 3 + litellm/exceptions.py | 25 ++ .../vector_store_pre_call_hook.py | 378 ++++++++++-------- litellm/types/vector_stores.py | 13 +- .../test_vector_store_pre_call_hook.py | 199 +++++++++ 5 files changed, 460 insertions(+), 158 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 41a3789ab0d..76daf771c4b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1360,6 +1360,7 @@ from .exceptions import ( InvalidRequestError, BadRequestError, ImageFetchError, + VectorStoreSearchError, NotFoundError, PermissionDeniedError, RateLimitError, @@ -1461,9 +1462,11 @@ from .vector_stores.vector_store_registry import ( VectorStoreRegistry, VectorStoreIndexRegistry, ) +from .types.vector_stores import VectorStoreSearchFailureMode vector_store_registry: Optional[VectorStoreRegistry] = None vector_store_index_registry: Optional[VectorStoreIndexRegistry] = None +vector_store_search_failure_mode: VectorStoreSearchFailureMode = "annotate" ### RAG ### from . import rag diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 286f7528896..4387370e58c 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -10,12 +10,14 @@ ## LiteLLM versions of the OpenAI Exception Types import enum +from collections.abc import Sequence from typing import Any, Final import httpx import openai from litellm.types.utils import LiteLLMCommonStrings +from litellm.types.vector_stores import VectorStoreSearchFailure class RateLimitErrorCategory(str, enum.Enum): @@ -288,6 +290,29 @@ class ImageFetchError(BadRequestError): ) +VECTOR_STORE_SEARCH_FAILED_CODE: Final = "vector_store_search_failed" + + +class VectorStoreSearchError(BadRequestError): + def __init__( + self, + failures: Sequence[VectorStoreSearchFailure], + model: str | None = None, + llm_provider: str | None = None, + ) -> None: + self.failures: Final[tuple[VectorStoreSearchFailure, ...]] = tuple(failures) + detail: Final = "; ".join(f"{failure['vector_store_id']}: {failure['error']}" for failure in self.failures) + super().__init__( + message=( + "The request could not be grounded in every configured vector store. " + f"{len(self.failures)} vector store search(es) failed: {detail}" + ), + model=model, + llm_provider=llm_provider, + body={"type": "invalid_request_error", "code": VECTOR_STORE_SEARCH_FAILED_CODE}, + ) + + class UnprocessableEntityError(openai.UnprocessableEntityError): def __init__( self, diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 12ff38ce4ba..d5c893c4b20 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -5,20 +5,21 @@ This hook is called before making an LLM request when a vector store is configur It searches the vector store for relevant context and appends it to the messages. """ -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, assert_never, cast import litellm import litellm.vector_stores from litellm._logging import verbose_logger +from litellm.exceptions import VectorStoreSearchError from litellm.integrations.custom_logger import CustomLogger from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.utils import CallTypes, StandardCallbackDynamicParams from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, - VectorStoreResultContent, + VectorStoreSearchFailure, VectorStoreSearchResponse, VectorStoreSearchResult, ) @@ -30,6 +31,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +SEARCH_FAILURES_FIELD: Final = "vector_store_search_failures" + class ProxyRuntime(Protocol): def llm_router(self) -> "Router | None": ... @@ -54,11 +57,31 @@ class ProxyServerRuntime: return prisma_client +@dataclass(frozen=True, slots=True) +class SearchSucceeded: + response: VectorStoreSearchResponse + + +@dataclass(frozen=True, slots=True) +class SearchFailed: + failure: VectorStoreSearchFailure + + +SearchOutcome = SearchSucceeded | SearchFailed + + +@dataclass(frozen=True, slots=True) +class VectorStoreAugmentation: + messages: tuple[AllMessageValues, ...] + search_results: tuple[VectorStoreSearchResponse, ...] + failures: tuple[VectorStoreSearchFailure, ...] + + class VectorStorePreCallHook(CustomLogger): CONTENT_PREFIX_STRING = "Context:\n\n" """ Custom logger that handles vector store searches before LLM calls. - + When a vector store is configured, this hook: 1. Extracts the query from the last user message 2. Calls litellm.vector_stores.search() to get relevant context @@ -101,100 +124,152 @@ class VectorStorePreCallHook(CustomLogger): Returns: Tuple of (model, modified_messages, non_default_params) """ + requested_vector_store_ids: Final = _requested_vector_store_ids(non_default_params) try: - # Check if vector store is configured - if litellm.vector_store_registry is None: - return model, messages, non_default_params - - prisma_client: Final = self.proxy_runtime.prisma_client() - llm_router: Final = self.proxy_runtime.llm_router() - - # Use database fallback to ensure synchronization across instances - vector_stores_to_run: list[ - LiteLLM_ManagedVectorStore - ] = await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback( + augmentation: VectorStoreAugmentation | None = await self._augment_messages( + messages=messages, non_default_params=non_default_params, tools=tools, - prisma_client=prisma_client, + litellm_logging_obj=litellm_logging_obj, ) - - if not vector_stores_to_run: - return model, messages, non_default_params - - # Extract the query from the last user message - query: Final = self._extract_query_from_messages(messages) - - if not query: - verbose_logger.debug("No query found in messages for vector store search") - return model, messages, non_default_params - - modified_messages: list[AllMessageValues] = messages.copy() - all_search_results: Final[list[VectorStoreSearchResponse]] = [] - - for vector_store_to_run in vector_stores_to_run: - # Get vector store id from the vector store config - vector_store_id = vector_store_to_run.get("vector_store_id", "") - custom_llm_provider = vector_store_to_run.get("custom_llm_provider") - litellm_params_for_vector_store = vector_store_to_run.get("litellm_params", {}) or {} - request_litellm_params = litellm_logging_obj.model_call_details.get("litellm_params", {}) - request_metadata = ( - request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {} - ) - if llm_router is not None: - search_function = cast( # cast-ok: normalize router search callable - Callable[..., Awaitable[VectorStoreSearchResponse]], - llm_router.avector_store_search, - ) - else: - search_function = cast( # cast-ok: normalize SDK search callable - Callable[..., Awaitable[VectorStoreSearchResponse]], - litellm.vector_stores.asearch, - ) - try: - search_response = await search_function( - **{ - "vector_store_id": vector_store_id, - "query": query, - "custom_llm_provider": custom_llm_provider, - "metadata": request_metadata, - **litellm_params_for_vector_store, - }, - ) - except Exception as search_error: - verbose_logger.warning( - "Vector store search failed for vector_store_id=%s, continuing without its context: %s", - vector_store_id, - search_error, - ) - continue - - verbose_logger.debug("search_response: %s", search_response) - - # Store search results for later use in citations - all_search_results.append(search_response) - - # Process search results and append as context - modified_messages = self._append_search_results_to_messages( - messages=modified_messages, search_response=search_response - ) - - # Get the number of results for logging - num_results = 0 - num_results = len(search_response.get("data", []) or []) - verbose_logger.debug("Vector store search completed. Added context from %s results", num_results) - - # Store search results as-is (already in OpenAI-compatible format) - if litellm_logging_obj and all_search_results: - litellm_logging_obj.model_call_details["search_results"] = all_search_results - - return model, modified_messages, non_default_params - except Exception as e: - verbose_logger.exception("Error in VectorStorePreCallHook: %s", e) - # Return original parameters on error + verbose_logger.exception( + "Error in VectorStorePreCallHook for vector_store_ids=%s: %s", + requested_vector_store_ids, + e, + ) return model, messages, non_default_params - def _extract_query_from_messages(self, messages: list[AllMessageValues]) -> str | None: + if augmentation is None: + return model, messages, non_default_params + + for detail, value in ( + ("search_results", list(augmentation.search_results)), + (SEARCH_FAILURES_FIELD, augmentation.failures), + ): + if value: + litellm_logging_obj.model_call_details[detail] = value + + if augmentation.failures: + match litellm.vector_store_search_failure_mode: + case "error": + raise VectorStoreSearchError(failures=augmentation.failures, model=model) + case "annotate": + pass + case unreachable: + assert_never(unreachable) + + return model, list(augmentation.messages), non_default_params + + async def _augment_messages( + self, + messages: Sequence[AllMessageValues], + non_default_params: dict, + tools: list[dict] | None, + litellm_logging_obj: LiteLLMLoggingObj, + ) -> VectorStoreAugmentation | None: + if litellm.vector_store_registry is None: + return None + + prisma_client: Final = self.proxy_runtime.prisma_client() + llm_router: Final = self.proxy_runtime.llm_router() + + # Use database fallback to ensure synchronization across instances + vector_stores_to_run: Final[ + Sequence[LiteLLM_ManagedVectorStore] + ] = await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback( + non_default_params=non_default_params, + tools=tools, + prisma_client=prisma_client, + ) + + if not vector_stores_to_run: + return None + + query: Final = self._extract_query_from_messages(messages) + + if not query: + verbose_logger.debug("No query found in messages for vector store search") + return None + + request_litellm_params: Final = litellm_logging_obj.model_call_details.get("litellm_params", {}) + request_metadata: Final = ( + request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {} + ) + search_function: Final = ( + cast( # cast-ok: normalize router search callable + Callable[..., Awaitable[VectorStoreSearchResponse]], + llm_router.avector_store_search, + ) + if llm_router is not None + else cast( # cast-ok: normalize SDK search callable + Callable[..., Awaitable[VectorStoreSearchResponse]], + litellm.vector_stores.asearch, + ) + ) + + outcomes: Final = tuple( + [ + await self._search_one( + vector_store=vector_store_to_run, + query=query, + request_metadata=request_metadata, + search_function=search_function, + ) + for vector_store_to_run in vector_stores_to_run + ] + ) + search_results: Final = tuple(outcome.response for outcome in outcomes if isinstance(outcome, SearchSucceeded)) + failures: Final = tuple(outcome.failure for outcome in outcomes if isinstance(outcome, SearchFailed)) + + return VectorStoreAugmentation( + messages=self._messages_with_context(messages=messages, search_results=search_results), + search_results=search_results, + failures=failures, + ) + + async def _search_one( + self, + vector_store: LiteLLM_ManagedVectorStore, + query: str, + request_metadata: Mapping[str, object], + search_function: Callable[..., Awaitable[VectorStoreSearchResponse]], + ) -> SearchOutcome: + vector_store_id: Final = vector_store.get("vector_store_id", "") + custom_llm_provider: Final = vector_store.get("custom_llm_provider") + litellm_params_for_vector_store: Final = vector_store.get("litellm_params", {}) or {} + try: + search_response: Final = await search_function( + **{ + "vector_store_id": vector_store_id, + "query": query, + "custom_llm_provider": custom_llm_provider, + "metadata": request_metadata, + **litellm_params_for_vector_store, + }, + ) + except Exception as search_error: + verbose_logger.warning( + "Vector store search failed for vector_store_id=%s, continuing without its context: %s", + vector_store_id, + search_error, + ) + return SearchFailed( + failure=VectorStoreSearchFailure( + vector_store_id=vector_store_id, + custom_llm_provider=custom_llm_provider, + error=str(search_error), + ) + ) + + verbose_logger.debug( + "Vector store search completed for vector_store_id=%s. Added context from %s results", + vector_store_id, + len(search_response.get("data", []) or []), + ) + return SearchSucceeded(response=search_response) + + def _extract_query_from_messages(self, messages: Sequence[AllMessageValues]) -> str | None: """ Extract the query from the last user message. @@ -223,48 +298,40 @@ class VectorStorePreCallHook(CustomLogger): return None - def _append_search_results_to_messages( + def _messages_with_context( self, - messages: list[AllMessageValues], - search_response: VectorStoreSearchResponse, - ) -> list[AllMessageValues]: - """ - Append search results as context to the messages. + messages: Sequence[AllMessageValues], + search_results: Sequence[VectorStoreSearchResponse], + ) -> tuple[AllMessageValues, ...]: + context_messages: Final = tuple( + context_message + for search_response in search_results + if (context_message := self._context_message(search_response)) is not None + ) + if not context_messages: + return tuple(messages) + return (*messages[:-1], *context_messages, *messages[-1:]) - Args: - messages: Original list of messages - search_response: Response from vector store search - - Returns: - Modified list of messages with context appended - """ - search_response_data: Final[list[VectorStoreSearchResult] | None] = search_response.get("data") + def _context_message(self, search_response: VectorStoreSearchResponse) -> AllMessageValues | None: + """Build the context message for one vector store's results, or None when it returned nothing usable.""" + search_response_data: Final[Sequence[VectorStoreSearchResult] | None] = search_response.get("data") if not search_response_data: - return messages + return None - context_content = self.CONTENT_PREFIX_STRING + context_texts: Final = tuple( + content_text + for result in search_response_data + for content_item in (result.get("content") or ()) + if (content_text := content_item.get("text")) + ) + if not context_texts: + return None - for result in search_response_data: - result_content: list[VectorStoreResultContent] | None = result.get("content") - if result_content: - for content_item in result_content: - content_text: str | None = content_item.get("text") - if content_text: - context_content += content_text + "\n\n" - - # Only add context if we found any content - if context_content != "Context:\n\n": - # Create a copy of messages to avoid modifying the original - modified_messages: Final = messages.copy() - # Add context as a new message before the last user message - context_message: Final[ChatCompletionUserMessage] = { - "role": "user", - "content": context_content, - } - modified_messages.insert(-1, cast(AllMessageValues, context_message)) - return modified_messages - - return messages + context_message: Final[ChatCompletionUserMessage] = { + "role": "user", + "content": self.CONTENT_PREFIX_STRING + "".join(f"{text}\n\n" for text in context_texts), + } + return cast(AllMessageValues, context_message) async def async_post_call_success_deployment_hook( self, @@ -287,34 +354,29 @@ class VectorStorePreCallHook(CustomLogger): verbose_logger.debug("No litellm_logging_obj in request_data") return None - verbose_logger.debug("model_call_details keys: %s", list(litellm_logging_obj.model_call_details.keys())) - # Get search results from model_call_details (already in OpenAI format) - search_results: Final[list[VectorStoreSearchResponse] | None] = litellm_logging_obj.model_call_details.get( - "search_results" + search_results: Final[Sequence[VectorStoreSearchResponse] | None] = ( + litellm_logging_obj.model_call_details.get("search_results") + ) + search_failures: Final[Sequence[VectorStoreSearchFailure] | None] = ( + litellm_logging_obj.model_call_details.get(SEARCH_FAILURES_FIELD) ) - verbose_logger.debug("Search results found: %s", search_results is not None) - - if not search_results: - verbose_logger.debug("No search results found") + if not search_results and not search_failures: + verbose_logger.debug("No search results or search failures found") return None # Add search results to response object if hasattr(response, "choices") and response.choices: for choice in response.choices: if hasattr(choice, "message") and choice.message: - # Get existing provider_specific_fields or create new dict provider_fields = getattr(choice.message, "provider_specific_fields", None) or {} - - # Add search results (already in OpenAI-compatible format) - provider_fields["search_results"] = search_results - - # Set the provider_specific_fields + if search_results: + provider_fields["search_results"] = search_results + if search_failures: + provider_fields[SEARCH_FAILURES_FIELD] = search_failures setattr(choice.message, "provider_specific_fields", provider_fields) - verbose_logger.debug("Added %s search results to response", len(search_results)) - # Return modified response return response @@ -339,29 +401,24 @@ class VectorStorePreCallHook(CustomLogger): verbose_logger.debug("VectorStorePreCallHook.async_post_call_streaming_deployment_hook called") # Get search results from model_call_details (already in OpenAI format) - search_results: Final[list[VectorStoreSearchResponse] | None] = request_data.get("search_results") + search_results: Final[Sequence[VectorStoreSearchResponse] | None] = request_data.get("search_results") + search_failures: Final[Sequence[VectorStoreSearchFailure] | None] = request_data.get(SEARCH_FAILURES_FIELD) - verbose_logger.debug("Search results found for streaming chunk: %s", search_results is not None) - - if not search_results: - verbose_logger.debug("No search results found for streaming chunk") + if not search_results and not search_failures: + verbose_logger.debug("No search results or search failures found for streaming chunk") return response_chunk # Add search results to streaming chunk if hasattr(response_chunk, "choices") and response_chunk.choices: for choice in response_chunk.choices: if hasattr(choice, "delta") and choice.delta: - # Get existing provider_specific_fields or create new dict provider_fields = getattr(choice.delta, "provider_specific_fields", None) or {} - - # Add search results (already in OpenAI-compatible format) - provider_fields["search_results"] = search_results - - # Set the provider_specific_fields + if search_results: + provider_fields["search_results"] = search_results + if search_failures: + provider_fields[SEARCH_FAILURES_FIELD] = search_failures choice.delta.provider_specific_fields = provider_fields - verbose_logger.debug("Added %s search results to streaming chunk", len(search_results)) - # Return modified chunk return response_chunk @@ -369,3 +426,10 @@ class VectorStorePreCallHook(CustomLogger): verbose_logger.exception("Error adding search results to streaming chunk: %s", e) # Don't fail the request if search results fail to be added return response_chunk + + +def _requested_vector_store_ids(non_default_params: Mapping[str, object]) -> tuple[str, ...]: + requested: Final = non_default_params.get("vector_store_ids") + if not isinstance(requested, (list, tuple)): + return () + return tuple(str(vector_store_id) for vector_store_id in requested) diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 474c652ff3a..9051287b30a 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -4,7 +4,7 @@ from enum import Enum from typing import Any, Literal from pydantic import BaseModel -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict class SupportedVectorStoreIntegrations(str, Enum): @@ -95,6 +95,17 @@ class VectorStoreSearchResponse(TypedDict, total=False): data: list[VectorStoreSearchResult] | None +VectorStoreSearchFailureMode = Literal["annotate", "error"] + + +class VectorStoreSearchFailure(TypedDict): + """A configured vector store whose search failed, as reported back to the API caller""" + + vector_store_id: ReadOnly[str] + custom_llm_provider: ReadOnly[str | None] + error: ReadOnly[str] + + class VectorStoreSearchOptionalRequestParams(TypedDict, total=False): """TypedDict for Optional parameters supported by the vector store search API.""" diff --git a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py index ae5cffd8ab0..64ee9474e44 100644 --- a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py +++ b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py @@ -12,6 +12,15 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i VectorStorePreCallHook, ) from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ( + CallTypes, + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) from litellm.types.vector_stores import ( VectorStoreResultContent, VectorStoreSearchResponse, @@ -36,6 +45,18 @@ def _search_response(text: str) -> VectorStoreSearchResponse: ) +def _first_message(response: ModelResponse) -> Message: + choice = response.choices[0] + assert isinstance(choice, Choices) + return choice.message + + +@dataclass(frozen=True) +class ExplodingRegistry: + async def pop_vector_stores_to_run_with_db_fallback(self, **kwargs: object) -> list[LiteLLM_ManagedVectorStore]: + raise RuntimeError("the registry blew up") + + @dataclass class RecordingRouter: failing_vector_store_ids: frozenset[str] = frozenset() @@ -285,3 +306,181 @@ def test_the_default_runtime_follows_the_proxy_globals(monkeypatch: pytest.Monke assert runtime.llm_router() is router assert runtime.prisma_client() is prisma + + +@pytest.mark.asyncio +async def test_a_failing_vector_store_is_reported_back_to_the_caller( + registry_with: RegisterStores, +) -> None: + """Regression (LIT-6809): a silently dropped store left the caller with an un-augmented answer and no signal.""" + registry_with("vs-broken", "vs-healthy") + logging_obj = FakeLoggingObj({}) + + await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))) + ), + ["vs-broken", "vs-healthy"], + logging_obj, + ) + + response = ModelResponse(choices=[Choices(message=Message(content="an answer"))]) + await VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)).async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": logging_obj}, + response=response, + call_type=CallTypes.acompletion, + ) + + provider_specific_fields = _first_message(response).provider_specific_fields or {} + assert provider_specific_fields["vector_store_search_failures"] == ( + { + "vector_store_id": "vs-broken", + "custom_llm_provider": "bedrock", + "error": "litellm.BadRequestError: no healthy deployments for vs-broken", + }, + ) + assert len(provider_specific_fields["search_results"]) == 1 + + +@pytest.mark.asyncio +async def test_a_healthy_vector_store_alone_reports_no_failures(registry_with: RegisterStores) -> None: + registry_with("vs-healthy") + logging_obj = FakeLoggingObj({}) + + await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=RecordingRouter())), + ["vs-healthy"], + logging_obj, + ) + + response = ModelResponse(choices=[Choices(message=Message(content="an answer"))]) + await VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)).async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": logging_obj}, + response=response, + call_type=CallTypes.acompletion, + ) + + assert "vector_store_search_failures" not in (_first_message(response).provider_specific_fields or {}) + + +@pytest.mark.asyncio +async def test_a_failing_vector_store_is_reported_on_the_streaming_chunk(registry_with: RegisterStores) -> None: + registry_with("vs-broken") + logging_obj = FakeLoggingObj({}) + + await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))) + ), + ["vs-broken"], + logging_obj, + ) + + chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content="an answer"))]) + await VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=None) + ).async_post_call_streaming_deployment_hook( + request_data=logging_obj.model_call_details, + response_chunk=chunk, + call_type=CallTypes.acompletion, + ) + + assert (chunk.choices[0].delta.provider_specific_fields or {})["vector_store_search_failures"] == ( + { + "vector_store_id": "vs-broken", + "custom_llm_provider": "bedrock", + "error": "litellm.BadRequestError: no healthy deployments for vs-broken", + }, + ) + + +@pytest.mark.asyncio +async def test_error_mode_fails_the_request_instead_of_answering_without_the_knowledge_base( + registry_with: RegisterStores, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression (LIT-6809): opting in must turn an ungrounded answer into a 400 the caller can act on.""" + registry_with("vs-broken", "vs-healthy") + monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "error") + + with pytest.raises(litellm.VectorStoreSearchError) as raised: + await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime( + router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"})) + ) + ), + ["vs-broken", "vs-healthy"], + FakeLoggingObj({}), + ) + + assert raised.value.status_code == 400 + assert raised.value.failures == ( + { + "vector_store_id": "vs-broken", + "custom_llm_provider": "bedrock", + "error": "litellm.BadRequestError: no healthy deployments for vs-broken", + }, + ) + assert "vs-broken: litellm.BadRequestError: no healthy deployments for vs-broken" in raised.value.message + + +@pytest.mark.asyncio +async def test_error_mode_leaves_a_fully_healthy_request_alone( + registry_with: RegisterStores, + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry_with("vs-healthy") + monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "error") + + _, messages, _ = await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=RecordingRouter())), + ["vs-healthy"], + FakeLoggingObj({}), + ) + + assert messages[0]["content"] == "Context:\n\ncontext from vs-healthy\n\n" + + +@pytest.mark.asyncio +async def test_error_mode_does_not_swallow_the_raise_in_the_hooks_own_catch_all( + registry_with: RegisterStores, + monkeypatch: pytest.MonkeyPatch, + warnings: list[logging.LogRecord], +) -> None: + """Regression (LIT-6809): the catch-all around the hook must not turn the opted-in failure back into a 200.""" + registry_with("vs-broken") + monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "error") + + with pytest.raises(litellm.VectorStoreSearchError): + await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime( + router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"})) + ) + ), + ["vs-broken"], + FakeLoggingObj({}), + ) + + assert [record.levelname for record in warnings] == ["WARNING"] + + +@pytest.mark.asyncio +async def test_a_crash_outside_the_search_names_the_requested_vector_stores( + monkeypatch: pytest.MonkeyPatch, + warnings: list[logging.LogRecord], +) -> None: + """Regression (LIT-6809): the catch-all logged no store id, so an operator could not tell which store broke.""" + monkeypatch.setattr(litellm, "vector_store_registry", ExplodingRegistry()) + + _, messages, _ = await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)), + ["vs-one", "vs-two"], + FakeLoggingObj({}), + ) + + assert messages == [{"role": "user", "content": "what is litellm?"}] + assert [record.getMessage() for record in warnings] == [ + "Error in VectorStorePreCallHook for vector_store_ids=('vs-one', 'vs-two'): the registry blew up" + ] From 6499ca349f1cfe5bc4e92d6128a11c13aea3420c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:11:22 -0700 Subject: [PATCH 004/155] fix(vector-stores): import assert_never from typing_extensions for Python 3.10 --- .../vector_store_integrations/vector_store_pre_call_hook.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index d5c893c4b20..ed125b915b4 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -7,7 +7,9 @@ It searches the vector store for relevant context and appends it to the messages from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, Protocol, assert_never, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast + +from typing_extensions import assert_never import litellm import litellm.vector_stores From e464e749e00ea680dbc59702218210ea772a59e6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:23:28 -0700 Subject: [PATCH 005/155] fix(vector-stores): fall back to annotate when the failure mode is unrecognized litellm_settings keys are set on the litellm module with no allowlist, so a typo in vector_store_search_failure_mode reached assert_never and turned every vector-store request into a 500. Validate the configured value and fall back to the permissive default with a warning naming the supported modes. --- .../vector_store_pre_call_hook.py | 21 +++++++++++-- .../test_vector_store_pre_call_hook.py | 30 +++++++++++++++++++ 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index ed125b915b4..274cb496b8f 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -7,8 +7,9 @@ It searches the vector store for relevant context and appends it to the messages from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_args +from pydantic import TypeAdapter, ValidationError from typing_extensions import assert_never import litellm @@ -22,6 +23,7 @@ from litellm.types.utils import CallTypes, StandardCallbackDynamicParams from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, VectorStoreSearchFailure, + VectorStoreSearchFailureMode, VectorStoreSearchResponse, VectorStoreSearchResult, ) @@ -34,6 +36,8 @@ else: LiteLLMLoggingObj = Any SEARCH_FAILURES_FIELD: Final = "vector_store_search_failures" +_DEFAULT_FAILURE_MODE: Final[VectorStoreSearchFailureMode] = "annotate" +_FAILURE_MODE_ADAPTER: Final = TypeAdapter(VectorStoreSearchFailureMode) class ProxyRuntime(Protocol): @@ -153,7 +157,7 @@ class VectorStorePreCallHook(CustomLogger): litellm_logging_obj.model_call_details[detail] = value if augmentation.failures: - match litellm.vector_store_search_failure_mode: + match _configured_failure_mode(): case "error": raise VectorStoreSearchError(failures=augmentation.failures, model=model) case "annotate": @@ -435,3 +439,16 @@ def _requested_vector_store_ids(non_default_params: Mapping[str, object]) -> tup if not isinstance(requested, (list, tuple)): return () return tuple(str(vector_store_id) for vector_store_id in requested) + + +def _configured_failure_mode() -> VectorStoreSearchFailureMode: + try: + return _FAILURE_MODE_ADAPTER.validate_python(litellm.vector_store_search_failure_mode) + except ValidationError: + verbose_logger.warning( + "Unsupported vector_store_search_failure_mode=%r, falling back to %r. Supported modes: %s", + litellm.vector_store_search_failure_mode, + _DEFAULT_FAILURE_MODE, + ", ".join(get_args(VectorStoreSearchFailureMode)), + ) + return _DEFAULT_FAILURE_MODE diff --git a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py index 64ee9474e44..eae53becc14 100644 --- a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py +++ b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py @@ -425,6 +425,36 @@ async def test_error_mode_fails_the_request_instead_of_answering_without_the_kno assert "vs-broken: litellm.BadRequestError: no healthy deployments for vs-broken" in raised.value.message +@pytest.mark.asyncio +async def test_a_misspelled_failure_mode_annotates_instead_of_erroring_the_request( + registry_with: RegisterStores, + monkeypatch: pytest.MonkeyPatch, + warnings: list[logging.LogRecord], +) -> None: + """Regression (LIT-6809): litellm_settings takes any value, so a typo must not become a 500.""" + registry_with("vs-broken") + monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "erorr") + + logging_obj = FakeLoggingObj({}) + _, messages, _ = await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))) + ), + ["vs-broken"], + logging_obj, + ) + + assert messages[0]["content"] == "what is litellm?" + assert logging_obj.model_call_details["vector_store_search_failures"] == ( + { + "vector_store_id": "vs-broken", + "custom_llm_provider": "bedrock", + "error": "litellm.BadRequestError: no healthy deployments for vs-broken", + }, + ) + assert any("erorr" in record.getMessage() for record in warnings) + + @pytest.mark.asyncio async def test_error_mode_leaves_a_fully_healthy_request_alone( registry_with: RegisterStores, From 2647c2890fc6b8a32ab5601ecd8ed04b1cc0ebff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:25:32 -0700 Subject: [PATCH 006/155] fix(oci): keep the stream's own finish reason instead of a synthetic stop chunk The OCI wrapper overrides chunk_creator wholesale, so it never recorded the finish reason or marked the terminal chunk as sent. The shared end-of-stream finalizer then appended a synthetic chunk whose finish_reason was always stop, which downgraded a tool_calls completion for any client that reads the finish reason off the last chunk. --- litellm/llms/oci/chat/transformation.py | 12 +++-- .../oci/chat/test_oci_chat_transformation.py | 54 +++++++++++++++++++ 2 files changed, 63 insertions(+), 3 deletions(-) diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index b1167d447fc..8e4e41b4ac1 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -745,7 +745,13 @@ class OCIStreamWrapper(CustomStreamWrapper): # single-event case (terminal chunk carries the only copy of the text). self._cohere_text_emitted = False - def _with_stream_identity(self, parsed: ModelResponseStream) -> ModelResponseStream: + def _emit_chunk(self, parsed: ModelResponseStream) -> ModelResponseStream: + for choice in parsed.choices: + if getattr(choice.delta, "tool_calls", None): + self.tool_call = True + if choice.finish_reason is not None: + self.received_finish_reason = choice.finish_reason + self.sent_last_chunk = True return self.model_response_creator(chunk={"choices": parsed.choices}) def chunk_creator(self, chunk: Any) -> ModelResponseStream | None: @@ -780,8 +786,8 @@ class OCIStreamWrapper(CustomStreamWrapper): if getattr(choice.delta, "content", None): self._cohere_text_emitted = True break - return self._with_stream_identity(result) - return self._with_stream_identity(handle_generic_stream_chunk(dict_chunk)) + return self._emit_chunk(result) + return self._emit_chunk(handle_generic_stream_chunk(dict_chunk)) __all__ = [ diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py index e77b2c24d01..4c9bd29b337 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py @@ -2041,3 +2041,57 @@ class TestOCIStreamWrapperDoneSentinel: wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) with pytest.raises(OCIError, match="Chunk cannot be parsed as JSON"): wrapper.chunk_creator("data: [DONE] trailing garbage") + + +_GENERIC_TOOL_CALL_EVENT = ( + 'data: {"index":0,"message":{"role":"ASSISTANT","content":[],' + '"toolCalls":[{"type":"FUNCTION","id":"call_1","name":"get_weather","arguments":"{}"}]}}' +) +_GENERIC_TOOL_TERMINAL_EVENT = ( + 'data: {"index":0,"message":{"role":"ASSISTANT","content":[]},"finishReason":"TOOL_CALLS"}' +) + + +def _drain_stream(model: str, events: list[str]) -> list: + logging_obj = MagicMock() + logging_obj.model_call_details = {"custom_llm_provider": "oci", "litellm_params": {}} + wrapper = OCIStreamWrapper( + completion_stream=iter(events), + model=model, + custom_llm_provider="oci", + logging_obj=logging_obj, + ) + return list(wrapper) + + +class TestOCIStreamWrapperTerminalChunk: + """OCI's ``chunk_creator`` override bypasses the shared handler's + finish-reason bookkeeping, so the shared end-of-stream finalizer used to + append a synthetic ``stop`` chunk after OCI's own terminal chunk, silently + downgrading a ``tool_calls`` completion for any client that reads the + finish reason off the last chunk.""" + + def test_generic_tool_call_stream_ends_on_tool_calls(self): + chunks = _drain_stream( + _STREAM_GENERIC_MODEL, + [_GENERIC_TOOL_CALL_EVENT, _GENERIC_TOOL_TERMINAL_EVENT, "data: [DONE]"], + ) + + assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "tool_calls"] + assert len({chunk.id for chunk in chunks}) == 1 + + def test_generic_text_stream_emits_exactly_one_finish_reason(self): + chunks = _drain_stream( + _STREAM_GENERIC_MODEL, + [_GENERIC_TEXT_EVENT.format(text="1"), _GENERIC_TERMINAL_EVENT, "data: [DONE]"], + ) + + assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "stop"] + + def test_cohere_stream_emits_exactly_one_finish_reason(self): + chunks = _drain_stream( + _STREAM_COHERE_MODEL, + [_COHERE_TEXT_EVENT.format(text="123"), _COHERE_TERMINAL_EVENT], + ) + + assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "stop"] From b55833b1f9269753fa35f5f6e68fc9782b83936f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:40:00 -0700 Subject: [PATCH 007/155] fix(ai-gateway): build the release image again and cover it in CI The image had two independent breaks. The Dockerfile pinned rust 1.90 while the repo pins 1.98 in rust-toolchain.toml and never copied it in, so the first cargo call died on crates needing a newer rustc. The runtime stage then ran pip install on the root pyproject, which builds with maturin against the python-bridge crate, so metadata generation failed with no Cargo manifest and no Rust toolchain in that stage. Copy rust-toolchain.toml into the builder so every cargo call uses the pinned channel, build the wheel in the builder stage where cargo and python3-dev already live, and have the runtime stage install that artifact instead of compiling anything. Add the ai-gateway image job to the rust workflow so a broken build fails a PR instead of surfacing on a release. --- .github/workflows/test-rust.yml | 39 +++++++++++++++++++ litellm-rust/crates/ai-gateway/Dockerfile | 31 ++++++++++----- .../crates/ai-gateway/Dockerfile.dockerignore | 6 ++- 3 files changed, 66 insertions(+), 10 deletions(-) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 9b8b132df62..99de9a25fbe 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -133,3 +133,42 @@ jobs: - name: Test native route wheel run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl + + ai-gateway-image: + name: ai-gateway release image + runs-on: ubuntu-latest + timeout-minutes: 60 + permissions: + contents: read + + steps: + - name: Checkout repository + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Build the release image + run: docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway:${{ github.sha }} . + + - name: Start the gateway and wait for readiness + env: + IMAGE: litellm-ai-gateway:${{ github.sha }} + run: | + docker run -d --name ai-gateway -p 4001:4001 \ + -e LITELLM_MASTER_KEY=sk-ci-not-a-real-key \ + -e OPENAI_API_KEY=sk-ci-not-a-real-key \ + "$IMAGE" + for _ in $(seq 1 60); do + if curl -fsS http://127.0.0.1:4001/health/readiness; then + echo "gateway is serving readiness" + exit 0 + fi + sleep 2 + done + echo "gateway never became ready" >&2 + docker logs ai-gateway >&2 + exit 1 + + - name: Stop the gateway + if: always() + run: docker rm -f ai-gateway || true diff --git a/litellm-rust/crates/ai-gateway/Dockerfile b/litellm-rust/crates/ai-gateway/Dockerfile index 2bc3c05ad7e..370d6b3f08f 100644 --- a/litellm-rust/crates/ai-gateway/Dockerfile +++ b/litellm-rust/crates/ai-gateway/Dockerfile @@ -14,15 +14,20 @@ # ---- Chef ------------------------------------------------------------------- # cargo-chef caches the dependency build so only the gateway crate recompiles on # a source-only change. python3-dev is present in every rust stage because the -# `python-config` feature links libpython via pyo3 (even in the cook step). -FROM rust:1.90-slim-bookworm AS chef +# `python-config` feature links libpython via pyo3 (even in the cook step), and +# python3-pip builds the litellm wheel in the builder stage. +FROM rust:1.98-slim-bookworm AS chef ENV PYO3_PYTHON=python3.11 +# rustup reads rust-toolchain.toml from any parent of the working directory, so +# copying it in is what keeps every cargo call below on the repo's pinned +# channel rather than on whatever the base image happens to ship. +COPY rust-toolchain.toml /build/rust-toolchain.toml +WORKDIR /build/litellm-rust RUN apt-get update \ && apt-get install -y --no-install-recommends \ - python3 python3-dev pkg-config libssl-dev clang \ + python3 python3-dev python3-pip pkg-config libssl-dev clang \ && rm -rf /var/lib/apt/lists/* \ && cargo install cargo-chef --locked --version 0.1.77 -WORKDIR /build/litellm-rust # ---- Planner ---------------------------------------------------------------- # Produce the dependency recipe from the rust workspace manifests + Cargo.lock. @@ -43,6 +48,14 @@ RUN cargo chef cook --locked --release \ COPY litellm-rust/ . RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gateway --features server,python-config +# The root pyproject builds with maturin against litellm-rust/crates/python-bridge, +# so the wheel is built here, next to the crate sources and the cargo toolchain, +# and the runtime stage installs the artifact instead of compiling anything. +COPY pyproject.toml README.md LICENSE /build/ +COPY litellm/ /build/litellm/ +COPY enterprise/ /build/enterprise/ +RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /build/dist /build + # ---- Runtime ---------------------------------------------------------------- # python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3 # 3.11 ABI so the embedded interpreter links and imports cleanly. @@ -56,11 +69,11 @@ RUN apt-get update \ WORKDIR /app # Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so -# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. Copy the -# package + packaging metadata, then pip install the proxy extra. -COPY pyproject.toml README.md LICENSE ./ -COPY litellm/ ./litellm/ -RUN pip install --no-cache-dir ".[proxy]" +# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. +COPY --from=builder /build/dist/*.whl /tmp/wheels/ +RUN wheel="$(ls /tmp/wheels/litellm-*.whl)" \ + && pip install --no-cache-dir "${wheel}[proxy]" \ + && rm -rf /tmp/wheels # The compiled gateway binary (pure-Rust realtime hot path; Python is load-time # only). diff --git a/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore b/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore index 030ee6a37c5..4fec13d18db 100644 --- a/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore +++ b/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore @@ -9,13 +9,17 @@ # Strategy: ignore everything, then re-include only what the build needs: # - litellm/ (pip install . needs the full package + proxy reader) # - litellm-rust/ (the rust workspace; Cargo.lock + crate sources) -# - pyproject.toml / README.md / LICENSE (packaging metadata for pip install) +# - enterprise/ (litellm/proxy/enterprise symlinks into it; maturin walks it) +# - pyproject.toml / README.md / LICENSE (packaging metadata for the wheel build) +# - rust-toolchain.toml (the pinned channel every cargo call in the build uses) * # --- re-include the build inputs --- !litellm/ !litellm-rust/ +!enterprise/ !pyproject.toml +!rust-toolchain.toml !README.md !LICENSE From fbb799e240eec1e4d8c41d9c905d3699b7bfaad7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:44:20 -0700 Subject: [PATCH 008/155] ci: drop the ai-gateway image from the coverage allowlist now that a job builds it --- .github/ci-coverage-allowlist.yml | 5 ----- 1 file changed, 5 deletions(-) diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index f7a785b3b80..810265f4b40 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -96,11 +96,6 @@ dockerfiles: and lint workflows already exercise that output, so building the image adds no signal about it paths: - ui/Dockerfile - - reason: >- - The Rust gateway ships as its own chart and package with a separate release pipeline, so its - image is not part of this repo's Python image set - paths: - - litellm-rust/crates/ai-gateway/Dockerfile - reason: >- An example image under cookbook/ that is documentation rather than a shipped artifact paths: From 1d1eb4264f2ea33beeda9f4a4b84c45ce35ebb0e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:49:40 -0700 Subject: [PATCH 009/155] build(ai-gateway): keep the committed enterprise wheels out of the build context --- litellm-rust/crates/ai-gateway/Dockerfile.dockerignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore b/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore index 4fec13d18db..c8e789b5e8c 100644 --- a/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore +++ b/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore @@ -26,6 +26,8 @@ # --- prune heavy / irrelevant subpaths back out of the re-included trees --- # Rust build artifacts (huge; regenerated in the builder). **/target/ +# Committed python distribution artifacts; the wheel build does not read them. +enterprise/dist/ # Python caches and compiled bytecode. **/__pycache__/ **/*.pyc From dc28f0eb4b9aeee5933368055daeb3a09d5fbab1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:54:09 -0700 Subject: [PATCH 010/155] fix(vector-stores): report search failures on the responses API surface too --- .../vector_store_pre_call_hook.py | 11 +++- .../test_vector_store_pre_call_hook.py | 55 ++++++++++++++++++- 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 274cb496b8f..849e87fae40 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -17,7 +17,11 @@ import litellm.vector_stores from litellm._logging import verbose_logger from litellm.exceptions import VectorStoreSearchError from litellm.integrations.custom_logger import CustomLogger -from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionUserMessage, + ResponsesAPIResponse, +) from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.utils import CallTypes, StandardCallbackDynamicParams from litellm.types.vector_stores import ( @@ -372,6 +376,11 @@ class VectorStorePreCallHook(CustomLogger): verbose_logger.debug("No search results or search failures found") return None + if isinstance(response, ResponsesAPIResponse): + if search_failures: + setattr(response, SEARCH_FAILURES_FIELD, list(search_failures)) + return response + # Add search results to response object if hasattr(response, "choices") and response.choices: for choice in response.choices: diff --git a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py index eae53becc14..f1f9f7c3f3f 100644 --- a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py +++ b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py @@ -11,7 +11,7 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i ProxyServerRuntime, VectorStorePreCallHook, ) -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, Choices, @@ -363,6 +363,59 @@ async def test_a_healthy_vector_store_alone_reports_no_failures(registry_with: R assert "vector_store_search_failures" not in (_first_message(response).provider_specific_fields or {}) +@pytest.mark.asyncio +async def test_a_failing_vector_store_is_reported_on_the_responses_api_response( + registry_with: RegisterStores, +) -> None: + """Regression (LIT-6809): /v1/responses answered 200 with no sign the knowledge base was missing.""" + registry_with("vs-broken") + logging_obj = FakeLoggingObj({}) + + await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))) + ), + ["vs-broken"], + logging_obj, + ) + + response = ResponsesAPIResponse(id="resp-lit6809", created_at=0, output=[]) + await VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)).async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": logging_obj}, + response=response, + call_type=CallTypes.aresponses, + ) + + assert response.model_dump()["vector_store_search_failures"] == [ + { + "vector_store_id": "vs-broken", + "custom_llm_provider": "bedrock", + "error": "litellm.BadRequestError: no healthy deployments for vs-broken", + } + ] + + +@pytest.mark.asyncio +async def test_a_healthy_vector_store_leaves_the_responses_api_response_alone(registry_with: RegisterStores) -> None: + registry_with("vs-healthy") + logging_obj = FakeLoggingObj({}) + + await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=RecordingRouter())), + ["vs-healthy"], + logging_obj, + ) + + response = ResponsesAPIResponse(id="resp-lit6809", created_at=0, output=[]) + await VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)).async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": logging_obj}, + response=response, + call_type=CallTypes.aresponses, + ) + + assert "vector_store_search_failures" not in response.model_dump() + + @pytest.mark.asyncio async def test_a_failing_vector_store_is_reported_on_the_streaming_chunk(registry_with: RegisterStores) -> None: registry_with("vs-broken") From 376f74c058cdaa00cf5abcda1ba83fb23615c801 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:26:45 -0700 Subject: [PATCH 011/155] refactor(vector_stores): bind the failure mode before matching on it basedpyright counts a named capture after an exhaustive match under reportUnnecessaryComparison, which pushed the merged tree one over the budget; the wildcard case with assert_never on the bound subject is the shape the rest of the codebase uses --- .../vector_store_pre_call_hook.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 849e87fae40..636e076a229 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -161,13 +161,14 @@ class VectorStorePreCallHook(CustomLogger): litellm_logging_obj.model_call_details[detail] = value if augmentation.failures: - match _configured_failure_mode(): + failure_mode: Final = _configured_failure_mode() + match failure_mode: case "error": raise VectorStoreSearchError(failures=augmentation.failures, model=model) case "annotate": pass - case unreachable: - assert_never(unreachable) + case _: + assert_never(failure_mode) return model, list(augmentation.messages), non_default_params From c4d09a31e0631f03c87fe27c1fd6dcfce16e2c6c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:20:07 -0700 Subject: [PATCH 012/155] fix(vector_stores): default the search-count debug log to a tuple so LIT002 stays within budget --- .../vector_store_integrations/vector_store_pre_call_hook.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 636e076a229..b2243060c6c 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -276,7 +276,7 @@ class VectorStorePreCallHook(CustomLogger): verbose_logger.debug( "Vector store search completed for vector_store_id=%s. Added context from %s results", vector_store_id, - len(search_response.get("data", []) or []), + len(search_response.get("data") or ()), ) return SearchSucceeded(response=search_response) From 025b9cb75168233ef06f17078e7159d6b3ae6f08 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:52:10 -0700 Subject: [PATCH 013/155] build(ai-gateway): build the sibling wheels from the repo so the image never waits on PyPI --- litellm-rust/crates/ai-gateway/Dockerfile | 16 +++++++++++++--- .../crates/ai-gateway/Dockerfile.dockerignore | 3 +++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/Dockerfile b/litellm-rust/crates/ai-gateway/Dockerfile index 370d6b3f08f..72ac25ce1d6 100644 --- a/litellm-rust/crates/ai-gateway/Dockerfile +++ b/litellm-rust/crates/ai-gateway/Dockerfile @@ -51,10 +51,15 @@ RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gatewa # The root pyproject builds with maturin against litellm-rust/crates/python-bridge, # so the wheel is built here, next to the crate sources and the cargo toolchain, # and the runtime stage installs the artifact instead of compiling anything. +# litellm[proxy] pins litellm-enterprise and litellm-proxy-extras to the versions +# in this repo, and those hit PyPI hours after every version bump merges, so both +# wheels are built from the repo too instead of being resolved from PyPI. COPY pyproject.toml README.md LICENSE /build/ COPY litellm/ /build/litellm/ COPY enterprise/ /build/enterprise/ -RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /build/dist /build +COPY litellm-proxy-extras/ /build/litellm-proxy-extras/ +RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /build/dist \ + /build /build/enterprise /build/litellm-proxy-extras # ---- Runtime ---------------------------------------------------------------- # python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3 @@ -69,10 +74,15 @@ RUN apt-get update \ WORKDIR /app # Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so -# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. +# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. The two +# sibling wheels come from the builder as well, so the pins in litellm[proxy] +# resolve against them and never wait on a PyPI publish. COPY --from=builder /build/dist/*.whl /tmp/wheels/ RUN wheel="$(ls /tmp/wheels/litellm-*.whl)" \ - && pip install --no-cache-dir "${wheel}[proxy]" \ + && pip install --no-cache-dir \ + /tmp/wheels/litellm_enterprise-*.whl \ + /tmp/wheels/litellm_proxy_extras-*.whl \ + "${wheel}[proxy]" \ && rm -rf /tmp/wheels # The compiled gateway binary (pure-Rust realtime hot path; Python is load-time diff --git a/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore b/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore index c8e789b5e8c..d1386ff684d 100644 --- a/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore +++ b/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore @@ -10,6 +10,7 @@ # - litellm/ (pip install . needs the full package + proxy reader) # - litellm-rust/ (the rust workspace; Cargo.lock + crate sources) # - enterprise/ (litellm/proxy/enterprise symlinks into it; maturin walks it) +# - litellm-proxy-extras/ (built into a wheel alongside enterprise/ for litellm[proxy]) # - pyproject.toml / README.md / LICENSE (packaging metadata for the wheel build) # - rust-toolchain.toml (the pinned channel every cargo call in the build uses) * @@ -18,6 +19,7 @@ !litellm/ !litellm-rust/ !enterprise/ +!litellm-proxy-extras/ !pyproject.toml !rust-toolchain.toml !README.md @@ -28,6 +30,7 @@ **/target/ # Committed python distribution artifacts; the wheel build does not read them. enterprise/dist/ +litellm-proxy-extras/dist/ # Python caches and compiled bytecode. **/__pycache__/ **/*.pyc From d1f76e81d25b47887e85204ad6a24c103e8c79c0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 17:14:32 -0700 Subject: [PATCH 014/155] ci(ai-gateway): build the release image in its own workflow and assert the config loads --- .github/workflows/ai-gateway-image.yml | 73 ++++++++++++++++++++++++++ .github/workflows/test-rust.yml | 39 -------------- 2 files changed, 73 insertions(+), 39 deletions(-) create mode 100644 .github/workflows/ai-gateway-image.yml diff --git a/.github/workflows/ai-gateway-image.yml b/.github/workflows/ai-gateway-image.yml new file mode 100644 index 00000000000..3f690f566b0 --- /dev/null +++ b/.github/workflows/ai-gateway-image.yml @@ -0,0 +1,73 @@ +name: ai-gateway image + +on: + push: + paths: + - "litellm-rust/**" + - "litellm/**" + - "enterprise/**" + - "litellm-proxy-extras/**" + - "pyproject.toml" + - "rust-toolchain.toml" + - ".github/workflows/ai-gateway-image.yml" + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" + paths: + - "litellm-rust/**" + - "litellm/**" + - "enterprise/**" + - "litellm-proxy-extras/**" + - "pyproject.toml" + - "rust-toolchain.toml" + - ".github/workflows/ai-gateway-image.yml" + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + ai-gateway-image: + name: ai-gateway release image + runs-on: ubuntu-latest + timeout-minutes: 60 + permissions: + contents: read + steps: + - name: Checkout repository + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + - name: Build the release image + run: docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway:${{ github.sha }} . + - name: Start the gateway and wait for readiness + env: + IMAGE: litellm-ai-gateway:${{ github.sha }} + run: | + docker run -d --name ai-gateway -p 4001:4001 \ + -e LITELLM_MASTER_KEY=sk-ci-not-a-real-key \ + -e OPENAI_API_KEY=sk-ci-not-a-real-key \ + "$IMAGE" + for _ in $(seq 1 60); do + if curl -fsS http://127.0.0.1:4001/health/readiness; then + echo "gateway is serving readiness" + exit 0 + fi + sleep 2 + done + echo "gateway never became ready" >&2 + docker logs ai-gateway >&2 + exit 1 + - name: Assert the gateway loaded the baked config + run: | + docker logs ai-gateway 2>&1 | tee gateway.log + grep 'via python config reader' gateway.log + - name: Stop the gateway + if: always() + run: docker rm -f ai-gateway || true diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 99de9a25fbe..9b8b132df62 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -133,42 +133,3 @@ jobs: - name: Test native route wheel run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl - - ai-gateway-image: - name: ai-gateway release image - runs-on: ubuntu-latest - timeout-minutes: 60 - permissions: - contents: read - - steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Build the release image - run: docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway:${{ github.sha }} . - - - name: Start the gateway and wait for readiness - env: - IMAGE: litellm-ai-gateway:${{ github.sha }} - run: | - docker run -d --name ai-gateway -p 4001:4001 \ - -e LITELLM_MASTER_KEY=sk-ci-not-a-real-key \ - -e OPENAI_API_KEY=sk-ci-not-a-real-key \ - "$IMAGE" - for _ in $(seq 1 60); do - if curl -fsS http://127.0.0.1:4001/health/readiness; then - echo "gateway is serving readiness" - exit 0 - fi - sleep 2 - done - echo "gateway never became ready" >&2 - docker logs ai-gateway >&2 - exit 1 - - - name: Stop the gateway - if: always() - run: docker rm -f ai-gateway || true From 8c5b29519e6609d743913bc8d3c85070fc5cdf16 Mon Sep 17 00:00:00 2001 From: Jon Walton Date: Wed, 9 Sep 2026 18:15:53 +0800 Subject: [PATCH 015/155] fix(proxy): emit internal user budget alerts --- .../SlackAlerting/slack_alerting.py | 2 +- litellm/proxy/auth/auth_checks.py | 13 ++ .../SlackAlerting/test_slack_alerting.py | 15 ++ .../proxy/auth/test_auth_checks.py | 131 +++++++++++++++++- 4 files changed, 158 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index dc41c7dadc8..caac8e888fd 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -673,7 +673,7 @@ class SlackAlerting(CustomBatchLogger): Create a standard message for a budget alert """ _all_fields_as_dict: Final[dict[str, object]] = user_info.model_dump(exclude_none=True) - _all_fields_as_dict.pop("token") + _all_fields_as_dict.pop("token", None) msg = "" for k, v in _all_fields_as_dict.items(): if isinstance(v, Litellm_EntityType): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dc693317de0..d26474778d9 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1020,6 +1020,19 @@ async def common_checks( fallback_spend=user_object.spend or 0.0, max_budget=user_budget, ) + call_info: Final = CallInfo( + spend=user_spend, + max_budget=user_budget, + user_id=user_object.user_id, + user_email=user_object.user_email, + event_group=Litellm_EntityType.USER, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="user_budget", + user_info=call_info, + ) + ) if math.isfinite(user_budget) and user_spend >= user_budget: raise litellm.BudgetExceededError( current_cost=user_spend, diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 55e2dcdc270..44dda57dd27 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -45,6 +45,21 @@ class TestSlackAlerting(unittest.TestCase): result = self.slack_alerting._get_percent_of_max_budget_left(user_info) self.assertEqual(result, -0.2) + def test_get_user_info_str_omits_absent_token_for_user_alert(self): + user_info = CallInfo( + spend=85.0, + max_budget=100.0, + user_id="user-1", + user_email="person@example.com", + event_group=Litellm_EntityType.USER, + ) + + result = self.slack_alerting._get_user_info_str(user_info) + + self.assertIn("*user_id:* `user-1`", result) + self.assertIn("*user_email:* `person@example.com`", result) + self.assertNotIn("*token:*", result) + def test_get_event_and_event_message_max_budget(self): # Initial setup with no event event = None diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 2284a05b2e9..1fc2e0b3dc1 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,7 +1,7 @@ import asyncio import json from types import SimpleNamespace -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch if TYPE_CHECKING: @@ -29,6 +29,7 @@ from litellm.proxy._types import ( ProxyException, SSOUserDefinedValues, UserAPIKeyAuth, + WebhookEvent, ) from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, @@ -52,6 +53,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.constants import ( DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, @@ -5375,6 +5377,127 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): assert "User=u1" in str(over.value) +async def _run_internal_user_budget_alert( + *, + spend: float, +) -> tuple[AsyncMock, litellm.BudgetExceededError | None]: + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user: Final = LiteLLM_UserTable( + user_id="user-1", + user_email="person@example.com", + spend=0.0, + max_budget=100.0, + ) + token: Final = UserAPIKeyAuth(token="hashed-key-1", user_id="user-1") + slack_alerting: Final = SlackAlerting(alerting=["webhook"]) + send_alert: Final = AsyncMock() + alert_finished: Final = asyncio.Event() + + async def _get_spend( + counter_key: str, + fallback_spend: float, + max_budget: float | None = None, + **kwargs: object, + ) -> float: + assert counter_key == "spend:user:user-1" + assert fallback_spend == 0.0 + assert max_budget == 100.0 + return spend + + async def _budget_alerts( + *, + type: Literal["user_budget"], + user_info: CallInfo, + ) -> None: + try: + await slack_alerting.budget_alerts(type=type, user_info=user_info) + finally: + alert_finished.set() + + proxy_logging_obj: Final = MagicMock(budget_alerts=_budget_alerts) + + async def _check() -> bool: + return await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=proxy_logging_obj, + valid_token=token, + request=MagicMock(spec=Request), + ) + + async def _check_for_error() -> litellm.BudgetExceededError | None: + if spend < 100.0: + assert await _check() is True + return None + + with pytest.raises(litellm.BudgetExceededError) as raised: + await _check() + return raised.value + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: common_checks has no database seam + patch("litellm.proxy.proxy_server.get_current_spend", _get_spend), # test-quality-ok: common_checks imports it locally + patch.object(slack_alerting, "send_alert", send_alert), + ): + error: Final = await _check_for_error() + await asyncio.wait_for(alert_finished.wait(), timeout=1.0) + + return send_alert, error + + +@pytest.mark.asyncio +async def test_common_checks_internal_user_budget_below_threshold_does_not_emit_alert(): + send_alert, error = await _run_internal_user_budget_alert(spend=84.0) + + assert error is None + send_alert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_common_checks_internal_user_budget_emits_user_threshold_event(): + send_alert, error = await _run_internal_user_budget_alert(spend=85.0) + + assert error is None + send_alert.assert_awaited_once() + event: Final = send_alert.await_args.kwargs["user_info"] + assert isinstance(event, WebhookEvent) + assert event.event == "threshold_crossed" + assert event.event_group == Litellm_EntityType.USER + assert event.user_id == "user-1" + assert event.user_email == "person@example.com" + assert event.spend == 85.0 + assert event.max_budget == 100.0 + assert event.token is None + assert event.key_alias is None + assert event.team_id is None + assert event.organization_id is None + + +@pytest.mark.asyncio +async def test_common_checks_internal_user_budget_emits_crossed_event_and_rejects(): + send_alert, error = await _run_internal_user_budget_alert(spend=100.0) + + assert error is not None + assert error.current_cost == 100.0 + assert error.max_budget == 100.0 + send_alert.assert_awaited_once() + event: Final = send_alert.await_args.kwargs["user_info"] + assert isinstance(event, WebhookEvent) + assert event.event == "budget_crossed" + assert event.event_group == Litellm_EntityType.USER + assert event.user_id == "user-1" + assert event.user_email == "person@example.com" + + @pytest.mark.asyncio async def test_common_checks_personal_user_budget_skipped_for_team_key(): """A user's personal max_budget does not apply to a team-scoped key. @@ -5398,6 +5521,9 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key(): async def _no_membership(*args, **kwargs): return None + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5412,11 +5538,12 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key(): general_settings={}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) assert result is True + proxy_logging_obj.budget_alerts.assert_not_awaited() @pytest.mark.asyncio From 6d6659a81bd1b6d92c5e8c43383a37b95711aa68 Mon Sep 17 00:00:00 2001 From: Jon Walton Date: Wed, 9 Sep 2026 18:39:36 +0800 Subject: [PATCH 016/155] test(proxy): harden user budget alert coverage --- .../proxy/auth/test_auth_checks.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 1fc2e0b3dc1..f9e25e9beaf 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5356,6 +5356,9 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5370,10 +5373,11 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): general_settings={}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) + await asyncio.sleep(0) assert "User=u1" in str(over.value) @@ -5412,6 +5416,7 @@ async def _run_internal_user_budget_alert( type: Literal["user_budget"], user_info: CallInfo, ) -> None: + assert type == "user_budget" try: await slack_alerting.budget_alerts(type=type, user_info=user_info) finally: @@ -5568,6 +5573,9 @@ async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag async def _no_membership(*args, **kwargs): return None + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5583,10 +5591,11 @@ async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag general_settings={"apply_user_budget_to_team_keys": True}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) + await asyncio.sleep(0) assert "ExceededBudget: User=u1" in str(exc_info.value) @@ -5603,6 +5612,9 @@ async def test_common_checks_personal_user_budget_still_enforced_on_personal_key async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5617,10 +5629,11 @@ async def test_common_checks_personal_user_budget_still_enforced_on_personal_key general_settings={"apply_user_budget_to_team_keys": True}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) + await asyncio.sleep(0) @pytest.mark.parametrize( From 1425c71c10634bc86a2966ac4cfc977b7b9b39de Mon Sep 17 00:00:00 2001 From: jesus Date: Wed, 9 Sep 2026 19:00:24 +0000 Subject: [PATCH 017/155] fix(mcp): log upstream request method, body and response on MCP tool-list and OAuth2 token failures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_experimental/mcp_server/mcp_debug.py | 65 +++++++++++++++++++ .../mcp_server/mcp_server_manager.py | 12 +++- .../client_credentials.py | 6 ++ .../mcp_server/test_mcp_debug.py | 58 +++++++++++++++++ .../test_mcp_oauth_passthrough_tools.py | 27 +++++++- 5 files changed, 165 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index 6aaa00ae415..7a04901c2ca 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -85,11 +85,14 @@ Usage with curl:: http://localhost:4000/mcp/atlassian_mcp """ +import re from typing import TYPE_CHECKING, Final +import httpx from starlette.types import Message, Send from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker +from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -127,6 +130,10 @@ class MCPDebug: @staticmethod def _mask(value: str | None) -> str: """Mask a single value for safe display in headers.""" + return MCPDebug.mask_secret(value) + + @staticmethod + def mask_secret(value: str | None) -> str: if not value: return "(none)" return MCPDebug._masker._mask_value(value) @@ -311,3 +318,61 @@ class MCPDebug: server_url=server_url, server_auth_type=server_auth_type, ) + + +_BODY_PREVIEW_CHARS: Final = 512 +_SENSITIVE_HEADER_NAMES: Final = frozenset({"authorization", "proxy-authorization", "cookie", "x-api-key"}) +_SENSITIVE_BODY_FIELD: Final = re.compile( + r'(?P"?(?:client_secret|client_assertion|refresh_token|access_token|id_token|password|code)"?\s*[=:]\s*"?)' + r'(?P[^&"\s,}]+)' +) + + +def _mask_body_match(match: re.Match[str]) -> str: + return f"{match.group('key')}{MCPDebug.mask_secret(match.group('value'))}" + + +def _preview(raw: bytes) -> str: + text: Final = _SENSITIVE_BODY_FIELD.sub(_mask_body_match, raw.decode("utf-8", errors="replace")) + return ( + text + if len(text) <= _BODY_PREVIEW_CHARS + else f"{text[:_BODY_PREVIEW_CHARS]}...(+{len(text) - _BODY_PREVIEW_CHARS} chars)" + ) + + +def _masked_headers(headers: httpx.Headers) -> str: + return ", ".join( + f"{name}={MCPDebug.mask_secret(value) if name.lower() in _SENSITIVE_HEADER_NAMES else value}" + for name, value in headers.items() + ) + + +def _request_body_preview(request: httpx.Request) -> str: + try: + return _preview(request.content) or "(empty)" + except httpx.RequestNotRead: + return "(streamed, not captured)" + + +def _response_body_preview(response: httpx.Response) -> str: + try: + return _preview(response.content) or "(empty)" + except httpx.ResponseNotRead: + return "(not read)" + + +def describe_upstream_http_failure(exc: BaseException) -> str | None: + """One line per upstream ``httpx.Response`` in the exception tree: the request method, URL, + masked request headers and JSON-RPC body that were sent, plus the status and body that came back. + ``None`` when the failure never reached an HTTP response (DNS, refused connection, timeout).""" + lines: Final = tuple( + f"{response.request.method} {response.request.url} -> HTTP {response.status_code} {response.reason_phrase}" + f" | request headers: {_masked_headers(response.request.headers)}" + f" | request body: {_request_body_preview(response.request)}" + f" | response body: {_response_body_preview(response)}" + for current in iter_exception_tree(exc) + for response in (getattr(current, "response", None),) + if isinstance(response, httpx.Response) + ) + return "\n".join(lines) or None diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d7f238f142c..1247623f491 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -80,6 +80,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( raise_classified_list_failure, upstream_auth_challenge, ) +from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPPerUserTokenCache, mcp_per_user_token_cache, @@ -1353,6 +1354,11 @@ def _extract_upstream_auth_failure( return upstream_auth_challenge(exc) +def _upstream_failure_suffix(exc: BaseException) -> str: + detail: Final = describe_upstream_http_failure(exc) + return f"\n upstream exchange: {detail}" if detail else "" + + def _obo_retry_applies(server: MCPServer, subject_token: str | None) -> bool: """Whether an upstream 401/403 should invalidate the minted credential and retry once. @@ -4259,7 +4265,9 @@ class MCPServerManager: except MCPServerListError: raise except Exception as e: - verbose_logger.warning("Failed to get tools from server %s: %s", server.name, e) + verbose_logger.warning( + "Failed to get tools from server %s: %s%s", server.name, e, _upstream_failure_suffix(e) + ) raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge) async def get_prompts_from_server( @@ -5004,7 +5012,7 @@ class MCPServerManager: verbose_logger.warning("Connection error while listing tools from %s: %s", server_name, e) raise MCPServerListError(ServerListFault(tag="unreachable"), server_name) from e except Exception as e: - verbose_logger.warning("Error listing tools from %s: %s", server_name, e) + verbose_logger.warning("Error listing tools from %s: %s%s", server_name, e, _upstream_failure_suffix(e)) raise_classified_list_failure(e, server_name) _SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024 diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index ad18d1bb10f..f962148c7ff 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -37,6 +37,8 @@ import httpx from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError from typing_extensions import assert_never +from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( InMemoryTokenCacheBackend, OAuthToken, @@ -110,6 +112,10 @@ async def post_client_credentials_grant( ) except httpx.HTTPStatusError as status_err: status_code: Final = status_err.response.status_code + verbose_logger.warning( + "OAuth2 client_credentials token request denied:\n upstream exchange: %s", + describe_upstream_http_failure(status_err), + ) return TokenEndpointDenied(status_code=status_code, detail=f"token endpoint returned HTTP {status_code}") except Exception as exc: # noqa: BLE001 # any transport failure is the same outcome: unreachable return TokenEndpointUnreachable(detail=str(exc)) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index d299239f68e..b37e738579c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -5,9 +5,12 @@ Tests for MCPDebug — MCP OAuth2 debug response headers. import asyncio from unittest.mock import MagicMock +import httpx + from litellm.proxy._experimental.mcp_server.mcp_debug import ( MCP_DEBUG_REQUEST_HEADER, MCPDebug, + describe_upstream_http_failure, ) @@ -269,3 +272,58 @@ class TestWrapSendWithDebugHeaders: asyncio.run(wrapped(body_msg)) assert captured[0] == body_msg + + +class TestDescribeUpstreamHttpFailure: + @staticmethod + def _status_error(*, body: bytes, response_body: bytes | None = None) -> httpx.HTTPStatusError: + request = httpx.Request( + "POST", + "https://upstream.example/apis/mcp", + headers={"Authorization": "Bearer secret-token-abcdef0123456789", "Content-Type": "application/json"}, + content=body, + ) + response = ( + httpx.Response(500, request=request, content=response_body) + if response_body is not None + else httpx.Response(500, request=request, stream=httpx.ByteStream(b'{"error":"boom"}')) + ) + return httpx.HTTPStatusError("500", request=request, response=response) + + def test_includes_method_url_status_and_request_body(self): + exc = self._status_error( + body=b'{"method":"initialize","jsonrpc":"2.0","id":0}', + response_body=b'{"error":"boom"}', + ) + described = describe_upstream_http_failure(exc) + assert described is not None + assert "POST https://upstream.example/apis/mcp -> HTTP 500" in described + assert '{"method":"initialize"' in described + assert 'response body: {"error":"boom"}' in described + + def test_masks_authorization_header_and_secret_body_fields(self): + exc = self._status_error( + body=b"grant_type=client_credentials&client_id=abc&client_secret=super-secret-value-1234", + response_body=b"{}", + ) + described = describe_upstream_http_failure(exc) + assert described is not None + assert "secret-token-abcdef0123456789" not in described + assert "super-secret-value-1234" not in described + assert "client_id=abc" in described + assert "client_secret=" in described + + def test_reports_unread_streamed_response_body(self): + described = describe_upstream_http_failure(self._status_error(body=b"{}")) + assert described is not None + assert "response body: (not read)" in described + + def test_finds_response_behind_cause_chain(self): + wrapper = RuntimeError("token minting failed") + wrapper.__cause__ = self._status_error(body=b"{}", response_body=b'{"error":"invalid_client"}') + described = describe_upstream_http_failure(wrapper) + assert described is not None + assert "invalid_client" in described + + def test_returns_none_without_http_response(self): + assert describe_upstream_http_failure(ConnectionError("refused")) is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index 9667224de98..c7f3ea47fb3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -1,5 +1,6 @@ """Unit tests for MCP OAuth passthrough tool-fetch behavior.""" +import logging import sys from unittest.mock import AsyncMock, MagicMock @@ -11,7 +12,7 @@ if sys.version_info < (3, 11): from exceptiongroup import ExceptionGroup -from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError +from litellm.proxy._experimental.mcp_server.exceptions import MCPServerListError, MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, _extract_upstream_auth_failure, @@ -434,3 +435,27 @@ async def test_aggregate_with_single_accessible_server_still_absorbs(): assert listing.tools == [] assert listing.outcomes["delegate_docs"].tag == "auth_required" + + +@pytest.mark.asyncio +async def test_fetch_tools_logs_upstream_request_details_on_500(caplog): + manager = MCPServerManager() + request = httpx.Request( + "POST", + "https://upstream/apis/mcp", + headers={"Authorization": "Bearer upstream-token-0123456789"}, + content=b'{"method":"initialize","jsonrpc":"2.0","id":0}', + ) + response = httpx.Response(500, request=request) + mock_client = MagicMock() + mock_client.list_tools = AsyncMock( + side_effect=httpx.HTTPStatusError("500", request=request, response=response) + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + with pytest.raises(MCPServerListError): + await manager._fetch_tools_with_timeout(mock_client, "sample_docs") + + assert "POST https://upstream/apis/mcp -> HTTP 500" in caplog.text + assert '"method":"initialize"' in caplog.text + assert "upstream-token-0123456789" not in caplog.text From 058ff8c63c222d52392814b0fda84d5b030a84ec Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 15:33:44 -0700 Subject: [PATCH 018/155] test(e2e): keep answering while every Redis command times out Add tests/e2e/router/test_redis_timeout_e2e.py against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config.yml: a real Redis with socket_timeout 0.001, so every command times out and the circuit breaker opens, plus a primary deployment that always fails and falls back to a healthy one, so every request carries retry breadcrumbs into cost tracking. The test drives twenty chat requests through the proxy and asserts each answers within ten seconds, the last third is no slower than the first, /health/liveliness stays fast, and every request still reaches the spend log. Gate it behind the redis_timeout marker and E2E_REDIS_TIMEOUT, exclude it from the per-PR e2e-changed selector, register the reliability.circuit_breaker.redis_timeout.stays_responsive cell, and run it as its own job in the weekly load anomaly workflow with a Postgres and Valkey service. Against a v1.100.0 proxy the run wedges the worker: requests time out and liveliness stops answering (LIT-6780). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- .github/e2e-stack/select_tests.py | 1 + .github/workflows/weekly_load_anomaly.yml | 84 +++++++++++++++++ tests/e2e/CLAUDE.md | 2 +- tests/e2e/CONTRIBUTING.md | 2 +- tests/e2e/conftest.py | 9 +- tests/e2e/coverage_registry/reliability.yaml | 1 + tests/e2e/e2e_config.py | 5 +- tests/e2e/gateway/redis_timeout_ci_config.yml | 29 ++++++ tests/e2e/pytest.ini | 1 + tests/e2e/router/conftest.py | 21 +++-- tests/e2e/router/test_redis_timeout_e2e.py | 93 +++++++++++++++++++ 11 files changed, 234 insertions(+), 14 deletions(-) create mode 100644 tests/e2e/gateway/redis_timeout_ci_config.yml create mode 100644 tests/e2e/router/test_redis_timeout_e2e.py diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index a62358f81ff..10f26bd3b8a 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -8,6 +8,7 @@ UNSUPPORTED: Final = re.compile( r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$" r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$" r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$" + r"|^tests/e2e/router/test_redis_timeout_e2e\.py$" ) HARNESS: Final = re.compile( r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$" diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml index 3e1fca89645..6e564cf7e78 100644 --- a/.github/workflows/weekly_load_anomaly.yml +++ b/.github/workflows/weekly_load_anomaly.yml @@ -83,3 +83,87 @@ jobs: - name: Show proxy log on failure if: failure() run: tail -n 300 proxy.log + + redis-timeout-e2e: + if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 30 + services: + postgres: + image: postgres:16.6 + env: + POSTGRES_USER: llmproxy + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U llmproxy" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + valkey: + image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa + ports: + - 6379:6379 + options: >- + --health-cmd "valkey-cli ping" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + LITELLM_MASTER_KEY: sk-redis-timeout-e2e + LITELLM_LOG: WARNING + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + + - name: Install dependencies + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy + + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + + - name: Generate Prisma client + run: | + uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Start the proxy with a Redis that times out on every command + run: | + nohup uv run --no-sync litellm --config tests/e2e/gateway/redis_timeout_ci_config.yml --port 4000 > proxy.log 2>&1 & + for _ in $(seq 1 90); do + if curl -fs http://localhost:4000/health/liveliness > /dev/null; then + exit 0 + fi + sleep 2 + done + echo "proxy never became live" + tail -n 100 proxy.log + exit 1 + + - name: Run the Redis timeout e2e test + env: + E2E_REDIS_TIMEOUT: "1" + LITELLM_PROXY_URL: http://localhost:4000 + run: | + uv run --no-sync pytest tests/e2e/router/test_redis_timeout_e2e.py -v --tb=short -rA + + - name: Show proxy log on failure + if: failure() + run: tail -n 300 proxy.log diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 89c04208d65..53342f66b67 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -17,7 +17,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion lists and executes the server's tools with the stored per-user token - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns) +- `router/` - routing and reliability behavior (fallbacks, cooldowns). Also holds `test_redis_timeout_e2e.py`, which needs a proxy booted from `gateway/redis_timeout_ci_config.yml` (real Redis, `socket_timeout: 0.001`, so every command times out); it is marked `redis_timeout`, deselected unless `E2E_REDIS_TIMEOUT` is set, excluded from the per-PR check, and driven by `.github/workflows/weekly_load_anomaly.yml` - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What remains here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`) and markerless harness unit tests for the Locust/session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 1183096b81e..ccd27df6f46 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -65,7 +65,7 @@ A couple of logging destinations are configured on the proxy rather than by the ### The pull request check -Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, and `guardrails/test_presidio_masking_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, `guardrails/test_presidio_masking_e2e.py`, and `router/test_redis_timeout_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start, and the Redis timeout test needs a proxy whose Redis times out on every command (`gateway/redis_timeout_ci_config.yml`), which the weekly load workflow boots Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A file whose tests are all marked skip therefore cannot pass this check, so unskip at least one of them, or add the file to `UNSUPPORTED` in `select_tests.py` with the reason, before changing one. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index e1b987cbfd9..ff3d2dcc6bd 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -20,16 +20,14 @@ from datetime import datetime, timezone import pytest import requests - from e2e_config import CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, PROXY_BASE_URL from e2e_db import RESET_OPT_IN_ENV, reset_spend_logs, run_spend_log_cleanup from fixture_mode import fixture_mode_collection_error, fixture_report_lines -from provider_edge import replay_leftover_error from junit_properties import attach_result_properties from lifecycle import ProxyClientProvider, ResourceManager +from provider_edge import replay_leftover_error from proxy_client import ProxyClient, build_proxy_client - _E2E_TEST_RAN = pytest.StashKey[bool]() _CALL_PASSED = pytest.StashKey[bool]() @@ -60,6 +58,11 @@ def pytest_configure(config: pytest.Config) -> None: "markers", "managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set", ) + config.addinivalue_line( + "markers", + "redis_timeout: needs a proxy booted from gateway/redis_timeout_ci_config.yml whose Redis times out on every " + "command; deselected unless E2E_REDIS_TIMEOUT is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 6b69677d490..d62f0105931 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -31,6 +31,7 @@ - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} +- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 691335ffdd5..6cb59ea4e90 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -13,7 +13,6 @@ from pathlib import Path from typing import Final from dotenv import load_dotenv - from fixture_mode import deterministic_marker, parse_fixture_mode from provider_edge import provider_edge_api_base @@ -144,6 +143,7 @@ LOAD_MIN_CONCURRENCY_EFFICIENCY = float(os.environ.get("E2E_LOAD_MIN_CONCURRENCY WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" +REDIS_TIMEOUT_OPT_IN_ENV = "E2E_REDIS_TIMEOUT" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) @@ -181,8 +181,7 @@ def datadog_mcp_url(*, toolsets: str = "core") -> str: site = ( os.environ.get("DD_SITE", DD_SITE) or "datadoghq.com" ).strip().removeprefix("https://").removeprefix("http://").rstrip("/") - if site.startswith("app."): - site = site[len("app.") :] + site = site.removeprefix("app.") host = "mcp.datadoghq.com" if site in ("", "datadoghq.com") else f"mcp.{site}" base = f"https://{host}/v1/mcp" return f"{base}?toolsets={toolsets}" if toolsets else base diff --git a/tests/e2e/gateway/redis_timeout_ci_config.yml b/tests/e2e/gateway/redis_timeout_ci_config.yml new file mode 100644 index 00000000000..69ab2ee14dc --- /dev/null +++ b/tests/e2e/gateway/redis_timeout_ci_config.yml @@ -0,0 +1,29 @@ +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + store_model_in_db: true + +litellm_settings: + cache: true + cache_params: + type: redis + host: 127.0.0.1 + port: 6379 + socket_timeout: 0.001 + +router_settings: + num_retries: 1 + fallbacks: + - redis-timeout-primary: + - redis-timeout-backup + +model_list: + - model_name: redis-timeout-primary + litellm_params: + model: openai/gpt-5-mini + api_key: sk-redis-timeout-primary-not-used + mock_response: "litellm.InternalServerError" + - model_name: redis-timeout-backup + litellm_params: + model: openai/gpt-5-mini + api_key: sk-redis-timeout-backup-not-used + mock_response: "ok" diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index c3f8865f218..f69d5d058f3 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -9,3 +9,4 @@ markers = load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set + redis_timeout: needs a proxy booted from gateway/redis_timeout_ci_config.yml whose Redis times out on every command; deselected unless E2E_REDIS_TIMEOUT is set diff --git a/tests/e2e/router/conftest.py b/tests/e2e/router/conftest.py index 98501f9bd7c..57fde729f41 100644 --- a/tests/e2e/router/conftest.py +++ b/tests/e2e/router/conftest.py @@ -10,13 +10,12 @@ proxy does not already list it (compose has it in static config; stage does not) from __future__ import annotations +import os from collections.abc import Iterator import pytest -from requests import RequestException - from complexity_router_client import ComplexityRouterClient, build_client -from proxy_client import ProxyClient +from e2e_config import REDIS_TIMEOUT_OPT_IN_ENV from e2e_http import NoBody, Success from lifecycle import ResourceManager from models import ( @@ -26,6 +25,8 @@ from models import ( LiteLLMParamsBody, ModelsListResponse, ) +from proxy_client import ProxyClient +from requests import RequestException ROUTER_MODEL = "complexity-smart-router" ROUTER_PARAMS = LiteLLMParamsBody( @@ -45,6 +46,16 @@ ROUTER_PARAMS = LiteLLMParamsBody( ROUTER_KEY_MODELS = [ROUTER_MODEL, "gpt-5.5", "claude-haiku-4-5"] +def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: + if os.environ.get(REDIS_TIMEOUT_OPT_IN_ENV): + return + deselected = [item for item in items if item.get_closest_marker("redis_timeout") is not None] + if not deselected: + return + config.hook.pytest_deselected(items=deselected) + items[:] = [item for item in items if item.get_closest_marker("redis_timeout") is None] + + @pytest.fixture(scope="session") def client(proxy: ProxyClient) -> ComplexityRouterClient: return build_client(proxy) @@ -120,8 +131,6 @@ def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] # @pytest.fixture def complexity_key(resources: ResourceManager, client: ComplexityRouterClient) -> str: """Per-test key allowed to call the complexity router and its tier backends.""" - key = client.proxy.generate_key( - KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-router") - ) + key = client.proxy.generate_key(KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-router")) resources.defer(lambda: client.proxy.delete_key(key)) return key diff --git a/tests/e2e/router/test_redis_timeout_e2e.py b/tests/e2e/router/test_redis_timeout_e2e.py new file mode 100644 index 00000000000..eeb5224e051 --- /dev/null +++ b/tests/e2e/router/test_redis_timeout_e2e.py @@ -0,0 +1,93 @@ +"""Live e2e: the proxy keeps answering while every Redis command times out. + +Runs only against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config.yml, which +points cache_params at a real Redis with socket_timeout 0.001 so every command times out and +the circuit breaker opens. Each request fails its primary deployment, retries, falls back to the +backup and succeeds, so it carries retry breadcrumbs; its cost tracking then fails on the spend +counter increment and stringifies the request metadata into a failed-tracking alert. On v1.100.0 +that string doubled per request until the worker hung (LIT-6780). Deselected unless +E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. +""" + +from __future__ import annotations + +import time +from typing import Final + +import pytest +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from e2e_http import NoBody, Success +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody + +pytestmark = [pytest.mark.e2e, pytest.mark.redis_timeout] + +PRIMARY_MODEL: Final = "redis-timeout-primary" +BACKUP_MODEL: Final = "redis-timeout-backup" +REQUESTS: Final = 20 +MAX_SECONDS_PER_REQUEST: Final = 10.0 +MAX_LATENCY_GROWTH_RATIO: Final = 3.0 +MAX_LIVELINESS_SECONDS: Final = 2.0 + + +class TestRedisTimeout: + @pytest.mark.covers( + "reliability.circuit_breaker.redis_timeout.stays_responsive", + exercised_on=["chat_completions"], + ) + def test_retries_under_redis_timeouts_keep_answering( + self, client: ComplexityRouterClient, resources: ResourceManager + ) -> None: + proxy = client.proxy + key = proxy.generate_key( + KeyGenerateBody(models=[PRIMARY_MODEL, BACKUP_MODEL], key_alias=f"e2e-redis-timeout-{unique_marker()}") + ) + resources.defer(lambda: proxy.delete_key(key)) + + latencies: list[float] = [] + for request_number in range(1, REQUESTS + 1): + started = time.monotonic() + result = proxy.transport.post( + "/chat/completions", + headers=proxy.transport.bearer(key), + json=ChatBody( + model=PRIMARY_MODEL, + messages=[ChatMessage(role="user", content=f"redis timeout {unique_marker()} {request_number}")], + max_tokens=5, + ), + response_type=ChatResponse, + timeout=MAX_SECONDS_PER_REQUEST, + ) + elapsed = time.monotonic() - started + assert isinstance(result, Success), ( + f"request {request_number} failed after {elapsed:.1f}s with Redis timing out: {result}; " + f"earlier requests took {[round(seconds, 2) for seconds in latencies]}" + ) + assert result.data.choices, f"request {request_number}: fallback to {BACKUP_MODEL} returned no choices" + assert elapsed < MAX_SECONDS_PER_REQUEST, ( + f"request {request_number} took {elapsed:.1f}s with Redis timing out; " + f"earlier requests took {[round(seconds, 2) for seconds in latencies]}" + ) + latencies.append(elapsed) + + third = REQUESTS // 3 + early = sum(latencies[:third]) / third + late = sum(latencies[-third:]) / third + assert late <= max(early * MAX_LATENCY_GROWTH_RATIO, 0.5), ( + f"per-request latency grew from {early:.2f}s to {late:.2f}s across {REQUESTS} requests " + "while Redis timed out; the proxy is paying more for each failed request" + ) + + started = time.monotonic() + probe = proxy.transport.probe("/health/liveliness", params=NoBody()) + liveliness_seconds = time.monotonic() - started + assert probe.healthy, f"/health/liveliness returned {probe.status_code} after the Redis timeout loop" + assert liveliness_seconds < MAX_LIVELINESS_SECONDS, ( + f"/health/liveliness took {liveliness_seconds:.1f}s after the loop; the worker is stalled" + ) + + rows = proxy.poll_logs_for_key(key, min_rows=REQUESTS) + assert len(rows) >= REQUESTS, ( + f"only {len(rows)} of {REQUESTS} requests reached the spend log; a Redis outage must not lose spend rows" + ) From d9e822d2562ddd716d8f6a5e86d3d8fedcb30df1 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 15:51:11 -0700 Subject: [PATCH 019/155] ci: run the Redis timeout e2e test from its own weekly workflow It is a functional e2e test, not a load test, so give it its own workflow instead of a job inside the load anomaly run. It keeps the Saturday 12:00 UTC cadence and manual dispatch, and boots the timeout-config proxy with Postgres and Valkey services exactly as before. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- .github/workflows/test-e2e-redis-timeout.yml | 94 ++++++++++++++++++++ .github/workflows/weekly_load_anomaly.yml | 84 ----------------- tests/e2e/CLAUDE.md | 2 +- tests/e2e/CONTRIBUTING.md | 2 +- 4 files changed, 96 insertions(+), 86 deletions(-) create mode 100644 .github/workflows/test-e2e-redis-timeout.yml diff --git a/.github/workflows/test-e2e-redis-timeout.yml b/.github/workflows/test-e2e-redis-timeout.yml new file mode 100644 index 00000000000..a956ed79590 --- /dev/null +++ b/.github/workflows/test-e2e-redis-timeout.yml @@ -0,0 +1,94 @@ +name: "Weekly Redis Timeout E2E" + +on: + schedule: + - cron: "0 12 * * 6" + workflow_dispatch: + +permissions: + contents: read + +jobs: + redis-timeout-e2e: + if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 30 + services: + postgres: + image: postgres:16.6 + env: + POSTGRES_USER: llmproxy + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U llmproxy" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + valkey: + image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa + ports: + - 6379:6379 + options: >- + --health-cmd "valkey-cli ping" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + LITELLM_MASTER_KEY: sk-redis-timeout-e2e + LITELLM_LOG: WARNING + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + + - name: Install dependencies + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy + + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + + - name: Generate Prisma client + run: | + uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Start the proxy with a Redis that times out on every command + run: | + nohup uv run --no-sync litellm --config tests/e2e/gateway/redis_timeout_ci_config.yml --port 4000 > proxy.log 2>&1 & + for _ in $(seq 1 90); do + if curl -fs http://localhost:4000/health/liveliness > /dev/null; then + exit 0 + fi + sleep 2 + done + echo "proxy never became live" + tail -n 100 proxy.log + exit 1 + + - name: Run the Redis timeout e2e test + env: + E2E_REDIS_TIMEOUT: "1" + LITELLM_PROXY_URL: http://localhost:4000 + run: | + uv run --no-sync pytest tests/e2e/router/test_redis_timeout_e2e.py -v --tb=short -rA + + - name: Show proxy log on failure + if: failure() + run: tail -n 300 proxy.log diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml index 6e564cf7e78..3e1fca89645 100644 --- a/.github/workflows/weekly_load_anomaly.yml +++ b/.github/workflows/weekly_load_anomaly.yml @@ -83,87 +83,3 @@ jobs: - name: Show proxy log on failure if: failure() run: tail -n 300 proxy.log - - redis-timeout-e2e: - if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 30 - services: - postgres: - image: postgres:16.6 - env: - POSTGRES_USER: llmproxy - POSTGRES_PASSWORD: dbpassword9090 - POSTGRES_DB: litellm - ports: - - 5432:5432 - options: >- - --health-cmd "pg_isready -U llmproxy" - --health-interval 5s - --health-timeout 5s - --health-retries 10 - valkey: - image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa - ports: - - 6379:6379 - options: >- - --health-cmd "valkey-cli ping" - --health-interval 5s - --health-timeout 5s - --health-retries 10 - env: - DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm - LITELLM_MASTER_KEY: sk-redis-timeout-e2e - LITELLM_LOG: WARNING - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Cache the Rust build - uses: ./.github/actions/cache-cargo-build - - - name: Install dependencies - run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy - - - name: Cache Prisma binaries - uses: ./.github/actions/cache-prisma-binaries - - - name: Generate Prisma client - run: | - uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - - name: Start the proxy with a Redis that times out on every command - run: | - nohup uv run --no-sync litellm --config tests/e2e/gateway/redis_timeout_ci_config.yml --port 4000 > proxy.log 2>&1 & - for _ in $(seq 1 90); do - if curl -fs http://localhost:4000/health/liveliness > /dev/null; then - exit 0 - fi - sleep 2 - done - echo "proxy never became live" - tail -n 100 proxy.log - exit 1 - - - name: Run the Redis timeout e2e test - env: - E2E_REDIS_TIMEOUT: "1" - LITELLM_PROXY_URL: http://localhost:4000 - run: | - uv run --no-sync pytest tests/e2e/router/test_redis_timeout_e2e.py -v --tb=short -rA - - - name: Show proxy log on failure - if: failure() - run: tail -n 300 proxy.log diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 53342f66b67..76e8550cb9a 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -17,7 +17,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion lists and executes the server's tools with the stored per-user token - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns). Also holds `test_redis_timeout_e2e.py`, which needs a proxy booted from `gateway/redis_timeout_ci_config.yml` (real Redis, `socket_timeout: 0.001`, so every command times out); it is marked `redis_timeout`, deselected unless `E2E_REDIS_TIMEOUT` is set, excluded from the per-PR check, and driven by `.github/workflows/weekly_load_anomaly.yml` +- `router/` - routing and reliability behavior (fallbacks, cooldowns). Also holds `test_redis_timeout_e2e.py`, which needs a proxy booted from `gateway/redis_timeout_ci_config.yml` (real Redis, `socket_timeout: 0.001`, so every command times out); it is marked `redis_timeout`, deselected unless `E2E_REDIS_TIMEOUT` is set, excluded from the per-PR check, and driven weekly by `.github/workflows/test-e2e-redis-timeout.yml` - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What remains here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`) and markerless harness unit tests for the Locust/session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index ccd27df6f46..7a859965c80 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -65,7 +65,7 @@ A couple of logging destinations are configured on the proxy rather than by the ### The pull request check -Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, `guardrails/test_presidio_masking_e2e.py`, and `router/test_redis_timeout_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start, and the Redis timeout test needs a proxy whose Redis times out on every command (`gateway/redis_timeout_ci_config.yml`), which the weekly load workflow boots +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, `guardrails/test_presidio_masking_e2e.py`, and `router/test_redis_timeout_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start, and the Redis timeout test needs a proxy whose Redis times out on every command (`gateway/redis_timeout_ci_config.yml`), which `.github/workflows/test-e2e-redis-timeout.yml` boots on a weekly schedule Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A file whose tests are all marked skip therefore cannot pass this check, so unskip at least one of them, or add the file to `UNSUPPORTED` in `select_tests.py` with the reason, before changing one. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch From 02351da51f14d644660c4a5583968389752af46b Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 15:51:35 -0700 Subject: [PATCH 020/155] test(e2e): register the Redis timeout cell at tier P1 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- tests/e2e/coverage_registry/reliability.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index d62f0105931..38427b647b9 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -31,7 +31,7 @@ - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} -- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} +- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} From 939ac04a932baae63f2f8b9440ee43ed70570763 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 16:03:36 -0700 Subject: [PATCH 021/155] test(e2e): cover /v1/responses in the Redis timeout test and fail the primary for real The Responses path returns a mock for any mock_response string, so the InternalServerError sentinel never failed there. Point the primary deployment's api_base at a closed port instead, which fails every endpoint the same way, then parametrize the test over /chat/completions and /v1/responses and register both on the coverage cell. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- tests/e2e/coverage_registry/reliability.yaml | 2 +- tests/e2e/gateway/redis_timeout_ci_config.yml | 2 +- tests/e2e/router/test_redis_timeout_e2e.py | 105 +++++++++++++----- 3 files changed, 81 insertions(+), 28 deletions(-) diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 38427b647b9..3b378b56870 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -31,7 +31,7 @@ - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} -- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} +- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions, responses], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} diff --git a/tests/e2e/gateway/redis_timeout_ci_config.yml b/tests/e2e/gateway/redis_timeout_ci_config.yml index 69ab2ee14dc..735285adb36 100644 --- a/tests/e2e/gateway/redis_timeout_ci_config.yml +++ b/tests/e2e/gateway/redis_timeout_ci_config.yml @@ -21,7 +21,7 @@ model_list: litellm_params: model: openai/gpt-5-mini api_key: sk-redis-timeout-primary-not-used - mock_response: "litellm.InternalServerError" + api_base: http://127.0.0.1:1 - model_name: redis-timeout-backup litellm_params: model: openai/gpt-5-mini diff --git a/tests/e2e/router/test_redis_timeout_e2e.py b/tests/e2e/router/test_redis_timeout_e2e.py index eeb5224e051..d10f32b6482 100644 --- a/tests/e2e/router/test_redis_timeout_e2e.py +++ b/tests/e2e/router/test_redis_timeout_e2e.py @@ -1,25 +1,29 @@ """Live e2e: the proxy keeps answering while every Redis command times out. Runs only against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config.yml, which -points cache_params at a real Redis with socket_timeout 0.001 so every command times out and -the circuit breaker opens. Each request fails its primary deployment, retries, falls back to the -backup and succeeds, so it carries retry breadcrumbs; its cost tracking then fails on the spend -counter increment and stringifies the request metadata into a failed-tracking alert. On v1.100.0 -that string doubled per request until the worker hung (LIT-6780). Deselected unless -E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. +points cache_params at a real Redis with socket_timeout 0.001 so commands time out and the +circuit breaker opens. Each request fails its primary deployment, whose api_base is a closed +port, retries, falls back to the backup and succeeds, so it carries retry breadcrumbs; its cost +tracking then fails on the spend counter increment and stringifies the request metadata into a +failed-tracking alert. On v1.100.0 that string doubled per request until the worker hung +(LIT-6780). Deselected unless E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. """ from __future__ import annotations import time +from collections.abc import Callable +from dataclasses import dataclass from typing import Final import pytest from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker -from e2e_http import NoBody, Success +from e2e_http import NoBody, Result, Success from lifecycle import ResourceManager from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody +from proxy_client import ProxyClient +from pydantic import BaseModel pytestmark = [pytest.mark.e2e, pytest.mark.redis_timeout] @@ -31,42 +35,90 @@ MAX_LATENCY_GROWTH_RATIO: Final = 3.0 MAX_LIVELINESS_SECONDS: Final = 2.0 +class ResponsesBody(BaseModel): + model: str + input: str + max_output_tokens: int = 5 + + +class ResponsesObject(BaseModel): + id: str | None = None + status: str | None = None + output: list[object] = [] + + +@dataclass(frozen=True, slots=True) +class Endpoint: + name: str + send: Callable[[ProxyClient, str, str], Result[BaseModel]] + served: Callable[[BaseModel], bool] + + +def _send_chat(proxy: ProxyClient, key: str, marker: str) -> Result[BaseModel]: + return proxy.transport.post( + "/chat/completions", + headers=proxy.transport.bearer(key), + json=ChatBody(model=PRIMARY_MODEL, messages=[ChatMessage(role="user", content=marker)], max_tokens=5), + response_type=ChatResponse, + timeout=MAX_SECONDS_PER_REQUEST, + ) + + +def _send_responses(proxy: ProxyClient, key: str, marker: str) -> Result[BaseModel]: + return proxy.transport.post( + "/v1/responses", + headers=proxy.transport.bearer(key), + json=ResponsesBody(model=PRIMARY_MODEL, input=marker), + response_type=ResponsesObject, + timeout=MAX_SECONDS_PER_REQUEST, + ) + + +ENDPOINTS: Final = ( + Endpoint( + name="chat_completions", + send=_send_chat, + served=lambda data: isinstance(data, ChatResponse) and bool(data.choices), + ), + Endpoint( + name="responses", + send=_send_responses, + served=lambda data: isinstance(data, ResponsesObject) and bool(data.output), + ), +) + + class TestRedisTimeout: + @pytest.mark.parametrize("endpoint", ENDPOINTS, ids=[endpoint.name for endpoint in ENDPOINTS]) @pytest.mark.covers( "reliability.circuit_breaker.redis_timeout.stays_responsive", - exercised_on=["chat_completions"], + exercised_on=["chat_completions", "responses"], ) def test_retries_under_redis_timeouts_keep_answering( - self, client: ComplexityRouterClient, resources: ResourceManager + self, client: ComplexityRouterClient, resources: ResourceManager, endpoint: Endpoint ) -> None: proxy = client.proxy key = proxy.generate_key( - KeyGenerateBody(models=[PRIMARY_MODEL, BACKUP_MODEL], key_alias=f"e2e-redis-timeout-{unique_marker()}") + KeyGenerateBody( + models=[PRIMARY_MODEL, BACKUP_MODEL], key_alias=f"e2e-redis-timeout-{endpoint.name}-{unique_marker()}" + ) ) resources.defer(lambda: proxy.delete_key(key)) latencies: list[float] = [] for request_number in range(1, REQUESTS + 1): started = time.monotonic() - result = proxy.transport.post( - "/chat/completions", - headers=proxy.transport.bearer(key), - json=ChatBody( - model=PRIMARY_MODEL, - messages=[ChatMessage(role="user", content=f"redis timeout {unique_marker()} {request_number}")], - max_tokens=5, - ), - response_type=ChatResponse, - timeout=MAX_SECONDS_PER_REQUEST, - ) + result = endpoint.send(proxy, key, f"redis timeout {unique_marker()} {request_number}") elapsed = time.monotonic() - started assert isinstance(result, Success), ( - f"request {request_number} failed after {elapsed:.1f}s with Redis timing out: {result}; " + f"{endpoint.name} request {request_number} failed after {elapsed:.1f}s with Redis timing out: {result}; " f"earlier requests took {[round(seconds, 2) for seconds in latencies]}" ) - assert result.data.choices, f"request {request_number}: fallback to {BACKUP_MODEL} returned no choices" + assert endpoint.served(result.data), ( + f"{endpoint.name} request {request_number}: fallback to {BACKUP_MODEL} returned no output" + ) assert elapsed < MAX_SECONDS_PER_REQUEST, ( - f"request {request_number} took {elapsed:.1f}s with Redis timing out; " + f"{endpoint.name} request {request_number} took {elapsed:.1f}s with Redis timing out; " f"earlier requests took {[round(seconds, 2) for seconds in latencies]}" ) latencies.append(elapsed) @@ -75,7 +127,7 @@ class TestRedisTimeout: early = sum(latencies[:third]) / third late = sum(latencies[-third:]) / third assert late <= max(early * MAX_LATENCY_GROWTH_RATIO, 0.5), ( - f"per-request latency grew from {early:.2f}s to {late:.2f}s across {REQUESTS} requests " + f"{endpoint.name} per-request latency grew from {early:.2f}s to {late:.2f}s across {REQUESTS} requests " "while Redis timed out; the proxy is paying more for each failed request" ) @@ -89,5 +141,6 @@ class TestRedisTimeout: rows = proxy.poll_logs_for_key(key, min_rows=REQUESTS) assert len(rows) >= REQUESTS, ( - f"only {len(rows)} of {REQUESTS} requests reached the spend log; a Redis outage must not lose spend rows" + f"only {len(rows)} of {REQUESTS} {endpoint.name} requests reached the spend log; " + "a Redis outage must not lose spend rows" ) From 2d7e5999b78f726c97a7142231de566d788d4fc0 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 16:26:03 -0700 Subject: [PATCH 022/155] test(e2e): pause Redis writes for the run and prove the timeouts from /metrics A loopback Redis answers many commands inside the 1 ms socket timeout, so nothing guaranteed the failure path ran. The test now holds the proxy's Redis in CLIENT PAUSE WRITE for its duration, so every write the proxy sends, the spend counter increment included, outlives the timeout, and lifts the pause in teardown. Reads stay live so the control connection can do that. Enable the prometheus callback in the gateway config and assert from /metrics that the proxy counted at least the breaker's five timeouts and that, during each case, it saw fresh timeouts, a breaker transition, or an open breaker rejecting every call. The open breaker is the state a customer's worker sits in, and cost tracking fails on every request either way. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- .github/workflows/test-e2e-redis-timeout.yml | 2 + tests/e2e/gateway/redis_timeout_ci_config.yml | 2 + tests/e2e/router/test_redis_timeout_e2e.py | 60 +++++++++++++++++-- 3 files changed, 60 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test-e2e-redis-timeout.yml b/.github/workflows/test-e2e-redis-timeout.yml index a956ed79590..9c318a1346c 100644 --- a/.github/workflows/test-e2e-redis-timeout.yml +++ b/.github/workflows/test-e2e-redis-timeout.yml @@ -86,6 +86,8 @@ jobs: env: E2E_REDIS_TIMEOUT: "1" LITELLM_PROXY_URL: http://localhost:4000 + REDIS_HOST: 127.0.0.1 + REDIS_PORT: "6379" run: | uv run --no-sync pytest tests/e2e/router/test_redis_timeout_e2e.py -v --tb=short -rA diff --git a/tests/e2e/gateway/redis_timeout_ci_config.yml b/tests/e2e/gateway/redis_timeout_ci_config.yml index 735285adb36..375a51f6127 100644 --- a/tests/e2e/gateway/redis_timeout_ci_config.yml +++ b/tests/e2e/gateway/redis_timeout_ci_config.yml @@ -3,6 +3,8 @@ general_settings: store_model_in_db: true litellm_settings: + callbacks: ["prometheus"] + require_auth_for_metrics_endpoint: false cache: true cache_params: type: redis diff --git a/tests/e2e/router/test_redis_timeout_e2e.py b/tests/e2e/router/test_redis_timeout_e2e.py index d10f32b6482..e99bd2fbe80 100644 --- a/tests/e2e/router/test_redis_timeout_e2e.py +++ b/tests/e2e/router/test_redis_timeout_e2e.py @@ -1,8 +1,10 @@ """Live e2e: the proxy keeps answering while every Redis command times out. Runs only against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config.yml, which -points cache_params at a real Redis with socket_timeout 0.001 so commands time out and the -circuit breaker opens. Each request fails its primary deployment, whose api_base is a closed +points cache_params at a real Redis with socket_timeout 0.001. The test holds that Redis in +CLIENT PAUSE WRITE for its duration, so every write the proxy sends, the spend counter increment +included, hangs past the timeout, and it proves the degradation was real from the breaker metrics on /metrics: fresh timeouts, a breaker transition, +or an already-open breaker rejecting every call, which is the state a customer's worker sits in. Each request fails its primary deployment, whose api_base is a closed port, retries, falls back to the backup and succeeds, so it carries retry breadcrumbs; its cost tracking then fails on the spend counter increment and stringifies the request metadata into a failed-tracking alert. On v1.100.0 that string doubled per request until the worker hung @@ -11,12 +13,15 @@ failed-tracking alert. On v1.100.0 that string doubled per request until the wor from __future__ import annotations +import os +import re import time -from collections.abc import Callable +from collections.abc import Callable, Iterator from dataclasses import dataclass from typing import Final import pytest +import redis from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker from e2e_http import NoBody, Result, Success @@ -33,6 +38,11 @@ REQUESTS: Final = 20 MAX_SECONDS_PER_REQUEST: Final = 10.0 MAX_LATENCY_GROWTH_RATIO: Final = 3.0 MAX_LIVELINESS_SECONDS: Final = 2.0 +REDIS_PAUSE_MS: Final = 600_000 +BREAKER_FAILURE_THRESHOLD: Final = 5 +TIMEOUT_FAILURES_RE: Final = re.compile( + r'^litellm_redis_circuit_breaker_failures_total\{failure_class="timeout"\} ([0-9.e+]+)$', re.M +) class ResponsesBody(BaseModel): @@ -86,6 +96,32 @@ ENDPOINTS: Final = ( served=lambda data: isinstance(data, ResponsesObject) and bool(data.output), ), ) +BREAKER_OPEN_RE: Final = re.compile(r'^litellm_redis_circuit_breaker_state\{state="open"\} ([0-9.e+]+)$', re.M) +BREAKER_TRANSITIONS_RE: Final = re.compile( + r'^litellm_redis_circuit_breaker_transitions_total\{state="[a-z_]+"\} ([0-9.e+]+)$', re.M +) + + +@pytest.fixture +def paused_redis() -> Iterator[None]: + """Hold the proxy's Redis in CLIENT PAUSE WRITE so every write it sends outlives the 1 ms socket + timeout. A loopback Redis otherwise answers many commands inside that budget. Reads stay live so + this control connection can lift the pause in teardown.""" + host = os.environ.get("REDIS_HOST") + port = os.environ.get("REDIS_PORT") + assert host and port, "REDIS_HOST and REDIS_PORT must name the Redis the proxy under test uses" + control = redis.Redis(host=host, port=int(port), socket_timeout=5) + control.client_pause(REDIS_PAUSE_MS, all=False) # pyright: ignore[reportUnknownMemberType] # redis-py stubs return Any + try: + yield + finally: + control.client_unpause() # pyright: ignore[reportUnknownMemberType] # redis-py stubs return Any + control.close() + + +def _metric(proxy: ProxyClient, pattern: re.Pattern[str]) -> float: + body = proxy.probe("/metrics", params=NoBody()).body + return sum(float(match.group(1)) for match in pattern.finditer(body)) class TestRedisTimeout: @@ -95,9 +131,11 @@ class TestRedisTimeout: exercised_on=["chat_completions", "responses"], ) def test_retries_under_redis_timeouts_keep_answering( - self, client: ComplexityRouterClient, resources: ResourceManager, endpoint: Endpoint + self, client: ComplexityRouterClient, resources: ResourceManager, endpoint: Endpoint, paused_redis: None ) -> None: proxy = client.proxy + timeouts_before = _metric(proxy, TIMEOUT_FAILURES_RE) + transitions_before = _metric(proxy, BREAKER_TRANSITIONS_RE) key = proxy.generate_key( KeyGenerateBody( models=[PRIMARY_MODEL, BACKUP_MODEL], key_alias=f"e2e-redis-timeout-{endpoint.name}-{unique_marker()}" @@ -139,6 +177,20 @@ class TestRedisTimeout: f"/health/liveliness took {liveliness_seconds:.1f}s after the loop; the worker is stalled" ) + timeouts_total = _metric(proxy, TIMEOUT_FAILURES_RE) + timeouts = timeouts_total - timeouts_before + transitions = _metric(proxy, BREAKER_TRANSITIONS_RE) - transitions_before + breaker_open = _metric(proxy, BREAKER_OPEN_RE) >= 1 + assert timeouts_total >= BREAKER_FAILURE_THRESHOLD, ( + f"the proxy counted only {timeouts_total:.0f} Redis timeouts in its lifetime; the write-paused Redis " + "never made its spend counter writes time out, so this run proved nothing" + ) + assert timeouts >= REQUESTS or transitions >= 1 or breaker_open, ( + f"during {REQUESTS} {endpoint.name} requests the breaker counted {timeouts:.0f} new timeouts, " + f"{transitions:.0f} state transitions, and ended {'open' if breaker_open else 'closed'}; " + "Redis was healthy for this case, so it proved nothing" + ) + rows = proxy.poll_logs_for_key(key, min_rows=REQUESTS) assert len(rows) >= REQUESTS, ( f"only {len(rows)} of {REQUESTS} {endpoint.name} requests reached the spend log; " From 2e4461ec458756c940f3a8863ea805bc44c33907 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 16:33:18 -0700 Subject: [PATCH 023/155] test(e2e): register the Redis timeout deployments through /model/new The e2e directive has every test create its deployments through the management API and delete them on teardown. Drop the static model_list from the gateway config; the test now registers the closed-port primary and the mock backup itself, and the fallback map stays in router_settings where proxy-level config belongs. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- tests/e2e/gateway/redis_timeout_ci_config.yml | 12 ----------- tests/e2e/router/test_redis_timeout_e2e.py | 20 ++++++++++++++++--- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/tests/e2e/gateway/redis_timeout_ci_config.yml b/tests/e2e/gateway/redis_timeout_ci_config.yml index 375a51f6127..9bcf65162d6 100644 --- a/tests/e2e/gateway/redis_timeout_ci_config.yml +++ b/tests/e2e/gateway/redis_timeout_ci_config.yml @@ -17,15 +17,3 @@ router_settings: fallbacks: - redis-timeout-primary: - redis-timeout-backup - -model_list: - - model_name: redis-timeout-primary - litellm_params: - model: openai/gpt-5-mini - api_key: sk-redis-timeout-primary-not-used - api_base: http://127.0.0.1:1 - - model_name: redis-timeout-backup - litellm_params: - model: openai/gpt-5-mini - api_key: sk-redis-timeout-backup-not-used - mock_response: "ok" diff --git a/tests/e2e/router/test_redis_timeout_e2e.py b/tests/e2e/router/test_redis_timeout_e2e.py index e99bd2fbe80..58d10f1d8ea 100644 --- a/tests/e2e/router/test_redis_timeout_e2e.py +++ b/tests/e2e/router/test_redis_timeout_e2e.py @@ -4,8 +4,8 @@ Runs only against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config. points cache_params at a real Redis with socket_timeout 0.001. The test holds that Redis in CLIENT PAUSE WRITE for its duration, so every write the proxy sends, the spend counter increment included, hangs past the timeout, and it proves the degradation was real from the breaker metrics on /metrics: fresh timeouts, a breaker transition, -or an already-open breaker rejecting every call, which is the state a customer's worker sits in. Each request fails its primary deployment, whose api_base is a closed -port, retries, falls back to the backup and succeeds, so it carries retry breadcrumbs; its cost +or an already-open breaker rejecting every call, which is the state a customer's worker sits in. The test registers two deployments through /model/new: a primary whose api_base is a closed port +and a backup that answers with a mock. Each request fails the primary, retries, falls back and succeeds, so it carries retry breadcrumbs; its cost tracking then fails on the spend counter increment and stringifies the request metadata into a failed-tracking alert. On v1.100.0 that string doubled per request until the worker hung (LIT-6780). Deselected unless E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. @@ -26,7 +26,7 @@ from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker from e2e_http import NoBody, Result, Success from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody +from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody, LiteLLMParamsBody from proxy_client import ProxyClient from pydantic import BaseModel @@ -34,6 +34,8 @@ pytestmark = [pytest.mark.e2e, pytest.mark.redis_timeout] PRIMARY_MODEL: Final = "redis-timeout-primary" BACKUP_MODEL: Final = "redis-timeout-backup" +BACKING_MODEL: Final = "openai/gpt-5-mini" +CLOSED_PORT_API_BASE: Final = "http://127.0.0.1:1" REQUESTS: Final = 20 MAX_SECONDS_PER_REQUEST: Final = 10.0 MAX_LATENCY_GROWTH_RATIO: Final = 3.0 @@ -134,6 +136,18 @@ class TestRedisTimeout: self, client: ComplexityRouterClient, resources: ResourceManager, endpoint: Endpoint, paused_redis: None ) -> None: proxy = client.proxy + primary_id = proxy.create_model( + PRIMARY_MODEL, + LiteLLMParamsBody( + model=BACKING_MODEL, api_key="sk-redis-timeout-primary-not-used", api_base=CLOSED_PORT_API_BASE + ), + ) + resources.defer(lambda: proxy.delete_model(primary_id)) + backup_id = proxy.create_model( + BACKUP_MODEL, + LiteLLMParamsBody(model=BACKING_MODEL, api_key="sk-redis-timeout-backup-not-used", mock_response="ok"), + ) + resources.defer(lambda: proxy.delete_model(backup_id)) timeouts_before = _metric(proxy, TIMEOUT_FAILURES_RE) transitions_before = _metric(proxy, BREAKER_TRANSITIONS_RE) key = proxy.generate_key( From 1213d7d39929500883fee4c3823f8dfbc073fdb8 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 17:17:33 -0700 Subject: [PATCH 024/155] test(e2e): cover /embeddings and assert memory and fallbacks in the Redis timeout test Add an embeddings case with its own closed-port primary and mock backup (the fallback map in the gateway config gains the pair; LiteLLMParamsBody.mock_response accepts the list an embedding mock needs). Assert from /metrics that the proxy's resident memory grows by no more than 200 MB across each case where the process collector reports it (Linux), that the router counted a successful fallback for every request, and that every spend row is a success. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- tests/e2e/coverage_registry/reliability.yaml | 2 +- tests/e2e/gateway/redis_timeout_ci_config.yml | 2 + tests/e2e/models.py | 2 +- tests/e2e/router/test_redis_timeout_e2e.py | 142 +++++++++++++----- 4 files changed, 110 insertions(+), 38 deletions(-) diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 3b378b56870..a42e4b2779a 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -31,7 +31,7 @@ - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} -- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions, responses], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} +- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions, responses, embeddings], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} diff --git a/tests/e2e/gateway/redis_timeout_ci_config.yml b/tests/e2e/gateway/redis_timeout_ci_config.yml index 9bcf65162d6..fd283dbc7df 100644 --- a/tests/e2e/gateway/redis_timeout_ci_config.yml +++ b/tests/e2e/gateway/redis_timeout_ci_config.yml @@ -17,3 +17,5 @@ router_settings: fallbacks: - redis-timeout-primary: - redis-timeout-backup + - redis-timeout-embed-primary: + - redis-timeout-embed-backup diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 62810e6cfd9..f3c111bdf06 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -920,7 +920,7 @@ class LiteLLMParamsBody(BaseModel): auto_router_default_model: str | None = None auto_router_embedding_model: str | None = None tags: list[str] | None = None - mock_response: str | None = None + mock_response: str | list[float] | None = None timeout: float | None = None tpm: int | None = None weight: int | None = None diff --git a/tests/e2e/router/test_redis_timeout_e2e.py b/tests/e2e/router/test_redis_timeout_e2e.py index 58d10f1d8ea..3f1984c0696 100644 --- a/tests/e2e/router/test_redis_timeout_e2e.py +++ b/tests/e2e/router/test_redis_timeout_e2e.py @@ -3,12 +3,12 @@ Runs only against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config.yml, which points cache_params at a real Redis with socket_timeout 0.001. The test holds that Redis in CLIENT PAUSE WRITE for its duration, so every write the proxy sends, the spend counter increment -included, hangs past the timeout, and it proves the degradation was real from the breaker metrics on /metrics: fresh timeouts, a breaker transition, -or an already-open breaker rejecting every call, which is the state a customer's worker sits in. The test registers two deployments through /model/new: a primary whose api_base is a closed port -and a backup that answers with a mock. Each request fails the primary, retries, falls back and succeeds, so it carries retry breadcrumbs; its cost -tracking then fails on the spend counter increment and stringifies the request metadata into a -failed-tracking alert. On v1.100.0 that string doubled per request until the worker hung -(LIT-6780). Deselected unless E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. +included, hangs past the timeout. For each endpoint it registers two deployments through +/model/new: a primary whose api_base is a closed port and a backup that answers with a mock. Each request fails the primary, +retries, falls back and succeeds, so it carries retry breadcrumbs; its cost tracking then fails on +the spend counter increment and stringifies the request metadata into a failed-tracking alert. On +v1.100.0 that string doubled per request until the worker hung (LIT-6780). Deselected unless +E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. """ from __future__ import annotations @@ -26,25 +26,29 @@ from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker from e2e_http import NoBody, Result, Success from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody, LiteLLMParamsBody +from models import ChatBody, ChatMessage, ChatResponse, EmbedBody, KeyGenerateBody, LiteLLMParamsBody from proxy_client import ProxyClient from pydantic import BaseModel pytestmark = [pytest.mark.e2e, pytest.mark.redis_timeout] -PRIMARY_MODEL: Final = "redis-timeout-primary" -BACKUP_MODEL: Final = "redis-timeout-backup" -BACKING_MODEL: Final = "openai/gpt-5-mini" -CLOSED_PORT_API_BASE: Final = "http://127.0.0.1:1" REQUESTS: Final = 20 MAX_SECONDS_PER_REQUEST: Final = 10.0 MAX_LATENCY_GROWTH_RATIO: Final = 3.0 MAX_LIVELINESS_SECONDS: Final = 2.0 +MAX_RSS_GROWTH_BYTES: Final = 200 * 1024 * 1024 REDIS_PAUSE_MS: Final = 600_000 BREAKER_FAILURE_THRESHOLD: Final = 5 +CLOSED_PORT_API_BASE: Final = "http://127.0.0.1:1" TIMEOUT_FAILURES_RE: Final = re.compile( r'^litellm_redis_circuit_breaker_failures_total\{failure_class="timeout"\} ([0-9.e+]+)$', re.M ) +BREAKER_OPEN_RE: Final = re.compile(r'^litellm_redis_circuit_breaker_state\{state="open"\} ([0-9.e+]+)$', re.M) +BREAKER_TRANSITIONS_RE: Final = re.compile( + r'^litellm_redis_circuit_breaker_transitions_total\{state="[a-z_]+"\} ([0-9.e+]+)$', re.M +) +FALLBACKS_RE: Final = re.compile(r"^litellm_deployment_successful_fallbacks_total\{[^}]*\} ([0-9.e+]+)$", re.M) +RSS_RE: Final = re.compile(r"^process_resident_memory_bytes ([0-9.e+]+)$", re.M) class ResponsesBody(BaseModel): @@ -59,48 +63,95 @@ class ResponsesObject(BaseModel): output: list[object] = [] +class EmbeddingsObject(BaseModel): + model: str | None = None + data: list[object] = [] + + @dataclass(frozen=True, slots=True) class Endpoint: + """One endpoint's deployments and request shape.""" + name: str - send: Callable[[ProxyClient, str, str], Result[BaseModel]] + primary: str + backup: str + primary_params: LiteLLMParamsBody + backup_params: LiteLLMParamsBody + send: Callable[[ProxyClient, str, str, str], Result[BaseModel]] served: Callable[[BaseModel], bool] -def _send_chat(proxy: ProxyClient, key: str, marker: str) -> Result[BaseModel]: +def _send_chat(proxy: ProxyClient, key: str, model: str, marker: str) -> Result[BaseModel]: return proxy.transport.post( "/chat/completions", headers=proxy.transport.bearer(key), - json=ChatBody(model=PRIMARY_MODEL, messages=[ChatMessage(role="user", content=marker)], max_tokens=5), + json=ChatBody(model=model, messages=[ChatMessage(role="user", content=marker)], max_tokens=5), response_type=ChatResponse, timeout=MAX_SECONDS_PER_REQUEST, ) -def _send_responses(proxy: ProxyClient, key: str, marker: str) -> Result[BaseModel]: +def _send_responses(proxy: ProxyClient, key: str, model: str, marker: str) -> Result[BaseModel]: return proxy.transport.post( "/v1/responses", headers=proxy.transport.bearer(key), - json=ResponsesBody(model=PRIMARY_MODEL, input=marker), + json=ResponsesBody(model=model, input=marker), response_type=ResponsesObject, timeout=MAX_SECONDS_PER_REQUEST, ) +def _send_embeddings(proxy: ProxyClient, key: str, model: str, marker: str) -> Result[BaseModel]: + return proxy.transport.post( + "/embeddings", + headers=proxy.transport.bearer(key), + json=EmbedBody(model=model, input=marker), + response_type=EmbeddingsObject, + timeout=MAX_SECONDS_PER_REQUEST, + ) + + +_CHAT_PRIMARY: Final = LiteLLMParamsBody( + model="openai/gpt-5-mini", api_key="sk-redis-timeout-primary-not-used", api_base=CLOSED_PORT_API_BASE +) +_CHAT_BACKUP: Final = LiteLLMParamsBody( + model="openai/gpt-5-nano", api_key="sk-redis-timeout-backup-not-used", mock_response="ok" +) +_EMBED_PRIMARY: Final = LiteLLMParamsBody( + model="openai/text-embedding-3-small", api_key="sk-redis-timeout-primary-not-used", api_base=CLOSED_PORT_API_BASE +) +_EMBED_BACKUP: Final = LiteLLMParamsBody( + model="openai/text-embedding-3-large", api_key="sk-redis-timeout-backup-not-used", mock_response=[0.1, 0.2, 0.3] +) + ENDPOINTS: Final = ( Endpoint( name="chat_completions", + primary="redis-timeout-primary", + backup="redis-timeout-backup", + primary_params=_CHAT_PRIMARY, + backup_params=_CHAT_BACKUP, send=_send_chat, served=lambda data: isinstance(data, ChatResponse) and bool(data.choices), ), Endpoint( name="responses", + primary="redis-timeout-primary", + backup="redis-timeout-backup", + primary_params=_CHAT_PRIMARY, + backup_params=_CHAT_BACKUP, send=_send_responses, served=lambda data: isinstance(data, ResponsesObject) and bool(data.output), ), -) -BREAKER_OPEN_RE: Final = re.compile(r'^litellm_redis_circuit_breaker_state\{state="open"\} ([0-9.e+]+)$', re.M) -BREAKER_TRANSITIONS_RE: Final = re.compile( - r'^litellm_redis_circuit_breaker_transitions_total\{state="[a-z_]+"\} ([0-9.e+]+)$', re.M + Endpoint( + name="embeddings", + primary="redis-timeout-embed-primary", + backup="redis-timeout-embed-backup", + primary_params=_EMBED_PRIMARY, + backup_params=_EMBED_BACKUP, + send=_send_embeddings, + served=lambda data: isinstance(data, EmbeddingsObject) and bool(data.data), + ), ) @@ -126,48 +177,51 @@ def _metric(proxy: ProxyClient, pattern: re.Pattern[str]) -> float: return sum(float(match.group(1)) for match in pattern.finditer(body)) +def _rss_bytes(proxy: ProxyClient) -> float | None: + """The proxy's resident memory from the Prometheus process collector, which reads /proc and + so reports on Linux only; None where the metric is absent.""" + body = proxy.probe("/metrics", params=NoBody()).body + match = RSS_RE.search(body) + return float(match.group(1)) if match else None + + class TestRedisTimeout: @pytest.mark.parametrize("endpoint", ENDPOINTS, ids=[endpoint.name for endpoint in ENDPOINTS]) @pytest.mark.covers( "reliability.circuit_breaker.redis_timeout.stays_responsive", - exercised_on=["chat_completions", "responses"], + exercised_on=["chat_completions", "responses", "embeddings"], ) def test_retries_under_redis_timeouts_keep_answering( self, client: ComplexityRouterClient, resources: ResourceManager, endpoint: Endpoint, paused_redis: None ) -> None: proxy = client.proxy - primary_id = proxy.create_model( - PRIMARY_MODEL, - LiteLLMParamsBody( - model=BACKING_MODEL, api_key="sk-redis-timeout-primary-not-used", api_base=CLOSED_PORT_API_BASE - ), - ) + primary_id = proxy.create_model(endpoint.primary, endpoint.primary_params) resources.defer(lambda: proxy.delete_model(primary_id)) - backup_id = proxy.create_model( - BACKUP_MODEL, - LiteLLMParamsBody(model=BACKING_MODEL, api_key="sk-redis-timeout-backup-not-used", mock_response="ok"), - ) + backup_id = proxy.create_model(endpoint.backup, endpoint.backup_params) resources.defer(lambda: proxy.delete_model(backup_id)) - timeouts_before = _metric(proxy, TIMEOUT_FAILURES_RE) - transitions_before = _metric(proxy, BREAKER_TRANSITIONS_RE) key = proxy.generate_key( KeyGenerateBody( - models=[PRIMARY_MODEL, BACKUP_MODEL], key_alias=f"e2e-redis-timeout-{endpoint.name}-{unique_marker()}" + models=[endpoint.primary, endpoint.backup], + key_alias=f"e2e-redis-timeout-{endpoint.name}-{unique_marker()}", ) ) resources.defer(lambda: proxy.delete_key(key)) + timeouts_before = _metric(proxy, TIMEOUT_FAILURES_RE) + transitions_before = _metric(proxy, BREAKER_TRANSITIONS_RE) + fallbacks_before = _metric(proxy, FALLBACKS_RE) + rss_before = _rss_bytes(proxy) latencies: list[float] = [] for request_number in range(1, REQUESTS + 1): started = time.monotonic() - result = endpoint.send(proxy, key, f"redis timeout {unique_marker()} {request_number}") + result = endpoint.send(proxy, key, endpoint.primary, f"redis timeout {unique_marker()} {request_number}") elapsed = time.monotonic() - started assert isinstance(result, Success), ( f"{endpoint.name} request {request_number} failed after {elapsed:.1f}s with Redis timing out: {result}; " f"earlier requests took {[round(seconds, 2) for seconds in latencies]}" ) assert endpoint.served(result.data), ( - f"{endpoint.name} request {request_number}: fallback to {BACKUP_MODEL} returned no output" + f"{endpoint.name} request {request_number}: fallback to {endpoint.backup} returned no output" ) assert elapsed < MAX_SECONDS_PER_REQUEST, ( f"{endpoint.name} request {request_number} took {elapsed:.1f}s with Redis timing out; " @@ -191,6 +245,20 @@ class TestRedisTimeout: f"/health/liveliness took {liveliness_seconds:.1f}s after the loop; the worker is stalled" ) + if rss_before is not None: + rss_after = _rss_bytes(proxy) + assert rss_after is not None + assert rss_after - rss_before <= MAX_RSS_GROWTH_BYTES, ( + f"proxy RSS grew {(rss_after - rss_before) / 2**20:.0f} MB across {REQUESTS} {endpoint.name} requests " + "with Redis timing out; on v1.100.0 this path grew by gigabytes" + ) + + fallbacks = _metric(proxy, FALLBACKS_RE) - fallbacks_before + assert fallbacks >= REQUESTS, ( + f"only {fallbacks:.0f} successful fallbacks were counted across {REQUESTS} {endpoint.name} requests; " + "the closed-port primary did not fail every request, so the retry path was not exercised" + ) + timeouts_total = _metric(proxy, TIMEOUT_FAILURES_RE) timeouts = timeouts_total - timeouts_before transitions = _metric(proxy, BREAKER_TRANSITIONS_RE) - transitions_before @@ -210,3 +278,5 @@ class TestRedisTimeout: f"only {len(rows)} of {REQUESTS} {endpoint.name} requests reached the spend log; " "a Redis outage must not lose spend rows" ) + failed_rows = [row.status for row in rows if row.status not in (None, "success")] + assert not failed_rows, f"{len(failed_rows)} {endpoint.name} spend rows are not successes: {failed_rows[:3]}" From 8b965880efd8e28321c1db74850fbb2b6f59242c Mon Sep 17 00:00:00 2001 From: dclark Date: Thu, 10 Sep 2026 11:29:37 +0100 Subject: [PATCH 025/155] fix(proxy): prevent spend counter double counting --- litellm/proxy/db/spend_counter_reseed.py | 24 ++++++++-- .../proxy/db/test_spend_counter_reseed.py | 48 +++++++++++++++++++ 2 files changed, 69 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index e35b1c8c82b..c6060a47f84 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -245,7 +245,16 @@ class SpendCounterReseed: value=current_value, ) else: - await spend_counter_cache.async_increment_cache(key=counter_key, value=db_spend, refresh_ttl=True) + # Repair/reservations can populate the counter during the DB read. + # Seed a floor without adding the database balance again. + # No await between read/compare/write: atomic within this worker. + cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + current_value = float(db_spend) + if cached is not None: + current_value = max(current_value, float(cached)) + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=current_value + ) except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", @@ -438,11 +447,20 @@ class SpendCounterReseed: value=current_value, ) else: - await spend_counter_cache.async_increment_cache(key=counter_key, value=window_spend) + # Repair/reservations can populate the counter during the DB read. + # Seed a floor without adding the database balance again. + # No await between read/compare/write: atomic within this worker. + cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + current_value = float(window_spend) + if cached is not None: + current_value = max(current_value, float(cached)) + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=current_value + ) except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced_window: failed to warm counter %s", counter_key, ) raise - return window_spend + return current_value 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 0361d5cfe8f..a0cf9128cc9 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -15,6 +15,7 @@ from typing import Final import pytest from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed @@ -270,6 +271,53 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): assert prisma.db.litellm_spendlogs.call_count == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("window", [False, True], ids=["primary", "window"]) +@pytest.mark.parametrize("concurrent_spend", [989.01459411, 995.0, 900.0]) +async def test_cold_reseed_does_not_add_database_spend_to_concurrent_cache( + monkeypatch: pytest.MonkeyPatch, + window: bool, + concurrent_spend: float, +): + """A repair/reservation write may populate the counter while DB read runs. + + Reseeding must establish the larger value, not increment the concurrent + value by the same authoritative spend a second time. + """ + cache = DualCache(in_memory_cache=InMemoryCache()) + counter_key = ( + "spend:team:team-1:window:1d" if window else "spend:user:user-1" + ) + db_spend = 989.01459411 + + async def read_db(*args, **kwargs): + cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend) + return db_spend + + if window: + monkeypatch.setattr(SpendCounterReseed, "window_from_db", staticmethod(read_db)) + result = await SpendCounterReseed.coalesced_window( + prisma_client=None, + spend_counter_cache=cache, + counter_key=counter_key, + entity_type="Team", + entity_id="team-1", + window_duration="1d", + window_start=WINDOW_START, + ) + else: + monkeypatch.setattr(SpendCounterReseed, "from_db", staticmethod(read_db)) + result = await SpendCounterReseed.coalesced( + prisma_client=None, + spend_counter_cache=cache, + counter_key=counter_key, + ) + + expected = max(db_spend, concurrent_spend) + assert cache.in_memory_cache.get_cache(key=counter_key) == expected + assert result == expected + + @pytest.mark.asyncio async def test_end_user_from_db_reads_the_end_user_row_by_user_id(): prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0)) From a39c35878569f74c50de3295320457dd7bd6adb2 Mon Sep 17 00:00:00 2001 From: dclark Date: Thu, 10 Sep 2026 13:44:08 +0100 Subject: [PATCH 026/155] fix(proxy): preserve local spend adjustments during reseeding --- litellm/proxy/db/spend_counter_reseed.py | 37 +++--- litellm/proxy/proxy_server.py | 12 +- .../proxy/db/test_spend_counter_reseed.py | 117 ++++++++++++------ 3 files changed, 102 insertions(+), 64 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index c6060a47f84..0131a67db8b 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -104,6 +104,15 @@ class SpendCounterReseed: SpendCounterReseed._locks.popitem(last=False) return lock + @staticmethod + async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None: + """Apply local deltas after an in-flight reseed establishes the spend balance.""" + lock: Final = await SpendCounterReseed._get_lock(counter_key) + async with lock: + return await spend_counter_cache.async_increment_cache( + key=counter_key, value=increment, local_only=True, refresh_ttl=True + ) + @staticmethod async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: """ @@ -245,16 +254,10 @@ class SpendCounterReseed: value=current_value, ) else: - # Repair/reservations can populate the counter during the DB read. - # Seed a floor without adding the database balance again. - # No await between read/compare/write: atomic within this worker. - cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - current_value = float(db_spend) - if cached is not None: - current_value = max(current_value, float(cached)) - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=current_value - ) + cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + seeded_spend: Final = max(db_spend, float(cached_spend)) if cached_spend is not None else db_spend + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend) + return seeded_spend except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", @@ -447,16 +450,12 @@ class SpendCounterReseed: value=current_value, ) else: - # Repair/reservations can populate the counter during the DB read. - # Seed a floor without adding the database balance again. - # No await between read/compare/write: atomic within this worker. - cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - current_value = float(window_spend) - if cached is not None: - current_value = max(current_value, float(cached)) - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=current_value + cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + seeded_spend: Final = ( + max(window_spend, float(cached_spend)) if cached_spend is not None else window_spend ) + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend) + return seeded_spend except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced_window: failed to warm counter %s", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9269fd48e6c..e8608661c48 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3330,10 +3330,8 @@ async def _increment_spend_counter_cache(counter_key: str, increment: float): ) return current_value - return await spend_counter_cache.async_increment_cache( - key=counter_key, - value=increment, - refresh_ttl=True, + return await SpendCounterReseed.increment_in_memory( + spend_counter_cache=spend_counter_cache, counter_key=counter_key, increment=increment ) @@ -3356,10 +3354,8 @@ async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncreme redis_cache: Final = spend_counter_cache.redis_cache if redis_cache is None: for item in pending: - await spend_counter_cache.async_increment_cache( - key=item.counter_key, - value=item.increment, - refresh_ttl=True, + await SpendCounterReseed.increment_in_memory( + spend_counter_cache=spend_counter_cache, counter_key=item.counter_key, increment=item.increment ) return ttl: Final = redis_cache.get_ttl() 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 a0cf9128cc9..69ec174903b 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -8,6 +8,7 @@ allowed to run: only when the row is missing or belongs to an older window. from __future__ import annotations import asyncio +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final @@ -79,6 +80,39 @@ def _row(window_start: datetime, spend: float) -> SimpleNamespace: return SimpleNamespace(window_start=window_start, spend=spend) +class _PausedSpendTable: + def __init__(self, spend: float) -> None: + self.spend: Final = spend + self.read_started: Final = asyncio.Event() + self.resume_read: Final = asyncio.Event() + + async def find_unique(self, where: Mapping[str, object]) -> SimpleNamespace: + self.read_started.set() + await self.resume_read.wait() + return _row(WINDOW_START, self.spend) + + +async def _reseed_with_paused_table( + table: _PausedSpendTable, cache: DualCache, counter_key: str, window: bool +) -> float | None: + prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_usertable=table, litellm_budgetwindowspend=table)) + if window: + return await SpendCounterReseed.coalesced_window( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + entity_type="Team", + entity_id="team-1", + window_duration="1d", + window_start=WINDOW_START, + ) + return await SpendCounterReseed.coalesced( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + ) + + @pytest.mark.asyncio async def test_window_from_table_reads_row_by_primary_key(): """The lookup must use the table's own entity_type values ("key"), not the @@ -275,49 +309,59 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): @pytest.mark.parametrize("window", [False, True], ids=["primary", "window"]) @pytest.mark.parametrize("concurrent_spend", [989.01459411, 995.0, 900.0]) async def test_cold_reseed_does_not_add_database_spend_to_concurrent_cache( - monkeypatch: pytest.MonkeyPatch, window: bool, concurrent_spend: float, -): - """A repair/reservation write may populate the counter while DB read runs. +) -> None: + cache: Final = DualCache(in_memory_cache=InMemoryCache()) + counter_key: Final = "spend:team:team-1:window:1d" if window else "spend:user:user-1" + db_spend: Final = 989.01459411 + table: Final = _PausedSpendTable(db_spend) + reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window)) - Reseeding must establish the larger value, not increment the concurrent - value by the same authoritative spend a second time. - """ - cache = DualCache(in_memory_cache=InMemoryCache()) - counter_key = ( - "spend:team:team-1:window:1d" if window else "spend:user:user-1" - ) - db_spend = 989.01459411 + await asyncio.wait_for(table.read_started.wait(), timeout=5) + cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend) + table.resume_read.set() + result: Final = await asyncio.wait_for(reseed_task, timeout=5) - async def read_db(*args, **kwargs): - cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend) - return db_spend - - if window: - monkeypatch.setattr(SpendCounterReseed, "window_from_db", staticmethod(read_db)) - result = await SpendCounterReseed.coalesced_window( - prisma_client=None, - spend_counter_cache=cache, - counter_key=counter_key, - entity_type="Team", - entity_id="team-1", - window_duration="1d", - window_start=WINDOW_START, - ) - else: - monkeypatch.setattr(SpendCounterReseed, "from_db", staticmethod(read_db)) - result = await SpendCounterReseed.coalesced( - prisma_client=None, - spend_counter_cache=cache, - counter_key=counter_key, - ) - - expected = max(db_spend, concurrent_spend) + expected: Final = max(db_spend, concurrent_spend) assert cache.in_memory_cache.get_cache(key=counter_key) == expected assert result == expected +@pytest.mark.asyncio +@pytest.mark.parametrize("window", [False, True], ids=["primary", "window"]) +@pytest.mark.parametrize("batch", [False, True], ids=["single_increment", "batch_increment"]) +@pytest.mark.parametrize("increment", [5.0, -5.0], ids=["charge", "refund"]) +async def test_cold_reseed_preserves_concurrent_local_increment( + monkeypatch: pytest.MonkeyPatch, window: bool, batch: bool, increment: float +) -> None: + from litellm.proxy import proxy_server + + cache: Final = DualCache(in_memory_cache=InMemoryCache()) + counter_key: Final = ( + f"spend:team:concurrent-{batch}-{increment}:window:1d" + if window + else f"spend:user:concurrent-{batch}-{increment}" + ) + table: Final = _PausedSpendTable(100.0) + monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) + reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window)) + await asyncio.wait_for(table.read_started.wait(), timeout=5) + + increment_task: Final = asyncio.create_task( + proxy_server._apply_spend_counter_increments( + pending=(proxy_server._PendingSpendIncrement(counter_key=counter_key, increment=increment),) + ) + if batch + else proxy_server._increment_spend_counter_cache(counter_key=counter_key, increment=increment) + ) + await asyncio.sleep(0) + table.resume_read.set() + await asyncio.wait_for(asyncio.gather(reseed_task, increment_task), timeout=5) + + assert cache.in_memory_cache.get_cache(key=counter_key) == 100.0 + increment + + @pytest.mark.asyncio async def test_end_user_from_db_reads_the_end_user_row_by_user_id(): prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0)) @@ -352,8 +396,7 @@ async def test_end_user_from_db_ignores_other_counter_kinds_without_touching_the @pytest.mark.asyncio async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_error(): assert ( - await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42") - is None + await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42") is None ) assert ( await SpendCounterReseed.end_user_from_db( From 60c7dd8348cb1634784efd908afa046233858585 Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 10 Sep 2026 19:13:49 +0000 Subject: [PATCH 027/155] fix(model_prices): add deepseek-flash and gpt-live-1, bill DeepSeek legacy flash aliases at Flash rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 77 +++++++++++++++---- model_prices_and_context_window.json | 77 +++++++++++++++---- tests/test_litellm/test_utils.py | 31 ++++++-- 3 files changed, 146 insertions(+), 39 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d8160a3d842..7f3e0565f88 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -56459,17 +56459,43 @@ "supports_reasoning": true, "source": "https://serverless.tensormesh.ai/v1/models/openrouter" }, - "deepseek-v4-flash": { + "deepseek-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "deepseek-v4-flash": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, + "litellm_provider": "deepseek", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -56487,15 +56513,15 @@ }, "deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -56539,15 +56565,15 @@ }, "deepseek/deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -56565,15 +56591,15 @@ }, "deepseek/deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -56929,6 +56955,23 @@ ], "supports_audio_input": true }, + "gpt-live-1": { + "input_cost_per_second": 0.0008333333333333334, + "litellm_provider": "openai", + "mode": "realtime", + "source": "https://developers.openai.com/api/docs/models/gpt-live-1", + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true + }, "gpt-realtime-translate": { "input_cost_per_second": 0.0005666666666666667, "litellm_provider": "openai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d8160a3d842..7f3e0565f88 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -56459,17 +56459,43 @@ "supports_reasoning": true, "source": "https://serverless.tensormesh.ai/v1/models/openrouter" }, - "deepseek-v4-flash": { + "deepseek-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "deepseek-v4-flash": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, + "litellm_provider": "deepseek", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -56487,15 +56513,15 @@ }, "deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -56539,15 +56565,15 @@ }, "deepseek/deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -56565,15 +56591,15 @@ }, "deepseek/deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -56929,6 +56955,23 @@ ], "supports_audio_input": true }, + "gpt-live-1": { + "input_cost_per_second": 0.0008333333333333334, + "litellm_provider": "openai", + "mode": "realtime", + "source": "https://developers.openai.com/api/docs/models/gpt-live-1", + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true + }, "gpt-realtime-translate": { "input_cost_per_second": 0.0005666666666666667, "litellm_provider": "openai", diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8b186be43e5..34b7f263dcd 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4071,7 +4071,7 @@ def test_deepseek_v4_models_in_cost_map(): configured in model_prices_and_context_window.json. Prices sourced from https://api-docs.deepseek.com/quick_start/pricing: - - deepseek-v4-flash: $0.44/M input, $1.32/M output + - deepseek-v4-flash: $0.30/M input, $1.20/M output - deepseek-v4-pro: $1.32/M input, $3.96/M output Closes https://github.com/BerriAI/litellm/issues/26709 @@ -4085,7 +4085,7 @@ def test_deepseek_v4_models_in_cost_map(): # --- bare model names --- for key, expected_input, expected_output, expected_cache in [ - ("deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), + ("deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09), ("deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), ]: info = model_cost.get(key) @@ -4101,7 +4101,7 @@ def test_deepseek_v4_models_in_cost_map(): # --- provider-prefixed names --- for key, expected_input, expected_output, expected_cache in [ - ("deepseek/deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), + ("deepseek/deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09), ("deepseek/deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), ]: info = model_cost.get(key) @@ -4129,7 +4129,7 @@ def test_deepseek_v4_models_in_backup_cost_map(): # --- bare model names --- for key, expected_input, expected_output, expected_cache in [ - ("deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), + ("deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09), ("deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), ]: info = model_cost.get(key) @@ -4143,7 +4143,7 @@ def test_deepseek_v4_models_in_backup_cost_map(): # --- provider-prefixed names --- for key, expected_input, expected_output, expected_cache in [ - ("deepseek/deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), + ("deepseek/deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09), ("deepseek/deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), ]: info = model_cost.get(key) @@ -4155,6 +4155,27 @@ def test_deepseek_v4_models_in_backup_cost_map(): assert info["cache_read_input_token_cost"] == expected_cache +def test_deepseek_flash_completion_cost(): + from litellm.types.utils import ModelResponse + + response = ModelResponse( + model="deepseek-flash", + usage=Usage( + prompt_tokens=1_000_000, + completion_tokens=1_000_000, + total_tokens=2_000_000, + ), + ) + + cost = litellm.completion_cost( + completion_response=response, + model="deepseek-flash", + custom_llm_provider="deepseek", + ) + + assert cost == pytest.approx(1.50, abs=1e-9) + + _FIREWORKS_MODELS = [ ( "accounts/fireworks/models/glm-5p2", From 896f0dff2e0f59d8d1d2ac8cc0c75c1ee611f881 Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 10 Sep 2026 19:14:28 +0000 Subject: [PATCH 028/155] fix(model_prices): add provider-prefixed deepseek/deepseek-flash Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 26 +++++++++++++++++++ model_prices_and_context_window.json | 26 +++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7f3e0565f88..91a24bfae2e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -56563,6 +56563,32 @@ "supports_tool_choice": true, "supports_vision": false }, + "deepseek/deepseek-flash": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, + "litellm_provider": "deepseek", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "deepseek/deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 6e-09, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7f3e0565f88..91a24bfae2e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -56563,6 +56563,32 @@ "supports_tool_choice": true, "supports_vision": false }, + "deepseek/deepseek-flash": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, + "litellm_provider": "deepseek", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "deepseek/deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 6e-09, From ace6de88a4cb3dc052908f394f64c17a73e5b816 Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 10 Sep 2026 19:31:11 +0000 Subject: [PATCH 029/155] fix(tests): load local model costs for DeepSeek flash regression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 34b7f263dcd..a2d430e8807 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4155,6 +4155,7 @@ def test_deepseek_v4_models_in_backup_cost_map(): assert info["cache_read_input_token_cost"] == expected_cache +@pytest.mark.usefixtures("local_model_cost_map") def test_deepseek_flash_completion_cost(): from litellm.types.utils import ModelResponse From 8deb465346317baf550acb948ff2c72bdf8b9315 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:18:52 -0700 Subject: [PATCH 030/155] fix(realtime): dial Azure's GA realtime upstream for GA clients Azure realtime defaulted to the beta upstream whenever realtime_protocol was not configured, so a GA client's session.update (session.type, output_modalities, nested audio) was forwarded unchanged to /openai/realtime and Azure rejected it with "Unknown parameter: 'session.type'" on gpt-realtime and gpt-realtime-1.5. The unset default now follows the client the way the OpenAI handler already does: a client that sends OpenAI-Beta: realtime=v1 keeps the beta upstream, any other client gets /openai/v1/realtime. An explicit realtime_protocol in litellm_params or LITELLM_AZURE_REALTIME_PROTOCOL still wins. --- litellm/realtime_api/main.py | 22 ++++++-- tests/test_litellm/realtime_api/test_main.py | 59 +++++++++++++++++++- 2 files changed, 74 insertions(+), 7 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index b824a5928c6..5310efba1c4 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -14,6 +14,7 @@ from litellm.constants import ( request_timeout, ) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.realtime_streaming import client_sent_openai_beta_realtime_header from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.xai.common_utils import XAIModelInfo @@ -413,14 +414,14 @@ async def _arealtime( api_version = api_version or litellm_params.api_version or "2024-10-01-preview" - realtime_protocol = ( + configured_realtime_protocol: Final = ( kwargs.get("realtime_protocol") or litellm_params.get("realtime_protocol") or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL") ) - if realtime_protocol is None and (query_params or {}).get("intent") == "transcription": - realtime_protocol = "GA" - realtime_protocol = realtime_protocol or "beta" + realtime_protocol: Final = _azure_realtime_protocol_for_client( + configured_realtime_protocol, query_params=query_params, websocket=websocket + ) resolved_azure_ad_token: Final = ( None if api_key else get_azure_ad_token(GenericLiteLLMParams(**kwargs, azure_ad_token=azure_ad_token)) ) @@ -576,6 +577,19 @@ def _is_transcription_only_realtime_model(model: str, custom_llm_provider: str) _TRANSCRIPTION_QUERY_PARAMS: Final[RealtimeQueryParams] = {"intent": "transcription"} +def _azure_realtime_protocol_for_client( + configured_protocol: object, + *, + query_params: RealtimeQueryParams | None, + websocket: "WebSocket", +) -> str: + if isinstance(configured_protocol, str) and configured_protocol: + return configured_protocol + if (query_params or {}).get("intent") == "transcription": + return "GA" + return "beta" if client_sent_openai_beta_realtime_header(websocket) else "GA" + + def _azure_realtime_health_protocol( model: str, realtime_protocol: str | None, model_params: Mapping[str, object] ) -> tuple[str, RealtimeQueryParams | None]: diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 761e87ac764..8a9abe819e5 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -327,7 +327,60 @@ async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai api_key="fake-key", litellm_logging_obj=FakeLogging(), ) - assert connect.url == ( - "wss://my-project.services.ai.azure.com/openai/realtime" - "?api-version=2024-10-01-preview&deployment=gpt-realtime-mini" + assert connect.url == "wss://my-project.services.ai.azure.com/openai/v1/realtime?model=gpt-realtime-mini" + + +class _ClientWebSocketWithHeaders: + def __init__(self, headers: tuple[tuple[bytes, bytes], ...]) -> None: + self.scope: Final = {"headers": headers} + + +_GA_CLIENT: Final = _ClientWebSocketWithHeaders(headers=()) +_BETA_CLIENT: Final = _ClientWebSocketWithHeaders(headers=((b"openai-beta", b"realtime=v1"),)) + + +async def _azure_backend_url_dialed_for(websocket: _ClientWebSocketWithHeaders, **kwargs: object) -> str | None: + connect: Final = _ConnectThatStopsAfterCapturingTheUrl() + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="azure/gpt-realtime", + websocket=websocket, + api_base="https://my-endpoint.openai.azure.com", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + **kwargs, + ) + return connect.url + + +@pytest.mark.asyncio +async def test_arealtime_azure_ga_client_without_beta_header_dials_the_ga_upstream(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + assert ( + await _azure_backend_url_dialed_for(_GA_CLIENT) + == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-realtime" + ) + + +@pytest.mark.asyncio +async def test_arealtime_azure_beta_header_client_keeps_the_beta_upstream(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + assert await _azure_backend_url_dialed_for(_BETA_CLIENT) == ( + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" + ) + + +@pytest.mark.asyncio +async def test_arealtime_azure_explicit_beta_protocol_wins_over_a_ga_client(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + assert await _azure_backend_url_dialed_for(_GA_CLIENT, realtime_protocol="beta") == ( + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" + ) + + +@pytest.mark.asyncio +async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypatch): + monkeypatch.setenv("LITELLM_AZURE_REALTIME_PROTOCOL", "beta") + assert await _azure_backend_url_dialed_for(_GA_CLIENT) == ( + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" ) From 8fe2094a5579e048c5d79f374f4a6ad7d7261326 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:46:13 -0700 Subject: [PATCH 031/155] refactor(realtime): move the Azure protocol picker into the Azure realtime handler --- litellm/llms/azure/realtime/handler.py | 23 +++++++++++++++++++++-- litellm/realtime_api/main.py | 18 ++---------------- 2 files changed, 23 insertions(+), 18 deletions(-) diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index e9913f0108d..88813c21cd9 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -6,17 +6,23 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. from collections.abc import Mapping from types import MappingProxyType -from typing import Any, Final, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast from litellm._logging import _redact_string, verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from ....litellm_core_utils.realtime_streaming import RealTimeStreaming +from ....litellm_core_utils.realtime_streaming import ( + RealTimeStreaming, + client_sent_openai_beta_realtime_header, +) from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..azure import AzureChatCompletion +if TYPE_CHECKING: + from fastapi import WebSocket + # BACKEND_WS_URL = "ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01" @@ -31,6 +37,19 @@ async def forward_messages(client_ws: Any, backend_ws: Any): pass +def azure_realtime_protocol_for_client( + configured_protocol: object, + *, + query_params: RealtimeQueryParams | None, + websocket: "WebSocket", +) -> str: + if isinstance(configured_protocol, str) and configured_protocol: + return configured_protocol + if (query_params or {}).get("intent") == "transcription": + return "GA" + return "beta" if client_sent_openai_beta_realtime_header(websocket) else "GA" + + class _ProxyClientWebSocket(Protocol): """Client-facing websocket handle: this path only closes it after a failed handshake.""" diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 5310efba1c4..d42e7e18b75 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -14,7 +14,6 @@ from litellm.constants import ( request_timeout, ) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.litellm_core_utils.realtime_streaming import client_sent_openai_beta_realtime_header from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.xai.common_utils import XAIModelInfo @@ -34,7 +33,7 @@ from litellm.utils import ProviderConfigManager from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ..llms.azure.common_utils import get_azure_ad_token -from ..llms.azure.realtime.handler import AzureOpenAIRealtime +from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client from ..llms.bedrock.realtime.handler import BedrockRealtime from ..llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..llms.openai.realtime.handler import OpenAIRealtime @@ -419,7 +418,7 @@ async def _arealtime( or litellm_params.get("realtime_protocol") or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL") ) - realtime_protocol: Final = _azure_realtime_protocol_for_client( + realtime_protocol: Final = azure_realtime_protocol_for_client( configured_realtime_protocol, query_params=query_params, websocket=websocket ) resolved_azure_ad_token: Final = ( @@ -577,19 +576,6 @@ def _is_transcription_only_realtime_model(model: str, custom_llm_provider: str) _TRANSCRIPTION_QUERY_PARAMS: Final[RealtimeQueryParams] = {"intent": "transcription"} -def _azure_realtime_protocol_for_client( - configured_protocol: object, - *, - query_params: RealtimeQueryParams | None, - websocket: "WebSocket", -) -> str: - if isinstance(configured_protocol, str) and configured_protocol: - return configured_protocol - if (query_params or {}).get("intent") == "transcription": - return "GA" - return "beta" if client_sent_openai_beta_realtime_header(websocket) else "GA" - - def _azure_realtime_health_protocol( model: str, realtime_protocol: str | None, model_params: Mapping[str, object] ) -> tuple[str, RealtimeQueryParams | None]: From 4a3950cf678fa1b65acfc468c61ae3dcbb67b472 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:57:48 -0700 Subject: [PATCH 032/155] refactor(realtime): type the Azure protocol picker with the streaming module's websocket protocol --- litellm/litellm_core_utils/realtime_streaming.py | 8 ++++---- litellm/llms/azure/realtime/handler.py | 8 +++----- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 75046f2cf87..4923bdda305 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -90,12 +90,12 @@ class _ResponseDoneBody(TypedDict, total=False): output: ReadOnly[Sequence[Mapping[str, object]]] -class _ScopedWebSocket(Protocol): +class ScopedWebSocket(Protocol): @property def scope(self) -> _ASGIScope: ... -class _ClientWebSocket(_ScopedWebSocket, Protocol): +class _ClientWebSocket(ScopedWebSocket, Protocol): async def send_text(self, data: str) -> None: ... async def receive_text(self) -> str: ... async def close(self, code: int = 1000, reason: str | None = None) -> None: ... @@ -1149,7 +1149,7 @@ class RealTimeStreaming: ) @staticmethod - def _detect_beta_header(websocket: _ScopedWebSocket) -> bool: + def _detect_beta_header(websocket: ScopedWebSocket) -> bool: """Return True if the client sent 'OpenAI-Beta: realtime=v1'. Checks the raw ASGI scope headers so it works for both FastAPI WebSocket @@ -1584,6 +1584,6 @@ class RealTimeStreaming: verbose_logger.debug("Could not relay the upstream close to the client: %s", e) -def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool: +def client_sent_openai_beta_realtime_header(websocket: ScopedWebSocket) -> bool: """True when the client WebSocket includes ``OpenAI-Beta: realtime=v1``.""" return RealTimeStreaming._detect_beta_header(websocket) diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 88813c21cd9..146915dd6fd 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -6,7 +6,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. from collections.abc import Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from typing import Any, Final, Protocol, cast from litellm._logging import _redact_string, verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES @@ -15,14 +15,12 @@ from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ....litellm_core_utils.realtime_streaming import ( RealTimeStreaming, + ScopedWebSocket, client_sent_openai_beta_realtime_header, ) from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..azure import AzureChatCompletion -if TYPE_CHECKING: - from fastapi import WebSocket - # BACKEND_WS_URL = "ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01" @@ -41,7 +39,7 @@ def azure_realtime_protocol_for_client( configured_protocol: object, *, query_params: RealtimeQueryParams | None, - websocket: "WebSocket", + websocket: ScopedWebSocket, ) -> str: if isinstance(configured_protocol, str) and configured_protocol: return configured_protocol From be7dce30a3ab832bdd104b2d2f15bace5d622b4d Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 10 Sep 2026 15:16:24 -0700 Subject: [PATCH 033/155] ci(e2e): pin the Redis timeout workflow's Postgres image by digest Co-Authored-By: Claude Code --- .github/workflows/test-e2e-redis-timeout.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-e2e-redis-timeout.yml b/.github/workflows/test-e2e-redis-timeout.yml index 9c318a1346c..501af0f240f 100644 --- a/.github/workflows/test-e2e-redis-timeout.yml +++ b/.github/workflows/test-e2e-redis-timeout.yml @@ -15,7 +15,7 @@ jobs: timeout-minutes: 30 services: postgres: - image: postgres:16.6 + image: postgres:16.6@sha256:557fea37a744d5f4c8faab304b0a90858b53ab119735a88c131fd19dab802f36 env: POSTGRES_USER: llmproxy POSTGRES_PASSWORD: dbpassword9090 From 857b9ad7d344b2cb6d76f0f0b4072b464ad3b297 Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 11 Sep 2026 01:08:20 +0000 Subject: [PATCH 034/155] fix(ui): jump straight to the last Request Logs page instead of advancing one page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_management_endpoints.py | 11 +++-- .../test_spend_management_endpoints.py | 42 +++++++++++++++++++ .../view_logs/RequestLogsPanel.test.tsx | 38 +++++++++++++++++ .../components/view_logs/RequestLogsPanel.tsx | 7 ++-- 4 files changed, 91 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index da79328fa59..4a5995167f7 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2944,7 +2944,10 @@ async def _ui_session_grouped_spend_logs( next ``page_size`` sessions ordered by ``(MAX(startTime), session_key, api_key)``, resumed from the ``session_cursor`` keyset ``'||'`` instead of an OFFSET, so - page depth does not degrade the query plan. Each session is represented + page depth does not degrade the query plan. A request for ``page > 1`` + without a cursor (the UI jumping straight to the last page, or back to a + page it never walked through) falls back to ``OFFSET (page - 1) * + page_size``, bounded by the capped total. Each session is represented by its newest non-MCP row, enriched by ``_build_ui_spend_logs_response`` exactly like the flat listing, and the response carries ``next_session_cursor`` / ``has_more`` while ``total`` counts sessions @@ -2963,6 +2966,8 @@ async def _ui_session_grouped_spend_logs( ) cursor_params: Final[tuple[object, ...]] = cursor if cursor else () limit_index: Final = next_param_index + len(cursor_params) + offset_params: Final[tuple[int, ...]] = ((page - 1) * page_size,) if cursor is None and page > 1 else () + offset_clause: Final = f"OFFSET ${limit_index + 1}" if offset_params else "" page_query: Final = f""" SELECT {_SESSION_KEY_EXPR} AS session_key, @@ -2973,10 +2978,10 @@ async def _ui_session_grouped_spend_logs( GROUP BY {_SESSION_GROUP_KEY_SQL} {having_clause} ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction} - LIMIT ${limit_index} + LIMIT ${limit_index} {offset_clause} """ page_rows: Final[Sequence[_SessionPageRow]] = await _query_raw( - prisma_client, page_query, *sql_params, *cursor_params, page_size + 1 + prisma_client, page_query, *sql_params, *cursor_params, page_size + 1, *offset_params ) has_more: Final = len(page_rows) > page_size diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 671a8ae63fc..5052cf3c085 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -6742,6 +6742,48 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor(client, monkeypatch): + """page > 1 with no session_cursor (the UI's last-page jump) skips (page - 1) * page_size sessions by OFFSET.""" + page_rows = [_session_page_row("sess-3", "2026-08-29 06:00:00")] + reps = [_session_representative_row("req-3", "sess-3")] + mock_prisma = _session_grouped_mock_prisma(page_rows, 60, reps) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page": 3, + "page_size": 25, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["page"] == 3 + assert data["has_more"] is False + assert [row["request_id"] for row in data["data"]] == ["req-3"] + + page_query_call = mock_prisma.db.query_raw.await_args_list[0] + page_query_sql = page_query_call.args[0] + assert "HAVING" not in page_query_sql + assert "OFFSET" in page_query_sql + assert page_query_call.args[-2:] == (26, 50), "LIMIT page_size + 1 then OFFSET (page - 1) * page_size" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort( client, monkeypatch diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index 295446186b1..e760d0b8dc3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -258,6 +258,44 @@ describe("RequestLogsPanel", () => { }); }); + it("jumps straight to the last page without a cursor when the last-page button is clicked", async () => { + const firstPage = Array.from({ length: 25 }, (_, index) => logEntry({ request_id: `req-${index}` })); + const lastPage = Array.from({ length: 10 }, (_, index) => logEntry({ request_id: `req-last-${index}` })); + vi.mocked(uiSpendLogsCall).mockImplementation(async ({ page }) => + page === 3 + ? { + data: lastPage, + total: 60, + page: 3, + page_size: 25, + total_pages: 3, + next_session_cursor: null, + has_more: false, + } + : { + data: firstPage, + total: 60, + page: 1, + page_size: 25, + total_pages: 3, + next_session_cursor: "2026-07-07 09:50:13|key-1|sess-1", + has_more: true, + }, + ); + renderPanel(); + + await waitFor(() => expect(row("req-0")).not.toBeNull()); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 3"); + fireEvent.click(screen.getByTestId("pagination-last")); + + await waitFor(() => expect(row("req-last-0")).not.toBeNull()); + expect(lastCall()?.page).toBe(3); + expect(lastCall()?.params?.session_cursor).toBeUndefined(); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 3 of 3"); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 51-60 of 60"); + expect(vi.mocked(uiSpendLogsCall).mock.calls.filter(([options]) => options.page === 2)).toHaveLength(0); + }); + it("drops the cursor and returns to the first page when a filter changes", async () => { const firstPage = Array.from({ length: 50 }, (_, index) => logEntry({ request_id: `req-${index}` })); vi.mocked(uiSpendLogsCall).mockResolvedValue({ diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx index 6e984297bf2..4f39bb3b79b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx @@ -209,15 +209,14 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, setPagination({ ...requested, pageIndex: 0 }); return; } - if (requested.pageIndex <= pagination.pageIndex) { + if (requested.pageIndex !== pagination.pageIndex + 1) { setPagination(requested); return; } const nextCursor = filteredLogs.next_session_cursor; if (!nextCursor || logsQuery.isPlaceholderData) return; - const nextPageIndex = pagination.pageIndex + 1; - setSessionCursors((previous) => ({ ...previous, [nextPageIndex]: nextCursor })); - setPagination({ ...requested, pageIndex: nextPageIndex }); + setSessionCursors((previous) => ({ ...previous, [requested.pageIndex]: nextCursor })); + setPagination(requested); }, [usesSessionCursor, pagination, filteredLogs.next_session_cursor, logsQuery.isPlaceholderData], ); From a9bd86b3711d2ca88c8b60be3e7d7b5677cebcfb Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 11 Sep 2026 01:37:43 +0000 Subject: [PATCH 035/155] feat(ui): search, sort and role filter for the team member table Rebuild the shared member table on DataTable so admins can search members by name, email or user id, sort by name, email, role, budget and spend, and filter by role. /team/info now returns each member's user_alias so the table can show a human-readable name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 5 + .../management_endpoints/team_endpoints.py | 60 ++-- .../test_team_endpoints.py | 110 +++--- .../common_components/MemberTable.test.tsx | 215 ++++++++++++ .../common_components/MemberTable.tsx | 329 +++++++++++++----- .../src/components/networking.tsx | 1 + .../organization/organization_view.tsx | 19 +- .../shared/DataTable/DataTable.test.tsx | 22 ++ .../components/shared/DataTable/DataTable.tsx | 10 + .../src/components/team/TeamMemberTab.tsx | 18 +- 10 files changed, 595 insertions(+), 194 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2cbff128635..22591eab651 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4358,7 +4358,12 @@ class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase): agent_ids: tuple[str, ...] = () +class TeamInfoMember(Member): + user_alias: str | None = None + + class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): + members_with_roles: tuple[TeamInfoMember, ...] = () team_member_budget_table: LiteLLM_BudgetTableFull | None = None # Resources inherited from access groups (separate from direct assignments) access_group_models: list[str] | None = None diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c050368b3fe..e9e37540dd8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -61,6 +61,7 @@ from litellm.proxy._types import ( SpecialProxyStrings, TeamAccessGroupModelGrant, TeamAddMemberResponse, + TeamInfoMember, TeamInfoResponseObject, TeamInfoResponseObjectTeamTable, TeamListResponseObject, @@ -4292,37 +4293,35 @@ async def _add_team_member_budget_table( return team_info_response_object -async def _hydrate_member_emails( +async def _hydrate_member_user_details( prisma_client: PrismaClient, members: Sequence[Member], -) -> tuple[Member, ...]: - """Fill in ``user_email`` for roster entries that were stored without one. - - ``members_with_roles`` is a denormalized snapshot written at add-time, so an entry - stored with ``user_email=None`` keeps that null even once the user row has an email. - Look the missing ones up in ``LiteLLM_UserTable`` (one indexed query) and fill them - in. A stored email is never overwritten - the snapshot stays the source of truth - wherever it has a value. - """ - missing_user_ids: Final = frozenset(m.user_id for m in members if not m.user_email and m.user_id is not None) - if not missing_user_ids: - return tuple(members) - - user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await _user_db(prisma_client).find_many( - where={ # mutable-ok: Prisma query filters are dict-shaped - "user_id": { # mutable-ok: Prisma query filters are dict-shaped - "in": sorted(missing_user_ids) +) -> tuple[TeamInfoMember, ...]: + """Attach ``user_alias`` and fill in a missing ``user_email`` from ``LiteLLM_UserTable`` in one query.""" + user_ids: Final = frozenset(m.user_id for m in members if m.user_id is not None) + user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = ( + await _user_db(prisma_client).find_many( + where={ # mutable-ok: Prisma query filters are dict-shaped + "user_id": { # mutable-ok: Prisma query filters are dict-shaped + "in": sorted(user_ids) + } } - } + ) + if user_ids + else () ) - email_by_user_id: Final = MappingProxyType({u.user_id: u.user_email for u in user_rows if u.user_email}) + user_by_id: Final = MappingProxyType({u.user_id: u for u in user_rows}) - return tuple( - m.model_copy(update={"user_email": email_by_user_id[m.user_id]}) # mutable-ok: pydantic update payload - if not m.user_email and m.user_id is not None and m.user_id in email_by_user_id - else m - for m in members - ) + def hydrate(m: Member) -> TeamInfoMember: + user_row: Final = user_by_id.get(m.user_id) if m.user_id is not None else None + return TeamInfoMember( + role=m.role, + user_id=m.user_id, + user_email=m.user_email or (user_row.user_email if user_row is not None else None), + user_alias=user_row.user_alias if user_row is not None else None, + ) + + return tuple(hydrate(m) for m in members) async def _resolve_team_access_group_resources( @@ -4462,17 +4461,12 @@ async def team_info( # Resolve resources inherited from access groups resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) - # Fill in emails the add-time roster snapshot never captured - hydrated_members: Final = await _hydrate_member_emails( + hydrated_members: Final = await _hydrate_member_user_details( prisma_client=prisma_client, members=resolved_team_info.members_with_roles, ) hydrated_team_info: Final = resolved_team_info.model_copy( - update={ # mutable-ok: pydantic update payload - # list(), not the tuple: model_copy skips validation, so the field has - # to be handed the list[Member] the response model declares. - "members_with_roles": list(hydrated_members) # mutable-ok: declared list[Member] - } + update={"members_with_roles": hydrated_members} # mutable-ok: pydantic update payload ) response_object: Final = TeamInfoResponseObject( diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 2f6561046b1..2fb496d6231 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -11220,13 +11220,14 @@ async def test_team_info_returns_model_aliases(): @pytest.mark.asyncio -async def test_team_info_hydrates_member_emails_from_the_user_table(): - """/team/info must fill in emails missing from the members_with_roles snapshot. +async def test_team_info_hydrates_member_names_and_emails_from_the_user_table(): + """/team/info must attach each member's display name and fill in emails missing + from the members_with_roles snapshot. - members_with_roles is written at add-time, so a member added by user_id alone - carries user_email=None forever. Without this join the Admin UI's member table - shows "-" for a user that has an email on their user row. A stored email is left - exactly as-is. + members_with_roles is written at add-time, so it never carries user_alias and a + member added by user_id alone carries user_email=None forever. Without this join + the Admin UI's member table can only show emails. A stored email is left exactly + as-is. """ from fastapi import Request @@ -11246,13 +11247,8 @@ async def test_team_info_hydrates_member_emails_from_the_user_table(): find_many = AsyncMock( return_value=[ - LiteLLM_UserTable( - user_id="no-email-on-roster", - user_email="real@example.com", - max_budget=None, - spend=0.0, - models=[], - ) + _user_row("no-email-on-roster", "real@example.com", "Real Person"), + _user_row("already-stored", "current@example.com", "Stored Person"), ] ) @@ -11270,12 +11266,12 @@ async def test_team_info_hydrates_member_emails_from_the_user_table(): ) members = response["team_info"].members_with_roles - assert [(m.user_id, m.user_email) for m in members] == [ - ("no-email-on-roster", "real@example.com"), - ("already-stored", "stored@example.com"), + assert [(m.user_id, m.user_email, m.user_alias) for m in members] == [ + ("no-email-on-roster", "real@example.com", "Real Person"), + ("already-stored", "stored@example.com", "Stored Person"), ] - # only the member actually missing an email is looked up - assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["no-email-on-roster"]}} + find_many.assert_awaited_once() + assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["already-stored", "no-email-on-roster"]}} @pytest.mark.asyncio @@ -12472,89 +12468,93 @@ async def test_resolve_existing_member_user_ids_skips_the_query_when_no_user_ids repo.return_value.table.find_many.assert_not_awaited() -def _user_row(user_id: str, user_email: str | None) -> LiteLLM_UserTable: +def _user_row(user_id: str, user_email: str | None, user_alias: str | None = None) -> LiteLLM_UserTable: return LiteLLM_UserTable( - user_id=user_id, user_email=user_email, max_budget=None, spend=0.0, models=[] + user_id=user_id, user_email=user_email, user_alias=user_alias, max_budget=None, spend=0.0, models=[] ) @pytest.mark.asyncio -async def test_hydrate_member_emails_fills_in_emails_the_roster_snapshot_never_captured(): - """A member added by user_id alone has user_email=None on the stored roster entry. - - /team/info has to fill it in from the user row, or the UI renders "-" for a user - that plainly has an email. +async def test_hydrate_member_user_details_attaches_alias_and_fills_in_missing_email(): + """The stored roster never carries a display name, and a member added by user_id + alone has user_email=None. /team/info has to fill both in from the user row so the + UI can show and search by a human-readable name instead of only an email. """ - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details - find_many = AsyncMock(return_value=[_user_row("by-id", "found@example.com")]) + find_many = AsyncMock(return_value=[_user_row("by-id", "found@example.com", "Found Person")]) with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: repo.return_value.table.find_many = find_many - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), members=[Member(user_id="by-id", role="admin")], ) - assert [(m.user_id, m.user_email, m.role) for m in hydrated] == [("by-id", "found@example.com", "admin")] + assert [(m.user_id, m.user_email, m.user_alias, m.role) for m in hydrated] == [ + ("by-id", "found@example.com", "Found Person", "admin") + ] find_many.assert_awaited_once() assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["by-id"]}} @pytest.mark.asyncio -async def test_hydrate_member_emails_never_overwrites_a_stored_email(): - """The snapshot wins wherever it has a value - hydration only fills blanks. - - Overwriting would be a real behavior change to /team/info; filling a null is not. - """ - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails - - find_many = AsyncMock(return_value=[_user_row("has-email", "current@example.com")]) +async def test_hydrate_member_user_details_never_overwrites_a_stored_email(): + """The snapshot wins wherever it has a value - hydration only fills blanks.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: - repo.return_value.table.find_many = find_many + repo.return_value.table.find_many = AsyncMock( + return_value=[_user_row("has-email", "current@example.com", "Current Name")] + ) - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), members=[Member(user_id="has-email", user_email="stored@example.com", role="user")], ) - assert hydrated[0].user_email == "stored@example.com" - # nothing was missing, so no round-trip either - find_many.assert_not_awaited() + assert (hydrated[0].user_email, hydrated[0].user_alias) == ("stored@example.com", "Current Name") @pytest.mark.asyncio -async def test_hydrate_member_emails_leaves_members_alone_when_the_user_row_has_no_email(): - """A user row with no email leaves the member as-is rather than inventing one.""" - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails +async def test_hydrate_member_user_details_leaves_blanks_when_the_user_row_is_bare_or_missing(): + """A user row with no email or alias, or no user row at all, must not invent values.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: - repo.return_value.table.find_many = AsyncMock(return_value=[_user_row("no-email", None)]) + repo.return_value.table.find_many = AsyncMock(return_value=[_user_row("bare", None)]) - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), - members=[Member(user_id="no-email", role="user"), Member(user_email="e@example.com", role="user")], + members=[ + Member(user_id="bare", role="user"), + Member(user_id="deleted", user_email="gone@example.com", role="user"), + Member(user_email="e@example.com", role="user"), + ], ) - assert [m.user_email for m in hydrated] == [None, "e@example.com"] + assert [(m.user_id, m.user_email, m.user_alias) for m in hydrated] == [ + ("bare", None, None), + ("deleted", "gone@example.com", None), + (None, "e@example.com", None), + ] @pytest.mark.asyncio -async def test_hydrate_member_emails_skips_the_query_when_every_member_has_one(): - """No blanks means /team/info pays for no extra query.""" - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails +async def test_hydrate_member_user_details_skips_the_query_when_no_member_has_a_user_id(): + """Email-only roster entries give nothing to look up, so /team/info pays for no query.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: repo.return_value.table.find_many = AsyncMock() - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), - members=[Member(user_id="a", user_email="a@example.com", role="user")], + members=[Member(user_email="a@example.com", role="user")], ) - assert hydrated[0].user_email == "a@example.com" + assert [(m.user_email, m.user_alias) for m in hydrated] == [("a@example.com", None)] repo.return_value.table.find_many.assert_not_awaited() diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx new file mode 100644 index 00000000000..a16f11eab60 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx @@ -0,0 +1,215 @@ +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; + +import type { Member } from "@/components/networking"; + +import { renderWithProviders } from "../../../tests/test-utils"; +import MemberTable, { MemberTableColumn, memberRoleOptions } from "./MemberTable"; + +const MEMBERS: Member[] = [ + { user_id: "u-zed", user_email: "zed@example.com", user_alias: "Zed Ortiz", role: "user" }, + { user_id: "u-nameless", user_email: "mystery@example.com", user_alias: null, role: "user" }, + { user_id: "u-amy", user_email: "amy@example.com", user_alias: "amy chen", role: "admin" }, + { user_id: "u-bob", user_email: null, user_alias: "Bob Lee", role: "user" }, +]; + +const BUDGETS: Record = { "u-zed": 50, "u-nameless": null, "u-amy": 1000, "u-bob": 5 }; + +const budgetColumn: MemberTableColumn = { + title: "Budget", + key: "budget", + sortValue: (member) => BUDGETS[member.user_id ?? ""] ?? null, + render: (member) => {BUDGETS[member.user_id ?? ""] ?? "Unlimited"}, +}; + +const renderTable = (overrides: Partial> = {}) => { + const props = { + members: MEMBERS, + canEdit: true, + onEdit: vi.fn(), + onDelete: vi.fn(), + extraColumns: [budgetColumn], + ...overrides, + }; + renderWithProviders(); + return props; +}; + +const rowIds = (): (string | null)[] => + Array.from(document.querySelectorAll("tbody tr[data-row-id]")).map((row) => row.getAttribute("data-row-id")); + +const search = (value: string) => fireEvent.change(screen.getByTestId("datatable-search"), { target: { value } }); + +describe("MemberTable display", () => { + it("shows each member's name, falling back to a dash when there is none", () => { + renderTable(); + + expect(within(screen.getByRole("row", { name: /zed@example\.com/ })).getByText("Zed Ortiz")).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: /^name/i })).toBeInTheDocument(); + const cells = within(screen.getByRole("row", { name: /mystery@example\.com/ })).getAllByRole("cell"); + expect(cells[0]).toHaveTextContent("-"); + }); + + it("orders members by name with nameless members last by default", () => { + renderTable(); + + expect(rowIds()).toEqual(["u-amy", "u-bob", "u-zed", "u-nameless"]); + }); + + it("reports the full member count", () => { + renderTable(); + + expect(screen.getByText("4 Members")).toBeInTheDocument(); + }); +}); + +describe("MemberTable search", () => { + it("matches on name case-insensitively", async () => { + renderTable(); + + search("ZED"); + + await waitFor(() => expect(rowIds()).toEqual(["u-zed"])); + }); + + it("matches on email", async () => { + renderTable(); + + search("mystery@"); + + await waitFor(() => expect(rowIds()).toEqual(["u-nameless"])); + }); + + it("matches on user id", async () => { + renderTable(); + + search("u-bob"); + + await waitFor(() => expect(rowIds()).toEqual(["u-bob"])); + }); + + it("still searches names when the first member has no name", async () => { + renderTable({ members: [MEMBERS[1], MEMBERS[0]] }); + + search("ortiz"); + + await waitFor(() => expect(rowIds()).toEqual(["u-zed"])); + }); + + it("does not match on role", async () => { + renderTable(); + + search("admin"); + + await waitFor(() => expect(rowIds()).toEqual([])); + expect(screen.getByText("No members match your search or filters")).toBeInTheDocument(); + }); +}); + +describe("MemberTable sorting", () => { + it("sorts by email with missing emails last in both directions", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-user_email")); + expect(rowIds()).toEqual(["u-amy", "u-nameless", "u-zed", "u-bob"]); + + await user.click(screen.getByTestId("sort-header-user_email")); + expect(rowIds()).toEqual(["u-zed", "u-nameless", "u-amy", "u-bob"]); + }); + + it("sorts by role", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-role")); + expect(rowIds()[0]).toBe("u-amy"); + + await user.click(screen.getByTestId("sort-header-role")); + expect(rowIds()[3]).toBe("u-amy"); + }); + + it("sorts an extra column numerically by its sort value with blanks last", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-budget")); + expect(rowIds()).toEqual(["u-bob", "u-zed", "u-amy", "u-nameless"]); + + await user.click(screen.getByTestId("sort-header-budget")); + expect(rowIds()).toEqual(["u-amy", "u-zed", "u-bob", "u-nameless"]); + }); + + it("flips name order on the second click", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-user_alias")); + expect(rowIds()).toEqual(["u-zed", "u-bob", "u-amy", "u-nameless"]); + }); + + it("leaves extra columns without a sort value unsortable", () => { + renderTable({ + extraColumns: [{ title: "Rate Limits", key: "rate_limits", render: () => No Limits }], + }); + + expect(screen.getByRole("columnheader", { name: "Rate Limits" })).toBeInTheDocument(); + expect(screen.queryByTestId("sort-header-rate_limits")).not.toBeInTheDocument(); + }); +}); + +describe("MemberTable role filter", () => { + it("shows only members with the chosen role and clears on reset", async () => { + const user = userEvent.setup(); + renderTable({ roleColumnTitle: "Team Role" }); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + await user.click(screen.getByTestId("filter-role")); + await user.click(await screen.findByRole("option", { name: "admin" })); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(rowIds()).toEqual(["u-amy"])); + expect(screen.getByTestId("filter-chip-role")).toHaveTextContent("Team Role"); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + await user.click(screen.getByTestId("filter-drawer-reset")); + + await waitFor(() => expect(rowIds()).toHaveLength(4)); + }); + + it("offers the roles present in the roster", () => { + expect(memberRoleOptions(MEMBERS)).toEqual(["admin", "user"]); + expect(memberRoleOptions([{ user_id: "x", role: "" }])).toEqual([]); + }); +}); + +describe("MemberTable actions", () => { + it("passes the clicked member to onEdit and onDelete", async () => { + const user = userEvent.setup(); + const { onEdit, onDelete } = renderTable(); + const row = screen.getByRole("row", { name: /amy@example\.com/ }); + + await user.click(within(row).getByTestId("edit-member")); + await user.click(within(row).getByTestId("delete-member")); + + expect(onEdit).toHaveBeenCalledWith(MEMBERS[2]); + expect(onDelete).toHaveBeenCalledWith(MEMBERS[2]); + }); + + it("hides delete for members the caller excludes", () => { + renderTable({ showDeleteForMember: (member) => member.role !== "admin" }); + + expect(screen.getAllByTestId("delete-member")).toHaveLength(3); + expect( + within(screen.getByRole("row", { name: /amy@example\.com/ })).queryByTestId("delete-member"), + ).not.toBeInTheDocument(); + }); + + it("shows the empty text when there are no members at all", () => { + renderTable({ members: [], emptyText: "No members found" }); + + expect(screen.getByText("No members found")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx index 8a0b7671078..e54994c1ab7 100644 --- a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx @@ -1,17 +1,29 @@ -import { SimpleTooltip } from "@/components/ui/tooltip"; +import type { ColumnDef, ColumnFiltersState } from "@tanstack/react-table"; +import { Crown, Info, User, UserPlus } from "lucide-react"; +import React, { useState } from "react"; + import { Member } from "@/components/networking"; +import { + DataTable, + DataTableFilterDrawer, + DataTableFilterField, + DataTableSortHeader, + DataTableToolbar, +} from "@/components/shared/DataTable"; import { StatusBadge } from "@/components/shared/table_cells"; import { Button } from "@/components/ui/button"; -import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; -import { Crown, Info, User, UserPlus } from "lucide-react"; -import React from "react"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { SimpleTooltip } from "@/components/ui/tooltip"; + import TableIconActionButton from "./IconActionButton/TableIconActionButtons/TableIconActionButton"; +export type MemberTableSortValue = string | number | null | undefined; + export interface MemberTableColumn { title: React.ReactNode; - key: React.Key; - dataIndex?: keyof Member; - render?: (value: Member[keyof Member], member: Member, index: number) => React.ReactNode; + key: string; + render: (member: Member) => React.ReactNode; + sortValue?: (member: Member) => MemberTableSortValue; } export interface MemberTableProps { @@ -27,12 +39,152 @@ export interface MemberTableProps { emptyText?: string; } -const extraColumnCell = (column: MemberTableColumn, member: Member, index: number): React.ReactNode => { - const value = column.dataIndex ? member[column.dataIndex] : undefined; - return column.render ? column.render(value, member, index) : value; +const ALL_ROLES = "all"; + +export const memberRowId = (member: Member): string => member.user_id ?? member.user_email ?? JSON.stringify(member); + +export const memberRoleOptions = (members: readonly Member[]): string[] => + Array.from(new Set(members.map((member) => member.role).filter((role) => role !== ""))).sort(); + +const isAdminRole = (role: string): boolean => { + const normalized = role.toLowerCase(); + return normalized === "admin" || normalized === "org_admin"; }; -const STICKY_ACTIONS_CLASS = "sticky right-0 w-[120px] bg-background"; +function RoleHeaderTitle({ title, tooltip }: { title: string; tooltip?: string }) { + if (tooltip === undefined) return <>{title}; + return ( + + {title} + + + + + ); +} + +const ACTIONS_COLUMN_WIDTH = 120; + +interface MemberColumnDeps { + canEdit: boolean; + onEdit: (member: Member) => void; + onDelete: (member: Member) => void; + roleColumnTitle: string; + roleTooltip?: string; + extraColumns: MemberTableColumn[]; + showDeleteForMember?: (member: Member) => boolean; +} + +const extraColumnDef = (column: MemberTableColumn): ColumnDef => { + const { sortValue } = column; + if (sortValue === undefined) { + return { + id: column.key, + header: () => {column.title}, + enableSorting: false, + enableGlobalFilter: false, + cell: ({ row }) => column.render(row.original), + }; + } + return { + id: column.key, + accessorFn: (member) => sortValue(member) ?? undefined, + header: ({ column: tableColumn }) => , + sortDescFirst: false, + sortUndefined: "last", + enableGlobalFilter: false, + cell: ({ row }) => column.render(row.original), + }; +}; + +const buildColumns = ({ + canEdit, + onEdit, + onDelete, + roleColumnTitle, + roleTooltip, + extraColumns, + showDeleteForMember, +}: MemberColumnDeps): ColumnDef[] => [ + { + id: "user_alias", + accessorFn: (member) => member.user_alias || undefined, + header: ({ column }) => , + sortingFn: "text", + sortUndefined: "last", + enableGlobalFilter: true, + meta: { title: "Name" }, + cell: ({ row }) => row.original.user_alias || -, + }, + { + id: "user_email", + accessorFn: (member) => member.user_email || undefined, + header: ({ column }) => , + sortingFn: "text", + sortUndefined: "last", + enableGlobalFilter: true, + meta: { title: "User Email" }, + cell: ({ row }) => row.original.user_email || "-", + }, + { + id: "user_id", + accessorFn: (member) => member.user_id ?? undefined, + header: "User ID", + enableSorting: false, + enableGlobalFilter: true, + cell: ({ row }) => + row.original.user_id === "default_user_id" ? ( + + ) : ( + row.original.user_id || "-" + ), + }, + { + id: "role", + accessorFn: (member) => member.role, + header: ({ column }) => ( + } /> + ), + sortingFn: "text", + filterFn: "equalsString", + enableGlobalFilter: false, + meta: { title: roleColumnTitle }, + cell: ({ row }) => ( + + {isAdminRole(row.original.role) ? : } + {row.original.role || "-"} + + ), + }, + ...extraColumns.map(extraColumnDef), + { + id: "actions", + header: "Actions", + size: ACTIONS_COLUMN_WIDTH, + enableSorting: false, + enableGlobalFilter: false, + meta: { pinned: "right" }, + cell: ({ row }) => + canEdit ? ( + + onEdit(row.original)} + /> + {(!showDeleteForMember || showDeleteForMember(row.original)) && ( + onDelete(row.original)} + /> + )} + + ) : null, + }, +]; export default function MemberTable({ members, @@ -46,90 +198,89 @@ export default function MemberTable({ showDeleteForMember, emptyText, }: MemberTableProps) { + const [globalFilter, setGlobalFilter] = useState(""); + const [columnFilters, setColumnFilters] = useState([]); + const [filtersOpen, setFiltersOpen] = useState(false); + + const columnDeps: MemberColumnDeps = { + canEdit, + onEdit, + onDelete, + roleColumnTitle, + roleTooltip, + extraColumns, + showDeleteForMember, + }; + const columns = buildColumns(columnDeps); + const roleFilterItems = [ + { value: ALL_ROLES, label: "All Roles" }, + ...memberRoleOptions(members).map((role) => ({ value: role, label: role })), + ]; + + const isNarrowed = globalFilter !== "" || columnFilters.length > 0; + return (
{members.length} Member{members.length !== 1 ? "s" : ""} - - - - User Email - User ID - - {roleTooltip ? ( - - {roleColumnTitle} - - - - - ) : ( - roleColumnTitle + + {isNarrowed ? "No members match your search or filters" : emptyText ?? "No data"} + + } + toolbar={(table) => ( + <> + setFiltersOpen(true)} + showViewOptions={false} + /> + + {({ get, set }) => ( + + + )} - - {extraColumns.map((column) => ( - {column.title} - ))} - Actions - - - - {members.length === 0 ? ( - - - {emptyText ?? "No data"} - - - ) : ( - members.map((member, memberIndex) => ( - - {member.user_email || "-"} - - {member.user_id === "default_user_id" ? ( - - ) : ( - member.user_id || "-" - )} - - - - {member.role?.toLowerCase() === "admin" || member.role?.toLowerCase() === "org_admin" ? ( - - ) : ( - - )} - {member.role || "-"} - - - {extraColumns.map((column) => ( - {extraColumnCell(column, member, memberIndex)} - ))} - - {canEdit ? ( - - onEdit(member)} - /> - {(!showDeleteForMember || showDeleteForMember(member)) && ( - onDelete(member)} - /> - )} - - ) : null} - - - )) - )} - -
+ + + )} + /> {onAddMember && canEdit && ( @@ -71,6 +75,14 @@ export function IdCell({ ); + const idElement = linked ? ( + + {value} + + ) : ( + unlinkedElement + ); + const withTooltip = ; if (!copyable) { @@ -94,3 +106,21 @@ export function IdCell({ ); } + +interface IdLinkProps extends React.ComponentPropsWithoutRef<"a"> { + href: string; + dataTestId?: string; +} + +const IdLink = React.forwardRef(function IdLink( + { href, dataTestId, children, ...props }, + ref, +) { + const handleClick = useEntityLinkClick(href); + + return ( + + {children} + + ); +}); From 06964e5603ecce13b0e4aabcbcab5d1382c4c6e5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 11 Sep 2026 17:48:57 -0700 Subject: [PATCH 137/155] feat(ui): link the User ID, Created By and Deleted By cells on Deleted Keys (#40750) * feat(ui): link the User ID, Created By and Deleted By cells on Deleted Keys All three columns rendered as plain text, so auditing a deleted key meant copying an id into the Users page search box. Route them through IdentityCell with userDetailHref, which keeps the proxy admin placeholder unlinked. User Email and Team Alias stay as they are: the deleted key table has no column for either, so the API never populates them. Claude-Session: https://claude.ai/code/session_01NfwfQhamRNnSqgXMUjf3h4 * test(ui): mount a router mock for the Deleted Keys page test The page test renders the table, and the newly linked cells call useRouter, which throws without an App Router mounted. Matches how the other 35 test files in the suite stub next/navigation. Claude-Session: https://claude.ai/code/session_01NfwfQhamRNnSqgXMUjf3h4 --- .../DeletedKeysPage/DeletedKeysPage.test.tsx | 2 ++ .../DeletedKeysTable.test.tsx | 18 +++++++++++++++++ .../DeletedKeysTableColumns.tsx | 20 +++++++++++++++---- 3 files changed, 36 insertions(+), 4 deletions(-) diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx index ee2f42ad85b..31cd407a5e6 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx @@ -5,6 +5,8 @@ import { renderWithProviders } from "../../../tests/test-utils"; import DeletedKeysPage from "./DeletedKeysPage"; import { useDeletedKeys, DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ useDeletedKeys: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx index 7e30ef2c135..dd8007be264 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx @@ -5,6 +5,8 @@ import { renderWithProviders } from "../../../../tests/test-utils"; import { DeletedKeysTable } from "./DeletedKeysTable"; import { DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + const makeDeletedKey = (overrides: Partial = {}): DeletedKeyResponse => ({ token: "sk-1234567890abcdef", @@ -86,3 +88,19 @@ it("should show the empty state when there are no deleted keys", () => { expect(screen.getByText("No deleted keys found")).toBeInTheDocument(); }); + +it("links the owner, creator and deleter cells to their user detail pages", () => { + renderWithProviders(); + + expect(screen.getByRole("link", { name: "user-1" })).toHaveAttribute("href", "/ui/users?user=user-1"); + expect(screen.getByRole("link", { name: "creator-1" })).toHaveAttribute("href", "/ui/users?user=creator-1"); + expect(screen.getByRole("link", { name: "deleter-1" })).toHaveAttribute("href", "/ui/users?user=deleter-1"); +}); + +it("leaves the default_user_id placeholder unlinked", () => { + const placeholderKey = makeDeletedKey({ user_id: "default_user_id", created_by: "default_user_id" }); + renderWithProviders(); + + expect(screen.getAllByText("default_user_id")).toHaveLength(2); + expect(screen.queryByRole("link", { name: "default_user_id" })).not.toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx index aa7d6380cc3..32185f867b4 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx @@ -3,8 +3,9 @@ import { ColumnDef } from "@tanstack/react-table"; import { DataTableSortHeader } from "@/components/shared/DataTable"; -import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells"; +import { DateCell, IdCell, IdentityCell, MoneyCell } from "@/components/shared/table_cells"; import { DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { userDetailHref } from "@/utils/entityLinks"; function TruncatedTextCell({ value }: { value: string | null | undefined }) { if (!value) { @@ -17,6 +18,17 @@ function TruncatedTextCell({ value }: { value: string | null | undefined }) { ); } +function UserLinkCell({ userId }: { userId: string | null | undefined }) { + if (!userId) { + return -; + } + return ( + + + + ); +} + export const getDeletedKeysTableColumns = (): ColumnDef[] => [ { id: "token", @@ -89,7 +101,7 @@ export const getDeletedKeysTableColumns = (): ColumnDef[] => header: "User ID", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => , }, { id: "created_at", @@ -107,7 +119,7 @@ export const getDeletedKeysTableColumns = (): ColumnDef[] => header: "Created By", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => , }, { id: "deleted_at", @@ -125,6 +137,6 @@ export const getDeletedKeysTableColumns = (): ColumnDef[] => header: "Deleted By", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => , }, ]; From c53f72c764f9c1ae9015df667e350b130ebb8a9f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 11 Sep 2026 17:49:01 -0700 Subject: [PATCH 138/155] feat(ui): link the Created By cell on the Prompts page (#40753) The column was plain muted text, so finding out who owns a prompt meant copying the id into the Users page search box. Route it through IdentityCell with userDetailHref, which keeps the proxy admin placeholder unlinked. Claude-Session: https://claude.ai/code/session_01NfwfQhamRNnSqgXMUjf3h4 --- .../prompts/_components/PromptTable.test.tsx | 11 +++++++++++ .../prompts/_components/PromptTableColumns.tsx | 12 ++++++++++-- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx index edbb897fb2f..009abf93d1f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx @@ -10,6 +10,8 @@ vi.mock("@/components/networking", () => ({ modelHubCall: vi.fn().mockResolvedValue({ data: [] }), })); +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + const mockPrompts: PromptSpec[] = [ { prompt_id: "prompt-newer", @@ -53,6 +55,15 @@ describe("PromptTable", () => { } }); + it("links the Created By cell to the creator's detail page, leaving the placeholder unlinked", () => { + const prompts = [mockPrompts[0], { ...mockPrompts[1], created_by: "default_user_id" }]; + render(); + + expect(screen.getByRole("link", { name: "user-1" })).toHaveAttribute("href", "/ui/users?user=user-1"); + expect(screen.getByText("default_user_id")).toBeInTheDocument(); + expect(screen.queryByRole("link", { name: "default_user_id" })).not.toBeInTheDocument(); + }); + it("should display the empty state when data is empty", () => { render(); expect(screen.getByText("No prompts yet")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx index f927a6d1486..db0392bbd6e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx @@ -17,6 +17,7 @@ import { } from "@/components/ui/dropdown-menu"; import { cn } from "@/lib/cva.config"; import { copyToClipboard } from "@/utils/dataUtils"; +import { userDetailHref } from "@/utils/entityLinks"; import { extractModel, getProviderFromModelHub, ModelGroupInfo } from "./prompt_utils"; @@ -191,9 +192,16 @@ export const getPromptTableColumns = ({ enableSorting: false, cell: ({ row }) => { const createdBy = row.original.created_by; + if (!createdBy) { + return -; + } return ( - - {createdBy || "-"} + + ); }, From b27d2cce77cd1d0ba9662ccb5b2fe9640471f2c6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 11 Sep 2026 17:49:07 -0700 Subject: [PATCH 139/155] feat(ui): link the Organization and Deleted By cells on Deleted Teams (#40751) * feat(ui): link the Organization and Deleted By cells on Deleted Teams Both columns rendered as plain text, so tracing a deleted team back to its org or to whoever removed it meant copying an id into another page's search box. Route them through IdentityCell with orgDetailHref and userDetailHref. Team ID stays unlinked because the team itself is gone. Claude-Session: https://claude.ai/code/session_01NfwfQhamRNnSqgXMUjf3h4 * test(ui): mount a router mock for the Deleted Teams page test The page test renders the table, and the newly linked cells call useRouter, which throws without an App Router mounted. Matches how the other 35 test files in the suite stub next/navigation. Claude-Session: https://claude.ai/code/session_01NfwfQhamRNnSqgXMUjf3h4 --- .../DeletedTeamsPage.test.tsx | 2 ++ .../DeletedTeamsTable.test.tsx | 20 +++++++++++++ .../DeletedTeamsTableColumns.tsx | 30 ++++++++++++------- 3 files changed, 41 insertions(+), 11 deletions(-) diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx index 952d8764463..d2d25f5e47e 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx @@ -5,6 +5,8 @@ import { renderWithProviders } from "../../../tests/test-utils"; import DeletedTeamsPage from "./DeletedTeamsPage"; import { useDeletedTeams, DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ useDeletedTeams: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx index e166f6b0d1b..837fb583f30 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx @@ -4,6 +4,8 @@ import { renderWithProviders } from "../../../../tests/test-utils"; import { DeletedTeamsTable } from "./DeletedTeamsTable"; import { DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + const makeDeletedTeam = (overrides: Partial = {}): DeletedTeam => ({ team_id: "team-1", team_alias: "Test Team", @@ -81,3 +83,21 @@ it("renders the shared pagination footer with the server row count", () => { expect(screen.getByTestId("pagination-prev")).toBeEnabled(); expect(screen.getByTestId("pagination-next")).toBeDisabled(); }); + +it("links the organization and deleted by cells, leaving the deleted team id unlinked", () => { + renderWithProviders( + , + ); + + expect(screen.getByRole("link", { name: "org-1" })).toHaveAttribute("href", "/ui/organizations?org=org-1"); + expect(screen.getByRole("link", { name: "user-1" })).toHaveAttribute("href", "/ui/users?user=user-1"); + expect(screen.queryByRole("link", { name: "team-1" })).not.toBeInTheDocument(); +}); + +it("leaves the default_user_id placeholder unlinked in the deleted by cell", () => { + const team = makeDeletedTeam({ deleted_by: "default_user_id", organization_id: null }); + renderWithProviders(); + + expect(screen.getByText("default_user_id")).toBeInTheDocument(); + expect(screen.queryByRole("link", { name: "default_user_id" })).not.toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx index e36077fd2c3..172f0417027 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx @@ -3,8 +3,20 @@ import { ColumnDef } from "@tanstack/react-table"; import { DataTableSortHeader } from "@/components/shared/DataTable"; -import { DateCell, IdCell, ModelsCell, MoneyCell } from "@/components/shared/table_cells"; +import { DateCell, IdCell, IdentityCell, ModelsCell, MoneyCell } from "@/components/shared/table_cells"; import { DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { orgDetailHref, userDetailHref } from "@/utils/entityLinks"; + +function EntityCell({ value, href }: { value: string | null | undefined; href: string | undefined }) { + if (!value) { + return -; + } + return ( + + + + ); +} export const getDeletedTeamsTableColumns = (): ColumnDef[] => [ { @@ -78,7 +90,10 @@ export const getDeletedTeamsTableColumns = (): ColumnDef[] => [ header: "Organization", size: 150, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => { + const orgId = row.original.organization_id; + return ; + }, }, { id: "deleted_at", @@ -97,15 +112,8 @@ export const getDeletedTeamsTableColumns = (): ColumnDef[] => [ size: 120, enableSorting: false, cell: ({ row }) => { - const value = row.original.deleted_by; - if (!value) { - return -; - } - return ( - - {value} - - ); + const deletedBy = row.original.deleted_by; + return ; }, }, ]; From 115535c3add1f79bf8b5056663c1456a0be5e1cf Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 12 Sep 2026 00:49:26 +0000 Subject: [PATCH 140/155] test(model_prices): cover Fireworks DeepSeek V4.1 costs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_fireworks_serverless_model_costs.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index b8d36b996df..e3404e477b7 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -14,6 +14,8 @@ import os import pytest +from litellm import completion_cost +from litellm.types.utils import Choices, Message, ModelResponse, Usage from litellm.utils import get_model_info @@ -56,6 +58,20 @@ def test_bare_fireworks_ids_resolve_through_prefixed_entries(): assert info["max_output_tokens"] == expected["max_output_tokens"] +def test_deepseek_v4p1_flash_twin_costs(): + for model in ( + "fireworks_ai/deepseek-v4p1-flash", + "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash", + ): + response = ModelResponse( + model=model, + choices=[Choices(index=0, message=Message(role="assistant", content="ok"))], + usage=Usage(prompt_tokens=1000, completion_tokens=1000, total_tokens=2000), + ) + cost = completion_cost(completion_response=response, model=model) + assert cost == pytest.approx(8.8e-04) + + TWIN_PINNED_PRICES = { "deepseek-v4-flash-0731": { "input_cost_per_token": 2.2e-07, From 05d2c316f5d389979ed60ba26cd0c2ef71ec3692 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:53:51 -0700 Subject: [PATCH 141/155] refactor(mcp): reuse in-memory discovery storage --- litellm/caching/in_memory_cache.py | 13 ++-- .../mcp_server/mcp_server_manager.py | 67 +++++++++---------- .../caching/test_in_memory_cache.py | 24 +++++++ .../mcp_server/test_mcp_server_manager.py | 51 +++++++++++++- 4 files changed, 112 insertions(+), 43 deletions(-) diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 4058a3d72dd..56c9147e066 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -13,6 +13,7 @@ import json import sys import threading import time +from collections.abc import Callable from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: @@ -34,6 +35,7 @@ class InMemoryCache(BaseCache): default_ttl: int | None = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute max_size_per_item: int | None = 1024, # 1MB = 1024KB + clock: Callable[[], float] | None = None, ): """ max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default @@ -49,6 +51,7 @@ class InMemoryCache(BaseCache): self.ttl_dict: dict = {} self.expiration_heap: list[tuple[float, str]] = [] self._increment_lock = threading.Lock() + self._clock = clock if clock is not None else lambda: time.time() def check_value_size(self, value: Any): """ @@ -91,7 +94,7 @@ class InMemoryCache(BaseCache): """ Check if a specific key is expired """ - return key in self.ttl_dict and time.time() > self.ttl_dict[key] + return key in self.ttl_dict and self._clock() > self.ttl_dict[key] def _remove_key(self, key: str) -> None: """ @@ -113,7 +116,7 @@ class InMemoryCache(BaseCache): - 3. the size of in-memory cache is bounded """ - current_time: Final = time.time() + current_time: Final = self._clock() # Step 1: Remove expired or outdated items while self.expiration_heap: @@ -147,7 +150,7 @@ class InMemoryCache(BaseCache): Check if ttl is set for a key """ ttl_time: Final = self.ttl_dict.get(key) - if ttl_time is None or float(ttl_time) < time.time(): # if ttl is not set, allow override + if ttl_time is None or float(ttl_time) < self._clock(): # if ttl is not set, allow override return True else: return False @@ -167,10 +170,10 @@ class InMemoryCache(BaseCache): self.cache_dict[key] = value if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl if "ttl" in kwargs and kwargs["ttl"] is not None: - self.ttl_dict[key] = time.time() + float(kwargs["ttl"]) + self.ttl_dict[key] = self._clock() + float(kwargs["ttl"]) heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) else: - self.ttl_dict[key] = time.time() + self.default_ttl + self.ttl_dict[key] = self._clock() + self.default_ttl heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) async def async_set_cache(self, key, value, **kwargs): diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ca81e8191f7..35dfeb9ee74 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -51,6 +51,7 @@ from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( MCP_CLIENT_TIMEOUT, MCP_HEALTH_CHECK_TIMEOUT, @@ -195,7 +196,6 @@ if TYPE_CHECKING: from mcp.shared.context import RequestContext from mcp.types import CreateMessageRequestParams - from litellm.caching.caching import InMemoryCache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.mcp_server.mcp_toolset import MCPToolset @@ -1686,21 +1686,29 @@ _DISCOVERY_CACHE_LIMIT: Final = 1024 @dataclass(frozen=True, slots=True) class _DiscoveryEntry(Generic[_DiscoveryItem]): - expires_at: float items: tuple[_DiscoveryItem, ...] class _DiscoveryCache(Generic[_DiscoveryItem]): def __init__(self, ttl: float, clock: Callable[[], float]) -> None: self._ttl = ttl - self._clock = clock - self._entries: Mapping[_DiscoveryKey, _DiscoveryEntry[_DiscoveryItem]] = MappingProxyType({}) - self._pending: Mapping[_DiscoveryKey, asyncio.Task[list[_DiscoveryItem]]] = MappingProxyType({}) - self._waiters: Mapping[asyncio.Task[list[_DiscoveryItem]], int] = MappingProxyType({}) + self._entries = InMemoryCache(max_size_in_memory=_DISCOVERY_CACHE_LIMIT, clock=clock) + self._pending: dict[ + _DiscoveryKey, asyncio.Task[list[_DiscoveryItem]] + ] = {} # mutable-ok: constant-time fetch registration + self._waiters: dict[asyncio.Task[list[_DiscoveryItem]], int] = {} # mutable-ok: constant-time waiter accounting def invalidate(self, server_id: str) -> None: - self._entries = MappingProxyType({key: entry for key, entry in self._entries.items() if key[0] != server_id}) - self._pending = MappingProxyType({key: task for key, task in self._pending.items() if key[0] != server_id}) + prefix: Final = f"[{json.dumps(server_id)}," + keys: Final = cast( # cast-ok: private cache contains only JSON string keys + "tuple[str, ...]", tuple(self._entries.cache_dict) + ) + for entry_key in keys: + if entry_key.startswith(prefix): + self._entries.delete_cache(entry_key) + for key in tuple(self._pending): + if key[0] == server_id: + self._pending.pop(key) @staticmethod def _observe_completion(task: asyncio.Task[list[_DiscoveryItem]]) -> None: @@ -1712,8 +1720,10 @@ class _DiscoveryCache(Generic[_DiscoveryItem]): ) -> tuple[_DiscoveryItem, ...]: if self._ttl <= 0: return tuple(await fetch()) - entry: Final = self._entries.get(key) - if entry is not None and entry.expires_at > self._clock(): + entry: Final = cast( # cast-ok: private cache contains only entries for this item type + "_DiscoveryEntry[_DiscoveryItem] | None", self._entries.get_cache(json.dumps(key)) + ) + if entry is not None: return tuple(item.model_copy(deep=True) for item in entry.items) pending: Final = self._pending.get(key) if pending is not None: @@ -1721,28 +1731,24 @@ class _DiscoveryCache(Generic[_DiscoveryItem]): if len(self._pending) >= _DISCOVERY_CACHE_LIMIT: return tuple(await fetch()) task: Final = asyncio.create_task(self._fetch(key, fetch)) - self._pending = MappingProxyType({**self._pending, key: task}) + self._pending[key] = task task.add_done_callback(self._observe_completion) return await self._await_fetch(key, task) async def _await_fetch( self, key: _DiscoveryKey, task: asyncio.Task[list[_DiscoveryItem]] ) -> tuple[_DiscoveryItem, ...]: - self._waiters = MappingProxyType({**self._waiters, task: self._waiters.get(task, 0) + 1}) + self._waiters[task] = self._waiters.get(task, 0) + 1 try: return tuple(item.model_copy(deep=True) for item in await asyncio.shield(task)) finally: remaining: Final = self._waiters[task] - 1 if remaining: - self._waiters = MappingProxyType({**self._waiters, task: remaining}) + self._waiters[task] = remaining else: - self._waiters = MappingProxyType( - {pending: count for pending, count in self._waiters.items() if pending is not task} - ) + self._waiters.pop(task) if self._pending.get(key) is task: - self._pending = MappingProxyType( - {entry_key: pending for entry_key, pending in self._pending.items() if entry_key != key} - ) + self._pending.pop(key) if not task.done(): task.cancel() @@ -1752,28 +1758,15 @@ class _DiscoveryCache(Generic[_DiscoveryItem]): try: items: Final = await fetch() if self._pending.get(key) is asyncio.current_task(): - now: Final = self._clock() - live_entries: Final = tuple( - (entry_key, entry) for entry_key, entry in self._entries.items() if entry.expires_at > now - ) - self._entries = MappingProxyType( - { - entry_key: entry - for entry_key, entry in ( - *live_entries[-(_DISCOVERY_CACHE_LIMIT - 1) :], - ( - key, - _DiscoveryEntry(now + self._ttl, tuple(item.model_copy(deep=True) for item in items)), - ), - ) - } + self._entries.set_cache( + json.dumps(key), + _DiscoveryEntry(tuple(item.model_copy(deep=True) for item in items)), + ttl=self._ttl, ) return items finally: if self._pending.get(key) is asyncio.current_task(): - self._pending = MappingProxyType( - {entry_key: task for entry_key, task in self._pending.items() if entry_key != key} - ) + self._pending.pop(key) def _mcp_discovery_cache_ttl() -> float: diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/test_litellm/caching/test_in_memory_cache.py index 85e8308ae91..40ad4f0c6f0 100644 --- a/tests/test_litellm/caching/test_in_memory_cache.py +++ b/tests/test_litellm/caching/test_in_memory_cache.py @@ -250,3 +250,27 @@ def test_in_memory_cache_prunes_expired_heap_entries_below_capacity(): assert len(in_memory_cache.cache_dict) == 5 assert len(in_memory_cache.ttl_dict) == 5 assert len(in_memory_cache.expiration_heap) == 5 + + +def test_in_memory_cache_injected_clock_controls_expiry_and_eviction() -> None: + class Clock: + now = 0.0 + + def __call__(self) -> float: + return self.now + + clock = Clock() + cache = InMemoryCache(max_size_in_memory=2, default_ttl=60, clock=clock) + cache.set_cache("first", "original", ttl=10) + clock.now = 9.0 + cache.set_cache("second", "survivor") + assert cache.get_cache("first") == "original" + clock.now = 10.001 + assert cache.get_cache("first") is None + cache.set_cache("third", "replacement") + assert cache.get_cache("second") == "survivor" + clock.now = 69.001 + cache.set_cache("fourth", "new") + assert cache.get_cache("second") is None + assert cache.get_cache("third") == "replacement" + assert cache.get_cache("fourth") == "new" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index f1343bd143d..6630e184a16 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13095,7 +13095,7 @@ async def test_discovery_cache_reuses_raw_results_and_expires(kind: str) -> None clock.now = 59.999 assert (await operation(server, None))[0].name == "discovery-example" assert upstream.initializes == 1 - clock.now = 60.0 + clock.now = 60.001 assert (await operation(server, None))[0].name == "discovery-example" assert upstream.initializes == 2 @@ -13383,3 +13383,52 @@ async def test_discovery_resolves_stored_oauth_for_the_requesting_user() -> None assert store.calls == (("requesting-user", "discovery"), ("requesting-user", "discovery")) assert upstream.initializes == 1 assert ("prompts/list", "Bearer stored-token") in upstream.requests + + +@pytest.mark.asyncio +async def test_discovery_cache_evicts_results_at_capacity() -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + + async def original() -> list[Prompt]: + return [Prompt(name="original")] + + async def refetched() -> list[Prompt]: + return [Prompt(name="refetched")] + + for index in range(1025): + assert (await cache.get((f"server-{index:04}", None), original))[0].name == "original" + assert (await cache.get(("server-1024", None), refetched))[0].name == "original" + assert (await cache.get(("server-0000", None), refetched))[0].name == "refetched" + + +@pytest.mark.asyncio +async def test_discovery_cache_invalidation_preserves_other_servers_and_pending_fetches() -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def original() -> list[Prompt]: + return [Prompt(name="original")] + + async def blocked() -> list[Prompt]: + entered.set() + await release.wait() + return [Prompt(name="pending")] + + async def refetched() -> list[Prompt]: + return [Prompt(name="refetched")] + + assert (await cache.get(("server", None), original))[0].name == "original" + assert (await cache.get(("server-extra", None), original))[0].name == "original" + task: Final = asyncio.create_task(cache.get(("other", None), blocked)) + await asyncio.wait_for(entered.wait(), timeout=5) + cache.invalidate("server") + release.set() + assert (await asyncio.wait_for(task, timeout=5))[0].name == "pending" + assert (await cache.get(("other", None), refetched))[0].name == "pending" + assert (await cache.get(("server-extra", None), refetched))[0].name == "original" + assert (await cache.get(("server", None), refetched))[0].name == "refetched" From 6c07876dcfe9ce74401526c3a29657848ffb4afe Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 12 Sep 2026 01:00:11 +0000 Subject: [PATCH 142/155] test(model_prices): use local cost map for Fireworks cost coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_fireworks_serverless_model_costs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index e3404e477b7..1303f46e8fa 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -58,7 +58,7 @@ def test_bare_fireworks_ids_resolve_through_prefixed_entries(): assert info["max_output_tokens"] == expected["max_output_tokens"] -def test_deepseek_v4p1_flash_twin_costs(): +def test_deepseek_v4p1_flash_twin_costs(local_model_cost_map): for model in ( "fireworks_ai/deepseek-v4p1-flash", "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash", From 9d31de2f20d444d02674c8ec3b8f25a2cfa3257e Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:08:27 -0700 Subject: [PATCH 143/155] fix(mcp): bound discovery cache result bytes --- .../mcp_server/mcp_server_manager.py | 34 ++++++++++--------- .../mcp_server/test_mcp_server_manager.py | 25 ++++++++++---- 2 files changed, 37 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 35dfeb9ee74..9211d8ab003 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -46,7 +46,7 @@ from mcp.types import ( ResourceTemplate, ) from mcp.types import Tool as MCPTool -from pydantic import AnyUrl, BaseModel +from pydantic import AnyUrl, BaseModel, TypeAdapter from typing_extensions import ReadOnly import litellm @@ -1684,15 +1684,13 @@ _DiscoveryKey: TypeAlias = tuple[str, str | None] _DISCOVERY_CACHE_LIMIT: Final = 1024 -@dataclass(frozen=True, slots=True) -class _DiscoveryEntry(Generic[_DiscoveryItem]): - items: tuple[_DiscoveryItem, ...] - - class _DiscoveryCache(Generic[_DiscoveryItem]): - def __init__(self, ttl: float, clock: Callable[[], float]) -> None: + def __init__( + self, ttl: float, clock: Callable[[], float], adapter: TypeAdapter[tuple[_DiscoveryItem, ...]] + ) -> None: self._ttl = ttl - self._entries = InMemoryCache(max_size_in_memory=_DISCOVERY_CACHE_LIMIT, clock=clock) + self._adapter = adapter + self._entries = InMemoryCache(max_size_in_memory=_DISCOVERY_CACHE_LIMIT, max_size_per_item=64, clock=clock) self._pending: dict[ _DiscoveryKey, asyncio.Task[list[_DiscoveryItem]] ] = {} # mutable-ok: constant-time fetch registration @@ -1720,11 +1718,9 @@ class _DiscoveryCache(Generic[_DiscoveryItem]): ) -> tuple[_DiscoveryItem, ...]: if self._ttl <= 0: return tuple(await fetch()) - entry: Final = cast( # cast-ok: private cache contains only entries for this item type - "_DiscoveryEntry[_DiscoveryItem] | None", self._entries.get_cache(json.dumps(key)) - ) + entry: Final[object] = self._entries.get_cache(json.dumps(key)) if entry is not None: - return tuple(item.model_copy(deep=True) for item in entry.items) + return self._adapter.validate_python(entry) pending: Final = self._pending.get(key) if pending is not None: return await self._await_fetch(key, pending) @@ -1760,7 +1756,7 @@ class _DiscoveryCache(Generic[_DiscoveryItem]): if self._pending.get(key) is asyncio.current_task(): self._entries.set_cache( json.dumps(key), - _DiscoveryEntry(tuple(item.model_copy(deep=True) for item in items)), + self._adapter.dump_json(tuple(items)), ttl=self._ttl, ) return items @@ -1909,9 +1905,15 @@ class MCPServerManager: token_exchanger=build_token_exchanger(), ) discovery_ttl: Final = _mcp_discovery_cache_ttl() - self._prompt_discovery_cache = _DiscoveryCache[Prompt](discovery_ttl, discovery_clock) - self._resource_discovery_cache = _DiscoveryCache[Resource](discovery_ttl, discovery_clock) - self._template_discovery_cache = _DiscoveryCache[ResourceTemplate](discovery_ttl, discovery_clock) + self._prompt_discovery_cache = _DiscoveryCache[Prompt]( + discovery_ttl, discovery_clock, TypeAdapter(tuple[Prompt, ...]) + ) + self._resource_discovery_cache = _DiscoveryCache[Resource]( + discovery_ttl, discovery_clock, TypeAdapter(tuple[Resource, ...]) + ) + self._template_discovery_cache = _DiscoveryCache[ResourceTemplate]( + discovery_ttl, discovery_clock, TypeAdapter(tuple[ResourceTemplate, ...]) + ) self.registry: dict[str, MCPServer] = {} self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe) self.config_mcp_servers: dict[str, MCPServer] = {} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 6630e184a16..d2987c5112e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -30,7 +30,7 @@ from mcp.types import ( TextResourceContents, ) from mcp.types import Tool as MCPTool -from pydantic import AnyUrl +from pydantic import AnyUrl, TypeAdapter from litellm.constants import MCP_METADATA_TIMEOUT from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -13224,7 +13224,7 @@ def test_discovery_cache_keys_isolate_user_dependent_auth(auth_type: MCPAuth) -> async def test_discovery_cache_retries_cancelled_fetches() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache - cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) async def cancelled() -> list[Prompt]: raise asyncio.CancelledError() @@ -13241,7 +13241,7 @@ async def test_discovery_cache_retries_cancelled_fetches() -> None: async def test_discovery_cache_cancels_fetch_when_last_waiter_leaves() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache - cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) entered: Final = asyncio.Event() stopped: Final = asyncio.Event() release: Final = asyncio.Event() @@ -13270,7 +13270,7 @@ async def test_discovery_cache_cancels_fetch_when_last_waiter_leaves() -> None: async def test_discovery_cache_bounds_detached_fetches_without_dropping_results() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache - cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) entered: Final[asyncio.Queue[None]] = asyncio.Queue() release: Final = asyncio.Event() @@ -13389,7 +13389,7 @@ async def test_discovery_resolves_stored_oauth_for_the_requesting_user() -> None async def test_discovery_cache_evicts_results_at_capacity() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache - cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) async def original() -> list[Prompt]: return [Prompt(name="original")] @@ -13407,7 +13407,7 @@ async def test_discovery_cache_evicts_results_at_capacity() -> None: async def test_discovery_cache_invalidation_preserves_other_servers_and_pending_fetches() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache - cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) entered: Final = asyncio.Event() release: Final = asyncio.Event() @@ -13432,3 +13432,16 @@ async def test_discovery_cache_invalidation_preserves_other_servers_and_pending_ assert (await cache.get(("other", None), refetched))[0].name == "pending" assert (await cache.get(("server-extra", None), refetched))[0].name == "original" assert (await cache.get(("server", None), refetched))[0].name == "refetched" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("description", ("x" * 96_000, "é" * 40_000), ids=("ascii", "unicode")) +async def test_discovery_cache_returns_oversized_results_without_retaining_them(description: str) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) + fetch: Final = AsyncMock(return_value=[Prompt(name="large", description=description)]) + for _ in range(2): + result: Final = await cache.get(("server", None), fetch) + assert result[0].description == description + assert fetch.await_count == 2 From 105dc7710959e63264c99f0f68c581781e254494 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:14:21 -0700 Subject: [PATCH 144/155] fix(realtime): probe Azure's GA realtime upstream in health checks when no protocol is pinned --- litellm/realtime_api/main.py | 8 ++--- tests/test_litellm/realtime_api/test_main.py | 36 +++++++++++++++----- 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index d42e7e18b75..44c47af57f4 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -586,9 +586,7 @@ def _azure_realtime_health_protocol( configured: Final = configured_raw if isinstance(configured_raw, str) else None if configured is not None: return configured, query_params - if query_params is not None: - return "GA", query_params - return "beta", None + return "GA", query_params def _realtime_health_check_auth_headers( @@ -621,8 +619,8 @@ async def _realtime_health_check( api_key: str - api key custom_llm_provider: str - custom llm provider realtime_protocol: Optional[str] - protocol version ("GA"/"v1" for GA path, "beta" for beta path); - None resolves it for Azure from model_params/env, with transcription-only models probing GA - plus intent=transcription the way real calls do + None resolves it for Azure from model_params/env and otherwise probes GA, the upstream a client + without the OpenAI-Beta header is bridged to, with transcription-only models adding intent=transcription Returns: bool - True if connection is successful, False otherwise diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 8a9abe819e5..0827bbcdc38 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -266,7 +266,8 @@ def test_transcription_only_detection_rejects_speech_model(local_model_cost_map) @pytest.mark.asyncio -async def test_azure_health_check_keeps_beta_path_for_speech_model(): +async def test_azure_health_check_probes_the_ga_upstream_for_an_unconfigured_speech_model(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) connect = _CapturingConnect() with patch("websockets.connect", connect): assert await realtime_main._realtime_health_check( @@ -276,14 +277,18 @@ async def test_azure_health_check_keeps_beta_path_for_speech_model(): api_base="https://my-endpoint.openai.azure.com", api_version="2024-10-01-preview", ) - assert connect.url == ( - "wss://my-endpoint.openai.azure.com/openai/realtime" - "?api-version=2024-10-01-preview&deployment=gpt-4o-realtime-preview" - ) + assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview" + + +_AZURE_BETA_HEALTH_URL: Final = ( + "wss://my-endpoint.openai.azure.com/openai/realtime" + "?api-version=2024-10-01-preview&deployment=gpt-4o-realtime-preview" +) @pytest.mark.asyncio -async def test_azure_health_check_honors_deployment_realtime_protocol(): +async def test_azure_health_check_honors_deployment_realtime_protocol(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) connect = _CapturingConnect() with patch("websockets.connect", connect): assert await realtime_main._realtime_health_check( @@ -292,9 +297,24 @@ async def test_azure_health_check_honors_deployment_realtime_protocol(): api_key="fake-key", api_base="https://my-endpoint.openai.azure.com", api_version="2024-10-01-preview", - model_params={"realtime_protocol": "GA"}, + model_params={"realtime_protocol": "beta"}, ) - assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview" + assert connect.url == _AZURE_BETA_HEALTH_URL + + +@pytest.mark.asyncio +async def test_azure_health_check_honors_env_realtime_protocol(monkeypatch): + monkeypatch.setenv("LITELLM_AZURE_REALTIME_PROTOCOL", "beta") + connect = _CapturingConnect() + with patch("websockets.connect", connect): + assert await realtime_main._realtime_health_check( + model="gpt-4o-realtime-preview", + custom_llm_provider="azure", + api_key="fake-key", + api_base="https://my-endpoint.openai.azure.com", + api_version="2024-10-01-preview", + ) + assert connect.url == _AZURE_BETA_HEALTH_URL class _ConnectThatStopsAfterCapturingTheUrl: From f84f986b4ee571c8b3db18a560e007801379a933 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:16:33 -0700 Subject: [PATCH 145/155] fix(guardrails): keep post_call guardrail info on streamed chat completions (#40806) * fix(guardrails): sync logging_obj guardrail info on every record so post_call entries survive streamed chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(guardrails): hoist regression test imports to module scope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 1 + .../integrations/test_custom_guardrail.py | 61 ++++++++++++++++++- 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index bb54767edef..8a976a966a6 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1217,6 +1217,7 @@ class CustomGuardrail(CustomLogger): _, metadata_bucket = get_or_create_metadata_bucket(request_data) _append_guardrail_info(metadata_bucket) + _sync_guardrail_info_to_logging_obj(request_data, request_data.get("litellm_logging_obj")) _guardrail_self_recorded.set(True) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 1644d78ae37..bb4822eae57 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,4 +1,5 @@ import asyncio +import datetime as dt from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional from unittest.mock import AsyncMock @@ -9,9 +10,16 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy._types import CallTypes, UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks, Mode -from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail +from litellm.types.utils import ( + Choices, + GenericGuardrailAPIInputs, + GuardrailTracingDetail, + Message, + ModelResponse, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -2345,6 +2353,57 @@ class TestUndecoratedApplyGuardrailIsLogged: assert _Labelled.seen_label == "docs-style" + @pytest.mark.asyncio + async def test_post_call_recorded_outside_decorator_reaches_standard_logging_object(self): + """LIT-7608 regression: the auto-wrapped pre_call apply_guardrail copies the request bucket + into logging_obj.litellm_params["metadata"]. A post_call entry recorded later without the + decorator (the Bedrock streaming hook) must not be shadowed by that stale copy.""" + messages: Final = [{"role": "user", "content": "hello there"}] + litellm_metadata: Final[dict] = {"user_api_key_user_id": "u1"} + logging_obj: Final = Logging( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=messages, + stream=True, + call_type=CallTypes.acompletion.value, + start_time=dt.datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + litellm_params={"litellm_metadata": litellm_metadata}, + optional_params={}, + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + custom_llm_provider="bedrock", + ) + request_data: Final = { + "model": "bedrock-haiku", + "messages": messages, + "litellm_metadata": litellm_metadata, + "litellm_logging_obj": logging_obj, + } + guardrail: Final = _UndecoratedGuardrail(guardrail_name="bedrock-pre", event_hook=GuardrailEventHooks.pre_call) + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello there"]), + request_data=request_data, + input_type="request", + logging_obj=logging_obj, + ) + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"action": "NONE"}, + request_data=request_data, + guardrail_status="success", + event_type=GuardrailEventHooks.post_call, + ) + await logging_obj.async_success_handler( + result=ModelResponse(choices=[Choices(message=Message(role="assistant", content="general kenobi"))]), + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + ) + + entries: Final = logging_obj.model_call_details["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_mode"] for e in entries] == ["pre_call", "post_call"] + class _ApplyOnlyObserver(CustomGuardrail): """Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook.""" From cf97b757a4356b02baef6b88b5b216ed01dcb515 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 11 Sep 2026 18:37:08 -0700 Subject: [PATCH 146/155] fix(ui): preserve cleared shared select values (#40795) Preserve explicit null when shared selectors clear and adapt affected forms, validation, and request payloads. Clear stale dependent relationships and retain required-selection checks. Document project detachment, user model-budget clearing, and routing-compression clearing as deferred follow-ups. --- .../add_agent_form.integration.test.tsx | 33 ++++++++++++++- .../agents/_components/add_agent_form.tsx | 16 ++++---- .../_components/ShadowEvalStartForm.tsx | 17 ++++---- .../_components/pricing_calculator/index.tsx | 4 +- .../_components/pricing_calculator/types.ts | 4 +- .../_components/TeamGuardrailsTab.tsx | 5 ++- .../components/AllModelsTable.tsx | 4 +- .../panels/AddModelPanel.tsx | 6 ++- ...I.test.tsx => ChatUI.integration.test.tsx} | 16 ++++++++ .../playground/components/chat_ui/ChatUI.tsx | 30 +++++++++----- .../components/chat_ui/EndpointSelector.tsx | 4 +- .../components/chat_ui/SessionManagement.tsx | 2 +- .../compareUI/components/ModelSelector.tsx | 3 +- ...x => add_policy_form.integration.test.tsx} | 0 .../policies/_components/add_policy_form.tsx | 14 +++---- .../_components/ai_suggestion_modal.tsx | 6 +-- .../_components/pipeline_flow_builder.tsx | 2 +- .../_components/template_parameter_modal.tsx | 6 +-- .../ProjectModals/EditProjectModal.tsx | 6 +-- .../ProjectModals/ProjectBaseForm.tsx | 6 +-- .../ProjectModals/projectFormSchema.ts | 10 +++-- .../prompt_editor_view/ModelConfigCard.tsx | 4 +- .../prompt_editor_view/PromptEditorHeader.tsx | 4 +- .../_components/prompt_editor_view/types.ts | 2 +- .../prompt_editor_view/utils.test.ts | 16 ++++++++ .../_components/prompt_editor_view/utils.ts | 4 +- .../DefaultUserSettingsForm.tsx | 11 +++-- .../default-user-settings/mapper.test.ts | 8 ++-- .../default-user-settings/mapper.ts | 11 +++-- .../default-user-settings/schema.ts | 12 ++++-- .../_components/view_users/UsersTable.tsx | 4 +- .../src/components/CreateUserButton.tsx | 2 +- .../MCPSemanticFilterSettings.tsx | 2 +- .../MCPSemanticFilterTestPanel.tsx | 4 +- .../semanticFilterTestUtils.ts | 6 +-- .../toolSearchForm.test.ts | 1 + .../MCPToolSearchSettings/toolSearchForm.ts | 6 +-- .../Fallbacks/FallbackGroupConfig.tsx | 10 ++--- ui/litellm-dashboard/src/components/Teams.tsx | 4 +- .../src/components/TeamsPage/TeamsTable.tsx | 2 +- .../VirtualKeysPage/VirtualKeysTable.tsx | 4 +- ....tsx => AddModelForm.integration.test.tsx} | 2 +- .../src/components/add_model/AddModelForm.tsx | 15 ++++--- .../add_model/ClassificationMethodConfig.tsx | 3 +- .../add_model/ComplexityRouterConfig.tsx | 2 +- .../add_model/CompressionControls.tsx | 8 ++-- ... RouterConfigBuilder.integration.test.tsx} | 2 +- .../add_model/RouterConfigBuilder.test.ts | 10 +++++ .../add_model/RouterConfigBuilder.tsx | 23 +++++++---- .../add_model/SemanticKeywordMatching.tsx | 4 +- .../add_model/add_auto_router_tab.tsx | 19 +++++---- .../add_model/litellm_model_name.tsx | 10 ++--- .../add_model/provider_specific_fields.tsx | 3 +- .../common_components/ModelAliasManager.tsx | 14 +++++-- .../common_components/ModelSelector.tsx | 16 ++++---- .../OrganizationDropdown.tsx | 4 +- .../common_components/ProjectDropdown.tsx | 4 +- .../common_components/UserDropdown.tsx | 4 +- .../common_components/team_dropdown.tsx | 8 ++-- ...=> user_search_modal.integration.test.tsx} | 18 +++++--- .../common_components/user_search_modal.tsx | 28 +++++++++---- .../edit_auto_router_modal.tsx | 8 ++-- .../BudgetFallbacksEditor.tsx | 4 +- .../key_team_helpers/ModelMaxBudgetEditor.tsx | 4 +- .../components/model_add/CredentialModal.tsx | 8 ++-- .../model_add/credential_form_helpers.ts | 6 +-- .../components/organisms/createKeyPayload.ts | 2 + .../create_key_button.integration.test.tsx | 32 ++++++++++++++- .../organisms/create_key_button.tsx | 24 +++++------ .../src/components/provider_info_helpers.tsx | 2 +- ...aginatedSearchSelect.integration.test.tsx} | 41 ++++++++++++++----- .../shared/PaginatedSearchSelect.tsx | 10 ++--- ....tsx => SearchSelect.integration.test.tsx} | 32 ++++++++++++--- .../src/components/shared/SearchSelect.tsx | 10 ++--- .../src/components/team/TeamInfo.tsx | 7 +++- ...tsx => key_edit_view.integration.test.tsx} | 32 +++++++++++++-- .../components/templates/key_edit_view.tsx | 15 ++++--- .../view_logs/RequestLogsFilters.tsx | 10 ++--- 78 files changed, 492 insertions(+), 263 deletions(-) rename ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/{ChatUI.test.tsx => ChatUI.integration.test.tsx} (95%) rename ui/litellm-dashboard/src/app/(dashboard)/policies/_components/{add_policy_form.test.tsx => add_policy_form.integration.test.tsx} (100%) rename ui/litellm-dashboard/src/components/add_model/{AddModelForm.test.tsx => AddModelForm.integration.test.tsx} (99%) rename ui/litellm-dashboard/src/components/add_model/{RouterConfigBuilder.test.tsx => RouterConfigBuilder.integration.test.tsx} (99%) create mode 100644 ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts rename ui/litellm-dashboard/src/components/common_components/{user_search_modal.test.tsx => user_search_modal.integration.test.tsx} (94%) rename ui/litellm-dashboard/src/components/shared/{PaginatedSearchSelect.test.tsx => PaginatedSearchSelect.integration.test.tsx} (89%) rename ui/litellm-dashboard/src/components/shared/{SearchSelect.test.tsx => SearchSelect.integration.test.tsx} (76%) rename ui/litellm-dashboard/src/components/templates/{key_edit_view.test.tsx => key_edit_view.integration.test.tsx} (98%) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx index dad8599e967..ebc97891744 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx @@ -1,11 +1,11 @@ import React from "react"; -import { render, screen, waitFor, within } from "@testing-library/react"; +import { screen, waitFor, within } from "@testing-library/react"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import AddAgentForm from "./add_agent_form"; import * as networking from "@/components/networking"; import type { AgentCreateInfo } from "@/components/networking"; -import { chooseSelectOption } from "../../../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders as render } from "../../../../../tests/test-utils"; vi.mock("@/components/networking", () => ({ createAgentCall: vi.fn(), @@ -351,4 +351,33 @@ describe("AddAgentForm submit payload", () => { expect(await screen.findByText("Agent Created!")).toBeInTheDocument(); expect(within(screen.getByText("Agent Created!").parentElement!).getByText("created-agent")).toBeInTheDocument(); }); + it("blocks creation after clearing the existing key and assigns the reselected key", async () => { + vi.mocked(networking.keyListCall).mockResolvedValue({ + keys: [{ token: "key-maple", key_alias: "Maple key" }], + }); + const user = userEvent.setup(); + renderForm(); + await user.type(await screen.findByLabelText("Agent Name"), "key-selection-agent"); + await user.type(screen.getByLabelText("Display Name"), "Key selection"); + await user.type(screen.getByPlaceholderText("Describe what this agent does..."), "d"); + for (let step = 0; step < 3; step++) { + await user.click(screen.getByRole("button", { name: /^Next/ })); + } + await user.click(screen.getByRole("radio", { name: "Assign an existing key" })); + const keySelector = await screen.findByPlaceholderText("Search by key name…"); + await chooseSelectOption(user, keySelector, "Maple key"); + await user.click(screen.getByRole("button", { name: "Clear" })); + await user.click(screen.getByRole("button", { name: /Create Agent/ })); + expect(networking.createAgentCall).not.toHaveBeenCalled(); + expect(networking.keyUpdateCall).not.toHaveBeenCalled(); + await chooseSelectOption(user, keySelector, "Maple key"); + await user.click(screen.getByRole("button", { name: /Create Agent/ })); + await waitFor(() => + expect(networking.keyUpdateCall).toHaveBeenCalledWith("tok", { + key: "key-maple", + agent_id: "agent-1", + }), + ); + expect(networking.createAgentCall).toHaveBeenCalledTimes(1); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx index 108bae977e1..e71fed40209 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx @@ -338,6 +338,11 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok return; } + if (keyAssignOption === "existing_key" && !selectedExistingKey) { + toast.error("Please select an existing key to assign"); + return; + } + setIsSubmitting(true); try { const isValid = await form.trigger(); @@ -406,12 +411,7 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok selectedTeamId, ); setCreatedKeyValue(keyResponse.key || null); - } else if (keyAssignOption === "existing_key") { - if (!selectedExistingKey) { - toast.error("Please select an existing key to assign"); - setIsSubmitting(false); - return; - } + } else if (keyAssignOption === "existing_key" && selectedExistingKey) { await keyUpdateCall(accessToken, { key: selectedExistingKey, agent_id: agentId, @@ -963,8 +963,8 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok setSelectedExistingKey(value || null)} + value={selectedExistingKey} + onValueChange={setSelectedExistingKey} options={existingKeys.map((k) => ({ label: k.key_alias || k.token?.slice(0, 12) + "…", value: k.token, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx index d85d26a21a8..06e332d2cfc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx @@ -170,8 +170,8 @@ interface StartFormValidityInputs { models: string[]; routerNames: string[]; direction: ShadowEvalDirection; - baselineModel: string; - judgeModel: string; + baselineModel: string | null; + judgeModel: string | null; percentage: string; maxBudget: string; } @@ -181,13 +181,13 @@ const startFormValidity = (inputs: StartFormValidityInputs) => { const percentageValid = parsedPct >= 0.1 && parsedPct <= 100; const parsedMaxBudget = Number.parseFloat(inputs.maxBudget); const maxBudgetValid = parsedMaxBudget >= 0.01 && parsedMaxBudget <= 10000; - const baselinePicked = inputs.direction === "forward" || inputs.baselineModel !== ""; + const baselinePicked = inputs.direction === "forward" || Boolean(inputs.baselineModel); const targetsPicked = inputs.apiKeyIds.length + inputs.teamIds.length + inputs.userIds.length > 0; const routerCountValid = inputs.routerNames.length >= 1 && inputs.routerNames.length <= MAX_ROUTERS; const routersMatchDirection = inputs.direction === "forward" || inputs.routerNames.length === 1; const routersValid = routerCountValid && routersMatchDirection; const scopeValid = routersValid && (inputs.direction === "reverse" || inputs.models.length <= MAX_MODELS); - const modelsPicked = scopeValid && inputs.judgeModel !== "" && baselinePicked; + const modelsPicked = scopeValid && Boolean(inputs.judgeModel) && baselinePicked; const filled = targetsPicked && modelsPicked; const boundsValid = percentageValid && maxBudgetValid; const valid = Boolean(inputs.accessToken) && filled && boundsValid; @@ -201,7 +201,7 @@ interface StartBodyInputs { models: string[]; routerNames: string[]; direction: ShadowEvalDirection; - baselineModel: string; + baselineModel: string | null; shadowPercentage: number; durationDays: number; maxBudget: number; @@ -215,7 +215,7 @@ const buildStartBody = (inputs: StartBodyInputs) => ({ models: inputs.direction === "forward" ? inputs.models : [], router_names: inputs.routerNames, direction: inputs.direction, - ...(inputs.direction === "reverse" ? { baseline_model: inputs.baselineModel } : {}), + ...(inputs.direction === "reverse" ? { baseline_model: inputs.baselineModel ?? undefined } : {}), shadow_percentage: inputs.shadowPercentage, duration_days: inputs.durationDays, max_budget: inputs.maxBudget, @@ -230,10 +230,10 @@ export const StartForm: React.FC = () => { const [models, setModels] = useState([]); const [routerNames, setRouterNames] = useState([]); const [direction, setDirection] = useState("forward"); - const [baselineModel, setBaselineModel] = useState(""); + const [baselineModel, setBaselineModel] = useState(null); const [percentage, setPercentage] = useState("10"); const [durationDays, setDurationDays] = useState("7"); - const [judgeModel, setJudgeModel] = useState(""); + const [judgeModel, setJudgeModel] = useState(null); const [maxBudget, setMaxBudget] = useState("10"); const { data: autoRouters } = useAutoRouters(); const configuredGroups = usePlainModelGroups(); @@ -286,6 +286,7 @@ export const StartForm: React.FC = () => { }; const { parsedPct, parsedMaxBudget, percentageValid, maxBudgetValid, valid } = startFormValidity(validityInputs); const handleStart = () => { + if (!valid || !judgeModel) return; const bodyInputs: StartBodyInputs = { apiKeyIds, teamIds, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx index f3bd74260ad..b0612412059 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx @@ -15,7 +15,7 @@ const generateId = () => `entry-${Date.now()}-${Math.random().toString(36).subst const createDefaultEntry = (): ModelEntry => ({ id: generateId(), - model: "", + model: null, input_tokens: 1000, output_tokens: 500, num_requests_per_day: undefined, @@ -28,7 +28,7 @@ const PricingCalculator: React.FC = ({ accessToken, mode const { debouncedFetchForEntry, removeEntry, getMultiModelResult } = useMultiCostEstimate(accessToken); const handleEntryChange = useCallback( - (id: string, field: keyof ModelEntry, value: string | number | undefined) => { + (id: string, field: keyof ModelEntry, value: string | number | null | undefined) => { setEntries((prev) => { const updated = prev.map((entry) => (entry.id === id ? { ...entry, [field]: value } : entry)); const changedEntry = updated.find((e) => e.id === id); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts index 250857a74f5..859d2f3e8d9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts @@ -4,7 +4,7 @@ export interface PricingCalculatorProps { } export interface PricingFormValues { - model: string; + model: string | null; input_tokens: number; output_tokens: number; num_requests_per_day?: number; @@ -13,7 +13,7 @@ export interface PricingFormValues { export interface ModelEntry { id: string; - model: string; + model: string | null; input_tokens: number; output_tokens: number; num_requests_per_day?: number; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx index 1de4e697f64..d45cfc3fe7d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -48,7 +48,10 @@ const GUARDRAIL_MODES = [ ] as const; const submitGuardrailSchema = z.object({ - team_id: z.string().min(1, "Select a team"), + team_id: z + .string() + .nullable() + .pipe(z.string({ error: "Select a team" }).min(1, "Select a team")), guardrail_name: z.string().min(1, "Enter a guardrail name"), mode: z.string().min(1, "Select a mode"), api_base: z.string().min(1, "Enter the API base URL").refine(isValidUrl, "Must be a valid URL"), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx index 5c7dbb18428..8482d0832c3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx @@ -278,7 +278,7 @@ export function AllModelsTable({ options={modelGroupOptions} value={(get(MODEL_NAME_COLUMN_ID) as string) ?? ALL_MODEL_GROUPS_VALUE} onValueChange={(value) => - set(MODEL_NAME_COLUMN_ID, value === ALL_MODEL_GROUPS_VALUE ? undefined : value) + set(MODEL_NAME_COLUMN_ID, value === ALL_MODEL_GROUPS_VALUE ? undefined : value ?? undefined) } placeholder="Filter by Public Model Name" emptyText="No models found" @@ -289,7 +289,7 @@ export function AllModelsTable({ options={accessGroupOptions} value={(get(ACCESS_GROUPS_COLUMN_ID) as string) ?? ALL_MODEL_GROUPS_VALUE} onValueChange={(value) => - set(ACCESS_GROUPS_COLUMN_ID, value === ALL_MODEL_GROUPS_VALUE ? undefined : value) + set(ACCESS_GROUPS_COLUMN_ID, value === ALL_MODEL_GROUPS_VALUE ? undefined : value ?? undefined) } placeholder="Filter by Model Access Group" emptyText="No model access groups found" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx index 1443c065d9b..1cdce04d07c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx @@ -26,7 +26,7 @@ export default function AddModelPanel() { const { data: modelCostMapData } = useModelCostMap(); const { data: credentialsResponse } = useCredentials(); const { data: teams } = useTeams(); - const [selectedProvider, setSelectedProvider] = useState(Providers.Anthropic); + const [selectedProvider, setSelectedProvider] = useState(Providers.Anthropic); const [providerModels, setProviderModels] = useState([]); const [showAdvancedSettings, setShowAdvancedSettings] = useState(false); @@ -57,7 +57,9 @@ export default function AddModelPanel() { selectedProvider={selectedProvider} setSelectedProvider={setSelectedProvider} providerModels={providerModels} - setProviderModelsFn={(provider) => setProviderModels(getProviderModels(provider, modelCostMapData))} + setProviderModelsFn={(provider) => + setProviderModels(provider === null ? [] : getProviderModels(provider, modelCostMapData)) + } getPlaceholder={getPlaceholder} showAdvancedSettings={showAdvancedSettings} setShowAdvancedSettings={setShowAdvancedSettings} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx similarity index 95% rename from ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx index bf6b092a2b0..984996351df 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx @@ -153,6 +153,7 @@ describe("ChatUI", () => { }); it("should allow the user to select a model", async () => { + const user = userEvent.setup(); render( { await waitFor(() => { expect(screen.getAllByText("Model 1").length).toBeGreaterThan(0); }); + await user.click(screen.getByRole("option", { name: "Model 1Mode: chat" })); + expect(screen.getByPlaceholderText("Select a Model")).toHaveValue("Model 1"); + + await user.click(screen.getAllByRole("button", { name: "Clear" })[0]); + const input = screen.getByPlaceholderText("Describe the image you want to generate..."); + fireEvent.change(input, { target: { value: "Contract endpoint check" } }); + expect(screen.getByRole("button", { name: "Send message" })).toBeDisabled(); + fireEvent.keyDown(input, { key: "Enter", code: "Enter" }); + expect(input).toHaveValue("Contract endpoint check"); + expect(makeOpenAIChatCompletionRequest).not.toHaveBeenCalled(); + expect(sessionStorage.getItem("endpointType")).toBeNull(); + + await selectComboboxOption("Select an endpoint", "/v1/chat/completions"); + await selectComboboxOption("Select a Model", "Model 1"); + expect(screen.getByRole("button", { name: "Send message" })).toBeEnabled(); }); it("shows only endpoint-compatible models when chat endpoint is selected", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index eca9e2323ae..ed8679cfdc1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -196,17 +196,17 @@ const ChatUI: React.FC = ({ () => sessionStorage.getItem("customProxyBaseUrl") || "", ); const [inputMessage, setInputMessage] = useState(""); - const [selectedModel, setSelectedModel] = useState(simplified ? fixedModel : undefined); + const [selectedModel, setSelectedModel] = useState(simplified ? fixedModel : null); const [showCustomModelInput, setShowCustomModelInput] = useState(false); const [modelInfo, setModelInfo] = useState([]); const [isLoadingModels, setIsLoadingModels] = useState(false); const [modelLoadError, setModelLoadError] = useState(false); const [agentInfo, setAgentInfo] = useState([]); - const [selectedAgent, setSelectedAgent] = useState(undefined); + const [selectedAgent, setSelectedAgent] = useState(null); const debouncedSetSelectedModel = useDebouncedCallback((value: string) => setSelectedModel(value), { wait: CUSTOM_MODEL_DEBOUNCE_WAIT_MS, }); - const [endpointType, setEndpointType] = useState( + const [endpointType, setEndpointType] = useState( () => sessionStorage.getItem("endpointType") || EndpointType.CHAT, ); const [isLoading, setIsLoading] = useState(false); @@ -327,7 +327,7 @@ const ChatUI: React.FC = ({ }; useEffect(() => { - if (isGetCodeModalVisible) { + if (isGetCodeModalVisible && endpointType !== null) { const code = generateCodeSnippet({ apiKeySource, accessToken, @@ -342,7 +342,7 @@ const ChatUI: React.FC = ({ mcpServers, mcpServerToolRestrictions, endpointType, - selectedModel, + selectedModel: selectedModel ?? undefined, selectedSdk, selectedVoice, proxySettings, @@ -376,7 +376,8 @@ const ChatUI: React.FC = ({ } catch { // Storage full or unavailable — non-critical, skip persisting. } - sessionStorage.setItem("endpointType", endpointType); + if (endpointType === null) sessionStorage.removeItem("endpointType"); + else sessionStorage.setItem("endpointType", endpointType); sessionStorage.setItem("selectedTags", JSON.stringify(selectedTags)); sessionStorage.setItem("selectedVectorStores", JSON.stringify(selectedVectorStores)); sessionStorage.setItem("selectedGuardrails", JSON.stringify(selectedGuardrails)); @@ -493,7 +494,7 @@ const ChatUI: React.FC = ({ setAgentInfo(agents); // Clear selection if current agent not in list if (selectedAgent && !agents.some((a) => a.agent_name === selectedAgent)) { - setSelectedAgent(undefined); + setSelectedAgent(null); } } catch (error) { console.error("Error fetching agents:", error); @@ -616,10 +617,11 @@ const ChatUI: React.FC = ({ setUploadedAudio(file); }; - const handleEndpointChange = (value: string) => { + const handleEndpointChange = (value: string | null) => { setEndpointType(value); - setSelectedModel(undefined); - setSelectedAgent(undefined); + setGeneratedCode(""); + setSelectedModel(null); + setSelectedAgent(null); setShowCustomModelInput(false); setSelectedMCPDirectTool(undefined); if (value === EndpointType.MCP) { @@ -710,6 +712,11 @@ const ChatUI: React.FC = ({ }; const handleSendMessage = async () => { + if (endpointType === null) { + toast.fromError("Please select an endpoint before sending a request"); + return; + } + if (inputMessage.trim() === "" && endpointType !== EndpointType.TRANSCRIPTION && endpointType !== EndpointType.MCP) return; @@ -1152,7 +1159,7 @@ const ChatUI: React.FC = ({ toast.success("Chat history cleared."); }; - const onModelChange = (value: string) => { + const onModelChange = (value: string | null) => { setSelectedModel(value); setShowCustomModelInput(value === "custom"); @@ -1210,6 +1217,7 @@ const ChatUI: React.FC = ({ : "Describe the image you want to generate..."; const sendDisabled = + endpointType === null || isLoading || (endpointType === EndpointType.MCP ? !(selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" && selectedMCPDirectTool) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx index 644bab5b5a4..6b286c403a8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx @@ -3,8 +3,8 @@ import React from "react"; import { ENDPOINT_OPTIONS } from "./chatConstants"; interface EndpointSelectorProps { - endpointType: string; // Accept string to avoid type conflicts - onEndpointChange: (value: string) => void; + endpointType: string | null; + onEndpointChange: (value: string | null) => void; className?: string; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx index 35c8e839159..4dbf9168897 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx @@ -7,7 +7,7 @@ import { Switch } from "@/components/ui/switch"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; interface SessionManagementProps { - endpointType: string; + endpointType: string | null; responsesSessionId: string | null; useApiSessionManagement: boolean; onToggleSessionManagement: (useApi: boolean) => void; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx index 5659875cf82..8ecc0e0c8bb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx @@ -22,7 +22,8 @@ export function ModelSelector({ value, onChange, models, loading, disabled }: Mo const selectValue = isAddingCustom ? "__custom__" : value || undefined; - const handleSelectChange = (selected: string) => { + const handleSelectChange = (selected: string | null) => { + if (selected === null) return; if (selected === "__custom__") { setIsAddingCustom(true); if (value && !options.includes(value)) { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.integration.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.integration.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx index 2830b0fda12..d14165dd849 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx @@ -44,10 +44,10 @@ const policyShape = { .min(1, "Please enter a policy name") .regex(/^[a-zA-Z0-9_-]+$/, "Policy name can only contain letters, numbers, hyphens, and underscores"), description: z.string(), - inherit: z.string(), + inherit: z.string().nullable(), guardrails_add: z.array(z.string()), guardrails_remove: z.array(z.string()), - model_condition: z.string(), + model_condition: z.string().nullable(), }; const policySchema = z.object(policyShape); @@ -57,19 +57,19 @@ type PolicyFormValues = z.infer; const EMPTY_VALUES: PolicyFormValues = { policy_name: "", description: "", - inherit: "", + inherit: null, guardrails_add: [], guardrails_remove: [], - model_condition: "", + model_condition: null, }; const toFormValues = (policy: Policy): PolicyFormValues => ({ policy_name: policy.policy_name, description: policy.description ?? "", - inherit: policy.inherit ?? "", + inherit: policy.inherit ?? null, guardrails_add: policy.guardrails_add || [], guardrails_remove: policy.guardrails_remove || [], - model_condition: policy.condition?.model ?? "", + model_condition: policy.condition?.model ?? null, }); const buildPolicyRequest = (values: PolicyFormValues): PolicyCreateRequest | PolicyUpdateRequest => ({ @@ -529,7 +529,7 @@ const AddPolicyForm: React.FC = ({ {...control} id={id} ref={ref} - value={value} + value={value ?? ""} onChange={onChange} placeholder="Leave empty to apply to all models (e.g., gpt-4.* or bedrock/claude-.*)" /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx index 32e777949d0..c4f76b07a4e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx @@ -68,7 +68,7 @@ const AiSuggestionModal: React.FC = ({ const [suggestions, setSuggestions] = useState(null); const [explanation, setExplanation] = useState(null); const [selectedIds, setSelectedIds] = useState>(new Set()); - const [selectedModel, setSelectedModel] = useState(undefined); + const [selectedModel, setSelectedModel] = useState(null); const [availableModels, setAvailableModels] = useState([]); const [isLoadingModels, setIsLoadingModels] = useState(false); // Test panel state @@ -114,7 +114,7 @@ const AiSuggestionModal: React.FC = ({ setSuggestions(null); setExplanation(null); setSelectedIds(new Set()); - setSelectedModel(undefined); + setSelectedModel(null); setShowTestPanel(false); setTestInputText(""); setIsTestLoading(false); @@ -837,7 +837,7 @@ const AiSuggestionModal: React.FC = ({ ({ label: m, value: m }))} value={selectedModel} - onValueChange={(value) => setSelectedModel(value || undefined)} + onValueChange={setSelectedModel} placeholder={isLoadingModels ? "Loading models..." : "Select a model to analyze your requirements"} emptyText="No models found" disabled={isLoadingModels} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/pipeline_flow_builder.tsx index 8651be39f0d..da77b348028 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/pipeline_flow_builder.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/pipeline_flow_builder.tsx @@ -343,7 +343,7 @@ const StepCard: React.FC = ({ onChange({ guardrail: value })} + onValueChange={(value) => onChange({ guardrail: value ?? undefined })} placeholder="Select a guardrail" emptyText="No guardrails found" /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx index 410d4ea8467..b90379b149d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx @@ -46,7 +46,7 @@ const TemplateParameterModal: React.FC = ({ }) => { const [parameterValues, setParameterValues] = useState>({}); const [competitorMode, setCompetitorMode] = useState<"ai" | "manual">("ai"); - const [selectedModel, setSelectedModel] = useState(undefined); + const [selectedModel, setSelectedModel] = useState(null); const [availableModels, setAvailableModels] = useState([]); const [isLoadingModels, setIsLoadingModels] = useState(false); const [competitorTags, setCompetitorTags] = useState([]); @@ -72,7 +72,7 @@ const TemplateParameterModal: React.FC = ({ }); setParameterValues(initial); setCompetitorMode("ai"); - setSelectedModel(undefined); + setSelectedModel(null); setCompetitorTags([]); setVariationsMap({}); setIsGenerating(false); @@ -297,7 +297,7 @@ const TemplateParameterModal: React.FC = ({ ({ label: m, value: m }))} value={selectedModel} - onValueChange={(value) => setSelectedModel(value || undefined)} + onValueChange={setSelectedModel} placeholder={isLoadingModels ? "Loading models..." : "Select a model to generate names"} emptyText="No models found" disabled={isLoadingModels} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx index 5c1a443d6c5..77f28b05ea5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx @@ -10,7 +10,7 @@ import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUpdateProject, ProjectUpdateParams } from "@/app/(dashboard)/hooks/projects/useUpdateProject"; import { ProjectBaseForm } from "./ProjectBaseForm"; -import { projectFormSchema, type ProjectFormValues } from "./projectFormSchema"; +import { projectFormSchema, type ProjectFormValues, type ProjectSubmitValues } from "./projectFormSchema"; import { buildProjectUpdateParams } from "./projectFormUtils"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; @@ -58,7 +58,7 @@ export const toFormValues = (project: ProjectResponse): ProjectFormValues => { return { project_alias: project.project_alias ?? "", - team_id: project.team_id ?? "", + team_id: project.team_id ?? null, description: project.description ?? "", models: project.models ?? [], max_budget: project.litellm_budget_table?.max_budget ?? undefined, @@ -81,7 +81,7 @@ function EditProjectForm({ project, onClose, onSuccess }: Omit { - const submitted: ProjectFormValues = advancedEverOpened + const submitted: ProjectSubmitValues = advancedEverOpened ? values : { ...values, guardrails: undefined, modelLimits: undefined, metadata: undefined }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx index fed1a88fe98..1ad8a1e4953 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx @@ -5,7 +5,7 @@ import { useFieldArray, useWatch, type UseFormReturn } from "react-hook-form"; import { ChevronDown, CircleAlert, Minus, Plus } from "lucide-react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { ALL_TEAM_MODELS, type ProjectFormValues } from "./projectFormSchema"; +import { ALL_TEAM_MODELS, type ProjectFormValues, type ProjectSubmitValues } from "./projectFormSchema"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { Team } from "@/components/key_team_helpers/key_list"; import { fetchTeamModels } from "@/components/organisms/create_key_button"; @@ -32,7 +32,7 @@ const toOptionalNumber = (raw: string): number | undefined => { }; interface ProjectBaseFormProps { - form: UseFormReturn; + form: UseFormReturn; advancedOpen: boolean; onAdvancedOpenChange: (open: boolean) => void; } @@ -94,7 +94,7 @@ export function ProjectBaseForm({ form, advancedOpen, onAdvancedOpenChange }: Pr } }, [selectedTeam, accessToken, userId, userRole]); - const handleTeamChange = (teamId: string) => { + const handleTeamChange = (teamId: string | null) => { const team = teams?.find((t) => t.team_id === teamId) ?? null; setSelectedTeam(team); form.setValue("models", []); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts index 6bfaa831bba..d4c85d89616 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts @@ -16,7 +16,10 @@ const modelLimitSchema = z.object({ export const projectFormSchema = z .object({ project_alias: z.string().min(1, "Please enter a project name"), - team_id: z.string().min(1, "Please select a team"), + team_id: z + .string() + .nullable() + .pipe(z.string({ error: "Please select a team" }).min(1, "Please select a team")), description: z.string().optional(), models: z.array(z.string()), max_budget: z.number().optional(), @@ -48,11 +51,12 @@ export const projectFormSchema = z }); }); -export type ProjectFormValues = z.output; +export type ProjectFormValues = z.input; +export type ProjectSubmitValues = z.output; export const emptyProjectFormValues: ProjectFormValues = { project_alias: "", - team_id: "", + team_id: null, description: undefined, models: [], max_budget: undefined, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/ModelConfigCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/ModelConfigCard.tsx index a14e45de99a..e9b526c9703 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/ModelConfigCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/ModelConfigCard.tsx @@ -6,11 +6,11 @@ import { SettingsIcon } from "lucide-react"; import ModelSelector from "@/components/common_components/ModelSelector"; interface ModelConfigCardProps { - model: string; + model: string | null; temperature?: number; maxTokens?: number; accessToken: string | null; - onModelChange: (model: string) => void; + onModelChange: (model: string | null) => void; onTemperatureChange: (temp: number) => void; onMaxTokensChange: (tokens: number) => void; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx index 04cac01365a..0f001979c28 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx @@ -21,7 +21,7 @@ interface PromptEditorHeaderProps { editMode?: boolean; onShowHistory?: () => void; version?: string | null; - promptModel?: string; + promptModel?: string | null; promptVariables?: Record; accessToken: string | null; proxySettings?: { @@ -85,7 +85,7 @@ const PromptEditorHeader: React.FC = ({
{ expect(result).toContain("output:"); expect(result).toContain("format: text"); expect(result).toContain("User: Hello world"); + const cleared = convertToDotPrompt({ ...prompt, model: null }); + expect(cleared).toBe(result.replace("model: gpt-4\n", "")); }); it("should include config parameters when set", () => { @@ -203,6 +205,20 @@ describe("convertToDotPrompt", () => { }); describe("parseExistingPrompt", () => { + it("should keep saved prompts with missing or blank models unassigned", () => { + for (const modelLine of ["", "model: \n"]) { + const prompt = parseExistingPrompt({ + prompt_spec: { + prompt_id: "unassigned-prompt", + litellm_params: { dotprompt_content: `---\n${modelLine}temperature: 0\n---\nUser: Keep this message` }, + }, + }); + + expect(prompt.model).toBeNull(); + expect(convertToDotPrompt(prompt)).not.toMatch(/^model:/m); + } + }); + it("should parse basic dotprompt content", () => { const apiResponse = { prompt_spec: { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/utils.ts b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/utils.ts index ef327d93495..3d768e930a9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/utils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/utils.ts @@ -23,7 +23,7 @@ export const extractVariables = (prompt: PromptType): string[] => { export const convertToDotPrompt = (prompt: PromptType): string => { const variables = extractVariables(prompt); - let result = `---\nmodel: ${prompt.model}\n`; + let result = prompt.model ? `---\nmodel: ${prompt.model}\n` : "---\n"; // Add temperature if set if (prompt.config.temperature !== undefined) { @@ -237,7 +237,7 @@ export const parseExistingPrompt = (apiResponse: any): PromptType => { return { name: baseName, - model: parsedFrontmatter.model || "gpt-4o", + model: parsedFrontmatter.model || null, config: parsedFrontmatter.config, tools: parsedFrontmatter.tools, developerMessage: parsedBody.developerMessage, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx index 269f3b0af39..0239844851f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx @@ -20,7 +20,12 @@ import { useZodForm } from "@/lib/forms/useZodForm"; import { fetchClient } from "@/lib/http/api"; import { buildBody, settingsToForm, type DefaultInternalUserParams, type InternalUserSettings } from "./mapper"; -import { defaultUserSettingsSchema, EMPTY_TEAM_ROW, type DefaultUserSettingsFormValues } from "./schema"; +import { + defaultUserSettingsSchema, + EMPTY_TEAM_ROW, + type DefaultUserSettingsFormValues, + type DefaultUserSettingsSubmitValues, +} from "./schema"; const NO_RESET = "never"; @@ -63,7 +68,7 @@ interface RoleOption { description: string; } -type SettingsControl = Control; +type SettingsControl = Control; const TeamPickerField = ({ control, index }: { control: SettingsControl; index: number }) => { const [search, setSearch] = React.useState(""); @@ -233,7 +238,7 @@ const SettingsForm = ({ initialValues, roleOptions, updateSettings, onCancel, on const { isDirty } = form.formState; const mutation = useMutation({ - mutationFn: (values: DefaultUserSettingsFormValues) => updateSettings(buildBody(values)), + mutationFn: (values: DefaultUserSettingsSubmitValues) => updateSettings(buildBody(values)), onSuccess: (_result, values) => { toast.success("Default user settings updated successfully"); queryClient.invalidateQueries({ queryKey: SETTINGS_QUERY_KEY }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts index e8b350332c8..f4a7f001480 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from "vitest"; import { buildBody, settingsToForm } from "./mapper"; -import type { DefaultUserSettingsFormValues } from "./schema"; +import type { DefaultUserSettingsSubmitValues } from "./schema"; const CONFIGURED_SETTINGS = { user_role: "internal_user", @@ -59,13 +59,13 @@ describe("settingsToForm", () => { it("degrades an unrecognisable team entry to a blank row instead of throwing", () => { expect(settingsToForm({ teams: [{ max_budget_in_team: 5 }, 7] }).teams).toStrictEqual([ - { team_id: "", max_budget_in_team: "", user_role: "user" }, - { team_id: "", max_budget_in_team: "", user_role: "user" }, + { team_id: null, max_budget_in_team: "", user_role: "user" }, + { team_id: null, max_budget_in_team: "", user_role: "user" }, ]); }); }); -const formValues = (overrides: Partial = {}): DefaultUserSettingsFormValues => ({ +const formValues = (overrides: Partial = {}): DefaultUserSettingsSubmitValues => ({ user_role: "internal_user", max_budget: "100", budget_duration: "30d", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts index 50365081afb..ac528542ae2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts @@ -2,7 +2,12 @@ import { z } from "zod/v4"; import type { components } from "@/lib/http/schema"; -import { EMPTY_TEAM_ROW, type DefaultTeamRowValues, type DefaultUserSettingsFormValues } from "./schema"; +import { + EMPTY_TEAM_ROW, + type DefaultTeamRowValues, + type DefaultUserSettingsFormValues, + type DefaultUserSettingsSubmitValues, +} from "./schema"; export type InternalUserSettings = components["schemas"]["InternalUserSettingsResponse"]; export type DefaultInternalUserParams = components["schemas"]["DefaultInternalUserParams"]; @@ -60,13 +65,13 @@ const textOrNull = (raw: string): string | null => (raw.trim() === "" ? null : r const listOrNull = (items: readonly T[]): T[] | null => (items.length === 0 ? null : [...items]); -const toTeamBody = (team: DefaultTeamRowValues): DefaultTeamBody => ({ +const toTeamBody = (team: DefaultUserSettingsSubmitValues["teams"][number]): DefaultTeamBody => ({ team_id: team.team_id, max_budget_in_team: numberOrNull(team.max_budget_in_team), user_role: team.user_role, }); -export const buildBody = (values: DefaultUserSettingsFormValues): DefaultInternalUserParams => ({ +export const buildBody = (values: DefaultUserSettingsSubmitValues): DefaultInternalUserParams => ({ user_role: asDefaultUserRole(values.user_role), max_budget: numberOrNull(values.max_budget), budget_duration: textOrNull(values.budget_duration), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts index 7309e3745da..cd66c3addc8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts @@ -10,14 +10,17 @@ const amountOrEmpty = z ); const defaultTeamRowSchema = z.object({ - team_id: z.string().min(1, "Select a team"), + team_id: z + .string() + .nullable() + .pipe(z.string({ error: "Select a team" }).min(1, "Select a team")), max_budget_in_team: amountOrEmpty, user_role: z.enum(["user", "admin"]), }); -export type DefaultTeamRowValues = z.output; +export type DefaultTeamRowValues = z.input; -export const EMPTY_TEAM_ROW: DefaultTeamRowValues = { team_id: "", max_budget_in_team: "", user_role: "user" }; +export const EMPTY_TEAM_ROW: DefaultTeamRowValues = { team_id: null, max_budget_in_team: "", user_role: "user" }; const defaultUserSettingsShape = { user_role: z.string(), @@ -41,4 +44,5 @@ export const defaultUserSettingsSchema = z.object(defaultUserSettingsShape).supe ); }); -export type DefaultUserSettingsFormValues = z.output; +export type DefaultUserSettingsFormValues = z.input; +export type DefaultUserSettingsSubmitValues = z.output; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx index 875df74652c..8fa76a64af6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx @@ -192,7 +192,7 @@ export function UsersTable({ set("user_role", value)} + onValueChange={(value) => set("user_role", value ?? undefined)} placeholder="Select a role…" emptyText="No roles found" /> @@ -201,7 +201,7 @@ export function UsersTable({ set("team", value)} + onValueChange={(value) => set("team", value ?? undefined)} placeholder="Select a team…" emptyText="No teams found" /> diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.tsx index 0f7c356b8cc..83a72e50f5c 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.tsx @@ -59,7 +59,7 @@ interface UISettings { interface CreateUserFormValues { user_email?: string; user_role: string; - team_id?: string; + team_id?: string | null; organization_ids?: string[]; metadata?: string; send_invite_email: boolean; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx index 390684d3739..f78790dea88 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx @@ -115,7 +115,7 @@ export default function MCPSemanticFilterSettings({ accessToken }: MCPSemanticFi // Test section state const [testQuery, setTestQuery] = useState(""); - const [testModel, setTestModel] = useState("gpt-4o"); + const [testModel, setTestModel] = useState("gpt-4o"); const [testResult, setTestResult] = useState(null); const [testError, setTestError] = useState(null); const [isTesting, setIsTesting] = useState(false); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterTestPanel.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterTestPanel.tsx index 5dc661391ac..82c699d9983 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterTestPanel.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterTestPanel.tsx @@ -11,8 +11,8 @@ interface MCPSemanticFilterTestPanelProps { accessToken: string | null; testQuery: string; setTestQuery: (value: string) => void; - testModel: string; - setTestModel: (value: string) => void; + testModel: string | null; + setTestModel: (value: string | null) => void; isTesting: boolean; onTest: () => void; filterEnabled: boolean; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/semanticFilterTestUtils.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/semanticFilterTestUtils.ts index d3e585e8052..c41b081da88 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/semanticFilterTestUtils.ts +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/semanticFilterTestUtils.ts @@ -32,7 +32,7 @@ export const runSemanticFilterTest = async ({ setTestError, }: { accessToken: string; - testModel: string; + testModel: string | null; testQuery: string; setIsTesting: (value: boolean) => void; setTestResult: (result: TestResult | null) => void; @@ -68,12 +68,12 @@ export const runSemanticFilterTest = async ({ } }; -export const getCurlCommand = (testModel: string, testQuery: string) => +export const getCurlCommand = (testModel: string | null, testQuery: string) => `curl --location 'http://localhost:4000/v1/responses' \\ --header 'Content-Type: application/json' \\ --header 'Authorization: Bearer sk-1234' \\ --data '{ - "model": "${testModel}", + "model": "${testModel ?? "YOUR_MODEL"}", "input": [ { "role": "user", diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.test.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.test.ts index 2f9032537c1..6d19c17c577 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.test.ts +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.test.ts @@ -51,6 +51,7 @@ describe("parseCoreTools", () => { describe("formToPayload", () => { it("sends null for a cleared embedding model so the proxy returns to keyword matching", () => { expect(formToPayload({ ...DEFAULT_FORM_VALUES, embedding_model: " " })).toEqual(KEYWORD_PAYLOAD); + expect(formToPayload({ ...DEFAULT_FORM_VALUES, embedding_model: null })).toEqual(KEYWORD_PAYLOAD); }); it("clamps top_k into the range the proxy accepts and lists core tools", () => { diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.ts index f10bfd249d7..d3535fe91a3 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.ts +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.ts @@ -1,7 +1,7 @@ import type { MCPToolSearchSettings } from "@/app/(dashboard)/hooks/mcpToolSearchSettings/useMCPToolSearchSettings"; export interface ToolSearchFormValues { - embedding_model: string; + embedding_model: string | null; top_k: number; similarity_threshold: number; core_tools_text: string; @@ -11,7 +11,7 @@ export const TOP_K_MIN = 1; export const TOP_K_MAX = 100; export const DEFAULT_FORM_VALUES: ToolSearchFormValues = { - embedding_model: "", + embedding_model: null, top_k: 5, similarity_threshold: 0, core_tools_text: "", @@ -42,7 +42,7 @@ export const storedValuesToForm = (values: Record): ToolSearchF }); export const formToPayload = (form: ToolSearchFormValues): MCPToolSearchSettings => ({ - embedding_model: form.embedding_model.trim() === "" ? null : form.embedding_model.trim(), + embedding_model: form.embedding_model?.trim() || null, top_k: clampTopK(form.top_k), similarity_threshold: form.similarity_threshold, core_tools: parseCoreTools(form.core_tools_text), diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx index a82e0fdfa88..3897a9d9fa2 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx @@ -32,12 +32,8 @@ export function FallbackGroupConfig({ // Filter available options for fallbacks (exclude primary only, allow already selected to be shown for deselection) const availableFallbackOptions = availableModels.filter((m) => m !== group.primaryModel); - const handlePrimaryChange = (value: string) => { - let newFallbacks = [...group.fallbackModels]; - // Remove from fallbacks if it was there - if (newFallbacks.includes(value)) { - newFallbacks = newFallbacks.filter((m) => m !== value); - } + const handlePrimaryChange = (value: string | null) => { + const newFallbacks = group.fallbackModels.filter((model) => model !== value); onChange({ ...group, primaryModel: value, @@ -76,7 +72,7 @@ export function FallbackGroupConfig({ ({ label: m, value: m }))} - value={group.primaryModel ?? ""} + value={group.primaryModel} onValueChange={handlePrimaryChange} placeholder="Select primary model" emptyText="No models found" diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index d2d18139893..f0d21cc5350 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -328,11 +328,11 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser }; const selectCreateTeamOrganization = ( - next: string, + next: string | null, currentOrganizationId: string | null, onChange: (organizationId: string | null) => void, ) => { - const nextOrganizationId = next === "" ? null : next; + const nextOrganizationId = next; if (nextOrganizationId === currentOrganizationId) return; onChange(nextOrganizationId); form.setValue("models", []); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx index 3fb19e522f3..bde1c7724c3 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx @@ -203,7 +203,7 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet set("org_id", value)} + onValueChange={(value) => set("org_id", value ?? undefined)} placeholder="Select an organization…" emptyText="No organizations found" /> diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index f28ae27b6bc..d0367cb4d5b 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -311,7 +311,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { set("team_id", value)} + onValueChange={(value) => set("team_id", value ?? undefined)} placeholder="Select a team…" emptyText="No teams found" /> @@ -320,7 +320,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { set("org_id", value)} + onValueChange={(value) => set("org_id", value ?? undefined)} placeholder="Select an organization…" emptyText="No organizations found" /> diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx rename to ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx index 3b879da2d68..536bc231008 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx @@ -164,7 +164,7 @@ const createTestProps = (userRole = "proxy_admin", userId = "user-1", isTeamAdmi handleOk: vi.fn().mockResolvedValue(true), setSelectedProvider: vi.fn(), setProviderModelsFn: vi.fn(), - getPlaceholder: vi.fn((provider: Providers) => `Enter ${provider} model name`), + getPlaceholder: vi.fn((provider: string) => `Enter ${provider} model name`), setShowAdvancedSettings: vi.fn(), selectedProvider: Providers.OpenAI, providerModels: ["gpt-4", "gpt-3.5-turbo"], diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 92951b68fd5..30b4911da3e 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -25,7 +25,6 @@ import { } from "../common_components/MountedFormField"; import type { Team } from "../key_team_helpers/key_list"; import { type CredentialItem, type ProviderCreateInfo, modelAvailableCall } from "../networking"; -import { Providers } from "../provider_info_helpers"; import { ProviderLogo } from "../molecules/models/ProviderLogo"; import AccessGroupTagsCombobox from "./AccessGroupTagsCombobox"; import AdvancedSettings from "./advanced_settings"; @@ -42,11 +41,11 @@ interface AddModelFormProps { registry: MountRegistry; mountedValues: () => MountedFormValues; handleOk: () => Promise; - selectedProvider: Providers; - setSelectedProvider: (provider: Providers) => void; + selectedProvider: string | null; + setSelectedProvider: (provider: string | null) => void; providerModels: string[]; - setProviderModelsFn: (provider: Providers) => void; - getPlaceholder: (provider: Providers) => string; + setProviderModelsFn: (provider: string | null) => void; + getPlaceholder: (provider: string) => string; showAdvancedSettings: boolean; setShowAdvancedSettings: (show: boolean) => void; teams: Team[] | null; @@ -140,7 +139,7 @@ const AddModelForm: React.FC = ({ [credentials], ); - const applyProviderSelection = (provider: Providers) => { + const applyProviderSelection = (provider: string | null) => { setSelectedProvider(provider); setProviderModelsFn(provider); form.setValue("model", []); @@ -227,10 +226,10 @@ const AddModelForm: React.FC = ({ options={providerOptions} emptyText={providerMetadataErrorText ?? "No providers found"} placeholder={isProviderMetadataLoading ? "Loading providers..." : "Select a provider"} - value={(control.value as string | undefined) ?? ""} + value={typeof control.value === "string" ? control.value : null} onValueChange={(value) => { control.onChange(value); - applyProviderSelection(value as Providers); + applyProviderSelection(value); }} /> )} diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 594c6d022ce..93ccf561387 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -328,7 +328,8 @@ const ClassificationMethodConfig: React.FC = ({ onChange(nextValue); }; - const handleClassifierModelChange = (model: string) => { + const handleClassifierModelChange = (model: string | null) => { + if (model === null) return; if (model === value.classifier_llm_config?.model) return; const { reasoning_effort: _reasoningEffort, ...classifierLlmConfig } = value.classifier_llm_config ?? { model: "", diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 97d6739aa22..febcde269f7 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -632,7 +632,7 @@ const ComplexityRouterConfig: React.FC = ({ // Clearing the select drops the key entirely rather than storing "", so an emptied pin reads as // "track the tiers" everywhere downstream instead of as a blank model name. - const handleDefaultModelChange = (model: string | undefined) => { + const handleDefaultModelChange = (model: string | null | undefined) => { onChange({ ...value, default_model: model || undefined }); }; diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx index c1817918f60..e67583febf3 100644 --- a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -42,8 +42,8 @@ const CompressionControls: React.FC = ({ value, onChan
onRoutingChange(value === "" ? undefined : value)} + value={routing} + onValueChange={(value) => onRoutingChange(value ?? undefined)} placeholder="Inherit from the request's own compression guardrails" emptyText="No compression guardrails found" aria-label="Routing decision compression" @@ -74,8 +74,8 @@ const CompressionControls: React.FC = ({ value, onChan
onModelChange(value === "" ? undefined : value)} + value={model} + onValueChange={(value) => onModelChange(value ?? undefined)} placeholder="None (no compression)" emptyText="No compression guardrails found" aria-label="Model call compression" diff --git a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.tsx b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.integration.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.tsx rename to ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.integration.test.tsx index ecd5abfa451..495fd207266 100644 --- a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.integration.test.tsx @@ -47,7 +47,7 @@ describe("RouterConfigBuilder", () => { expect(onChange).toHaveBeenCalledWith({ routes: [ expect.objectContaining({ - name: "", + name: null, utterances: [], description: "", score_threshold: 0.5, diff --git a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts new file mode 100644 index 00000000000..d681c911c46 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts @@ -0,0 +1,10 @@ +import { describe, expect, it } from "vitest"; +import { serializeRouterConfig } from "./RouterConfigBuilder"; + +describe("serializeRouterConfig", () => { + it("rejects a cleared route model and preserves selected route settings", () => { + expect(() => serializeRouterConfig({ routes: [{ name: null }] })).toThrow("Please select a model for every route"); + const config = { routes: [{ name: "model-silver", utterances: [], description: "", score_threshold: 0 }] }; + expect(JSON.parse(serializeRouterConfig(config))).toEqual(config); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx index 6dfb9f32091..9713d18872b 100644 --- a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx +++ b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx @@ -15,7 +15,7 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/comp interface Route { id: string; - model: string; + model: string | null; utterances: string[]; description: string; score_threshold: number; @@ -23,21 +23,28 @@ interface Route { interface SavedRoute { id?: string; - name?: string; - model?: string; + name?: string | null; + model?: string | null; utterances?: string[]; description?: string; score_threshold?: number; } -interface RouterConfig { +export interface RouterConfig { routes?: SavedRoute[]; } +export function serializeRouterConfig(config: RouterConfig | null): string { + if (config?.routes?.some((route) => !(route.name ?? route.model))) { + throw new Error("Please select a model for every route"); + } + return JSON.stringify(config); +} + interface RouterConfigBuilderProps { modelInfo: ModelGroup[]; - value?: RouterConfig; - onChange?: (config: any) => void; + value?: RouterConfig | null; + onChange?: (config: RouterConfig) => void; } interface UtteranceInputProps { @@ -136,7 +143,7 @@ const RouterConfigBuilder: React.FC = ({ modelInfo, va routeIds.push(id); return { id, - model: route.name || route.model || "", + model: route.name || route.model || null, utterances: route.utterances || [], description: route.description || "", score_threshold: route.score_threshold ?? 0.5, @@ -165,7 +172,7 @@ const RouterConfigBuilder: React.FC = ({ modelInfo, va const newRouteId = `route-${Date.now()}`; const updatedRoutes = [ ...routes, - { id: newRouteId, model: "", utterances: [], description: "", score_threshold: 0.5 }, + { id: newRouteId, model: null, utterances: [], description: "", score_threshold: 0.5 }, ]; setRoutes(updatedRoutes); updateConfig(updatedRoutes); diff --git a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx index 5684d32a319..0a5a9a5bfb6 100644 --- a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx +++ b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx @@ -61,7 +61,9 @@ const SemanticKeywordMatching: React.FC = ({ { + if (model !== null) onEmbeddingModelChange(model); + }} placeholder="Select an embedding model" emptyText="No embedding models found" aria-label="Embedding model" diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index d7c4a671317..632f4427a82 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -150,7 +150,10 @@ export const getSubmitBlockedReason = ( const autoRouterSchema = (requiresTeamScope: boolean) => z.object({ auto_router_name: z.string().min(1, "Auto router name is required"), - team_id: requiresTeamScope ? z.string().min(1, "Please select a team to continue") : z.string(), + team_id: z + .string() + .nullable() + .refine((teamId) => !requiresTeamScope || Boolean(teamId), "Please select a team to continue"), model_access_group: z.array(z.string()).optional(), }); @@ -158,12 +161,12 @@ type AddAutoRouterFormValues = z.infer>; const EMPTY_FORM_VALUES: AddAutoRouterFormValues = { auto_router_name: "", - team_id: "", + team_id: null, model_access_group: undefined, }; -const teamScopePayload = (requiresTeamScope: boolean, teamId: string): { team_id?: string } => - requiresTeamScope ? { team_id: teamId } : {}; +const teamScopePayload = (requiresTeamScope: boolean, teamId: string | null): { team_id?: string } => + requiresTeamScope && teamId ? { team_id: teamId } : {}; const BlockedReasonTooltip: React.FC<{ reason: string | null; children: React.ReactElement }> = ({ reason, @@ -458,7 +461,7 @@ const AddAutoRouterTab: React.FC = ({ const serverVerdict = await validateAutoRouterConfig( accessToken, complexityRouterConfigPayload as unknown as Record, - requiresTeamScope ? form.getValues("team_id") : undefined, + requiresTeamScope ? form.getValues("team_id") ?? undefined : undefined, ); const dryRunError = dryRunRejection(serverVerdict); if (dryRunError) { @@ -630,9 +633,7 @@ const AddAutoRouterTab: React.FC = ({ "Select the team this auto router belongs to. Only keys for this team will be able to call it.", )} > - {({ id, value, onChange }) => ( - onChange(next ?? "")} /> - )} + {({ id, value, onChange }) => } )} @@ -776,7 +777,7 @@ const AddAutoRouterTab: React.FC = ({ config={buildComplexityRouterConfig(complexityRouterConfigParams)} defaultModel={resolveComplexityDefaultModel(complexityRouterConfig, complexityRouterConfig.default_model)} routerName={watchedName} - teamId={requiresTeamScope ? watchedTeamId : undefined} + teamId={requiresTeamScope ? watchedTeamId ?? undefined : undefined} /> )} diff --git a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx index bfa2df8f54d..8cf4b077f39 100644 --- a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx +++ b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx @@ -8,9 +8,9 @@ import { MountedFormField, type MountedFormValues } from "../common_components/M import { Providers } from "../provider_info_helpers"; interface LiteLLMModelNameFieldProps { - selectedProvider: Providers; + selectedProvider: string | null; providerModels: string[]; - getPlaceholder: (provider: Providers) => string; + getPlaceholder: (provider: string) => string; } const LiteLLMModelNameField: React.FC = ({ @@ -123,7 +123,7 @@ const LiteLLMModelNameField: React.FC = ({ id={control.id} value={(control.value as string | undefined) ?? ""} onBlur={control.onBlur} - placeholder={getPlaceholder(selectedProvider)} + placeholder={selectedProvider === null ? "Select a provider first" : getPlaceholder(selectedProvider)} onChange={(event) => { control.onChange(event); if (selectedProvider === Providers.Azure) { @@ -147,7 +147,7 @@ const LiteLLMModelNameField: React.FC = ({ value: "custom", }, { - label: `All ${selectedProvider} Models (Wildcard)`, + label: `All ${selectedProvider ?? "provider"} Models (Wildcard)`, value: "all-wildcard", }, ...providerModels.map((model) => ({ @@ -163,7 +163,7 @@ const LiteLLMModelNameField: React.FC = ({ value={(control.value as string | undefined) ?? ""} onChange={control.onChange} onBlur={control.onBlur} - placeholder={getPlaceholder(selectedProvider)} + placeholder={selectedProvider === null ? "Select a provider first" : getPlaceholder(selectedProvider)} /> ) } diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index e1a47e9792d..4add7b918f7 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -18,7 +18,7 @@ import { provider_map, Providers } from "../provider_info_helpers"; import { labelWithHint } from "@/components/shared/form/LabelWithHint"; interface ProviderSpecificFieldsProps { - selectedProvider: Providers; + selectedProvider: string | null; } const readTextFile = (file: File, onLoaded: (contents: string) => void) => { @@ -168,6 +168,7 @@ const ProviderSpecificFields: React.FC = ({ selecte }, [cacheEntries]); const allFields = React.useMemo(() => { + if (selectedProvider === null) return []; // First try to resolve from the in-memory cache. We support both the // enum/display-name form and the raw provider slug (e.g. "petals"). const cachedFields = diff --git a/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx b/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx index c6e24d3b470..384fa44cfde 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx @@ -27,8 +27,13 @@ const ModelAliasManager: React.FC = ({ showExampleConfig = true, }) => { const [aliases, setAliases] = useState([]); - const [newAlias, setNewAlias] = useState({ aliasName: "", targetModel: "" }); - const [editingAlias, setEditingAlias] = useState(null); + const [newAlias, setNewAlias] = useState<{ aliasName: string; targetModel: string | null }>({ + aliasName: "", + targetModel: null, + }); + const [editingAlias, setEditingAlias] = useState< + (Omit & { targetModel: string | null }) | null + >(null); const aliasNameId = useId(); useEffect(() => { @@ -61,7 +66,7 @@ const ModelAliasManager: React.FC = ({ const updatedAliases = [...aliases, newAliasObj]; setAliases(updatedAliases); - setNewAlias({ aliasName: "", targetModel: "" }); + setNewAlias({ aliasName: "", targetModel: null }); // Convert array back to object format and notify parent const aliasObject: { [key: string]: string } = {}; @@ -94,7 +99,8 @@ const ModelAliasManager: React.FC = ({ return; } - const updatedAliases = aliases.map((alias) => (alias.id === editingAlias.id ? editingAlias : alias)); + const savedAlias: AliasItem = { ...editingAlias, targetModel: editingAlias.targetModel }; + const updatedAliases = aliases.map((alias) => (alias.id === savedAlias.id ? savedAlias : alias)); setAliases(updatedAliases); setEditingAlias(null); diff --git a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx index 6a6807de5c0..15b6b1644b6 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx @@ -9,9 +9,9 @@ const MODEL_SELECT_DEBOUNCE_MS = 500; interface ModelSelectorProps { accessToken: string; - value?: string; + value?: string | null; placeholder?: string; - onChange?: (value: string) => void; + onChange?: (value: string | null) => void; disabled?: boolean; style?: React.CSSProperties; className?: string; @@ -30,12 +30,12 @@ const ModelSelector: React.FC = ({ showLabel = true, labelText = "Select Model", }) => { - const [selectedModel, setSelectedModel] = useState(value); + const [selectedModel, setSelectedModel] = useState(value ?? null); const [showCustomModelInput, setShowCustomModelInput] = useState(false); const [modelInfo, setModelInfo] = useState([]); useEffect(() => { - setSelectedModel(value); + setSelectedModel(value ?? null); }, [value]); useEffect(() => { @@ -56,13 +56,13 @@ const ModelSelector: React.FC = ({ loadModels(); }, [accessToken]); - const onModelChange = (value: string) => { + const onModelChange = (value: string | null) => { if (value === "custom") { setShowCustomModelInput(true); - setSelectedModel(undefined); + setSelectedModel(null); } else { setShowCustomModelInput(false); - setSelectedModel(value); + setSelectedModel(value ?? null); if (onChange) { onChange(value); } @@ -71,7 +71,7 @@ const ModelSelector: React.FC = ({ const debouncedSelect = useDebouncedCallback( (value: string) => { - setSelectedModel(value); + setSelectedModel(value ?? null); onChange?.(value); }, { wait: MODEL_SELECT_DEBOUNCE_MS }, diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx index 663028b2d92..f2f1504fed5 100644 --- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx @@ -4,7 +4,7 @@ import { Organization } from "../networking"; interface OrganizationDropdownProps { organizations?: Organization[] | null; - value?: string; + value?: string | null; onChange?: (value: string | null) => void; disabled?: boolean; loading?: boolean; @@ -32,7 +32,7 @@ const OrganizationDropdown: React.FC = ({ sublabel: org.organization_id, }))} value={value} - onValueChange={(organizationId) => onChange?.(organizationId || null)} + onValueChange={(organizationId) => onChange?.(organizationId)} placeholder={placeholder} emptyText={loading ? "Loading organizations…" : "No organizations found"} disabled={disabled} diff --git a/ui/litellm-dashboard/src/components/common_components/ProjectDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/ProjectDropdown.tsx index 0b5a84f9589..f27e1a683dd 100644 --- a/ui/litellm-dashboard/src/components/common_components/ProjectDropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ProjectDropdown.tsx @@ -4,8 +4,8 @@ import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects"; interface ProjectDropdownProps { projects?: ProjectResponse[] | null; - value?: string; - onChange?: (value: string) => void; + value?: string | null; + onChange?: (value: string | null) => void; disabled?: boolean; loading?: boolean; /** When set, only show projects belonging to this team */ diff --git a/ui/litellm-dashboard/src/components/common_components/UserDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/UserDropdown.tsx index eef1a8233bc..c80f654bbd1 100644 --- a/ui/litellm-dashboard/src/components/common_components/UserDropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/UserDropdown.tsx @@ -47,8 +47,8 @@ const UserDropdown: React.FC = ({ value, onChange, disabled,
onChange(next === "" ? null : next)} + value={value} + onValueChange={onChange} onSearchChange={setSearch} onLoadMore={fetchNextPage} hasNextPage={hasNextPage} diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx index 7d385c2a3f7..e10a8d1e562 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx @@ -4,7 +4,7 @@ import { useInfiniteTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { Team } from "../key_team_helpers/key_list"; interface TeamDropdownProps { - value?: string; + value?: string | null; onChange?: (value: string | null) => void; /** Callback with the full Team object (or null on clear). */ onTeamSelect?: (team: Team | null) => void; @@ -46,8 +46,8 @@ const TeamDropdown: React.FC = ({ return result; }, [data]); - const handleChange = (teamId: string) => { - onChange?.(teamId || null); + const handleChange = (teamId: string | null) => { + onChange?.(teamId); if (onTeamSelect) { onTeamSelect(teamId ? teams.find((t) => t.team_id === teamId) ?? null : null); } @@ -61,7 +61,7 @@ const TeamDropdown: React.FC = ({ value: team.team_id, sublabel: team.team_id, }))} - value={value || undefined} + value={value} onValueChange={handleChange} onSearchChange={setSearch} onLoadMore={fetchNextPage} diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.integration.test.tsx similarity index 94% rename from ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx rename to ui/litellm-dashboard/src/components/common_components/user_search_modal.integration.test.tsx index 7ca2b530235..5769d8c2f88 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.integration.test.tsx @@ -1,4 +1,5 @@ -import { act, fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor, within } from "@testing-library/react"; +import { renderWithProviders as render } from "../../../tests/test-utils"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import UserSearchModal from "./user_search_modal"; @@ -101,15 +102,18 @@ describe("UserSearchModal submit payload", () => { await user.click(await screen.findByRole("option", { name: "picked@example.com" })); }; - it("submits every registered field, with the untouched identity fields undefined", async () => { + it("should block submission without a selected user and after clearing the paired identity", async () => { const { user, onSubmit } = setup(); + expect(save()).toBeDisabled(); + await searchByEmail(user, "pick"); + await user.click(screen.getAllByRole("button", { name: "Clear" })[0]); + expect(screen.getByLabelText("Email")).toHaveValue(""); + expect(screen.getByLabelText("User ID")).toHaveValue(""); + expect(save()).toBeDisabled(); await user.click(save()); - await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); - const values = onSubmit.mock.calls[0][0]; - expect(Object.keys(values).sort()).toEqual(["role", "user_email", "user_id"]); - expect(values).toStrictEqual({ user_email: undefined, user_id: undefined, role: "user" }); + expect(onSubmit).not.toHaveBeenCalled(); }); it("carries the picked user's email and id into the payload", async () => { @@ -130,6 +134,7 @@ describe("UserSearchModal submit payload", () => { const { onSubmit } = setup(); const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + await searchByEmail(user, "pick"); await user.click(screen.getByLabelText("Member Role")); await user.click(await screen.findByRole("option", { name: /^admin/ })); await user.click(save()); @@ -184,6 +189,7 @@ describe("UserSearchModal submit payload", () => { it("does not submit on Enter in any field, while the button still does", async () => { const { user, onSubmit } = setup(); + await searchByEmail(user, "pick"); await user.click(getEmailSearchInput()); await user.keyboard("{Enter}"); await user.click(screen.getByLabelText("User ID")); diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx index 20c45cd0d2e..04ac412397f 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx @@ -31,8 +31,8 @@ interface Role { } interface FormValues { - user_email: string | undefined; - user_id: string | undefined; + user_email: string | null | undefined; + user_id: string | null | undefined; role: string; } @@ -66,6 +66,8 @@ const UserSearchModal: React.FC = ({ }) => { const emptyValues: FormValues = { user_email: undefined, user_id: undefined, role: defaultRole }; const form = useForm({ defaultValues: emptyValues }); + const selectedUserId = form.watch("user_id"); + const selectedUserEmail = form.watch("user_email"); const [userOptions, setUserOptions] = useState([]); const [loading, setLoading] = useState(false); const [selectedField, setSelectedField] = useState<"user_email" | "user_id">("user_email"); @@ -143,19 +145,29 @@ const UserSearchModal: React.FC = ({ const renderUserSearch = ( fieldName: "user_email" | "user_id", placeholder: string, - controlProps: { id: string; value: string | undefined; onChange: (value: string | undefined) => void }, + controlProps: { + id: string; + value: string | null | undefined; + onChange: (value: string | null | undefined) => void; + }, testId?: string, ) => { const items = selectedField === fieldName ? userOptions : []; + const handleValueChange = (value: string | null) => { + if (value === null) { + form.setValue("user_email", null); + form.setValue("user_id", null); + return; + } + controlProps.onChange(value); + handleSelect(items.find((option) => option.value === value) ?? null); + }; return (
{ - controlProps.onChange(value === "" ? undefined : value); - handleSelect(items.find((option) => option.value === value) ?? null); - }} + onValueChange={handleValueChange} onSearchChange={(query: string) => handleSearch(query, fieldName)} autoHighlight="always" isLoading={loading} @@ -226,7 +238,7 @@ const UserSearchModal: React.FC = ({
- diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 8a62e86e842..86f18ee9b12 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -13,7 +13,7 @@ import AccessGroupTagsCombobox from "../add_model/AccessGroupTagsCombobox"; import ModelChoiceCombobox, { type ModelChoice } from "../add_model/ModelChoiceCombobox"; import { modelAvailableCall, modelPatchUpdateCall, validateAutoRouterConfig } from "../networking"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; -import RouterConfigBuilder from "../add_model/RouterConfigBuilder"; +import RouterConfigBuilder, { type RouterConfig, serializeRouterConfig } from "../add_model/RouterConfigBuilder"; import { hydrateTierModelParams } from "../add_model/complexity_router_tiers"; import { type ActiveTierSet, @@ -452,7 +452,7 @@ const EditAutoRouterModal: React.FC = ({ const [modelInfo, setModelInfo] = useState([]); const [showValidationErrors, setShowValidationErrors] = useState(false); const [editingTiers, setEditingTiers] = useState(false); - const [routerConfig, setRouterConfig] = useState(null); + const [routerConfig, setRouterConfig] = useState(null); const [customTechnicalKeywords, setCustomTechnicalKeywords] = useState([]); const [keywordTierRules, setKeywordTierRules] = useState([]); const [escalationKeywords, setEscalationKeywords] = useState([]); @@ -706,7 +706,7 @@ const EditAutoRouterModal: React.FC = ({ // Prepare the updated litellm_params const updatedLitellmParams = { ...modelData.litellm_params, - auto_router_config: JSON.stringify(routerConfig), + auto_router_config: serializeRouterConfig(routerConfig), auto_router_default_model: values.auto_router_default_model, auto_router_embedding_model: values.auto_router_embedding_model || undefined, }; @@ -745,7 +745,7 @@ const EditAutoRouterModal: React.FC = ({ })(); } catch (error) { console.error("Error updating auto router:", error); - toast.fromError("Failed to update auto router configuration"); + toast.fromError(error); } finally { setLoading(false); } diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx index c0c12356c6f..dbc270fa62e 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx @@ -96,10 +96,10 @@ export function BudgetFallbacksEditor({ value, onChange, availableModels }: Budg ({ label: m, value: m }))} - value={entry.primaryModel ?? ""} + value={entry.primaryModel} onValueChange={(v) => { const newFallbacks = entry.fallbackModels.filter((m) => m !== v); - updateEntry(entry.id, { primaryModel: v === "" ? null : v, fallbackModels: newFallbacks }); + updateEntry(entry.id, { primaryModel: v, fallbackModels: newFallbacks }); }} placeholder="Select model" emptyText="No models found" diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx index ef5a3e15b52..0fab5555343 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx @@ -153,8 +153,8 @@ export function ModelMaxBudgetEditor({ ({ label: model, value: model }))} - value={entry.model ?? ""} - onValueChange={(model) => updateEntry(entry.id, { model: model === "" ? null : model })} + value={entry.model} + onValueChange={(model) => updateEntry(entry.id, { model })} placeholder="Select model" emptyText="No models found" disabled={!premiumUser} diff --git a/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx index 8108cc1c20e..fdc02b4fc41 100644 --- a/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx @@ -42,7 +42,7 @@ export default function CredentialModal({ existingCredential = null, }: CredentialModalProps) { const isEdit = mode === "edit"; - const [selectedProvider, setSelectedProvider] = useState( + const [selectedProvider, setSelectedProvider] = useState( (existingCredential?.credential_info.custom_llm_provider as Providers) ?? Providers.OpenAI, ); @@ -110,7 +110,7 @@ export default function CredentialModal({ {(control) => ( { control.onChange(value); - resetCredentialFormOnProviderChange(formAdapter, value as Providers, setSelectedProvider); + resetCredentialFormOnProviderChange(formAdapter, value, setSelectedProvider); }} /> )} diff --git a/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts b/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts index 7190c539c81..e7b1e861811 100644 --- a/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts +++ b/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts @@ -1,5 +1,3 @@ -import { Providers } from "../provider_info_helpers"; - interface CredentialFormAdapter { getFieldValue: (field: string) => unknown; resetFields: () => void; @@ -25,8 +23,8 @@ interface CredentialFormAdapter { */ export function resetCredentialFormOnProviderChange( form: CredentialFormAdapter, - newProvider: Providers, - setSelectedProvider: (p: Providers) => void, + newProvider: string | null, + setSelectedProvider: (p: string | null) => void, ): void { const preservedName = form.getFieldValue("credential_name"); form.resetFields(); diff --git a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts index 37a973d5def..311e9357825 100644 --- a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts +++ b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts @@ -202,6 +202,8 @@ export const buildKeyCreatePayload = (input: KeyCreateInput): KeyPayloadResult = endpoint: input.keyOwner === "service_account" ? "service_account" : "standard", payload: { ...withoutKeys(values, dropped), + ...(values.organization_id === null && { organization_id: undefined }), + ...(values.project_id === null && { project_id: undefined }), ...(input.keyOwner === "you" && { user_id: input.userID }), ...(input.keyOwner === "agent" && { agent_id: input.selectedAgentId }), ...(input.autoRotationEnabled && { auto_rotate: true, rotation_interval: input.rotationInterval }), diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 2bf39bf4cd7..0d5d9f5ec8d 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -16,7 +16,7 @@ const state = vi.hoisted(() => ({ can: {} as Record, uiSettings: {} as Record, tags: {} as Record, - teams: [] as { team_id: string; team_alias: string; models: string[] }[], + teams: [] as { team_id: string; team_alias: string; models: string[]; organization_id?: string }[], organizations: [] as { organization_id: string; organization_alias: string }[], accessGroups: [] as { access_group_id: string; access_group_name: string }[], projects: [] as { project_id: string; project_alias: string; team_id?: string; models?: string[] }[], @@ -806,6 +806,36 @@ describe("CreateKey", () => { expect((await createdPayload()).organization_id).toBe("org-1"); }); + it("discards the old project and team when the organization changes", async () => { + state.uiSettings = { enable_projects_ui: true }; + state.organizations = [ + { organization_id: "scope-silver", organization_alias: "Silver" }, + { organization_id: "scope-copper", organization_alias: "Copper" }, + ]; + state.teams = [{ team_id: "group-maple", team_alias: "Maple", organization_id: "scope-silver", models: [] }]; + state.projects = [{ project_id: "project-orbit", project_alias: "Orbit", team_id: "group-maple", models: [] }]; + await openModal({ teams: state.teams as Team[] }); + await nameTheKey(); + await userEvent.click(await screen.findByLabelText("Organization")); + await userEvent.click(await screen.findByRole("option", { name: /Silver/ })); + await userEvent.click(await screen.findByLabelText("Project")); + await userEvent.click(await screen.findByRole("option", { name: /Orbit/ })); + await waitFor(() => expect(screen.getByLabelText("Team")).toHaveValue("Maple")); + expect(screen.getByLabelText("Team")).toBeDisabled(); + + await userEvent.click(screen.getByLabelText("Organization")); + await userEvent.click(await screen.findByRole("option", { name: /Copper/ })); + expect(screen.getByLabelText("Project")).toHaveValue(""); + expect(screen.getByLabelText("Team")).toHaveValue(""); + expect(screen.getByLabelText("Team")).toBeEnabled(); + await submit(); + + const payload = JSON.parse(JSON.stringify(await createdPayload())); + expect(payload.organization_id).toBe("scope-copper"); + expect(payload.team_id).toBeNull(); + expect(payload).not.toHaveProperty("project_id"); + }); + it("drops organization_id when the chosen organization is cleared again", async () => { state.organizations = [{ organization_id: "org-1", organization_alias: "Engineering" }]; await openModal(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 831cb0cf6a6..82580fd4667 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -592,35 +592,35 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp }; const changeOrganization = (write: FieldWrite) => (orgId: string | null) => { - write(orgId ?? undefined); + write(orgId); setSelectedOrganizationId(orgId); // Clear team and project when org changes setSelectedCreateKeyTeam(null); setSelectedProjectId(null); - form.setValue("team_id", undefined); - form.setValue("project_id", undefined); + form.setValue("team_id", null); + form.setValue("project_id", null); }; const selectTeam = (team: Team | null) => { setSelectedCreateKeyTeam(team); setSelectedProjectId(null); - form.setValue("project_id", undefined); + form.setValue("project_id", null); // Auto-populate org from team for non-admin users if (team?.organization_id) { setSelectedOrganizationId(team.organization_id); form.setValue("organization_id", team.organization_id); } else if (!team) { setSelectedOrganizationId(null); - form.setValue("organization_id", undefined); + form.setValue("organization_id", null); } }; - const changeProject = (write: FieldWrite) => (projectId: string) => { + const changeProject = (write: FieldWrite) => (projectId: string | null) => { write(projectId); if (!projectId) { setSelectedProjectId(null); setSelectedCreateKeyTeam(null); - form.setValue("team_id", undefined); + form.setValue("team_id", null); return; } setSelectedProjectId(projectId); @@ -756,8 +756,8 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp inputId="create-key-agent" placeholder="Select an agent" emptyText="No agents found" - value={selectedAgentId ?? undefined} - onValueChange={(value) => setSelectedAgentId(value === "" ? null : value)} + value={selectedAgentId} + onValueChange={setSelectedAgentId} options={agentsList.map((a) => ({ label: a.agent_name || a.agent_id, value: a.agent_id, @@ -783,7 +783,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp {(control) => ( = ({ team, teams, data, addKey, autoOp {(control) => ( = ({ team, teams, data, addKey, autoOp {(control) => ( { return providerPlaceholderMap[resolvedProvider] ?? "gpt-3.5-turbo"; }; -export const getProviderModels = (provider: Providers, modelMap: any): Array => { +export const getProviderModels = (provider: string, modelMap: any): Array => { let providerKey = provider; let custom_llm_provider = provider_map[providerKey]; diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx similarity index 89% rename from ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx rename to ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx index 04c8d3e9012..2b948ca8420 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; @@ -51,7 +51,7 @@ describe("PaginatedSearchSelect", () => { const onSearchChange = vi.fn(); function Controlled() { - const [value, setValue] = useState(""); + const [value, setValue] = useState(null); return ( { expect(onSearchChange).not.toHaveBeenCalled(); }); - it("still reports a cleared input so the unfiltered page comes back", async () => { + it("should keep a cleared selection empty after a late page arrives and reset the query", async () => { const user = userEvent.setup(); const onSearchChange = vi.fn(); - renderSelect({ onSearchChange, value: "alias-alpha" }); - - await user.click(document.querySelector('[data-slot="combobox-clear"]') as HTMLElement); - - await waitFor(() => expect(onSearchChange).toHaveBeenCalledWith("")); + const onValueChange = vi.fn(); + function Controlled({ options }: { options: SearchSelectOption[] }) { + const [value, setValue] = useState("alias-alpha"); + return ( + { + setValue(next); + onValueChange(next); + }} + /> + ); + } + const { rerender } = render(); + await user.click(screen.getByRole("button", { name: "Clear" })); + expect(onValueChange).toHaveBeenLastCalledWith(null); + rerender( ({ ...option }))} />); + await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("")); + expect(screen.getByRole("combobox")).toHaveValue(""); + expect(screen.queryByRole("button", { name: "Clear" })).not.toBeInTheDocument(); + await chooseSelectOption(user, screen.getByRole("combobox"), "alias-beta"); + expect(onValueChange).toHaveBeenLastCalledWith("alias-beta"); }); it("requests the next page once the list is scrolled near the bottom", async () => { @@ -147,7 +166,7 @@ describe("PaginatedSearchSelect", () => { function ServerBacked() { const [search, setSearch] = useState(""); - const [value, setValue] = useState("alias-alpha"); + const [value, setValue] = useState("alias-alpha"); const freshlyBuiltOptions = OPTIONS.filter((option) => option.label.includes(search)).map((option) => ({ ...option, })); @@ -236,7 +255,7 @@ describe("PaginatedSearchSelect", () => { function Refetching() { const [options, setOptions] = useState([{ label: "Beta Team", value: "team-2" }]); - const [value, setValue] = useState(""); + const [value, setValue] = useState(null); return ( <> { function ServerBacked() { const [search, setSearch] = useState(""); - const [value, setValue] = useState(""); + const [value, setValue] = useState(null); return ( option.label.includes(search))} diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx index 4b7ef7401c7..1314afdf2bf 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx @@ -17,8 +17,8 @@ import { usePaginatedCombobox } from "./usePaginatedCombobox"; interface PaginatedSearchSelectProps { options: SearchSelectOption[]; - value?: string; - onValueChange: (value: string) => void; + value?: string | null; + onValueChange: (value: string | null) => void; onSearchChange: (query: string) => void; onLoadMore?: () => void; hasNextPage?: boolean; @@ -86,7 +86,7 @@ export function PaginatedSearchSelect({ }; const selected = useMemo(() => { - if (value === undefined || value === "") return null; + if (value == null || value === "") return null; return ( options.find((option) => option.value === value) ?? (pickedOption?.value === value ? pickedOption : { label: value, value }) @@ -118,7 +118,7 @@ export function PaginatedSearchSelect({ inputValue={typedQuery ?? selected?.label ?? ""} onValueChange={(item: SearchSelectOption | null) => { setPickedOption(item); - onValueChange(item?.value ?? ""); + onValueChange(item?.value ?? null); }} onInputValueChange={(next, eventDetails) => handleTypedInput(next, eventDetails.reason)} onOpenChange={(nextOpen, eventDetails) => handleOpenChange(nextOpen, eventDetails.reason)} @@ -139,7 +139,7 @@ export function PaginatedSearchSelect({ onKeyDown={snapshotWholeSelection} onPaste={snapshotWholeSelection} placeholder={placeholder} - showClear={value !== undefined && value !== ""} + showClear={value != null && value !== ""} className={`w-full ${className ?? ""}`} /> diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx similarity index 76% rename from ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx rename to ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx index 62a510268dd..c981010dff9 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx @@ -1,4 +1,5 @@ -import { fireEvent, render, screen } from "@testing-library/react"; +import { fireEvent, renderWithProviders as render, screen } from "../../../tests/test-utils"; +import { useState } from "react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; @@ -36,11 +37,30 @@ describe("SearchSelect", () => { expect(screen.getByRole("combobox")).toHaveValue("Growth"); }); - it("shows a clear control only when a value is selected", () => { - const { rerender } = render(); - expect(document.querySelector('[data-slot="combobox-clear"]')).toBeNull(); - rerender(); - expect(document.querySelector('[data-slot="combobox-clear"]')).not.toBeNull(); + it("should clear to null and allow selecting again through the real control", async () => { + const onValueChange = vi.fn(); + const user = userEvent.setup(); + function Controlled() { + const [value, setValue] = useState(null); + return ( + { + setValue(next); + onValueChange(next); + }} + /> + ); + } + render(); + expect(screen.queryByRole("button", { name: "Clear" })).not.toBeInTheDocument(); + await chooseSelectOption(user, screen.getByRole("combobox"), "Growth"); + await user.click(screen.getByRole("button", { name: "Clear" })); + expect(onValueChange).toHaveBeenLastCalledWith(null); + expect(screen.getByRole("combobox")).toHaveValue(""); + await chooseSelectOption(user, screen.getByRole("combobox"), "Data Team"); + expect(onValueChange).toHaveBeenLastCalledWith("team-3"); }); it("filters the options client-side as you type", async () => { diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx index 53a5371d72d..21d38e9458c 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx @@ -21,8 +21,8 @@ export interface SearchSelectOption { interface SearchSelectProps { options: SearchSelectOption[]; - value?: string; - onValueChange: (value: string) => void; + value?: string | null; + onValueChange: (value: string | null) => void; placeholder?: string; emptyText?: string; disabled?: boolean; @@ -51,9 +51,7 @@ export function SearchSelect({ "aria-label": ariaLabel, }: SearchSelectProps) { const selected = - value === undefined || value === "" - ? null - : options.find((option) => option.value === value) ?? { label: value, value }; + value == null || value === "" ? null : options.find((option) => option.value === value) ?? { label: value, value }; const items = selected !== null && !options.some((option) => option.value === selected.value) ? [selected, ...options] : options; @@ -61,7 +59,7 @@ export function SearchSelect({ onValueChange(item?.value ?? "")} + onValueChange={(item: SearchSelectOption | null) => onValueChange(item?.value ?? null)} isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value} itemToStringLabel={(item: SearchSelectOption) => item.label} filter={matchesQuery} diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 988e2e082aa..8e3a17c2622 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -333,7 +333,10 @@ const teamUpdateFieldsSchema = z.object({ modelLimits: z .array( z.object({ - model: z.string().min(1, "Missing model"), + model: z + .string() + .nullable() + .refine((model) => Boolean(model), "Missing model"), tpm: z.number().nullish(), rpm: z.number().nullish(), }), @@ -1879,7 +1882,7 @@ const TeamInfoView: React.FC = ({ onChange(next === "" ? null : next)} + onValueChange={onChange} options={userOrganizations.map((org) => ({ value: org.organization_id ?? "", label: org.organization_alias || org.organization_id || "", diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx rename to ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index b2c2a381b42..7c2226f0369 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -1483,11 +1483,11 @@ describe("KeyEditView", () => { }); }); - it("submits organization_id as null after the organization is cleared", async () => { + it("clears the organization and its dependent team in the update payload", async () => { const onSubmit = vi.fn().mockResolvedValue(undefined); renderWithProviders( {}} onSubmit={onSubmit} accessToken="" @@ -1504,9 +1504,33 @@ describe("KeyEditView", () => { await userEvent.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => { - expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ organization_id: null })); + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ organization_id: null, team_id: null })); }); - expect(JSON.parse(JSON.stringify(onSubmit.mock.calls[0][0]))).toHaveProperty("organization_id", null); + expect(JSON.parse(JSON.stringify(onSubmit.mock.calls[0][0]))).toMatchObject({ + organization_id: null, + team_id: null, + }); + }); + + it("keeps project key relationships locked and omits unsupported project updates", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderWithProviders( + {}} + onSubmit={onSubmit} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + expect(await screen.findByRole("combobox", { name: "Organization" })).toBeDisabled(); + expect(screen.getByRole("combobox", { name: "Team ID" })).toBeDisabled(); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); + expect(onSubmit.mock.calls[0][0]).toMatchObject({ organization_id: "org-1", team_id: "group-maple" }); + expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("project_id"); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index e464dfe5008..f0d63f9671d 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -305,7 +305,7 @@ export function KeyEditView({ const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => { setField(orgId); setSelectedOrganizationId(orgId); - form.setValue("team_id", undefined); + form.setValue("team_id", null); }; const handleTeamChange = (setField: (value: string | null) => void, teamId: string | null) => { @@ -316,7 +316,7 @@ export function KeyEditView({ form.setValue("organization_id", selectedTeam.organization_id); } else if (!teamId) { setSelectedOrganizationId(null); - form.setValue("organization_id", undefined); + form.setValue("organization_id", null); } }; @@ -769,14 +769,15 @@ export function KeyEditView({ "Organization", "The organization this key belongs to. Selecting an organization filters the available teams.", )} + description={hasProject ? "Organization is locked because this key belongs to a project" : undefined} > {({ value, onChange, id }) => ( handleOrganizationChange(onChange, orgId)} /> )} @@ -786,15 +787,13 @@ export function KeyEditView({ control={form.control} name="team_id" label="Team ID" - description={ - enableProjectsUI && hasProject ? "Team is locked because this key belongs to a project" : undefined - } + description={hasProject ? "Team is locked because this key belongs to a project" : undefined} > {({ value, onChange, id }) => (