From 8a4942340e20550c60a8f24723fb69dc18e38da8 Mon Sep 17 00:00:00 2001 From: Dhruv Yadav Date: Mon, 6 Jul 2026 16:46:32 +0000 Subject: [PATCH 01/62] fix(streaming): use provider-reported usage cost for OpenRouter streams Port of #16162 by @dhruvyad onto litellm_internal_staging. OpenRouter sends a usage chunk (including a provider-reported cost field) after the finish_reason chunk. Previously the stream handler raised StopIteration on the first post-finish chunk, so that usage/cost never reached the assembled response and cost tracking fell back to token-based estimates. Carry usage.cost through chunk accumulation, preserve stripped usage in _hidden_params, and propagate the provider cost into _hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] so the cost calculator uses it. --- .../streaming_chunk_builder_utils.py | 40 ++++- .../litellm_core_utils/streaming_handler.py | 53 +++++- .../streaming_chunk_builder_utils.py | 1 + litellm/types/utils.py | 10 +- .../test_streaming_chunk_builder_utils.py | 36 ++++ .../test_streaming_handler.py | 163 ++++++++++++++++++ 6 files changed, 283 insertions(+), 20 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 38bc68f2f78..fbb11fa7a1f 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -467,6 +467,7 @@ class ChunkProcessor: cache_read_input_tokens: Optional[int] = None completion_tokens_details: Optional[CompletionTokensDetails] = None prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None + cost: Optional[float] = None if "prompt_tokens" in usage_chunk: prompt_tokens = usage_chunk.get("prompt_tokens", 0) or 0 @@ -476,6 +477,8 @@ class ChunkProcessor: cache_creation_input_tokens = usage_chunk.get("cache_creation_input_tokens") if "cache_read_input_tokens" in usage_chunk: cache_read_input_tokens = usage_chunk.get("cache_read_input_tokens") + if "cost" in usage_chunk: + cost = usage_chunk.get("cost") if hasattr(usage_chunk, "completion_tokens_details"): if isinstance(usage_chunk.completion_tokens_details, dict): completion_tokens_details = CompletionTokensDetails(**usage_chunk.completion_tokens_details) @@ -494,6 +497,7 @@ class ChunkProcessor: "cache_read_input_tokens": cache_read_input_tokens, "completion_tokens_details": completion_tokens_details, "prompt_tokens_details": prompt_tokens_details, + "cost": cost, } def count_reasoning_tokens(self, response: ModelResponse) -> Optional[int]: @@ -512,6 +516,22 @@ class ChunkProcessor: return reasoning_tokens + @staticmethod + def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse) -> Usage | None: + usage_chunk: Usage | dict[str, Any] | None = None + if hasattr(chunk, "usage") and chunk.usage is not None: + usage_chunk = chunk.usage + elif "usage" in chunk: + usage_chunk = chunk["usage"] + elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr( + chunk, "_hidden_params" + ): + usage_chunk = chunk._hidden_params.get("usage", None) + + if isinstance(usage_chunk, dict): + return Usage(**usage_chunk) + return usage_chunk + def _calculate_usage_per_chunk( self, chunks: List[Union[Dict[str, Any], ModelResponse]], @@ -548,18 +568,12 @@ class ChunkProcessor: # is last-wins, so without preserving this separately the 1h breakdown is # lost and 1h cache writes get billed at the 5m rate. cache_creation_token_details: Optional[CacheCreationTokenDetails] = None + cost: Optional[float] = None + for chunk in chunks: - usage_chunk: Optional[Usage] = None - if "usage" in chunk: - usage_chunk = chunk["usage"] - elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr( - chunk, "_hidden_params" - ): - usage_chunk = chunk._hidden_params.get("usage", None) + usage_chunk = self._extract_usage_chunk(chunk) if usage_chunk is not None: - if isinstance(usage_chunk, dict): - usage_chunk = Usage(**usage_chunk) usage_chunk_dict = self._usage_chunk_calculation_helper(usage_chunk) if usage_chunk_dict["prompt_tokens"] is not None and usage_chunk_dict["prompt_tokens"] > 0: prompt_tokens = usage_chunk_dict["prompt_tokens"] @@ -610,6 +624,9 @@ class ChunkProcessor: prompt_tokens_details, cache_creation_token_details ) + if usage_chunk_dict["cost"] is not None: + cost = usage_chunk_dict["cost"] + prompt_tokens_details = self._attach_cache_creation_token_details( prompt_tokens_details, cache_creation_token_details ) @@ -629,6 +646,7 @@ class ChunkProcessor: web_search_requests=web_search_requests, completion_tokens_details=completion_tokens_details, prompt_tokens_details=prompt_tokens_details, + cost=cost, ) @staticmethod @@ -727,6 +745,7 @@ class ChunkProcessor: prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = calculated_usage_per_chunk[ "prompt_tokens_details" ] + cost: Optional[float] = calculated_usage_per_chunk["cost"] try: returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages) @@ -784,6 +803,9 @@ class ChunkProcessor: else: returned_usage.prompt_tokens_details.web_search_requests = web_search_requests + if cost is not None: + setattr(returned_usage, "cost", cost) + # Return a new usage object with the new values returned_usage = Usage(**returned_usage.model_dump()) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 128ba0bf3ab..ba5bf11ebd6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -962,10 +962,11 @@ class CustomStreamWrapper: if self.custom_llm_provider == "bedrock" and "trace" in model_response: return model_response - # Default - return StopIteration - if hasattr(model_response, "usage"): - self.chunks.append(model_response) - raise StopIteration + # Don't raise StopIteration here - some providers (like OpenRouter) + # send usage/cost data in chunks after the finish_reason chunk + if hasattr(model_response, "usage") and model_response.usage is not None: + return model_response + return # flush any remaining holding chunk if len(self.holding_chunk) > 0: if model_response.choices[0].delta.content is None: @@ -1474,12 +1475,16 @@ class CustomStreamWrapper: self.tool_call = True + if hasattr(chunk, "usage") and chunk.usage is not None: + model_response.usage = chunk.usage + ## RETURN ARG - return self.return_processed_chunk_logic( + result = self.return_processed_chunk_logic( completion_obj=completion_obj, model_response=model_response, # type: ignore response_obj=response_obj, ) + return result except StopIteration: raise StopIteration @@ -1686,6 +1691,21 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = "tool_calls" return model_response + @staticmethod + def _propagate_usage_cost_to_hidden_params( + response: "ModelResponse", + ) -> None: + """ + If the assembled response carries a provider-reported cost on + usage.cost, copy it into _hidden_params so litellm's cost + calculator uses it instead of a token-based estimate. + """ + _usage = getattr(response, "usage", None) + if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None: + if "additional_headers" not in response._hidden_params: + response._hidden_params["additional_headers"] = {} + response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(_usage.cost) + def __next__(self) -> "ModelResponseStream": cache_hit = False if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response": @@ -1741,6 +1761,10 @@ class CustomStreamWrapper: # hasattr(response, "usage") is always True — must check # `is not None` to avoid running this path on every chunk. if getattr(response, "usage", None) is not None: + usage_to_preserve = response.usage + if usage_to_preserve: + response._hidden_params["usage"] = usage_to_preserve + obj_dict = response.model_dump() if "usage" in obj_dict: @@ -1789,6 +1813,8 @@ class CustomStreamWrapper: response = self.model_response_creator() if complete_streaming_response is not None: + self._propagate_usage_cost_to_hidden_params(complete_streaming_response) + setattr( response, "usage", @@ -1999,6 +2025,8 @@ class CustomStreamWrapper: response = self.model_response_creator() if complete_streaming_response is not None: + self._propagate_usage_cost_to_hidden_params(complete_streaming_response) + setattr( response, "usage", @@ -2228,12 +2256,16 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: """Assume most recent usage chunk has total usage uptil then.""" prompt_tokens: int = 0 completion_tokens: int = 0 + latest_usage_chunk = None + for chunk in chunks: if "usage" in chunk and chunk["usage"] is not None: - if "prompt_tokens" in chunk["usage"]: - prompt_tokens = chunk["usage"].get("prompt_tokens", 0) or 0 - if "completion_tokens" in chunk["usage"]: - completion_tokens = chunk["usage"].get("completion_tokens", 0) or 0 + usage = chunk["usage"] + latest_usage_chunk = usage + if "prompt_tokens" in usage: + prompt_tokens = usage.get("prompt_tokens", 0) or 0 + if "completion_tokens" in usage: + completion_tokens = usage.get("completion_tokens", 0) or 0 returned_usage_chunk = Usage( prompt_tokens=prompt_tokens, @@ -2241,6 +2273,9 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: total_tokens=prompt_tokens + completion_tokens, ) + if latest_usage_chunk and hasattr(latest_usage_chunk, "cost") and latest_usage_chunk.cost is not None: + returned_usage_chunk.cost = latest_usage_chunk.cost + return returned_usage_chunk diff --git a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py index a1f89dac5cf..c9f9d4e6baa 100644 --- a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py @@ -14,3 +14,4 @@ class UsagePerChunk(TypedDict): web_search_requests: Optional[int] completion_tokens_details: Optional[CompletionTokensDetails] prompt_tokens_details: Optional[PromptTokensDetailsWrapper] + cost: Optional[float] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 380621f88a8..d6e39815506 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1795,14 +1795,17 @@ class ModelResponseStream(ModelResponseBase): else: created = created + usage_to_set = None if "usage" in kwargs and kwargs["usage"] is not None: if isinstance(kwargs["usage"], dict): - kwargs["usage"] = Usage(**kwargs["usage"]) + usage_to_set = Usage(**kwargs["usage"]) + kwargs["usage"] = usage_to_set elif isinstance(kwargs["usage"], BaseModel): dump = ( kwargs["usage"].model_dump() if hasattr(kwargs["usage"], "model_dump") else kwargs["usage"].dict() ) - kwargs["usage"] = Usage(**dump) + usage_to_set = Usage(**dump) + kwargs["usage"] = usage_to_set kwargs["id"] = id kwargs["created"] = created @@ -1811,6 +1814,9 @@ class ModelResponseStream(ModelResponseBase): super().__init__(**kwargs) + if usage_to_set is not None: + self.usage = usage_to_set + def __contains__(self, key): # Define custom behavior for the 'in' operator return hasattr(self, key) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index ba95c90798e..be8c5a05601 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -956,3 +956,39 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks(): assert response.model_dump()["vertex_ai_grounding_metadata"] == [ {"webSearchQueries": ["test query"]} ] + + +def test_cost_field_in_usage_chunks(): + chunk1_usage = Usage(completion_tokens=1, prompt_tokens=10, total_tokens=11) + chunk1 = ModelResponseStream( + id="chatcmpl-1", + created=1745513206, + model="openrouter/claude", + choices=[ + StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi")) + ], + usage=chunk1_usage, + ) + + chunk2_usage = Usage( + completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 + ) + chunk2 = ModelResponseStream( + id="chatcmpl-1", + created=1745513207, + model="openrouter/claude", + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) + ], + usage=chunk2_usage, + ) + + processor = ChunkProcessor(chunks=[chunk1, chunk2]) + usage = processor.calculate_usage( + chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi" + ) + + assert hasattr(usage, "cost") + assert usage.cost == 0.00025 + assert usage.prompt_tokens == 10 + assert usage.completion_tokens == 5 diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index e430ce3b084..6108d853286 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1393,6 +1393,169 @@ def test_has_any_special_delta_attributes( assert result is False +def test_calculate_total_usage_with_cost(): + from litellm.litellm_core_utils.streaming_handler import calculate_total_usage + + chunk1_usage = Usage(completion_tokens=1, prompt_tokens=10, total_tokens=11) + chunk1 = ModelResponseStream( + id="test-1", + created=1745513206, + model="openrouter/test", + choices=[ + StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi")) + ], + usage=chunk1_usage, + ) + + chunk2_usage = Usage( + completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 + ) + chunk2 = ModelResponseStream( + id="test-1", + created=1745513207, + model="openrouter/test", + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) + ], + usage=chunk2_usage, + ) + + usage = calculate_total_usage([chunk1, chunk2]) + + assert hasattr(usage, "cost") + assert usage.cost == 0.00025 + assert usage.prompt_tokens == 10 + assert usage.completion_tokens == 5 + + +@pytest.mark.asyncio +async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Logging): + from litellm.utils import ModelResponseListIterator + + chunk1 = ModelResponseStream( + id="chatcmpl-or", + created=1742056047, + model="openrouter/claude", + choices=[ + StreamingChoices( + finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant") + ) + ], + usage=None, + ) + chunk2 = ModelResponseStream( + id="chatcmpl-or", + created=1742056048, + model="openrouter/claude", + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) + ], + usage=None, + ) + chunk3_usage = Usage( + completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 + ) + chunk3 = ModelResponseStream( + id="chatcmpl-or", + created=1742056049, + model="openrouter/claude", + choices=[ + StreamingChoices(finish_reason=None, index=0, delta=Delta(content="")) + ], + usage=chunk3_usage, + ) + + completion_stream = ModelResponseListIterator( + model_responses=[chunk1, chunk2, chunk3] + ) + response = CustomStreamWrapper( + completion_stream=completion_stream, + model="openrouter/claude", + custom_llm_provider="openrouter", + logging_obj=logging_obj, + stream_options={"include_usage": True}, + ) + + collected_chunks = [] + async for chunk in response: + collected_chunks.append(chunk) + + usage_chunks = [c for c in collected_chunks if hasattr(c, "usage") and c.usage] + assert len(usage_chunks) > 0 + assert hasattr(usage_chunks[-1].usage, "cost") + assert usage_chunks[-1].usage.cost == 0.00025 + + +def test_openrouter_streaming_cost_propagates_to_hidden_params(): + """ + Verify that provider-reported cost from usage.cost flows into + _hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] + on the complete streaming response, so litellm's cost calculator uses it. + """ + import litellm + + chunk1 = ModelResponseStream( + id="chatcmpl-or", + created=1742056047, + model="openrouter/claude", + choices=[ + StreamingChoices( + finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant") + ) + ], + usage=None, + ) + chunk2 = ModelResponseStream( + id="chatcmpl-or", + created=1742056048, + model="openrouter/claude", + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) + ], + usage=None, + ) + chunk3 = ModelResponseStream( + id="chatcmpl-or", + created=1742056049, + model="openrouter/claude", + choices=[ + StreamingChoices(finish_reason=None, index=0, delta=Delta(content="")) + ], + usage=Usage( + completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 + ), + ) + + # Build the complete response as stream_chunk_builder does + complete_response = litellm.stream_chunk_builder( + chunks=[chunk1, chunk2, chunk3], + messages=[{"role": "user", "content": "test"}], + ) + + assert complete_response is not None + assert hasattr(complete_response.usage, "cost") + assert complete_response.usage.cost == 0.00025 + + # Use the real propagation method from CustomStreamWrapper + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) + + assert "additional_headers" in complete_response._hidden_params + assert ( + complete_response._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] + == 0.00025 + ) + + # Verify the cost calculator would pick this up + from litellm.cost_calculator import get_response_cost_from_hidden_params + + provider_cost = get_response_cost_from_hidden_params( + complete_response._hidden_params + ) + assert provider_cost == 0.00025 + + def test_handle_special_delta_attributes( initialized_custom_stream_wrapper: CustomStreamWrapper, ): From 21c126464b8d4a6e21a0008a6d73c3cdf0a4f767 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 6 Jul 2026 17:03:18 +0000 Subject: [PATCH 02/62] fix(streaming): include ModelResponseStream in _extract_usage_chunk annotation --- litellm/litellm_core_utils/streaming_chunk_builder_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fbb11fa7a1f..d52d9849310 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -517,7 +517,7 @@ class ChunkProcessor: return reasoning_tokens @staticmethod - def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse) -> Usage | None: + def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None: usage_chunk: Usage | dict[str, Any] | None = None if hasattr(chunk, "usage") and chunk.usage is not None: usage_chunk = chunk.usage From 214130e82e8e4085a0f3085f3a0d41f96cf5ebd4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 6 Jul 2026 17:52:44 +0000 Subject: [PATCH 03/62] fix(streaming): preserve provider cost when usage chunk is a dict --- .../litellm_core_utils/streaming_handler.py | 10 +++++++-- .../test_streaming_handler.py | 21 +++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index ba5bf11ebd6..2d640f79d4a 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2273,8 +2273,14 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: total_tokens=prompt_tokens + completion_tokens, ) - if latest_usage_chunk and hasattr(latest_usage_chunk, "cost") and latest_usage_chunk.cost is not None: - returned_usage_chunk.cost = latest_usage_chunk.cost + if latest_usage_chunk is not None: + latest_cost = ( + latest_usage_chunk.get("cost") + if isinstance(latest_usage_chunk, dict) + else getattr(latest_usage_chunk, "cost", None) + ) + if latest_cost is not None: + returned_usage_chunk.cost = latest_cost return returned_usage_chunk diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 6108d853286..41c72c9593e 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1428,6 +1428,27 @@ def test_calculate_total_usage_with_cost(): assert usage.completion_tokens == 5 +def test_calculate_total_usage_with_dict_usage_cost(): + """Regression: dict-shaped `usage` with a `cost` key must still surface + provider cost even though `hasattr` on a dict does not consult its keys.""" + from litellm.litellm_core_utils.streaming_handler import calculate_total_usage + + chunk = { + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "cost": 0.00025, + } + } + + usage = calculate_total_usage([chunk]) + + assert usage.prompt_tokens == 10 + assert usage.completion_tokens == 5 + assert getattr(usage, "cost", None) == 0.00025 + + @pytest.mark.asyncio async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Logging): from litellm.utils import ModelResponseListIterator From 799a559871f2463bde4984794b71f5eb2133337d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 10 Jul 2026 16:49:19 -0700 Subject: [PATCH 04/62] fix(guardrails): show YAML-defined guardrails in the Guardrail Monitor The /guardrails/usage/{overview,detail,logs} endpoints resolved guardrails only from the litellm_guardrailstable Prisma table, so guardrails defined in config.yaml (which live only in IN_MEMORY_GUARDRAIL_HANDLER) were invisible: detail 404'd, overview omitted them or rendered them as Custom/Guardrail orphans, and logs missed their logical-name alias. Add config-owned accessors (list_config_guardrails, get_config_guardrail_by_id) to the in-memory handler and use them in the usage endpoints, mirroring the union/fallback already used by list_guardrails_v2 and get_guardrail_info. Also preserve guardrail_info when storing a config guardrail (type/description were dropped at initialize time) and read the Prisma-row / dict / LitellmParams shapes uniformly. Resolves LIT-2529 --- .../proxy/guardrails/guardrail_registry.py | 22 ++ litellm/proxy/guardrails/usage_endpoints.py | 69 +++-- .../guardrails/test_guardrail_registry.py | 46 +++- .../proxy/guardrails/test_init_guardrails.py | 35 ++- .../proxy/guardrails/test_usage_endpoints.py | 239 ++++++++++++++++++ 5 files changed, 366 insertions(+), 45 deletions(-) create mode 100644 tests/test_litellm/proxy/guardrails/test_usage_endpoints.py diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 8962073fe7a..47a2f112396 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -478,6 +478,7 @@ class InMemoryGuardrailHandler: guardrail_id=guardrail.get("guardrail_id"), guardrail_name=guardrail["guardrail_name"], litellm_params=litellm_params, + guardrail_info=guardrail.get("guardrail_info"), ) # store references to the guardrail in memory @@ -596,6 +597,27 @@ class InMemoryGuardrailHandler: """ return self._sources.get(guardrail_id) + def list_config_guardrails(self) -> List[Guardrail]: + """ + List in-memory guardrails owned by config.yaml. + + DB-sourced entries are excluded: a read surface that also queries the DB + would double-count live ones, and a DB-sourced entry that's missing from + the DB is stale (deleted on another pod, awaiting reconciliation here). + """ + return [g for gid, g in self.IN_MEMORY_GUARDRAILS.items() if self._sources.get(gid) == "config"] + + def get_config_guardrail_by_id(self, guardrail_id: str) -> Optional[Guardrail]: + """ + Get a config-owned in-memory guardrail by its ID, or None. + + Mirrors the fallback in get_guardrail_info: a DB-sourced in-memory entry + that missed the DB lookup is stale and must not be surfaced. + """ + if self._sources.get(guardrail_id) != "config": + return None + return self.IN_MEMORY_GUARDRAILS.get(guardrail_id) + def reconcile_db_guardrails(self, db_guardrail_ids: Set[str]) -> List[str]: """ Drop in-memory entries that originated from the DB but are no longer diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index e03bdbb95d2..c63526e4833 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -20,6 +20,7 @@ from litellm.repositories.table_repositories import ( SpendLogGuardrailIndexRepository, SpendLogsRepository, ) +from litellm.types.guardrails import LitellmParams router = APIRouter() @@ -137,10 +138,26 @@ def _chart_from_metrics(metrics: Any) -> List[Dict[str, Any]]: return [{"date": d, "passed": v["passed"], "blocked": v["blocked"]} for d, v in sorted(chart_by_date.items())] +def _get_guardrail_field(g: Any, field: str) -> Any: + """Read `field` off a guardrail whether it's a Prisma row (attr) or a dict/TypedDict (key).""" + if isinstance(g, dict): + return g.get(field) + return getattr(g, field, None) + + +def _to_dict(value: Any) -> Dict[str, Any]: + """Coerce a LitellmParams / guardrail_info value into a plain dict.""" + if isinstance(value, LitellmParams): + return value.model_dump(exclude_none=True) + if isinstance(value, dict): + return value + return {} + + def _get_guardrail_attrs(g: Any) -> tuple[Any, str]: """Get (guardrail_id, display_name) from guardrail - handles Prisma model or dict.""" - gid = getattr(g, "guardrail_id", None) or (g.get("guardrail_id") if isinstance(g, dict) else None) - name = getattr(g, "guardrail_name", None) or (g.get("guardrail_name") if isinstance(g, dict) else None) + gid = _get_guardrail_field(g, "guardrail_id") + name = _get_guardrail_field(g, "guardrail_name") return gid, (name or gid or "") @@ -163,9 +180,9 @@ def _guardrail_overview_rows( break req, blocked = a["requests"], a["blocked"] fail_rate = (100.0 * blocked / req) if req else 0.0 - litellm_params = (g.litellm_params or {}) if isinstance(g.litellm_params, dict) else {} + litellm_params = _to_dict(_get_guardrail_field(g, "litellm_params")) provider = str(litellm_params.get("guardrail", "Unknown")) - guardrail_info = (g.guardrail_info or {}) if isinstance(g.guardrail_info, dict) else {} + guardrail_info = _to_dict(_get_guardrail_field(g, "guardrail_info")) gtype = str(guardrail_info.get("type", "Guardrail")) prev_fail = 0.0 for k in lookup_keys: @@ -262,9 +279,15 @@ async def guardrails_usage_overview( end = end_date or now.strftime("%Y-%m-%d") start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + try: - # Guardrails from DB - guardrails = await GuardrailsRepository(prisma_client).table.find_many() + db_guardrails = await GuardrailsRepository(prisma_client).table.find_many() + seen_ids = {gid for g in db_guardrails if (gid := _get_guardrail_field(g, "guardrail_id")) is not None} + config_guardrails = [ + g for g in IN_MEMORY_GUARDRAIL_HANDLER.list_config_guardrails() if g.get("guardrail_id") not in seen_ids + ] + guardrails: List[Any] = [*db_guardrails, *config_guardrails] # Daily metrics in range metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( @@ -321,16 +344,18 @@ async def guardrails_usage_detail( end = end_date or now.strftime("%Y-%m-%d") start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") - guardrail = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) - if not guardrail: + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + guardrail: Any = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) + if guardrail is None: + guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_config_guardrail_by_id(guardrail_id=guardrail_id) + if guardrail is None: from fastapi import HTTPException raise HTTPException(status_code=404, detail="Guardrail not found") # Metrics are keyed by logical name (from spend log metadata), not UUID - logical_id = getattr(guardrail, "guardrail_name", None) or ( - guardrail.get("guardrail_name") if isinstance(guardrail, dict) else None - ) + logical_id = _get_guardrail_field(guardrail, "guardrail_name") metric_ids = [i for i in (logical_id, guardrail_id) if i] metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( @@ -367,17 +392,9 @@ async def guardrails_usage_detail( {"date": d, "passed": v["passed"], "blocked": v["blocked"], "score": None} for d, v in sorted(ts_by_date.items()) ] - _litellm_params = getattr(guardrail, "litellm_params", None) or ( - guardrail.get("litellm_params") if isinstance(guardrail, dict) else None - ) - litellm_params = _litellm_params if isinstance(_litellm_params, dict) else {} - _guardrail_info = getattr(guardrail, "guardrail_info", None) or ( - guardrail.get("guardrail_info") if isinstance(guardrail, dict) else None - ) - guardrail_info = _guardrail_info if isinstance(_guardrail_info, dict) else {} - _guardrail_name = getattr(guardrail, "guardrail_name", None) or ( - guardrail.get("guardrail_name") if isinstance(guardrail, dict) else None - ) + litellm_params = _to_dict(_get_guardrail_field(guardrail, "litellm_params")) + guardrail_info = _to_dict(_get_guardrail_field(guardrail, "guardrail_info")) + _guardrail_name = _get_guardrail_field(guardrail, "guardrail_name") return UsageDetailResponse( guardrail_id=guardrail_id, @@ -548,11 +565,15 @@ async def guardrails_usage_logs( # Query by both so we match regardless of which was written. effective_guardrail_ids: List[str] = [guardrail_id] if guardrail_id else [] if guardrail_id: - guardrail = await GuardrailsRepository(prisma_client).table.find_unique( + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + guardrail: Any = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) + if guardrail is None: + guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_config_guardrail_by_id(guardrail_id=guardrail_id) if guardrail: - logical_name = getattr(guardrail, "guardrail_name", None) + logical_name = _get_guardrail_field(guardrail, "guardrail_name") if logical_name and logical_name not in effective_guardrail_ids: effective_guardrail_ids.append(logical_name) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 0ef9ad857f9..26feddadf79 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -44,9 +44,7 @@ def test_update_in_memory_guardrail(): "123", Guardrail( guardrail_name="test-guardrail", - litellm_params=LitellmParams( - guardrail="test-guardrail", mode="pre_call", default_on=True - ), + litellm_params=LitellmParams(guardrail="test-guardrail", mode="pre_call", default_on=True), ), ) @@ -56,10 +54,7 @@ def test_update_in_memory_guardrail(): ) is True ) - assert ( - handler.guardrail_id_to_custom_guardrail["123"].event_hook - is GuardrailEventHooks.pre_call - ) + assert handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail: @@ -135,6 +130,34 @@ def test_delete_in_memory_guardrail_clears_source_marker(): assert handler.get_source("a") is None +def test_list_config_guardrails_excludes_db_sourced(): + """LIT-2529: read surfaces union DB rows with config guardrails; db-sourced + in-memory entries would double-count (or resurrect stale ones), so exclude them.""" + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg", name="config-one") + handler._sources["cfg"] = "config" + handler.IN_MEMORY_GUARDRAILS["db"] = _make_guardrail("db", name="db-one") + handler._sources["db"] = "db" + + config_guardrails = handler.list_config_guardrails() + + assert [g["guardrail_id"] for g in config_guardrails] == ["cfg"] + + +def test_get_config_guardrail_by_id_returns_config_only(): + """LIT-2529: the detail/logs fallback must return config-owned guardrails and + treat a db-sourced (stale) or missing id as a miss.""" + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg", name="config-one") + handler._sources["cfg"] = "config" + handler.IN_MEMORY_GUARDRAILS["db"] = _make_guardrail("db", name="db-one") + handler._sources["db"] = "db" + + assert handler.get_config_guardrail_by_id("cfg")["guardrail_name"] == "config-one" + assert handler.get_config_guardrail_by_id("db") is None + assert handler.get_config_guardrail_by_id("missing") is None + + def test_initialize_guardrail_early_return_updates_source_marker(): """ When initialize_guardrail is called for a guardrail that already exists @@ -152,9 +175,7 @@ def test_initialize_guardrail_early_return_updates_source_marker(): g = Guardrail( guardrail_id="collide", guardrail_name="bedrock", - litellm_params=LitellmParams( - guardrail="bedrock", mode="pre_call", default_on=False - ), + litellm_params=LitellmParams(guardrail="bedrock", mode="pre_call", default_on=False), ) handler.initialize_guardrail(guardrail=g, source="config") @@ -331,10 +352,7 @@ def test_repeated_db_sync_does_not_accumulate_runner_instances(): def distinct_runner_instances() -> int: seen = set() for callback in litellm.logging_callback_manager._get_all_callbacks(): - if ( - isinstance(callback, CustomGuardrail) - and getattr(callback, "guardrail_name", None) == name - ): + if isinstance(callback, CustomGuardrail) and getattr(callback, "guardrail_name", None) == name: seen.add(id(callback)) return len(seen) diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index a511229942a..83593c20110 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -5,9 +5,7 @@ from unittest.mock import MagicMock, patch import pytest -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -36,8 +34,31 @@ def test_initialize_presidio_guardrail(): ) assert result["guardrail_name"] == "test_presidio_guardrail" - assert ( - result["litellm_params"].guardrail - == SupportedGuardrailIntegrations.PRESIDIO.value - ) + assert result["litellm_params"].guardrail == SupportedGuardrailIntegrations.PRESIDIO.value assert result["litellm_params"].mode == "pre_call" + + +def test_initialize_guardrail_preserves_guardrail_info(): + """ + Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the + stored in-memory Guardrail. Dropping it left the Guardrail Monitor's usage + endpoints unable to render type/description for YAML-defined guardrails. + """ + test_guardrail = { + "guardrail_name": "test_presidio_with_info", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value, + "mode": "pre_call", + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + }, + "guardrail_info": {"type": "PII", "description": "masks PII"}, + } + + guardrail_handler = InMemoryGuardrailHandler() + result = guardrail_handler.initialize_guardrail(guardrail=test_guardrail) + + assert result is not None + assert result["guardrail_info"] == {"type": "PII", "description": "masks PII"} + stored = guardrail_handler.IN_MEMORY_GUARDRAILS[result["guardrail_id"]] + assert stored["guardrail_info"] == {"type": "PII", "description": "masks PII"} diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py new file mode 100644 index 00000000000..bf7b1b3b238 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -0,0 +1,239 @@ +""" +Tests for the /guardrails/usage/* endpoints backing the dashboard Guardrail Monitor. + +Regression (LIT-2529): guardrails defined in config.yaml live only in +IN_MEMORY_GUARDRAIL_HANDLER, so the monitor's overview/detail/logs endpoints — +which read the litellm_guardrailstable Prisma table — could not see them: +detail 404'd, overview omitted them (or rendered them as Custom/Guardrail +orphans), and logs missed their logical-name alias. +""" + +import os +import sys +from datetime import datetime +from typing import Any, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.proxy.guardrails.usage_endpoints import ( + guardrails_usage_detail, + guardrails_usage_logs, + guardrails_usage_overview, +) +from litellm.types.guardrails import Guardrail, LitellmParams + +ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) +# Query() defaults don't resolve to None when the handler is called directly. +START, END = "2026-04-20", "2026-04-27" + + +def _config_handler(*guardrails: Guardrail) -> InMemoryGuardrailHandler: + """A real handler seeded with config-sourced YAML guardrails (no callbacks).""" + handler = InMemoryGuardrailHandler() + for g in guardrails: + gid = g["guardrail_id"] + handler.IN_MEMORY_GUARDRAILS[gid] = g + handler._sources[gid] = "config" + return handler + + +def _yaml_guardrail( + guardrail_id: str = "yaml-1", + name: str = "yaml-pii", + provider: str = "presidio", + info: Optional[dict] = None, +) -> Guardrail: + return Guardrail( + guardrail_id=guardrail_id, + guardrail_name=name, + litellm_params=LitellmParams(guardrail=provider, mode="pre_call"), + guardrail_info=info if info is not None else {"type": "PII", "description": "yaml-defined"}, + ) + + +def _db_row(guardrail_id: str = "db-1", name: str = "db-guard", provider: str = "aim") -> Any: + """A Prisma-style row: attribute access, litellm_params/guardrail_info as plain dicts.""" + row = MagicMock(spec=["guardrail_id", "guardrail_name", "litellm_params", "guardrail_info"]) + row.guardrail_id = guardrail_id + row.guardrail_name = name + row.litellm_params = {"guardrail": provider, "mode": "pre_call"} + row.guardrail_info = {"type": "ContentSafety", "description": "db-defined"} + return row + + +def _metric(guardrail_id: str, date: str = "2026-04-25", requests: int = 10, passed: int = 8, blocked: int = 2) -> Any: + m = MagicMock() + m.guardrail_id = guardrail_id + m.date = date + m.requests_evaluated = requests + m.passed_count = passed + m.blocked_count = blocked + m.flagged_count = 0 + return m + + +def _prisma( + *, + find_many=None, + find_unique=None, + metrics=None, + index_find_many=None, +) -> MagicMock: + client = MagicMock() + db = client.db + db.litellm_guardrailstable.find_many = AsyncMock(return_value=find_many or []) + db.litellm_guardrailstable.find_unique = AsyncMock(return_value=find_unique) + db.litellm_dailyguardrailmetrics.find_many = AsyncMock(return_value=metrics or []) + db.litellm_spendlogguardrailindex.find_many = AsyncMock(return_value=index_find_many or []) + db.litellm_spendlogguardrailindex.count = AsyncMock(return_value=0) + db.litellm_spendlogs.find_many = AsyncMock(return_value=[]) + return client + + +def _patches(prisma: MagicMock, handler: InMemoryGuardrailHandler): + return ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", handler), + ) + + +# ---- detail ----------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_detail_returns_yaml_guardrail_when_db_misses(): + prisma = _prisma(find_unique=None) + handler = _config_handler(_yaml_guardrail()) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_detail( + guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN + ) + assert resp.guardrail_id == "yaml-1" + assert resp.guardrail_name == "yaml-pii" + assert resp.provider == "presidio" # coerced from the LitellmParams pydantic model + assert resp.type == "PII" # from guardrail_info + assert resp.description == "yaml-defined" + + +@pytest.mark.asyncio +async def test_detail_404_when_neither_db_nor_config(): + prisma = _prisma(find_unique=None) + handler = _config_handler() # empty + p1, p2 = _patches(prisma, handler) + with p1, p2, pytest.raises(HTTPException) as exc: + await guardrails_usage_detail(guardrail_id="ghost", start_date=START, end_date=END, user_api_key_dict=ADMIN) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_detail_does_not_surface_db_sourced_in_memory_entry(): + """A stale in-memory entry (source=db, gone from DB) must 404, not resurface.""" + prisma = _prisma(find_unique=None) + handler = InMemoryGuardrailHandler() + stale = _yaml_guardrail(guardrail_id="stale-1", name="stale") + handler.IN_MEMORY_GUARDRAILS["stale-1"] = stale + handler._sources["stale-1"] = "db" + p1, p2 = _patches(prisma, handler) + with p1, p2, pytest.raises(HTTPException) as exc: + await guardrails_usage_detail(guardrail_id="stale-1", start_date=START, end_date=END, user_api_key_dict=ADMIN) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_detail_db_row_still_resolves(): + prisma = _prisma(find_unique=_db_row(guardrail_id="db-1", provider="aim")) + handler = _config_handler() + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_detail( + guardrail_id="db-1", start_date=START, end_date=END, user_api_key_dict=ADMIN + ) + assert resp.provider == "aim" + assert resp.type == "ContentSafety" + + +# ---- overview --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_overview_includes_yaml_guardrail_with_no_metrics(): + """The core bug: a YAML guardrail with zero metrics must still appear as a row.""" + prisma = _prisma(find_many=[]) # no DB guardrails + handler = _config_handler(_yaml_guardrail()) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + rows = [r for r in resp.rows if r.id == "yaml-1"] + assert len(rows) == 1 + assert rows[0].name == "yaml-pii" + assert rows[0].provider == "presidio" + assert rows[0].type == "PII" + assert rows[0].requestsEvaluated == 0 + + +@pytest.mark.asyncio +async def test_overview_yaml_metrics_matched_by_logical_name(): + """Daily metrics are keyed by logical name; the YAML row must pick them up.""" + prisma = _prisma( + find_many=[], + metrics=[_metric("yaml-pii", requests=10, blocked=2)], # keyed by name, not uuid + ) + handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii")) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + rows = [r for r in resp.rows if r.id == "yaml-uuid"] + assert len(rows) == 1 + assert rows[0].requestsEvaluated == 10 + assert rows[0].failRate == 20.0 + # must not also emit an orphan row keyed by the logical name + assert [r for r in resp.rows if r.id == "yaml-pii"] == [] + + +@pytest.mark.asyncio +async def test_overview_excludes_db_sourced_in_memory_entry(): + """union must not resurrect a stale db-sourced in-memory guardrail.""" + prisma = _prisma(find_many=[]) + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["cfg"] = _yaml_guardrail(guardrail_id="cfg", name="cfg-guard") + handler._sources["cfg"] = "config" + handler.IN_MEMORY_GUARDRAILS["stale"] = _yaml_guardrail(guardrail_id="stale", name="stale-guard") + handler._sources["stale"] = "db" + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + ids = {r.id for r in resp.rows} + assert "cfg" in ids + assert "stale" not in ids + + +# ---- logs ------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_logs_resolves_config_guardrail_logical_name(): + """The index query must include the YAML guardrail's logical name alias.""" + prisma = _prisma(find_unique=None) + handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii")) + p1, p2 = _patches(prisma, handler) + with p1, p2: + await guardrails_usage_logs( + guardrail_id="yaml-uuid", + policy_id=None, + page=1, + page_size=50, + action=None, + start_date=START, + end_date=END, + user_api_key_dict=ADMIN, + ) + where = prisma.db.litellm_spendlogguardrailindex.find_many.call_args.kwargs["where"] + assert where["guardrail_id"] == {"in": ["yaml-uuid", "yaml-pii"]} From 8b31b200ca03d2c67b725440c11e78d2276de401 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 10 Jul 2026 17:24:36 -0700 Subject: [PATCH 05/62] refactor(guardrails): generalize _to_dict to any pydantic model Addresses PR review: _to_dict special-cased LitellmParams and returned an empty dict for any other pydantic model. Switch to isinstance(value, BaseModel) so it coerces any pydantic model uniformly (BaseModel is already imported for the response models), which also drops the now-unused LitellmParams import. Behavior is unchanged for the current call sites. --- litellm/proxy/guardrails/usage_endpoints.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index c63526e4833..f56b22ddd49 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -20,7 +20,6 @@ from litellm.repositories.table_repositories import ( SpendLogGuardrailIndexRepository, SpendLogsRepository, ) -from litellm.types.guardrails import LitellmParams router = APIRouter() @@ -146,8 +145,8 @@ def _get_guardrail_field(g: Any, field: str) -> Any: def _to_dict(value: Any) -> Dict[str, Any]: - """Coerce a LitellmParams / guardrail_info value into a plain dict.""" - if isinstance(value, LitellmParams): + """Coerce a pydantic model (e.g. LitellmParams) / dict value into a plain dict.""" + if isinstance(value, BaseModel): return value.model_dump(exclude_none=True) if isinstance(value, dict): return value From 26a7f6dc9da0f64071591203ee3f1dae449f67d7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 14 Jul 2026 07:21:45 -0700 Subject: [PATCH 06/62] fix(streaming): surface upstream connection resets instead of empty 200 streams --- .../litellm_core_utils/streaming_handler.py | 216 +++++++++--------- .../llms/custom_httpx/aiohttp_transport.py | 25 +- .../test_streaming_handler.py | 112 +++++++++ .../custom_httpx/test_aiohttp_transport.py | 101 +++++--- 4 files changed, 304 insertions(+), 150 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 128ba0bf3ab..977536803f5 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1974,97 +1974,7 @@ class CustomStreamWrapper: self.chunks.append(processed_chunk) return processed_chunk except (StopAsyncIteration, StopIteration): - if self.sent_last_chunk is True: - # log the final chunk with accurate streaming values - try: - complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages, - logging_obj=self.logging_obj, - ) - except Exception as e: - # see sync __next__: a raise from stream_chunk_builder inside this - # except handler escapes __anext__ and drops the request from SpendLogs. - # Recover best-effort usage from the raw chunks so cost is still tracked - verbose_logger.warning( - "stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.", - str(e), - ) - try: - complete_streaming_response = self.model_response_creator( - chunk={"usage": calculate_total_usage(chunks=self.chunks)} - ) - except Exception: - complete_streaming_response = None - - response = self.model_response_creator() - if complete_streaming_response is not None: - setattr( - response, - "usage", - getattr(complete_streaming_response, "usage"), - ) - try: - _copy = complete_streaming_response.model_copy(deep=True) - except RuntimeError: - _copy = complete_streaming_response.model_copy() - asyncio.create_task( - self.async_cache_streaming_response( - processed_chunk=_copy, - cache_hit=cache_hit, - ) - ) - # Update hidden_params with final usage from - # stream_chunk_builder (see sync __next__ for full comment). - if ( - self.stream_options is None - and complete_streaming_response is not None - and self._last_returned_hidden_params is not None - ): - final_usage = getattr(complete_streaming_response, "usage", None) - if final_usage is not None: - self._last_returned_hidden_params["usage"] = final_usage - - if self.sent_stream_usage is False and self.send_stream_usage is True: - self.sent_stream_usage = True - return response - - _deferred_cb = getattr( - self.logging_obj, - "_on_deferred_stream_complete", - None, - ) - if _deferred_cb is not None: - # Proxy has post-call guardrails. Store the assembled - # response so the outer streaming consumer - # (ProxyLogging.async_post_call_streaming_iterator_hook) - # can fire the deferred callback AFTER all guardrail - # end-of-stream blocks complete. Scheduling here via - # create_task would race with unified_guardrail's - # end-of-stream block for short-stream providers. - self.logging_obj._deferred_stream_complete_args = ( # type: ignore[attr-defined] - complete_streaming_response, - cache_hit, - ) - else: - # prefer_async_handlers routes CustomLogger to async_success_handler - # when consumers use ``async for`` on sync-SDK streams. Legacy string - # callbacks still run via executor.submit inside dispatch_success_handlers. - asyncio.create_task( - self.logging_obj.dispatch_success_handlers( - complete_streaming_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) - ) - - raise StopAsyncIteration # Re-raise StopIteration - else: - self.sent_last_chunk = True - processed_chunk = self.finish_reason_handler() - return processed_chunk + return await self._finalize_completed_stream(cache_hit=cache_hit) except httpx.TimeoutException as e: # if httpx read timeout error occues traceback_exception = traceback.format_exc() ## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT @@ -2079,20 +1989,120 @@ class CustomStreamWrapper: # Handle any exceptions that might occur during streaming asyncio.create_task(self.logging_obj.async_failure_handler(e, traceback_exception)) self._handle_stream_fallback_error(e) + except (httpx.ReadError, httpx.RemoteProtocolError) as e: + if self.received_finish_reason is None: + self._log_stream_failure_and_raise(e) + return await self._finalize_completed_stream(cache_hit=cache_hit) except Exception as e: - traceback_exception = traceback.format_exc() - if self.logging_obj is not None: - self._record_partial_usage_for_failure() - ## LOGGING - threading.Thread( - target=self.logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() # log response - # Handle any exceptions that might occur during streaming - asyncio.create_task( - self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore + self._log_stream_failure_and_raise(e) + + async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream": + if self.sent_last_chunk is True: + # log the final chunk with accurate streaming values + try: + complete_streaming_response = litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, ) - self._handle_stream_fallback_error(e) + except Exception as e: + # see sync __next__: a raise from stream_chunk_builder inside this + # except handler escapes __anext__ and drops the request from SpendLogs. + # Recover best-effort usage from the raw chunks so cost is still tracked + verbose_logger.warning( + "stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.", + str(e), + ) + try: + complete_streaming_response = self.model_response_creator( + chunk={"usage": calculate_total_usage(chunks=self.chunks)} + ) + except Exception: + complete_streaming_response = None + + response = self.model_response_creator() + if complete_streaming_response is not None: + setattr( + response, + "usage", + getattr(complete_streaming_response, "usage"), + ) + try: + _copy = complete_streaming_response.model_copy(deep=True) + except RuntimeError: + _copy = complete_streaming_response.model_copy() + asyncio.create_task( + self.async_cache_streaming_response( + processed_chunk=_copy, + cache_hit=cache_hit, + ) + ) + # Update hidden_params with final usage from + # stream_chunk_builder (see sync __next__ for full comment). + if ( + self.stream_options is None + and complete_streaming_response is not None + and self._last_returned_hidden_params is not None + ): + final_usage = getattr(complete_streaming_response, "usage", None) + if final_usage is not None: + self._last_returned_hidden_params["usage"] = final_usage + + if self.sent_stream_usage is False and self.send_stream_usage is True: + self.sent_stream_usage = True + return response + + _deferred_cb = getattr( + self.logging_obj, + "_on_deferred_stream_complete", + None, + ) + if _deferred_cb is not None: + # Proxy has post-call guardrails. Store the assembled + # response so the outer streaming consumer + # (ProxyLogging.async_post_call_streaming_iterator_hook) + # can fire the deferred callback AFTER all guardrail + # end-of-stream blocks complete. Scheduling here via + # create_task would race with unified_guardrail's + # end-of-stream block for short-stream providers. + self.logging_obj._deferred_stream_complete_args = ( # type: ignore[attr-defined] + complete_streaming_response, + cache_hit, + ) + else: + # prefer_async_handlers routes CustomLogger to async_success_handler + # when consumers use ``async for`` on sync-SDK streams. Legacy string + # callbacks still run via executor.submit inside dispatch_success_handlers. + asyncio.create_task( + self.logging_obj.dispatch_success_handlers( + complete_streaming_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + ) + + raise StopAsyncIteration # Re-raise StopIteration + else: + self.sent_last_chunk = True + processed_chunk = self.finish_reason_handler() + return processed_chunk + + def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn: + traceback_exception = traceback.format_exc() + if self.logging_obj is not None: + self._record_partial_usage_for_failure() + ## LOGGING + threading.Thread( + target=self.logging_obj.failure_handler, + args=(e, traceback_exception), + ).start() # log response + # Handle any exceptions that might occur during streaming + asyncio.create_task( + self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore + ) + self._handle_stream_fallback_error(e) def _record_partial_usage_for_failure(self) -> None: """ diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 3172d3667e1..adac6a1b276 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -85,29 +85,12 @@ class AiohttpResponseStream(httpx.AsyncByteStream): try: async for chunk in self._aiohttp_response.content.iter_chunked(self.CHUNK_SIZE): yield chunk - except ( - aiohttp.ClientPayloadError, - aiohttp.client_exceptions.ClientPayloadError, - ) as e: - # Handle incomplete transfers more gracefully - # Log the error but don't re-raise if we've already yielded some data - verbose_logger.debug(f"Transfer incomplete, but continuing: {e}") - # If the error is due to incomplete transfer encoding, we can still - # return what we've received so far, similar to how httpx handles it - return except RuntimeError as e: - # Some providers (e.g., SSE streams) may close the connection - # causing aiohttp StreamReader to raise a generic RuntimeError - # with message "Connection closed.". Treat this as a graceful - # end-of-stream so downstream consumers don't error. - if "Connection closed" in str(e): - verbose_logger.debug("Upstream closed streaming connection; ending iterator gracefully") - return - raise + if "Connection closed" not in str(e): + raise + raise httpx.ReadError(str(e)) from e except aiohttp.http_exceptions.TransferEncodingError as e: - # Handle transfer encoding errors gracefully - verbose_logger.debug(f"Transfer encoding error, but continuing: {e}") - return + raise httpx.ReadError(str(e)) from e except Exception: # For other exceptions, use the normal mapping with map_aiohttp_exceptions(): diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index e430ce3b084..e66e9f07d26 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -3059,3 +3059,115 @@ async def test_stream_chunk_builder_raise_and_usage_recovery_failure_does_not_cr chunks = [c async for c in response] assert len(chunks) > 0 + + +class TransportErrorAfterChunksIterator: + """Yields the given chunks, then raises the given exception once, then StopAsyncIteration.""" + + def __init__(self, model_responses, exception): + self.model_responses = model_responses + self.exception = exception + self.index = 0 + self.raised = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index < len(self.model_responses): + chunk = self.model_responses[self.index] + self.index += 1 + return chunk + if not self.raised: + self.raised = True + raise self.exception + raise StopAsyncIteration + + +def _reset_test_chunk(content: Optional[str] = None, finish_reason: Optional[str] = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-reset-test", + created=1783458104, + model="stub-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=content), + finish_reason=finish_reason, + ) + ], + ) + + +@pytest.mark.asyncio +async def test_transport_read_error_after_finish_reason_ends_stream_gracefully( + logging_obj: Logging, +): + """A trailing connection reset after the provider's finish chunk must not fail the stream.""" + import httpx + + completion_stream = TransportErrorAfterChunksIterator( + model_responses=[ + _reset_test_chunk(content="Hello"), + _reset_test_chunk(finish_reason="stop"), + ], + exception=httpx.ReadError("Response payload is not completed"), + ) + response = CustomStreamWrapper( + completion_stream=completion_stream, + model="hosted_vllm/stub-model", + custom_llm_provider="hosted_vllm", + logging_obj=logging_obj, + ) + + chunks = [chunk async for chunk in response] + + finish_reasons = [ + chunk.choices[0].finish_reason + for chunk in chunks + if chunk.choices and chunk.choices[0].finish_reason + ] + contents = [ + chunk.choices[0].delta.content + for chunk in chunks + if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content + ] + assert finish_reasons == ["stop"] + assert contents == ["Hello"] + + +@pytest.mark.asyncio +async def test_transport_read_error_before_finish_reason_raises(logging_obj: Logging): + """A connection reset before any finish chunk must surface, never end as a clean stop. + + Regression test for silent empty/truncated HTTP 200 streams: the aiohttp + transport used to swallow mid-stream connection resets, so the wrapper saw a + clean end-of-stream and fabricated finish_reason "stop". + """ + import httpx + + from litellm.exceptions import MidStreamFallbackError + + completion_stream = TransportErrorAfterChunksIterator( + model_responses=[_reset_test_chunk(content="Hel")], + exception=httpx.ReadError("Response payload is not completed"), + ) + response = CustomStreamWrapper( + completion_stream=completion_stream, + model="hosted_vllm/stub-model", + custom_llm_provider="hosted_vllm", + logging_obj=logging_obj, + ) + + received = [] + with pytest.raises(MidStreamFallbackError): + async for chunk in response: + received.append(chunk) + + fabricated_finish_reasons = [ + chunk.choices[0].finish_reason + for chunk in received + if chunk.choices and chunk.choices[0].finish_reason + ] + assert fabricated_finish_reasons == [] diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py index 0c5e386c438..d1bc356662f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py @@ -81,7 +81,7 @@ class MockContent: def __init__(self, chunks=None, exception_to_raise=None, exception_at_chunk=None): self.chunks = chunks or [b"chunk1", b"chunk2", b"chunk3"] self.exception_to_raise = exception_to_raise - self.exception_at_chunk = exception_at_chunk or (len(self.chunks) - 1) + self.exception_at_chunk = exception_at_chunk if exception_at_chunk is not None else (len(self.chunks) - 1) self.chunk_index = 0 async def iter_chunked(self, chunk_size): @@ -107,15 +107,11 @@ async def test_aiohttp_response_stream_normal_flow(): @pytest.mark.asyncio -async def test_transfer_encoding_error_no_httpx_read_error(): - """Test that TransferEncodingError doesn't get converted to httpx.ReadError""" - - # Create a TransferEncodingError wrapped in ClientPayloadError (like in real scenarios) +async def test_client_payload_error_mid_stream_raises_read_error(): + """A connection reset mid-body must surface as httpx.ReadError, not truncate silently""" transfer_error = aiohttp.http_exceptions.TransferEncodingError( message="400, message: Not enough data for satisfy transfer length header." ) - - # Wrap it in ClientPayloadError as aiohttp does client_payload_error = aiohttp.ClientPayloadError( "Response payload is not completed" ) @@ -124,47 +120,100 @@ async def test_transfer_encoding_error_no_httpx_read_error(): mock_response = MockAiohttpResponse( content_chunks=[b"chunk1", b"chunk2", b"chunk3"], exception_to_raise=client_payload_error, - exception_at_chunk=1, # Error occurs at chunk 1 + exception_at_chunk=1, ) stream = AiohttpResponseStream(mock_response) # type: ignore received_chunks = [] - # This should NOT raise httpx.ReadError or any other exception - # It should handle the error gracefully and just return what was received - async for chunk in stream: - received_chunks.append(chunk) - print(f"received_chunks: {received_chunks}") + with pytest.raises(httpx.ReadError): + async for chunk in stream: + received_chunks.append(chunk) - # Should have received the first chunk before the error assert received_chunks == [b"chunk1"] - assert len(received_chunks) == 1 + assert mock_response.closed is True @pytest.mark.asyncio -async def test_client_payload_error_graceful_handling(): - """Test that ClientPayloadError is handled gracefully without stacktrace""" - # Create a ClientPayloadError directly +async def test_client_payload_error_before_first_chunk_raises_read_error(): + """A connection reset before any body byte must surface, not yield an empty 200 body""" client_error = aiohttp.client_exceptions.ClientPayloadError( "Response payload is not completed" ) mock_response = MockAiohttpResponse( - content_chunks=[b"data1", b"data2", b"data3"], + content_chunks=[b"data1", b"data2"], exception_to_raise=client_error, - exception_at_chunk=2, # Error occurs at chunk 2 + exception_at_chunk=0, ) stream = AiohttpResponseStream(mock_response) # type: ignore received_chunks = [] - # This should handle the error gracefully without raising - async for chunk in stream: - received_chunks.append(chunk) + with pytest.raises(httpx.ReadError): + async for chunk in stream: + received_chunks.append(chunk) - # Should have received chunks before the error - assert received_chunks == [b"data1", b"data2"] - assert len(received_chunks) == 2 + assert received_chunks == [] + assert mock_response.closed is True + + +@pytest.mark.asyncio +async def test_connection_closed_runtime_error_raises_read_error(): + """aiohttp's bare RuntimeError('Connection closed.') must surface as httpx.ReadError""" + mock_response = MockAiohttpResponse( + content_chunks=[b"data1", b"data2"], + exception_to_raise=RuntimeError("Connection closed."), + exception_at_chunk=1, + ) + + stream = AiohttpResponseStream(mock_response) # type: ignore + received_chunks = [] + + with pytest.raises(httpx.ReadError): + async for chunk in stream: + received_chunks.append(chunk) + + assert received_chunks == [b"data1"] + assert mock_response.closed is True + + +@pytest.mark.asyncio +async def test_unrelated_runtime_error_propagates_unmapped(): + """RuntimeErrors other than 'Connection closed' must propagate untouched""" + mock_response = MockAiohttpResponse( + content_chunks=[b"data1"], + exception_to_raise=RuntimeError("something else broke"), + exception_at_chunk=0, + ) + + stream = AiohttpResponseStream(mock_response) # type: ignore + + with pytest.raises(RuntimeError, match="something else broke"): + async for _ in stream: + pass + + +@pytest.mark.asyncio +async def test_transfer_encoding_error_raises_read_error(): + """A raw TransferEncodingError mid-body must surface as httpx.ReadError""" + mock_response = MockAiohttpResponse( + content_chunks=[b"data1", b"data2"], + exception_to_raise=aiohttp.http_exceptions.TransferEncodingError( + message="Not enough data to satisfy transfer length header." + ), + exception_at_chunk=1, + ) + + stream = AiohttpResponseStream(mock_response) # type: ignore + received_chunks = [] + + with pytest.raises(httpx.ReadError): + async for chunk in stream: + received_chunks.append(chunk) + + assert received_chunks == [b"data1"] + assert mock_response.closed is True @pytest.mark.asyncio From ae2462b369e4c55e19d28b2dbe417dc370b189ef Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 18:26:01 -0700 Subject: [PATCH 07/62] fix(mcp): index authed request-time tools missing from the semantic filter startup index The semantic tool filter builds its index by listing every MCP server without per-user credentials, so servers needing per-user auth (interactive OAuth tokens, user-scoped env vars) reject the anonymous tools/list and contribute zero routes. Request-time expansion resolves that auth, so filter_tools received tools the router could never select: an empty index failed open N->N (past 128 tools OpenAI rejects the request outright) and a partial index matched only unavailable tools, stripping every tool from the request. filter_tools now syncs missing tools into the router before matching (building the router when absent) behind an asyncio lock so each tool embeds once, and falls back to the full tool list when matches map to no available tool, consistent with the zero-match fallback. Context-window overflows keep failing closed. --- .../mcp_server/semantic_tool_filter.py | 78 ++++++++- .../mcp_server/test_semantic_tool_filter.py | 148 ++++++++++++++++++ 2 files changed, 219 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index e12c6cdbd56..59098ef6b0c 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -4,6 +4,7 @@ Semantic MCP Tool Filtering using semantic-router Filters MCP tools semantically for /chat/completions and /responses endpoints. """ +import asyncio from typing import TYPE_CHECKING, Any, Dict, List, Optional from litellm._logging import verbose_logger @@ -76,6 +77,7 @@ class SemanticMCPToolFilter: self.tool_router: Optional["SemanticRouter"] = None self.context_window_error: Optional[str] = None self._tool_map: Dict[str, Any] = {} # MCPTool objects or OpenAI function dicts + self._index_sync_lock = asyncio.Lock() async def build_router_from_mcp_registry(self) -> None: """Build semantic router from all MCP tools in the registry (no auth checks).""" @@ -182,6 +184,55 @@ class SemanticMCPToolFilter: return raise + def _tools_missing_from_index(self, tools: list[Any]) -> dict[str, Any]: + """Map name -> tool for every named tool not yet in the semantic index.""" + return { + name: tool + for name, tool in ((self._extract_tool_info(t)[0], t) for t in tools) + if name and name not in self._tool_map + } + + async def _ensure_tools_indexed(self, available_tools: list[Any]) -> None: + """ + Index request-time tools the startup build never saw. + + The startup index lists every registered MCP server WITHOUT per-user + credentials, so servers requiring per-user auth (interactive OAuth + tokens, user-scoped env vars) contribute zero routes. Tools reaching + the filter came through an authenticated expansion; without indexing + them here they can never be selected, so requests either bypass + filtering entirely (N->N) or lose every tool to unrelated matches. + """ + from semantic_router.routers.base import Route + + if not self._tools_missing_from_index(available_tools): + return + + async with self._index_sync_lock: + missing = self._tools_missing_from_index(available_tools) + if not missing: + return + + if self.tool_router is None: + self._build_router(list(missing.values())) + return + + descriptions = {name: self._extract_tool_info(tool)[1] for name, tool in missing.items()} + routes = [ + Route( + name=name, + description=description, + utterances=[description], + score_threshold=self.similarity_threshold, + ) + for name, description in descriptions.items() + ] + await self.tool_router.aadd(routes) + self._tool_map.update(missing) + verbose_logger.info( + f"Semantic tool filter indexed {len(routes)} request-time tools missing from the startup index" + ) + async def filter_tools( self, query: str, @@ -216,13 +267,21 @@ class SemanticMCPToolFilter: if not query or not query.strip(): return available_tools - # Router should be built on startup - if not, something went wrong - if self.tool_router is None: - verbose_logger.warning("Router not initialized - was build_router_from_mcp_registry() called on startup?") - return available_tools - # Run semantic filtering try: + await self._ensure_tools_indexed(available_tools) + + if self.context_window_error is not None: + raise SemanticToolFilterContextWindowError( + embedding_model=self.embedding_model, + stage="the MCP tool descriptions during semantic router build", + original_error=self.context_window_error, + ) + + if self.tool_router is None: + verbose_logger.warning("Semantic router could not be built from the request's tools") + return available_tools + limit = top_k or self.top_k matches = self.tool_router(text=query, limit=limit) matched_tool_names = self._extract_tool_names_from_matches(matches) @@ -230,8 +289,13 @@ class SemanticMCPToolFilter: if not matched_tool_names: return available_tools - return self._get_tools_by_names(matched_tool_names, available_tools) + filtered_tools = self._get_tools_by_names(matched_tool_names, available_tools) + if not filtered_tools: + return available_tools + return filtered_tools + except SemanticToolFilterContextWindowError: + raise except Exception as e: if _is_context_window_error(e): verbose_logger.error( @@ -240,7 +304,7 @@ class SemanticMCPToolFilter: ) raise SemanticToolFilterContextWindowError( embedding_model=self.embedding_model, - stage="the user query", + stage="the user query or the MCP tool descriptions being indexed", original_error=str(e), ) from e verbose_logger.error(f"Semantic tool filter failed: {e}", exc_info=True) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 9bc0a525326..f2327f977f1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -1664,3 +1664,151 @@ def test_is_context_window_error_detection_variants(): assert _is_context_window_error(ValueError("Invalid 'input[0]': maximum input length is 8192 tokens.")) assert not _is_context_window_error(ValueError("A generic API error occurred.")) assert not _is_context_window_error(None) + + +def _make_keyword_embedding_router(recorded_inputs): + """ + Mock litellm Router whose embeddings are deterministic keyword one-hots: + texts mentioning linear/issue/ticket embed to [1, 0], everything else to + [0, 1]. Lets tests assert real similarity ranking through the actual + semantic-router index. Every embedding input batch is appended to + recorded_inputs. + """ + from litellm.types.utils import Embedding, EmbeddingResponse + + def _vector(text): + lowered = text.lower() + if "linear" in lowered or "issue" in lowered or "ticket" in lowered: + return [1.0, 0.0] + return [0.0, 1.0] + + def mock_embedding_sync(*args, **kwargs): + texts = kwargs["input"] + recorded_inputs.append(list(texts)) + return EmbeddingResponse( + data=[Embedding(embedding=_vector(t), index=i, object="embedding") for i, t in enumerate(texts)], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync(*args, **kwargs) + + mock_router = Mock() + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + return mock_router + + +def _make_keyword_filter(recorded_inputs, top_k: int = 3): + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + + return SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=_make_keyword_embedding_router(recorded_inputs), + top_k=top_k, + similarity_threshold=0.3, + enabled=True, + ) + + +def _linear_issue_tool(): + return MCPTool( + name="linear_stub-get_issue", + description="Get a Linear issue (ticket) by its identifier such as LIT-1234", + inputSchema={"type": "object"}, + ) + + +def _linear_list_tool(): + return MCPTool( + name="linear_stub-list_issues", + description="List Linear issues (tickets) in the workspace", + inputSchema={"type": "object"}, + ) + + +def _weather_tool(): + return MCPTool( + name="weather_stub-get_weather", + description="Get the current weather conditions for a city", + inputSchema={"type": "object"}, + ) + + +@pytest.mark.asyncio +async def test_filter_indexes_request_tools_when_startup_index_is_empty(): + """ + Regression test: the startup index is built by listing every MCP server + WITHOUT per-user credentials, so a gateway whose servers all require + per-user auth (e.g. interactive OAuth) starts with an empty index + (tool_router is None). filter_tools then failed open and returned all N + tools unfiltered (customer-visible as an N->N header and, past 128 tools, + an OpenAI 400 "tools array too long"). The authed request-time tools must + instead be indexed on first sight so filtering actually runs. + """ + filter_instance = _make_keyword_filter([]) + assert filter_instance.tool_router is None + + tools = [_linear_issue_tool(), _weather_tool()] + filtered = await filter_instance.filter_tools( + query="what is Linear ticket LIT-3794 about", + available_tools=tools, + ) + + assert [t.name for t in filtered] == ["linear_stub-get_issue"] + print("✅ Empty startup index is built from authed request-time tools") + + +@pytest.mark.asyncio +async def test_filter_indexes_tools_missing_from_partial_index(): + """ + Regression test: servers whose tools/list needs per-user auth contribute + zero routes to the startup index while anonymously listable servers are + indexed. Tools reaching the filter through the authed request-time + expansion must be added to the existing router (and only embedded once; + repeat requests embed just the query). + """ + recorded_inputs = [] + filter_instance = _make_keyword_filter(recorded_inputs) + filter_instance._build_router([_weather_tool()]) + assert filter_instance.tool_router is not None + + tools = [_linear_issue_tool(), _weather_tool()] + query = "what is Linear ticket LIT-3794 about" + + filtered = await filter_instance.filter_tools(query=query, available_tools=tools) + assert [t.name for t in filtered] == ["linear_stub-get_issue"] + + calls_after_first = len(recorded_inputs) + filtered_again = await filter_instance.filter_tools(query=query, available_tools=tools) + assert [t.name for t in filtered_again] == ["linear_stub-get_issue"] + assert len(recorded_inputs) == calls_after_first + 1 + + print("✅ Partial startup index is completed from request-time tools, embedding each tool once") + + +@pytest.mark.asyncio +async def test_filter_fails_open_when_matches_are_not_in_available_tools(): + """ + Regression test: when the semantic router's matches are all tools that are + NOT in the request's available_tools (an index/request mismatch), the + filter returned an empty list, stripping every tool from the request and + breaking it outright (observed live as a 3->0 header followed by a + provider 400). It must fail open with the full tool list instead, matching + the zero-match fallback. + """ + filter_instance = _make_keyword_filter([]) + filter_instance._build_router([_weather_tool()]) + + tools = [_linear_issue_tool(), _linear_list_tool()] + filtered = await filter_instance.filter_tools( + query="current weather in San Francisco", + available_tools=tools, + ) + + assert [t.name for t in filtered] == ["linear_stub-get_issue", "linear_stub-list_issues"] + print("✅ Matches outside available_tools fail open instead of dropping every tool") From a6390abefa84666d0719e20866db6464c4a83453 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 18:43:46 -0700 Subject: [PATCH 08/62] refactor(mcp): log lazily built semantic index and avoid per-request dict allocation in the fast path --- .../_experimental/mcp_server/semantic_tool_filter.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 59098ef6b0c..75eaf162c86 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -184,6 +184,10 @@ class SemanticMCPToolFilter: return raise + def _has_tools_missing_from_index(self, tools: list[Any]) -> bool: + """Allocation-free check for any named tool not yet in the semantic index.""" + return any(name and name not in self._tool_map for name in (self._extract_tool_info(t)[0] for t in tools)) + def _tools_missing_from_index(self, tools: list[Any]) -> dict[str, Any]: """Map name -> tool for every named tool not yet in the semantic index.""" return { @@ -205,7 +209,7 @@ class SemanticMCPToolFilter: """ from semantic_router.routers.base import Route - if not self._tools_missing_from_index(available_tools): + if not self._has_tools_missing_from_index(available_tools): return async with self._index_sync_lock: @@ -215,6 +219,10 @@ class SemanticMCPToolFilter: if self.tool_router is None: self._build_router(list(missing.values())) + if self.tool_router is not None: + verbose_logger.info( + f"Semantic tool filter indexed {len(missing)} request-time tools missing from the startup index" + ) return descriptions = {name: self._extract_tool_info(tool)[1] for name, tool in missing.items()} From 8ca809d426ce889fea106c6c10ff6e13fce6cbdb Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 19:34:17 -0700 Subject: [PATCH 09/62] fix(mcp): scope semantic filter matching and indexing errors to the requesting call Match candidates are restricted to the request's own tool names via route_filter, so routes learned from other principals' listings cannot displace the caller's tools from top_k. Lazy indexing now uses the async aadd flow exclusively; an embedding failure, including a context-window overflow on an oversized description, raises for the requesting call only and never writes the shared context_window_error, so one request cannot poison the filter for every user on the worker. The router is also sized to the configured top_k, which the semantic-router index layer otherwise silently caps at its default of 5. --- .../mcp_server/semantic_tool_filter.py | 51 ++++++--- .../mcp_server/test_semantic_tool_filter.py | 101 ++++++++++++++++++ 2 files changed, 135 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 75eaf162c86..37d3e0aedea 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -206,9 +206,19 @@ class SemanticMCPToolFilter: the filter came through an authenticated expansion; without indexing them here they can never be selected, so requests either bypass filtering entirely (N->N) or lose every tool to unrelated matches. + + Runs async-only (no synchronous embedding on the request path) and + never writes shared error state: an embedding failure here raises and + is scoped to the requesting call, so one request's oversized tool + description cannot poison the filter for other users on the worker. """ + from semantic_router.routers import SemanticRouter from semantic_router.routers.base import Route + from litellm.router_strategy.auto_router.litellm_encoder import ( + LiteLLMRouterEncoder, + ) + if not self._has_tools_missing_from_index(available_tools): return @@ -217,14 +227,6 @@ class SemanticMCPToolFilter: if not missing: return - if self.tool_router is None: - self._build_router(list(missing.values())) - if self.tool_router is not None: - verbose_logger.info( - f"Semantic tool filter indexed {len(missing)} request-time tools missing from the startup index" - ) - return - descriptions = {name: self._extract_tool_info(tool)[1] for name, tool in missing.items()} routes = [ Route( @@ -235,7 +237,23 @@ class SemanticMCPToolFilter: ) for name, description in descriptions.items() ] - await self.tool_router.aadd(routes) + + if self.tool_router is None: + router = SemanticRouter( + routes=[], + encoder=LiteLLMRouterEncoder( + litellm_router_instance=self.router_instance, + model_name=self.embedding_model, + score_threshold=self.similarity_threshold, + ), + auto_sync="local", + top_k=self.top_k, + ) + await router.aadd(routes) + self.tool_router = router + else: + await self.tool_router.aadd(routes) + self._tool_map.update(missing) verbose_logger.info( f"Semantic tool filter indexed {len(routes)} request-time tools missing from the startup index" @@ -279,19 +297,18 @@ class SemanticMCPToolFilter: try: await self._ensure_tools_indexed(available_tools) - if self.context_window_error is not None: - raise SemanticToolFilterContextWindowError( - embedding_model=self.embedding_model, - stage="the MCP tool descriptions during semantic router build", - original_error=self.context_window_error, - ) - if self.tool_router is None: verbose_logger.warning("Semantic router could not be built from the request's tools") return available_tools + available_names = [name for name in (self._extract_tool_info(t)[0] for t in available_tools) if name] + if not available_names: + return available_tools + limit = top_k or self.top_k - matches = self.tool_router(text=query, limit=limit) + if self.tool_router.top_k < limit: + self.tool_router.top_k = limit + matches = self.tool_router(text=query, limit=limit, route_filter=available_names) matched_tool_names = self._extract_tool_names_from_matches(matches) if not matched_tool_names: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index f2327f977f1..0f392a54b4c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -1678,6 +1678,8 @@ def _make_keyword_embedding_router(recorded_inputs): def _vector(text): lowered = text.lower() + if "kanban" in lowered: + return [0.6, 0.8] if "linear" in lowered or "issue" in lowered or "ticket" in lowered: return [1.0, 0.0] return [0.0, 1.0] @@ -1812,3 +1814,102 @@ async def test_filter_fails_open_when_matches_are_not_in_available_tools(): assert [t.name for t in filtered] == ["linear_stub-get_issue", "linear_stub-list_issues"] print("✅ Matches outside available_tools fail open instead of dropping every tool") + + +@pytest.mark.asyncio +async def test_request_time_context_window_error_is_request_scoped(): + """ + Regression test: an oversized tool description hitting the embedding + context window while lazily indexing request-time tools must fail only + the requesting call. Previously the lazy path reused the startup build + and recorded the overflow in the shared context_window_error, after + which EVERY user's MCP requests on the worker were blocked with a 400 + until restart (index poisoning via a single request). + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticToolFilterContextWindowError, + ) + + state = {"raise_context_error": True} + filter_instance = _make_context_window_filter(state) + tools = [ + MCPTool(name="tool_a", description="Tool A", inputSchema={"type": "object"}), + MCPTool(name="tool_b", description="Tool B", inputSchema={"type": "object"}), + ] + + with pytest.raises(SemanticToolFilterContextWindowError): + await filter_instance.filter_tools(query="send an email", available_tools=tools) + + assert filter_instance.context_window_error is None + assert filter_instance.tool_router is None + + state["raise_context_error"] = False + filtered = await filter_instance.filter_tools(query="send an email", available_tools=tools) + + assert len(filtered) > 0 + assert filter_instance.context_window_error is None + assert filter_instance.tool_router is not None + print("✅ Request-time context window overflow is scoped to the request, not the worker") + + +@pytest.mark.asyncio +async def test_foreign_index_routes_cannot_displace_available_tools(): + """ + Regression test: routes indexed from OTHER principals' tool listings must + not occupy the match candidate set for this request. Previously the + router matched over the whole shared index, so foreign routes that + embedded closer to the query displaced the caller's own tools from + top_k, degrading results to the fail-open list (or, before the + empty-result guard, stripping every tool). Matching is now scoped to the + request's own tool names via route_filter. + """ + filter_instance = _make_keyword_filter([], top_k=1) + foreign_tools = [ + MCPTool( + name=f"other_user-linear_tool_{i}", + description=f"Get a Linear issue variant {i}", + inputSchema={"type": "object"}, + ) + for i in range(6) + ] + filter_instance._build_router(foreign_tools) + + my_kanban = MCPTool( + name="mine-kanban_board", + description="Manage kanban board cards", + inputSchema={"type": "object"}, + ) + filtered = await filter_instance.filter_tools( + query="what is Linear ticket LIT-3794 about", + available_tools=[my_kanban, _weather_tool()], + ) + + assert [t.name for t in filtered] == ["mine-kanban_board"] + print("✅ Foreign index routes cannot displace the caller's own tools") + + +@pytest.mark.asyncio +async def test_top_k_above_router_default_is_respected(): + """ + Regression test: semantic-router's SemanticRouter defaults to top_k=5 at + the index-query layer, silently capping any configured filter top_k + above 5 regardless of the limit passed to __call__. The router must be + sized (and resized) to honor the configured top_k. + """ + filter_instance = _make_keyword_filter([], top_k=6) + tools = [ + MCPTool( + name=f"linear_stub-tool_{i}", + description=f"Work with Linear issues part {i}", + inputSchema={"type": "object"}, + ) + for i in range(6) + ] + + filtered = await filter_instance.filter_tools( + query="Linear ticket work", + available_tools=tools, + ) + + assert len(filtered) == 6 + print("✅ Configured top_k above the semantic-router default of 5 is honored") From 66f012a06bb448f459bb8f7c90a044abfaabc8d3 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 18:16:45 -0700 Subject: [PATCH 10/62] fix(mcp): discover missing OAuth scopes and token_url when authorization_url is set manually --- .../mcp_server/mcp_server_manager.py | 5 +- .../mcp_server/test_mcp_server_manager.py | 60 ++++++++++++++++++- 2 files changed, 62 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e6e265abb61..fda6ae9bef5 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1447,8 +1447,9 @@ class MCPServerManager: auth_type = cast(MCPAuthType, mcp_server.auth_type) server_url = mcp_server.url + has_all_upstream_oauth_fields = bool(mcp_server.authorization_url and mcp_server.token_url and scopes) needs_discovery = bool(server_url) and ( - (auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and not mcp_server.authorization_url) + (auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and not has_all_upstream_oauth_fields) or self._obo_needs_endpoint_discovery( auth_type, mcp_server.token_exchange_endpoint @@ -1467,7 +1468,7 @@ class MCPServerManager: if needs_discovery and mcp_oauth_metadata is None: verbose_logger.warning( "MCP OAuth discovery yielded no metadata for server %s (%s); " - "OAuth endpoints stay unresolved until a rebuild succeeds", + "OAuth endpoints/scopes stay unresolved until a rebuild succeeds", mcp_server.server_id, server_url, ) 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 adcfff6fe9d..6af86b9bda3 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 @@ -1026,7 +1026,6 @@ class TestMCPServerManager: """The gateway's relayed authorize flow (used by the browser-only Authorize) needs the upstream's authorization_url on the registry entry, and these rows never persist one, so the DB build must discover it the same way oauth2 rows do.""" - from types import SimpleNamespace manager = MCPServerManager() row = LiteLLM_MCPServerTable( @@ -1053,6 +1052,65 @@ class TestMCPServerManager: assert built.authorization_url == "https://idp.example.com/authorize" assert built.token_url == "https://idp.example.com/token" + @pytest.mark.asyncio + async def test_build_from_table_discovers_scopes_when_authorization_url_is_manual(self): + """An admin-typed authorization_url must not switch off discovery for the fields left + blank: without the scopes_supported backfill the authorize redirect goes out scope-less + and IdPs like Google hard-fail it with 400 "Missing required parameter: scope".""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="manual-auth-url-1", + alias="manual_auth_url", + description="manual authorization_url, blank scopes", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/manual-authorize", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/discovered-authorize", + token_url="https://idp.example.com/token", + registration_url=None, + scopes=["calendar.read", "calendar.write"], + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)) as mock_discovery: + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + mock_discovery.assert_awaited_once() + assert built.authorization_url == "https://idp.example.com/manual-authorize" + assert built.token_url == "https://idp.example.com/token" + assert built.scopes == ["calendar.read", "calendar.write"] + + @pytest.mark.asyncio + async def test_build_from_table_skips_discovery_when_all_upstream_oauth_fields_present(self): + """A fully hand-configured server (authorization_url, token_url, and scopes all set) has + nothing left for discovery to fill, so the build must not fetch upstream metadata.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="fully-manual-1", + alias="fully_manual", + description="all upstream oauth fields set by the admin", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/manual-authorize", + token_url="https://idp.example.com/manual-token", + credentials={"scopes": ["calendar.read"]}, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)) as mock_discovery: + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + mock_discovery.assert_not_awaited() + assert built.authorization_url == "https://idp.example.com/manual-authorize" + assert built.token_url == "https://idp.example.com/manual-token" + assert built.scopes == ["calendar.read"] + async def _capture_subject_token(self, call) -> Optional[str]: """Run a manager method (via ``call(manager)``) and return the subject_token it threaded into ``_create_mcp_client``.""" From 56655218d0329c15f2e7241f4114d2e09fb6b755 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 14 Jul 2026 19:49:09 -0700 Subject: [PATCH 11/62] refactor(ui): migrate vector stores, prompts, and skills tables onto shared DataTable Rewrites the three Batch A tables from hand-rolled TanStack + tremor renderers into thin DataTable consumers with a separate ColumnDef module each, matching the Guardrails and Tags migrations. Row actions move into a per-row overflow menu (edit/copy/delete for vector stores; copy for everyone plus admin-gated delete for prompts and skills). The vector stores parent gains an isLoading flag resolved on every exit path so the table shows the shared skeleton instead of flashing the empty state. Table files are renamed to PascalCase and stale eslint bulk-suppressions for the rewritten files are pruned. --- ui/litellm-dashboard/eslint-suppressions.json | 21 - .../prompts/_components/PromptTable.test.tsx | 104 +++++ .../prompts/_components/PromptTable.tsx | 89 ++++ .../_components/PromptTableColumns.tsx | 220 +++++++++ .../prompts/_components/index.test.tsx | 51 +++ .../(dashboard)/prompts/_components/index.tsx | 5 +- .../prompts/_components/prompt_table.tsx | 284 ------------ .../prompts/_components/prompt_utils.tsx | 2 +- .../ClaudeCodePluginsPanel.test.tsx | 50 +++ .../_components/ClaudeCodePluginsPanel.tsx | 9 +- .../skills/_components/PluginTable.test.tsx | 114 +++++ .../skills/_components/PluginTable.tsx | 60 +++ .../skills/_components/PluginTableColumns.tsx | 180 ++++++++ .../skills/_components/plugin_table.tsx | 245 ---------- .../_components/VectorStoreTable.test.tsx | 420 +++--------------- .../_components/VectorStoreTable.tsx | 238 ++-------- .../_components/VectorStoreTableColumns.tsx | 211 +++++++++ .../vector-stores/_components/index.test.tsx | 53 +++ .../vector-stores/_components/index.tsx | 9 +- 19 files changed, 1254 insertions(+), 1111 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_table.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTableColumns.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/skills/_components/plugin_table.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTableColumns.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 67139d2c932..44d4a111abd 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1098,14 +1098,6 @@ "count": 2 } }, - "src/app/(dashboard)/prompts/_components/prompt_table.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/router-settings/_components/general_settings.tsx": { "no-nested-ternary": { "count": 3 @@ -1153,14 +1145,6 @@ "count": 1 } }, - "src/app/(dashboard)/skills/_components/plugin_table.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/tag-management/_components/TagTable.tsx": { "no-restricted-imports": { "count": 1 @@ -1325,11 +1309,6 @@ "count": 1 } }, - "src/app/(dashboard)/vector-stores/_components/VectorStoreTable.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/vector-stores/_components/index.tsx": { "no-restricted-imports": { "count": 1 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 new file mode 100644 index 00000000000..52efd6407f8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx @@ -0,0 +1,104 @@ +import { render, screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { PromptSpec } from "@/components/networking"; + +import PromptTable from "./PromptTable"; + +vi.mock("@/components/networking", () => ({ + modelHubCall: vi.fn().mockResolvedValue({ data: [] }), +})); + +const mockPrompts: PromptSpec[] = [ + { + prompt_id: "prompt-newer", + litellm_params: { prompt_id: "prompt-newer" }, + prompt_info: { prompt_type: "dotprompt" }, + created_at: "2025-01-15T10:30:00Z", + updated_at: "2025-01-15T11:00:00Z", + environment: "production", + created_by: "user-1", + }, + { + prompt_id: "prompt-older", + litellm_params: { prompt_id: "prompt-older" }, + prompt_info: { prompt_type: "dotprompt" }, + created_at: "2024-01-10T09:15:00Z", + updated_at: "2024-01-12T14:20:00Z", + }, +]; + +const mockOnPromptClick = vi.fn(); +const mockOnDeleteClick = vi.fn(); + +const defaultProps = { + promptsList: mockPrompts, + isLoading: false, + onPromptClick: mockOnPromptClick, + onDeleteClick: mockOnDeleteClick, + accessToken: null, + isAdmin: true, +}; + +describe("PromptTable", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render every column header", () => { + render(); + for (const header of ["Prompt ID", "Model", "Created At", "Updated At", "Environment", "Created By", "Type"]) { + expect(screen.getByText(header)).toBeInTheDocument(); + } + }); + + it("should display the empty state when data is empty", () => { + render(); + expect(screen.getByText("No prompts yet")).toBeInTheDocument(); + }); + + it("should sort by created date descending by default", () => { + render(); + const rows = screen.getAllByRole("row").slice(1); + expect(within(rows[0]).getByText("prompt-newer")).toBeInTheDocument(); + expect(within(rows[1]).getByText("prompt-older")).toBeInTheDocument(); + }); + + it("should call onPromptClick when the prompt ID is clicked", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByRole("button", { name: "prompt-newer" })); + expect(mockOnPromptClick).toHaveBeenCalledWith("prompt-newer"); + }); + + it("should label the environment and default missing environments to development", () => { + render(); + expect(screen.getByText("production")).toBeInTheDocument(); + expect(screen.getByText("development")).toBeInTheDocument(); + }); + + it("should delete a prompt through the actions menu when admin", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("prompt-actions-prompt-newer")); + await user.click(await screen.findByTestId("prompt-action-delete")); + expect(mockOnDeleteClick).toHaveBeenCalledWith("prompt-newer", "prompt-newer"); + }); + + it("should copy the prompt ID through the actions menu", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("prompt-actions-prompt-newer")); + await user.click(await screen.findByTestId("prompt-action-copy")); + expect(await window.navigator.clipboard.readText()).toBe("prompt-newer"); + }); + + it("should hide the delete action for non-admins but keep copy available", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("prompt-actions-prompt-newer")); + expect(await screen.findByTestId("prompt-action-copy")).toBeInTheDocument(); + expect(screen.queryByTestId("prompt-action-delete")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx new file mode 100644 index 00000000000..47d4f64f254 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx @@ -0,0 +1,89 @@ +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import React, { useEffect, useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; +import { modelHubCall, PromptSpec } from "@/components/networking"; + +import { getPromptTableColumns } from "./PromptTableColumns"; +import { ModelGroupInfo } from "./prompt_utils"; + +interface PromptTableProps { + promptsList: PromptSpec[]; + isLoading: boolean; + onPromptClick?: (id: string) => void; + onDeleteClick?: (id: string, name: string) => void; + accessToken: string | null; + isAdmin: boolean; +} + +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +function EmptyState() { + return ( +
+
+ +
+
No prompts yet
+
Add a prompt to start managing reusable templates.
+
+ ); +} + +const PromptTable: React.FC = ({ + promptsList, + isLoading, + onPromptClick, + onDeleteClick, + accessToken, + isAdmin, +}) => { + const [sorting, setSorting] = useState(DEFAULT_SORTING); + const [modelHubData, setModelHubData] = useState>(new Map()); + + useEffect(() => { + const fetchModelHubData = async () => { + if (!accessToken) return; + + try { + const response = await modelHubCall(accessToken); + if (response?.data) { + const modelMap = new Map(); + response.data.forEach((model: ModelGroupInfo) => { + modelMap.set(model.model_group, model); + }); + setModelHubData(modelMap); + } + } catch (error) { + console.error("Error fetching model hub data:", error); + } + }; + + fetchModelHubData(); + }, [accessToken]); + + const columns = useMemo( + () => getPromptTableColumns({ modelHubData, isAdmin, onPromptClick, onDeleteClick }), + [modelHubData, isAdmin, onPromptClick, onDeleteClick], + ); + + return ( + prompt.prompt_id || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + isLoading={isLoading} + loadingMessage="Loading prompts…" + noDataMessage={} + size="compact" + /> + ); +}; + +export default PromptTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx new file mode 100644 index 00000000000..ae584ef6df6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx @@ -0,0 +1,220 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { Copy, MoreHorizontal, Trash2 } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { CellTooltip, DateCell, IdentityCell, StatusBadge, StatusTone } from "@/components/shared/table_cells"; +import { PromptSpec } from "@/components/networking"; +import { getProviderLogoAndName } from "@/components/provider_info_helpers"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; +import { copyToClipboard } from "@/utils/dataUtils"; + +import { extractModel, getProviderFromModelHub, ModelGroupInfo } from "./prompt_utils"; + +const ENVIRONMENT_TONE: Record = { + production: "error", + staging: "warning", + development: "success", +}; + +function PromptModelCell({ prompt, modelHubData }: { prompt: PromptSpec; modelHubData: Map }) { + const model = extractModel(prompt); + if (!model) { + return -; + } + + const provider = getProviderFromModelHub(model, modelHubData); + const { logo } = provider ? getProviderLogoAndName(provider) : { logo: "" }; + + return ( + + {logo ? ( + { + (event.currentTarget as HTMLImageElement).style.display = "none"; + }} + /> + ) : ( + + {provider?.charAt(0) || "-"} + + )} + {model} + + } + /> + ); +} + +interface PromptRowActionsProps { + prompt: PromptSpec; + isAdmin: boolean; + onDeleteClick?: (id: string, name: string) => void; +} + +function PromptRowActions({ prompt, isAdmin, onDeleteClick }: PromptRowActionsProps) { + return ( + + + + + + void copyToClipboard(prompt.prompt_id, "Prompt ID copied")} + > + + Copy prompt ID + + {isAdmin && ( + <> + + onDeleteClick?.(prompt.prompt_id, prompt.prompt_id || "Unknown Prompt")} + > + + Delete + + + )} + + + ); +} + +interface PromptTableColumnsDeps { + modelHubData: Map; + isAdmin: boolean; + onPromptClick?: (id: string) => void; + onDeleteClick?: (id: string, name: string) => void; +} + +export const getPromptTableColumns = ({ + modelHubData, + isAdmin, + onPromptClick, + onDeleteClick, +}: PromptTableColumnsDeps): ColumnDef[] => [ + { + id: "prompt_id", + accessorKey: "prompt_id", + meta: { title: "Prompt ID" }, + header: ({ column }) => , + size: 220, + enableSorting: true, + cell: ({ row }) => ( + onPromptClick(row.original.prompt_id) : undefined} + /> + ), + }, + { + id: "model", + meta: { title: "Model" }, + header: "Model", + size: 200, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "created_at", + accessorKey: "created_at", + sortingFn: "datetime", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 160, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "updated_at", + accessorKey: "updated_at", + sortingFn: "datetime", + meta: { title: "Updated At" }, + header: ({ column }) => , + size: 160, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "environment", + accessorKey: "environment", + meta: { title: "Environment", skeleton: "badge" }, + header: "Environment", + size: 130, + enableSorting: false, + cell: ({ row }) => { + const environment = row.original.environment || "development"; + return ; + }, + }, + { + id: "created_by", + accessorKey: "created_by", + meta: { title: "Created By" }, + header: "Created By", + size: 160, + enableSorting: false, + cell: ({ row }) => { + const createdBy = row.original.created_by; + return ( + + {createdBy || "-"} + + ); + }, + }, + { + id: "prompt_type", + accessorKey: "prompt_info.prompt_type", + meta: { title: "Type" }, + header: "Type", + size: 140, + enableSorting: false, + cell: ({ row }) => { + const promptType = row.original.prompt_info.prompt_type; + return ( + + {promptType} + + ); + }, + }, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx new file mode 100644 index 00000000000..99c58e2b98f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx @@ -0,0 +1,51 @@ +import { render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { getPromptsList } from "@/components/networking"; + +import PromptsPanel from "./index"; + +vi.mock("@/components/networking", () => ({ + getPromptsList: vi.fn(), + deletePromptCall: vi.fn(), +})); + +vi.mock("./PromptTable", () => ({ + __esModule: true, + default: ({ isLoading }: { isLoading: boolean }) => ( +
{isLoading ? "table-loading" : "table-loaded"}
+ ), +})); + +vi.mock("./prompt_info", () => ({ __esModule: true, default: () => null })); +vi.mock("./add_prompt_form", () => ({ __esModule: true, default: () => null })); +vi.mock("./prompt_editor_view", () => ({ __esModule: true, default: () => null })); + +const mockGetPromptsList = vi.mocked(getPromptsList); + +describe("PromptsPanel loading state", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should resolve the loading state when accessToken is null instead of showing the skeleton forever", async () => { + render(); + expect(await screen.findByText("table-loaded")).toBeInTheDocument(); + expect(mockGetPromptsList).not.toHaveBeenCalled(); + }); + + it("should show the loading state until the prompt fetch settles", async () => { + let resolveFetch: (value: { prompts: never[] }) => void = () => {}; + mockGetPromptsList.mockReturnValue( + new Promise((resolve) => { + resolveFetch = resolve; + }), + ); + render(); + expect(screen.getByText("table-loading")).toBeInTheDocument(); + + resolveFetch({ prompts: [] }); + expect(await screen.findByText("table-loaded")).toBeInTheDocument(); + expect(mockGetPromptsList).toHaveBeenCalledWith("sk-test", undefined); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx index 6e0e13d5181..de461ebd86d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx @@ -3,7 +3,7 @@ import React, { useState, useEffect } from "react"; import { Button } from "@tremor/react"; import { Modal, Select } from "antd"; import { getPromptsList, PromptSpec, ListPromptsResponse, deletePromptCall } from "@/components/networking"; -import PromptTable from "./prompt_table"; +import PromptTable from "./PromptTable"; import PromptInfoView from "./prompt_info"; import AddPromptForm from "./add_prompt_form"; import PromptEditorView from "./prompt_editor_view"; @@ -17,7 +17,7 @@ interface PromptsProps { const PromptsPanel: React.FC = ({ accessToken, userRole }) => { const [promptsList, setPromptsList] = useState([]); - const [isLoading, setIsLoading] = useState(false); + const [isLoading, setIsLoading] = useState(true); const [selectedEnvironment, setSelectedEnvironment] = useState(undefined); const [selectedPromptId, setSelectedPromptId] = useState(null); const [isAddModalVisible, setIsAddModalVisible] = useState(false); @@ -32,6 +32,7 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { const fetchPrompts = async () => { if (!accessToken) { + setIsLoading(false); return; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_table.tsx deleted file mode 100644 index 51a03d19e83..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_table.tsx +++ /dev/null @@ -1,284 +0,0 @@ -import React, { useState, useEffect } from "react"; -import { Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Button } from "@tremor/react"; -import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon, TrashIcon } from "@heroicons/react/outline"; -import { Tooltip } from "antd"; -import { PromptSpec, modelHubCall } from "@/components/networking"; -import { DateCell, IdCell } from "@/components/shared/table_cells"; -import { - ColumnDef, - flexRender, - getCoreRowModel, - getSortedRowModel, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { getProviderLogoAndName } from "@/components/provider_info_helpers"; -import { extractModel, getProviderFromModelHub } from "./prompt_utils"; - -interface PromptTableProps { - promptsList: PromptSpec[]; - isLoading: boolean; - onPromptClick?: (id: string) => void; - onDeleteClick?: (id: string, name: string) => void; - accessToken: string | null; - isAdmin: boolean; -} - -interface ModelGroupInfo { - model_group: string; - providers: string[]; - [key: string]: any; -} - -const PromptTable: React.FC = ({ - promptsList, - isLoading, - onPromptClick, - onDeleteClick, - accessToken, - isAdmin, -}) => { - const [sorting, setSorting] = useState([{ id: "created_at", desc: true }]); - const [modelHubData, setModelHubData] = useState>(new Map()); - - useEffect(() => { - const fetchModelHubData = async () => { - if (!accessToken) return; - - try { - const response = await modelHubCall(accessToken); - if (response?.data) { - const modelMap = new Map(); - response.data.forEach((model: ModelGroupInfo) => { - modelMap.set(model.model_group, model); - }); - setModelHubData(modelMap); - } - } catch (error) { - console.error("Error fetching model hub data:", error); - } - }; - - fetchModelHubData(); - }, [accessToken]); - - const columns: ColumnDef[] = [ - { - header: "Prompt ID", - accessorKey: "prompt_id", - cell: (info: any) => , - }, - { - header: "Model", - accessorKey: "model", - cell: ({ row }) => { - const prompt = row.original; - const model = extractModel(prompt); - - if (!model) { - return -; - } - - const provider = getProviderFromModelHub(model, modelHubData); - const { logo } = getProviderLogoAndName(provider || ""); - - return ( - -
- {/* Provider Icon */} -
- {provider && logo ? ( - {`${provider} { - const target = e.currentTarget as HTMLImageElement; - const parent = target.parentElement; - if (!parent || !parent.contains(target)) { - return; - } - - try { - const fallbackDiv = document.createElement("div"); - fallbackDiv.className = - "w-4 h-4 rounded-full bg-gray-200 flex items-center justify-center text-xs"; - fallbackDiv.textContent = provider?.charAt(0) || "-"; - parent.replaceChild(fallbackDiv, target); - } catch (error) { - console.error("Failed to replace provider logo fallback:", error); - } - }} - /> - ) : ( -
-
- )} -
- - {/* Model Name */} - {model} -
- - ); - }, - }, - { - header: "Created At", - accessorKey: "created_at", - cell: ({ row }) => , - }, - { - header: "Updated At", - accessorKey: "updated_at", - cell: ({ row }) => , - }, - { - header: "Environment", - accessorKey: "environment", - cell: ({ row }) => { - const prompt = row.original; - const env = prompt.environment || "development"; - const colorMap: Record = { - production: "text-red-600 bg-red-50", - staging: "text-yellow-600 bg-yellow-50", - development: "text-green-600 bg-green-50", - }; - return ( - {env} - ); - }, - }, - { - header: "Created By", - accessorKey: "created_by", - cell: ({ row }) => { - const prompt = row.original; - return {prompt.created_by || "-"}; - }, - }, - { - header: "Type", - accessorKey: "prompt_info.prompt_type", - cell: ({ row }) => { - const prompt = row.original; - return ( - - {prompt.prompt_info.prompt_type} - - ); - }, - }, - ...(isAdmin - ? [ - { - header: "Actions", - id: "actions", - enableSorting: false, - cell: ({ row }: any) => { - const prompt = row.original; - const promptName = prompt.prompt_id || "Unknown Prompt"; - - return ( -
- -
- ); - }, - }, - ] - : []), - ]; - - const table = useReactTable({ - data: promptsList, - columns, - state: { - sorting, - }, - onSortingChange: setSorting, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - enableSorting: true, - }); - - return ( -
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - -
-
- {header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())} -
-
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
-
- - ))} - - ))} - - - {isLoading ? ( - - -
-

Loading...

-
- - - ) : promptsList.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No prompts found

-
- - - )} - -
-
-
- ); -}; - -export default PromptTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_utils.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_utils.tsx index 6e0bd096af0..5c21a7d9ac6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_utils.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_utils.tsx @@ -1,7 +1,7 @@ import { PromptSpec } from "@/components/networking"; import { getVersionNumber } from "./prompt_editor_view/utils"; -interface ModelGroupInfo { +export interface ModelGroupInfo { model_group: string; providers: string[]; [key: string]: any; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.test.tsx new file mode 100644 index 00000000000..52f3dc21b7a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.test.tsx @@ -0,0 +1,50 @@ +import { render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { getClaudeCodePluginsList } from "@/components/networking"; + +import ClaudeCodePluginsPanel from "./ClaudeCodePluginsPanel"; + +vi.mock("@/components/networking", () => ({ + getClaudeCodePluginsList: vi.fn(), + deleteClaudeCodePlugin: vi.fn(), +})); + +vi.mock("./PluginTable", () => ({ + __esModule: true, + default: ({ isLoading }: { isLoading: boolean }) => ( +
{isLoading ? "table-loading" : "table-loaded"}
+ ), +})); + +vi.mock("./add_plugin_form", () => ({ __esModule: true, default: () => null })); +vi.mock("@/components/claude_code_plugins/skill_detail", () => ({ __esModule: true, default: () => null })); + +const mockGetClaudeCodePluginsList = vi.mocked(getClaudeCodePluginsList); + +describe("ClaudeCodePluginsPanel loading state", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should resolve the loading state when accessToken is null instead of showing the skeleton forever", async () => { + render(); + expect(await screen.findByText("table-loaded")).toBeInTheDocument(); + expect(mockGetClaudeCodePluginsList).not.toHaveBeenCalled(); + }); + + it("should show the loading state until the skills fetch settles", async () => { + let resolveFetch: (value: { plugins: never[]; count: number }) => void = () => {}; + mockGetClaudeCodePluginsList.mockReturnValue( + new Promise((resolve) => { + resolveFetch = resolve; + }), + ); + render(); + expect(screen.getByText("table-loading")).toBeInTheDocument(); + + resolveFetch({ plugins: [], count: 0 }); + expect(await screen.findByText("table-loaded")).toBeInTheDocument(); + expect(mockGetClaudeCodePluginsList).toHaveBeenCalledWith("sk-test", false); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx index 178cd36857f..e1dd325def1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx @@ -3,7 +3,7 @@ import { Button } from "@tremor/react"; import { Modal } from "antd"; import { getClaudeCodePluginsList, deleteClaudeCodePlugin } from "@/components/networking"; import AddPluginForm from "./add_plugin_form"; -import PluginTable from "./plugin_table"; +import PluginTable from "./PluginTable"; import SkillDetail from "@/components/claude_code_plugins/skill_detail"; import { isAdminRole } from "@/utils/roles"; import NotificationsManager from "@/components/molecules/notifications_manager"; @@ -17,7 +17,7 @@ interface ClaudeCodePluginsPanelProps { const ClaudeCodePluginsPanel: React.FC = ({ accessToken, userRole }) => { const [pluginsList, setPluginsList] = useState([]); const [isAddModalVisible, setIsAddModalVisible] = useState(false); - const [isLoading, setIsLoading] = useState(false); + const [isLoading, setIsLoading] = useState(true); const [isDeleting, setIsDeleting] = useState(false); const [pluginToDelete, setPluginToDelete] = useState<{ name: string; @@ -28,7 +28,10 @@ const ClaudeCodePluginsPanel: React.FC = ({ accessT const isAdmin = userRole ? isAdminRole(userRole) : false; const fetchPlugins = async () => { - if (!accessToken) return; + if (!accessToken) { + setIsLoading(false); + return; + } setIsLoading(true); try { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx new file mode 100644 index 00000000000..3c9d46e26fd --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx @@ -0,0 +1,114 @@ +import { render, screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { Plugin } from "@/components/claude_code_plugins/types"; + +import PluginTable from "./PluginTable"; + +const mockPlugins: Plugin[] = [ + { + id: "plugin-id-newer", + name: "newer-skill", + version: "1.2.0", + description: "A skill for testing", + source: { source: "github", repo: "org/newer-skill" }, + category: "development", + enabled: true, + created_at: "2025-01-15T10:30:00Z", + }, + { + id: "plugin-id-older", + name: "older-skill", + source: { source: "github", repo: "org/older-skill" }, + enabled: false, + created_at: "2024-01-10T09:15:00Z", + }, +]; + +const mockOnDeleteClick = vi.fn(); +const mockOnPluginClick = vi.fn(); + +const defaultProps = { + pluginsList: mockPlugins, + isLoading: false, + onDeleteClick: mockOnDeleteClick, + accessToken: "sk-test", + isAdmin: true, + onPluginClick: mockOnPluginClick, +}; + +describe("PluginTable", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render every column header", () => { + render(); + for (const header of ["Skill Name", "Version", "Description", "Category", "Public", "Created At"]) { + expect(screen.getByText(header)).toBeInTheDocument(); + } + }); + + it("should display the empty state when data is empty", () => { + render(); + expect(screen.getByText("No skills found")).toBeInTheDocument(); + }); + + it("should sort by created date descending by default", () => { + render(); + const rows = screen.getAllByRole("row").slice(1); + expect(within(rows[0]).getByText("newer-skill")).toBeInTheDocument(); + expect(within(rows[1]).getByText("older-skill")).toBeInTheDocument(); + }); + + it("should call onPluginClick with the plugin ID when the skill name is clicked", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByRole("button", { name: "newer-skill" })); + expect(mockOnPluginClick).toHaveBeenCalledWith("plugin-id-newer"); + }); + + it("should call onPluginClick with the plugin ID when the row is clicked", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByText("A skill for testing")); + expect(mockOnPluginClick).toHaveBeenCalledWith("plugin-id-newer"); + }); + + it("should badge the category and fall back to Uncategorized", () => { + render(); + expect(screen.getByText("development")).toBeInTheDocument(); + expect(screen.getByText("Uncategorized")).toBeInTheDocument(); + }); + + it("should show whether the skill is public", () => { + render(); + expect(screen.getByText("Yes")).toBeInTheDocument(); + expect(screen.getByText("No")).toBeInTheDocument(); + }); + + it("should delete a skill through the actions menu when admin", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("plugin-actions-newer-skill")); + await user.click(await screen.findByTestId("plugin-action-delete")); + expect(mockOnDeleteClick).toHaveBeenCalledWith("newer-skill", "newer-skill"); + }); + + it("should copy the skill ID through the actions menu", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("plugin-actions-newer-skill")); + await user.click(await screen.findByTestId("plugin-action-copy")); + expect(await window.navigator.clipboard.readText()).toBe("plugin-id-newer"); + }); + + it("should hide the delete action for non-admins but keep copy available", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("plugin-actions-newer-skill")); + expect(await screen.findByTestId("plugin-action-copy")).toBeInTheDocument(); + expect(screen.queryByTestId("plugin-action-delete")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx new file mode 100644 index 00000000000..9078124c7f5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx @@ -0,0 +1,60 @@ +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; +import { Plugin } from "@/components/claude_code_plugins/types"; + +import { getPluginTableColumns } from "./PluginTableColumns"; + +interface PluginTableProps { + pluginsList: Plugin[]; + isLoading: boolean; + onDeleteClick: (pluginName: string, displayName: string) => void; + accessToken: string | null; + isAdmin: boolean; + onPluginClick: (pluginId: string) => void; +} + +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +function EmptyState() { + return ( +
+
+ +
+
No skills found
+
Add one to get started.
+
+ ); +} + +const PluginTable: React.FC = ({ pluginsList, isLoading, onDeleteClick, isAdmin, onPluginClick }) => { + const [sorting, setSorting] = useState(DEFAULT_SORTING); + + const columns = useMemo( + () => getPluginTableColumns({ isAdmin, onPluginClick, onDeleteClick }), + [isAdmin, onPluginClick, onDeleteClick], + ); + + return ( + plugin.id || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + onRowClick={(plugin) => onPluginClick(plugin.id)} + isLoading={isLoading} + loadingMessage="Loading skills…" + noDataMessage={} + size="compact" + /> + ); +}; + +export default PluginTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTableColumns.tsx new file mode 100644 index 00000000000..95c9924b375 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTableColumns.tsx @@ -0,0 +1,180 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { Copy, MoreHorizontal, Trash2 } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell, StatusBadge } from "@/components/shared/table_cells"; +import { getCategoryBadgeColor } from "@/components/claude_code_plugins/helpers"; +import { Plugin } from "@/components/claude_code_plugins/types"; +import { Badge } from "@/components/ui/badge"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; +import { copyToClipboard } from "@/utils/dataUtils"; + +const CATEGORY_BADGE_CLASS: Record, string> = { + blue: "border-blue-200 bg-blue-50 text-blue-600", + green: "border-green-200 bg-green-50 text-green-600", + purple: "border-purple-200 bg-purple-50 text-purple-600", + red: "border-red-200 bg-red-50 text-red-600", + orange: "border-orange-200 bg-orange-50 text-orange-600", + yellow: "border-yellow-200 bg-yellow-50 text-yellow-600", + gray: "border-gray-200 bg-gray-50 text-gray-600", +}; + +function PluginCategoryBadge({ category }: { category?: string }) { + return ( + + {category || "Uncategorized"} + + ); +} + +interface PluginRowActionsProps { + plugin: Plugin; + isAdmin: boolean; + onDeleteClick: (pluginName: string, displayName: string) => void; +} + +function PluginRowActions({ plugin, isAdmin, onDeleteClick }: PluginRowActionsProps) { + return ( + + + + + + void copyToClipboard(plugin.id, "Skill ID copied")} + > + + Copy skill ID + + {isAdmin && ( + <> + + onDeleteClick(plugin.name, plugin.name)} + > + + Delete + + + )} + + + ); +} + +interface PluginTableColumnsDeps { + isAdmin: boolean; + onPluginClick: (pluginId: string) => void; + onDeleteClick: (pluginName: string, displayName: string) => void; +} + +export const getPluginTableColumns = ({ + isAdmin, + onPluginClick, + onDeleteClick, +}: PluginTableColumnsDeps): ColumnDef[] => [ + { + id: "name", + accessorKey: "name", + meta: { title: "Skill Name" }, + header: ({ column }) => , + size: 220, + enableSorting: true, + cell: ({ row }) => ( + onPluginClick(row.original.id)} + /> + ), + }, + { + id: "version", + accessorKey: "version", + meta: { title: "Version" }, + header: "Version", + size: 100, + enableSorting: false, + cell: ({ row }) => {row.original.version || "N/A"}, + }, + { + id: "description", + accessorKey: "description", + meta: { title: "Description" }, + header: "Description", + size: 300, + enableSorting: false, + cell: ({ row }) => { + const description = row.original.description; + return ( + + {description || "No description"} + + ); + }, + }, + { + id: "category", + accessorKey: "category", + meta: { title: "Category", skeleton: "badge" }, + header: "Category", + size: 150, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "enabled", + accessorKey: "enabled", + meta: { title: "Public", skeleton: "badge" }, + header: "Public", + size: 100, + enableSorting: false, + cell: ({ row }) => ( + + ), + }, + { + id: "created_at", + accessorKey: "created_at", + sortingFn: "datetime", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 160, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/plugin_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/plugin_table.tsx deleted file mode 100644 index eb1c495374a..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/plugin_table.tsx +++ /dev/null @@ -1,245 +0,0 @@ -import { CopyOutlined } from "@ant-design/icons"; -import { ChevronDownIcon, ChevronUpIcon, SwitchVerticalIcon, TrashIcon } from "@heroicons/react/outline"; -import { - ColumnDef, - flexRender, - getCoreRowModel, - getSortedRowModel, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { Badge, Button, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow } from "@tremor/react"; -import { Tooltip } from "antd"; -import React, { useState } from "react"; -import { DateCell, IdCell, StatusBadge } from "@/components/shared/table_cells"; -import NotificationsManager from "@/components/molecules/notifications_manager"; -import { getCategoryBadgeColor } from "@/components/claude_code_plugins/helpers"; -import { Plugin } from "@/components/claude_code_plugins/types"; - -interface PluginTableProps { - pluginsList: Plugin[]; - isLoading: boolean; - onDeleteClick: (pluginName: string, displayName: string) => void; - accessToken: string | null; - isAdmin: boolean; - onPluginClick: (pluginId: string) => void; -} - -const PluginTable: React.FC = ({ - pluginsList, - isLoading, - onDeleteClick, - accessToken, - isAdmin, - onPluginClick, -}) => { - const [sorting, setSorting] = useState([{ id: "created_at", desc: true }]); - - const copyToClipboard = (text: string) => { - navigator.clipboard.writeText(text); - NotificationsManager.success("Copied to clipboard!"); - }; - - const columns: ColumnDef[] = [ - { - header: "Skill Name", - accessorKey: "name", - cell: ({ row }) => { - const plugin = row.original; - return ( -
- onPluginClick(plugin.id)} /> - - { - e.stopPropagation(); - copyToClipboard(plugin.id); - }} - className="cursor-pointer text-gray-500 hover:text-blue-500 text-xs" - /> - -
- ); - }, - }, - { - header: "Version", - accessorKey: "version", - cell: ({ row }) => { - const version = row.original.version || "N/A"; - return {version}; - }, - }, - { - header: "Description", - accessorKey: "description", - cell: ({ row }) => { - const description = row.original.description || "No description"; - return ( - - {description} - - ); - }, - }, - { - header: "Category", - accessorKey: "category", - cell: ({ row }) => { - const category = row.original.category; - if (!category) { - return ( - - Uncategorized - - ); - } - const badgeColor = getCategoryBadgeColor(category); - return ( - - {category} - - ); - }, - }, - { - header: "Public", - accessorKey: "enabled", - cell: ({ row }) => { - const plugin = row.original; - return ; - }, - }, - { - header: "Created At", - accessorKey: "created_at", - cell: ({ row }) => , - }, - ...(isAdmin - ? [ - { - header: "Actions", - id: "actions", - enableSorting: false, - cell: ({ row }: any) => { - const plugin = row.original; - - return ( -
- -
- ); - }, - }, - ] - : []), - ]; - - const table = useReactTable({ - data: pluginsList, - columns, - state: { - sorting, - }, - onSortingChange: setSorting, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - enableSorting: true, - }); - - return ( -
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - -
-
- {header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())} -
- {header.column.getCanSort() && ( -
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
- )} -
- - ))} - - ))} - - - {isLoading ? ( - - -
-

Loading...

-
- - - ) : pluginsList && pluginsList.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - onPluginClick(row.original.id)} - > - {row.getVisibleCells().map((cell) => ( - - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No skills found. Add one to get started.

-
- - - )} - -
-
-
- ); -}; - -export default PluginTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.test.tsx index 45ca40152fa..7446d0efe3c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.test.tsx @@ -1,90 +1,46 @@ -import { render, screen } from "@testing-library/react"; +import { render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import VectorStoreTable from "./VectorStoreTable"; + import { VectorStore } from "@/components/vector_store_management/types"; -// Mock dependencies -const mockGetProviderLogoAndName = vi.fn(); -const mockTableIconActionButton = vi.fn(); +import VectorStoreTable from "./VectorStoreTable"; vi.mock("@/components/provider_info_helpers", () => ({ - getProviderLogoAndName: (...args: any[]) => mockGetProviderLogoAndName(...args), -})); - -vi.mock("@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton", () => ({ - default: (props: any) => { - mockTableIconActionButton(props); - return ( - - ); + getProviderLogoAndName: (provider: string) => { + const providerMap: Record = { + openai: { displayName: "OpenAI", logo: "/openai-logo.png" }, + azure: { displayName: "Azure", logo: "/azure-logo.png" }, + }; + return providerMap[provider] || { displayName: provider, logo: "" }; }, })); -// Mock Tremor components to avoid complex styling issues -vi.mock("@tremor/react", () => ({ - Table: ({ children, ...props }: any) => {children}
, - TableHead: ({ children, ...props }: any) => {children}, - TableBody: ({ children, ...props }: any) => {children}, - TableRow: ({ children, ...props }: any) => {children}, - TableHeaderCell: ({ children, ...props }: any) => {children}, - TableCell: ({ children, ...props }: any) => {children}, -})); - -// Mock antd Tooltip -vi.mock("antd", () => ({ - Tooltip: ({ children, title }: any) => ( -
- {children} -
- ), -})); - -// Mock Heroicons -vi.mock("@heroicons/react/outline", () => ({ - ChevronDownIcon: (props: any) =>
, - ChevronUpIcon: (props: any) =>
, - SwitchVerticalIcon: (props: any) =>
, -})); - -// Test data const mockVectorStores: VectorStore[] = [ { - vector_store_id: "short-id", + vector_store_id: "vs-newer", custom_llm_provider: "openai", vector_store_name: "My OpenAI Store", vector_store_description: "A store for OpenAI vectors", + vector_store_metadata: { + ingested_files: [ + { filename: "a.pdf", ingested_at: "2024-01-15T10:00:00Z" }, + { filename: "b.pdf", ingested_at: "2024-01-15T10:00:00Z" }, + ], + }, created_at: "2024-01-15T10:30:00Z", updated_at: "2024-01-15T11:00:00Z", - created_by: "user-1", - updated_by: "user-1", }, { - vector_store_id: "very-long-vector-store-id-that-should-be-truncated", + vector_store_id: "vs-older", custom_llm_provider: "azure", - vector_store_name: undefined, // Test missing name - vector_store_description: "A store for Azure vectors with a very long description that should show a tooltip", + vector_store_name: undefined, + vector_store_description: undefined, created_at: "2024-01-10T09:15:00Z", updated_at: "2024-01-12T14:20:00Z", }, - { - vector_store_id: "store-3", - custom_llm_provider: "pg_vector", - vector_store_name: "PostgreSQL Store", - vector_store_description: undefined, // Test missing description - created_at: "2024-01-05T08:00:00Z", - updated_at: "2024-01-08T16:45:00Z", - }, ]; -// Mock functions const mockOnView = vi.fn(); const mockOnEdit = vi.fn(); const mockOnDelete = vi.fn(); @@ -96,319 +52,71 @@ const defaultProps = { onDelete: mockOnDelete, }; -// Helper function to render component -const renderComponent = (props = {}) => { - return render(); -}; - describe("VectorStoreTable", () => { beforeEach(() => { vi.clearAllMocks(); - - // Setup default mock returns for getProviderLogoAndName - mockGetProviderLogoAndName.mockImplementation((provider: string) => { - const providerMap: Record = { - openai: { displayName: "OpenAI", logo: "/openai-logo.png" }, - azure: { displayName: "Azure", logo: "/azure-logo.png" }, - pg_vector: { displayName: "PostgreSQL Vector", logo: "/pg-logo.png" }, - }; - return providerMap[provider] || { displayName: provider, logo: "" }; - }); }); - describe("Rendering", () => { - it("should render the table with data", () => { - renderComponent(); - expect(screen.getByRole("table")).toBeInTheDocument(); - }); - - it("should render table headers", () => { - renderComponent(); - expect(screen.getByText("Vector Store ID")).toBeInTheDocument(); - expect(screen.getByText("Name")).toBeInTheDocument(); - expect(screen.getByText("Description")).toBeInTheDocument(); - expect(screen.getByText("Provider")).toBeInTheDocument(); - expect(screen.getByText("Created At")).toBeInTheDocument(); - expect(screen.getByText("Updated At")).toBeInTheDocument(); - // Check that we have the expected number of header cells (7 data + 1 actions) - const headers = screen.getAllByRole("columnheader"); - expect(headers).toHaveLength(8); - }); - - it("should render all vector store rows", () => { - renderComponent(); - expect(screen.getAllByRole("row")).toHaveLength(mockVectorStores.length + 1); // +1 for header row - }); - - it("should render empty state when no data", () => { - renderComponent({ data: [] }); - expect(screen.getByText("No vector stores found")).toBeInTheDocument(); - }); + it("should render every column header", () => { + render(); + for (const header of ["Vector Store ID", "Name", "Description", "Files", "Provider", "Created At", "Updated At"]) { + expect(screen.getByText(header)).toBeInTheDocument(); + } }); - describe("Vector Store ID Column", () => { - it("should render short vector store IDs fully", () => { - renderComponent(); - expect(screen.getByText("short-id")).toBeInTheDocument(); - }); - - it("should truncate long vector store IDs", () => { - renderComponent(); - const idButton = screen.getByText("very-long-vector-store-id-that-should-be-truncated"); - expect(idButton).toHaveClass("truncate", "max-w-[15ch]"); - }); - - it("should make vector store ID clickable", async () => { - const user = userEvent.setup(); - renderComponent(); - const idButton = screen.getByText("short-id"); - await user.click(idButton); - expect(mockOnView).toHaveBeenCalledWith("short-id"); - }); - - it("should have correct styling for vector store ID button", () => { - renderComponent(); - const idButton = screen.getByText("short-id").closest("button"); - expect(idButton).toHaveClass("font-mono", "text-blue-500", "bg-blue-50", "hover:bg-blue-100"); - }); + it("should display the empty state when data is empty", () => { + render(); + expect(screen.getByText("No vector stores")).toBeInTheDocument(); }); - describe("Name Column", () => { - it("should render vector store name", () => { - renderComponent(); - expect(screen.getByText("My OpenAI Store")).toBeInTheDocument(); - }); - - it("should render fallback for missing name", () => { - renderComponent(); - const fallbackElements = screen.getAllByText("-"); - expect(fallbackElements.length).toBe(5); // One for missing name, one for missing description, three for missing files (one per store) - }); - - it("should wrap name in tooltip", () => { - renderComponent(); - const tooltips = screen.getAllByTestId("tooltip"); - const nameTooltip = tooltips.find((t) => t.getAttribute("data-title") === "My OpenAI Store"); - expect(nameTooltip).toBeInTheDocument(); - }); + it("should sort by created date descending by default", () => { + render(); + const rows = screen.getAllByRole("row").slice(1); + expect(within(rows[0]).getByText("vs-newer")).toBeInTheDocument(); + expect(within(rows[1]).getByText("vs-older")).toBeInTheDocument(); }); - describe("Description Column", () => { - it("should render vector store description", () => { - renderComponent(); - expect(screen.getByText("A store for OpenAI vectors")).toBeInTheDocument(); - }); - - it("should render fallback for missing description", () => { - renderComponent(); - const fallbackElements = screen.getAllByText("-"); - expect(fallbackElements.length).toBe(5); // One for missing name, one for missing description, three for missing files (one per store) - }); - - it("should wrap description in tooltip", () => { - renderComponent(); - const tooltips = screen.getAllByTestId("tooltip"); - const descTooltip = tooltips.find( - (t) => - t.getAttribute("data-title") === - "A store for Azure vectors with a very long description that should show a tooltip", - ); - expect(descTooltip).toBeInTheDocument(); - }); + it("should call onView when the vector store ID is clicked", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByRole("button", { name: "vs-newer" })); + expect(mockOnView).toHaveBeenCalledWith("vs-newer"); }); - describe("Provider Column", () => { - it("should render provider display name", () => { - renderComponent(); - expect(screen.getByText("OpenAI")).toBeInTheDocument(); - expect(screen.getByText("Azure")).toBeInTheDocument(); - expect(screen.getByText("PostgreSQL Vector")).toBeInTheDocument(); - }); - - it("should render provider logo when available", () => { - renderComponent(); - const logos = screen.getAllByRole("img"); - expect(logos).toHaveLength(3); // All providers have logos in our mock - expect(logos[0]).toHaveAttribute("src", "/openai-logo.png"); - expect(logos[0]).toHaveAttribute("alt", "OpenAI"); - }); - - it("should call getProviderLogoAndName for each provider", () => { - renderComponent(); - expect(mockGetProviderLogoAndName).toHaveBeenCalledWith("openai"); - expect(mockGetProviderLogoAndName).toHaveBeenCalledWith("azure"); - expect(mockGetProviderLogoAndName).toHaveBeenCalledWith("pg_vector"); - }); + it("should render provider display names", () => { + render(); + expect(screen.getByText("OpenAI")).toBeInTheDocument(); + expect(screen.getByText("Azure")).toBeInTheDocument(); }); - describe("Date Columns", () => { - it("should render created at dates", () => { - renderComponent(); - const dateElements = screen.getAllByText(/Jan \d+, 2024/); - expect(dateElements.length).toBe(6); // 3 created_at + 3 updated_at dates - }); - - it("should render updated at dates", () => { - renderComponent(); - const dateElements = screen.getAllByText(/Jan \d+, 2024/); - expect(dateElements.length).toBe(6); // 3 created_at + 3 updated_at dates - }); + it("should summarize ingested files and fall back to a dash without files", () => { + render(); + expect(screen.getByText("2 files")).toBeInTheDocument(); + const olderRow = screen.getAllByRole("row").slice(1)[1]; + expect(within(olderRow).getAllByText("-").length).toBeGreaterThan(0); }); - describe("Actions Column", () => { - it("should render edit and delete action buttons for each row", () => { - renderComponent(); - expect(screen.getAllByTestId("action-button-edit")).toHaveLength(mockVectorStores.length); - expect(screen.getAllByTestId("action-button-delete")).toHaveLength(mockVectorStores.length); - }); - - it("should call onEdit when edit button is clicked", async () => { - const user = userEvent.setup(); - renderComponent(); - const editButtons = screen.getAllByTestId("action-button-edit"); - await user.click(editButtons[0]); - expect(mockOnEdit).toHaveBeenCalledWith("short-id"); - }); - - it("should call onDelete when delete button is clicked", async () => { - const user = userEvent.setup(); - renderComponent(); - const deleteButtons = screen.getAllByTestId("action-button-delete"); - await user.click(deleteButtons[0]); - expect(mockOnDelete).toHaveBeenCalledWith("short-id"); - }); - - it("should pass correct props to TableIconActionButton", () => { - renderComponent(); - expect(mockTableIconActionButton).toHaveBeenCalledWith( - expect.objectContaining({ - variant: "Edit", - tooltipText: "Edit vector store", - onClick: expect.any(Function), - }), - ); - expect(mockTableIconActionButton).toHaveBeenCalledWith( - expect.objectContaining({ - variant: "Delete", - tooltipText: "Delete vector store", - onClick: expect.any(Function), - }), - ); - }); + it("should edit a vector store through the actions menu", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("vector-store-actions-vs-newer")); + await user.click(await screen.findByTestId("vector-store-action-edit")); + expect(mockOnEdit).toHaveBeenCalledWith("vs-newer"); }); - describe("Sorting", () => { - it("should initialize with created_at descending sort", () => { - renderComponent(); - // The table should initialize with sorting state - expect(screen.getByTestId("chevron-down")).toBeInTheDocument(); - }); - - it("should render sort icons for sortable columns", () => { - renderComponent(); - // Should have sort icons for Created At and Updated At columns - const sortIcons = screen.getAllByTestId(/^chevron-(up|down)$|^switch-vertical$/); - expect(sortIcons.length).toBeGreaterThan(0); - }); - - it("should make header cells clickable for sorting", () => { - renderComponent(); - const headerCells = screen.getAllByRole("columnheader"); - const sortableHeaders = headerCells.filter((cell) => cell.textContent !== ""); - expect(sortableHeaders.length).toBeGreaterThan(0); - }); - - it("should show ascending icon when sorted ascending", () => { - renderComponent(); - // Initially shows descending, but we can test the logic by checking the icons are present - expect(screen.getByTestId("chevron-down")).toBeInTheDocument(); - }); + it("should delete a vector store through the actions menu", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("vector-store-actions-vs-newer")); + await user.click(await screen.findByTestId("vector-store-action-delete")); + expect(mockOnDelete).toHaveBeenCalledWith("vs-newer"); }); - describe("Styling and Layout", () => { - it("should apply correct CSS classes to table container", () => { - renderComponent(); - const tableContainer = screen.getByRole("table").parentElement?.parentElement; - expect(tableContainer).toHaveClass("rounded-lg", "custom-border", "relative"); - }); - - it("should apply overflow styling to table wrapper", () => { - renderComponent(); - const tableWrapper = screen.getByRole("table").parentElement; - expect(tableWrapper).toHaveClass("overflow-x-auto"); - }); - - it("should apply sticky styling to actions column", () => { - renderComponent(); - const headerCells = screen.getAllByRole("columnheader"); - const actionsHeader = headerCells[headerCells.length - 1]; - expect(actionsHeader).toHaveClass("sticky", "right-0", "bg-white"); - }); - - it("should apply sticky styling to action cells", () => { - renderComponent(); - const rows = screen.getAllByRole("row").slice(1); // Skip header row - rows.forEach((row) => { - const cells = row.querySelectorAll("td"); - const lastCell = cells[cells.length - 1]; - expect(lastCell).toHaveClass("sticky", "right-0", "bg-white"); - }); - }); - }); - - describe("Table Row Styling", () => { - it("should apply correct height to table rows", () => { - renderComponent(); - const rows = screen.getAllByRole("row").slice(1); // Skip header row - rows.forEach((row) => { - expect(row).toHaveClass("h-8"); - }); - }); - - it("should apply correct cell padding and styling", () => { - renderComponent(); - const cells = screen.getAllByRole("cell"); - cells.forEach((cell) => { - expect(cell).toHaveClass("py-0.5", "max-h-8", "overflow-hidden", "text-ellipsis", "whitespace-nowrap"); - }); - }); - }); - - describe("Empty State", () => { - it("should render single row with centered message when no data", () => { - renderComponent({ data: [] }); - const rows = screen.getAllByRole("row"); - expect(rows).toHaveLength(2); // Header + empty state row - expect(screen.getByText("No vector stores found")).toBeInTheDocument(); - }); - - it("should span all columns in empty state", () => { - renderComponent({ data: [] }); - const emptyCell = screen.getByText("No vector stores found").closest("td"); - expect(emptyCell).toHaveAttribute("colSpan", "8"); // 7 data columns + 1 actions column - }); - }); - - describe("Data Edge Cases", () => { - it("should handle vector stores with minimal data", () => { - const minimalData: VectorStore[] = [ - { - vector_store_id: "minimal", - custom_llm_provider: "test", - created_at: "2024-01-01T00:00:00Z", - updated_at: "2024-01-01T00:00:00Z", - }, - ]; - - renderComponent({ data: minimalData }); - expect(screen.getByText("minimal")).toBeInTheDocument(); - expect(screen.getAllByText("-")).toHaveLength(3); // Name, description, and files fallbacks - }); - - it("should handle single vector store", () => { - const singleData = [mockVectorStores[0]]; - renderComponent({ data: singleData }); - expect(screen.getAllByRole("row")).toHaveLength(2); // Header + 1 data row - }); + it("should copy the vector store ID through the actions menu", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("vector-store-actions-vs-newer")); + await user.click(await screen.findByTestId("vector-store-action-copy")); + expect(await window.navigator.clipboard.readText()).toBe("vs-newer"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.tsx index 074af873601..2f8508dc7c6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.tsx @@ -1,215 +1,57 @@ -import { ChevronDownIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; -import { - ColumnDef, - flexRender, - getCoreRowModel, - getSortedRowModel, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow } from "@tremor/react"; -import { Tooltip } from "antd"; -import React from "react"; -import { DateCell, IdCell } from "@/components/shared/table_cells"; -import TableIconActionButton from "@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; -import { getProviderLogoAndName } from "@/components/provider_info_helpers"; +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; import { VectorStore } from "@/components/vector_store_management/types"; +import { getVectorStoreTableColumns } from "./VectorStoreTableColumns"; + interface VectorStoreTableProps { data: VectorStore[]; onView: (vectorStoreId: string) => void; onEdit: (vectorStoreId: string) => void; onDelete: (vectorStoreId: string) => void; + isLoading?: boolean; } -const VectorStoreTable: React.FC = ({ data, onView, onEdit, onDelete }) => { - const [sorting, setSorting] = React.useState([{ id: "created_at", desc: true }]); - - const columns: ColumnDef[] = [ - { - header: "Vector Store ID", - accessorKey: "vector_store_id", - cell: ({ row }) => , - }, - { - header: "Name", - accessorKey: "vector_store_name", - cell: ({ row }) => { - const vectorStore = row.original; - return ( - - {vectorStore.vector_store_name || "-"} - - ); - }, - }, - { - header: "Description", - accessorKey: "vector_store_description", - cell: ({ row }) => { - const vectorStore = row.original; - return ( - - {vectorStore.vector_store_description || "-"} - - ); - }, - }, - { - header: "Files", - accessorKey: "vector_store_metadata", - cell: ({ row }) => { - const vectorStore = row.original; - const ingestedFiles = vectorStore.vector_store_metadata?.ingested_files || []; - - if (ingestedFiles.length === 0) { - return -; - } - - const filenames = ingestedFiles.map((file) => file.filename || file.file_url || "Unknown").join(", "); - - const displayText = - ingestedFiles.length === 1 - ? ingestedFiles[0].filename || ingestedFiles[0].file_url || "1 file" - : `${ingestedFiles.length} files`; - - return ( - - {displayText} - - ); - }, - }, - { - header: "Provider", - accessorKey: "custom_llm_provider", - cell: ({ row }) => { - const vectorStore = row.original; - const { displayName, logo } = getProviderLogoAndName(vectorStore.custom_llm_provider); - return ( -
- {logo && {displayName}} - {displayName} -
- ); - }, - }, - { - header: "Created At", - accessorKey: "created_at", - sortingFn: "datetime", - cell: ({ row }) => , - }, - { - header: "Updated At", - accessorKey: "updated_at", - sortingFn: "datetime", - cell: ({ row }) => , - }, - { - id: "actions", - header: "", - cell: ({ row }) => { - const vectorStore = row.original; - return ( -
- onEdit(vectorStore.vector_store_id)} - /> - onDelete(vectorStore.vector_store_id)} - /> -
- ); - }, - }, - ]; - - const table = useReactTable({ - data, - columns, - state: { - sorting, - }, - onSortingChange: setSorting, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - enableSorting: true, - }); +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; +function EmptyState() { return ( -
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - -
-
- {header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())} -
- {header.id !== "actions" && ( -
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
- )} -
- - ))} - - ))} - - - {table.getRowModel().rows.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No vector stores found

-
- - - )} - -
+
+
+ +
+
No vector stores
+
+ Connect a vector store to enable retrieval-augmented generation.
); +} + +const VectorStoreTable: React.FC = ({ data, onView, onEdit, onDelete, isLoading = false }) => { + const [sorting, setSorting] = useState(DEFAULT_SORTING); + + const columns = useMemo(() => getVectorStoreTableColumns({ onView, onEdit, onDelete }), [onView, onEdit, onDelete]); + + return ( + vectorStore.vector_store_id || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + isLoading={isLoading} + loadingMessage="Loading vector stores…" + noDataMessage={} + size="compact" + /> + ); }; export default VectorStoreTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTableColumns.tsx new file mode 100644 index 00000000000..cf162578177 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTableColumns.tsx @@ -0,0 +1,211 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { Copy, MoreHorizontal, Pencil, Trash2 } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { CellTooltip, DateCell, IdentityCell } from "@/components/shared/table_cells"; +import { getProviderLogoAndName } from "@/components/provider_info_helpers"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { VectorStore } from "@/components/vector_store_management/types"; +import { cn } from "@/lib/cva.config"; +import { copyToClipboard } from "@/utils/dataUtils"; + +function VectorStoreProviderCell({ provider }: { provider: string }) { + const { displayName, logo } = getProviderLogoAndName(provider); + return ( +
+ {logo ? ( + { + (event.currentTarget as HTMLImageElement).style.display = "none"; + }} + /> + ) : null} + {displayName} +
+ ); +} + +function VectorStoreFilesCell({ vectorStore }: { vectorStore: VectorStore }) { + const ingestedFiles = vectorStore.vector_store_metadata?.ingested_files || []; + if (ingestedFiles.length === 0) { + return -; + } + + const filenames = ingestedFiles.map((file) => file.filename || file.file_url || "Unknown").join(", "); + const displayText = + ingestedFiles.length === 1 + ? ingestedFiles[0].filename || ingestedFiles[0].file_url || "1 file" + : `${ingestedFiles.length} files`; + + return ( + {displayText}} + /> + ); +} + +interface VectorStoreRowActionsProps { + vectorStore: VectorStore; + onEdit: (vectorStoreId: string) => void; + onDelete: (vectorStoreId: string) => void; +} + +function VectorStoreRowActions({ vectorStore, onEdit, onDelete }: VectorStoreRowActionsProps) { + return ( + + + + + + onEdit(vectorStore.vector_store_id)}> + + Edit + + void copyToClipboard(vectorStore.vector_store_id, "Vector store ID copied")} + > + + Copy vector store ID + + + onDelete(vectorStore.vector_store_id)} + > + + Delete + + + + ); +} + +interface VectorStoreTableColumnsDeps { + onView: (vectorStoreId: string) => void; + onEdit: (vectorStoreId: string) => void; + onDelete: (vectorStoreId: string) => void; +} + +export const getVectorStoreTableColumns = ({ + onView, + onEdit, + onDelete, +}: VectorStoreTableColumnsDeps): ColumnDef[] => [ + { + id: "vector_store_id", + accessorKey: "vector_store_id", + meta: { title: "Vector Store ID" }, + header: ({ column }) => , + size: 220, + enableSorting: true, + cell: ({ row }) => ( + onView(row.original.vector_store_id)} + /> + ), + }, + { + id: "vector_store_name", + accessorKey: "vector_store_name", + meta: { title: "Name" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => { + const name = row.original.vector_store_name; + return ( + + {name || "-"} + + ); + }, + }, + { + id: "vector_store_description", + accessorKey: "vector_store_description", + meta: { title: "Description" }, + header: "Description", + size: 280, + enableSorting: false, + cell: ({ row }) => { + const description = row.original.vector_store_description; + return ( + + {description || "-"} + + ); + }, + }, + { + id: "files", + meta: { title: "Files" }, + header: "Files", + size: 160, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "provider", + accessorKey: "custom_llm_provider", + meta: { title: "Provider" }, + header: "Provider", + size: 160, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "created_at", + accessorKey: "created_at", + sortingFn: "datetime", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 150, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "updated_at", + accessorKey: "updated_at", + sortingFn: "datetime", + meta: { title: "Updated At" }, + header: ({ column }) => , + size: 150, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx new file mode 100644 index 00000000000..2931372f384 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx @@ -0,0 +1,53 @@ +import { render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { vectorStoreListCall } from "@/components/networking"; + +import VectorStoreManagement from "./index"; + +vi.mock("@/components/networking", () => ({ + vectorStoreListCall: vi.fn(), + vectorStoreDeleteCall: vi.fn(), + credentialListCall: vi.fn(), +})); + +vi.mock("./VectorStoreTable", () => ({ + __esModule: true, + default: ({ isLoading }: { isLoading?: boolean }) => ( +
{isLoading ? "table-loading" : "table-loaded"}
+ ), +})); + +vi.mock("./VectorStoreForm", () => ({ __esModule: true, default: () => null })); +vi.mock("./vector_store_info", () => ({ __esModule: true, default: () => null })); +vi.mock("./CreateVectorStore", () => ({ __esModule: true, default: () => null })); +vi.mock("./TestVectorStoreTab", () => ({ __esModule: true, default: () => null })); + +const mockVectorStoreListCall = vi.mocked(vectorStoreListCall); + +describe("VectorStoreManagement loading state", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should resolve the loading state when accessToken is null instead of showing the skeleton forever", async () => { + render(); + expect(await screen.findByText("table-loaded")).toBeInTheDocument(); + expect(mockVectorStoreListCall).not.toHaveBeenCalled(); + }); + + it("should show the loading state until the vector store fetch settles", async () => { + let resolveFetch: (value: { data: never[] }) => void = () => {}; + mockVectorStoreListCall.mockReturnValue( + new Promise((resolve) => { + resolveFetch = resolve; + }), + ); + render(); + expect(screen.getByText("table-loading")).toBeInTheDocument(); + + resolveFetch({ data: [] }); + expect(await screen.findByText("table-loaded")).toBeInTheDocument(); + expect(mockVectorStoreListCall).toHaveBeenCalledWith("sk-test"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx index 565e645f7b5..5d3f81f0275 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx @@ -36,6 +36,7 @@ interface VectorStoreProps { const VectorStoreManagement: React.FC = ({ accessToken, userID, userRole }) => { const [vectorStores, setVectorStores] = useState([]); + const [isLoadingVectorStores, setIsLoadingVectorStores] = useState(true); const [isCreateModalVisible, setIsCreateModalVisible] = useState(false); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [vectorStoreToDelete, setVectorStoreToDelete] = useState(null); @@ -46,13 +47,18 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID const [isDeleting, setIsDeleting] = useState(false); const fetchVectorStores = async () => { - if (!accessToken) return; + if (!accessToken) { + setIsLoadingVectorStores(false); + return; + } try { const response = await vectorStoreListCall(accessToken); setVectorStores(response.data || []); } catch (error) { console.error("Error fetching vector stores:", error); NotificationsManager.fromBackend("Error fetching vector stores: " + error); + } finally { + setIsLoadingVectorStores(false); } }; @@ -181,6 +187,7 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID Date: Tue, 14 Jul 2026 19:52:43 -0700 Subject: [PATCH 12/62] refactor(ui): drop whole-row navigation on the skills table Only the name cell and the overflow menu act on a row, matching the unified table pattern; the previous table navigated on any row click --- .../app/(dashboard)/skills/_components/PluginTable.test.tsx | 4 ++-- .../src/app/(dashboard)/skills/_components/PluginTable.tsx | 1 - 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx index 3c9d46e26fd..dd81a96d80a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx @@ -69,11 +69,11 @@ describe("PluginTable", () => { expect(mockOnPluginClick).toHaveBeenCalledWith("plugin-id-newer"); }); - it("should call onPluginClick with the plugin ID when the row is clicked", async () => { + it("should not navigate when clicking elsewhere in the row", async () => { const user = userEvent.setup(); render(); await user.click(screen.getByText("A skill for testing")); - expect(mockOnPluginClick).toHaveBeenCalledWith("plugin-id-newer"); + expect(mockOnPluginClick).not.toHaveBeenCalled(); }); it("should badge the category and fall back to Uncategorized", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx index 9078124c7f5..7cbe53e609c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx @@ -48,7 +48,6 @@ const PluginTable: React.FC = ({ pluginsList, isLoading, onDel sortingMode="client" sorting={sorting} onSortingChange={setSorting} - onRowClick={(plugin) => onPluginClick(plugin.id)} isLoading={isLoading} loadingMessage="Loading skills…" noDataMessage={} From 46939ea7a75945c89eae84b92ed3e717499683fc Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 14 Jul 2026 20:00:28 -0700 Subject: [PATCH 13/62] refactor(ui): drop unused accessToken prop from PluginTable --- .../(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx | 1 - .../src/app/(dashboard)/skills/_components/PluginTable.test.tsx | 1 - .../src/app/(dashboard)/skills/_components/PluginTable.tsx | 1 - 3 files changed, 3 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx index e1dd325def1..5a638f9ae79 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx @@ -98,7 +98,6 @@ const ClaudeCodePluginsPanel: React.FC = ({ accessT pluginsList={pluginsList} isLoading={isLoading} onDeleteClick={handleDeleteClick} - accessToken={accessToken} isAdmin={isAdmin} onPluginClick={(id) => { const skill = pluginsList.find((p) => p.id === id); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx index dd81a96d80a..66a4d7e2524 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.test.tsx @@ -33,7 +33,6 @@ const defaultProps = { pluginsList: mockPlugins, isLoading: false, onDeleteClick: mockOnDeleteClick, - accessToken: "sk-test", isAdmin: true, onPluginClick: mockOnPluginClick, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx index 7cbe53e609c..c581b0dfdeb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx @@ -13,7 +13,6 @@ interface PluginTableProps { pluginsList: Plugin[]; isLoading: boolean; onDeleteClick: (pluginName: string, displayName: string) => void; - accessToken: string | null; isAdmin: boolean; onPluginClick: (pluginId: string) => void; } From df60e36d07bd39d765c1d737ba98313a3214b06b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 15 Jul 2026 08:33:50 -0700 Subject: [PATCH 14/62] fix(responses): end stream cleanly on transport error after terminal event --- litellm/responses/streaming_iterator.py | 12 ++ .../responses/test_streaming_iterator.py | 143 ++++++++++++++++-- 2 files changed, 140 insertions(+), 15 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index eb78e6f9c8d..357a7ecefe6 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -718,6 +718,12 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): except StopAsyncIteration: # Normal end of stream - don't log as failure raise + except (httpx.ReadError, httpx.RemoteProtocolError) as e: + self.finished = True + if self.completed_response is None: + self._handle_failure(e) + raise + raise StopAsyncIteration from e except httpx.HTTPError as e: # Handle HTTP errors self.finished = True @@ -794,6 +800,12 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): except StopIteration: # Normal end of stream - don't log as failure raise + except (httpx.ReadError, httpx.RemoteProtocolError) as e: + self.finished = True + if self.completed_response is None: + self._handle_failure(e) + raise + raise StopIteration from e except httpx.HTTPError as e: # Handle HTTP errors self.finished = True diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 3e4b07fce18..5b0f40fdf27 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -5,13 +5,18 @@ completion_start_time = end_time.""" import json from datetime import datetime +from typing import Optional from unittest.mock import Mock +import httpx import pytest from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig -from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator +from litellm.responses.streaming_iterator import ( + ResponsesAPIStreamingIterator, + SyncResponsesAPIStreamingIterator, +) from litellm.types.llms.openai import ( ResponseCompletedEvent, ResponsesAPIResponse, @@ -23,19 +28,7 @@ def _sse_event(payload: dict) -> bytes: return f"data: {json.dumps(payload)}\n\n".encode("utf-8") -def _make_iterator( - *, - sse_events: list[bytes], - logging_obj: LiteLLMLoggingObj, -) -> ResponsesAPIStreamingIterator: - async def aiter_bytes(): - for evt in sse_events: - yield evt - - mock_response = Mock() - mock_response.headers = {} - mock_response.aiter_bytes = aiter_bytes - +def _mock_config() -> Mock: mock_config = Mock(spec=BaseResponsesAPIConfig) mock_responses_api_response = Mock(spec=ResponsesAPIResponse) mock_responses_api_response.id = "resp_ttft" @@ -52,17 +45,68 @@ def _make_iterator( return stub mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +def _make_iterator( + *, + sse_events: list[bytes], + logging_obj: LiteLLMLoggingObj, + trailing_error: Optional[Exception] = None, +) -> ResponsesAPIStreamingIterator: + async def aiter_bytes(): + for evt in sse_events: + yield evt + if trailing_error is not None: + raise trailing_error + + mock_response = Mock() + mock_response.headers = {} + mock_response.aiter_bytes = aiter_bytes return ResponsesAPIStreamingIterator( response=mock_response, model="gpt-4o-mini", - responses_api_provider_config=mock_config, + responses_api_provider_config=_mock_config(), logging_obj=logging_obj, litellm_metadata={}, custom_llm_provider="openai", ) +def _make_sync_iterator( + *, + sse_events: list[bytes], + logging_obj: LiteLLMLoggingObj, + trailing_error: Optional[Exception] = None, +) -> SyncResponsesAPIStreamingIterator: + def iter_bytes(): + for evt in sse_events: + yield evt + if trailing_error is not None: + raise trailing_error + + mock_response = Mock() + mock_response.headers = {} + mock_response.iter_bytes = iter_bytes + + return SyncResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-4o-mini", + responses_api_provider_config=_mock_config(), + logging_obj=logging_obj, + litellm_metadata={}, + custom_llm_provider="openai", + ) + + +def _logging_obj_stub() -> Mock: + logging_obj = Mock(spec=LiteLLMLoggingObj) + logging_obj.completion_start_time = None + logging_obj.model_call_details = {"litellm_params": {}} + return logging_obj + + @pytest.mark.asyncio async def test_responses_streaming_stamps_completion_start_time_on_first_chunk(): """Without the fix, `logging_obj.completion_start_time` stays None across the @@ -122,3 +166,72 @@ async def test_responses_streaming_does_not_reset_prior_completion_start_time(): logging_obj._update_completion_start_time.assert_not_called() assert logging_obj.completion_start_time == prior + + +_COMPLETE_STREAM_EVENTS = [ + _sse_event({"type": "response.created"}), + _sse_event({"type": "response.output_text.delta", "delta": "hi"}), + _sse_event({"type": "response.completed"}), +] + +_TRAILING_ERRORS = [ + httpx.ReadError("Response payload is not completed"), + httpx.RemoteProtocolError("peer closed connection without sending complete message body"), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) +async def test_transport_error_after_completed_event_ends_stream_cleanly(trailing_error): + """A sloppy connection close after `response.completed` must not turn a + complete stream into an error (regression guard for the transport no longer + swallowing ClientPayloadError/TransferEncodingError).""" + iterator = _make_iterator( + sse_events=_COMPLETE_STREAM_EVENTS, + logging_obj=_logging_obj_stub(), + trailing_error=trailing_error, + ) + + seen = [event.type async for event in iterator] + + assert ResponsesAPIStreamEvents.RESPONSE_COMPLETED in seen + + +@pytest.mark.asyncio +async def test_transport_error_before_completed_event_raises(): + """A connection lost before any terminal event is a real failure and must + surface, not end the stream as if it completed.""" + iterator = _make_iterator( + sse_events=_COMPLETE_STREAM_EVENTS[:-1], + logging_obj=_logging_obj_stub(), + trailing_error=httpx.ReadError("Response payload is not completed"), + ) + + with pytest.raises(httpx.ReadError): + async for _ in iterator: + pass + + +@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) +def test_sync_transport_error_after_completed_event_ends_stream_cleanly(trailing_error): + iterator = _make_sync_iterator( + sse_events=_COMPLETE_STREAM_EVENTS, + logging_obj=_logging_obj_stub(), + trailing_error=trailing_error, + ) + + seen = [event.type for event in iterator] + + assert ResponsesAPIStreamEvents.RESPONSE_COMPLETED in seen + + +def test_sync_transport_error_before_completed_event_raises(): + iterator = _make_sync_iterator( + sse_events=_COMPLETE_STREAM_EVENTS[:-1], + logging_obj=_logging_obj_stub(), + trailing_error=httpx.ReadError("Response payload is not completed"), + ) + + with pytest.raises(httpx.ReadError): + for _ in iterator: + pass From e14da657c80f72f32b42ab210c16a0f826745f3e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 15 Jul 2026 09:38:30 -0700 Subject: [PATCH 15/62] docs(proxy): document mcp_rpm_limit in update_team docstring --- litellm/proxy/management_endpoints/team_endpoints.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 797f4600857..d267a3cac69 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1621,6 +1621,7 @@ async def update_team( - allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team. - model_rpm_limit: Optional[Dict[str, int]] - The RPM (Requests Per Minute) limit per model for this team. Example: {"gpt-4": 100, "gpt-3.5-turbo": 200} - model_tpm_limit: Optional[Dict[str, int]] - The TPM (Tokens Per Minute) limit per model for this team. Example: {"gpt-4": 10000, "gpt-3.5-turbo": 20000} + - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. Example - update team TPM Limit - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. - secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 80e9b41852f..9f6d7410e77 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -13810,6 +13810,7 @@ export interface paths { * - allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team. * - model_rpm_limit: Optional[Dict[str, int]] - The RPM (Requests Per Minute) limit per model for this team. Example: {"gpt-4": 100, "gpt-3.5-turbo": 200} * - model_tpm_limit: Optional[Dict[str, int]] - The TPM (Tokens Per Minute) limit per model for this team. Example: {"gpt-4": 10000, "gpt-3.5-turbo": 20000} + * - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. * Example - update team TPM Limit * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. * - secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview) From a81c6ce350c2fd38d43a15c22edc5c4078f19675 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 10:22:24 -0700 Subject: [PATCH 16/62] fix(mcp): reject discovered token endpoints uncorroborated by the manual authorization_url Discovery is rooted at the MCP resource, so a compromised upstream can advertise an attacker-run authorization server. When authorization_url is manually configured and another field is blank, the per-field merge would combine the trusted authorize endpoint with the advertised token_url, and the gateway would redeem authorization codes (with the stored client secret and PKCE verifier) at that endpoint, then persist it. Discovered token_url and registration_url are now accepted only when the same metadata document advertises an authorization_endpoint matching the configured value (scheme+host+path). Scope backfill is unaffected. Applies to both the DB and config build paths. --- .../mcp_server/mcp_server_manager.py | 90 +++++++++-- .../mcp_server/test_mcp_server_manager.py | 144 +++++++++++++++++- 2 files changed, 221 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index fda6ae9bef5..284bf6a0926 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -213,6 +213,57 @@ def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serv new_server.scopes = previous_server.scopes +def _normalized_authorize_endpoint(url: str) -> str: + parsed = urlparse(url) + return f"{parsed.scheme.lower()}://{parsed.netloc.lower()}{parsed.path.rstrip('/')}" + + +def _gate_discovered_endpoints_against_manual_authorization_url( + metadata: MCPOAuthMetadata | None, + manual_authorization_url: str | None, + server_identifier: str, + is_dcr_bridge: bool, +) -> MCPOAuthMetadata | None: + """Refuse discovered token/registration endpoints the pinned authorize endpoint cannot vouch for. + + Discovery is rooted at the MCP resource (RFC 9728), so a compromised upstream can advertise an + attacker-run authorization server. When the admin manually configured ``authorization_url``, + filling a blank ``token_url`` from that advertisement recreates the RFC 9700 mix-up attack at + configuration time: users sign in at the trusted authorize endpoint while the gateway redeems + the code, with the stored client secret and PKCE verifier, at the attacker's token endpoint. + Endpoints from one metadata document are only trustworthy together, so the discovered + ``token_url`` and ``registration_url`` are accepted only when that same document's + ``authorization_endpoint`` matches the pinned value (scheme+host+path; query and trailing slash + are not identity). Scope discovery stays ungated: scopes steer the redirect to the trusted + authorize endpoint and carry no credentials. + """ + if metadata is None or not manual_authorization_url: + return metadata + if not metadata.token_url and not metadata.registration_url: + return metadata + if metadata.authorization_url and _normalized_authorize_endpoint( + manual_authorization_url + ) == _normalized_authorize_endpoint(metadata.authorization_url): + return metadata + bridge_note = ( + " The discovered registration_url is rejected with it, so this dcr_bridge server stays on the" + " short-circuit registration arm." + if is_dcr_bridge and metadata.registration_url + else "" + ) + verbose_logger.warning( + "MCP OAuth discovery for server %s advertised authorization_endpoint %s, which does not match the " + "manually configured authorization_url %s; rejecting the discovered token_url/registration_url so " + "authorization codes and client credentials only go to endpoints vouched for by the configured " + "authorization server. Configure Token URL manually if the mismatch is intentional.%s", + server_identifier, + _normalized_authorize_endpoint(metadata.authorization_url) if metadata.authorization_url else "", + _normalized_authorize_endpoint(manual_authorization_url), + bridge_note, + ) + return metadata.model_copy(update={"token_url": None, "registration_url": None}) + + def invalidate_user_env_vars_cache(user_id: str, server_id: str) -> None: """Drop a cached entry after the user stores or clears their env var values so the next request reads the fresh value instead of a stale one.""" @@ -1041,20 +1092,31 @@ class MCPServerManager: else: mcp_oauth_metadata = None + gated_oauth_metadata = ( + _gate_discovered_endpoints_against_manual_authorization_url( + mcp_oauth_metadata, + server_config.get("authorization_url"), + server_name or server_id, + bool(server_config.get("dcr_bridge")), + ) + if auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES + else mcp_oauth_metadata + ) + # Filter blank scopes (e.g. YAML ``scopes: [""]``) the same way the DB-build path does, so # an all-blank list normalizes to None rather than a ``("",)`` tuple that skips the # entra_obo fail-closed scope precondition and POSTs an empty scope to the IdP. resolved_scopes = self._extract_scopes(server_config.get("scopes")) or ( - mcp_oauth_metadata.scopes if mcp_oauth_metadata else None + gated_oauth_metadata.scopes if gated_oauth_metadata else None ) resolved_authorization_url = server_config.get("authorization_url") or ( - mcp_oauth_metadata.authorization_url if mcp_oauth_metadata else None + gated_oauth_metadata.authorization_url if gated_oauth_metadata else None ) resolved_token_url = server_config.get("token_url") or ( - mcp_oauth_metadata.token_url if mcp_oauth_metadata else None + gated_oauth_metadata.token_url if gated_oauth_metadata else None ) resolved_registration_url = server_config.get("registration_url") or ( - mcp_oauth_metadata.registration_url if mcp_oauth_metadata else None + gated_oauth_metadata.registration_url if gated_oauth_metadata else None ) config_oauth2_flow = server_config.get("oauth2_flow", None) @@ -1472,8 +1534,18 @@ class MCPServerManager: mcp_server.server_id, server_url, ) + gated_oauth_metadata = ( + _gate_discovered_endpoints_against_manual_authorization_url( + mcp_oauth_metadata, + mcp_server.authorization_url, + mcp_server.server_id, + bool(getattr(mcp_server, "dcr_bridge", None)), + ) + if auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES + else mcp_oauth_metadata + ) - resolved_scopes = scopes or (mcp_oauth_metadata.scopes if mcp_oauth_metadata else None) + resolved_scopes = scopes or (gated_oauth_metadata.scopes if gated_oauth_metadata else None) new_server = MCPServer( server_id=mcp_server.server_id, @@ -1493,9 +1565,9 @@ class MCPServerManager: client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)), scopes=resolved_scopes, - authorization_url=mcp_server.authorization_url or getattr(mcp_oauth_metadata, "authorization_url", None), - token_url=mcp_server.token_url or getattr(mcp_oauth_metadata, "token_url", None), - registration_url=mcp_server.registration_url or getattr(mcp_oauth_metadata, "registration_url", None), + authorization_url=mcp_server.authorization_url or getattr(gated_oauth_metadata, "authorization_url", None), + token_url=mcp_server.token_url or getattr(gated_oauth_metadata, "token_url", None), + registration_url=mcp_server.registration_url or getattr(gated_oauth_metadata, "registration_url", None), token_endpoint_auth_method=( credentials_dict.get("token_endpoint_auth_method") if credentials_dict else None ), @@ -1555,7 +1627,7 @@ class MCPServerManager: existing_authorization_url=mcp_server.authorization_url, existing_token_url=mcp_server.token_url, existing_scopes=scopes, - metadata=mcp_oauth_metadata, + metadata=gated_oauth_metadata, ) return new_server 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 6af86b9bda3..8091675090e 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 @@ -356,6 +356,55 @@ class TestMCPServerManager: assert server.oauth2_flow == "authorization_code" assert server.needs_user_oauth_token is True + @pytest.mark.asyncio + async def test_load_servers_from_config_rejects_discovered_token_url_on_authorization_endpoint_mismatch(self): + """The config loader always runs discovery and or-merges per field, so a yaml server with a + manual authorization_url has the same config-time mix-up exposure as a DB row: a discovered + token_url from a document advertising a different authorize endpoint must not be combined + with the pinned one. Scopes still backfill.""" + manager = MCPServerManager() + + metadata = MCPOAuthMetadata( + authorization_url="https://attacker.example.com/authorize", + token_url="https://attacker.example.com/token", + scopes=["read"], + ) + config = self._oauth2_config( + oauth2_flow="authorization_code", + authorization_url="https://idp.example.com/authorize", + token_url=None, + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + await manager.load_servers_from_config(config) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.authorization_url == "https://idp.example.com/authorize" + assert server.token_url is None + assert server.scopes == ["read"] + + @pytest.mark.asyncio + async def test_load_servers_from_config_fills_token_url_when_metadata_corroborates_manual_authorization_url(self): + """Corroborated metadata keeps the self-heal on the config path: when the discovered + document advertises the same authorize endpoint the admin pinned, its token_url fills the + blank field.""" + manager = MCPServerManager() + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["read"], + ) + config = self._oauth2_config( + oauth2_flow="authorization_code", + authorization_url="https://idp.example.com/authorize/", + token_url=None, + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + await manager.load_servers_from_config(config) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.token_url == "https://idp.example.com/token" + @pytest.mark.asyncio async def test_load_servers_from_config_non_oauth2_needs_no_flow(self): manager = MCPServerManager() @@ -1056,7 +1105,9 @@ class TestMCPServerManager: async def test_build_from_table_discovers_scopes_when_authorization_url_is_manual(self): """An admin-typed authorization_url must not switch off discovery for the fields left blank: without the scopes_supported backfill the authorize redirect goes out scope-less - and IdPs like Google hard-fail it with 400 "Missing required parameter: scope".""" + and IdPs like Google hard-fail it with 400 "Missing required parameter: scope". Scope + backfill works even when the advertised authorization_endpoint differs from the manual + value, because scopes only steer the redirect to the trusted authorize endpoint.""" manager = MCPServerManager() row = LiteLLM_MCPServerTable( server_id="manual-auth-url-1", @@ -1081,9 +1132,90 @@ class TestMCPServerManager: mock_discovery.assert_awaited_once() assert built.authorization_url == "https://idp.example.com/manual-authorize" - assert built.token_url == "https://idp.example.com/token" + assert built.token_url is None assert built.scopes == ["calendar.read", "calendar.write"] + @pytest.mark.asyncio + async def test_build_from_table_fills_endpoints_when_metadata_corroborates_manual_authorization_url(self): + """A discovered token_url is only trusted next to a manual authorization_url when the same + metadata document advertises that authorize endpoint, and the comparison must tolerate + formatting-only differences (host case, trailing slash, query params like ?prompt=consent) + so hand-copied URLs still self-heal.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="manual-auth-url-2", + alias="manual_auth_url_match", + description="manual authorization_url matching discovery, blank token_url", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://IDP.example.com/authorize/?prompt=consent", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + scopes=["read"], + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + assert built.authorization_url == "https://IDP.example.com/authorize/?prompt=consent" + assert built.token_url == "https://idp.example.com/token" + assert built.registration_url == "https://idp.example.com/register" + assert built.scopes == ["read"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "advertised_authorization_url", + ["https://attacker.example.com/authorize", None], + ) + async def test_build_from_table_rejects_discovered_token_url_on_authorization_endpoint_mismatch( + self, advertised_authorization_url + ): + """Resource-rooted discovery lets a compromised upstream advertise its own authorization + server. With a manual authorization_url pinned, accepting that document's token_url would + send the authorization code, stored client secret, and PKCE verifier to the attacker's + token endpoint (config-time RFC 9700 mix-up), and the persist hook would make the hostile + endpoint durable. Both the in-memory merge and the persisted metadata must drop the + uncorroborated token_url and registration_url.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="manual-auth-url-3", + alias="manual_auth_url_mismatch", + description="manual authorization_url, hostile discovery document", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/authorize", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + metadata = MCPOAuthMetadata( + authorization_url=advertised_authorization_url, + token_url="https://attacker.example.com/token", + registration_url="https://attacker.example.com/register", + scopes=["read"], + ) + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)), + patch.object(manager, "_persist_discovered_oauth_endpoints", new=AsyncMock()) as mock_persist, + ): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + assert built.authorization_url == "https://idp.example.com/authorize" + assert built.token_url is None + assert built.registration_url is None + assert built.scopes == ["read"] + persisted_metadata = mock_persist.await_args.kwargs["metadata"] + assert persisted_metadata.token_url is None + assert persisted_metadata.registration_url is None + assert persisted_metadata.scopes == ["read"] + @pytest.mark.asyncio async def test_build_from_table_skips_discovery_when_all_upstream_oauth_fields_present(self): """A fully hand-configured server (authorization_url, token_url, and scopes all set) has @@ -2310,6 +2442,10 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_load_servers_from_config_overrides_discovery_metadata(self): + """Config values win per field. The discovered token_url/registration_url do NOT fill the + blanks here: the document advertises a different authorization_endpoint than the manually + configured one, so combining its endpoints with the pinned authorize URL would be the + config-time mix-up the discovery gate exists to prevent.""" manager = MCPServerManager() discovered_metadata = MCPOAuthMetadata( @@ -2343,8 +2479,8 @@ class TestMCPServerManager: server = next(iter(manager.config_mcp_servers.values())) assert server.scopes == ["config"] # config overrides discovery assert server.authorization_url == "https://config.example.com/auth" - assert server.token_url == "https://discovered.example.com/token" - assert server.registration_url == "https://discovered.example.com/register" + assert server.token_url is None + assert server.registration_url is None @pytest.mark.asyncio async def test_load_servers_from_config_filters_blank_scopes(self): From 447d50fa4041c3d6dcbfdaa6fa4b6f89ea2eaaee Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 10:39:04 -0700 Subject: [PATCH 17/62] fix(mcp): enforce the OAuth endpoint trust rule at carry-forward too, elide default port The corroboration check belongs to adopting a token_url from any non-manual source, not to discovery alone. Carry-forward is the other such source: it copied a prior registry entry's token_url/registration_url onto a rebuild whose authorization_url had been re-pointed to a different server, reviving an uncorroborated token endpoint the discovery gate would reject. Both sites now share one predicate, _endpoints_corroborate_authorization_url: previous endpoints carry forward only when the previous authorization_url corroborates the authorize endpoint the build will use (absent -> the previous one is adopted too, a consistent group; else it must match). Endpoint comparison now elides the default port so :443 and formatting-only differences still match. --- .../mcp_server/mcp_server_manager.py | 95 ++++++++++++------ .../mcp_server/test_mcp_server_manager.py | 96 +++++++++++++++++++ 2 files changed, 160 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 284bf6a0926..b84c5601560 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -186,6 +186,46 @@ _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: tuple[MCPAuth, ...] = ( ) +def _normalized_authorize_endpoint(url: str) -> str: + """Compare authorize endpoints on scheme, host, and path only. The default port is elided and + the host is lowercased so ``https://IDP.example.com:443/authorize/`` and + ``https://idp.example.com/authorize`` are the same identity; query and trailing slash are not.""" + parsed = urlparse(url) + scheme = parsed.scheme.lower() + host = (parsed.hostname or "").lower() + default_port = {"https": 443, "http": 80}.get(scheme) + try: + port = parsed.port + except ValueError: + port = None + authority = host if port is None or port == default_port else f"{host}:{port}" + return f"{scheme}://{authority}{parsed.path.rstrip('/')}" + + +def _endpoints_corroborate_authorization_url( + source_authorization_url: str | None, + trusted_authorization_url: str | None, +) -> bool: + """Whether a source's ``token_url``/``registration_url`` may be paired with a trusted authorize + endpoint. This is the single trust rule for adopting OAuth endpoints from any non-manual source. + + Discovery is rooted at the MCP resource (RFC 9728), so a compromised upstream can advertise an + attacker-run authorization server. When ``authorization_url`` is admin-pinned, pairing it with a + ``token_url`` from a different source is the RFC 9700 authorization-server mix-up: the user signs + in at the trusted authorize endpoint while the gateway redeems the code, with the stored client + secret and PKCE verifier, at the attacker's token endpoint. Endpoints are trustworthy together + only when they share an authorization server, so a source's endpoints are adopted only when the + same source advertised an ``authorization_endpoint`` matching the pinned value. With no pinned + value (``trusted_authorization_url is None``) there is nothing to protect: the authorize endpoint + comes from the same source as the token endpoint, so they corroborate each other by construction. + """ + if trusted_authorization_url is None: + return True + return bool(source_authorization_url) and _normalized_authorize_endpoint( + source_authorization_url + ) == _normalized_authorize_endpoint(trusted_authorization_url) + + def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_server: MCPServer | None) -> None: """Keep the last known good OAuth endpoints when a rebuild's re-discovery comes back empty. @@ -193,57 +233,50 @@ def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serv during re-discovery downgrades a working server (``authorization_url`` set) to a broken one (``None``, /authorize 400s) with no configuration change. Mirrors the ``short_prefix`` carry-forward. Skipped when the server's ``url`` or ``auth_type`` changed, since the previous - endpoints may then belong to a different upstream. ``registration_url`` IS carried here even - though ``_persist_discovered_oauth_endpoints`` refuses to write it to the row: carrying only - restores the same in-memory value the previous build already ran with, while persisting it - would flip ``_dcr_bridge_relays_client_registration`` (which keys off the stored column) for - dcr_bridge servers that never had one configured. + endpoints may then belong to a different upstream. ``registration_url`` IS carried even though + ``_persist_discovered_oauth_endpoints`` refuses to write it to the row: carrying only restores + the same in-memory value the previous build already ran with, while persisting it would flip + ``_dcr_bridge_relays_client_registration`` (which keys off the stored column) for dcr_bridge + servers that never had one configured. + + Carry-forward is a non-manual endpoint source, so the same trust rule as discovery applies: the + previous ``token_url``/``registration_url`` are carried only when the previous + ``authorization_url`` corroborates the authorize endpoint this build will use, i.e. when the + incoming build has no pinned authorize endpoint (``None`` -> we adopt the previous one too, a + consistent group) or pins the same one. An admin re-pointing ``authorization_url`` to a different + server must not keep serving the old token endpoint. """ if previous_server is None: return if previous_server.url != new_server.url or previous_server.auth_type != new_server.auth_type: return + may_carry_endpoints = _endpoints_corroborate_authorization_url( + previous_server.authorization_url, new_server.authorization_url + ) if new_server.authorization_url is None and previous_server.authorization_url: new_server.authorization_url = previous_server.authorization_url - if new_server.token_url is None and previous_server.token_url: + if may_carry_endpoints and new_server.token_url is None and previous_server.token_url: new_server.token_url = previous_server.token_url - if new_server.registration_url is None and previous_server.registration_url: + if may_carry_endpoints and new_server.registration_url is None and previous_server.registration_url: new_server.registration_url = previous_server.registration_url if not new_server.scopes and previous_server.scopes: new_server.scopes = previous_server.scopes -def _normalized_authorize_endpoint(url: str) -> str: - parsed = urlparse(url) - return f"{parsed.scheme.lower()}://{parsed.netloc.lower()}{parsed.path.rstrip('/')}" - - def _gate_discovered_endpoints_against_manual_authorization_url( metadata: MCPOAuthMetadata | None, manual_authorization_url: str | None, server_identifier: str, is_dcr_bridge: bool, ) -> MCPOAuthMetadata | None: - """Refuse discovered token/registration endpoints the pinned authorize endpoint cannot vouch for. - - Discovery is rooted at the MCP resource (RFC 9728), so a compromised upstream can advertise an - attacker-run authorization server. When the admin manually configured ``authorization_url``, - filling a blank ``token_url`` from that advertisement recreates the RFC 9700 mix-up attack at - configuration time: users sign in at the trusted authorize endpoint while the gateway redeems - the code, with the stored client secret and PKCE verifier, at the attacker's token endpoint. - Endpoints from one metadata document are only trustworthy together, so the discovered - ``token_url`` and ``registration_url`` are accepted only when that same document's - ``authorization_endpoint`` matches the pinned value (scheme+host+path; query and trailing slash - are not identity). Scope discovery stays ungated: scopes steer the redirect to the trusted - authorize endpoint and carry no credentials. + """Apply :func:`_endpoints_corroborate_authorization_url` to freshly discovered metadata, the + other non-manual endpoint source. Drops the discovered ``token_url``/``registration_url`` (never + the scopes, which carry no credentials) when they cannot be vouched for by the pinned authorize + endpoint, logging why so an intentional mismatch can be resolved by setting Token URL by hand. """ - if metadata is None or not manual_authorization_url: + if metadata is None or (not metadata.token_url and not metadata.registration_url): return metadata - if not metadata.token_url and not metadata.registration_url: - return metadata - if metadata.authorization_url and _normalized_authorize_endpoint( - manual_authorization_url - ) == _normalized_authorize_endpoint(metadata.authorization_url): + if _endpoints_corroborate_authorization_url(metadata.authorization_url, manual_authorization_url): return metadata bridge_note = ( " The discovered registration_url is rejected with it, so this dcr_bridge server stays on the" @@ -258,7 +291,7 @@ def _gate_discovered_endpoints_against_manual_authorization_url( "authorization server. Configure Token URL manually if the mismatch is intentional.%s", server_identifier, _normalized_authorize_endpoint(metadata.authorization_url) if metadata.authorization_url else "", - _normalized_authorize_endpoint(manual_authorization_url), + _normalized_authorize_endpoint(manual_authorization_url) if manual_authorization_url else "", bridge_note, ) return metadata.model_copy(update={"token_url": None, "registration_url": None}) 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 8091675090e..229b1bfb014 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 @@ -5287,6 +5287,102 @@ class TestMCPServerTimestamps: _carry_forward_resolved_oauth_endpoints(new_server=explicit, previous_server=previous) assert explicit.authorization_url == "https://configured.example.com/auth" + def test_carry_forward_does_not_revive_token_url_across_authorization_url_change(self): + """Carry-forward is a non-manual endpoint source, so it obeys the same trust rule as + discovery: a previous token_url/registration_url belongs to the previous authorization + server, so it must not be pinned to a NEW authorization_url the admin re-pointed to. Without + this, re-pointing authorize to server B while the same MCP url keeps serving A's token + endpoint recreates the RFC 9700 mix-up, durably, and the discovery gate alone cannot catch + it because the stale endpoint comes from the registry, not from discovery.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _carry_forward_resolved_oauth_endpoints, + ) + + previous = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp-a.example.com/authorize", + token_url="https://idp-a.example.com/token", + registration_url="https://idp-a.example.com/register", + ) + repointed = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp-b.example.com/authorize", + ) + + _carry_forward_resolved_oauth_endpoints(new_server=repointed, previous_server=previous) + + assert repointed.authorization_url == "https://idp-b.example.com/authorize" + assert repointed.token_url is None + assert repointed.registration_url is None + + def test_carry_forward_restores_endpoints_when_authorization_url_unchanged(self): + """The last-known-good path still works: a rebuild whose discovery blipped (no authorize + endpoint) adopts the previous authorize endpoint AND its token endpoint together as a + consistent group, and a rebuild that re-pins the same authorize endpoint (formatting aside) + keeps carrying the corroborated token endpoint.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _carry_forward_resolved_oauth_endpoints, + ) + + def previous() -> MCPServer: + return MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + ) + + blipped = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url=None, + ) + _carry_forward_resolved_oauth_endpoints(new_server=blipped, previous_server=previous()) + assert blipped.authorization_url == "https://idp.example.com/authorize" + assert blipped.token_url == "https://idp.example.com/token" + assert blipped.registration_url == "https://idp.example.com/register" + + same_authorize = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://IDP.example.com:443/authorize/", + ) + _carry_forward_resolved_oauth_endpoints(new_server=same_authorize, previous_server=previous()) + assert same_authorize.token_url == "https://idp.example.com/token" + assert same_authorize.registration_url == "https://idp.example.com/register" + + def test_normalized_authorize_endpoint_treats_default_port_and_slash_as_identity(self): + """The corroboration check must not fail on formatting-only differences an IdP legitimately + emits: default port, trailing slash, host case, and query string are not identity, but a + non-default port is.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _normalized_authorize_endpoint, + ) + + canonical = _normalized_authorize_endpoint("https://idp.example.com/authorize") + assert _normalized_authorize_endpoint("https://idp.example.com:443/authorize") == canonical + assert _normalized_authorize_endpoint("https://IDP.example.com/authorize/") == canonical + assert _normalized_authorize_endpoint("https://idp.example.com/authorize?prompt=consent") == canonical + assert _normalized_authorize_endpoint("https://idp.example.com:8443/authorize") != canonical + def test_build_mcp_server_table_preserves_timestamps(self): """_build_mcp_server_table must use the MCPServer's stored timestamps, not datetime.now().""" manager = MCPServerManager() From 8650f6c7d352ca079fe82fba3a84f1f5adcf3c91 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 10:57:40 -0700 Subject: [PATCH 18/62] fix(mcp): bound pinned-config discovery to the corroborated authorization server, scopes included Provenance is a property of the whole discovered metadata document, not per field. Waving scopes through while gating endpoints left a second inflation vector: a compromised upstream advertises broad scopes via the resource metadata (RFC 9728 / WWW-Authenticate), the gateway requests them from the trusted authorization server, and the resulting token flows back to the upstream. Both that and the token-endpoint mix-up are now one rule: when authorization_url is admin-pinned, discovered token_url/registration_url are kept only if the document corroborates the pin, and scopes come from the authorization server's own scopes_supported (a new authorization_server_scopes field, trusted tier) rather than the resource-advertised scopes. A document that does not corroborate backfills nothing. Blank (empty-string) authorization_url is treated as unpinned so the merge and the gate agree. Carry-forward, the other non-manual source, drops the same three across an authorization_url change. --- .../mcp_server/mcp_server_manager.py | 51 ++++--- .../types/mcp_server/mcp_server_manager.py | 9 ++ .../mcp_server/test_mcp_server_manager.py | 137 ++++++++++++++---- 3 files changed, 147 insertions(+), 50 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b84c5601560..2ace0d9b16a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -219,7 +219,7 @@ def _endpoints_corroborate_authorization_url( value (``trusted_authorization_url is None``) there is nothing to protect: the authorize endpoint comes from the same source as the token endpoint, so they corroborate each other by construction. """ - if trusted_authorization_url is None: + if not (trusted_authorization_url and trusted_authorization_url.strip()): return True return bool(source_authorization_url) and _normalized_authorize_endpoint( source_authorization_url @@ -240,43 +240,53 @@ def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serv servers that never had one configured. Carry-forward is a non-manual endpoint source, so the same trust rule as discovery applies: the - previous ``token_url``/``registration_url`` are carried only when the previous + previous ``token_url``/``registration_url``/``scopes`` are carried only when the previous ``authorization_url`` corroborates the authorize endpoint this build will use, i.e. when the incoming build has no pinned authorize endpoint (``None`` -> we adopt the previous one too, a consistent group) or pins the same one. An admin re-pointing ``authorization_url`` to a different - server must not keep serving the old token endpoint. + server must not keep serving the old server's token endpoint or granted scopes. """ if previous_server is None: return if previous_server.url != new_server.url or previous_server.auth_type != new_server.auth_type: return - may_carry_endpoints = _endpoints_corroborate_authorization_url( + may_carry = _endpoints_corroborate_authorization_url( previous_server.authorization_url, new_server.authorization_url ) if new_server.authorization_url is None and previous_server.authorization_url: new_server.authorization_url = previous_server.authorization_url - if may_carry_endpoints and new_server.token_url is None and previous_server.token_url: + if may_carry and new_server.token_url is None and previous_server.token_url: new_server.token_url = previous_server.token_url - if may_carry_endpoints and new_server.registration_url is None and previous_server.registration_url: + if may_carry and new_server.registration_url is None and previous_server.registration_url: new_server.registration_url = previous_server.registration_url - if not new_server.scopes and previous_server.scopes: + if may_carry and not new_server.scopes and previous_server.scopes: new_server.scopes = previous_server.scopes -def _gate_discovered_endpoints_against_manual_authorization_url( +def _restrict_discovery_to_corroborated_authorization_server( metadata: MCPOAuthMetadata | None, manual_authorization_url: str | None, server_identifier: str, is_dcr_bridge: bool, ) -> MCPOAuthMetadata | None: - """Apply :func:`_endpoints_corroborate_authorization_url` to freshly discovered metadata, the - other non-manual endpoint source. Drops the discovered ``token_url``/``registration_url`` (never - the scopes, which carry no credentials) when they cannot be vouched for by the pinned authorize - endpoint, logging why so an intentional mismatch can be resolved by setting Token URL by hand. + """Bound what freshly discovered metadata may backfill into a manually pinned config. + + Discovery is rooted at the MCP resource, so provenance is a property of the whole metadata + document, not per field: a compromised upstream can advertise both an attacker ``token_endpoint`` + (the RFC 9700 mix-up) and inflated ``scopes`` (tricking the user into granting a broader token + that then flows to the upstream). Both are closed by one rule. When ``authorization_url`` is + admin-pinned, the discovered ``token_url`` and ``registration_url`` are kept only if the document + corroborates the pin (its ``authorization_endpoint`` matches), and scopes are taken from the + authorization server's own ``scopes_supported`` (``authorization_server_scopes``, trusted tier) + rather than the resource-advertised ``scopes`` a compromised upstream controls. A document that + does not corroborate backfills nothing. With no pin there is no trust anchor to protect and the + authorize endpoint comes from the same chain as everything else, so discovery is returned as-is. """ - if metadata is None or (not metadata.token_url and not metadata.registration_url): + if metadata is None or not (manual_authorization_url and manual_authorization_url.strip()): return metadata if _endpoints_corroborate_authorization_url(metadata.authorization_url, manual_authorization_url): + return metadata.model_copy(update={"scopes": metadata.authorization_server_scopes}) + if not metadata.token_url and not metadata.registration_url and not metadata.scopes: return metadata bridge_note = ( " The discovered registration_url is rejected with it, so this dcr_bridge server stays on the" @@ -286,15 +296,15 @@ def _gate_discovered_endpoints_against_manual_authorization_url( ) verbose_logger.warning( "MCP OAuth discovery for server %s advertised authorization_endpoint %s, which does not match the " - "manually configured authorization_url %s; rejecting the discovered token_url/registration_url so " - "authorization codes and client credentials only go to endpoints vouched for by the configured " - "authorization server. Configure Token URL manually if the mismatch is intentional.%s", + "manually configured authorization_url %s; rejecting the discovered token_url/registration_url/scopes " + "so authorization codes, client credentials, and granted scopes only follow the configured " + "authorization server. Configure Token URL and Scopes manually if the mismatch is intentional.%s", server_identifier, _normalized_authorize_endpoint(metadata.authorization_url) if metadata.authorization_url else "", - _normalized_authorize_endpoint(manual_authorization_url) if manual_authorization_url else "", + _normalized_authorize_endpoint(manual_authorization_url), bridge_note, ) - return metadata.model_copy(update={"token_url": None, "registration_url": None}) + return metadata.model_copy(update={"token_url": None, "registration_url": None, "scopes": None}) def invalidate_user_env_vars_cache(user_id: str, server_id: str) -> None: @@ -1126,7 +1136,7 @@ class MCPServerManager: mcp_oauth_metadata = None gated_oauth_metadata = ( - _gate_discovered_endpoints_against_manual_authorization_url( + _restrict_discovery_to_corroborated_authorization_server( mcp_oauth_metadata, server_config.get("authorization_url"), server_name or server_id, @@ -1568,7 +1578,7 @@ class MCPServerManager: server_url, ) gated_oauth_metadata = ( - _gate_discovered_endpoints_against_manual_authorization_url( + _restrict_discovery_to_corroborated_authorization_server( mcp_oauth_metadata, mcp_server.authorization_url, mcp_server.server_id, @@ -3365,6 +3375,7 @@ class MCPServerManager: authorization_url=data.get("authorization_endpoint"), token_url=data.get("token_endpoint"), registration_url=data.get("registration_endpoint"), + authorization_server_scopes=scopes, ) if any( diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 82ff15303f3..79a523f7c8b 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -17,9 +17,18 @@ MCPInfo = Dict[str, Any] class MCPOAuthMetadata(BaseModel): scopes: Optional[List[str]] = None + """Effective scopes, resource-preferred: the RFC 9728 protected-resource advertisement or the + WWW-Authenticate challenge when the resource supplied one, else the authorization server's + ``scopes_supported``. A compromised resource server can influence this, so it must not expand a + manually pinned ``authorization_url`` (see ``authorization_server_scopes``).""" authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + authorization_server_scopes: Optional[List[str]] = None + """The ``scopes_supported`` enumerated by the authorization-server metadata document itself + (RFC 8414), independent of anything the resource server advertised. This is the only scope + source trusted to backfill a manually pinned ``authorization_url``, because it shares provenance + with the ``authorization_endpoint`` used to corroborate that pin.""" from_origin_fallback: bool = False """True when the metadata came from guessing the resource origin as its authorization server rather than from an RFC 9728/8414-advertised document. Guessed endpoints are 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 229b1bfb014..d14ae561f71 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 @@ -357,17 +357,18 @@ class TestMCPServerManager: assert server.needs_user_oauth_token is True @pytest.mark.asyncio - async def test_load_servers_from_config_rejects_discovered_token_url_on_authorization_endpoint_mismatch(self): + async def test_load_servers_from_config_rejects_uncorroborated_discovery_including_scopes(self): """The config loader always runs discovery and or-merges per field, so a yaml server with a - manual authorization_url has the same config-time mix-up exposure as a DB row: a discovered - token_url from a document advertising a different authorize endpoint must not be combined - with the pinned one. Scopes still backfill.""" + manual authorization_url has the same config-time mix-up exposure as a DB row: a document + advertising a different authorize endpoint backfills nothing, neither its token_url nor its + scopes.""" manager = MCPServerManager() metadata = MCPOAuthMetadata( authorization_url="https://attacker.example.com/authorize", token_url="https://attacker.example.com/token", - scopes=["read"], + scopes=["read", "admin"], + authorization_server_scopes=["read", "admin"], ) config = self._oauth2_config( oauth2_flow="authorization_code", @@ -380,19 +381,20 @@ class TestMCPServerManager: server = next(iter(manager.config_mcp_servers.values())) assert server.authorization_url == "https://idp.example.com/authorize" assert server.token_url is None - assert server.scopes == ["read"] + assert server.scopes is None @pytest.mark.asyncio async def test_load_servers_from_config_fills_token_url_when_metadata_corroborates_manual_authorization_url(self): """Corroborated metadata keeps the self-heal on the config path: when the discovered document advertises the same authorize endpoint the admin pinned, its token_url fills the - blank field.""" + blank field and scopes come from the authorization server's own scopes_supported.""" manager = MCPServerManager() metadata = MCPOAuthMetadata( authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", - scopes=["read"], + scopes=["read", "admin"], + authorization_server_scopes=["read"], ) config = self._oauth2_config( oauth2_flow="authorization_code", @@ -404,6 +406,34 @@ class TestMCPServerManager: server = next(iter(manager.config_mcp_servers.values())) assert server.token_url == "https://idp.example.com/token" + assert server.scopes == ["read"] + + @pytest.mark.asyncio + async def test_load_servers_from_config_blank_authorization_url_is_not_a_pin(self): + """A blank (empty-string) authorization_url is not a trust anchor, so discovery backfills + the whole set — authorize endpoint, token_url, and its resource-preferred scopes — from the + same chain, exactly as if the field had been omitted. The corroboration gate must treat + empty-string as unpinned so it does not strand the token_url the merge still fills.""" + manager = MCPServerManager() + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["read"], + authorization_server_scopes=["read"], + ) + config = self._oauth2_config( + oauth2_flow="authorization_code", + authorization_url="", + token_url=None, + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + await manager.load_servers_from_config(config) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.authorization_url == "https://idp.example.com/authorize" + assert server.token_url == "https://idp.example.com/token" + assert server.scopes == ["read"] @pytest.mark.asyncio async def test_load_servers_from_config_non_oauth2_needs_no_flow(self): @@ -1102,12 +1132,12 @@ class TestMCPServerManager: assert built.token_url == "https://idp.example.com/token" @pytest.mark.asyncio - async def test_build_from_table_discovers_scopes_when_authorization_url_is_manual(self): - """An admin-typed authorization_url must not switch off discovery for the fields left - blank: without the scopes_supported backfill the authorize redirect goes out scope-less - and IdPs like Google hard-fail it with 400 "Missing required parameter: scope". Scope - backfill works even when the advertised authorization_endpoint differs from the manual - value, because scopes only steer the redirect to the trusted authorize endpoint.""" + async def test_build_from_table_backfills_scopes_from_authorization_server_not_resource(self): + """When authorization_url is admin-pinned, scopes backfill from the authorization server's + own scopes_supported (trusted tier), never from the resource-advertised scopes a compromised + upstream controls. Here the corroborating document carries an inflated resource `scopes` + (`admin`) alongside the real authorization_server_scopes; only the latter may be requested, + otherwise a hostile resource could trick the user into granting a broader token.""" manager = MCPServerManager() row = LiteLLM_MCPServerTable( server_id="manual-auth-url-1", @@ -1116,24 +1146,24 @@ class TestMCPServerManager: url="https://up.example.com/mcp", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, - authorization_url="https://idp.example.com/manual-authorize", + authorization_url="https://idp.example.com/authorize", created_at=datetime.now(), updated_at=datetime.now(), ) metadata = MCPOAuthMetadata( - authorization_url="https://idp.example.com/discovered-authorize", + authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", - registration_url=None, - scopes=["calendar.read", "calendar.write"], + scopes=["read", "admin"], + authorization_server_scopes=["read", "write"], ) with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)) as mock_discovery: built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) mock_discovery.assert_awaited_once() - assert built.authorization_url == "https://idp.example.com/manual-authorize" - assert built.token_url is None - assert built.scopes == ["calendar.read", "calendar.write"] + assert built.authorization_url == "https://idp.example.com/authorize" + assert built.token_url == "https://idp.example.com/token" + assert built.scopes == ["read", "write"] @pytest.mark.asyncio async def test_build_from_table_fills_endpoints_when_metadata_corroborates_manual_authorization_url(self): @@ -1159,6 +1189,7 @@ class TestMCPServerManager: token_url="https://idp.example.com/token", registration_url="https://idp.example.com/register", scopes=["read"], + authorization_server_scopes=["read"], ) with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) @@ -1173,15 +1204,15 @@ class TestMCPServerManager: "advertised_authorization_url", ["https://attacker.example.com/authorize", None], ) - async def test_build_from_table_rejects_discovered_token_url_on_authorization_endpoint_mismatch( + async def test_build_from_table_rejects_uncorroborated_discovery_including_scopes( self, advertised_authorization_url ): """Resource-rooted discovery lets a compromised upstream advertise its own authorization - server. With a manual authorization_url pinned, accepting that document's token_url would - send the authorization code, stored client secret, and PKCE verifier to the attacker's - token endpoint (config-time RFC 9700 mix-up), and the persist hook would make the hostile - endpoint durable. Both the in-memory merge and the persisted metadata must drop the - uncorroborated token_url and registration_url.""" + server. With a manual authorization_url pinned, a document that does not corroborate it + backfills nothing: accepting its token_url would send the code, client secret, and PKCE + verifier to the attacker (config-time RFC 9700 mix-up), and accepting its scopes would let + the upstream inflate the granted token. Both the in-memory merge and the persisted metadata + must drop the uncorroborated token_url, registration_url, and scopes.""" manager = MCPServerManager() row = LiteLLM_MCPServerTable( server_id="manual-auth-url-3", @@ -1199,7 +1230,8 @@ class TestMCPServerManager: authorization_url=advertised_authorization_url, token_url="https://attacker.example.com/token", registration_url="https://attacker.example.com/register", - scopes=["read"], + scopes=["read", "admin"], + authorization_server_scopes=["read", "admin"], ) with ( patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)), @@ -1210,11 +1242,11 @@ class TestMCPServerManager: assert built.authorization_url == "https://idp.example.com/authorize" assert built.token_url is None assert built.registration_url is None - assert built.scopes == ["read"] + assert built.scopes is None persisted_metadata = mock_persist.await_args.kwargs["metadata"] assert persisted_metadata.token_url is None assert persisted_metadata.registration_url is None - assert persisted_metadata.scopes == ["read"] + assert persisted_metadata.scopes is None @pytest.mark.asyncio async def test_build_from_table_skips_discovery_when_all_upstream_oauth_fields_present(self): @@ -2317,6 +2349,48 @@ class TestMCPServerManager: assert result.scopes == ["api://some-scope/.default"] assert result.from_origin_fallback is False + @pytest.mark.asyncio + async def test_descovery_metadata_preserves_authorization_server_scopes_under_resource_override(self): + """The effective `scopes` field is resource-preferred (RFC 9728 / WWW-Authenticate), but the + authorization server's own scopes_supported must survive on `authorization_server_scopes` so + the pinned-config backfill can request the trusted-tier scopes instead of resource-advertised + ones. This is the provenance split the scope-inflation defense depends on.""" + manager = MCPServerManager() + + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + authorization_server_metadata = MCPOAuthMetadata( + scopes=["as.read", "as.write"], + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + authorization_server_scopes=["as.read", "as.write"], + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ), + patch.object( + manager, + "_attempt_well_known_discovery", + AsyncMock(return_value=(["https://idp.example.com"], ["resource.only"])), + ), + patch.object( + manager, + "_fetch_authorization_server_metadata", + AsyncMock(return_value=authorization_server_metadata), + ), + ): + result = await manager._descovery_metadata("https://up.example.com/mcp") + + assert result is not None + assert result.scopes == ["resource.only"] + assert result.authorization_server_scopes == ["as.read", "as.write"] + @pytest.mark.asyncio async def test_fetch_single_authorization_server_metadata_supports_azure_issuer_path( self, @@ -2357,6 +2431,9 @@ class TestMCPServerManager: assert result.authorization_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" assert result.token_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" assert result.scopes == ["api://some-scope/.default"] + # The authorization server's own scopes_supported is retained under a dedicated field so a + # later resource-scope override cannot erase the trusted-tier value used to backfill a pin. + assert result.authorization_server_scopes == ["api://some-scope/.default"] @pytest.mark.asyncio async def test_fetch_single_authorization_server_metadata_derives_azure_metadata( From feedab214ee4cdbb19b0a24200004f9e5db1f080 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 11:12:54 -0700 Subject: [PATCH 19/62] fix(mcp): normalize blank OAuth endpoint fields to None at build entry points A whitespace-only authorization_url was truthy to the row/config merges and has_all check but blank to the corroboration gate, so discovery and carry-forward adopted token_url/registration_url/scopes as if unpinned while the broken whitespace value was still used for redirects. Rather than add another strip() at each site, the pinned authorization_url/token_url/ registration_url are normalized once per build path (DB and config) via _blank_to_none, so the merge, has_all gate, discovery gate, persist hook, and carry-forward all see a single notion of blank. Empty and whitespace pins now behave identically to an omitted field. --- .../mcp_server/mcp_server_manager.py | 51 +++++++++++++------ .../mcp_server/test_mcp_server_manager.py | 46 ++++++++++++++--- 2 files changed, 75 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 2ace0d9b16a..e011c8f8e88 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -186,6 +186,21 @@ _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: tuple[MCPAuth, ...] = ( ) +def _blank_to_none(value: str | None) -> str | None: + """Collapse an absent, empty, or whitespace-only string to ``None``. + + OAuth endpoint fields are consumed by truthiness-based merges (``row or discovered``) and by the + corroboration gate. A whitespace-only value is truthy to ``or`` but is not a usable endpoint, so + without this the merge would keep the blank value for redirects while the gate treats it as + unpinned and backfills the other fields, yielding a broken half-discovered config. Normalizing + the pinned fields once, at each build entry point, gives every downstream consumer a single + notion of "blank" so those code paths cannot disagree. + """ + if not isinstance(value, str): + return None + return value.strip() or None + + def _normalized_authorize_endpoint(url: str) -> str: """Compare authorize endpoints on scheme, host, and path only. The default port is elided and the host is lowercased so ``https://IDP.example.com:443/authorize/`` and @@ -1120,12 +1135,15 @@ class MCPServerManager: ) auth_type = server_config.get("auth_type", None) + manual_authorization_url = _blank_to_none(server_config.get("authorization_url")) + manual_token_url = _blank_to_none(server_config.get("token_url")) + manual_registration_url = _blank_to_none(server_config.get("registration_url")) if server_url and ( auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES or self._obo_needs_endpoint_discovery( auth_type, server_config.get("token_exchange_endpoint"), - server_config.get("token_url"), + manual_token_url, ) ): mcp_oauth_metadata = await self._descovery_metadata( @@ -1138,7 +1156,7 @@ class MCPServerManager: gated_oauth_metadata = ( _restrict_discovery_to_corroborated_authorization_server( mcp_oauth_metadata, - server_config.get("authorization_url"), + manual_authorization_url, server_name or server_id, bool(server_config.get("dcr_bridge")), ) @@ -1152,13 +1170,11 @@ class MCPServerManager: resolved_scopes = self._extract_scopes(server_config.get("scopes")) or ( gated_oauth_metadata.scopes if gated_oauth_metadata else None ) - resolved_authorization_url = server_config.get("authorization_url") or ( + resolved_authorization_url = manual_authorization_url or ( gated_oauth_metadata.authorization_url if gated_oauth_metadata else None ) - resolved_token_url = server_config.get("token_url") or ( - gated_oauth_metadata.token_url if gated_oauth_metadata else None - ) - resolved_registration_url = server_config.get("registration_url") or ( + resolved_token_url = manual_token_url or (gated_oauth_metadata.token_url if gated_oauth_metadata else None) + resolved_registration_url = manual_registration_url or ( gated_oauth_metadata.registration_url if gated_oauth_metadata else None ) @@ -1552,14 +1568,17 @@ class MCPServerManager: auth_type = cast(MCPAuthType, mcp_server.auth_type) server_url = mcp_server.url - has_all_upstream_oauth_fields = bool(mcp_server.authorization_url and mcp_server.token_url and scopes) + manual_authorization_url = _blank_to_none(mcp_server.authorization_url) + manual_token_url = _blank_to_none(mcp_server.token_url) + manual_registration_url = _blank_to_none(mcp_server.registration_url) + has_all_upstream_oauth_fields = bool(manual_authorization_url and manual_token_url and scopes) needs_discovery = bool(server_url) and ( (auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and not has_all_upstream_oauth_fields) or self._obo_needs_endpoint_discovery( auth_type, mcp_server.token_exchange_endpoint or (credentials_dict.get("token_exchange_endpoint") if credentials_dict else None), - mcp_server.token_url, + manual_token_url, ) ) mcp_oauth_metadata = ( @@ -1580,7 +1599,7 @@ class MCPServerManager: gated_oauth_metadata = ( _restrict_discovery_to_corroborated_authorization_server( mcp_oauth_metadata, - mcp_server.authorization_url, + manual_authorization_url, mcp_server.server_id, bool(getattr(mcp_server, "dcr_bridge", None)), ) @@ -1608,9 +1627,9 @@ class MCPServerManager: client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)), scopes=resolved_scopes, - authorization_url=mcp_server.authorization_url or getattr(gated_oauth_metadata, "authorization_url", None), - token_url=mcp_server.token_url or getattr(gated_oauth_metadata, "token_url", None), - registration_url=mcp_server.registration_url or getattr(gated_oauth_metadata, "registration_url", None), + authorization_url=manual_authorization_url or getattr(gated_oauth_metadata, "authorization_url", None), + token_url=manual_token_url or getattr(gated_oauth_metadata, "token_url", None), + registration_url=manual_registration_url or getattr(gated_oauth_metadata, "registration_url", None), token_endpoint_auth_method=( credentials_dict.get("token_endpoint_auth_method") if credentials_dict else None ), @@ -1661,14 +1680,14 @@ class MCPServerManager: await self._persist_discovered_obo_token_url( server_id=mcp_server.server_id, auth_type=auth_type, - existing_token_url=mcp_server.token_url, + existing_token_url=manual_token_url, discovered_token_url=new_server.token_url, ) await self._persist_discovered_oauth_endpoints( server_id=mcp_server.server_id, auth_type=auth_type, - existing_authorization_url=mcp_server.authorization_url, - existing_token_url=mcp_server.token_url, + existing_authorization_url=manual_authorization_url, + existing_token_url=manual_token_url, existing_scopes=scopes, metadata=gated_oauth_metadata, ) 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 d14ae561f71..989281c27e9 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 @@ -409,11 +409,13 @@ class TestMCPServerManager: assert server.scopes == ["read"] @pytest.mark.asyncio - async def test_load_servers_from_config_blank_authorization_url_is_not_a_pin(self): - """A blank (empty-string) authorization_url is not a trust anchor, so discovery backfills - the whole set — authorize endpoint, token_url, and its resource-preferred scopes — from the - same chain, exactly as if the field had been omitted. The corroboration gate must treat - empty-string as unpinned so it does not strand the token_url the merge still fills.""" + @pytest.mark.parametrize("blank_authorization_url", ["", " "]) + async def test_load_servers_from_config_blank_authorization_url_is_not_a_pin(self, blank_authorization_url): + """A blank authorization_url — empty or whitespace-only — is not a trust anchor, so discovery + backfills the whole set (authorize endpoint, token_url, and its resource-preferred scopes) + from the same chain, exactly as if the field had been omitted. The merge and the corroboration + gate must agree that blank means unpinned; a whitespace value that the merge kept for redirects + while the gate treated as unpinned would strand a broken half-discovered config.""" manager = MCPServerManager() metadata = MCPOAuthMetadata( @@ -424,7 +426,7 @@ class TestMCPServerManager: ) config = self._oauth2_config( oauth2_flow="authorization_code", - authorization_url="", + authorization_url=blank_authorization_url, token_url=None, ) with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): @@ -1165,6 +1167,38 @@ class TestMCPServerManager: assert built.token_url == "https://idp.example.com/token" assert built.scopes == ["read", "write"] + @pytest.mark.asyncio + async def test_build_from_table_whitespace_authorization_url_is_not_a_pin(self): + """A whitespace-only authorization_url on the row must not be kept for redirects while the + gate treats it as unpinned. It is normalized to unpinned everywhere, so the built server + takes the discovered authorize endpoint, token_url, and scopes as one consistent group + rather than serving the whitespace value with half-discovered fields.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="whitespace-auth-url", + alias="whitespace_auth_url", + description="whitespace authorization_url is not a pin", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url=" ", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["read"], + authorization_server_scopes=["read"], + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + assert built.authorization_url == "https://idp.example.com/authorize" + assert built.token_url == "https://idp.example.com/token" + assert built.scopes == ["read"] + @pytest.mark.asyncio async def test_build_from_table_fills_endpoints_when_metadata_corroborates_manual_authorization_url(self): """A discovered token_url is only trusted next to a manual authorization_url when the same From 2b9b681e124c27dd5af1064f8c8100a74ccdc2f5 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 15 Jul 2026 13:48:27 -0700 Subject: [PATCH 20/62] fix(ui/chat): resolve chat routes at render time so navigation works under server_root_path CHAT_ROUTES was built once at module load via migratedHref, capturing an empty server root path before the UI-config bootstrap sets it. Under SERVER_ROOT_PATH every chat route came out unprefixed, so router.push hard-navigated to a 404; during a send that page unload also aborted the streaming request, so the first message of a new conversation never rendered and only a second send appeared to work. Compute the routes at render time via getChatRoutes(), update the send URL with a shallow history.pushState instead of a router navigation, and source the active conversation id from the hook's local state so it propagates without a router round-trip --- ui/litellm-dashboard/src/app/chat/page.tsx | 9 ++-- .../chat/ChatShell.serverRootPath.test.ts | 43 ++++++++++++++++++ .../src/components/chat/ChatShell.tsx | 45 ++++++++++--------- .../src/components/chat/useChatHistory.ts | 2 + .../src/contexts/ChatShellContext.tsx | 7 +-- 5 files changed, 77 insertions(+), 29 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/chat/ChatShell.serverRootPath.test.ts diff --git a/ui/litellm-dashboard/src/app/chat/page.tsx b/ui/litellm-dashboard/src/app/chat/page.tsx index ae86ca6e926..b6dccef47c6 100644 --- a/ui/litellm-dashboard/src/app/chat/page.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.tsx @@ -10,7 +10,7 @@ import { Button } from "@/components/ui/button"; import MessageManager from "@/components/molecules/message_manager"; import { useRouter } from "next/navigation"; import { useChatShell } from "@/contexts/ChatShellContext"; -import { CHAT_ROUTES } from "@/components/chat/ChatShell"; +import { getChatRoutes } from "@/components/chat/ChatShell"; import ChatMessages from "@/components/chat/ChatMessages"; import MCPConnectPicker from "@/components/chat/MCPConnectPicker"; import { fetchAvailableModels } from "@/components/llm_calls/fetch_models"; @@ -87,7 +87,7 @@ export default function ChatConversationPage() { const streamScrollLock = useRef(null); useEffect(() => { - if (staleId) router.replace(CHAT_ROUTES.chats); + if (staleId) router.replace(getChatRoutes().chats); }, [staleId, router]); // Load models @@ -140,7 +140,7 @@ export default function ChatConversationPage() { if (!convId) { convId = createConversation(model); setResponsesSessionId(null); // new conversation starts a fresh session - router.push(`${CHAT_ROUTES.chats}?id=${convId}`); + window.history.pushState(null, "", `${window.location.pathname}?id=${convId}`); } appendMessage(convId, { role: "user", content: trimmed }); @@ -248,7 +248,6 @@ export default function ChatConversationPage() { createConversation, appendMessage, updateLastAssistantMessage, - router, isStreaming, responsesSessionId, ], @@ -529,7 +528,7 @@ export default function ChatConversationPage() { Chat with 100+ LLMs + MCP tools; authenticate once, use them here.{" "} @@ -85,32 +88,32 @@ const ChatShell: React.FC = ({ children }) => { } label="Chats" - onClick={() => router.push(CHAT_ROUTES.chats)} + onClick={() => router.push(routes.chats)} active={isChatsRoute} /> } label="Integrations" - onClick={() => router.push(CHAT_ROUTES.integrations)} - active={pathname === CHAT_ROUTES.integrations} + onClick={() => router.push(routes.integrations)} + active={pathname === routes.integrations} /> } label="Credentials" - onClick={() => router.push(CHAT_ROUTES.credentials)} - active={pathname === CHAT_ROUTES.credentials} + onClick={() => router.push(routes.credentials)} + active={pathname === routes.credentials} /> } label="API Keys" - onClick={() => router.push(CHAT_ROUTES.apiKeys)} - active={pathname === CHAT_ROUTES.apiKeys} + onClick={() => router.push(routes.apiKeys)} + active={pathname === routes.apiKeys} /> } label="Usage" - onClick={() => router.push(CHAT_ROUTES.usage)} - active={pathname === CHAT_ROUTES.usage} + onClick={() => router.push(routes.usage)} + active={pathname === routes.usage} />
@@ -120,10 +123,10 @@ const ChatShell: React.FC = ({ children }) => { router.push(`${CHAT_ROUTES.chats}?id=${id}`)} + onSelect={(id) => router.push(`${routes.chats}?id=${id}`)} onDelete={(id) => { deleteConversation(id); - if (id === activeConversationId) router.push(CHAT_ROUTES.chats); + if (id === activeConversationId) router.push(routes.chats); }} onRename={renameConversation} /> diff --git a/ui/litellm-dashboard/src/components/chat/useChatHistory.ts b/ui/litellm-dashboard/src/components/chat/useChatHistory.ts index 3b6d350520b..8e6c3ce5877 100644 --- a/ui/litellm-dashboard/src/components/chat/useChatHistory.ts +++ b/ui/litellm-dashboard/src/components/chat/useChatHistory.ts @@ -52,6 +52,7 @@ export function useChatHistory( ): { conversations: Conversation[]; activeConversation: Conversation | null; + currentActiveId: string | null; storageUnavailable: boolean; staleId: boolean; createConversation: (model: string) => string; @@ -208,6 +209,7 @@ export function useChatHistory( return { conversations, activeConversation, + currentActiveId, storageUnavailable, staleId, createConversation, diff --git a/ui/litellm-dashboard/src/contexts/ChatShellContext.tsx b/ui/litellm-dashboard/src/contexts/ChatShellContext.tsx index 96d352ade6f..b0f590d2394 100644 --- a/ui/litellm-dashboard/src/contexts/ChatShellContext.tsx +++ b/ui/litellm-dashboard/src/contexts/ChatShellContext.tsx @@ -57,12 +57,13 @@ export function ChatShellProvider({ children, }: ChatShellProviderProps) { const searchParams = useSearchParams(); - const activeConversationId = searchParams.get("id"); + const urlConversationId = searchParams.get("id"); const [selectedMCPServers, setSelectedMCPServers] = useState([]); const { conversations, activeConversation, + currentActiveId, storageUnavailable, staleId, createConversation, @@ -71,7 +72,7 @@ export function ChatShellProvider({ truncateFromMessage, deleteConversation, renameConversation, - } = useChatHistory(activeConversationId, userId); + } = useChatHistory(urlConversationId, userId); return ( Date: Wed, 15 Jul 2026 16:23:26 -0700 Subject: [PATCH 21/62] test(claude_code): rename misleading REPO_ROOT to SUITE_ROOT in test_v0_layout --- .../_builder_unit_tests/test_v0_layout.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py b/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py index b1745008fac..a3569ebdb49 100644 --- a/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py +++ b/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py @@ -16,8 +16,8 @@ from pathlib import Path import pytest import yaml -REPO_ROOT = Path(__file__).resolve().parents[1] -MANIFEST_PATH = REPO_ROOT / "manifest.yaml" +SUITE_ROOT = Path(__file__).resolve().parents[1] +MANIFEST_PATH = SUITE_ROOT / "manifest.yaml" # The PRD's "Features in v0" section, in row order. EXPECTED_FEATURE_IDS = [ @@ -90,14 +90,14 @@ def test_manifest_every_feature_has_human_readable_name(manifest): @pytest.mark.parametrize("feature_id", EXPECTED_FEATURE_IDS) def test_feature_directory_exists(feature_id): - feature_dir = REPO_ROOT / feature_id + feature_dir = SUITE_ROOT / feature_id assert feature_dir.is_dir(), f"missing feature directory: {feature_dir}" @pytest.mark.parametrize("feature_id", EXPECTED_FEATURE_IDS) @pytest.mark.parametrize("provider", EXPECTED_PROVIDERS) def test_per_provider_test_file_exists(feature_id, provider): - test_file = REPO_ROOT / feature_id / f"test_{provider}.py" + test_file = SUITE_ROOT / feature_id / f"test_{provider}.py" assert test_file.is_file(), f"missing per-provider test file: {test_file}" @@ -106,7 +106,7 @@ def test_feature_directory_has_init_file(feature_id): """Each feature directory needs an __init__.py so pytest collects the per-provider test files as a package — matches the layout established by `basic_messaging_non_streaming/`.""" - init_file = REPO_ROOT / feature_id / "__init__.py" + init_file = SUITE_ROOT / feature_id / "__init__.py" assert init_file.is_file(), f"missing __init__.py: {init_file}" @@ -117,7 +117,7 @@ def test_feature_directory_has_init_file(feature_id): # a broken post-v0 directory still fails CI. @pytest.mark.parametrize("feature_id", ALL_FEATURE_IDS) def test_every_manifest_feature_has_directory(feature_id): - feature_dir = REPO_ROOT / feature_id + feature_dir = SUITE_ROOT / feature_id assert feature_dir.is_dir(), ( f"manifest declares {feature_id!r} but {feature_dir} is missing — " "feature_id MUST match its on-disk directory (see manifest.yaml header)." @@ -126,7 +126,7 @@ def test_every_manifest_feature_has_directory(feature_id): @pytest.mark.parametrize("feature_id", ALL_FEATURE_IDS) def test_every_manifest_feature_has_init_file(feature_id): - init_file = REPO_ROOT / feature_id / "__init__.py" + init_file = SUITE_ROOT / feature_id / "__init__.py" assert init_file.is_file(), f"missing __init__.py: {init_file}" @@ -137,7 +137,7 @@ def test_every_manifest_feature_has_per_provider_test_file(feature_id, provider) backed by a per-provider test file. Without this check, a missing file silently becomes a `not_tested` cell in the published matrix rather than a CI failure surfacing the layout drift.""" - test_file = REPO_ROOT / feature_id / f"test_{provider}.py" + test_file = SUITE_ROOT / feature_id / f"test_{provider}.py" assert test_file.is_file(), f"missing per-provider test file: {test_file}" @@ -151,7 +151,7 @@ def test_per_provider_test_file_imports_and_parametrizes_three_models( use plain aliases or per-provider-suffixed aliases (e.g. `claude-opus-4-7-bedrock-invoke`), so we check for the tier substrings rather than exact alias names.""" - text = (REPO_ROOT / feature_id / f"test_{provider}.py").read_text() + text = (SUITE_ROOT / feature_id / f"test_{provider}.py").read_text() for tier in ("haiku-4-5", "sonnet-4-6", "opus-4-7"): assert ( tier in text @@ -171,7 +171,7 @@ def test_azure_test_file_drives_the_proxy(feature_id): that wraps them — both shapes drive the proxy, and we don't want this layout pin to block legitimate de-duplication of test bodies. """ - text = (REPO_ROOT / feature_id / "test_azure.py").read_text() + text = (SUITE_ROOT / feature_id / "test_azure.py").read_text() assert "run_claude" in text or "run_basic_messaging_cell" in text, ( f"{feature_id}/test_azure.py must drive the claude CLI via run_claude() " "or a shared helper that wraps it; the not_applicable stub was removed " From e4a6516b4916eaae239ba9b147ae71b28b7a0022 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 16:27:15 -0700 Subject: [PATCH 22/62] fix(mcp): keep scope selection resource-driven, not authorization-server-driven MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reverts the over-correction that restricted a pinned-authorization_url server's discovered scopes to the authorization server's own scopes_supported. Per the MCP authorization spec Scope Selection Strategy and RFC 9700 §2.3, the scopes a client requests are resource-driven: the WWW-Authenticate 401 challenge scope, else the RFC 9728 protected-resource scopes_supported. The authorization server's RFC 8414 scopes_supported is a non-exhaustive capability list (the server MAY omit supported scopes) and is never the selection source; scope inflation by a compromised resource is bounded by the authorization server and user consent (RFC 6749 §3.3), not by the client restricting the request. The corroboration gate now rejects only the uncorroborated token_url/registration_url (the RFC 9700 endpoint mix-up) and leaves scopes untouched. Removes the now-unused authorization_server_scopes field. --- .../mcp_server/mcp_server_manager.py | 37 ++++----- .../types/mcp_server/mcp_server_manager.py | 15 ++-- .../mcp_server/test_mcp_server_manager.py | 77 +++++++++---------- 3 files changed, 60 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e011c8f8e88..8b1d00c2855 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -284,24 +284,26 @@ def _restrict_discovery_to_corroborated_authorization_server( server_identifier: str, is_dcr_bridge: bool, ) -> MCPOAuthMetadata | None: - """Bound what freshly discovered metadata may backfill into a manually pinned config. + """Reject discovered token/registration endpoints a manually pinned authorize endpoint cannot + vouch for (the RFC 9700 authorization-server mix-up). - Discovery is rooted at the MCP resource, so provenance is a property of the whole metadata - document, not per field: a compromised upstream can advertise both an attacker ``token_endpoint`` - (the RFC 9700 mix-up) and inflated ``scopes`` (tricking the user into granting a broader token - that then flows to the upstream). Both are closed by one rule. When ``authorization_url`` is - admin-pinned, the discovered ``token_url`` and ``registration_url`` are kept only if the document - corroborates the pin (its ``authorization_endpoint`` matches), and scopes are taken from the - authorization server's own ``scopes_supported`` (``authorization_server_scopes``, trusted tier) - rather than the resource-advertised ``scopes`` a compromised upstream controls. A document that - does not corroborate backfills nothing. With no pin there is no trust anchor to protect and the - authorize endpoint comes from the same chain as everything else, so discovery is returned as-is. + Discovery is rooted at the MCP resource, so a compromised upstream can advertise an attacker + ``token_endpoint``: with ``authorization_url`` admin-pinned but ``token_url`` blank, the merge + would pair the trusted authorize endpoint with that attacker token endpoint, and the gateway would + post the authorization code and client secret there. So the discovered ``token_url`` and + ``registration_url`` are kept only if the document corroborates the pin (its + ``authorization_endpoint`` matches). ``scopes`` are deliberately NOT gated here: per the MCP + authorization spec Scope Selection Strategy and RFC 9700 §2.3, the scopes a client requests are + resource-driven (the WWW-Authenticate challenge or the RFC 9728 protected-resource + ``scopes_supported``), and scope inflation by a compromised resource is bounded by the + authorization server and user consent (RFC 6749 §3.3), not by the client second-guessing the + request. With no pin there is no trust anchor to protect, so discovery is returned as-is. """ if metadata is None or not (manual_authorization_url and manual_authorization_url.strip()): return metadata if _endpoints_corroborate_authorization_url(metadata.authorization_url, manual_authorization_url): - return metadata.model_copy(update={"scopes": metadata.authorization_server_scopes}) - if not metadata.token_url and not metadata.registration_url and not metadata.scopes: + return metadata + if not metadata.token_url and not metadata.registration_url: return metadata bridge_note = ( " The discovered registration_url is rejected with it, so this dcr_bridge server stays on the" @@ -311,15 +313,15 @@ def _restrict_discovery_to_corroborated_authorization_server( ) verbose_logger.warning( "MCP OAuth discovery for server %s advertised authorization_endpoint %s, which does not match the " - "manually configured authorization_url %s; rejecting the discovered token_url/registration_url/scopes " - "so authorization codes, client credentials, and granted scopes only follow the configured " - "authorization server. Configure Token URL and Scopes manually if the mismatch is intentional.%s", + "manually configured authorization_url %s; rejecting the discovered token_url/registration_url so " + "authorization codes and client credentials only follow the configured authorization server. " + "Configure Token URL manually if the mismatch is intentional.%s", server_identifier, _normalized_authorize_endpoint(metadata.authorization_url) if metadata.authorization_url else "", _normalized_authorize_endpoint(manual_authorization_url), bridge_note, ) - return metadata.model_copy(update={"token_url": None, "registration_url": None, "scopes": None}) + return metadata.model_copy(update={"token_url": None, "registration_url": None}) def invalidate_user_env_vars_cache(user_id: str, server_id: str) -> None: @@ -3394,7 +3396,6 @@ class MCPServerManager: authorization_url=data.get("authorization_endpoint"), token_url=data.get("token_endpoint"), registration_url=data.get("registration_endpoint"), - authorization_server_scopes=scopes, ) if any( diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 79a523f7c8b..e5c726296b2 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -17,18 +17,15 @@ MCPInfo = Dict[str, Any] class MCPOAuthMetadata(BaseModel): scopes: Optional[List[str]] = None - """Effective scopes, resource-preferred: the RFC 9728 protected-resource advertisement or the - WWW-Authenticate challenge when the resource supplied one, else the authorization server's - ``scopes_supported``. A compromised resource server can influence this, so it must not expand a - manually pinned ``authorization_url`` (see ``authorization_server_scopes``).""" + """Resource-driven scopes for the authorization request: the RFC 9728 protected-resource + ``scopes_supported``, or the ``scope`` from the WWW-Authenticate 401 challenge when the resource + supplied one, else the authorization server's ``scopes_supported``. This is the scope value a + client requests per the MCP authorization spec Scope Selection Strategy; scope minimization and + inflation control are the authorization server's and user's job at consent (RFC 6749 §3.3), not + the client's.""" authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None - authorization_server_scopes: Optional[List[str]] = None - """The ``scopes_supported`` enumerated by the authorization-server metadata document itself - (RFC 8414), independent of anything the resource server advertised. This is the only scope - source trusted to backfill a manually pinned ``authorization_url``, because it shares provenance - with the ``authorization_endpoint`` used to corroborate that pin.""" from_origin_fallback: bool = False """True when the metadata came from guessing the resource origin as its authorization server rather than from an RFC 9728/8414-advertised document. Guessed endpoints are 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 989281c27e9..bba00ed1819 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 @@ -357,18 +357,18 @@ class TestMCPServerManager: assert server.needs_user_oauth_token is True @pytest.mark.asyncio - async def test_load_servers_from_config_rejects_uncorroborated_discovery_including_scopes(self): - """The config loader always runs discovery and or-merges per field, so a yaml server with a - manual authorization_url has the same config-time mix-up exposure as a DB row: a document - advertising a different authorize endpoint backfills nothing, neither its token_url nor its - scopes.""" + async def test_load_servers_from_config_rejects_uncorroborated_endpoints_but_keeps_resource_scopes(self): + """A yaml server with a manual authorization_url has the same config-time mix-up exposure as a + DB row: a document advertising a different authorize endpoint has its token_url rejected. The + resource-driven scopes are kept, because scope selection is resource-driven (MCP Scope + Selection Strategy) and scope inflation is bounded by the authorization server at consent, not + by dropping scopes when an endpoint mismatches.""" manager = MCPServerManager() metadata = MCPOAuthMetadata( authorization_url="https://attacker.example.com/authorize", token_url="https://attacker.example.com/token", scopes=["read", "admin"], - authorization_server_scopes=["read", "admin"], ) config = self._oauth2_config( oauth2_flow="authorization_code", @@ -381,20 +381,20 @@ class TestMCPServerManager: server = next(iter(manager.config_mcp_servers.values())) assert server.authorization_url == "https://idp.example.com/authorize" assert server.token_url is None - assert server.scopes is None + assert server.scopes == ["read", "admin"] @pytest.mark.asyncio async def test_load_servers_from_config_fills_token_url_when_metadata_corroborates_manual_authorization_url(self): - """Corroborated metadata keeps the self-heal on the config path: when the discovered - document advertises the same authorize endpoint the admin pinned, its token_url fills the - blank field and scopes come from the authorization server's own scopes_supported.""" + """Corroborated metadata keeps the self-heal on the config path: when the discovered document + advertises the same authorize endpoint the admin pinned, its token_url fills the blank field + and scopes come through resource-driven (the discovered document's resource-preferred scopes), + not the authorization server's own capability list.""" manager = MCPServerManager() metadata = MCPOAuthMetadata( authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", scopes=["read", "admin"], - authorization_server_scopes=["read"], ) config = self._oauth2_config( oauth2_flow="authorization_code", @@ -406,7 +406,7 @@ class TestMCPServerManager: server = next(iter(manager.config_mcp_servers.values())) assert server.token_url == "https://idp.example.com/token" - assert server.scopes == ["read"] + assert server.scopes == ["read", "admin"] @pytest.mark.asyncio @pytest.mark.parametrize("blank_authorization_url", ["", " "]) @@ -422,7 +422,6 @@ class TestMCPServerManager: authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", scopes=["read"], - authorization_server_scopes=["read"], ) config = self._oauth2_config( oauth2_flow="authorization_code", @@ -1134,12 +1133,13 @@ class TestMCPServerManager: assert built.token_url == "https://idp.example.com/token" @pytest.mark.asyncio - async def test_build_from_table_backfills_scopes_from_authorization_server_not_resource(self): - """When authorization_url is admin-pinned, scopes backfill from the authorization server's - own scopes_supported (trusted tier), never from the resource-advertised scopes a compromised - upstream controls. Here the corroborating document carries an inflated resource `scopes` - (`admin`) alongside the real authorization_server_scopes; only the latter may be requested, - otherwise a hostile resource could trick the user into granting a broader token.""" + async def test_build_from_table_backfills_resource_driven_scopes_for_pinned_authorization_url(self): + """When authorization_url is admin-pinned and corroborated, scopes backfill as the + resource-driven value (the WWW-Authenticate challenge scope, else the RFC 9728 + protected-resource scopes_supported), per the MCP authorization spec Scope Selection Strategy. + The client does not restrict scopes to the authorization server's own scopes_supported; scope + minimization and inflation control are the authorization server's and user's job at consent + (RFC 6749 §3.3).""" manager = MCPServerManager() row = LiteLLM_MCPServerTable( server_id="manual-auth-url-1", @@ -1157,7 +1157,6 @@ class TestMCPServerManager: authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", scopes=["read", "admin"], - authorization_server_scopes=["read", "write"], ) with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)) as mock_discovery: built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) @@ -1165,7 +1164,7 @@ class TestMCPServerManager: mock_discovery.assert_awaited_once() assert built.authorization_url == "https://idp.example.com/authorize" assert built.token_url == "https://idp.example.com/token" - assert built.scopes == ["read", "write"] + assert built.scopes == ["read", "admin"] @pytest.mark.asyncio async def test_build_from_table_whitespace_authorization_url_is_not_a_pin(self): @@ -1190,7 +1189,6 @@ class TestMCPServerManager: authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", scopes=["read"], - authorization_server_scopes=["read"], ) with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) @@ -1223,7 +1221,6 @@ class TestMCPServerManager: token_url="https://idp.example.com/token", registration_url="https://idp.example.com/register", scopes=["read"], - authorization_server_scopes=["read"], ) with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) @@ -1238,15 +1235,17 @@ class TestMCPServerManager: "advertised_authorization_url", ["https://attacker.example.com/authorize", None], ) - async def test_build_from_table_rejects_uncorroborated_discovery_including_scopes( + async def test_build_from_table_rejects_uncorroborated_endpoints_but_keeps_resource_scopes( self, advertised_authorization_url ): """Resource-rooted discovery lets a compromised upstream advertise its own authorization - server. With a manual authorization_url pinned, a document that does not corroborate it - backfills nothing: accepting its token_url would send the code, client secret, and PKCE - verifier to the attacker (config-time RFC 9700 mix-up), and accepting its scopes would let - the upstream inflate the granted token. Both the in-memory merge and the persisted metadata - must drop the uncorroborated token_url, registration_url, and scopes.""" + server. With a manual authorization_url pinned, a document that does not corroborate it has + its token_url and registration_url dropped: accepting them would send the code, client secret, + and PKCE verifier to the attacker (config-time RFC 9700 mix-up). The resource-driven scopes + are kept, because scope selection is resource-driven (MCP Scope Selection Strategy) and scope + inflation is bounded by the authorization server at consent (RFC 6749 §3.3), not by dropping + scopes on an endpoint mismatch. Both the in-memory merge and the persisted metadata drop only + the uncorroborated endpoints.""" manager = MCPServerManager() row = LiteLLM_MCPServerTable( server_id="manual-auth-url-3", @@ -1265,7 +1264,6 @@ class TestMCPServerManager: token_url="https://attacker.example.com/token", registration_url="https://attacker.example.com/register", scopes=["read", "admin"], - authorization_server_scopes=["read", "admin"], ) with ( patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)), @@ -1276,11 +1274,11 @@ class TestMCPServerManager: assert built.authorization_url == "https://idp.example.com/authorize" assert built.token_url is None assert built.registration_url is None - assert built.scopes is None + assert built.scopes == ["read", "admin"] persisted_metadata = mock_persist.await_args.kwargs["metadata"] assert persisted_metadata.token_url is None assert persisted_metadata.registration_url is None - assert persisted_metadata.scopes is None + assert persisted_metadata.scopes == ["read", "admin"] @pytest.mark.asyncio async def test_build_from_table_skips_discovery_when_all_upstream_oauth_fields_present(self): @@ -2384,11 +2382,11 @@ class TestMCPServerManager: assert result.from_origin_fallback is False @pytest.mark.asyncio - async def test_descovery_metadata_preserves_authorization_server_scopes_under_resource_override(self): - """The effective `scopes` field is resource-preferred (RFC 9728 / WWW-Authenticate), but the - authorization server's own scopes_supported must survive on `authorization_server_scopes` so - the pinned-config backfill can request the trusted-tier scopes instead of resource-advertised - ones. This is the provenance split the scope-inflation defense depends on.""" + async def test_descovery_metadata_scopes_are_resource_driven(self): + """The effective `scopes` are resource-driven: the RFC 9728 protected-resource advertisement + (or WWW-Authenticate challenge) overrides the authorization server's own scopes_supported. This + is the MCP Scope Selection Strategy: the client requests what the resource needs, not the AS's + full capability list.""" manager = MCPServerManager() mock_response = MagicMock() @@ -2400,7 +2398,6 @@ class TestMCPServerManager: scopes=["as.read", "as.write"], authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", - authorization_server_scopes=["as.read", "as.write"], ) with ( @@ -2423,7 +2420,6 @@ class TestMCPServerManager: assert result is not None assert result.scopes == ["resource.only"] - assert result.authorization_server_scopes == ["as.read", "as.write"] @pytest.mark.asyncio async def test_fetch_single_authorization_server_metadata_supports_azure_issuer_path( @@ -2465,9 +2461,6 @@ class TestMCPServerManager: assert result.authorization_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" assert result.token_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" assert result.scopes == ["api://some-scope/.default"] - # The authorization server's own scopes_supported is retained under a dedicated field so a - # later resource-scope override cannot erase the trusted-tier value used to backfill a pin. - assert result.authorization_server_scopes == ["api://some-scope/.default"] @pytest.mark.asyncio async def test_fetch_single_authorization_server_metadata_derives_azure_metadata( From 056f85a71eab22b470ec62df0778e12d805cdf65 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 15 Jul 2026 16:52:21 -0700 Subject: [PATCH 23/62] feat(ui): show exact license expiration date in usage cards --- .../src/components/SidebarUsageCard.test.tsx | 23 ++++++++++++++ .../src/components/SidebarUsageCard.tsx | 15 ++-------- .../src/components/UsageIndicator.test.tsx | 30 ++++++++++++++++++- .../src/components/UsageIndicator.tsx | 18 ++--------- .../src/utils/licenseUtils.test.ts | 30 ++++++++++++++++++- .../src/utils/licenseUtils.ts | 8 +++++ 6 files changed, 94 insertions(+), 30 deletions(-) diff --git a/ui/litellm-dashboard/src/components/SidebarUsageCard.test.tsx b/ui/litellm-dashboard/src/components/SidebarUsageCard.test.tsx index 3244906ed8e..71ed3094a20 100644 --- a/ui/litellm-dashboard/src/components/SidebarUsageCard.test.tsx +++ b/ui/litellm-dashboard/src/components/SidebarUsageCard.test.tsx @@ -128,6 +128,29 @@ describe("SidebarUsageCard", () => { expect(container.querySelector('[data-slot="meter"]')).toBeNull(); }); + it("shows the exact license expiration date as the subtitle instead of time remaining", async () => { + mockUseLicenseInfo.mockReturnValue(licenseResult({ ...ACTIVE_LICENSE, expiration_date: "2099-12-31" })); + + renderWithClient( {}} />); + + expect(await screen.findByText("Expires Dec 31, 2099")).toBeInTheDocument(); + expect(screen.queryByText(/(day|days|month|months) remaining/)).not.toBeInTheDocument(); + }); + + it("shows the exact date as the subtitle when the license is expired", async () => { + mockUseLicenseInfo.mockReturnValue(licenseResult({ ...ACTIVE_LICENSE, expiration_date: "2020-01-01" })); + + renderWithClient( {}} />); + + expect(await screen.findByText("Expired Jan 1, 2020")).toBeInTheDocument(); + }); + + it("falls back to Active plan when the license has no expiration date", async () => { + renderWithClient( {}} />); + + expect(await screen.findByText("Active plan")).toBeInTheDocument(); + }); + it("shows a collapsed rail button that expands the sidebar", async () => { const onExpandRail = vi.fn(); const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/SidebarUsageCard.tsx b/ui/litellm-dashboard/src/components/SidebarUsageCard.tsx index 4fb0e1dac21..2a6d6b43f38 100644 --- a/ui/litellm-dashboard/src/components/SidebarUsageCard.tsx +++ b/ui/litellm-dashboard/src/components/SidebarUsageCard.tsx @@ -1,6 +1,6 @@ import { useDisableUsageIndicator } from "@/app/(dashboard)/hooks/useDisableUsageIndicator"; import { useLicenseInfo } from "@/app/(dashboard)/hooks/license/useLicenseInfo"; -import { getDaysUntilExpiration } from "@/utils/licenseUtils"; +import { formatExpirationStatus } from "@/utils/licenseUtils"; import { Button } from "@/components/ui/button"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { Meter, MeterIndicator, MeterLabel, MeterTrack } from "@/components/ui/meter"; @@ -20,16 +20,6 @@ interface MeterData { total: number; } -const formatExpiration = (daysRemaining: number | null): string => { - if (daysRemaining === null) return "No expiration"; - if (daysRemaining < 0) return "Expired"; - if (daysRemaining === 0) return "Expires today"; - if (daysRemaining === 1) return "1 day remaining"; - if (daysRemaining < 30) return `${daysRemaining} days remaining`; - if (daysRemaining < 60) return "1 month remaining"; - return `${Math.floor(daysRemaining / 30)} months remaining`; -}; - const meterTone = (pct: number): "default" | "warning" | "over" => { if (pct > 100) return "over"; if (pct >= 80) return "warning"; @@ -104,8 +94,7 @@ export default function SidebarUsageCard({ accessToken, collapsed, onExpandRail ); } - const daysUntilExpiration = licenseInfo?.expiration_date ? getDaysUntilExpiration(licenseInfo.expiration_date) : null; - const subtitle = licenseInfo?.expiration_date ? formatExpiration(daysUntilExpiration) : "Active plan"; + const subtitle = licenseInfo?.expiration_date ? formatExpirationStatus(licenseInfo.expiration_date) : "Active plan"; const meters = buildMeters(data); return ( diff --git a/ui/litellm-dashboard/src/components/UsageIndicator.test.tsx b/ui/litellm-dashboard/src/components/UsageIndicator.test.tsx index ad27fbcd74f..c587ebde57f 100644 --- a/ui/litellm-dashboard/src/components/UsageIndicator.test.tsx +++ b/ui/litellm-dashboard/src/components/UsageIndicator.test.tsx @@ -14,9 +14,19 @@ vi.mock("@/app/(dashboard)/hooks/useDisableUsageIndicator", () => ({ useDisableUsageIndicator: vi.fn(() => false), })); -import { getRemainingUsers } from "./networking"; +import { getLicenseInfo, getRemainingUsers } from "./networking"; +import type { LicenseInfo } from "./networking"; const mockGetRemainingUsers = vi.mocked(getRemainingUsers); +const mockGetLicenseInfo = vi.mocked(getLicenseInfo); + +const licenseWithExpiry = (expiration_date: string): LicenseInfo => ({ + has_license: true, + license_type: "enterprise", + expiration_date, + allowed_features: [], + limits: { max_users: null, max_teams: null }, +}); const renderWithClient = (ui: React.ReactElement) => { const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); @@ -36,6 +46,7 @@ describe("UsageIndicator", () => { beforeEach(() => { vi.clearAllMocks(); mockGetRemainingUsers.mockResolvedValue(DEFAULT_USAGE_DATA); + mockGetLicenseInfo.mockResolvedValue(null); }); it("should render when given access token and usage data loads", async () => { @@ -126,6 +137,23 @@ describe("UsageIndicator", () => { expect(screen.getByText("Over limit")).toBeInTheDocument(); }); + it("should show the exact license expiration date instead of time remaining", async () => { + mockGetLicenseInfo.mockResolvedValue(licenseWithExpiry("2099-12-31")); + + renderWithClient(); + + expect(await screen.findByText("Expires Dec 31, 2099")).toBeInTheDocument(); + expect(screen.queryByText(/(day|days|month|months) remaining/)).not.toBeInTheDocument(); + }); + + it("should show the exact date when the license is expired", async () => { + mockGetLicenseInfo.mockResolvedValue(licenseWithExpiry("2020-01-01")); + + renderWithClient(); + + expect(await screen.findByText("Expired Jan 1, 2020")).toBeInTheDocument(); + }); + it("should render nothing when accessToken is null", () => { renderWithClient(); diff --git a/ui/litellm-dashboard/src/components/UsageIndicator.tsx b/ui/litellm-dashboard/src/components/UsageIndicator.tsx index 6e7b5e9ec60..6ff4008c2cf 100644 --- a/ui/litellm-dashboard/src/components/UsageIndicator.tsx +++ b/ui/litellm-dashboard/src/components/UsageIndicator.tsx @@ -15,7 +15,7 @@ import { useEffect, useState } from "react"; import { getRemainingUsers } from "./networking"; import { cn } from "@/lib/cva.config"; -import { getDaysUntilExpiration } from "@/utils/licenseUtils"; +import { formatExpirationStatus, getDaysUntilExpiration } from "@/utils/licenseUtils"; import { useLicenseInfo } from "@/app/(dashboard)/hooks/license/useLicenseInfo"; interface UsageIndicatorProps { @@ -32,18 +32,6 @@ interface UsageData { total_teams_remaining: number | null; } -// Format expiration for display -const formatExpirationDisplay = (daysRemaining: number | null): string => { - if (daysRemaining === null) return "No expiration"; - if (daysRemaining < 0) return "Expired"; - if (daysRemaining === 0) return "Expires today"; - if (daysRemaining === 1) return "1 day remaining"; - if (daysRemaining < 30) return `${daysRemaining} days remaining`; - if (daysRemaining < 60) return "1 month remaining"; - const months = Math.floor(daysRemaining / 30); - return `${months} months remaining`; -}; - export default function UsageIndicator({ accessToken, width = 220 }: UsageIndicatorProps) { const disableUsageIndicator = useDisableUsageIndicator(); const [isExpanded, setIsExpanded] = useState(false); @@ -292,7 +280,7 @@ export default function UsageIndicator({ accessToken, width = 220 }: UsageIndica ) : isLicenseExpiringSoon ? ( ) : null} - {formatExpirationDisplay(daysUntilExpiration)} + {formatExpirationStatus(licenseInfo.expiration_date)}
)} @@ -529,7 +517,7 @@ export default function UsageIndicator({ accessToken, width = 220 }: UsageIndica isLicenseExpiringSoon && "text-yellow-600", )} > - {formatExpirationDisplay(daysUntilExpiration)} + {formatExpirationStatus(licenseInfo.expiration_date)}
{licenseInfo.license_type && ( diff --git a/ui/litellm-dashboard/src/utils/licenseUtils.test.ts b/ui/litellm-dashboard/src/utils/licenseUtils.test.ts index 717b8f0d90d..8a489c30c2d 100644 --- a/ui/litellm-dashboard/src/utils/licenseUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/licenseUtils.test.ts @@ -1,5 +1,11 @@ import { describe, it, expect } from "vitest"; -import { type LicenseExpiryTier, formatExpiryDate, getDaysUntilExpiration, getLicenseExpiryTier } from "./licenseUtils"; +import { + type LicenseExpiryTier, + formatExpirationStatus, + formatExpiryDate, + getDaysUntilExpiration, + getLicenseExpiryTier, +} from "./licenseUtils"; const NOW = new Date("2026-07-08T00:00:00Z"); @@ -59,3 +65,25 @@ describe("formatExpiryDate", () => { expect(formatExpiryDate("bogus")).toBe("bogus"); }); }); + +describe("formatExpirationStatus", () => { + it("shows the exact date for a future expiration", () => { + expect(formatExpirationStatus("2026-08-07", NOW)).toBe("Expires Aug 7, 2026"); + }); + + it("still reads as upcoming on the expiration day itself", () => { + expect(formatExpirationStatus("2026-07-08", NOW)).toBe("Expires Jul 8, 2026"); + }); + + it("shows the exact date for a past expiration", () => { + expect(formatExpirationStatus("2026-07-07", NOW)).toBe("Expired Jul 7, 2026"); + }); + + it("returns No expiration for a null date", () => { + expect(formatExpirationStatus(null, NOW)).toBe("No expiration"); + }); + + it("returns No expiration for an unparseable date", () => { + expect(formatExpirationStatus("not-a-date", NOW)).toBe("No expiration"); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/licenseUtils.ts b/ui/litellm-dashboard/src/utils/licenseUtils.ts index 57acad85508..4ba39dabe0b 100644 --- a/ui/litellm-dashboard/src/utils/licenseUtils.ts +++ b/ui/litellm-dashboard/src/utils/licenseUtils.ts @@ -49,3 +49,11 @@ export const formatExpiryDate = (expirationDate: string): string => { } return date.toLocaleDateString("en-US", EXPIRY_DATE_FORMAT); }; + +export const formatExpirationStatus = (expirationDate: string | null, now: Date = new Date()): string => { + const days = getDaysUntilExpiration(expirationDate, now); + if (expirationDate === null || days === null) { + return "No expiration"; + } + return days < 0 ? `Expired ${formatExpiryDate(expirationDate)}` : `Expires ${formatExpiryDate(expirationDate)}`; +}; From 8e73ff057fb7f8aeac1e2054996ac6bec48de91a Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 13:37:10 -0700 Subject: [PATCH 24/62] =?UTF-8?q?feat(mcp):=20issuer-anchored=20OAuth=20di?= =?UTF-8?q?scovery=20(RFC=208414=20=C2=A73.3)=20as=20the=20trust=20anchor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an admin-configured issuer to MCP servers. When set, OAuth metadata is fetched from the issuer's own origin and adopted only when the document self-attests that same issuer (RFC 8414 §3.3), making token_endpoint, registration_endpoint, and scopes authoritative for the pinned issuer instead of a document the MCP resource server chose. This closes the mix-up where a compromised resource echoes a pinned authorization_url to smuggle its own token endpoint and inflated scopes past the corroboration gate. Discovery is same-authority against the issuer origin, fails closed on a §3.3 mismatch, and does not fall back to resource-rooted discovery. Rows without an issuer keep the existing corroboration-gate behavior unchanged. Backend + schema only; UI field and live-proxy proof follow. --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/models/mcp_server.py | 1 + litellm/proxy/_experimental/mcp_server/db.py | 2 + .../mcp_server/mcp_server_manager.py | 113 +++++++++--- .../mcp_server/rest_endpoints.py | 1 + litellm/proxy/_types.py | 2 + litellm/proxy/schema.prisma | 1 + .../types/mcp_server/mcp_server_manager.py | 1 + schema.prisma | 1 + .../mcp_server/test_mcp_server_manager.py | 162 ++++++++++++++++++ 11 files changed, 260 insertions(+), 27 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260715000000_add_issuer_to_mcp_server_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260715000000_add_issuer_to_mcp_server_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260715000000_add_issuer_to_mcp_server_table/migration.sql new file mode 100644 index 00000000000..f7f23e6a55e --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260715000000_add_issuer_to_mcp_server_table/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "issuer" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a23cecc3911..f842bf13da9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -325,6 +325,7 @@ model LiteLLM_MCPServerTable { command String? args String[] @default([]) env Json? @default("{}") + issuer String? authorization_url String? token_url String? registration_url String? diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py index af2efa822b0..23b26bd8e89 100644 --- a/litellm/models/mcp_server.py +++ b/litellm/models/mcp_server.py @@ -79,6 +79,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): command: Optional[str] = None args: List[str] = Field(default_factory=list) env: Dict[str, str] = Field(default_factory=dict) + issuer: Optional[str] = None authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 97cefb3f2cb..ff018a1a021 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -48,6 +48,7 @@ if TYPE_CHECKING: _AUTH_FLOW_SCOPED_FIELDS: frozenset = frozenset( { + "issuer", "authorization_url", "token_url", "registration_url", @@ -1181,6 +1182,7 @@ def mcp_oauth_token_identity(server: object) -> tuple[object, ...]: getattr(server, "spec_path", None), getattr(server, "auth_type", None), getattr(server, "oauth2_flow", None), + getattr(server, "issuer", None), getattr(server, "authorization_url", None), getattr(server, "token_url", None), getattr(server, "registration_url", None), diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8b1d00c2855..9c5d3bab5e9 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -217,6 +217,17 @@ def _normalized_authorize_endpoint(url: str) -> str: return f"{scheme}://{authority}{parsed.path.rstrip('/')}" +def _issuer_matches(claimed_issuer: object, configured_issuer: str) -> bool: + """RFC 8414 §3.3 issuer equality between the metadata document's self-attested ``issuer`` and the + admin-configured issuer, tolerant only of URL-insignificant differences (scheme/host case, the + default port, a trailing slash). A non-string or empty claimed issuer never matches, so a + document that omits ``issuer`` fails closed under issuer-anchored discovery. + """ + if not isinstance(claimed_issuer, str) or not claimed_issuer: + return False + return _normalized_authorize_endpoint(claimed_issuer) == _normalized_authorize_endpoint(configured_issuer) + + def _endpoints_corroborate_authorization_url( source_authorization_url: str | None, trusted_authorization_url: str | None, @@ -1137,34 +1148,41 @@ class MCPServerManager: ) auth_type = server_config.get("auth_type", None) + manual_issuer = _blank_to_none(server_config.get("issuer")) manual_authorization_url = _blank_to_none(server_config.get("authorization_url")) manual_token_url = _blank_to_none(server_config.get("token_url")) manual_registration_url = _blank_to_none(server_config.get("registration_url")) - if server_url and ( - auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES + is_discovery_auth_type = auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES + use_issuer_anchor = manual_issuer is not None and is_discovery_auth_type + should_discover = bool(server_url) and ( + is_discovery_auth_type or self._obo_needs_endpoint_discovery( auth_type, server_config.get("token_exchange_endpoint"), manual_token_url, ) - ): + ) + if not should_discover: + mcp_oauth_metadata = None + elif manual_issuer is not None and is_discovery_auth_type: + mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer) + else: mcp_oauth_metadata = await self._descovery_metadata( server_url=server_url, - allow_origin_fallback=auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, + allow_origin_fallback=is_discovery_auth_type, ) - else: - mcp_oauth_metadata = None - gated_oauth_metadata = ( - _restrict_discovery_to_corroborated_authorization_server( + if use_issuer_anchor: + gated_oauth_metadata = mcp_oauth_metadata + elif is_discovery_auth_type: + gated_oauth_metadata = _restrict_discovery_to_corroborated_authorization_server( mcp_oauth_metadata, manual_authorization_url, server_name or server_id, bool(server_config.get("dcr_bridge")), ) - if auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES - else mcp_oauth_metadata - ) + else: + gated_oauth_metadata = mcp_oauth_metadata # Filter blank scopes (e.g. YAML ``scopes: [""]``) the same way the DB-build path does, so # an all-blank list normalizes to None rather than a ``("",)`` tuple that skips the @@ -1227,6 +1245,7 @@ class MCPServerManager: client_secret=server_config.get("client_secret", None), oauth2_flow=self._explicit_oauth2_flow(config_oauth2_flow), scopes=resolved_scopes, + issuer=manual_issuer, authorization_url=resolved_authorization_url, token_url=resolved_token_url, registration_url=resolved_registration_url, @@ -1570,12 +1589,15 @@ class MCPServerManager: auth_type = cast(MCPAuthType, mcp_server.auth_type) server_url = mcp_server.url + manual_issuer = _blank_to_none(mcp_server.issuer) manual_authorization_url = _blank_to_none(mcp_server.authorization_url) manual_token_url = _blank_to_none(mcp_server.token_url) manual_registration_url = _blank_to_none(mcp_server.registration_url) + is_discovery_auth_type = auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES + use_issuer_anchor = manual_issuer is not None and is_discovery_auth_type has_all_upstream_oauth_fields = bool(manual_authorization_url and manual_token_url and scopes) needs_discovery = bool(server_url) and ( - (auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and not has_all_upstream_oauth_fields) + (is_discovery_auth_type and not has_all_upstream_oauth_fields) or self._obo_needs_endpoint_discovery( auth_type, mcp_server.token_exchange_endpoint @@ -1583,31 +1605,33 @@ class MCPServerManager: manual_token_url, ) ) - mcp_oauth_metadata = ( - await self._descovery_metadata( + if not needs_discovery: + mcp_oauth_metadata = None + elif manual_issuer is not None and is_discovery_auth_type: + mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer) + else: + mcp_oauth_metadata = await self._descovery_metadata( server_url=server_url, # type: ignore[arg-type] - allow_origin_fallback=auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, + allow_origin_fallback=is_discovery_auth_type, ) - if needs_discovery - else None - ) - if needs_discovery and mcp_oauth_metadata is None: + if needs_discovery and not use_issuer_anchor and mcp_oauth_metadata is None: verbose_logger.warning( "MCP OAuth discovery yielded no metadata for server %s (%s); " "OAuth endpoints/scopes stay unresolved until a rebuild succeeds", mcp_server.server_id, server_url, ) - gated_oauth_metadata = ( - _restrict_discovery_to_corroborated_authorization_server( + if use_issuer_anchor: + gated_oauth_metadata = mcp_oauth_metadata + elif is_discovery_auth_type: + gated_oauth_metadata = _restrict_discovery_to_corroborated_authorization_server( mcp_oauth_metadata, manual_authorization_url, mcp_server.server_id, bool(getattr(mcp_server, "dcr_bridge", None)), ) - if auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES - else mcp_oauth_metadata - ) + else: + gated_oauth_metadata = mcp_oauth_metadata resolved_scopes = scopes or (gated_oauth_metadata.scopes if gated_oauth_metadata else None) @@ -1629,6 +1653,7 @@ class MCPServerManager: client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)), scopes=resolved_scopes, + issuer=manual_issuer, authorization_url=manual_authorization_url or getattr(gated_oauth_metadata, "authorization_url", None), token_url=manual_token_url or getattr(gated_oauth_metadata, "token_url", None), registration_url=manual_registration_url or getattr(gated_oauth_metadata, "registration_url", None), @@ -3337,8 +3362,28 @@ class MCPServerManager: return metadata return None + async def _fetch_issuer_anchored_oauth_metadata(self, issuer: str) -> Optional[MCPOAuthMetadata]: + """RFC 8414 issuer-anchored discovery. + + Fetch authorization-server metadata from the admin-configured issuer's own origin and adopt + it only when the document self-attests that same issuer (RFC 8414 §3.3). Because the trust + anchor is the pinned issuer rather than anything the MCP resource server advertises, the + resulting token_endpoint/registration_endpoint/scopes are authoritative for that issuer and + cannot be substituted by a compromised resource. Fails closed (returns None) on a §3.3 + mismatch or a fetch failure. The issuer is passed as its own ``server_url`` so the fetch is + treated as same-authority and is not subject to the resource-scoped SSRF shortcut. + """ + metadata = await self._fetch_single_authorization_server_metadata(issuer, issuer, require_issuer=issuer) + if metadata is None: + verbose_logger.warning( + "MCP OAuth issuer-anchored discovery for issuer %s yielded no metadata whose issuer " + "matched (RFC 8414 §3.3); OAuth endpoints/scopes stay unresolved until a rebuild succeeds", + issuer, + ) + return metadata + async def _fetch_single_authorization_server_metadata( - self, issuer_url: str, server_url: str + self, issuer_url: str, server_url: str, require_issuer: Optional[str] = None ) -> Optional[MCPOAuthMetadata]: try: parsed = urlparse(issuer_url) @@ -3382,15 +3427,27 @@ class MCPServerManager: ) continue - scopes = self._extract_scopes(data.get("scopes_supported")) + claimed_issuer = data.get("issuer") verbose_logger.debug( "Authorization server metadata from %s: issuer=%s grant_types_supported=%s " "token_endpoint_auth_methods_supported=%s", url, - data.get("issuer"), + claimed_issuer, data.get("grant_types_supported"), data.get("token_endpoint_auth_methods_supported"), ) + if require_issuer is not None and not _issuer_matches(claimed_issuer, require_issuer): + verbose_logger.warning( + "MCP OAuth issuer-anchored discovery: metadata at %s self-attests issuer %r, which " + "does not match the configured issuer %r (RFC 8414 §3.3); rejecting so a compromised " + "resource cannot substitute an attacker authorization server", + url, + claimed_issuer, + require_issuer, + ) + continue + + scopes = self._extract_scopes(data.get("scopes_supported")) metadata = MCPOAuthMetadata( scopes=scopes, authorization_url=data.get("authorization_endpoint"), @@ -3408,6 +3465,8 @@ class MCPServerManager: ): return metadata + if require_issuer is not None: + return None return self._build_azure_authorization_server_metadata(parsed) @staticmethod diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 111fde86ea0..7ca4923337c 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1138,6 +1138,7 @@ if MCP_AVAILABLE: static_headers=request.static_headers, client_id=client_id, client_secret=client_secret, + issuer=request.issuer, token_url=request.token_url, scopes=scopes, authorization_url=request.authorization_url, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5e3ea4b7dcb..053f78a3698 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1263,6 +1263,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): command: Optional[str] = None args: List[str] = Field(default_factory=list) env: Dict[str, str] = Field(default_factory=dict) + issuer: Optional[str] = None authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None @@ -1368,6 +1369,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): command: Optional[str] = None args: List[str] = Field(default_factory=list) env: Dict[str, str] = Field(default_factory=dict) + issuer: Optional[str] = None authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a23cecc3911..f842bf13da9 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -325,6 +325,7 @@ model LiteLLM_MCPServerTable { command String? args String[] @default([]) env Json? @default("{}") + issuer String? authorization_url String? token_url String? registration_url String? diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 801436c774a..c892824fb54 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -60,6 +60,7 @@ class MCPServer(BaseModel): # OAuth-specific fields client_id: Optional[str] = None client_secret: Optional[str] = None + issuer: Optional[str] = None scopes: Optional[List[str]] = None authorization_url: Optional[str] = None token_url: Optional[str] = None diff --git a/schema.prisma b/schema.prisma index a23cecc3911..f842bf13da9 100644 --- a/schema.prisma +++ b/schema.prisma @@ -325,6 +325,7 @@ model LiteLLM_MCPServerTable { command String? args String[] @default([]) env Json? @default("{}") + issuer String? authorization_url String? token_url String? registration_url String? 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 bba00ed1819..58773642692 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 @@ -1230,6 +1230,80 @@ class TestMCPServerManager: assert built.registration_url == "https://idp.example.com/register" assert built.scopes == ["read"] + @pytest.mark.asyncio + async def test_build_from_table_issuer_anchor_ignores_resource_rooted_discovery(self): + """When an admin configures an issuer, discovery is anchored on that issuer's own metadata + (RFC 8414), not on the resource-rooted RFC 9728 chain a compromised MCP resource controls. + The resource-rooted _descovery_metadata must not run at all, and the authoritative token_url, + registration_url, and scopes come from the issuer document. This closes the mix-up where a + compromised resource echoes the pinned authorize endpoint to smuggle its own token_url.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="issuer-anchored-1", + alias="issuer_anchored", + description="issuer configured, blank endpoints", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + issuer="https://idp.example.com", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + authoritative = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + scopes=["read", "write"], + authorization_server_scopes=["read", "write"], + ) + resource_rooted = AsyncMock(return_value=MCPOAuthMetadata(token_url="https://attacker.example.com/steal")) + with ( + patch.object(manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=authoritative)) as anchored, + patch.object(manager, "_descovery_metadata", new=resource_rooted), + ): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + anchored.assert_awaited_once_with("https://idp.example.com") + resource_rooted.assert_not_awaited() + assert built.issuer == "https://idp.example.com" + assert built.authorization_url == "https://idp.example.com/authorize" + assert built.token_url == "https://idp.example.com/token" + assert built.registration_url == "https://idp.example.com/register" + assert built.scopes == ["read", "write"] + + @pytest.mark.asyncio + async def test_build_from_table_issuer_anchor_fails_closed_without_falling_back_to_resource(self): + """A configured issuer whose metadata does not validate (RFC 8414 §3.3 mismatch or fetch + failure) yields None from the anchored fetch. The build must adopt nothing and must NOT fall + back to resource-rooted discovery, or the fail-closed guarantee would be defeated by the very + resource the issuer anchor exists to distrust.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="issuer-anchored-2", + alias="issuer_anchored_failclosed", + description="issuer configured, upstream fails validation", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + issuer="https://idp.example.com", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + resource_rooted = AsyncMock(return_value=MCPOAuthMetadata(token_url="https://attacker.example.com/steal")) + with ( + patch.object(manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "_descovery_metadata", new=resource_rooted), + ): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + resource_rooted.assert_not_awaited() + assert built.issuer == "https://idp.example.com" + assert built.token_url is None + assert built.registration_url is None + assert built.scopes is None + @pytest.mark.asyncio @pytest.mark.parametrize( "advertised_authorization_url", @@ -2462,6 +2536,78 @@ class TestMCPServerManager: assert result.token_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" assert result.scopes == ["api://some-scope/.default"] + @staticmethod + def _issuer_doc_response_builder(well_known_url: str, document: dict): + def build_response(url: str, **kwargs): + mock_response = MagicMock() + if url == well_known_url: + mock_response.json.return_value = document + mock_response.raise_for_status = MagicMock() + else: + request = httpx.Request("GET", url) + response_obj = httpx.Response(status_code=404, request=request) + mock_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError("not found", request=request, response=response_obj) + ) + return mock_response + + return build_response + + @pytest.mark.asyncio + async def test_fetch_single_authorization_server_metadata_adopts_document_with_matching_issuer(self): + """RFC 8414 §3.3: under require_issuer, a document that self-attests the same issuer it was + fetched from is authoritative and its endpoints and scopes are adopted.""" + manager = MCPServerManager() + issuer = "https://idp.example.com" + build_response = self._issuer_doc_response_builder( + f"{issuer}/.well-known/oauth-authorization-server", + { + "issuer": issuer, + "authorization_endpoint": "https://idp.example.com/authorize", + "token_endpoint": "https://idp.example.com/token", + "scopes_supported": ["read", "write"], + }, + ) + mock_client = MagicMock() + mock_client.get = AsyncMock(side_effect=build_response) + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + result = await manager._fetch_single_authorization_server_metadata(issuer, issuer, require_issuer=issuer) + + assert result is not None + assert result.authorization_url == "https://idp.example.com/authorize" + assert result.token_url == "https://idp.example.com/token" + assert result.scopes == ["read", "write"] + + @pytest.mark.asyncio + async def test_fetch_single_authorization_server_metadata_rejects_issuer_mismatch(self): + """RFC 8414 §3.3 fail-closed: a document self-attesting a DIFFERENT issuer than the one it was + fetched from is rejected even though it carries valid-looking endpoints, so a compromised + resource cannot point the issuer-anchored fetch at an attacker authorization server that + smuggles its own token_endpoint and inflated scopes.""" + manager = MCPServerManager() + issuer = "https://idp.example.com" + build_response = self._issuer_doc_response_builder( + f"{issuer}/.well-known/oauth-authorization-server", + { + "issuer": "https://attacker.example.com", + "authorization_endpoint": "https://idp.example.com/authorize", + "token_endpoint": "https://attacker.example.com/steal", + "scopes_supported": ["admin"], + }, + ) + mock_client = MagicMock() + mock_client.get = AsyncMock(side_effect=build_response) + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + result = await manager._fetch_single_authorization_server_metadata(issuer, issuer, require_issuer=issuer) + + assert result is None + @pytest.mark.asyncio async def test_fetch_single_authorization_server_metadata_derives_azure_metadata( self, @@ -5487,6 +5633,22 @@ class TestMCPServerTimestamps: assert _normalized_authorize_endpoint("https://idp.example.com/authorize?prompt=consent") == canonical assert _normalized_authorize_endpoint("https://idp.example.com:8443/authorize") != canonical + def test_issuer_matches_rfc8414_section_3_3(self): + """Issuer equality tolerates only URL-insignificant differences (scheme/host case, default + port, a trailing slash). A different host, a non-string, an empty string, or a None issuer + never matches, so a document that omits issuer fails closed under issuer-anchored discovery.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _issuer_matches + + assert _issuer_matches("https://mcp.slack.com", "https://mcp.slack.com") + assert _issuer_matches("https://MCP.slack.com/", "https://mcp.slack.com") + assert _issuer_matches("https://mcp.slack.com:443", "https://mcp.slack.com") + assert _issuer_matches("https://login.example.com/tenant/v2.0", "https://login.example.com/tenant/v2.0") + assert not _issuer_matches("https://attacker.example.com", "https://mcp.slack.com") + assert not _issuer_matches("https://login.example.com/other/v2.0", "https://login.example.com/tenant/v2.0") + assert not _issuer_matches(None, "https://mcp.slack.com") + assert not _issuer_matches("", "https://mcp.slack.com") + assert not _issuer_matches(123, "https://mcp.slack.com") + def test_build_mcp_server_table_preserves_timestamps(self): """_build_mcp_server_table must use the MCPServer's stored timestamps, not datetime.now().""" manager = MCPServerManager() From 031922eec648fd5d5dc863d13833c89c1d769bc6 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 14:06:51 -0700 Subject: [PATCH 25/62] feat(mcp): discover-by-default issuer (trust-on-first-use) + stale-clear on url/auth_type change + UI Discover the issuer from the upstream and persist it trust-on-first-use (fill-empty-only, frozen thereafter), so admins do not have to type it; an admin-configured issuer always wins and is never overwritten. Re-pointing the server url now clears the discovered issuer and endpoints (matching the existing auth_type-change clearing) so a new upstream re-discovers instead of anchoring on the previous upstream's issuer. Adds the issuer to the per-user OAuth token identity so re-pointing it purges stale tokens. Surfaces the issuer as an optional, auto-discovered, overridable field in the create and edit MCP server forms. C901 gate shows +1 vs staging; that is inherited from the #33317 stack base (delta 0 against --- litellm/proxy/_experimental/mcp_server/db.py | 10 ++-- .../mcp_server/mcp_server_manager.py | 13 +++++- .../types/mcp_server/mcp_server_manager.py | 5 ++ .../mcp_server/test_mcp_partial_update.py | 41 +++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 46 ++++++++++++++++++- .../_components/OAuthFormFields.tsx | 11 +++++ .../_components/create_mcp_server.tsx | 3 +- .../_components/mcp_server_edit.tsx | 24 +++++++++- .../src/components/mcp_tools/types.tsx | 2 + 9 files changed, 147 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index ff018a1a021..22c531b5c79 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -698,13 +698,14 @@ async def update_mcp_server( # of being reset to a schema default (transport=sse, allow_all_keys=False...). data_dict = _prepare_mcp_server_data(data, exclude_unset=True, fields_set=fields_set) - # Pre-fetch existing record once if we need it for auth_type or credential logic + # Pre-fetch existing record once if we need it for auth_type, url, or credential logic existing = None has_credentials = "credentials" in data_dict and data_dict["credentials"] is not None # An explicit token-exchange column write (set or clear) also migrates the # legacy blob copies below, so the existing row is needed for those updates. explicit_te_write = bool(_TOKEN_EXCHANGE_COLUMN_FIELDS & data_dict.keys()) - if data.auth_type or has_credentials or explicit_te_write: + url_provided = "url" in data_dict and data_dict["url"] is not None + if data.auth_type or has_credentials or explicit_te_write or url_provided: existing = await MCPServerRepository(prisma_client).table.find_unique(where={"server_id": data.server_id}) auth_type_changed = bool( @@ -712,12 +713,15 @@ async def update_mcp_server( and existing and _credential_auth_class(existing.auth_type) != _credential_auth_class(data.auth_type) ) + # A url change re-points the server at a potentially different upstream, so any discovered or + # trust-on-first-use OAuth endpoints/issuer belong to the old upstream and must re-discover. + url_changed = bool(url_provided and existing and existing.url != data_dict["url"]) # Clear stale credentials when auth_type changes but no new credentials provided if auth_type_changed and "credentials" not in data_dict: data_dict["credentials"] = None - if auth_type_changed: + if auth_type_changed or url_changed: data_dict.update({field: None for field in _AUTH_FLOW_SCOPED_FIELDS if field not in data_dict}) # An explicit column write that does not touch credentials must still migrate diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 9c5d3bab5e9..42c9921ecf7 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1713,6 +1713,7 @@ class MCPServerManager: await self._persist_discovered_oauth_endpoints( server_id=mcp_server.server_id, auth_type=auth_type, + existing_issuer=manual_issuer, existing_authorization_url=manual_authorization_url, existing_token_url=manual_token_url, existing_scopes=scopes, @@ -1760,6 +1761,7 @@ class MCPServerManager: *, server_id: str, auth_type: MCPAuthType | None, + existing_issuer: str | None, existing_authorization_url: str | None, existing_token_url: str | None, existing_scopes: list[str] | None, @@ -1782,6 +1784,9 @@ class MCPServerManager: return if metadata is None or metadata.from_origin_fallback: return + issuer_update = ( + {"issuer": metadata.discovered_issuer} if metadata.discovered_issuer and not existing_issuer else {} + ) authorization_url_update = ( {"authorization_url": metadata.authorization_url} if metadata.authorization_url and not existing_authorization_url @@ -1789,7 +1794,12 @@ class MCPServerManager: ) token_url_update = {"token_url": metadata.token_url} if metadata.token_url and not existing_token_url else {} scopes_update = {"credentials": {"scopes": metadata.scopes}} if metadata.scopes and not existing_scopes else {} - updates: dict[str, object] = {**authorization_url_update, **token_url_update, **scopes_update} + updates: dict[str, object] = { + **issuer_update, + **authorization_url_update, + **token_url_update, + **scopes_update, + } if not updates: return from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # db.py imports this module at load @@ -3453,6 +3463,7 @@ class MCPServerManager: authorization_url=data.get("authorization_endpoint"), token_url=data.get("token_endpoint"), registration_url=data.get("registration_endpoint"), + discovered_issuer=claimed_issuer if isinstance(claimed_issuer, str) and claimed_issuer else None, ) if any( diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index c892824fb54..cf3eed20f5b 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -26,6 +26,11 @@ class MCPOAuthMetadata(BaseModel): authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + discovered_issuer: Optional[str] = None + """The ``issuer`` the authorization-server metadata document self-attests (RFC 8414). Persisted + trust-on-first-use as the server's ``issuer`` when none is configured, so that later rebuilds + anchor discovery on it (RFC 8414 §3.3) and a subsequently compromised resource cannot re-point + it. Never overwrites an admin-configured issuer.""" from_origin_fallback: bool = False """True when the metadata came from guessing the resource origin as its authorization server rather than from an RFC 9728/8414-advertised document. Guessed endpoints are diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index 41d61c8f508..637ee409fbe 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -203,6 +203,7 @@ async def test_auth_type_switch_clears_stale_flow_scoped_fields(): data_dict = await _run_update_with_existing(data, existing_auth_type="oauth2") for stale_field in ( + "issuer", "authorization_url", "token_url", "registration_url", @@ -217,6 +218,46 @@ async def test_auth_type_switch_clears_stale_flow_scoped_fields(): assert _credentials_cleared(data_dict["credentials"]) +@pytest.mark.asyncio +async def test_url_change_clears_stale_discovered_oauth_fields(): + """Re-pointing the server url at a potentially different upstream must clear the discovered or + trust-on-first-use OAuth issuer and endpoints, so the new upstream re-discovers instead of + anchoring on the previous upstream's issuer (RFC 8414 §3.3 against a stale anchor).""" + mock_prisma = _mock_prisma() + existing = MagicMock() + existing.auth_type = "oauth2" + existing.url = "https://old.example.com/mcp" + existing.credentials = None + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest(server_id="my-test-server", url="https://new.example.com/mcp") + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + assert data_dict["url"] == "https://new.example.com/mcp" + for stale_field in ("issuer", "authorization_url", "token_url", "registration_url"): + assert data_dict[stale_field] is None, f"{stale_field} must be cleared on url change" + + +@pytest.mark.asyncio +async def test_unchanged_url_does_not_clear_discovered_oauth_fields(): + """A partial update that resends the same url (or omits it) must not clear the discovered OAuth + fields, so a routine save does not force needless re-discovery.""" + mock_prisma = _mock_prisma() + existing = MagicMock() + existing.auth_type = "oauth2" + existing.url = "https://same.example.com/mcp" + existing.credentials = None + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest(server_id="my-test-server", url="https://same.example.com/mcp") + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + for preserved_field in ("issuer", "authorization_url", "token_url", "registration_url"): + assert preserved_field not in data_dict, f"{preserved_field} must not be cleared when url is unchanged" + + @pytest.mark.asyncio async def test_auth_type_switch_keeps_explicitly_provided_flow_fields(): """Fields explicitly provided alongside the auth_type switch must survive it.""" 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 58773642692..3dcd05e5936 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 @@ -1255,7 +1255,6 @@ class TestMCPServerManager: token_url="https://idp.example.com/token", registration_url="https://idp.example.com/register", scopes=["read", "write"], - authorization_server_scopes=["read", "write"], ) resource_rooted = AsyncMock(return_value=MCPOAuthMetadata(token_url="https://attacker.example.com/steal")) with ( @@ -5340,6 +5339,7 @@ class TestMCPServerTimestamps: await manager._persist_discovered_oauth_endpoints( server_id="s", auth_type=MCPAuth.api_key, + existing_issuer=None, existing_authorization_url=None, existing_token_url=None, existing_scopes=None, @@ -5348,6 +5348,7 @@ class TestMCPServerTimestamps: await manager._persist_discovered_oauth_endpoints( server_id="s", auth_type=MCPAuth.oauth2, + existing_issuer=None, existing_authorization_url=None, existing_token_url=None, existing_scopes=None, @@ -5356,6 +5357,7 @@ class TestMCPServerTimestamps: await manager._persist_discovered_oauth_endpoints( server_id="s", auth_type=MCPAuth.oauth2, + existing_issuer=None, existing_authorization_url=None, existing_token_url=None, existing_scopes=None, @@ -5364,6 +5366,7 @@ class TestMCPServerTimestamps: await manager._persist_discovered_oauth_endpoints( server_id="s", auth_type=MCPAuth.oauth2, + existing_issuer=None, existing_authorization_url="https://configured.example.com/authorize", existing_token_url="https://configured.example.com/token", existing_scopes=["configured"], @@ -5389,6 +5392,7 @@ class TestMCPServerTimestamps: await manager._persist_discovered_oauth_endpoints( server_id="s", auth_type=MCPAuth.oauth2, + existing_issuer=None, existing_authorization_url=None, existing_token_url="https://configured.example.com/token", existing_scopes=None, @@ -5405,6 +5409,46 @@ class TestMCPServerTimestamps: assert persisted.credentials == {"scopes": ["s1"]} assert "token_url" not in persisted.fields_set() + @pytest.mark.asyncio + async def test_persist_discovered_oauth_endpoints_writes_discovered_issuer_trust_on_first_use(self): + """A server with no configured issuer records the discovered issuer trust-on-first-use, so the + next rebuild anchors discovery on it (RFC 8414 §3.3) instead of re-trusting the resource. When + an issuer is already set (admin-typed or a prior discovery), it is never overwritten.""" + manager = MCPServerManager() + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + discovered_issuer="https://idp.example.com", + ) + + update_mcp_server_mock = AsyncMock() + with ( + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=update_mcp_server_mock), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + ): + await manager._persist_discovered_oauth_endpoints( + server_id="s", + auth_type=MCPAuth.oauth2, + existing_issuer=None, + existing_authorization_url=None, + existing_token_url=None, + existing_scopes=None, + metadata=metadata, + ) + await manager._persist_discovered_oauth_endpoints( + server_id="s", + auth_type=MCPAuth.oauth2, + existing_issuer="https://admin-configured.example.com", + existing_authorization_url="https://admin-configured.example.com/authorize", + existing_token_url="https://admin-configured.example.com/token", + existing_scopes=["cfg"], + metadata=metadata, + ) + + assert update_mcp_server_mock.await_count == 1 + persisted = update_mcp_server_mock.call_args.kwargs["data"] + assert persisted.issuer == "https://idp.example.com" + @pytest.mark.asyncio async def test_build_mcp_server_from_table_skips_persistence_for_temporary_servers(self): """The session endpoint builds temporary servers whose server_id has no DB row; with diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx index fc72981edd0..76e4342f52e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx @@ -167,6 +167,17 @@ const OAuthFormFields: React.FC = ({ > + @@ -1588,6 +1607,7 @@ const MCPServerEdit: React.FC = ({ oauthFlowTypeValue ?? oauth2FlowToFormValue(mcpServer.oauth2_flow) ?? OAUTH_FLOW.INTERACTIVE, static_headers: currentStaticHeaders ?? mcpServer.static_headers, credentials: currentCredentials, + issuer: currentIssuer ?? mcpServer.issuer, authorization_url: currentAuthorizationUrl ?? mcpServer.authorization_url, token_url: currentTokenUrl ?? mcpServer.token_url, registration_url: currentRegistrationUrl ?? mcpServer.registration_url, diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 04766aad7b4..dd7ed2bfbe4 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -81,6 +81,7 @@ export const getOAuthAuthorizationIdentity = (values: Record): client_id: credentials.client_id ?? null, client_secret: credentials.client_secret ?? null, scopes: credentials.scopes ?? null, + issuer: values.issuer ?? null, authorization_url: values.authorization_url ?? null, token_url: values.token_url ?? null, registration_url: values.registration_url ?? null, @@ -341,6 +342,7 @@ export interface MCPServer { transport?: string | null; auth_type?: string | null; oauth2_flow?: string | null; + issuer?: string | null; authorization_url?: string | null; token_url?: string | null; registration_url?: string | null; From b3af125078bcc2572986db0dbf77bd21380dc0d8 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 16:36:52 -0700 Subject: [PATCH 26/62] fix(mcp): keep scopes resource-driven under a pinned issuer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The issuer anchor is for the token/registration endpoints only (the RFC 9700 mix-up). Scope selection stays resource-driven per the MCP authorization spec Scope Selection Strategy: _fetch_issuer_anchored_oauth_metadata now validates the issuer document (RFC 8414 §3.3) for the endpoints and separately fetches the resource's advertised scopes (WWW-Authenticate challenge, else RFC 9728 scopes_supported) for the scope value, instead of using the issuer document's own scopes_supported. The resource can influence only the requested scope, which the authorization server and user consent bound (RFC 6749 §3.3), never the token endpoint. --- .../mcp_server/mcp_server_manager.py | 36 ++++++++----- .../mcp_server/test_mcp_server_manager.py | 53 +++++++++++++++---- 2 files changed, 68 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 42c9921ecf7..70180d5d469 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1165,7 +1165,7 @@ class MCPServerManager: if not should_discover: mcp_oauth_metadata = None elif manual_issuer is not None and is_discovery_auth_type: - mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer) + mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer, server_url) else: mcp_oauth_metadata = await self._descovery_metadata( server_url=server_url, @@ -1608,7 +1608,10 @@ class MCPServerManager: if not needs_discovery: mcp_oauth_metadata = None elif manual_issuer is not None and is_discovery_auth_type: - mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer) + mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata( + manual_issuer, + server_url, # type: ignore[arg-type] + ) else: mcp_oauth_metadata = await self._descovery_metadata( server_url=server_url, # type: ignore[arg-type] @@ -3372,25 +3375,34 @@ class MCPServerManager: return metadata return None - async def _fetch_issuer_anchored_oauth_metadata(self, issuer: str) -> Optional[MCPOAuthMetadata]: - """RFC 8414 issuer-anchored discovery. + async def _fetch_issuer_anchored_oauth_metadata(self, issuer: str, server_url: str) -> Optional[MCPOAuthMetadata]: + """RFC 8414 issuer-anchored discovery for the OAuth endpoints, with resource-driven scopes. Fetch authorization-server metadata from the admin-configured issuer's own origin and adopt - it only when the document self-attests that same issuer (RFC 8414 §3.3). Because the trust - anchor is the pinned issuer rather than anything the MCP resource server advertises, the - resulting token_endpoint/registration_endpoint/scopes are authoritative for that issuer and - cannot be substituted by a compromised resource. Fails closed (returns None) on a §3.3 - mismatch or a fetch failure. The issuer is passed as its own ``server_url`` so the fetch is - treated as same-authority and is not subject to the resource-scoped SSRF shortcut. + its ``token_endpoint``/``registration_endpoint`` only when the document self-attests that same + issuer (RFC 8414 §3.3). Because the trust anchor is the pinned issuer rather than anything the + MCP resource advertises, the endpoints are authoritative for that issuer and cannot be + substituted by a compromised resource. Fails closed (returns None) on a §3.3 mismatch or a + fetch failure. The issuer is passed as its own ``server_url`` so the endpoint fetch is treated + as same-authority and is not subject to the resource-scoped SSRF shortcut. + + Scopes are NOT taken from the issuer document. Per the MCP authorization spec Scope Selection + Strategy and RFC 9728, the scopes a client requests are resource-driven (the WWW-Authenticate + challenge or the protected-resource ``scopes_supported``), so the resource's advertised scopes + are fetched separately and used; the resource can influence only the requested scope, which + the authorization server and user consent bound (RFC 6749 §3.3), never the token endpoint. """ metadata = await self._fetch_single_authorization_server_metadata(issuer, issuer, require_issuer=issuer) if metadata is None: verbose_logger.warning( "MCP OAuth issuer-anchored discovery for issuer %s yielded no metadata whose issuer " - "matched (RFC 8414 §3.3); OAuth endpoints/scopes stay unresolved until a rebuild succeeds", + "matched (RFC 8414 §3.3); OAuth endpoints stay unresolved until a rebuild succeeds", issuer, ) - return metadata + return None + resource_metadata = await self._descovery_metadata(server_url, allow_origin_fallback=False) + resource_scopes = resource_metadata.scopes if resource_metadata else None + return metadata.model_copy(update={"scopes": resource_scopes}) async def _fetch_single_authorization_server_metadata( self, issuer_url: str, server_url: str, require_issuer: Optional[str] = None 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 3dcd05e5936..75f77c1ec6d 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 @@ -1231,12 +1231,12 @@ class TestMCPServerManager: assert built.scopes == ["read"] @pytest.mark.asyncio - async def test_build_from_table_issuer_anchor_ignores_resource_rooted_discovery(self): - """When an admin configures an issuer, discovery is anchored on that issuer's own metadata - (RFC 8414), not on the resource-rooted RFC 9728 chain a compromised MCP resource controls. - The resource-rooted _descovery_metadata must not run at all, and the authoritative token_url, - registration_url, and scopes come from the issuer document. This closes the mix-up where a - compromised resource echoes the pinned authorize endpoint to smuggle its own token_url.""" + async def test_build_from_table_uses_issuer_anchored_endpoints_when_issuer_configured(self): + """When an admin configures an issuer, the build takes its endpoints from the issuer-anchored + fetch (RFC 8414 §3.3) rather than the resource-rooted corroboration path. The build path does + not call _descovery_metadata directly; the issuer-anchored helper is responsible for combining + issuer endpoints with resource-driven scopes internally, and is invoked with the server url so + it can fetch those scopes.""" manager = MCPServerManager() row = LiteLLM_MCPServerTable( server_id="issuer-anchored-1", @@ -1250,7 +1250,7 @@ class TestMCPServerManager: updated_at=datetime.now(), ) - authoritative = MCPOAuthMetadata( + resolved = MCPOAuthMetadata( authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", registration_url="https://idp.example.com/register", @@ -1258,12 +1258,12 @@ class TestMCPServerManager: ) resource_rooted = AsyncMock(return_value=MCPOAuthMetadata(token_url="https://attacker.example.com/steal")) with ( - patch.object(manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=authoritative)) as anchored, + patch.object(manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=resolved)) as anchored, patch.object(manager, "_descovery_metadata", new=resource_rooted), ): built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) - anchored.assert_awaited_once_with("https://idp.example.com") + anchored.assert_awaited_once_with("https://idp.example.com", "https://up.example.com/mcp") resource_rooted.assert_not_awaited() assert built.issuer == "https://idp.example.com" assert built.authorization_url == "https://idp.example.com/authorize" @@ -1271,6 +1271,41 @@ class TestMCPServerManager: assert built.registration_url == "https://idp.example.com/register" assert built.scopes == ["read", "write"] + @pytest.mark.asyncio + async def test_fetch_issuer_anchored_metadata_takes_endpoints_from_issuer_scopes_from_resource(self): + """The issuer-anchored helper adopts token_endpoint/registration_endpoint from the pinned + issuer's own §3.3-validated document, but the scopes are resource-driven: it fetches the + resource's advertised scopes and uses those, not the issuer document's scopes_supported. This + keeps endpoint trust anchored on the issuer while scope selection stays resource-driven per the + MCP Scope Selection Strategy.""" + manager = MCPServerManager() + + issuer_document = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + scopes=["as.everything"], + ) + resource_document = MCPOAuthMetadata(scopes=["resource.read"]) + with ( + patch.object( + manager, "_fetch_single_authorization_server_metadata", new=AsyncMock(return_value=issuer_document) + ) as issuer_fetch, + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=resource_document)) as resource_fetch, + ): + result = await manager._fetch_issuer_anchored_oauth_metadata( + "https://idp.example.com", "https://up.example.com/mcp" + ) + + issuer_fetch.assert_awaited_once_with( + "https://idp.example.com", "https://idp.example.com", require_issuer="https://idp.example.com" + ) + resource_fetch.assert_awaited_once() + assert result is not None + assert result.token_url == "https://idp.example.com/token" + assert result.registration_url == "https://idp.example.com/register" + assert result.scopes == ["resource.read"] + @pytest.mark.asyncio async def test_build_from_table_issuer_anchor_fails_closed_without_falling_back_to_resource(self): """A configured issuer whose metadata does not validate (RFC 8414 §3.3 mismatch or fetch From daa15ba0747b66d4c5ae70ade65d293e00adb2b7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 15 Jul 2026 17:09:24 -0700 Subject: [PATCH 27/62] chore(ui): remove unmounted UsageIndicator component --- .../src/components/UsageIndicator.test.tsx | 193 ----- .../src/components/UsageIndicator.tsx | 681 ------------------ 2 files changed, 874 deletions(-) delete mode 100644 ui/litellm-dashboard/src/components/UsageIndicator.test.tsx delete mode 100644 ui/litellm-dashboard/src/components/UsageIndicator.tsx diff --git a/ui/litellm-dashboard/src/components/UsageIndicator.test.tsx b/ui/litellm-dashboard/src/components/UsageIndicator.test.tsx deleted file mode 100644 index ad27fbcd74f..00000000000 --- a/ui/litellm-dashboard/src/components/UsageIndicator.test.tsx +++ /dev/null @@ -1,193 +0,0 @@ -import React from "react"; -import { describe, it, expect, vi, beforeEach } from "vitest"; -import { render, screen, waitFor } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import UsageIndicator from "./UsageIndicator"; - -vi.mock("./networking", () => ({ - getRemainingUsers: vi.fn(), - getLicenseInfo: vi.fn().mockResolvedValue(null), -})); - -vi.mock("@/app/(dashboard)/hooks/useDisableUsageIndicator", () => ({ - useDisableUsageIndicator: vi.fn(() => false), -})); - -import { getRemainingUsers } from "./networking"; - -const mockGetRemainingUsers = vi.mocked(getRemainingUsers); - -const renderWithClient = (ui: React.ReactElement) => { - const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - return render({ui}); -}; - -const DEFAULT_USAGE_DATA = { - total_users: 100, - total_users_used: 1, - total_users_remaining: 99, - total_teams: null, - total_teams_used: 0, - total_teams_remaining: null, -}; - -describe("UsageIndicator", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockGetRemainingUsers.mockResolvedValue(DEFAULT_USAGE_DATA); - }); - - it("should render when given access token and usage data loads", async () => { - renderWithClient(); - - await screen.findByText("Usage"); - - expect(screen.getByText("Usage")).toBeInTheDocument(); - }); - - it("should not show Near limit when users usage is below 80% (1/100 -> 1%)", async () => { - renderWithClient(); - - await screen.findByText("Usage"); - - expect(screen.queryByText("Near limit")).not.toBeInTheDocument(); - }); - - it("should render nothing when both total_users and total_teams are null", async () => { - mockGetRemainingUsers.mockResolvedValue({ - total_users: null, - total_teams: null, - total_users_used: 520, - total_teams_used: 4, - total_teams_remaining: null, - total_users_remaining: null, - }); - - renderWithClient(); - - await waitFor(() => { - expect(screen.queryByText("Usage")).not.toBeInTheDocument(); - expect(screen.queryByText("Loading...")).not.toBeInTheDocument(); - }); - }); - - it("should show Near limit for Teams when at 80% usage (4/5)", async () => { - mockGetRemainingUsers.mockResolvedValue({ - total_users: null, - total_users_used: 0, - total_users_remaining: null, - total_teams: 5, - total_teams_used: 4, - total_teams_remaining: 1, - }); - - renderWithClient(); - - await screen.findByText("Usage"); - - expect(screen.getByText("Teams")).toBeInTheDocument(); - expect(screen.getByText("Near limit")).toBeInTheDocument(); - }); - - it("should show Over limit for Users when usage exceeds 100% (105/100)", async () => { - mockGetRemainingUsers.mockResolvedValue({ - total_users: 100, - total_users_used: 105, - total_users_remaining: -5, - total_teams: null, - total_teams_used: 0, - total_teams_remaining: null, - }); - - renderWithClient(); - - await screen.findByText("Usage"); - - expect(screen.getByText("Users")).toBeInTheDocument(); - expect(screen.getByText("Over limit")).toBeInTheDocument(); - }); - - it("should show Over limit for Teams when usage exceeds 100%", async () => { - mockGetRemainingUsers.mockResolvedValue({ - total_users: null, - total_users_used: 0, - total_users_remaining: null, - total_teams: 10, - total_teams_used: 12, - total_teams_remaining: -2, - }); - - renderWithClient(); - - await screen.findByText("Usage"); - - expect(screen.getByText("Teams")).toBeInTheDocument(); - expect(screen.getByText("Over limit")).toBeInTheDocument(); - }); - - it("should render nothing when accessToken is null", () => { - renderWithClient(); - - expect(mockGetRemainingUsers).not.toHaveBeenCalled(); - expect(screen.queryByText("Usage")).not.toBeInTheDocument(); - }); - - it("should render nothing when disableUsageIndicator is true", async () => { - const { useDisableUsageIndicator } = await import("@/app/(dashboard)/hooks/useDisableUsageIndicator"); - (useDisableUsageIndicator as ReturnType).mockReturnValue(true); - - renderWithClient(); - - await waitFor(() => { - expect(screen.queryByText("Usage")).not.toBeInTheDocument(); - }); - - (useDisableUsageIndicator as ReturnType).mockReturnValue(false); - }); - - it("should show Loading while fetching", () => { - mockGetRemainingUsers.mockImplementation(() => new Promise(() => {})); - - renderWithClient(); - - expect(screen.getByText("Loading...")).toBeInTheDocument(); - }); - - it("should show error message when fetch fails", async () => { - const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); - mockGetRemainingUsers.mockRejectedValue(new Error("Network error")); - - renderWithClient(); - - expect(await screen.findByText("Failed to load usage data")).toBeInTheDocument(); - - consoleSpy.mockRestore(); - }); - - it("should minimize when user clicks minimize button", async () => { - const user = userEvent.setup(); - renderWithClient(); - - await screen.findByText("Usage"); - - const minimizeButton = screen.getByTitle("Minimize"); - await user.click(minimizeButton); - - expect(screen.queryByText("Users")).not.toBeInTheDocument(); - expect(screen.getByTitle("Show usage details")).toBeInTheDocument(); - }); - - it("should restore from minimized when user clicks restore button", async () => { - const user = userEvent.setup(); - renderWithClient(); - - await screen.findByText("Usage"); - - await user.click(screen.getByTitle("Minimize")); - await user.click(screen.getByTitle("Show usage details")); - - expect(screen.getByText("Usage")).toBeInTheDocument(); - expect(screen.getByText("Users")).toBeInTheDocument(); - }); -}); diff --git a/ui/litellm-dashboard/src/components/UsageIndicator.tsx b/ui/litellm-dashboard/src/components/UsageIndicator.tsx deleted file mode 100644 index 6e7b5e9ec60..00000000000 --- a/ui/litellm-dashboard/src/components/UsageIndicator.tsx +++ /dev/null @@ -1,681 +0,0 @@ -import { useDisableUsageIndicator } from "@/app/(dashboard)/hooks/useDisableUsageIndicator"; -import { Badge } from "@tremor/react"; -import { - AlertTriangle, - Calendar, - ChevronDown, - ChevronUp, - Loader2, - Minus, - TrendingUp, - UserCheck, - Users, -} from "lucide-react"; -import { useEffect, useState } from "react"; -import { getRemainingUsers } from "./networking"; - -import { cn } from "@/lib/cva.config"; -import { getDaysUntilExpiration } from "@/utils/licenseUtils"; -import { useLicenseInfo } from "@/app/(dashboard)/hooks/license/useLicenseInfo"; - -interface UsageIndicatorProps { - accessToken: string | null; - width: number; -} - -interface UsageData { - total_users: number | null; - total_users_used: number; - total_users_remaining: number | null; - total_teams: number | null; - total_teams_used: number; - total_teams_remaining: number | null; -} - -// Format expiration for display -const formatExpirationDisplay = (daysRemaining: number | null): string => { - if (daysRemaining === null) return "No expiration"; - if (daysRemaining < 0) return "Expired"; - if (daysRemaining === 0) return "Expires today"; - if (daysRemaining === 1) return "1 day remaining"; - if (daysRemaining < 30) return `${daysRemaining} days remaining`; - if (daysRemaining < 60) return "1 month remaining"; - const months = Math.floor(daysRemaining / 30); - return `${months} months remaining`; -}; - -export default function UsageIndicator({ accessToken, width = 220 }: UsageIndicatorProps) { - const disableUsageIndicator = useDisableUsageIndicator(); - const [isExpanded, setIsExpanded] = useState(false); - const [isMinimized, setIsMinimized] = useState(false); - const [data, setData] = useState(null); - const [isLoading, setIsLoading] = useState(false); - const [error, setError] = useState(null); - - const licenseInfo = useLicenseInfo(accessToken).data ?? null; - - useEffect(() => { - const fetchData = async () => { - if (!accessToken) return; - - setIsLoading(true); - setError(null); - - try { - const usageResult = await getRemainingUsers(accessToken); - setData(usageResult); - } catch (err) { - console.error("Failed to fetch usage data:", err); - setError("Failed to load usage data"); - } finally { - setIsLoading(false); - } - }; - - fetchData(); - }, [accessToken]); - - // Calculate license expiration metrics - const daysUntilExpiration = licenseInfo?.expiration_date ? getDaysUntilExpiration(licenseInfo.expiration_date) : null; - const isLicenseExpired = daysUntilExpiration !== null && daysUntilExpiration < 0; - const isLicenseExpiringSoon = daysUntilExpiration !== null && daysUntilExpiration >= 0 && daysUntilExpiration < 30; - - // Calculate derived values from data - const getUsageMetrics = (data: UsageData | null) => { - if (!data) { - return { - isOverLimit: false, - isNearLimit: false, - usagePercentage: 0, - userMetrics: { - isOverLimit: false, - isNearLimit: false, - usagePercentage: 0, - }, - teamMetrics: { - isOverLimit: false, - isNearLimit: false, - usagePercentage: 0, - }, - }; - } - - // User metrics - const userUsagePercentage = data.total_users ? (data.total_users_used / data.total_users) * 100 : 0; - const userIsOverLimit = userUsagePercentage > 100; - const userIsNearLimit = userUsagePercentage >= 80 && userUsagePercentage <= 100; - - // Team metrics - const teamUsagePercentage = data.total_teams ? (data.total_teams_used / data.total_teams) * 100 : 0; - const teamIsOverLimit = teamUsagePercentage > 100; - const teamIsNearLimit = teamUsagePercentage >= 80 && teamUsagePercentage <= 100; - - // Combined status (worst case scenario) - const isOverLimit = userIsOverLimit || teamIsOverLimit; - const isNearLimit = (userIsNearLimit || teamIsNearLimit) && !isOverLimit; - const usagePercentage = Math.max(userUsagePercentage, teamUsagePercentage); - - return { - isOverLimit, - isNearLimit, - usagePercentage, - userMetrics: { - isOverLimit: userIsOverLimit, - isNearLimit: userIsNearLimit, - usagePercentage: userUsagePercentage, - }, - teamMetrics: { - isOverLimit: teamIsOverLimit, - isNearLimit: teamIsNearLimit, - usagePercentage: teamUsagePercentage, - }, - }; - }; - - const { isOverLimit, isNearLimit, usagePercentage, userMetrics, teamMetrics } = getUsageMetrics(data); - - // Include license status in overall status - const hasAnyIssue = isOverLimit || isNearLimit || isLicenseExpired || isLicenseExpiringSoon; - const hasError = isOverLimit || isLicenseExpired; - const hasWarning = (isNearLimit || isLicenseExpiringSoon) && !hasError; - - const getStatusColor = () => { - if (hasError) return "red"; - if (hasWarning) return "yellow"; - return "green"; - }; - - const getStatusIcon = () => { - if (hasError) return ; - if (hasWarning) return ; - return null; - }; - - // Minimized view - just a small restore button - const MinimizedView = () => { - return ( -
- -
- ); - }; - - // Sidebar/nav style component - const NavStyleView = () => { - if (isMinimized) { - return ; - } - - if (isLoading) { - return ( -
- - Loading... -
- ); - } - - if (error || !data) { - return ( -
-
- - {error || "No data"} -
- -
- ); - } - - return ( -
- {/* Main nav item style */} -
- - - {/* Minimize button */} - -
- - {/* Expanded details - simple and compact */} - {isExpanded && ( -
- {/* License expiration section */} - {licenseInfo?.has_license && licenseInfo.expiration_date && ( -
-
- - License -
-
- {isLicenseExpired ? ( - - ) : isLicenseExpiringSoon ? ( - - ) : null} - {formatExpirationDisplay(daysUntilExpiration)} -
-
- )} - - {/* Users section */} - {data.total_users !== null && ( -
-
- - - {data.total_users_used}/{data.total_users} - - users -
- - {/* User progress bar */} -
-
-
- - {(userMetrics.isOverLimit || userMetrics.isNearLimit) && ( -
- {userMetrics.isOverLimit ? ( - - ) : ( - - )} - Users {userMetrics.isOverLimit ? "Over Limit" : "Near Limit"} -
- )} -
- )} - - {/* Teams section */} - {data.total_teams !== null && ( -
-
- - - {data.total_teams_used}/{data.total_teams} - - teams -
- - {/* Team progress bar */} -
-
-
- - {(teamMetrics.isOverLimit || teamMetrics.isNearLimit) && ( -
- {teamMetrics.isOverLimit ? ( - - ) : ( - - )} - Teams {teamMetrics.isOverLimit ? "Over Limit" : "Near Limit"} -
- )} -
- )} -
- )} -
- ); - }; - - // Optimized CardStyleView for 220px width - const CardStyleView = () => { - if (isMinimized) { - return ( - - ); - } - - if (isLoading) { - return ( -
-
- - Loading... -
-
- ); - } - - if (error || !data) { - return ( -
-
-
- {error || "No data"} -
- -
-
- ); - } - - return ( -
-
-
- - Usage -
- -
- - {/* Compact stats optimized for 220px */} -
- {/* License expiration section */} - {licenseInfo?.has_license && licenseInfo.expiration_date && ( -
-
- - License - - {isLicenseExpired ? "Expired" : isLicenseExpiringSoon ? "Expiring soon" : "OK"} - -
-
- Status: - - {formatExpirationDisplay(daysUntilExpiration)} - -
- {licenseInfo.license_type && ( -
- Type: - {licenseInfo.license_type} -
- )} -
- )} - - {/* Users section */} - {data.total_users !== null && ( -
-
- - Users - - {userMetrics.isOverLimit ? "Over limit" : userMetrics.isNearLimit ? "Near limit" : "OK"} - -
-
- Used: - - {data.total_users_used}/{data.total_users} - -
-
- Remaining: - - {data.total_users_remaining} - -
-
- Usage: - {Math.round(userMetrics.usagePercentage)}% -
- - {/* User progress bar */} -
-
-
-
- )} - - {/* Teams section */} - {data.total_teams !== null && ( -
-
- - Teams - - {teamMetrics.isOverLimit ? "Over limit" : teamMetrics.isNearLimit ? "Near limit" : "OK"} - -
-
- Used: - - {data.total_teams_used}/{data.total_teams} - -
-
- Remaining: - - {data.total_teams_remaining} - -
-
- Usage: - {Math.round(teamMetrics.usagePercentage)}% -
- - {/* Team progress bar */} -
-
-
-
- )} -
-
- ); - }; - - // Don't render anything if disabled, no access token, or if both total_users and total_teams are null - if (disableUsageIndicator || !accessToken || (data?.total_users === null && data?.total_teams === null)) { - return null; - } - - // Fixed positioning with proper spacing from edges - return ( -
- -
- ); -} From edc38eab348d33d01a8eeb2347e985494c69c5eb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 15 Jul 2026 17:23:17 -0700 Subject: [PATCH 28/62] build(deps): update ddtrace to the 4.x line A single ddtrace constraint now covers every supported Python version, so this collapses the version split introduced in #33438. Also aligns the build_from_pip image pin and updates the type-only Tracer import to its current module path --- .../build_from_pip/Dockerfile.build_from_pip | 2 +- litellm/litellm_core_utils/dd_tracing.py | 2 +- pyproject.toml | 3 +- uv.lock | 129 +++--------------- 4 files changed, 25 insertions(+), 111 deletions(-) diff --git a/docker/build_from_pip/Dockerfile.build_from_pip b/docker/build_from_pip/Dockerfile.build_from_pip index bda742c71a9..372606a5b0f 100644 --- a/docker/build_from_pip/Dockerfile.build_from_pip +++ b/docker/build_from_pip/Dockerfile.build_from_pip @@ -36,7 +36,7 @@ RUN uv venv --python python && \ "opentelemetry-api==1.28.0" \ "opentelemetry-sdk==1.28.0" \ "opentelemetry-exporter-otlp==1.28.0" \ - "ddtrace==2.19.0" \ + "ddtrace==4.11.0" \ "sentry-sdk==2.21.0" \ "mangum==0.17.0" \ "azure-ai-contentsafety==1.0.0" \ diff --git a/litellm/litellm_core_utils/dd_tracing.py b/litellm/litellm_core_utils/dd_tracing.py index ae4f46c38bd..3a1bd72e1a5 100644 --- a/litellm/litellm_core_utils/dd_tracing.py +++ b/litellm/litellm_core_utils/dd_tracing.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Optional, Union from litellm.secret_managers.main import get_secret_bool if TYPE_CHECKING: - from ddtrace.tracer import Tracer as DD_TRACER + from ddtrace.trace import Tracer as DD_TRACER else: DD_TRACER = Any diff --git a/pyproject.toml b/pyproject.toml index 8d88344a3f7..d70e99c5775 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,8 +129,7 @@ proxy-runtime = [ "opentelemetry-sdk==1.28.0", "opentelemetry-exporter-otlp==1.28.0", "opentelemetry-instrumentation-fastapi==0.49b0", - "ddtrace>=2.19.0,<3.0; python_version < '3.14'", - "ddtrace>=4.0.0,<5.0; python_version >= '3.14'", + "ddtrace>=4.8.2,<5.0", "sentry-sdk>=2.21.0,<3.0", "mangum>=0.17.0,<1.0", "azure-ai-contentsafety>=1.0.0,<2.0", diff --git a/uv.lock b/uv.lock index 80ab2f42f67..2439d1488cd 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-07-13T00:19:39.570486Z" exclude-newer-span = "P3D" [manifest] @@ -222,9 +222,9 @@ name = "aiologic" version = "0.17.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "sniffio", marker = "python_full_version < '3.14'" }, + { name = "sniffio", marker = "python_full_version < '3.13'" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, - { name = "wrapt", marker = "python_full_version < '3.14'" }, + { name = "wrapt", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/53/a7/809482759f40079f4c4328c7318bf569ae25d457f5017aad30a1b9aafedc/aiologic-0.17.0.tar.gz", hash = "sha256:65aa058e858c94cd208badb188e7f00b54dcabb3ba85b34f794db98074d108b9", size = 251625, upload-time = "2026-06-14T12:24:35.367Z" } wheels = [ @@ -1047,7 +1047,7 @@ name = "coloredlogs" version = "15.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "humanfriendly" }, + { name = "humanfriendly", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/c7/eed8f27100517e8c0e6b923d5f0845d0cb99763da6fdee00478f91db7325/coloredlogs-15.0.1.tar.gz", hash = "sha256:7c991aa71a4577af2f82600d8f8f3a89f936baeaf9b50a9c197da014e5bf16b0", size = 278520, upload-time = "2021-06-11T10:22:45.202Z" } wheels = [ @@ -1420,7 +1420,7 @@ name = "culsans" version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiologic", marker = "python_full_version < '3.14'" }, + { name = "aiologic", marker = "python_full_version < '3.13'" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/d9/e3/49afa1bc180e0d28008ec6bcdf82a4072d1c7a41032b5b759b60814ca4b0/culsans-0.11.0.tar.gz", hash = "sha256:0b43d0d05dce6106293d114c86e3fb4bfc63088cfe8ff08ed3fe36891447fe33", size = 107546, upload-time = "2025-12-31T23:15:38.196Z" } @@ -1464,80 +1464,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c3/be/d0d44e092656fe7a06b55e6103cbce807cdbdee17884a5367c68c9860853/dataclasses_json-0.6.7-py3-none-any.whl", hash = "sha256:0dbf33f26c8d5305befd61b39d2b3414e8a407bedc2834dea9b8d642666fb40a", size = 28686, upload-time = "2024-06-09T16:20:16.715Z" }, ] -[[package]] -name = "ddtrace" -version = "2.19.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version == '3.13.*'", - "python_full_version == '3.12.*'", - "python_full_version == '3.11.*'", - "python_full_version < '3.11'", -] -dependencies = [ - { name = "bytecode", marker = "python_full_version < '3.14'" }, - { name = "envier", marker = "python_full_version < '3.14'" }, - { name = "legacy-cgi", marker = "python_full_version == '3.13.*'" }, - { name = "opentelemetry-api", marker = "python_full_version < '3.14'" }, - { name = "protobuf", marker = "python_full_version < '3.14'" }, - { name = "typing-extensions", marker = "python_full_version < '3.14'" }, - { name = "wrapt", marker = "python_full_version < '3.14'" }, - { name = "xmltodict", marker = "python_full_version < '3.14'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c3/06/417a8a9a8c89dc2fdb94c3acdb3f6f9da835e109c2a217fb5863d0d97df9/ddtrace-2.19.0.tar.gz", hash = "sha256:90d217b1906074881afd3e656a3cd1a630dd798bd25077254588c382a4075345", size = 8708460, upload-time = "2025-01-16T17:19:46.303Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0c/be/b3fc069ff2a20cc1d053b030268ba6999926232ce2195b4958486a9035ea/ddtrace-2.19.0-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:50d79ff042868b4d1d80b424285d755d9e0d466119399c2166a3894b178b85fd", size = 4412309, upload-time = "2025-01-16T17:16:10.884Z" }, - { url = "https://files.pythonhosted.org/packages/14/69/2d42669829c09eefbf4cbabb94dfe7615ee4610019ded28c9634411a574c/ddtrace-2.19.0-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:e17b3b8e1cadf23ed8e4466679a0cb1262aa00190bddcd0fc0f5f6f9a9c25480", size = 3051126, upload-time = "2025-01-16T17:16:15.377Z" }, - { url = "https://files.pythonhosted.org/packages/57/e0/82d3b5d474ea66e777c38e584053ee7f6ac923218642fdf4967857f48daa/ddtrace-2.19.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7804977b388fed1b1cbb0ef138100be923bf7afbe7d25357fcee07315a66cc8b", size = 6087687, upload-time = "2025-01-16T17:16:17.314Z" }, - { url = "https://files.pythonhosted.org/packages/54/ba/051ea8720695a8c0ecf7a5d9dcbc7000da18b2cd4efe27af69c09999832d/ddtrace-2.19.0-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c96ae2f074e422202f2b98a021b9fb864fc08bd5864eff27b0ec9da5919c0b1e", size = 2852443, upload-time = "2025-01-16T17:16:20.242Z" }, - { url = "https://files.pythonhosted.org/packages/d8/94/38f7706bbc1b3c010aab457a94edc07ac7145ebfd8ab01797634b60746cb/ddtrace-2.19.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a14662366c5c1d8898c057ef6820de85d71b3ada6fd89638c5ec4ba9d45c21b7", size = 6420509, upload-time = "2025-01-16T17:16:22.521Z" }, - { url = "https://files.pythonhosted.org/packages/37/f8/b900ffbdf85a06220ca04905caa066dfc1f60643c3d501e92fee08d32951/ddtrace-2.19.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:1638fff37abf61d16f3dbef009c45d4c33b962324b10f642fb7966d0055c28e9", size = 7073719, upload-time = "2025-01-16T17:16:24.518Z" }, - { url = "https://files.pythonhosted.org/packages/e9/a4/12ffed1870c6ecc638283163d11cb675c2840a46dbd741acb2568bf94a6a/ddtrace-2.19.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:3f72627f1887d628b025d227a642e5ae30884eedd6f6ef1afee02461bc19c95f", size = 3918050, upload-time = "2025-01-16T17:16:27.279Z" }, - { url = "https://files.pythonhosted.org/packages/29/35/d4c6a99df2a7ea6219c9b6390ad66ae88f63c86bc7ef84c2a3784f5b5798/ddtrace-2.19.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:d5b8ac83af10e567564d1f016d72549a8b83e06f4c6a3440f7f610ede0b51954", size = 7463892, upload-time = "2025-01-16T17:16:29.299Z" }, - { url = "https://files.pythonhosted.org/packages/49/7e/881b58c69d7e2316ccc99e8c3b1d4b4d382d5c704a9a0aebc571353b1413/ddtrace-2.19.0-cp310-cp310-win32.whl", hash = "sha256:17971717ad481c2273336957a8c2f328f2e7776f2065821c00332f33cdaa2053", size = 3120778, upload-time = "2025-01-16T17:16:31.278Z" }, - { url = "https://files.pythonhosted.org/packages/27/39/d5d92f7d0f6d3f98c708c514498562965bd8bfbea8234d1cf3a2ab9f245e/ddtrace-2.19.0-cp310-cp310-win_amd64.whl", hash = "sha256:41629eaa0e16367a45e5fe64b0bd969dd31eb2067e0224e48f149ea976ed5848", size = 3348128, upload-time = "2025-01-16T17:16:33.199Z" }, - { url = "https://files.pythonhosted.org/packages/a6/ec/ac70516f825aba5a5bea78cea568fef6a6c34b80c621bea70d3f9128d3f2/ddtrace-2.19.0-cp311-cp311-macosx_12_0_universal2.whl", hash = "sha256:e58123e8bce549aa159cc3748248987dd2ac63ba4c69c7f0b0d49c2d2c05d20a", size = 4414054, upload-time = "2025-01-16T17:16:36.181Z" }, - { url = "https://files.pythonhosted.org/packages/e7/73/4f0cb04aef8450f23fbe6fc0ba66868bc9e415830fecbe88b65c658866e0/ddtrace-2.19.0-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:1cc7b2b7e9396c17b0356f550b278d1adc5112a0da57e9052169a98b75cbdb66", size = 3052135, upload-time = "2025-01-16T17:16:38.102Z" }, - { url = "https://files.pythonhosted.org/packages/f7/20/0e8d2ef1b1d2c7b4f55b4d2e978bb143fafac36cdc456e8e521b9559c484/ddtrace-2.19.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cb3a4941c7604f0ee56713c207c1baf8984cb25e0acd9872512d2cd1cd9ef40e", size = 6093484, upload-time = "2025-01-16T17:16:41.192Z" }, - { url = "https://files.pythonhosted.org/packages/28/f8/af03509c93d91fc35b71c89b02e86635ef2c0d5c56379048c93b4b1d338b/ddtrace-2.19.0-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f906e6b05a66c85c7049076c69186b6208a9868086cb114e5cd784e5705d11ed", size = 2858353, upload-time = "2025-01-16T17:16:43.276Z" }, - { url = "https://files.pythonhosted.org/packages/dc/de/9062ccdd6b0bc00b15dc58bd7bb7ad1e27ab0c78cb4c9ad7218f7dd58106/ddtrace-2.19.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:544a41bba75c547c52595cec535d55ed3a8b109068121682ebe04e48a3af73d9", size = 6426319, upload-time = "2025-01-16T17:16:46.095Z" }, - { url = "https://files.pythonhosted.org/packages/c6/bc/2c8b9afa39c5b8370cb8587a8325bcf01ff6c5d87b07690ae764c3e02a9f/ddtrace-2.19.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d4c9c4dc015e285368ac29ca263d41ff9480b1df42c5599860c139007a11dd54", size = 7076423, upload-time = "2025-01-16T17:16:49.439Z" }, - { url = "https://files.pythonhosted.org/packages/18/9c/caa119adf66d4a6b0e7f7d0de8ed6ecfb18ff2249eb55aeed03f412233ce/ddtrace-2.19.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:7965330c03a4793d8bc71c702d49832b9900cf0e3b5f36e9d4c9037285f5fc73", size = 3919949, upload-time = "2025-01-16T17:16:51.98Z" }, - { url = "https://files.pythonhosted.org/packages/66/53/0d6b96db5c9ee6fdaf010adc968c7dadadbc5031d05e382dbb3088a20ea7/ddtrace-2.19.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:7b203be50ca19182120063a34ccffc34e3016555c8ceb5b1439ae61d7ef88ad0", size = 7470201, upload-time = "2025-01-16T17:16:54.498Z" }, - { url = "https://files.pythonhosted.org/packages/7b/d4/81b8df76e10dbcd85ad1f0bb17075d6334c1abcb4136ad9a03835880fabb/ddtrace-2.19.0-cp311-cp311-win32.whl", hash = "sha256:bed9aa688e7f0185f96407fe9bd20192e767aa812fdaac5552bf3edc4fa5182c", size = 3120927, upload-time = "2025-01-16T17:16:57.998Z" }, - { url = "https://files.pythonhosted.org/packages/60/73/3ea8f4ddcf3b451ca2523767262fd8d9df76aa1e0403932c5cf49ff73eab/ddtrace-2.19.0-cp311-cp311-win_amd64.whl", hash = "sha256:a40e64a96dbdb5b2124b54051c2de371b895175b62a97273b7f527be9721d8c2", size = 3352777, upload-time = "2025-01-16T17:16:59.998Z" }, - { url = "https://files.pythonhosted.org/packages/62/64/8c696adb83f2a1a5310d8f64094d8d76417928c136f1b2fc55bb912977ad/ddtrace-2.19.0-cp312-cp312-macosx_12_0_universal2.whl", hash = "sha256:8f5e6e0086717cc7c8fd1ad3da2ee7d5cb30ba3eb0d75ee79b070b310443d884", size = 4852896, upload-time = "2025-01-16T17:17:02.584Z" }, - { url = "https://files.pythonhosted.org/packages/ef/b9/2cd4347db133128429f60044e40600c9016a98e147b110d7020e8767ee60/ddtrace-2.19.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:aa38304a6b5c937154acd33dafdc8d1cbdc4c4879e135578515dd9be44241b2c", size = 3280741, upload-time = "2025-01-16T17:17:04.77Z" }, - { url = "https://files.pythonhosted.org/packages/85/a2/a94bd0e39657b45008cce9c33931f824f27a3db2da655b0b599c44d51617/ddtrace-2.19.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e77b42bd5a269f2bc1ad0ba8141987634c288ee96366ea9505aec8871ee5662f", size = 6062585, upload-time = "2025-01-16T17:17:07.095Z" }, - { url = "https://files.pythonhosted.org/packages/66/07/f655ede9fbf1c7de2a0a271687d0a31c39e4afc46102b1e73eac342298d8/ddtrace-2.19.0-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:56816cdda82b18e8e99ee60d0f70a5cdbd57fb54bdc9038ba01d779139db1fcc", size = 2827198, upload-time = "2025-01-16T17:17:11.978Z" }, - { url = "https://files.pythonhosted.org/packages/0f/9d/a193623a7d9a5226cd63ddbdc42250ef3e6d4b37bc77725ff06b2a9838c4/ddtrace-2.19.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9fd3237342fa7753c47161904bbb3bb691625a99a4b396069fd1db927d20a74c", size = 6398024, upload-time = "2025-01-16T17:17:14.303Z" }, - { url = "https://files.pythonhosted.org/packages/56/76/43c132d259d1fd710a5ece3abac4e6d7789626ce1ae66536ed0e22fc5361/ddtrace-2.19.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:713beebab7398310f0753e33234cb70b91b3eb387525a7cae5897d19471917f4", size = 7041348, upload-time = "2025-01-16T17:17:16.996Z" }, - { url = "https://files.pythonhosted.org/packages/0a/9e/f59030213600c58f87b4d5d814ded8b9453cbbfdc7c3a02a313f07c62db1/ddtrace-2.19.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:adaf1ef268c5bb3599f3a1e34b1089d77b6e323cd1ec31da482f794f64213aa6", size = 3886975, upload-time = "2025-01-16T17:17:20.77Z" }, - { url = "https://files.pythonhosted.org/packages/84/c6/626560e37f0024572456d7cc2cafb0ab61da22deb2f9b218231d43053325/ddtrace-2.19.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0c38952a4f4d1ed61d53bdc55acb9d931c355e1e5a90001b95576e99dcacfc11", size = 7434753, upload-time = "2025-01-16T17:17:24.304Z" }, - { url = "https://files.pythonhosted.org/packages/1f/2a/3c181fc7f2021ec05e95586ca8fa8236f1429adfabb6152bb950b79247a4/ddtrace-2.19.0-cp312-cp312-win32.whl", hash = "sha256:045773c382aada18feeb5584fdba9aa47ff660ac93a94b43b24434760c77802a", size = 3108737, upload-time = "2025-01-16T17:17:27.696Z" }, - { url = "https://files.pythonhosted.org/packages/cd/81/b000c6919d9cc204fead0069b3523d6a65d0da21a45a686a872b4201013e/ddtrace-2.19.0-cp312-cp312-win_amd64.whl", hash = "sha256:e96980c8e81831c7cb367b1ab066ba4cfaf389be1099c0f15985484a8de6d80d", size = 3343024, upload-time = "2025-01-16T17:17:29.98Z" }, - { url = "https://files.pythonhosted.org/packages/d9/55/32f7142cc96410a534868eb553ef9d238cf44d2cb10c2107cf880d9d42b9/ddtrace-2.19.0-cp313-cp313-macosx_12_0_universal2.whl", hash = "sha256:c3cccef7e15a561ad5e5699ca2f6045d7d1e655c487fd2d617b9541152c3f217", size = 4832649, upload-time = "2025-01-16T17:17:32.785Z" }, - { url = "https://files.pythonhosted.org/packages/1a/60/5d1e99cfa6bc29d13eae55fbe6b395138ce17b96d6e95c9c7b57c071d410/ddtrace-2.19.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:a13acabcf0fad276e55e9d35ade74aae91b9629367e06692788ee5ed484491d6", size = 3269614, upload-time = "2025-01-16T17:17:35.273Z" }, - { url = "https://files.pythonhosted.org/packages/5d/75/b3b00c1325d64ab1445a7965554b7a311842ae26f6995a0f64348c597848/ddtrace-2.19.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:724f3954e16bf66f0f45479ba3fffb4fa39fe4e43667c3085d9090b17ea9242d", size = 6016732, upload-time = "2025-01-16T17:17:37.783Z" }, - { url = "https://files.pythonhosted.org/packages/84/a7/ec1fa6f8ad7254f9baed33c37cb88abf0f324e2740792f8e5f24bd1fbfea/ddtrace-2.19.0-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:47bb8980fb8d711d96f66d8727196476bbdafe2cccc501c93b5352c5797a4422", size = 2815983, upload-time = "2025-01-16T17:17:40.305Z" }, - { url = "https://files.pythonhosted.org/packages/4a/1a/2b8102e738bc4ed335dd3a5bfeba21b554b73abc9ea8d1c47cb7f3ccfbe0/ddtrace-2.19.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:661341280a69d8ceb91e48f67cfa33adb3af901b09d329b8375b1a3ba04a68b7", size = 6351654, upload-time = "2025-01-16T17:17:44.094Z" }, - { url = "https://files.pythonhosted.org/packages/0e/30/56095f289ae7689cb789f2c85e0e227b02bd8de548c25ab0c952cc823051/ddtrace-2.19.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:facae23052586b171c47faecb622a0c8a15beeab0fb3af4d53367a749b1cbded", size = 6996885, upload-time = "2025-01-16T17:17:47.304Z" }, - { url = "https://files.pythonhosted.org/packages/2b/c0/c315500dbde69a4193665c964dd56e9be523b7e05979718140ad2c9a6821/ddtrace-2.19.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:4cb13fdb6587ff1c460b7e673d83979280efa0a5a5a4104fc92ebcaf0c6ca36e", size = 3881860, upload-time = "2025-01-16T17:17:52.081Z" }, - { url = "https://files.pythonhosted.org/packages/26/8c/e1a7043e562b5b29fb5d0930630a18078fecb1c30ca6776221ce0dab6f95/ddtrace-2.19.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:d36a16e8746cb38a143faa6e1cd10927bf4a482c29f4010afde4bd0f4bb89db4", size = 7390107, upload-time = "2025-01-16T17:17:54.826Z" }, -] - [[package]] name = "ddtrace" version = "4.11.0" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.14'", -] dependencies = [ - { name = "bytecode", marker = "python_full_version >= '3.14'" }, - { name = "envier", marker = "python_full_version >= '3.14'" }, - { name = "opentelemetry-api", marker = "python_full_version >= '3.14'" }, - { name = "wrapt", marker = "python_full_version >= '3.14'" }, + { name = "bytecode" }, + { name = "envier" }, + { name = "opentelemetry-api" }, + { name = "wrapt" }, ] sdist = { url = "https://files.pythonhosted.org/packages/81/51/a628f0177274bab5b67c93a1558fd222babb15286a427e4e8c1d65a265b2/ddtrace-4.11.0.tar.gz", hash = "sha256:260c5b46e80565f4fd08cec2650f707627fb57cd6f8951a2d8c9e6a02b490074", size = 2422158, upload-time = "2026-07-10T08:57:53.396Z" } wheels = [ @@ -3031,7 +2966,7 @@ name = "humanfriendly" version = "10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pyreadline3", marker = "sys_platform == 'win32'" }, + { name = "pyreadline3", marker = "python_full_version < '3.14' and sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/3f/2c29224acb2e2df4d2046e4c73ee2662023c58ff5b113c4c1adac0886c43/humanfriendly-10.0.tar.gz", hash = "sha256:6b0b831ce8f15f7300721aa49829fc4e83921a9a301cc7f606be6686a2288ddc", size = 360702, upload-time = "2021-09-17T21:40:43.31Z" } wheels = [ @@ -3795,15 +3730,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/a0/b91504515c1f9a299fc157967ffbd2f0321bce0516a3d5b89f6f4cad0355/lazy_object_proxy-1.12.0-pp39.pp310.pp311.graalpy311-none-any.whl", hash = "sha256:c3b2e0af1f7f77c4263759c4824316ce458fabe0fceadcd24ef8ca08b2d1e402", size = 15072, upload-time = "2025-08-22T13:50:05.498Z" }, ] -[[package]] -name = "legacy-cgi" -version = "2.6.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f4/9c/91c7d2c5ebbdf0a1a510bfa0ddeaa2fbb5b78677df5ac0a0aa51cf7125b0/legacy_cgi-2.6.4.tar.gz", hash = "sha256:abb9dfc7835772f7c9317977c63253fd22a7484b5c9bbcdca60a29dcce97c577", size = 24603, upload-time = "2025-10-27T05:20:05.395Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8c/7e/e7394eeb49a41cc514b3eb49020223666cbf40d86f5721c2f07871e6d84a/legacy_cgi-2.6.4-py3-none-any.whl", hash = "sha256:7e235ce58bf1e25d1fc9b2d299015e4e2cd37305eccafec1e6bac3fc04b878cd", size = 20035, upload-time = "2025-10-27T05:20:04.289Z" }, -] - [[package]] name = "litellm" version = "1.94.0" @@ -3887,8 +3813,7 @@ proxy-runtime = [ { name = "anthropic", extra = ["vertex"] }, { name = "azure-ai-contentsafety" }, { name = "azure-storage-file-datalake" }, - { name = "ddtrace", version = "2.19.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, - { name = "ddtrace", version = "4.11.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, + { name = "ddtrace" }, { name = "detect-secrets" }, { name = "google-cloud-aiplatform" }, { name = "google-genai" }, @@ -4023,8 +3948,7 @@ requires-dist = [ { name = "boto3", marker = "extra == 'proxy'", specifier = ">=1.43.1,<2.0" }, { name = "click", specifier = ">=8.0.0,<9.0" }, { name = "cryptography", marker = "extra == 'proxy'", specifier = ">=48.0.1,<49.0" }, - { name = "ddtrace", marker = "python_full_version >= '3.14' and extra == 'proxy-runtime'", specifier = ">=4.0.0,<5.0" }, - { name = "ddtrace", marker = "python_full_version < '3.14' and extra == 'proxy-runtime'", specifier = ">=2.19.0,<3.0" }, + { name = "ddtrace", marker = "extra == 'proxy-runtime'", specifier = ">=4.8.2,<5.0" }, { name = "detect-secrets", marker = "extra == 'proxy-runtime'", specifier = ">=1.5.0,<2.0" }, { name = "diskcache", marker = "extra == 'caching'", specifier = ">=5.6.3,<6.0" }, { name = "expression", marker = "extra == 'proxy'", specifier = ">=5.6.0,<6.0" }, @@ -4540,7 +4464,7 @@ version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/fd/15/76f86faa0902836cc133939732f7611ace68cf54148487a99c539c272dc8/ml_dtypes-0.4.1.tar.gz", hash = "sha256:fad5f2de464fd09127e49b7fd1252b9006fb43d2edc1ff112d390c324af5ca7a", size = 692594, upload-time = "2024-09-13T19:07:11.624Z" } wheels = [ @@ -7201,16 +7125,16 @@ name = "redisvl" version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "coloredlogs" }, - { name = "ml-dtypes" }, + { name = "coloredlogs", marker = "python_full_version < '3.14'" }, + { name = "ml-dtypes", marker = "python_full_version < '3.14'" }, { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "pydantic" }, - { name = "python-ulid" }, - { name = "pyyaml" }, - { name = "redis" }, - { name = "tabulate" }, - { name = "tenacity" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, + { name = "pydantic", marker = "python_full_version < '3.14'" }, + { name = "python-ulid", marker = "python_full_version < '3.14'" }, + { name = "pyyaml", marker = "python_full_version < '3.14'" }, + { name = "redis", marker = "python_full_version < '3.14'" }, + { name = "tabulate", marker = "python_full_version < '3.14'" }, + { name = "tenacity", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/21/33/ab14865a0b2a31b1d003c29e7e8ea3a7a2f2c8ecb24e58e58d606e1f031b/redisvl-0.4.1.tar.gz", hash = "sha256:fd6a36426ba94792c0efca20915c31232d4ee3cc58eb23794a62c142696401e6", size = 77688, upload-time = "2025-02-21T22:51:41.389Z" } wheels = [ @@ -9223,15 +9147,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a4/f5/10b68b7b1544245097b2a1b8238f66f2fc6dcaeb24ba5d917f52bd2eed4f/wsproto-1.3.2-py3-none-any.whl", hash = "sha256:61eea322cdf56e8cc904bd3ad7573359a242ba65688716b0710a5eb12beab584", size = 24405, upload-time = "2025-11-20T18:18:00.454Z" }, ] -[[package]] -name = "xmltodict" -version = "1.0.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/19/70/80f3b7c10d2630aa66414bf23d210386700aa390547278c789afa994fd7e/xmltodict-1.0.4.tar.gz", hash = "sha256:6d94c9f834dd9e44514162799d344d815a3a4faec913717a9ecbfa5be1bb8e61", size = 26124, upload-time = "2026-02-22T02:21:22.074Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/38/34/98a2f52245f4d47be93b580dae5f9861ef58977d73a79eb47c58f1ad1f3a/xmltodict-1.0.4-py3-none-any.whl", hash = "sha256:a4a00d300b0e1c59fc2bfccb53d7b2e88c32f200df138a0dd2229f842497026a", size = 13580, upload-time = "2026-02-22T02:21:21.039Z" }, -] - [[package]] name = "xxhash" version = "3.7.0" From da472722279e533fc3b5e569826b439c3243f17d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:33:46 -0700 Subject: [PATCH 29/62] test(reasoning_effort_grid): enable azure fable-5 and opus-4-8 grid cells --- .../reasoning_effort_grid/grid_spec.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py index 762606bb3a0..c47fddb1d8d 100644 --- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py +++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py @@ -212,12 +212,6 @@ AZURE_AI_MODELS: Tuple[ModelEntry, ...] = ( mode="adaptive", required_env=_AZURE_FOUNDRY_REQ, caps=_CAPS_XHIGH_MAX, - fail_reason=( - "claude-fable-5 has no deployment on the CI Microsoft Foundry " - "resource yet; Foundry returns DeploymentNotFound until someone " - "creates the fable-5 deployment, so this cell stays loud in CI. " - "Remove this fail_reason once the deployment exists." - ), ), ModelEntry( alias="azure-claude-opus-4-8", @@ -225,12 +219,6 @@ AZURE_AI_MODELS: Tuple[ModelEntry, ...] = ( mode="adaptive", required_env=_AZURE_FOUNDRY_REQ, caps=_CAPS_XHIGH_MAX, - fail_reason=( - "claude-opus-4-8 has no deployment on the CI Microsoft Foundry " - "resource yet; Foundry returns DeploymentNotFound until someone " - "creates the opus-4-8 deployment, so this cell stays loud in CI. " - "Remove this fail_reason once the deployment exists." - ), ), ModelEntry( alias="azure-claude-opus-4-7", From 032a2f2d76ef06cb71f982610503de73afd7bde0 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 17:44:21 -0700 Subject: [PATCH 30/62] fix(mcp): enforce the issuer trust anchor at every endpoint adoption site When an admin pins an issuer, RFC 8414 section 3.3 makes that issuer the sole authoritative source of the authorization and token endpoints, so a compromised or misconfigured upstream cannot smuggle a token endpoint by echoing the pinned authorize URL. The first cut enforced that only on the database build path; the carry-forward, persistence, config-load, serialization and sanitization paths could still restore or emit upstream-derived endpoints for an issuer-anchored server, which is the class of gap the review flagged. Every site now routes through one predicate. _endpoints_yield_to_issuer returns all-None whenever the issuer is the anchor, so both build paths, has_all_upstream_oauth_fields, needs_discovery and the endpoint merge defer to the issuer. _carry_forward_resolved_oauth_endpoints carries only scopes for an issuer-anchored server and fails closed on endpoints. _persist_discovered_oauth_endpoints skips endpoint writes under the anchor. The two table serializers round-trip the issuer and both non-admin sanitizers redact it. Scope selection stays resource-driven per the MCP authorization spec: _fetch_issuer_anchored_oauth_metadata takes endpoints from the issuer document and scopes from the resource document. The OAuth metadata resolution and corroboration gating for the database build path move into _resolve_table_oauth_metadata so build_mcp_server_from_table stays within the cyclomatic-complexity budget without changing behavior. Regression tests pin the invariant at each site: the issuer overrides stored endpoints even when they are populated, carry-forward does not restore endpoints under the anchor, persistence does not write endpoints under the anchor, a url or auth_type change clears stale issuer-scoped fields even when resubmitted unchanged, the Azure heuristic stays reachable under a required issuer, and anchored metadata takes endpoints from the issuer while scopes come from the resource --- litellm/proxy/_experimental/mcp_server/db.py | 12 +- .../mcp_server/mcp_server_manager.py | 169 +++++++++++++----- .../mcp_management_endpoints.py | 3 + .../mcp_server/test_mcp_partial_update.py | 31 ++++ .../mcp_server/test_mcp_server_manager.py | 148 +++++++++++++++ 5 files changed, 317 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 22c531b5c79..a21710892eb 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -722,7 +722,17 @@ async def update_mcp_server( data_dict["credentials"] = None if auth_type_changed or url_changed: - data_dict.update({field: None for field in _AUTH_FLOW_SCOPED_FIELDS if field not in data_dict}) + # Clear each auth-flow-scoped field that the caller either omitted (partial update) or + # resubmitted unchanged. The edit form re-sends every field, so a stale issuer/endpoint + # belonging to the old upstream would otherwise survive a url/auth_type change and win in the + # resolution merge; only a genuinely new submitted value is kept. + data_dict.update( + { + field: None + for field in _AUTH_FLOW_SCOPED_FIELDS + if field not in data_dict or data_dict[field] == getattr(existing, field, None) + } + ) # An explicit column write that does not touch credentials must still migrate # the row's legacy blob copies: lift values for columns the caller left diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 70180d5d469..a547f5b6920 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -201,6 +201,26 @@ def _blank_to_none(value: str | None) -> str | None: return value.strip() or None +def _endpoints_yield_to_issuer( + issuer: str | None, + is_discovery_auth_type: bool, + authorization_url: str | None, + token_url: str | None, + registration_url: str | None, +) -> tuple[str | None, str | None, str | None]: + """The single rule that makes an admin-configured ``issuer`` the sole authoritative endpoint + source (RFC 8414 §3.3): when it is set for a discovery auth type, the stored/manual + ``authorization_url``/``token_url``/``registration_url`` do not apply. They neither anchor nor + short-circuit discovery, never override the issuer document in the merge, and never substitute for + it when the issuer fetch fails (fail-closed). Returns the endpoint values that remain in force, + i.e. all ``None`` when issuer-anchored, else the inputs unchanged. Called at every resolution site + so the invariant holds in one place instead of being re-derived per merge. + """ + if issuer is not None and is_discovery_auth_type: + return None, None, None + return authorization_url, token_url, registration_url + + def _normalized_authorize_endpoint(url: str) -> str: """Compare authorize endpoints on scheme, host, and path only. The default port is elided and the host is lowercased so ``https://IDP.example.com:443/authorize/`` and @@ -271,11 +291,23 @@ def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serv incoming build has no pinned authorize endpoint (``None`` -> we adopt the previous one too, a consistent group) or pins the same one. An admin re-pointing ``authorization_url`` to a different server must not keep serving the old server's token endpoint or granted scopes. + + When an ``issuer`` is configured the endpoints must come solely from the §3.3-validated issuer + document, so carry-forward is skipped entirely for its endpoints: a failed issuer fetch leaves + them ``None`` and must stay ``None`` (fail-closed), never resurrected from the previous registry + entry. Scopes stay resource-driven and can still carry. """ if previous_server is None: return if previous_server.url != new_server.url or previous_server.auth_type != new_server.auth_type: return + if _blank_to_none(new_server.issuer): + # Endpoints come solely from the §3.3-validated issuer document; a failed fetch stays + # fail-closed and must not be resurrected from the previous entry. Only the resource-driven + # scopes carry as last-known-good. + if not new_server.scopes and previous_server.scopes: + new_server.scopes = previous_server.scopes + return may_carry = _endpoints_corroborate_authorization_url( previous_server.authorization_url, new_server.authorization_url ) @@ -1154,6 +1186,13 @@ class MCPServerManager: manual_registration_url = _blank_to_none(server_config.get("registration_url")) is_discovery_auth_type = auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES use_issuer_anchor = manual_issuer is not None and is_discovery_auth_type + manual_authorization_url, manual_token_url, manual_registration_url = _endpoints_yield_to_issuer( + manual_issuer, + is_discovery_auth_type, + manual_authorization_url, + manual_token_url, + manual_registration_url, + ) should_discover = bool(server_url) and ( is_discovery_auth_type or self._obo_needs_endpoint_discovery( @@ -1506,6 +1545,52 @@ class MCPServerManager: decrypt_global_env_var_values(env_vars_list) return env_vars_list + async def _resolve_table_oauth_metadata( + self, + *, + mcp_server: LiteLLM_MCPServerTable, + auth_type: MCPAuthType, + server_url: Optional[str], + manual_issuer: Optional[str], + manual_authorization_url: Optional[str], + manual_token_url: Optional[str], + is_discovery_auth_type: bool, + use_issuer_anchor: bool, + scopes: Optional[list[str]], + token_exchange_endpoint: Optional[str], + ) -> Optional[MCPOAuthMetadata]: + has_all_upstream_oauth_fields = bool(manual_authorization_url and manual_token_url and scopes) + needs_discovery = bool(server_url) and ( + (is_discovery_auth_type and not has_all_upstream_oauth_fields) + or self._obo_needs_endpoint_discovery(auth_type, token_exchange_endpoint, manual_token_url) + ) + if not needs_discovery: + mcp_oauth_metadata: Optional[MCPOAuthMetadata] = None + elif use_issuer_anchor and manual_issuer is not None: + mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer, server_url) + else: + mcp_oauth_metadata = await self._descovery_metadata( + server_url=server_url, # type: ignore[arg-type] + allow_origin_fallback=is_discovery_auth_type, + ) + if needs_discovery and not use_issuer_anchor and mcp_oauth_metadata is None: + verbose_logger.warning( + "MCP OAuth discovery yielded no metadata for server %s (%s); " + "OAuth endpoints/scopes stay unresolved until a rebuild succeeds", + mcp_server.server_id, + server_url, + ) + if use_issuer_anchor: + return mcp_oauth_metadata + if is_discovery_auth_type: + return _restrict_discovery_to_corroborated_authorization_server( + mcp_oauth_metadata, + manual_authorization_url, + mcp_server.server_id, + bool(getattr(mcp_server, "dcr_bridge", None)), + ) + return mcp_oauth_metadata + async def build_mcp_server_from_table( self, mcp_server: LiteLLM_MCPServerTable, @@ -1595,46 +1680,24 @@ class MCPServerManager: manual_registration_url = _blank_to_none(mcp_server.registration_url) is_discovery_auth_type = auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES use_issuer_anchor = manual_issuer is not None and is_discovery_auth_type - has_all_upstream_oauth_fields = bool(manual_authorization_url and manual_token_url and scopes) - needs_discovery = bool(server_url) and ( - (is_discovery_auth_type and not has_all_upstream_oauth_fields) - or self._obo_needs_endpoint_discovery( - auth_type, - mcp_server.token_exchange_endpoint - or (credentials_dict.get("token_exchange_endpoint") if credentials_dict else None), - manual_token_url, - ) + manual_authorization_url, manual_token_url, manual_registration_url = _endpoints_yield_to_issuer( + manual_issuer, is_discovery_auth_type, manual_authorization_url, manual_token_url, manual_registration_url + ) + token_exchange_endpoint = mcp_server.token_exchange_endpoint or ( + credentials_dict.get("token_exchange_endpoint") if credentials_dict else None + ) + gated_oauth_metadata = await self._resolve_table_oauth_metadata( + mcp_server=mcp_server, + auth_type=auth_type, + server_url=server_url, + manual_issuer=manual_issuer, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + is_discovery_auth_type=is_discovery_auth_type, + use_issuer_anchor=use_issuer_anchor, + scopes=scopes, + token_exchange_endpoint=token_exchange_endpoint, ) - if not needs_discovery: - mcp_oauth_metadata = None - elif manual_issuer is not None and is_discovery_auth_type: - mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata( - manual_issuer, - server_url, # type: ignore[arg-type] - ) - else: - mcp_oauth_metadata = await self._descovery_metadata( - server_url=server_url, # type: ignore[arg-type] - allow_origin_fallback=is_discovery_auth_type, - ) - if needs_discovery and not use_issuer_anchor and mcp_oauth_metadata is None: - verbose_logger.warning( - "MCP OAuth discovery yielded no metadata for server %s (%s); " - "OAuth endpoints/scopes stay unresolved until a rebuild succeeds", - mcp_server.server_id, - server_url, - ) - if use_issuer_anchor: - gated_oauth_metadata = mcp_oauth_metadata - elif is_discovery_auth_type: - gated_oauth_metadata = _restrict_discovery_to_corroborated_authorization_server( - mcp_oauth_metadata, - manual_authorization_url, - mcp_server.server_id, - bool(getattr(mcp_server, "dcr_bridge", None)), - ) - else: - gated_oauth_metadata = mcp_oauth_metadata resolved_scopes = scopes or (gated_oauth_metadata.scopes if gated_oauth_metadata else None) @@ -1721,6 +1784,7 @@ class MCPServerManager: existing_token_url=manual_token_url, existing_scopes=scopes, metadata=gated_oauth_metadata, + is_issuer_anchored=use_issuer_anchor, ) return new_server @@ -1769,6 +1833,7 @@ class MCPServerManager: existing_token_url: str | None, existing_scopes: list[str] | None, metadata: MCPOAuthMetadata | None, + is_issuer_anchored: bool = False, ) -> None: """Write freshly discovered OAuth endpoints back onto the DB row. @@ -1782,6 +1847,12 @@ class MCPServerManager: because ``_dcr_bridge_relays_client_registration`` keys off that column. Best-effort: a failed write re-discovers on the next build. Scopes go through ``update_mcp_server`` so they merge into the credentials blob without touching the stored client credentials. + + For an issuer-anchored server (``is_issuer_anchored``) the endpoints are re-derived from the + §3.3-validated issuer document on every build, so they are NOT persisted into the endpoint + columns: persisting them would make the next build see populated endpoints and treat them as + authoritative stored values, defeating the "endpoints come solely from the issuer" invariant. + Only the resource-driven scopes are persisted for such servers. """ if auth_type not in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: return @@ -1792,10 +1863,14 @@ class MCPServerManager: ) authorization_url_update = ( {"authorization_url": metadata.authorization_url} - if metadata.authorization_url and not existing_authorization_url + if metadata.authorization_url and not existing_authorization_url and not is_issuer_anchored + else {} + ) + token_url_update = ( + {"token_url": metadata.token_url} + if metadata.token_url and not existing_token_url and not is_issuer_anchored else {} ) - token_url_update = {"token_url": metadata.token_url} if metadata.token_url and not existing_token_url else {} scopes_update = {"credentials": {"scopes": metadata.scopes}} if metadata.scopes and not existing_scopes else {} updates: dict[str, object] = { **issuer_update, @@ -3375,7 +3450,9 @@ class MCPServerManager: return metadata return None - async def _fetch_issuer_anchored_oauth_metadata(self, issuer: str, server_url: str) -> Optional[MCPOAuthMetadata]: + async def _fetch_issuer_anchored_oauth_metadata( + self, issuer: str, server_url: Optional[str] + ) -> Optional[MCPOAuthMetadata]: """RFC 8414 issuer-anchored discovery for the OAuth endpoints, with resource-driven scopes. Fetch authorization-server metadata from the admin-configured issuer's own origin and adopt @@ -3400,7 +3477,9 @@ class MCPServerManager: issuer, ) return None - resource_metadata = await self._descovery_metadata(server_url, allow_origin_fallback=False) + resource_metadata = ( + await self._descovery_metadata(server_url, allow_origin_fallback=False) if server_url else None + ) resource_scopes = resource_metadata.scopes if resource_metadata else None return metadata.model_copy(update={"scopes": resource_scopes}) @@ -3488,8 +3567,6 @@ class MCPServerManager: ): return metadata - if require_issuer is not None: - return None return self._build_azure_authorization_server_metadata(parsed) @staticmethod @@ -5198,6 +5275,7 @@ class MCPServerManager: command=getattr(server, "command", None), args=getattr(server, "args", None) or [], env=getattr(server, "env", None) or {}, + issuer=server.issuer, authorization_url=server.authorization_url, token_url=server.token_url, registration_url=server.registration_url, @@ -5307,6 +5385,7 @@ class MCPServerManager: command=getattr(server, "command", None), args=getattr(server, "args", None) or [], env=getattr(server, "env", None) or {}, + issuer=server.issuer, authorization_url=server.authorization_url, token_url=server.token_url, registration_url=server.registration_url, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 288282dd08b..d920ee474cc 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -536,6 +536,7 @@ if MCP_AVAILABLE: sanitized.env = {} sanitized.command = None sanitized.args = [] + sanitized.issuer = None sanitized.authorization_url = None sanitized.token_url = None sanitized.registration_url = None @@ -581,6 +582,7 @@ if MCP_AVAILABLE: sanitized.teams = [] sanitized.env_vars = None + sanitized.issuer = None sanitized.authorization_url = None sanitized.token_url = None sanitized.registration_url = None @@ -686,6 +688,7 @@ if MCP_AVAILABLE: command=payload.command, args=payload.args, env=payload.env, + issuer=payload.issuer, authorization_url=payload.authorization_url, token_url=payload.token_url, registration_url=payload.registration_url, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index 637ee409fbe..0e53ebd76e1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -239,6 +239,37 @@ async def test_url_change_clears_stale_discovered_oauth_fields(): assert data_dict[stale_field] is None, f"{stale_field} must be cleared on url change" +@pytest.mark.asyncio +async def test_url_change_clears_stale_oauth_fields_even_when_resubmitted_unchanged(): + """The edit form re-sends every field, so a URL change arrives WITH the previous upstream's issuer + and endpoints in the payload. Those resubmitted-unchanged values are stale and must still clear + (otherwise they survive the url change and win in the resolution merge). A genuinely new value the + caller changed in the same submit is kept.""" + mock_prisma = _mock_prisma() + existing = MagicMock() + existing.auth_type = "oauth2" + existing.url = "https://old.example.com/mcp" + existing.credentials = None + existing.issuer = "https://old-idp.example.com" + existing.token_url = "https://old-idp.example.com/token" + existing.authorization_url = "https://old-idp.example.com/authorize" + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest( + server_id="my-test-server", + url="https://new.example.com/mcp", + issuer="https://old-idp.example.com", # resubmitted unchanged -> stale, must clear + token_url="https://old-idp.example.com/token", # resubmitted unchanged -> stale, must clear + authorization_url="https://new-idp.example.com/authorize", # genuinely changed -> kept + ) + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + assert data_dict["issuer"] is None + assert data_dict["token_url"] is None + assert data_dict["authorization_url"] == "https://new-idp.example.com/authorize" + + @pytest.mark.asyncio async def test_unchanged_url_does_not_clear_discovered_oauth_fields(): """A partial update that resends the same url (or omits it) must not clear the discovered OAuth 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 75f77c1ec6d..1f966627aea 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 @@ -1338,6 +1338,50 @@ class TestMCPServerManager: assert built.registration_url is None assert built.scopes is None + @pytest.mark.asyncio + async def test_build_from_table_issuer_anchor_overrides_stored_endpoints_even_when_populated(self): + """When an issuer is pinned, the endpoints come SOLELY from the §3.3-validated issuer document + and win over any stored/manual endpoint values, even a fully-populated row. Otherwise an + attacker who controls a stored token endpoint keeps receiving codes/secrets after an admin + pins a trusted issuer: `needs_discovery` must not short-circuit on populated fields, and the + issuer's endpoints must override the stored ones.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="issuer-anchored-populated", + alias="issuer_anchored_populated", + description="issuer set, but stale/hostile endpoints already stored", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + issuer="https://idp.example.com", + authorization_url="https://attacker.example.com/authorize", + token_url="https://attacker.example.com/steal", + credentials={"scopes": ["stale"]}, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + issuer_resolved = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["read"], + ) + with ( + patch.object( + manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=issuer_resolved) + ) as anchored, + patch.object(manager, "_persist_discovered_oauth_endpoints", new=AsyncMock()) as mock_persist, + ): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + anchored.assert_awaited_once_with("https://idp.example.com", "https://up.example.com/mcp") + assert built.authorization_url == "https://idp.example.com/authorize" + assert built.token_url == "https://idp.example.com/token" + assert built.token_url != "https://attacker.example.com/steal" + # The issuer-anchored endpoints are never persisted into the endpoint columns, so a later + # build cannot treat them as authoritative stored values. + assert mock_persist.await_args.kwargs["is_issuer_anchored"] is True + @pytest.mark.asyncio @pytest.mark.parametrize( "advertised_authorization_url", @@ -2669,6 +2713,37 @@ class TestMCPServerManager: assert result.authorization_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" assert result.token_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" + @pytest.mark.asyncio + async def test_azure_heuristic_reachable_under_require_issuer(self): + """Under issuer-anchored discovery (require_issuer set), an Entra issuer whose OIDC document + cannot be fetched still gets the deterministic Azure endpoint construction. The heuristic + derives the endpoints from the pinned issuer's own tenant URL, so it is authoritative-by- + construction and safe under require_issuer; only a non-Entra issuer stays fail-closed (None).""" + manager = MCPServerManager() + issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0" + + request = httpx.Request("GET", issuer) + response_obj = httpx.Response(status_code=404, request=request) + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError("not found", request=request, response=response_obj) + ) + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + azure = await manager._fetch_single_authorization_server_metadata(issuer, issuer, require_issuer=issuer) + non_entra = await manager._fetch_single_authorization_server_metadata( + "https://idp.example.com", "https://idp.example.com", require_issuer="https://idp.example.com" + ) + + assert azure is not None + assert azure.token_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" + assert non_entra is None + @pytest.mark.asyncio async def test_descovery_metadata_falls_back_to_origin_when_no_auth_servers(self): manager = MCPServerManager() @@ -5484,6 +5559,41 @@ class TestMCPServerTimestamps: persisted = update_mcp_server_mock.call_args.kwargs["data"] assert persisted.issuer == "https://idp.example.com" + @pytest.mark.asyncio + async def test_persist_discovered_oauth_endpoints_does_not_persist_endpoints_for_issuer_anchored(self): + """For an issuer-anchored server the endpoints are re-derived from the §3.3-validated issuer + document every build, so they must NOT be written into the endpoint columns: persisting them + would make the next build see populated endpoints and treat them as authoritative stored + values, defeating the issuer-only invariant. Only the resource-driven scopes are persisted.""" + manager = MCPServerManager() + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["read"], + ) + + update_mcp_server_mock = AsyncMock() + with ( + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=update_mcp_server_mock), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + ): + await manager._persist_discovered_oauth_endpoints( + server_id="s", + auth_type=MCPAuth.oauth2, + existing_issuer="https://idp.example.com", + existing_authorization_url=None, + existing_token_url=None, + existing_scopes=None, + metadata=metadata, + is_issuer_anchored=True, + ) + + update_mcp_server_mock.assert_awaited_once() + persisted = update_mcp_server_mock.call_args.kwargs["data"] + assert "authorization_url" not in persisted.fields_set() + assert "token_url" not in persisted.fields_set() + assert persisted.credentials == {"scopes": ["read"]} + @pytest.mark.asyncio async def test_build_mcp_server_from_table_skips_persistence_for_temporary_servers(self): """The session endpoint builds temporary servers whose server_id has no DB row; with @@ -5698,6 +5808,44 @@ class TestMCPServerTimestamps: assert same_authorize.token_url == "https://idp.example.com/token" assert same_authorize.registration_url == "https://idp.example.com/register" + def test_carry_forward_does_not_restore_endpoints_for_issuer_anchored_server(self): + """When an issuer is configured the endpoints come solely from the §3.3-validated issuer + document, so a failed issuer fetch (token_url None) must stay fail-closed. Carry-forward must + NOT resurrect the previous registry entry's token endpoint, or the very attacker-controlled + endpoint the issuer anchor distrusts would keep being served across rebuilds. Resource-driven + scopes still carry as last-known-good.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _carry_forward_resolved_oauth_endpoints, + ) + + previous = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + issuer="https://idp.example.com", + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + scopes=["read"], + ) + failed_rebuild = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + issuer="https://idp.example.com", + ) + + _carry_forward_resolved_oauth_endpoints(new_server=failed_rebuild, previous_server=previous) + + assert failed_rebuild.authorization_url is None + assert failed_rebuild.token_url is None + assert failed_rebuild.registration_url is None + assert failed_rebuild.scopes == ["read"] + def test_normalized_authorize_endpoint_treats_default_port_and_slash_as_identity(self): """The corroboration check must not fail on formatting-only differences an IdP legitimately emits: default port, trailing slash, host case, and query string are not identity, but a From 585f21aaf77b59b2827c8235cc19d480c9f3d644 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 15 Jul 2026 17:45:40 -0700 Subject: [PATCH 31/62] chore(ui): remove Hide Usage Indicator flag and hook --- ui/litellm-dashboard/eslint-suppressions.json | 11 - .../hooks/useDisableUsageIndicator.test.ts | 190 ------------------ .../hooks/useDisableUsageIndicator.ts | 33 --- .../Navbar/UserDropdown/UserDropdown.tsx | 19 -- .../SidebarAccountMenu.test.tsx | 4 - .../SidebarAccountMenu/SidebarAccountMenu.tsx | 9 - .../src/components/SidebarUsageCard.test.tsx | 4 - .../src/components/SidebarUsageCard.tsx | 4 +- 8 files changed, 1 insertion(+), 273 deletions(-) delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.test.ts delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.ts diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 05ad536b260..9b55bf6ca0d 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1655,17 +1655,6 @@ "count": 1 } }, - "src/components/UsageIndicator.tsx": { - "no-nested-ternary": { - "count": 4 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/static-components": { - "count": 1 - } - }, "src/components/activity_metrics.tsx": { "no-nested-ternary": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.test.ts deleted file mode 100644 index bd0e69c0de3..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.test.ts +++ /dev/null @@ -1,190 +0,0 @@ -import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; -import { act, renderHook, waitFor } from "@testing-library/react"; -import { useDisableUsageIndicator } from "./useDisableUsageIndicator"; -import { LOCAL_STORAGE_EVENT } from "@/utils/localStorageUtils"; - -describe("useDisableUsageIndicator", () => { - const STORAGE_KEY = "disableUsageIndicator"; - - beforeEach(() => { - localStorage.clear(); - vi.clearAllMocks(); - }); - - afterEach(() => { - localStorage.clear(); - }); - - it("should return false when localStorage is empty", () => { - const { result } = renderHook(() => useDisableUsageIndicator()); - - expect(result.current).toBe(false); - }); - - it("should return false when localStorage value is not 'true'", () => { - localStorage.setItem(STORAGE_KEY, "false"); - - const { result } = renderHook(() => useDisableUsageIndicator()); - - expect(result.current).toBe(false); - }); - - it("should return true when localStorage value is 'true'", () => { - localStorage.setItem(STORAGE_KEY, "true"); - - const { result } = renderHook(() => useDisableUsageIndicator()); - - expect(result.current).toBe(true); - }); - - it("should return false when localStorage value is an empty string", () => { - localStorage.setItem(STORAGE_KEY, ""); - - const { result } = renderHook(() => useDisableUsageIndicator()); - - expect(result.current).toBe(false); - }); - - it("should update when storage event fires for the correct key", async () => { - const { result } = renderHook(() => useDisableUsageIndicator()); - - expect(result.current).toBe(false); - - await act(async () => { - localStorage.setItem(STORAGE_KEY, "true"); - const storageEvent = new StorageEvent("storage", { - key: STORAGE_KEY, - newValue: "true", - }); - window.dispatchEvent(storageEvent); - }); - - await waitFor(() => { - expect(result.current).toBe(true); - }); - }); - - it("should not update when storage event fires for a different key", () => { - localStorage.setItem(STORAGE_KEY, "false"); - const { result } = renderHook(() => useDisableUsageIndicator()); - - expect(result.current).toBe(false); - - const storageEvent = new StorageEvent("storage", { - key: "otherKey", - newValue: "true", - }); - window.dispatchEvent(storageEvent); - - expect(result.current).toBe(false); - }); - - it("should update when custom LOCAL_STORAGE_EVENT fires for the correct key", async () => { - const { result } = renderHook(() => useDisableUsageIndicator()); - - expect(result.current).toBe(false); - - await act(async () => { - localStorage.setItem(STORAGE_KEY, "true"); - const customEvent = new CustomEvent(LOCAL_STORAGE_EVENT, { - detail: { key: STORAGE_KEY }, - }); - window.dispatchEvent(customEvent); - }); - - await waitFor(() => { - expect(result.current).toBe(true); - }); - }); - - it("should not update when custom LOCAL_STORAGE_EVENT fires for a different key", () => { - localStorage.setItem(STORAGE_KEY, "false"); - const { result } = renderHook(() => useDisableUsageIndicator()); - - expect(result.current).toBe(false); - - const customEvent = new CustomEvent(LOCAL_STORAGE_EVENT, { - detail: { key: "otherKey" }, - }); - window.dispatchEvent(customEvent); - - expect(result.current).toBe(false); - }); - - it("should update when localStorage changes from false to true via custom event", async () => { - localStorage.setItem(STORAGE_KEY, "false"); - const { result } = renderHook(() => useDisableUsageIndicator()); - - expect(result.current).toBe(false); - - await act(async () => { - localStorage.setItem(STORAGE_KEY, "true"); - const customEvent = new CustomEvent(LOCAL_STORAGE_EVENT, { - detail: { key: STORAGE_KEY }, - }); - window.dispatchEvent(customEvent); - }); - - await waitFor(() => { - expect(result.current).toBe(true); - }); - }); - - it("should update when localStorage changes from true to false via storage event", async () => { - localStorage.setItem(STORAGE_KEY, "true"); - const { result } = renderHook(() => useDisableUsageIndicator()); - - expect(result.current).toBe(true); - - await act(async () => { - localStorage.setItem(STORAGE_KEY, "false"); - const storageEvent = new StorageEvent("storage", { - key: STORAGE_KEY, - newValue: "false", - }); - window.dispatchEvent(storageEvent); - }); - - await waitFor(() => { - expect(result.current).toBe(false); - }); - }); - - it("should cleanup event listeners on unmount", () => { - const addEventListenerSpy = vi.spyOn(window, "addEventListener"); - const removeEventListenerSpy = vi.spyOn(window, "removeEventListener"); - - const { unmount } = renderHook(() => useDisableUsageIndicator()); - - expect(addEventListenerSpy).toHaveBeenCalledTimes(2); - expect(addEventListenerSpy).toHaveBeenCalledWith("storage", expect.any(Function)); - expect(addEventListenerSpy).toHaveBeenCalledWith(LOCAL_STORAGE_EVENT, expect.any(Function)); - - unmount(); - - expect(removeEventListenerSpy).toHaveBeenCalledTimes(2); - expect(removeEventListenerSpy).toHaveBeenCalledWith("storage", expect.any(Function)); - expect(removeEventListenerSpy).toHaveBeenCalledWith(LOCAL_STORAGE_EVENT, expect.any(Function)); - }); - - it("should handle multiple hooks independently", async () => { - const { result: result1 } = renderHook(() => useDisableUsageIndicator()); - const { result: result2 } = renderHook(() => useDisableUsageIndicator()); - - expect(result1.current).toBe(false); - expect(result2.current).toBe(false); - - await act(async () => { - localStorage.setItem(STORAGE_KEY, "true"); - const customEvent = new CustomEvent(LOCAL_STORAGE_EVENT, { - detail: { key: STORAGE_KEY }, - }); - window.dispatchEvent(customEvent); - }); - - await waitFor(() => { - expect(result1.current).toBe(true); - expect(result2.current).toBe(true); - }); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.ts deleted file mode 100644 index 7f4e2295090..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.ts +++ /dev/null @@ -1,33 +0,0 @@ -import { getLocalStorageItem, LOCAL_STORAGE_EVENT } from "@/utils/localStorageUtils"; -import { useSyncExternalStore } from "react"; - -function subscribe(callback: () => void) { - const onStorage = (e: StorageEvent) => { - if (e.key === "disableUsageIndicator") { - callback(); - } - }; - - const onCustom = (e: Event) => { - const { key } = (e as CustomEvent).detail; - if (key === "disableUsageIndicator") { - callback(); - } - }; - - window.addEventListener("storage", onStorage); - window.addEventListener(LOCAL_STORAGE_EVENT, onCustom); - - return () => { - window.removeEventListener("storage", onStorage); - window.removeEventListener(LOCAL_STORAGE_EVENT, onCustom); - }; -} - -function getSnapshot() { - return getLocalStorageItem("disableUsageIndicator") === "true"; -} - -export function useDisableUsageIndicator() { - return useSyncExternalStore(subscribe, getSnapshot); -} diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx index 983ab980d2f..a71fc1b97a8 100644 --- a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx @@ -2,7 +2,6 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useDisableBlogPosts } from "@/app/(dashboard)/hooks/useDisableBlogPosts"; import { useDisableBouncingIcon } from "@/app/(dashboard)/hooks/useDisableBouncingIcon"; import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPrompts"; -import { useDisableUsageIndicator } from "@/app/(dashboard)/hooks/useDisableUsageIndicator"; import { emitLocalStorageChange, getLocalStorageItem, @@ -72,7 +71,6 @@ interface UserDropdownProps { const UserDropdown: React.FC = ({ onLogout, variant = "navbar", collapsed = false }) => { const { userId, userEmail, userRole, premiumUser } = useAuthorized(); const disableShowPrompts = useDisableShowPrompts(); - const disableUsageIndicator = useDisableUsageIndicator(); const disableBlogPosts = useDisableBlogPosts(); const disableBouncingIcon = useDisableBouncingIcon(); const [disableShowNewBadge, setDisableShowNewBadge] = useState(false); @@ -165,23 +163,6 @@ const UserDropdown: React.FC = ({ onLogout, variant = "navbar aria-label="Toggle hide all prompts" /> - - Hide Usage Indicator - { - if (checked) { - setLocalStorageItem("disableUsageIndicator", "true"); - emitLocalStorageChange("disableUsageIndicator"); - } else { - removeLocalStorageItem("disableUsageIndicator"); - emitLocalStorageChange("disableUsageIndicator"); - } - }} - aria-label="Toggle hide usage indicator" - /> - Hide Blog Posts ({ useDisableShowPrompts: () => mockUseDisableShowPromptsImpl(), })); -vi.mock("@/app/(dashboard)/hooks/useDisableUsageIndicator", () => ({ - useDisableUsageIndicator: () => false, -})); - vi.mock("@/app/(dashboard)/hooks/useDisableBlogPosts", () => ({ useDisableBlogPosts: () => false, })); diff --git a/ui/litellm-dashboard/src/components/SidebarAccountMenu/SidebarAccountMenu.tsx b/ui/litellm-dashboard/src/components/SidebarAccountMenu/SidebarAccountMenu.tsx index 4842f25f36b..b7a16bcf09a 100644 --- a/ui/litellm-dashboard/src/components/SidebarAccountMenu/SidebarAccountMenu.tsx +++ b/ui/litellm-dashboard/src/components/SidebarAccountMenu/SidebarAccountMenu.tsx @@ -4,7 +4,6 @@ import { useDisableBlogPosts } from "@/app/(dashboard)/hooks/useDisableBlogPosts import { useDisableBouncingIcon } from "@/app/(dashboard)/hooks/useDisableBouncingIcon"; import { useDisableShowNewBadge } from "@/app/(dashboard)/hooks/useDisableShowNewBadge"; import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPrompts"; -import { useDisableUsageIndicator } from "@/app/(dashboard)/hooks/useDisableUsageIndicator"; import { emitLocalStorageChange, removeLocalStorageItem, setLocalStorageItem } from "@/utils/localStorageUtils"; import { navAccountDisplayName } from "@/components/Navbar/navDisplayName"; import CopyButton from "@/components/shared/CopyButton"; @@ -86,7 +85,6 @@ const SidebarAccountMenu: React.FC = ({ onLogout, colla const { data: healthData } = useHealthReadinessDetails(accessToken); const version = healthData?.litellm_version; const disableShowPrompts = useDisableShowPrompts(); - const disableUsageIndicator = useDisableUsageIndicator(); const disableBlogPosts = useDisableBlogPosts(); const disableBouncingIcon = useDisableBouncingIcon(); const disableShowNewBadge = useDisableShowNewBadge(); @@ -115,13 +113,6 @@ const SidebarAccountMenu: React.FC = ({ onLogout, colla checked: disableShowPrompts, onCheckedChange: (checked: boolean) => setFlag("disableShowPrompts", checked), }, - { - key: "disableUsageIndicator", - label: "Hide Usage Indicator", - ariaLabel: "Toggle hide usage indicator", - checked: disableUsageIndicator, - onCheckedChange: (checked: boolean) => setFlag("disableUsageIndicator", checked), - }, { key: "disableBlogPosts", label: "Hide Blog Posts", diff --git a/ui/litellm-dashboard/src/components/SidebarUsageCard.test.tsx b/ui/litellm-dashboard/src/components/SidebarUsageCard.test.tsx index 3244906ed8e..90d22accbc3 100644 --- a/ui/litellm-dashboard/src/components/SidebarUsageCard.test.tsx +++ b/ui/litellm-dashboard/src/components/SidebarUsageCard.test.tsx @@ -8,10 +8,6 @@ import type { LicenseInfo } from "./networking"; vi.mock("./networking", () => ({ getRemainingUsers: vi.fn() })); -vi.mock("@/app/(dashboard)/hooks/useDisableUsageIndicator", () => ({ - useDisableUsageIndicator: vi.fn(() => false), -})); - vi.mock("@/app/(dashboard)/hooks/license/useLicenseInfo", () => ({ useLicenseInfo: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/SidebarUsageCard.tsx b/ui/litellm-dashboard/src/components/SidebarUsageCard.tsx index 4fb0e1dac21..f3afc1987ee 100644 --- a/ui/litellm-dashboard/src/components/SidebarUsageCard.tsx +++ b/ui/litellm-dashboard/src/components/SidebarUsageCard.tsx @@ -1,4 +1,3 @@ -import { useDisableUsageIndicator } from "@/app/(dashboard)/hooks/useDisableUsageIndicator"; import { useLicenseInfo } from "@/app/(dashboard)/hooks/license/useLicenseInfo"; import { getDaysUntilExpiration } from "@/utils/licenseUtils"; import { Button } from "@/components/ui/button"; @@ -79,7 +78,6 @@ const buildMeters = (data: RemainingUsage | null): MeterData[] => { * design's Spend / API-request meters are intentionally omitted. */ export default function SidebarUsageCard({ accessToken, collapsed, onExpandRail }: SidebarUsageCardProps) { - const disableUsageIndicator = useDisableUsageIndicator(); const licenseInfo = useLicenseInfo(accessToken).data ?? null; const { data: usageData, isLoading } = useQuery(remainingUsersQuery(accessToken)); const data = usageData ?? null; @@ -87,7 +85,7 @@ export default function SidebarUsageCard({ accessToken, collapsed, onExpandRail const hasData = data !== null && (data.total_users !== null || data.total_teams !== null); const noUsableData = !isLoading && !hasData; const noLicensedUsage = !licenseInfo?.has_license || noUsableData; - if (disableUsageIndicator || !accessToken || noLicensedUsage) { + if (!accessToken || noLicensedUsage) { return null; } From fac43df9b95376cd7159e105415d3b0316bbbf07 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:46:00 -0700 Subject: [PATCH 32/62] fix(complexity_router): return empty dict from _classifier_call_metadata when metadata is absent (#33452) * fix(complexity_router): return empty dict from _classifier_call_metadata when metadata is absent The LLM classifier reads request_kwargs.get("litellm_metadata"), but the proxy stores request metadata under "metadata", so this returned None. _classifier_call_metadata then passed None straight through to the classifier acompletion call, which assumes a dict and blows up with 'NoneType' object has no attribute 'update'; the router swallowed it and silently fell back to heuristic scoring, so the configured LLM classifier never ran. Returning an empty dict keeps the classifier call well-formed. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): cover complexity-router LLM classifier routes over the proxy Add a live e2e regression for the complexity auto-router: a lexically simple but hard prompt ("Is P equal to NP?") is routed by the LLM classifier to the higher-tier anthropic backend, read back from the spend log's model. Before the metadata fix the classifier silently crashed and the router fell back to heuristic SIMPLE scoring on the openai backend, so this test fails pre-fix and passes post-fix. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../complexity_router/complexity_router.py | 8 +-- tests/e2e/coverage_registry/reliability.yaml | 1 + tests/e2e/docker-compose.yml | 17 +++++ tests/e2e/router/complexity_router_client.py | 20 ++++++ tests/e2e/router/conftest.py | 15 +++++ .../e2e/router/test_complexity_router_e2e.py | 62 +++++++++++++++++++ .../router_strategy/test_complexity_router.py | 10 +++ 7 files changed, 129 insertions(+), 4 deletions(-) create mode 100644 tests/e2e/router/complexity_router_client.py create mode 100644 tests/e2e/router/conftest.py create mode 100644 tests/e2e/router/test_complexity_router_e2e.py diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index eb0f74a58e7..e85987870e1 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -98,9 +98,9 @@ def _sanitize_user_api_key_auth(auth: Any) -> Any: return auth -def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any] | None: +def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]: if not metadata: - return metadata + return {} return { k: _sanitize_user_api_key_auth(v) if k == "user_api_key_auth" else v for k, v in metadata.items() @@ -763,8 +763,8 @@ class ComplexityRouter(CustomLogger): # embedding call. Forwarding it would let the embedding's cost callback finalize the # reservation, so the routed completion's own callback then skips incrementing the # key/team budget. Key/team attribution fields are preserved for spend logging. - metadata = _classifier_call_metadata(request_kwargs.get("metadata")) or {} - litellm_metadata = _classifier_call_metadata(request_kwargs.get("litellm_metadata")) or {} + metadata = _classifier_call_metadata(request_kwargs.get("metadata")) + litellm_metadata = _classifier_call_metadata(request_kwargs.get("litellm_metadata")) query_vector = ( await encoder.aencode_queries([user_message], metadata=metadata, litellm_metadata=litellm_metadata) )[0] diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 3746b029331..ba192c2912e 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -17,6 +17,7 @@ - {id: reliability.routing.cost_based.picks_lowest_cost, module: reliability, tier: P1, behavior: routing, variant: cost_based, assertions: [picks_lowest_cost], exercised_on: [chat_completions, messages], source: "router_strategy/lowest_cost.py", rationale: "Spend-aware routing"} - {id: reliability.routing.usage_based.picks_under_tpm, module: reliability, tier: P0, behavior: routing, variant: usage_based, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "router_strategy/lowest_tpm_rpm_v2.py", rationale: "Routes to lowest-TPM deployment; prevents over-allocation"} - {id: reliability.routing.least_busy.picks_lowest_traffic, module: reliability, tier: P1, behavior: routing, variant: least_busy, assertions: [picks_lowest_traffic], exercised_on: [chat_completions, messages], source: "router_strategy/least_busy.py", rationale: "Fewest in-flight requests"} +- {id: reliability.routing.complexity_llm_classifier.routes_by_llm_tier, module: reliability, tier: P1, behavior: routing, variant: complexity_llm_classifier, assertions: [routes_by_llm_tier], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py", fail_before_fix: proven, rationale: "v2 auto-router LLM complexity classifier runs over the proxy and routes by semantic tier instead of silently crashing on absent litellm_metadata and falling back to heuristic scoring"} - {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"} diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index e2bb6ca8933..b64f3d8dbfd 100644 --- a/tests/e2e/docker-compose.yml +++ b/tests/e2e/docker-compose.yml @@ -66,6 +66,23 @@ configs: model: openai/text-embedding-3-small api_key: os.environ/OPENAI_API_KEY + # v2 auto-router with the LLM complexity classifier. SIMPLE stays on the + # openai backend; every higher tier routes to the anthropic backend, so the + # served deployment (read back from the spend log's model) reveals whether + # the LLM classifier actually ran or silently fell back to heuristic scoring. + - model_name: complexity-smart-router + litellm_params: + model: auto_router/complexity_router + complexity_router_config: + classifier_type: llm + classifier_llm_config: + model: gpt-5.5 + tiers: + SIMPLE: gpt-5.5 + MEDIUM: claude-haiku-4-5 + COMPLEX: claude-haiku-4-5 + REASONING: claude-haiku-4-5 + services: litellm: image: ghcr.io/berriai/litellm:main-latest diff --git a/tests/e2e/router/complexity_router_client.py b/tests/e2e/router/complexity_router_client.py new file mode 100644 index 00000000000..929acbb3461 --- /dev/null +++ b/tests/e2e/router/complexity_router_client.py @@ -0,0 +1,20 @@ +"""Client for the complexity auto-router e2e tests. + +The suite drives the shared /chat/completions and spend-log reads on the Gateway, +so this client only carries the Gateway the shared lifecycle needs for cleanup. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from e2e_gateway import Gateway, build_gateway + + +@dataclass(frozen=True, slots=True) +class ComplexityRouterClient: + gateway: Gateway + + +def build_client() -> ComplexityRouterClient: + return ComplexityRouterClient(gateway=build_gateway()) diff --git a/tests/e2e/router/conftest.py b/tests/e2e/router/conftest.py new file mode 100644 index 00000000000..e8c05520b10 --- /dev/null +++ b/tests/e2e/router/conftest.py @@ -0,0 +1,15 @@ +"""Router suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +live in the parent tests/e2e/conftest.py. ComplexityRouterClient holds the shared +Gateway, so the `resources` fixture cleans up keys this suite creates. +""" + +import pytest + +from complexity_router_client import ComplexityRouterClient, build_client + + +@pytest.fixture(scope="session") +def client() -> ComplexityRouterClient: + return build_client() diff --git a/tests/e2e/router/test_complexity_router_e2e.py b/tests/e2e/router/test_complexity_router_e2e.py new file mode 100644 index 00000000000..88d79a9cac0 --- /dev/null +++ b/tests/e2e/router/test_complexity_router_e2e.py @@ -0,0 +1,62 @@ +"""Live e2e: the v2 auto-router's LLM complexity classifier actually runs over the +proxy and drives routing, instead of silently crashing and falling back to the +local heuristic scorer. + +The regression this guards (complexity_router.py `_classifier_call_metadata` +returning None when the request carries no `litellm_metadata`, which the classifier +sub-call then fed into a `.update`, raising `'NoneType' object has no attribute +'update'`) was invisible from the outside: the router caught the error and answered +from heuristic scoring, so every request still returned 200. The only tell is which +tier, and therefore which backend, served the request. + +`complexity-smart-router` (see the inline config in docker-compose.yml) pins SIMPLE +to the openai backend and every higher tier to the anthropic backend. "Is P equal +to NP?" is lexically trivial, so the heuristic scorer lands it in SIMPLE (openai), +but any competent LLM classifier reads it as a hard reasoning question and lands it +above SIMPLE (anthropic). The served deployment is read back from the spend log's +`model`, so anthropic proves the classifier ran and openai proves it silently fell +back - the exact failure before the fix. +""" + +import pytest + +from complexity_router_client import ComplexityRouterClient +from e2e_http import unwrap +from models import ChatBody, ChatMessage + +pytestmark = pytest.mark.e2e + +ROUTER_MODEL = "complexity-smart-router" +# Lexically simple (heuristic -> SIMPLE) but a hard reasoning question (LLM -> above SIMPLE). +LEXICALLY_SIMPLE_HARD_PROMPT = "Is P equal to NP?" +# SIMPLE tier backend; served only when the classifier silently falls back to heuristic. +HEURISTIC_TIER_MODEL = "openai/gpt-5.5" +# MEDIUM/COMPLEX/REASONING tier backend; served only when the LLM classifier runs. +LLM_TIER_MODEL = "anthropic/claude-haiku-4-5" + + +class TestComplexityRouterLlmClassifier: + @pytest.mark.covers("reliability.routing.complexity_llm_classifier.routes_by_llm_tier") + def test_llm_classifier_runs_and_routes_by_semantic_tier( + self, client: ComplexityRouterClient, scoped_key: str + ) -> None: + chat = unwrap( + client.gateway.chat( + scoped_key, + ChatBody( + model=ROUTER_MODEL, + messages=[ChatMessage(role="user", content=LEXICALLY_SIMPLE_HARD_PROMPT)], + max_tokens=16, + ), + ) + ) + assert chat.choices, f"router returned no choices: {chat}" + + rows = client.gateway.poll_logs_for_key(scoped_key, min_rows=1) + served = [row.model for row in rows] + assert served == [LLM_TIER_MODEL], ( + f"expected the request to be served by {LLM_TIER_MODEL!r} (the higher-tier " + f"backend the LLM classifier picks for a hard prompt), but the spend log shows " + f"{served!r}. {HEURISTIC_TIER_MODEL!r} means the LLM classifier silently failed " + f"and the router fell back to heuristic scoring (SIMPLE) - the pre-fix regression" + ) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 41dc7269372..3404b55f0db 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2387,6 +2387,16 @@ class TestSubCallMetadataSanitization: assert sanitized["user_api_key_auth"] is not None assert _get_budget_reservation_from_metadata(sanitized) is None + def test_returns_empty_dict_for_missing_metadata(self): + from litellm.router_strategy.complexity_router.complexity_router import ( + _classifier_call_metadata, + ) + + for absent in (None, {}): + result = _classifier_call_metadata(absent) + assert result == {} + assert isinstance(result, dict) + def test_sanitized_auth_keeps_access_group_fields_and_leaves_original_untouched(self): from litellm.proxy._types import UserAPIKeyAuth from litellm.router_strategy.complexity_router.complexity_router import ( From 17690dece2277ecca0b04ef0823f9a336c6f8b82 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 15 Jul 2026 18:13:22 -0700 Subject: [PATCH 33/62] test(ocr): use mistral-document-ai-2512 in azure_ai OCR tests --- tests/ocr_tests/test_ocr_azure_ai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/ocr_tests/test_ocr_azure_ai.py b/tests/ocr_tests/test_ocr_azure_ai.py index 172682175c5..acb44958fd9 100644 --- a/tests/ocr_tests/test_ocr_azure_ai.py +++ b/tests/ocr_tests/test_ocr_azure_ai.py @@ -23,7 +23,7 @@ class TestAzureAIOCR(BaseOCRTest): Return the base OCR call args for Azure AI. """ return { - "model": "azure_ai/mistral-document-ai-2505", + "model": "azure_ai/mistral-document-ai-2512", "api_key": os.getenv("AZURE_API_KEY"), "api_base": os.getenv("AZURE_API_BASE"), } From ad73f3a7a282560cf6c654512b2f12366f24b666 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 18:15:12 -0700 Subject: [PATCH 34/62] fix(mcp): keep issuer provenance consistent when it changes or is discovered Two lifecycle gaps let the issuer trust anchor drift out of sync with the endpoints it governs. Changing or clearing a previously pinned issuer left the authorization_url and token_url that were resolved under the old issuer in the row, so clearing the anchor could revive stale, possibly untrusted endpoints instead of re-discovering. And a build that discovered an issuer trust-on-first-use persisted it to the row while the returned in-memory server kept the issuer unset, so the registry and the row disagreed and the per-user OAuth token identity, which includes the issuer, differed between that build and the next rebuild and forced a spurious re-auth. update_mcp_server now treats a change to a previously pinned issuer the same as a url or auth_type change and clears the auth-flow-scoped endpoint fields that were resolved under it. The trigger fires only when an issuer was already pinned and is now changed or cleared, so establishing one for the first time, including the trust-on-first-use discovery write-back, does not wipe the fields it just resolved. Both build paths, build_mcp_server_from_table and load_servers_from_config, now construct the server with effective_issuer = manual_issuer or the discovered issuer, skipping an origin-fallback guess exactly as the persistence does, so the in-memory object always reflects what the row will hold. Regression tests pin each case: clearing and re-pointing a pinned issuer clear the stale endpoints, a first-time establish preserves the discovered fields, and a build reflects the discovered issuer while an origin-fallback guess is not reflected --- litellm/proxy/_experimental/mcp_server/db.py | 16 +++- .../mcp_server/mcp_server_manager.py | 16 +++- .../mcp_server/test_mcp_partial_update.py | 88 +++++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 60 +++++++++++++ 4 files changed, 176 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index a21710892eb..d55eb3ac014 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -61,6 +61,13 @@ _AUTH_FLOW_SCOPED_FIELDS: frozenset = frozenset( } ) + +def _blank_to_none(value: Optional[str]) -> Optional[str]: + if not isinstance(value, str): + return None + return value.strip() or None + + # Token-exchange settings with dedicated columns that also exist on # ``MCPCredentials`` as a legacy shape (rows and REST callers that predate the # columns). Every write lifts blob values into the columns and strips them from @@ -705,7 +712,8 @@ async def update_mcp_server( # legacy blob copies below, so the existing row is needed for those updates. explicit_te_write = bool(_TOKEN_EXCHANGE_COLUMN_FIELDS & data_dict.keys()) url_provided = "url" in data_dict and data_dict["url"] is not None - if data.auth_type or has_credentials or explicit_te_write or url_provided: + issuer_provided = "issuer" in data_dict + if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided: existing = await MCPServerRepository(prisma_client).table.find_unique(where={"server_id": data.server_id}) auth_type_changed = bool( @@ -716,12 +724,16 @@ async def update_mcp_server( # A url change re-points the server at a potentially different upstream, so any discovered or # trust-on-first-use OAuth endpoints/issuer belong to the old upstream and must re-discover. url_changed = bool(url_provided and existing and existing.url != data_dict["url"]) + old_issuer = _blank_to_none(getattr(existing, "issuer", None)) if existing else None + issuer_changed = bool( + issuer_provided and old_issuer is not None and _blank_to_none(data_dict.get("issuer")) != old_issuer + ) # Clear stale credentials when auth_type changes but no new credentials provided if auth_type_changed and "credentials" not in data_dict: data_dict["credentials"] = None - if auth_type_changed or url_changed: + if auth_type_changed or url_changed or issuer_changed: # Clear each auth-flow-scoped field that the caller either omitted (partial update) or # resubmitted unchanged. The edit form re-sends every field, so a stale issuer/endpoint # belonging to the old upstream would otherwise survive a url/auth_type change and win in the diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a547f5b6920..57cd755797b 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1236,6 +1236,12 @@ class MCPServerManager: resolved_registration_url = manual_registration_url or ( gated_oauth_metadata.registration_url if gated_oauth_metadata else None ) + discovered_issuer = ( + gated_oauth_metadata.discovered_issuer + if gated_oauth_metadata and not gated_oauth_metadata.from_origin_fallback + else None + ) + effective_issuer = manual_issuer or discovered_issuer config_oauth2_flow = server_config.get("oauth2_flow", None) if auth_type == MCPAuth.oauth2 and config_oauth2_flow not in ( @@ -1284,7 +1290,7 @@ class MCPServerManager: client_secret=server_config.get("client_secret", None), oauth2_flow=self._explicit_oauth2_flow(config_oauth2_flow), scopes=resolved_scopes, - issuer=manual_issuer, + issuer=effective_issuer, authorization_url=resolved_authorization_url, token_url=resolved_token_url, registration_url=resolved_registration_url, @@ -1700,6 +1706,12 @@ class MCPServerManager: ) resolved_scopes = scopes or (gated_oauth_metadata.scopes if gated_oauth_metadata else None) + discovered_issuer = ( + gated_oauth_metadata.discovered_issuer + if gated_oauth_metadata and not gated_oauth_metadata.from_origin_fallback + else None + ) + effective_issuer = manual_issuer or discovered_issuer new_server = MCPServer( server_id=mcp_server.server_id, @@ -1719,7 +1731,7 @@ class MCPServerManager: client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)), scopes=resolved_scopes, - issuer=manual_issuer, + issuer=effective_issuer, authorization_url=manual_authorization_url or getattr(gated_oauth_metadata, "authorization_url", None), token_url=manual_token_url or getattr(gated_oauth_metadata, "token_url", None), registration_url=manual_registration_url or getattr(gated_oauth_metadata, "registration_url", None), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index 0e53ebd76e1..968fafc0e0e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -270,6 +270,94 @@ async def test_url_change_clears_stale_oauth_fields_even_when_resubmitted_unchan assert data_dict["authorization_url"] == "https://new-idp.example.com/authorize" +@pytest.mark.asyncio +async def test_clearing_pinned_issuer_clears_stale_oauth_endpoints(): + """Clearing a previously pinned issuer must not revive the endpoints resolved under it. Under an + issuer anchor the endpoints come solely from the issuer document and are not persisted, but a row + that was resource-rooted before the pin can still hold stale authorization_url/token_url; clearing + the anchor without clearing those would let them win the resolution merge and be posted to without + fresh discovery (RFC 8414 §3.3 provenance).""" + mock_prisma = _mock_prisma() + existing = MagicMock() + existing.auth_type = "oauth2" + existing.url = "https://same.example.com/mcp" + existing.credentials = None + existing.issuer = "https://pinned-idp.example.com" + existing.token_url = "https://pinned-idp.example.com/token" + existing.authorization_url = "https://pinned-idp.example.com/authorize" + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest( + server_id="my-test-server", + issuer="", # admin clears the anchor; url and auth_type unchanged + token_url="https://pinned-idp.example.com/token", + authorization_url="https://pinned-idp.example.com/authorize", + ) + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + assert data_dict["token_url"] is None + assert data_dict["authorization_url"] is None + + +@pytest.mark.asyncio +async def test_repointing_pinned_issuer_clears_stale_endpoints_keeps_new_issuer(): + """Re-pointing the issuer to a different authorization server invalidates the old issuer's + endpoints while keeping the new issuer the admin submitted.""" + mock_prisma = _mock_prisma() + existing = MagicMock() + existing.auth_type = "oauth2" + existing.url = "https://same.example.com/mcp" + existing.credentials = None + existing.issuer = "https://old-idp.example.com" + existing.token_url = "https://old-idp.example.com/token" + existing.authorization_url = "https://old-idp.example.com/authorize" + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest( + server_id="my-test-server", + issuer="https://new-idp.example.com", + token_url="https://old-idp.example.com/token", # resubmitted stale -> must clear + authorization_url="https://old-idp.example.com/authorize", # resubmitted stale -> must clear + ) + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + assert data_dict["issuer"] == "https://new-idp.example.com" + assert data_dict["token_url"] is None + assert data_dict["authorization_url"] is None + + +@pytest.mark.asyncio +async def test_establishing_issuer_first_time_preserves_discovered_fields(): + """Establishing an issuer for the first time (None -> X), which is exactly what the trust-on-first-use + discovery write-back does, must NOT clear the endpoints or oauth2_flow it discovered in the same + write. Only an issuer that was already pinned and is now changed or cleared invalidates its + endpoints, so the discovery persist cannot wipe the fields it just resolved.""" + mock_prisma = _mock_prisma() + existing = MagicMock() + existing.auth_type = "oauth2" + existing.url = "https://same.example.com/mcp" + existing.credentials = None + existing.issuer = None + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest( + server_id="my-test-server", + issuer="https://discovered-idp.example.com", + authorization_url="https://discovered-idp.example.com/authorize", + token_url="https://discovered-idp.example.com/token", + oauth2_flow="authorization_code", + ) + await update_mcp_server(mock_prisma, data, "mcp_oauth_discovery") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + assert data_dict["issuer"] == "https://discovered-idp.example.com" + assert data_dict["authorization_url"] == "https://discovered-idp.example.com/authorize" + assert data_dict["token_url"] == "https://discovered-idp.example.com/token" + assert data_dict.get("oauth2_flow") == "authorization_code" + + @pytest.mark.asyncio async def test_unchanged_url_does_not_clear_discovered_oauth_fields(): """A partial update that resends the same url (or omits it) must not clear the discovered OAuth 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 1f966627aea..b16840741af 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 @@ -1166,6 +1166,66 @@ class TestMCPServerManager: assert built.token_url == "https://idp.example.com/token" assert built.scopes == ["read", "admin"] + @pytest.mark.asyncio + async def test_build_from_table_reflects_discovered_issuer_trust_on_first_use(self): + """An unpinned server resolves endpoints resource-rooted on first discovery and records the + discovered issuer trust-on-first-use. The returned in-memory server must carry that discovered + issuer so the registry matches what gets persisted to the row; otherwise the OAuth token + identity (which includes issuer) differs between this build and the next rebuild, forcing a + spurious re-auth. Endpoints and issuer come from the same authorization-server document, so + they are consistent.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="tofu-issuer-1", + alias="tofu_issuer", + description="unpinned, discovers its issuer", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["read"], + discovered_issuer="https://idp.example.com", + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + assert built.issuer == "https://idp.example.com" + assert built.authorization_url == "https://idp.example.com/authorize" + + @pytest.mark.asyncio + async def test_build_from_table_origin_fallback_issuer_is_not_reflected(self): + """An origin-fallback discovery is a guess that is deliberately never persisted, so the built + server must not claim an issuer the row will not hold; otherwise in-memory and DB would + disagree in the opposite direction.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="origin-fallback-1", + alias="origin_fallback", + description="unpinned, origin-fallback discovery", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + metadata = MCPOAuthMetadata( + authorization_url="https://up.example.com/authorize", + token_url="https://up.example.com/token", + discovered_issuer="https://up.example.com", + from_origin_fallback=True, + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + assert built.issuer is None + @pytest.mark.asyncio async def test_build_from_table_whitespace_authorization_url_is_not_a_pin(self): """A whitespace-only authorization_url on the row must not be kept for redirects while the From 3cea2431165247646f20535ca189b3e58c490824 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 15 Jul 2026 18:27:38 -0700 Subject: [PATCH 35/62] fix(key management): enforce minimum custom key length and mask short keys in key_name (#33462) * fix(key management): enforce minimum custom key length and mask short keys in key_name * fix(key management): validate new_key before assignment and sync generated schema docstrings * fix(key management): lower minimum custom key length default from 20 to 16 --- litellm/constants.py | 1 + .../litellm_core_utils/secret_redaction.py | 4 +- litellm/proxy/auth/auth_utils.py | 4 +- .../key_management_endpoints.py | 24 +++- .../proxy/auth/test_auth_utils.py | 11 +- .../test_key_management_endpoints.py | 110 +++++++++++++++++- tests/test_litellm/test_secret_redaction.py | 10 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 +- 8 files changed, 157 insertions(+), 15 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 715d57e594d..8e0a5cfe50f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1496,6 +1496,7 @@ MAX_TEAM_LIST_LIMIT = int(os.getenv("MAX_TEAM_LIST_LIMIT", 20)) MAX_POLICY_ESTIMATE_IMPACT_ROWS = int(os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000)) DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7)) LENGTH_OF_LITELLM_GENERATED_KEY = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16)) +MINIMUM_CUSTOM_KEY_LENGTH = int(os.getenv("MINIMUM_CUSTOM_KEY_LENGTH", 16)) SECRET_MANAGER_REFRESH_INTERVAL = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400)) LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ "default_internal_user_params", diff --git a/litellm/litellm_core_utils/secret_redaction.py b/litellm/litellm_core_utils/secret_redaction.py index b526068589d..455d0f00c35 100644 --- a/litellm/litellm_core_utils/secret_redaction.py +++ b/litellm/litellm_core_utils/secret_redaction.py @@ -9,6 +9,8 @@ secrets from strings without depending on the logging-configuration module. import re from typing import List +from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH + _REDACTED = "REDACTED" @@ -30,7 +32,7 @@ def _build_secret_patterns() -> "re.Pattern[str]": # Basic auth headers r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}", # OpenAI / Anthropic sk- prefixed keys - r"sk-[A-Za-z0-9\-_]{20,}", + rf"sk-[A-Za-z0-9\-_]{{{MINIMUM_CUSTOM_KEY_LENGTH - len('sk-')},}}", # Generic api_key / api-key / apikey (handles 'key': 'value' dict repr) r"(?:api[_-]?key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]{8,}", # x-api-key / api-key header values (handles 'key': 'value' dict repr) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 893e09ece6e..a610e44e69c 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -10,7 +10,7 @@ from fastapi import HTTPException, Request, status import litellm from litellm import Router, provider_list from litellm._logging import verbose_proxy_logger -from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS +from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HEADERS from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.proxy._types import * @@ -1533,4 +1533,6 @@ def get_model_from_request( def abbreviate_api_key(api_key: str) -> str: + if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH: + return "sk-..." return f"sk-...{api_key[-4:]}" diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index df311bed7b2..01f4e040e58 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -32,6 +32,7 @@ from litellm._uuid import uuid from litellm.constants import ( LENGTH_OF_LITELLM_GENERATED_KEY, LITELLM_PROXY_ADMIN_NAME, + MINIMUM_CUSTOM_KEY_LENGTH, UI_SESSION_TOKEN_TEAM_ID, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds @@ -1022,6 +1023,14 @@ async def _common_key_generation_helper( detail={"error": f"Invalid key format. LiteLLM Virtual Key must start with 'sk-'. Received: {_masked}"}, ) + if data.key is not None and len(data.key) < MINIMUM_CUSTOM_KEY_LENGTH: + raise HTTPException( + status_code=400, + detail={ + "error": f"Invalid key format. LiteLLM Virtual Key must be at least {MINIMUM_CUSTOM_KEY_LENGTH} characters long." + }, + ) + # check org key limits - done here to handle inheriting org id from team if data.organization_id is not None: from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -1474,7 +1483,7 @@ async def generate_key_fn( Parameters: - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - key_alias: Optional[str] - User defined key alias - - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. + - key: Optional[str] - User defined key value. Must start with 'sk-' and be at least 16 characters long. If not set, a 16-digit unique sk-key is created for you. - team_id: Optional[str] - The team id of the key - user_id: Optional[str] - The user id of the key - agent_id: Optional[str] - The agent id associated with the key. @@ -1688,7 +1697,7 @@ async def generate_service_account_key_fn( Parameters: - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - key_alias: Optional[str] - User defined key alias - - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. + - key: Optional[str] - User defined key value. Must start with 'sk-' and be at least 16 characters long. If not set, a 16-digit unique sk-key is created for you. - team_id: Optional[str] - The team id of the key - user_id: Optional[str] - [NON-FUNCTIONAL] THIS WILL BE IGNORED. The user id of the key - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. @@ -4356,7 +4365,6 @@ async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: if data and data.new_key is not None: # Reject custom key values if disabled by admin await _check_custom_key_allowed(data.new_key) - new_token = data.new_key if not data.new_key.startswith("sk-"): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -4364,6 +4372,12 @@ async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: "error": "New key must start with 'sk-'. This is to distinguish a key hash (used by litellm for logging / internal logic) from the actual key." }, ) + if len(data.new_key) < MINIMUM_CUSTOM_KEY_LENGTH: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": f"New key must be at least {MINIMUM_CUSTOM_KEY_LENGTH} characters long."}, + ) + new_token = data.new_key else: new_token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}" return new_token @@ -4470,7 +4484,7 @@ async def _execute_virtual_key_regeneration( new_token = await get_new_token(data=data) new_token_hash = hash_token(new_token) - new_token_key_name = f"sk-...{new_token[-4:]}" + new_token_key_name = abbreviate_api_key(api_key=new_token) update_data = {"token": new_token_hash, "key_name": new_token_key_name} non_default_values = {} @@ -4550,7 +4564,7 @@ async def regenerate_key_fn( - data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update - key: Optional[str] - The key to regenerate. - new_master_key: Optional[str] - The new master key to use, if key is the master key. - - new_key: Optional[str] - The new key to use, if key is not the master key. If both set, new_master_key will be used. + - new_key: Optional[str] - The new key to use, if key is not the master key. Must start with 'sk-' and be at least 16 characters long. If both set, new_master_key will be used. - key_alias: Optional[str] - User-friendly key alias - user_id: Optional[str] - User ID associated with key - team_id: Optional[str] - Team ID associated with key diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 042fc107f40..17ff700791f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -659,7 +659,16 @@ def test_get_model_from_request_ignores_session_model_on_non_realtime_routes(): def test_abbreviate_api_key(): - assert abbreviate_api_key("sk-test-1234") == "sk-...1234" + assert abbreviate_api_key("sk-test-1234-abcdefgh") == "sk-...efgh" + assert abbreviate_api_key("sk-abcdefghijklm") == "sk-...jklm" + + +def test_abbreviate_api_key_short_key_is_fully_masked(): + """Regression test for LIT-4355: for keys shorter than the enforced minimum, + showing the last 4 characters can reveal the entire key (sk-1234 -> sk-...1234).""" + assert abbreviate_api_key("sk-1234") == "sk-..." + assert abbreviate_api_key("sk-test-1234") == "sk-..." + assert abbreviate_api_key("") == "sk-..." def test_get_customer_user_header_returns_none_when_no_customer_role(): diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index aa24b0199ab..dffca3093fa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -563,7 +563,7 @@ async def test_generate_key_debug_log_never_contains_raw_token(monkeypatch, capl generate_key_fn, ) - raw_key = "sk-short-secret" + raw_key = "sk-short-secret-a1b2" with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): await generate_key_fn( data=GenerateKeyRequest(key=raw_key), @@ -1336,10 +1336,10 @@ async def test_get_new_token_with_valid_key(monkeypatch): ) # Test with valid new_key - data = RegenerateKeyRequest(new_key="sk-test123456789") + data = RegenerateKeyRequest(new_key="sk-test1234567890abc") result = await get_new_token(data) - assert result == "sk-test123456789" + assert result == "sk-test1234567890abc" @pytest.mark.asyncio @@ -1370,6 +1370,110 @@ async def test_get_new_token_with_invalid_key(monkeypatch): assert "New key must start with 'sk-'" in str(exc_info.value.detail) +@pytest.mark.asyncio +async def test_get_new_token_rejects_short_new_key(monkeypatch): + """Regression test for LIT-4355: a short custom key like sk-99 must be rejected, + otherwise the stored key_name (sk-...{last 4 chars}) reveals the entire key.""" + from unittest.mock import AsyncMock + + from fastapi import HTTPException + + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + get_new_token, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + + data = RegenerateKeyRequest(new_key="sk-99") + + with pytest.raises(HTTPException) as exc_info: + await get_new_token(data) + + assert exc_info.value.status_code == 400 + assert "at least 16 characters" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("short_key", ["sk-1234", "sk-abcdefghijkl"]) +async def test_generate_key_fn_rejects_short_custom_key(monkeypatch, short_key): + """Regression test for LIT-4355: /key/generate must reject custom keys shorter + than the minimum length (including the 15-char boundary); sk-1234 used to be + accepted and fully exposed via key_name.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles, ProxyException + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_fn, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + + assert len(short_key) < 16 + + with pytest.raises(ProxyException) as exc_info: + await generate_key_fn( + data=GenerateKeyRequest(key=short_key), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234" + ), + ) + + assert exc_info.value.code == "400" + assert "at least 16 characters" in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_generate_key_fn_accepts_custom_key_at_minimum_length(monkeypatch): + """Custom keys at exactly the minimum length (16 chars) are still accepted.""" + mock_prisma_client = AsyncMock() + mock_insert_data = AsyncMock( + return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None, object_permission=None) + ) + mock_prisma_client.insert_data = mock_insert_data + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_fn, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + + custom_key = "sk-abcdefghijklm" + assert len(custom_key) == 16 + + response = await generate_key_fn( + data=GenerateKeyRequest(key=custom_key), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234" + ), + ) + + assert response.key == custom_key + + @pytest.mark.asyncio async def test_check_custom_key_allowed_when_disabled(monkeypatch): """_check_custom_key_allowed raises 403 when disable_custom_api_keys is true.""" diff --git a/tests/test_litellm/test_secret_redaction.py b/tests/test_litellm/test_secret_redaction.py index 85430ba752b..7152eea7c9f 100644 --- a/tests/test_litellm/test_secret_redaction.py +++ b/tests/test_litellm/test_secret_redaction.py @@ -65,6 +65,16 @@ def test_redact_string_catches_secret_patterns(): assert redact_string(normal) == normal +def test_redact_string_catches_minimum_length_virtual_key(): + """Regression test for LIT-4355: keys at the enforced 16-char minimum + (MINIMUM_CUSTOM_KEY_LENGTH) must be treated as key-shaped by the scrubber.""" + minimum_length_key = "sk-abcdefghijklm" + assert len(minimum_length_key) == 16 + result = redact_string("msg: " + minimum_length_key) + assert minimum_length_key not in result + assert "REDACTED" in result + + def test_filter_redacts_secrets_in_logger_output(): def log_messages(): verbose_logger.debug("Key: " + SECRET) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 394f7be9859..87c760257d1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -6544,7 +6544,7 @@ export interface paths { * Parameters: * - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). * - key_alias: Optional[str] - User defined key alias - * - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. + * - key: Optional[str] - User defined key value. Must start with 'sk-' and be at least 16 characters long. If not set, a 16-digit unique sk-key is created for you. * - team_id: Optional[str] - The team id of the key * - user_id: Optional[str] - The user id of the key * - agent_id: Optional[str] - The agent id associated with the key. @@ -6765,7 +6765,7 @@ export interface paths { * - data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update * - key: Optional[str] - The key to regenerate. * - new_master_key: Optional[str] - The new master key to use, if key is the master key. - * - new_key: Optional[str] - The new key to use, if key is not the master key. If both set, new_master_key will be used. + * - new_key: Optional[str] - The new key to use, if key is not the master key. Must start with 'sk-' and be at least 16 characters long. If both set, new_master_key will be used. * - key_alias: Optional[str] - User-friendly key alias * - user_id: Optional[str] - User ID associated with key * - team_id: Optional[str] - Team ID associated with key @@ -6834,7 +6834,7 @@ export interface paths { * Parameters: * - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). * - key_alias: Optional[str] - User defined key alias - * - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. + * - key: Optional[str] - User defined key value. Must start with 'sk-' and be at least 16 characters long. If not set, a 16-digit unique sk-key is created for you. * - team_id: Optional[str] - The team id of the key * - user_id: Optional[str] - [NON-FUNCTIONAL] THIS WILL BE IGNORED. The user id of the key * - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. @@ -7024,7 +7024,7 @@ export interface paths { * - data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update * - key: Optional[str] - The key to regenerate. * - new_master_key: Optional[str] - The new master key to use, if key is the master key. - * - new_key: Optional[str] - The new key to use, if key is not the master key. If both set, new_master_key will be used. + * - new_key: Optional[str] - The new key to use, if key is not the master key. Must start with 'sk-' and be at least 16 characters long. If both set, new_master_key will be used. * - key_alias: Optional[str] - User-friendly key alias * - user_id: Optional[str] - User ID associated with key * - team_id: Optional[str] - Team ID associated with key From f6516e7be5256fa69f1ff29f7a1ec53e566f8ecc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 15 Jul 2026 21:41:45 -0700 Subject: [PATCH 36/62] fix(ui): stop sending the complexity-router pseudo-model to /health/test_connection (#33498) * fix(ui): stop sending the complexity-router pseudo-model to /health/test_connection Test Connection on a saved auto-router model sent the raw "auto_router/complexity_router" model string to the generic health-check endpoint, which always failed with "Unmapped LLM provider" since it's a routing-strategy config, not a real completion endpoint. Reuse the per-tier connection test already built for the Add Auto Router wizard: for complexity-router models, test each configured tier's underlying model group instead of the router pseudo-model. Semantic-type auto routers (auto_router_config) have no equivalent tier-based test yet, so the button is hidden for them instead of guaranteed to fail. * fix(ui): address review feedback on auto-router test connection fix Type the complexity-router config parsing instead of using `any`, use NotificationsManager.warning instead of fromBackend for the client-generated "no tiers configured" message, remove comments added in the previous commit, and also test the deployment's configured complexity_router_default_model as a fallback target when it isn't already covered by a configured tier (matches the fallback Router itself uses for unconfigured tiers). --- .../src/components/model_info_view.test.tsx | 155 ++++++++++++++++++ .../src/components/model_info_view.tsx | 113 ++++++++++++- 2 files changed, 260 insertions(+), 8 deletions(-) diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 2546601b4db..a496bb05b91 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -16,6 +16,7 @@ vi.mock("./molecules/notifications_manager", () => ({ success: vi.fn(), error: vi.fn(), info: vi.fn(), + warning: vi.fn(), fromBackend: vi.fn(), }, })); @@ -27,6 +28,7 @@ vi.mock("./networking", () => ({ getGuardrailsList: vi.fn(), tagListCall: vi.fn(), testConnectionRequest: vi.fn(), + testModelGroupConnection: vi.fn(), modelPatchUpdateCall: vi.fn(), modelDeleteCall: vi.fn(), credentialCreateCall: vi.fn(), @@ -52,6 +54,7 @@ const mockCredentialListCall = vi.mocked(networking.credentialListCall); const mockGetGuardrailsList = vi.mocked(networking.getGuardrailsList); const mockTagListCall = vi.mocked(networking.tagListCall); const mockTestConnectionRequest = vi.mocked(networking.testConnectionRequest); +const mockTestModelGroupConnection = vi.mocked(networking.testModelGroupConnection); const mockModelPatchUpdateCall = vi.mocked(networking.modelPatchUpdateCall); const mockModelDeleteCall = vi.mocked(networking.modelDeleteCall); const mockCredentialCreateCall = vi.mocked(networking.credentialCreateCall); @@ -732,6 +735,158 @@ describe("ModelInfoView", () => { }); }); + it("does not offer Test Connection for semantic auto router models (no tier-based test exists yet)", async () => { + const semanticAutoRouterModelData = { + ...defaultModelData, + litellm_params: { + ...defaultModelData.litellm_params, + auto_router_config: {}, + }, + }; + + mockUseModelsInfo.mockReturnValue({ + data: { + data: [semanticAutoRouterModelData], + }, + isLoading: false, + error: null, + }); + + render(, { wrapper }); + await waitFor(() => { + expect(screen.getByText("Model Settings")).toBeInTheDocument(); + }); + expect(screen.queryByTestId("test-connection-button")).not.toBeInTheDocument(); + }); + + it("tests each complexity tier's model group instead of sending the router pseudo-model to /health/test_connection (regression: raw test previously threw 'Unmapped LLM provider... model=complexity_router')", async () => { + const complexityRouterModelData = { + ...defaultModelData, + litellm_params: { + ...defaultModelData.litellm_params, + model: "auto_router/complexity_router", + complexity_router_config: { + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: ["gpt-4o"], COMPLEX: [], REASONING: [] }, + }, + }, + }; + + mockUseModelsInfo.mockReturnValue({ + data: { + data: [complexityRouterModelData], + }, + isLoading: false, + error: null, + }); + mockTestModelGroupConnection.mockResolvedValue({ status: "success" }); + + render(, { wrapper }); + const testConnectionButton = await screen.findByTestId("test-connection-button"); + await userEvent.click(testConnectionButton); + + await waitFor(() => { + expect(mockTestModelGroupConnection).toHaveBeenCalledWith("test-token", "gpt-4o-mini", "chat"); + expect(mockTestModelGroupConnection).toHaveBeenCalledWith("test-token", "gpt-4o", "chat"); + }); + expect(mockTestConnectionRequest).not.toHaveBeenCalled(); + }); + + it("also tests the configured default model when an unconfigured tier would fall back to it in production", async () => { + const complexityRouterModelData = { + ...defaultModelData, + litellm_params: { + ...defaultModelData.litellm_params, + model: "auto_router/complexity_router", + complexity_router_config: { + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] }, + }, + complexity_router_default_model: "gpt-4o", + }, + }; + + mockUseModelsInfo.mockReturnValue({ + data: { + data: [complexityRouterModelData], + }, + isLoading: false, + error: null, + }); + mockTestModelGroupConnection.mockResolvedValue({ status: "success" }); + + render(, { wrapper }); + const testConnectionButton = await screen.findByTestId("test-connection-button"); + await userEvent.click(testConnectionButton); + + await waitFor(() => { + expect(mockTestModelGroupConnection).toHaveBeenCalledWith("test-token", "gpt-4o-mini", "chat"); + expect(mockTestModelGroupConnection).toHaveBeenCalledWith("test-token", "gpt-4o", "chat"); + }); + }); + + it("does not duplicate the default model as a test target when it is already covered by a configured tier", async () => { + const complexityRouterModelData = { + ...defaultModelData, + litellm_params: { + ...defaultModelData.litellm_params, + model: "auto_router/complexity_router", + complexity_router_config: { + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: ["gpt-4o"], COMPLEX: [], REASONING: [] }, + }, + complexity_router_default_model: "gpt-4o", + }, + }; + + mockUseModelsInfo.mockReturnValue({ + data: { + data: [complexityRouterModelData], + }, + isLoading: false, + error: null, + }); + mockTestModelGroupConnection.mockResolvedValue({ status: "success" }); + + render(, { wrapper }); + const testConnectionButton = await screen.findByTestId("test-connection-button"); + await userEvent.click(testConnectionButton); + + await waitFor(() => { + expect(mockTestModelGroupConnection).toHaveBeenCalledWith("test-token", "gpt-4o", "chat"); + }); + expect(mockTestModelGroupConnection).toHaveBeenCalledTimes(2); + }); + + it("warns instead of erroring when no complexity tiers are configured to test", async () => { + const complexityRouterModelData = { + ...defaultModelData, + litellm_params: { + ...defaultModelData.litellm_params, + model: "auto_router/complexity_router", + complexity_router_config: { + tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, + }, + }, + }; + + mockUseModelsInfo.mockReturnValue({ + data: { + data: [complexityRouterModelData], + }, + isLoading: false, + error: null, + }); + + render(, { wrapper }); + const testConnectionButton = await screen.findByTestId("test-connection-button"); + await userEvent.click(testConnectionButton); + + await waitFor(() => { + expect(mockNotificationsManager.warning).toHaveBeenCalledWith( + "No complexity tiers are configured yet, so there is nothing to test.", + ); + }); + expect(mockTestModelGroupConnection).not.toHaveBeenCalled(); + }); + it("should display model access groups field", async () => { render(, { wrapper }); await waitFor(() => { diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index aaad6e5a474..5b7f82c44b4 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -23,6 +23,8 @@ import { CheckIcon, CopyIcon } from "lucide-react"; import { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../utils/dataUtils"; import { formItemValidateJSON, truncateString } from "../utils/textUtils"; +import AutoRouterConnectionTest from "./add_model/auto_router_connection_test"; +import { AutoRouterTestTarget, buildAutoRouterTestTargets } from "./add_model/build_auto_router_test_targets"; import CacheControlSettings from "./add_model/cache_control_settings"; import DeleteResourceModal from "./common_components/DeleteResourceModal"; import EditAutoRouterModal from "./edit_auto_router/edit_auto_router_modal"; @@ -68,6 +70,66 @@ const isMaskedSecret = (value: unknown): boolean => typeof value === "string" && const stripMaskedSecrets = (params: Record): Record => Object.fromEntries(Object.entries(params).filter(([, value]) => !isMaskedSecret(value))); +const normalizeTierModels = (value: unknown): string[] => { + if (Array.isArray(value)) return value; + if (typeof value === "string" && value) return [value]; + return []; +}; + +interface ComplexityRouterTierConfig { + tiers?: { + SIMPLE?: unknown; + MEDIUM?: unknown; + COMPLEX?: unknown; + REASONING?: unknown; + }; + semantic_keyword_matching?: boolean; + embedding_model?: string; +} + +interface ComplexityRouterModelData { + litellm_params?: { + complexity_router_config?: ComplexityRouterTierConfig | string; + complexity_router_default_model?: string; + }; +} + +const buildComplexityRouterTestTargets = ( + modelData: ComplexityRouterModelData | null | undefined, +): AutoRouterTestTarget[] => { + const rawConfig = modelData?.litellm_params?.complexity_router_config; + let config: ComplexityRouterTierConfig = {}; + if (typeof rawConfig === "string") { + try { + config = JSON.parse(rawConfig); + } catch { + config = {}; + } + } else if (rawConfig) { + config = rawConfig; + } + + const tierTargets = buildAutoRouterTestTargets({ + tiers: { + SIMPLE: normalizeTierModels(config.tiers?.SIMPLE), + MEDIUM: normalizeTierModels(config.tiers?.MEDIUM), + COMPLEX: normalizeTierModels(config.tiers?.COMPLEX), + REASONING: normalizeTierModels(config.tiers?.REASONING), + }, + semanticMatchingEnabled: Boolean(config.semantic_keyword_matching), + embeddingModel: config.embedding_model, + }); + + const defaultModel = modelData?.litellm_params?.complexity_router_default_model?.trim(); + if (!defaultModel || tierTargets.some((target) => target.modelGroup === defaultModel)) { + return tierTargets; + } + return [ + ...tierTargets, + { labels: ["Default (unconfigured tiers)"], modelGroup: defaultModel, mode: "chat" as const }, + ]; +}; + export default function ModelInfoView({ modelId, onClose, @@ -91,6 +153,9 @@ export default function ModelInfoView({ const [showCacheControl, setShowCacheControl] = useState(false); const [copiedStates, setCopiedStates] = useState>({}); const [isAutoRouterModalOpen, setIsAutoRouterModalOpen] = useState(false); + const [isAutoRouterTestModalOpen, setIsAutoRouterTestModalOpen] = useState(false); + const [autoRouterTestId, setAutoRouterTestId] = useState(0); + const [autoRouterTestTargets, setAutoRouterTestTargets] = useState([]); const [guardrailsList, setGuardrailsList] = useState([]); const [tagsList, setTagsList] = useState>({}); const [credentialsList, setCredentialsList] = useState([]); @@ -128,6 +193,9 @@ export default function ModelInfoView({ modelData?.litellm_params?.auto_router_config != null || modelData?.litellm_params?.complexity_router_config != null || modelData?.litellm_params?.model?.startsWith("auto_router/complexity_router"); + const isComplexityRouter = + modelData?.litellm_params?.complexity_router_config != null || + modelData?.litellm_params?.model?.startsWith("auto_router/complexity_router"); const usingExistingCredential = modelData?.litellm_params?.litellm_credential_name != null && @@ -435,6 +503,17 @@ export default function ModelInfoView({ const handleTestConnection = async () => { if (!accessToken) return; + if (isComplexityRouter) { + const targets = buildComplexityRouterTestTargets(localModelData ?? modelData); + if (targets.length === 0) { + NotificationsManager.warning("No complexity tiers are configured yet, so there is nothing to test."); + return; + } + setAutoRouterTestTargets(targets); + setAutoRouterTestId((id) => id + 1); + setIsAutoRouterTestModalOpen(true); + return; + } try { NotificationsManager.info("Testing connection..."); const response = await testConnectionRequest( @@ -536,14 +615,16 @@ export default function ModelInfoView({
- + {(!isAutoRouter || isComplexityRouter) && ( + + )} , + ]} + width={700} + > + {isAutoRouterTestModalOpen && accessToken && ( + + )} +
); } From cf90445574c718fbac653a36cd714e8b5999fdb3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 15 Jul 2026 21:46:02 -0700 Subject: [PATCH 37/62] feat(cli): add lite up/down to ambiently route Claude Code through the proxy (#33231) * feat(cli): add `lite up`/`lite down` to ambiently route Claude Code through the proxy Patches ~/.claude/settings.json in place (env.ANTHROPIC_BASE_URL + apiKeyHelper via `lite auth print-token`) so any `claude` session started afterward, from any terminal, routes through the local LiteLLM proxy with no wrapper command needed, unlike the existing `lite claude` subprocess-exec approach. Backs up the original file first and restores it on Ctrl-C/SIGTERM, or via `lite down` after an unclean exit. Cursor is not supported: no equivalent file-based config to patch. * feat(cli): add lite autoroute to QA complexity-based auto-routing against a real proxy (#33249) * feat(cli): add lite autoroute to QA complexity-based auto-routing against a real proxy Lets a customer try litellm's complexity_router against models they already have on their existing, unmodified production proxy, with no config.yaml edits and no new infra. lite autoroute configure discovers accessible models via /model_group/info and walks through tier assignment (plus optional LLM classifier / semantic matching / adaptive selection); every referenced model becomes its own litellm_proxy/ deployment forwarding back to the real proxy with the real key, so every actual call, routed completions, classifier calls, embedding calls, still lands on their real proxy. lite autoroute up launches that generated config as an ephemeral local proxy, patches ~/.claude/settings.json to point Claude Code at it, and streams routing decisions live; Ctrl-C/SIGTERM (or lite autoroute down after an unclean exit) restores everything. Also adds lite model-groups list (a thin CLI wrapper over the existing ModelGroupsManagementClient), and generalizes up.py's settings-backup/restore helpers to take explicit paths so this feature can reuse them instead of duplicating the logic. Depends on litellm_lite_up_down (#33231) for that generalization. * feat(cli): allow multiple models per autoroute tier complexity_router already supports a pool of models per tier (randomly picked per request; adaptive mode specifically needs a pool to choose within), but the configure wizard only ever let you assign one. Tiers are now a tuple of model names; the wizard prompt accepts comma-separated indices to pick more than one per tier. * feat(cli): fuzzy model picker and auto-route Claude Code to autorouter Numbered-index selection didn't scale past a handful of models, so switch the tier picker to InquirerPy's fzf-style fuzzy search. Also set ANTHROPIC_DEFAULT_{SONNET,HAIKU,OPUS}_MODEL to "autorouter" in Claude Code's settings, since Router resolves auto-router deployments by literal model name with no wildcard support, so a "*" catch-all model_name would never match real traffic. * feat(cli): allow installing lite CLI from source via LITELLM_CLI_REF Lets testers try an unreleased branch's CLI changes with the same curl-piped installer, instead of waiting for a PyPI release. * fix(ci): modernize type hints to clear ruff strict-rule budget * fix(ci): bump httplib2 and setuptools to patched versions Clears osv-scan findings for PYSEC-2026-3444 and PYSEC-2026-3447. * fix(cli): write autoroute's secret-bearing files with mode 0600 commands.py wrote config.yaml (embeds the real proxy key) and Claude Code's settings.json (embeds the ephemeral proxy's master key) with plain open(), landing at the umask-derived default (commonly 0644) until a later chmod call caught up. That window, and the missed case where settings.json already exists (chmod never ran at all there), left a credential-bearing file readable by another local account. secure_create() fixes the mode via fchmod on the fd before any content is written, covering both the brand-new-file and already-exists cases, and commands.py/wizard.py now route their sensitive writes through it. * docs(cli): warn that a stale Claude Code session can leak to a squatted port lite autoroute up's master key is embedded statically (unlike lite up's apiKeyHelper, resolved per request), so a Claude Code session still running after teardown keeps sending it, along with prompt content, to a now-unbound loopback port that another local account can bind. This is the same one-time-patch tradeoff lite up already accepts, just with a static secret instead of a re-resolved one -- document it in the README's Caveats section and surface it in the teardown message itself. * fix(cli): address greptile review feedback on autoroute PR - terminate the ephemeral proxy child process when its health check fails, instead of leaking an orphaned, unrecoverable process bound to the port - replace bare assert isinstance checks (no-ops under python -O) with click.ClickException in the model-groups list and configure wizard code paths - close launch_proxy's log file handle once the child process has inherited its fd, instead of leaking it - add build_generated_proxy_config to config.py's __all__ * fix(cli): close TOCTOU window in lite up's settings backup write write_backup wrote the backup (which can embed the original apiKeyHelper/settings content) with plain open() + a chmod call after the fact -- the same permissive-until-corrected window already fixed for autoroute's config.yaml and Claude settings writes, and missed entirely when the backup file already exists with broader permissions. Moves secure_create (atomic-enough 0600 via fchmod before any content is written) to up.py, the module both lite up and lite autoroute share, and has autoroute/process.py import it from there instead of keeping its own copy. * fix(cli): refuse autoroute up when a stale backup exists from a crash The pid-record check only catches a still-live duplicate process; a SIGKILL'd `up` leaves no live pid but does leave AUTOROUTE_BACKUP_PATH behind. Without this guard, a fresh `up` overwrote that backup with the currently-patched Claude settings instead of the true originals, so `down`/Ctrl-C would restore the wrong content permanently. up.py's `lite up` already guards the analogous case; mirror it here. * fix(cli): bind the ephemeral autoroute proxy to loopback only proxy_cli.py defaults --host to 0.0.0.0 when not passed explicitly. launch_proxy never passed it, so the ephemeral proxy -- despite every base_url in this module being built from 127.0.0.1 -- was actually reachable from other hosts on the network, including its unauthenticated-until-config-lands routes before the master key is wired in. * docs(cli): show curl install for the autoroute QA flow Points readers at scripts/install-cli.sh's curl one-liner instead of assuming uv/pip is already set up, and documents the LITELLM_CLI_REF override for trying an unreleased branch or commit. * fix(cli): surface a clean error on an empty or corrupt autoroute config A configure run killed between secure_create's O_TRUNC and the write completing leaves an empty config.yaml on disk. The next up read that via yaml.safe_load (None) into the generated-config TypeAdapter uncaught, surfacing a raw pydantic.ValidationError instead of pointing the user back at `lite autoroute configure`. * fix(cli): bind lite up's apiKeyHelper to the proxy it was started against _ensure_fresh_login only checked token freshness, not which proxy the cached token belonged to, and resolve_api_key_helper built a bare `lite auth print-token` command with no --base-url. A user logged into proxy A who ran `up --base-url proxy-b` (or LITELLM_PROXY_URL=proxy-b) would silently get proxy A's real token wired into Claude Code's apiKeyHelper; since apiKeyHelper is invoked bare, print-token's existing origin check never engaged, so proxy B -- attacker-controlled or not -- received every subsequent request's Authorization header carrying proxy A's credential. _ensure_fresh_login now requires the cached token's base_url to match before treating it as usable, forcing a fresh login for the selected proxy otherwise. resolve_api_key_helper now takes that base_url and threads it through as an explicit --base-url, so print-token's existing (but previously unreachable in the apiKeyHelper flow) base_url_explicit check actually enforces the match at request time too. * fix(cli): surface clean errors instead of raw tracebacks in lite up/down load_json_or_empty and read_backup both delegate to pydantic's validate_json, which raises ValidationError on invalid JSON or a non-object root -- neither up() nor down() caught it, so a corrupt settings or backup file surfaced an unformatted Python traceback instead of a clean CLI error. Both now convert to UpError, and down() (previously uncaught entirely) and up()'s teardown path now handle it. restore_claude_settings also gained a parent.mkdir guard before rewriting CLAUDE_SETTINGS_PATH: if ~/.claude/ was removed while `lite up` was running, the restore would crash before deleting the backup file, permanently stranding it and breaking every future `lite down`. * docs(cli): call out env-var auth for autoroute commands * fix(cli): clean up leaked proxy and surface clean errors in autoroute Three related gaps, all following an UpError getting raised somewhere that wasn't catching it yet: - up() left the just-launched ephemeral proxy running with no pid record if load_json_or_empty/write_backup/secure_create raised after the health check passed, mirroring the existing ProcessLaunchError cleanup for the health-check-failure branch. - _teardown() didn't catch restore_claude_settings raising UpError (e.g. a corrupt backup at stop time), which would otherwise escape to Click as an unhandled error in the normal-exit path, or print "Error in atexit" in the atexit path. up.py's own _restore_once handles the identical case the same way. - read_pid_record let a corrupt PID file surface a raw pydantic.ValidationError instead of a clean message, and did so in down(), the command specifically meant for crash recovery. down() now clears an unreadable pid record and continues cleanup instead of aborting, since a corrupt pid file must never block the one command meant to recover from exactly this kind of crash. * docs(cli): warn against running lite up and lite autoroute up together --- litellm/proxy/client/cli/README.md | 101 +++++ litellm/proxy/client/cli/commands/agents.py | 5 +- .../client/cli/commands/autoroute/__init__.py | 0 .../client/cli/commands/autoroute/commands.py | 196 +++++++++ .../client/cli/commands/autoroute/config.py | 237 +++++++++++ .../client/cli/commands/autoroute/process.py | 174 ++++++++ .../client/cli/commands/autoroute/settings.py | 46 ++ .../client/cli/commands/autoroute/wizard.py | 128 ++++++ .../proxy/client/cli/commands/model_groups.py | 57 +++ litellm/proxy/client/cli/commands/up.py | 283 +++++++++++++ litellm/proxy/client/cli/main.py | 10 + pyproject.toml | 4 +- scripts/install-cli.sh | 19 +- .../proxy/client/cli/autoroute/__init__.py | 0 .../client/cli/autoroute/test_commands.py | 300 +++++++++++++ .../proxy/client/cli/autoroute/test_config.py | 181 ++++++++ .../client/cli/autoroute/test_process.py | 139 ++++++ .../client/cli/autoroute/test_settings.py | 56 +++ .../proxy/client/cli/autoroute/test_wizard.py | 293 +++++++++++++ .../proxy/client/cli/test_auth_commands.py | 49 ++- .../client/cli/test_model_groups_commands.py | 114 +++++ .../proxy/client/cli/test_up_commands.py | 398 ++++++++++++++++++ uv.lock | 154 ++++--- 23 files changed, 2869 insertions(+), 75 deletions(-) create mode 100644 litellm/proxy/client/cli/commands/autoroute/__init__.py create mode 100644 litellm/proxy/client/cli/commands/autoroute/commands.py create mode 100644 litellm/proxy/client/cli/commands/autoroute/config.py create mode 100644 litellm/proxy/client/cli/commands/autoroute/process.py create mode 100644 litellm/proxy/client/cli/commands/autoroute/settings.py create mode 100644 litellm/proxy/client/cli/commands/autoroute/wizard.py create mode 100644 litellm/proxy/client/cli/commands/model_groups.py create mode 100644 litellm/proxy/client/cli/commands/up.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/__init__.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_commands.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_config.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_process.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_settings.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py create mode 100644 tests/test_litellm/proxy/client/cli/test_model_groups_commands.py create mode 100644 tests/test_litellm/proxy/client/cli/test_up_commands.py diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index ce13a906a36..be8905e11ca 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -471,6 +471,107 @@ The token minted by `lite login` is a short-lived, per-session agent credential, The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, and `lite opencode` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead. +### Route Every Claude Code Session Through the Proxy + +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. + +Two things need to already be true: you've run `lite login`, since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. + +```bash +lite login +litellm --config litellm/proxy/dev_config.yaml & +lite up +``` + +`lite up` runs in the foreground and blocks. Press Ctrl-C to stop it, which restores the original settings file and exits. If the process is ever killed uncleanly instead -- `kill -9`, a crash -- the settings file is left patched, and `lite down` is the manual recovery path: run it at any later point to restore from the same backup. + +This is a one-time file patch and restore, not a live traffic interceptor. A Claude Code session already running before `lite up` started keeps whatever `ANTHROPIC_BASE_URL` and token it loaded at its own startup, and a session still running when `lite up` stops keeps routing through the proxy until it exits; only sessions *started* while the patch is in effect are affected, and only *new* sessions after a restore go back to Anthropic directly. + +Cursor is not supported: it has no equivalent file-based config to hot-patch this way, since its model routing lives in its own app storage and is configured through its GUI. + +### QA Complexity-Based Auto-Routing Against Your Real Proxy + +`lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session. + +#### Install the CLI + +If you don't already have the `lite` command, install it with a single curl command -- no existing Python tooling required, `uv` is bootstrapped automatically if missing: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install-cli.sh | sh +``` + +This installs only `litellm[cli]`, the thin client (`lite`), not the full proxy server. To try an unreleased branch or commit instead of the latest PyPI release, set `LITELLM_CLI_REF`: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm//scripts/install-cli.sh | \ + LITELLM_CLI_REF= sh +``` + +Point the CLI at your real proxy and key before running any `lite model-groups` or `lite autoroute` command -- like every other command in this CLI, they read `LITELLM_PROXY_URL`/`LITELLM_PROXY_API_KEY` (or `--base-url`/`--api-key`), no `lite login` required: + +```bash +export LITELLM_PROXY_URL=http://localhost:4000 +export LITELLM_PROXY_API_KEY=sk-... +``` + +#### List Your Accessible Model Groups + +```bash +lite model-groups list [--format table|json] +``` + +Lists the model groups your key can reach on the proxy, via `/model_group/info`, along with each group's mode (`chat`, `embedding`, etc.) and per-token pricing. This is also what `lite autoroute configure` uses internally to discover what it can offer you. + +#### Configure the Auto-Router + +```bash +lite autoroute configure +``` + +An interactive wizard. It runs the same model-group discovery as above, splits the results into chat-capable and embedding-capable pools, and asks you to assign one or more models from the chat pool to each of the four complexity tiers -- SIMPLE, MEDIUM, COMPLEX, REASONING. Each tier's picker is a type-to-filter fuzzy search (fzf-style) rather than a scrollable numbered list, so it stays usable even with hundreds of model groups: type a substring to narrow the list, tab to toggle a model into the selection, enter to confirm (assigning more than one model to a tier is exactly when this matters -- complexity_router picks randomly among a tier's pool per request, and adaptive mode specifically depends on having more than one candidate to choose from). From there it optionally offers: classifying prompt complexity with an LLM (again picked from your discovered pool) instead of the free built-in heuristic scorer, semantic keyword matching for tier assignment (needs an embedding model from the pool), and adaptive (bandit-based) selection layered on top of tiering. + +The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key. + +You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.) + +You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first. + +#### Launch the Ephemeral Auto-Router Proxy + +```bash +lite autoroute up +``` + +Starts a local, throwaway litellm proxy on a random free port, running the config `configure` generated, with a freshly-minted random API key baked in for this session only (your real proxy key never leaves the generated config -- it only appears there, forwarding to your real proxy). It waits for the ephemeral proxy to report healthy, then patches `~/.claude/settings.json` the same way `lite up` does, except with a static `ANTHROPIC_AUTH_TOKEN` env var instead of an `apiKeyHelper`, since this key is short-lived and self-issued rather than something needing SSO refresh. Any `claude` session started afterward, from any terminal, routes through the ephemeral proxy. + +`lite autoroute up` runs in the foreground and streams the ephemeral proxy's own log file into your terminal, so you can watch its routing decisions -- which tier and model got picked for each request -- as you use Claude Code normally. Press Ctrl-C (or send SIGTERM) to stop it; this kills the child proxy process and restores your original Claude Code settings, in that order. + +#### Recover From an Unclean Shutdown + +```bash +lite autoroute down +``` + +If the `lite autoroute up` process dies uncleanly -- `kill -9`, a crash -- rather than being stopped with Ctrl-C, `down` is the manual recovery path: it kills any leftover ephemeral proxy process found via a recorded pid file and restores Claude Code's settings from whatever backup is on disk. + +#### Example + +```bash +lite autoroute configure +lite autoroute up +# use Claude Code as normal in another terminal; routing decisions stream live +lite autoroute down # only needed if `up` was killed uncleanly instead of Ctrl-C'd +``` + +#### Caveats + +Adaptive mode's learned state does not persist across `lite autoroute up` sessions -- there is no local database, so every session starts adaptive selection cold. A Claude Code session already running before `up` started, or still running when it stops, keeps whatever settings it loaded at its own startup; like `lite up`, this is a one-time file patch and restore, not a live traffic interceptor. Only Claude Code is supported, for the same reason as `lite up`: no other supported agent (for example Cursor) has an equivalent hot-patchable config file. + +A session that outlives `up` (or is still running the moment you stop it) keeps sending requests, master key included, to that now-freed loopback port until you restart it. Once the ephemeral proxy process exits, nothing stops another local account on the same machine from binding that same port and receiving those requests instead -- unlike `lite up`'s `apiKeyHelper`, which is re-resolved per request, `autoroute`'s master key is a static value, so whoever receives them gets a live-looking token along with the prompt content. Restart any Claude Code session before you consider the machine clean, run `lite autoroute down` promptly rather than leaving a stopped session's settings patched, and do not run `lite autoroute up` on a shared or multi-tenant host. + +Do not run `lite up` and `lite autoroute up` at the same time. Each patches `~/.claude/settings.json` and keeps its own separate backup, with no coordination between them: whichever one you stop or crash out of last is the one whose backup gets restored, which can silently leave the *other* mode's settings (a static master key and a now-dead loopback URL, or a stale `apiKeyHelper`) active. Run `lite down` or `lite autoroute down` (whichever applies) before switching to the other mode. + ## Environment Variables The CLI respects the following environment variables: diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index dfcedb70686..6b6252d8ecb 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -212,7 +212,7 @@ def _is_interactive() -> bool: return sys.stdin.isatty() -def _resolve_api_key(ctx: click.Context) -> str: +def resolve_api_key(ctx: click.Context) -> str: base_url = ctx.obj["base_url"] api_key = ctx.obj.get("api_key") if api_key: @@ -238,7 +238,7 @@ _SKIP_VERIFY_HELP = "Skip the pre-launch key check against the proxy." def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: base_url = ctx.obj["base_url"] started_interactive = _is_interactive() - api_key = _resolve_api_key(ctx) + api_key = resolve_api_key(ctx) display_name, _ = agent_profile(binary) click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}") @@ -288,5 +288,6 @@ __all__ = [ "agent_launch_args", "verify_proxy_key", "agent_profile", + "resolve_api_key", "AgentRunError", ] diff --git a/litellm/proxy/client/cli/commands/autoroute/__init__.py b/litellm/proxy/client/cli/commands/autoroute/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/client/cli/commands/autoroute/commands.py b/litellm/proxy/client/cli/commands/autoroute/commands.py new file mode 100644 index 00000000000..428daeefa60 --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/commands.py @@ -0,0 +1,196 @@ +import atexit +import json +import secrets +import signal +import threading +from types import FrameType + +import click +import yaml +from pydantic import JsonValue, TypeAdapter, ValidationError + +from ..up import CLAUDE_SETTINGS_PATH, UpError, load_json_or_empty, restore_claude_settings, write_backup +from ..up import BackupRecord as ClaudeBackupRecord +from .process import ( + AUTOROUTE_DIR, + CONFIG_PATH, + LOG_PATH, + PidRecord, + ProcessLaunchError, + allocate_free_port, + clear_pid_record, + is_running, + launch_proxy, + poll_liveliness, + read_pid_record, + secure_create, + stream_log, + terminate, + write_pid_record, +) +from .settings import merge_claude_settings_static_token +from .wizard import run_configure_wizard + +AUTOROUTE_BACKUP_PATH = AUTOROUTE_DIR / "claude_settings_backup.json" + +_GENERATED_CONFIG_ADAPTER = TypeAdapter(dict[str, JsonValue]) + + +def _mint_and_embed_master_key() -> str: + """Generate a fresh key for this session and write it into the generated config.yaml. + + Must go under general_settings, not litellm_settings -- the proxy server only ever + reads general_settings.master_key (proxy_server.py:4530) to authenticate requests. A + key placed under litellm_settings is silently ignored, leaving the ephemeral proxy with + no real auth: any request reaches it regardless of the token Claude Code sends. + """ + master_key = secrets.token_urlsafe(32) + with open(CONFIG_PATH, "r") as f: + try: + generated = _GENERATED_CONFIG_ADAPTER.validate_python(yaml.safe_load(f)) + except (yaml.YAMLError, ValidationError): + raise click.ClickException( + f"{CONFIG_PATH} is empty or corrupt. Run `lite autoroute configure` again to regenerate it." + ) + general_settings = generated.get("general_settings") + updated_settings: dict[str, JsonValue] = { + **(general_settings if isinstance(general_settings, dict) else {}), + "master_key": master_key, + } + updated: dict[str, JsonValue] = {**generated, "general_settings": updated_settings} + with secure_create(CONFIG_PATH) as f: + yaml.safe_dump(updated, f, sort_keys=False) + return master_key + + +@click.group(name="autoroute") +def autoroute_group() -> None: + """QA complexity-based auto-routing against models your key can already use""" + + +@autoroute_group.command("configure") +@click.pass_context +def configure(ctx: click.Context) -> None: + """Discover accessible models and generate an ephemeral auto-router config""" + run_configure_wizard(ctx) + + +@autoroute_group.command("up") +def up() -> None: + """Launch the ephemeral auto-router proxy and route Claude Code through it""" + if not CONFIG_PATH.exists(): + raise click.ClickException("No config found. Run `lite autoroute configure` first.") + + try: + existing_pid = read_pid_record() + except UpError as e: + raise click.ClickException(str(e)) + if existing_pid is not None and is_running(existing_pid.pid): + raise click.ClickException( + "An ephemeral proxy is already running (lite autoroute up looks already active). " + "Run `lite autoroute down` first." + ) + + if AUTOROUTE_BACKUP_PATH.exists(): + raise click.ClickException( + f"{AUTOROUTE_BACKUP_PATH} already exists -- `lite autoroute up` looks like it's already " + "running (or crashed without cleanup). Run `lite autoroute down` first." + ) + + master_key = _mint_and_embed_master_key() + port = allocate_free_port() + base_url = f"http://127.0.0.1:{port}" + process = launch_proxy(CONFIG_PATH, port, LOG_PATH) + write_pid_record(PidRecord(pid=process.pid, port=port, config_path=str(CONFIG_PATH), log_path=str(LOG_PATH))) + + try: + poll_liveliness(base_url, LOG_PATH, process) + except ProcessLaunchError as e: + terminate(process.pid) + clear_pid_record() + raise click.ClickException(str(e)) + + try: + original_existed = CLAUDE_SETTINGS_PATH.exists() + original_settings = load_json_or_empty(CLAUDE_SETTINGS_PATH) + write_backup( + ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None), + AUTOROUTE_BACKUP_PATH, + ) + merged = merge_claude_settings_static_token(original_settings, base_url, master_key) + CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True) + with secure_create(CLAUDE_SETTINGS_PATH) as f: + json.dump(merged, f, indent=2) + except UpError as e: + terminate(process.pid) + clear_pid_record() + raise click.ClickException(str(e)) + + click.echo(f"litellm: ephemeral auto-router proxy up at {base_url} (pid {process.pid})") + click.echo("Claude Code sessions started now will route through it. Press Ctrl-C to stop and restore.") + + stop_event = threading.Event() + restored = threading.Lock() + + def _teardown() -> None: + if not restored.acquire(blocking=False): + return + terminate(process.pid) + clear_pid_record() + try: + restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH) + except UpError as e: + # Runs from atexit/a signal handler too, outside Click's own exception + # handling -- raising here would only produce an unhandled-exception + # warning on stderr, not a clean message. + click.echo(str(e), err=True) + return + click.echo("\nStopped ephemeral proxy and restored Claude Code settings.") + click.echo( + f"Restart any Claude Code session still open from this session, or another local account could " + f"bind the now-free port {port} and receive its requests. Do not use `lite autoroute up` on a " + f"shared or multi-tenant host." + ) + + def _handle_signal(_signum: int, _frame: FrameType | None) -> None: + stop_event.set() + + signal.signal(signal.SIGINT, _handle_signal) + signal.signal(signal.SIGTERM, _handle_signal) + atexit.register(_teardown) + + log_thread = threading.Thread(target=stream_log, args=(LOG_PATH, stop_event), daemon=True) + log_thread.start() + + stop_event.wait() + _teardown() + + +@autoroute_group.command("down") +def down() -> None: + """Restore Claude Code settings and stop a leftover ephemeral proxy, if any""" + try: + record: PidRecord | None = read_pid_record() + except UpError as e: + # down is the crash-recovery path -- a corrupt pid record must not block it; clear the + # unusable record and keep going rather than leaving the user with no way to clean up. + click.echo(f"{e} Clearing it and continuing cleanup.", err=True) + record = None + if record is not None and is_running(record.pid): + terminate(record.pid) + click.echo(f"Stopped leftover ephemeral proxy (pid {record.pid}).") + clear_pid_record() + + try: + restored = restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH) + except UpError as e: + raise click.ClickException(str(e)) + if restored is None: + click.echo("Nothing to restore.") + elif restored.existed: + click.echo(f"Restored {CLAUDE_SETTINGS_PATH} to its original contents.") + else: + click.echo(f"Removed {CLAUDE_SETTINGS_PATH} (it did not exist before `lite autoroute up`).") + + +__all__ = ["autoroute_group"] diff --git a/litellm/proxy/client/cli/commands/autoroute/config.py b/litellm/proxy/client/cli/commands/autoroute/config.py new file mode 100644 index 00000000000..fe2c4e467e7 --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/config.py @@ -0,0 +1,237 @@ +from typing import Literal, Union + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter + +TIER_NAMES: tuple[str, ...] = ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING") +AUTOROUTER_MODEL_NAME = "autorouter" + + +class ConfigGenerationError(Exception): + """Raised when an AutorouteConfig references a model the discovery step didn't find.""" + + +class DiscoveredModel(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + mode: str = "chat" + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + + +class _RawModelGroup(BaseModel): + model_config = ConfigDict(extra="ignore") + + model_group: str + # Optional: some real deployments return an explicit `"mode": null` for models that + # were registered without a mode (seen for embedding models like voyage-4-large). + # ModelGroupInfo's own "chat" default (litellm/types/router.py) only applies when the + # key is missing entirely, not when it's present as null, so this must tolerate None. + mode: str | None = "chat" + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + + +_RAW_MODEL_GROUPS_ADAPTER = TypeAdapter(list[_RawModelGroup]) + + +def parse_discovered_models(raw: list[JsonValue]) -> tuple[DiscoveredModel, ...]: + """Validate a raw `/model_group/info` response into typed models.""" + parsed = _RAW_MODEL_GROUPS_ADAPTER.validate_python(raw) + return tuple( + DiscoveredModel( + name=group.model_group, + # A null mode means the server genuinely doesn't know what this model does; + # "unknown" (rather than guessing "chat") keeps it out of both chat_models() + # and embedding_models() instead of risking a wrong-mode deployment. + mode=group.mode or "unknown", + input_cost_per_token=group.input_cost_per_token, + output_cost_per_token=group.output_cost_per_token, + ) + for group in parsed + ) + + +def chat_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]: + return tuple(m for m in models if m.mode == "chat") + + +def embedding_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]: + return tuple(m for m in models if m.mode == "embedding") + + +class HeuristicClassifier(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["heuristic"] = "heuristic" + + +class LLMClassifier(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["llm"] = "llm" + model: str + timeout_ms: int = 3000 + + +ClassifierChoice = Union[HeuristicClassifier, LLMClassifier] + + +class NoSemanticMatching(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["none"] = "none" + + +class SemanticMatching(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["semantic"] = "semantic" + embedding_model: str + match_threshold: float = 0.5 + + +SemanticMatchingChoice = Union[NoSemanticMatching, SemanticMatching] + +# Satisfies complexity_router's "semantic matching requires non-empty keyword_tier_rules" +# invariant with a sane starting point; the generated config.yaml can be hand-edited afterward. +_DEFAULT_KEYWORD_TIER_RULES: tuple[dict[str, JsonValue], ...] = ( + {"keywords": ["hi", "hello", "thanks"], "tier": "SIMPLE"}, + {"keywords": ["explain", "how does"], "tier": "MEDIUM"}, + {"keywords": ["refactor", "implement", "debug"], "tier": "COMPLEX"}, + {"keywords": ["step by step", "think through", "prove"], "tier": "REASONING"}, +) + + +class AutorouteConfig(BaseModel): + model_config = ConfigDict(frozen=True) + + base_url: str + api_key: str + # Each tier maps to a pool of one or more models; complexity_router picks randomly among + # them per request (or, in adaptive mode, learns which to prefer within the pool). + tiers: dict[str, tuple[str, ...]] + default_model: str + classifier: ClassifierChoice = Field(default_factory=HeuristicClassifier) + semantic_matching: SemanticMatchingChoice = Field(default_factory=NoSemanticMatching) + adaptive: bool = False + + +def validate_config(config: AutorouteConfig, discovered: tuple[DiscoveredModel, ...]) -> None: + """Raise ConfigGenerationError if config references a model discovery didn't return.""" + chat_names: frozenset[str] = frozenset(m.name for m in chat_models(discovered)) + embedding_names: frozenset[str] = frozenset(m.name for m in embedding_models(discovered)) + + for tier, models in config.tiers.items(): + for model in models: + if model not in chat_names: + raise ConfigGenerationError(f"Tier {tier} references unknown chat model '{model}'") + + if config.default_model not in chat_names: + raise ConfigGenerationError(f"default_model '{config.default_model}' is not a known chat model") + + if isinstance(config.classifier, LLMClassifier) and config.classifier.model not in chat_names: + raise ConfigGenerationError(f"classifier model '{config.classifier.model}' is not a known chat model") + + if ( + isinstance(config.semantic_matching, SemanticMatching) + and config.semantic_matching.embedding_model not in embedding_names + ): + raise ConfigGenerationError( + f"embedding model '{config.semantic_matching.embedding_model}' is not a known embedding model" + ) + + +def _litellm_proxy_deployment(name: str, base_url: str, api_key: str) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"litellm_proxy/{name}", + "api_base": base_url, + "api_key": api_key, + }, + } + + +def build_generated_model_list(config: AutorouteConfig) -> list[JsonValue]: + """Build the model_list for the ephemeral proxy's config.yaml. + + Every real model referenced anywhere (tier targets, classifier, embedding) is deduplicated + to exactly one `litellm_proxy/` deployment forwarding to the customer's real proxy, + plus one `auto_router/complexity_router` deployment tying the tiers together. + """ + referenced_names = {model for models in config.tiers.values() for model in models} + referenced_names.add(config.default_model) + if isinstance(config.classifier, LLMClassifier): + referenced_names.add(config.classifier.model) + if isinstance(config.semantic_matching, SemanticMatching): + referenced_names.add(config.semantic_matching.embedding_model) + + proxy_deployments = [ + _litellm_proxy_deployment(name, config.base_url, config.api_key) for name in sorted(referenced_names) + ] + + complexity_router_config: dict[str, JsonValue] = { + "tiers": {tier: list(models) for tier, models in config.tiers.items()}, + "default_model": config.default_model, + } + if isinstance(config.classifier, LLMClassifier): + complexity_router_config["classifier_type"] = "llm" + complexity_router_config["classifier_llm_config"] = { + "model": config.classifier.model, + "timeout_ms": config.classifier.timeout_ms, + } + if isinstance(config.semantic_matching, SemanticMatching): + complexity_router_config["semantic_keyword_matching"] = True + complexity_router_config["embedding_model"] = config.semantic_matching.embedding_model + complexity_router_config["match_threshold"] = config.semantic_matching.match_threshold + complexity_router_config["keyword_tier_rules"] = list(_DEFAULT_KEYWORD_TIER_RULES) + if config.adaptive: + complexity_router_config["adaptive"] = True + + auto_router_litellm_params: dict[str, JsonValue] = { + "model": "auto_router/complexity_router", + "complexity_router_config": complexity_router_config, + } + # A bare "*" model_name looks like the obvious way to catch every request Claude Code + # might send regardless of which model it thinks it's using, but Router's auto-router + # registry is keyed by the literal requested model string (router.py:10711-10717), not + # resolved through pattern/wildcard matching first -- so a "*" entry here would only ever + # match a client that literally sends model="*", never an actual wildcard catch-all. Callers + # instead need to make Claude Code request this "autorouter" name directly (see + # ANTHROPIC_DEFAULT_*_MODEL in settings.py's merge_claude_settings_static_token). + return [ + *proxy_deployments, + {"model_name": AUTOROUTER_MODEL_NAME, "litellm_params": auto_router_litellm_params}, + ] + + +def build_generated_proxy_config(config: AutorouteConfig, master_key: str) -> dict[str, JsonValue]: + """Full config.yaml content for the ephemeral proxy, including its own auth key. + + master_key must live under general_settings, not litellm_settings -- the proxy server + only ever reads general_settings.master_key (proxy_server.py:4530) to authenticate + requests; a key placed under litellm_settings is silently ignored, leaving the proxy + with no real auth at all. + """ + return { + "model_list": build_generated_model_list(config), + "general_settings": {"master_key": master_key}, + } + + +__all__ = [ + "AUTOROUTER_MODEL_NAME", + "TIER_NAMES", + "AutorouteConfig", + "ClassifierChoice", + "ConfigGenerationError", + "DiscoveredModel", + "HeuristicClassifier", + "LLMClassifier", + "NoSemanticMatching", + "SemanticMatching", + "SemanticMatchingChoice", + "build_generated_model_list", + "build_generated_proxy_config", + "chat_models", + "embedding_models", + "parse_discovered_models", + "validate_config", +] diff --git a/litellm/proxy/client/cli/commands/autoroute/process.py b/litellm/proxy/client/cli/commands/autoroute/process.py new file mode 100644 index 00000000000..86fa584a223 --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/process.py @@ -0,0 +1,174 @@ +import contextlib +import json +import os +import signal +import socket +import subprocess +import sys +import threading +import time +from dataclasses import dataclass +from pathlib import Path + +import click +import requests +from pydantic import TypeAdapter, ValidationError + +from ..up import UpError, secure_create + +AUTOROUTE_DIR = Path.home() / ".litellm" / "autorouter" +CONFIG_PATH = AUTOROUTE_DIR / "config.yaml" +LOG_PATH = AUTOROUTE_DIR / "proxy.log" +PID_RECORD_PATH = AUTOROUTE_DIR / "proxy.pid.json" + + +class ProcessLaunchError(Exception): + """Raised when the ephemeral proxy subprocess fails to come up healthy.""" + + +@dataclass(frozen=True, slots=True) +class PidRecord: + pid: int + port: int + config_path: str + log_path: str + + +_PID_RECORD_ADAPTER = TypeAdapter(PidRecord) + + +def allocate_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def launch_proxy(config_path: Path, port: int, log_path: Path) -> "subprocess.Popen[bytes]": + log_path.parent.mkdir(parents=True, exist_ok=True) + with open(log_path, "w") as log_file: + return subprocess.Popen( + [ + sys.executable, + "-m", + "litellm.proxy.proxy_cli", + "--config", + str(config_path), + "--port", + str(port), + "--host", + "127.0.0.1", + ], + stdout=log_file, + stderr=subprocess.STDOUT, + ) + + +def _tail(log_path: Path, lines: int = 40) -> str: + if not log_path.exists(): + return "(no log output captured)" + return "\n".join(log_path.read_text(errors="replace").splitlines()[-lines:]) + + +def poll_liveliness(base_url: str, log_path: Path, process: "subprocess.Popen[bytes]", timeout: float = 30.0) -> None: + """Poll /health/liveliness until it responds, the process dies, or timeout elapses.""" + deadline = time.monotonic() + timeout + url = base_url.rstrip("/") + "/health/liveliness" + while time.monotonic() < deadline: + if process.poll() is not None: + raise ProcessLaunchError( + f"Ephemeral proxy exited early (code {process.returncode}). Last log lines:\n{_tail(log_path)}" + ) + with contextlib.suppress(requests.RequestException): + if requests.get(url, timeout=2).status_code == 200: + return + time.sleep(0.5) + raise ProcessLaunchError( + f"Ephemeral proxy never became healthy within {timeout}s. Last log lines:\n{_tail(log_path)}" + ) + + +def write_pid_record(record: PidRecord, path: Path | None = None) -> None: + resolved_path = path if path is not None else PID_RECORD_PATH + resolved_path.parent.mkdir(parents=True, exist_ok=True) + with open(resolved_path, "w") as f: + json.dump( + {"pid": record.pid, "port": record.port, "config_path": record.config_path, "log_path": record.log_path}, + f, + indent=2, + ) + + +def read_pid_record(path: Path | None = None) -> PidRecord | None: + resolved_path = path if path is not None else PID_RECORD_PATH + if not resolved_path.exists(): + return None + with open(resolved_path, "r") as f: + content = f.read() + try: + return _PID_RECORD_ADAPTER.validate_json(content) + except ValidationError: + raise UpError(f"{resolved_path} contains invalid or unexpected JSON; cannot proceed safely.") + + +def clear_pid_record(path: Path | None = None) -> None: + resolved_path = path if path is not None else PID_RECORD_PATH + resolved_path.unlink(missing_ok=True) + + +def is_running(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def terminate(pid: int, grace_period: float = 5.0) -> None: + """Terminate a process by pid, escalating from SIGTERM to SIGKILL if needed.""" + if not is_running(pid): + return + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGTERM) + deadline = time.monotonic() + grace_period + while time.monotonic() < deadline and is_running(pid): + time.sleep(0.2) + if is_running(pid): + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + + +def stream_log(log_path: Path, stop_event: threading.Event) -> None: + """Print new lines appended to log_path until stop_event is set. Blocks the calling thread.""" + while not log_path.exists() and not stop_event.is_set(): + time.sleep(0.1) + if stop_event.is_set() or not log_path.exists(): + return + with open(log_path, "r") as f: + while not stop_event.is_set(): + line = f.readline() + if line: + click.echo(line, nl=False) + else: + time.sleep(0.2) + + +__all__ = [ + "AUTOROUTE_DIR", + "CONFIG_PATH", + "LOG_PATH", + "PID_RECORD_PATH", + "PidRecord", + "ProcessLaunchError", + "allocate_free_port", + "clear_pid_record", + "is_running", + "launch_proxy", + "poll_liveliness", + "read_pid_record", + "secure_create", + "stream_log", + "terminate", + "write_pid_record", +] diff --git a/litellm/proxy/client/cli/commands/autoroute/settings.py b/litellm/proxy/client/cli/commands/autoroute/settings.py new file mode 100644 index 00000000000..4bed184eb34 --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/settings.py @@ -0,0 +1,46 @@ +from pydantic import JsonValue + +from .config import AUTOROUTER_MODEL_NAME + +ENV_KEY = "env" +API_KEY_HELPER_KEY = "apiKeyHelper" +ANTHROPIC_API_KEY_KEY = "ANTHROPIC_API_KEY" +ANTHROPIC_AUTH_TOKEN_KEY = "ANTHROPIC_AUTH_TOKEN" +ANTHROPIC_BASE_URL_KEY = "ANTHROPIC_BASE_URL" +# Force every one of Claude Code's own model tiers to request the auto-router by name. +# Router's auto-router registry is keyed by the literal requested model string +# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*" +# model_name can never work as a catch-all -- these overrides are what actually makes +# Claude Code send "autorouter" regardless of /model or its own version-specific defaults. +ANTHROPIC_DEFAULT_MODEL_ENV_KEYS = ( + "ANTHROPIC_DEFAULT_SONNET_MODEL", + "ANTHROPIC_DEFAULT_HAIKU_MODEL", + "ANTHROPIC_DEFAULT_OPUS_MODEL", +) + + +def merge_claude_settings_static_token( + settings: dict[str, JsonValue], base_url: str, auth_token: str +) -> dict[str, JsonValue]: + """Return a new settings dict wired to a local ephemeral proxy with a static token. + + Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real + remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key was just + minted for this session, so a plain env var is simpler and correct. Any existing + apiKeyHelper is cleared so it can't fight with the static token. + """ + raw_env = settings.get(ENV_KEY, {}) + base_env = raw_env if isinstance(raw_env, dict) else {} + env: dict[str, JsonValue] = { + **base_env, + ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), + ANTHROPIC_AUTH_TOKEN_KEY: auth_token, + **{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS}, + } + env.pop(ANTHROPIC_API_KEY_KEY, None) + merged: dict[str, JsonValue] = {**settings, ENV_KEY: env} + merged.pop(API_KEY_HELPER_KEY, None) + return merged + + +__all__ = ["merge_claude_settings_static_token"] diff --git a/litellm/proxy/client/cli/commands/autoroute/wizard.py b/litellm/proxy/client/cli/commands/autoroute/wizard.py new file mode 100644 index 00000000000..89b53d8f9fd --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/wizard.py @@ -0,0 +1,128 @@ +import sys +from pathlib import Path + +import click +import yaml +from InquirerPy import inquirer +from InquirerPy.base.control import Choice + +from .... import Client +from .config import ( + TIER_NAMES, + AutorouteConfig, + ConfigGenerationError, + DiscoveredModel, + HeuristicClassifier, + LLMClassifier, + NoSemanticMatching, + SemanticMatching, + build_generated_model_list, + chat_models, + embedding_models, + parse_discovered_models, + validate_config, +) +from .process import CONFIG_PATH, secure_create + + +def _is_interactive() -> bool: + return sys.stdin.isatty() + + +def _fuzzy_pick(models: tuple[DiscoveredModel, ...], prompt_label: str, multiselect: bool) -> list[str]: + """Type-to-filter picker over a (possibly huge) model pool, using InquirerPy's fzf-style fuzzy prompt. + + A plain numbered table + typed index does not scale past a handful of models -- proxies with + hundreds of model groups made that interaction unusable. This lets the user narrow the pool by + typing a substring instead of scrolling/counting. + + Assumes the caller already checked interactivity (run_configure_wizard does, once, up front) -- + checking here too would check the wrong thing under test, where InquirerPy is driven through its + own injected input/output rather than the real process stdin. + """ + choices = [Choice(value=model.name, name=model.name) for model in models] + toggle_hint = "tab to toggle, " if multiselect else "" + while True: + result = inquirer.fuzzy( + message=f"{prompt_label}: type to filter, {toggle_hint}enter to confirm", + choices=choices, + multiselect=multiselect, + max_height="70%", + ).execute() + selected = result if multiselect else [result] + if selected: + return selected + click.echo("Select at least one model.") + + +def _render_and_prompt_for_model(models: tuple[DiscoveredModel, ...], prompt_label: str) -> str: + return _fuzzy_pick(models, prompt_label, multiselect=False)[0] + + +def _render_and_prompt_for_models(models: tuple[DiscoveredModel, ...], prompt_label: str) -> tuple[str, ...]: + return tuple(_fuzzy_pick(models, prompt_label, multiselect=True)) + + +def run_configure_wizard(ctx: click.Context) -> Path: + """Discover the caller's accessible models, walk them through tier assignment, write config.""" + base_url = ctx.obj["base_url"] + api_key = ctx.obj["api_key"] + client = Client(base_url=base_url, api_key=api_key) + + raw_groups = client.model_groups.info() + if not isinstance(raw_groups, list): + raise click.ClickException( + f"Unexpected response from /model_group/info: expected a list, got {type(raw_groups).__name__}" + ) + discovered = parse_discovered_models(raw_groups) + chat_pool = chat_models(discovered) + embedding_pool = embedding_models(discovered) + + if not chat_pool: + raise click.ClickException("Your key has no chat-capable models available on this proxy.") + + if not _is_interactive(): + raise click.ClickException("`lite autoroute configure` requires an interactive terminal.") + + click.echo("Assign model(s) to each complexity tier (from what your key can access):") + tiers = {tier: _render_and_prompt_for_models(chat_pool, tier) for tier in TIER_NAMES} + default_model = tiers["MEDIUM"][0] + + classifier = HeuristicClassifier() + if click.confirm("\nUse an LLM classifier instead of the free heuristic scorer?", default=False): + classifier_model = _render_and_prompt_for_model(chat_pool, "LLM classifier") + classifier = LLMClassifier(model=classifier_model) + + semantic_matching = NoSemanticMatching() + if embedding_pool and click.confirm("\nEnable semantic keyword matching?", default=False): + embedding_model = _render_and_prompt_for_model(embedding_pool, "semantic embeddings") + semantic_matching = SemanticMatching(embedding_model=embedding_model) + + adaptive = click.confirm("\nEnable adaptive (bandit) selection on top of tiering?", default=False) + + config = AutorouteConfig( + base_url=base_url, + api_key=api_key, + tiers=tiers, + default_model=default_model, + classifier=classifier, + semantic_matching=semantic_matching, + adaptive=adaptive, + ) + try: + validate_config(config, discovered) + except ConfigGenerationError as e: + raise click.ClickException(str(e)) + + model_list = build_generated_model_list(config) + CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True) + with secure_create(CONFIG_PATH) as f: + yaml.safe_dump({"model_list": model_list}, f, sort_keys=False) + + click.echo(f"\nWrote {CONFIG_PATH}") + for tier, models in tiers.items(): + click.echo(f" {tier}: {', '.join(models)}") + return CONFIG_PATH + + +__all__ = ["run_configure_wizard"] diff --git a/litellm/proxy/client/cli/commands/model_groups.py b/litellm/proxy/client/cli/commands/model_groups.py new file mode 100644 index 00000000000..7de959a9b78 --- /dev/null +++ b/litellm/proxy/client/cli/commands/model_groups.py @@ -0,0 +1,57 @@ +from typing import Literal + +import click +import rich +import rich.table + +from ... import Client + + +def create_client(ctx: click.Context) -> Client: + return Client(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + + +@click.group(name="model-groups") +def model_groups() -> None: + """Inspect model groups your key can access on the proxy""" + + +@model_groups.command("list") +@click.option( + "--format", + "output_format", + type=click.Choice(["table", "json"]), + default="table", + help="Output format (table or json)", +) +@click.pass_context +def list_model_groups(ctx: click.Context, output_format: Literal["table", "json"]) -> None: + """List model groups accessible to your key, with mode and pricing""" + client = create_client(ctx) + groups = client.model_groups.info() + if not isinstance(groups, list): + raise click.ClickException( + f"Unexpected response from /model_group/info: expected a list, got {type(groups).__name__}" + ) + + if output_format == "json": + rich.print_json(data=groups) + return + + table = rich.table.Table(title="Accessible Model Groups") + table.add_column("Model", style="cyan") + table.add_column("Mode", style="green") + table.add_column("Input $/token", style="yellow") + table.add_column("Output $/token", style="yellow") + + for group in groups: + table.add_row( + str(group.get("model_group", "")), + str(group.get("mode", "chat")), + str(group.get("input_cost_per_token", "")), + str(group.get("output_cost_per_token", "")), + ) + rich.print(table) + + +__all__ = ["model_groups"] diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py new file mode 100644 index 00000000000..dc9157d7ca4 --- /dev/null +++ b/litellm/proxy/client/cli/commands/up.py @@ -0,0 +1,283 @@ +import atexit +import contextlib +import json +import os +import shlex +import shutil +import signal +import sys +import threading +from dataclasses import dataclass +from pathlib import Path +from types import FrameType +from typing import IO, Iterator, Mapping + +import click +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh + +from .agents import AgentRunError, resolve_api_key, verify_proxy_key +from .auth import load_token, login + +ENV_KEY = "env" +API_KEY_HELPER_KEY = "apiKeyHelper" +ANTHROPIC_BASE_URL_KEY = "ANTHROPIC_BASE_URL" +ANTHROPIC_API_KEY_KEY = "ANTHROPIC_API_KEY" + +CLAUDE_SETTINGS_PATH = Path.home() / ".claude" / "settings.json" +BACKUP_PATH = Path.home() / ".litellm" / "claude_settings_backup.json" + + +class UpError(Exception): + """Raised for any user-actionable failure while starting/stopping interception.""" + + +@dataclass(frozen=True, slots=True) +class BackupRecord: + """Snapshot of ~/.claude/settings.json taken right before `lite up` patches it.""" + + existed: bool + content: dict[str, JsonValue] | None + + +_SETTINGS_ADAPTER = TypeAdapter(dict[str, JsonValue]) +_BACKUP_RECORD_ADAPTER = TypeAdapter(BackupRecord) + + +def load_json_or_empty(path: Path) -> dict[str, JsonValue]: + if not path.exists(): + return {} + with open(path, "r") as f: + content = f.read() + if not content.strip(): + return {} + try: + return _SETTINGS_ADAPTER.validate_json(content) + except ValidationError: + raise UpError(f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely.") + + +def merge_claude_settings( + settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str +) -> dict[str, JsonValue]: + """Return a new settings dict wired to route Claude Code through the proxy. + + Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a + stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued + token (same reasoning as build_agent_env in agents.py). Every other key is + preserved untouched. + """ + raw_env = settings.get(ENV_KEY, {}) + base_env = raw_env if isinstance(raw_env, dict) else {} + env = {**base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/")} + env.pop(ANTHROPIC_API_KEY_KEY, None) + return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} + + +@contextlib.contextmanager +def secure_create(path: Path) -> Iterator[IO[str]]: + """Open path for writing with mode 0600 fixed up before any content is written. + + A plain `open(path, "w")` creates a *new* file at the umask-derived default (commonly 0644) + and leaves it world- or group-readable until a later `chmod` call catches up -- a real window + in which a file holding a credential is readable by another local account. Passing the mode to + `os.open` closes that window for a brand-new file, but `O_CREAT`'s mode argument is only + applied on creation: if the file already exists its old, broader permissions carry over + untouched. `os.fchmod` right after opening -- before a single byte of the new content is + written -- covers both cases. + """ + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + os.fchmod(fd, 0o600) + f: IO[str] = os.fdopen(fd, "w") + try: + yield f + finally: + f.close() + + +def write_backup(record: BackupRecord, backup_path: Path | None = None) -> None: + path = backup_path if backup_path is not None else BACKUP_PATH + path.parent.mkdir(exist_ok=True) + with secure_create(path) as f: + json.dump({"existed": record.existed, "content": record.content}, f, indent=2) + + +def read_backup(backup_path: Path | None = None) -> BackupRecord | None: + path = backup_path if backup_path is not None else BACKUP_PATH + if not path.exists(): + return None + with open(path, "r") as f: + content = f.read() + try: + return _BACKUP_RECORD_ADAPTER.validate_json(content) + except ValidationError: + raise UpError(f"{path} contains invalid or unexpected JSON; cannot restore from it safely.") + + +def restore_claude_settings(settings_path: Path | None = None, backup_path: Path | None = None) -> BackupRecord | None: + """Restore settings_path from the backup at backup_path, then delete the backup. + + Returns the restored record, or None if there was nothing to restore. + """ + resolved_settings_path = settings_path if settings_path is not None else CLAUDE_SETTINGS_PATH + resolved_backup_path = backup_path if backup_path is not None else BACKUP_PATH + record = read_backup(resolved_backup_path) + if record is None: + return None + if record.existed and record.content is not None: + resolved_settings_path.parent.mkdir(parents=True, exist_ok=True) + with open(resolved_settings_path, "w") as f: + json.dump(record.content, f, indent=2) + elif resolved_settings_path.exists(): + resolved_settings_path.unlink() + resolved_backup_path.unlink() + return record + + +def resolve_api_key_helper(base_url: str) -> str: + """Build the shell command Claude Code should run for its apiKeyHelper. + + Resolves `lite` to an absolute path so the helper works regardless of the + PATH visible to whatever subprocess Claude Code spawns it from. Passing + --base-url explicitly (rather than relying on the bare invocation Claude + Code would otherwise use) makes `print-token` enforce that the cached + token was actually issued for this proxy -- without it, a token minted + for a different, previously-logged-into proxy would be handed to + whichever server `up` currently points at. + """ + lite_path = shutil.which("lite") + if lite_path is None: + raise UpError( + "Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs " + "an absolute path to it, so `lite up` cannot continue." + ) + return f"{shlex.quote(lite_path)} auth print-token --base-url {shlex.quote(base_url)}" + + +def _ensure_fresh_login(ctx: click.Context) -> None: + base_url = ctx.obj["base_url"].rstrip("/") + token_data = load_token() + if token_data and token_data.get("base_url") == base_url and is_cli_token_fresh(token_data): + return + + if not sys.stdin.isatty(): + raise UpError( + "No fresh LiteLLM login found for this proxy. Run `lite login` first (apiKeyHelper " + "reads this token on every Claude Code request)." + ) + + click.echo("No fresh LiteLLM login found for this proxy; starting login...") + ctx.invoke(login) + token_data = load_token() + if not token_data or token_data.get("base_url") != base_url or not is_cli_token_fresh(token_data): + raise UpError("Login did not produce a usable token; cannot start `lite up`.") + + +def _restore_and_report() -> None: + record = restore_claude_settings() + if record is None: + click.echo("Nothing to restore.") + return + if record.existed: + click.echo(f"Restored {CLAUDE_SETTINGS_PATH} to its original contents.") + else: + click.echo(f"Removed {CLAUDE_SETTINGS_PATH} (it did not exist before `lite up`).") + + +@click.command(name="up") +@click.pass_context +def up(ctx: click.Context) -> None: + """Route every Claude Code session through your LiteLLM proxy until stopped. + + Patches ~/.claude/settings.json so Claude Code picks up the proxy on its own + next startup, from any terminal -- no need to launch it through `lite`. + Press Ctrl-C to stop and restore your original settings. Assumes the proxy + is already running (this does not start one for you). Cursor is not + supported: it has no equivalent file-based config to patch. + """ + base_url = ctx.obj["base_url"] + + try: + _ensure_fresh_login(ctx) + api_key = resolve_api_key(ctx) + verify_proxy_key(base_url, api_key) + + if BACKUP_PATH.exists(): + raise UpError( + f"{BACKUP_PATH} already exists -- `lite up` looks like it's already " + "running (or crashed without cleanup). Run `lite down` first." + ) + + api_key_helper = resolve_api_key_helper(base_url) + original_existed = CLAUDE_SETTINGS_PATH.exists() + original_settings = load_json_or_empty(CLAUDE_SETTINGS_PATH) + write_backup( + BackupRecord( + existed=original_existed, + content=original_settings if original_existed else None, + ) + ) + + CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True) + merged = merge_claude_settings(original_settings, base_url, api_key_helper) + with open(CLAUDE_SETTINGS_PATH, "w") as f: + json.dump(merged, f, indent=2) + except (AgentRunError, UpError) as e: + raise click.ClickException(str(e)) + + click.echo(f"litellm: routing Claude Code through proxy at {base_url.rstrip('/')}") + click.echo("Press Ctrl-C to stop and restore your original settings.") + + stop_event = threading.Event() + restored = threading.Lock() + + def _handle_signal(_signum: int, _frame: FrameType | None) -> None: + stop_event.set() + + def _restore_once() -> None: + if not restored.acquire(blocking=False): + return + try: + _restore_and_report() + except UpError as e: + # Runs from atexit/a signal handler, outside Click's own exception + # handling -- raising here would only produce an unhandled-exception + # warning on stderr, not a clean message. + click.echo(str(e), err=True) + + signal.signal(signal.SIGINT, _handle_signal) + signal.signal(signal.SIGTERM, _handle_signal) + atexit.register(_restore_once) + + stop_event.wait() + _restore_once() + + +@click.command(name="down") +def down() -> None: + """Restore ~/.claude/settings.json if a `lite up` session left it patched. + + Use this after a `lite up` process was killed uncleanly (e.g. `kill -9`) + instead of stopped with Ctrl-C. + """ + try: + _restore_and_report() + except UpError as e: + raise click.ClickException(str(e)) + + +__all__ = [ + "BACKUP_PATH", + "CLAUDE_SETTINGS_PATH", + "BackupRecord", + "UpError", + "down", + "load_json_or_empty", + "merge_claude_settings", + "read_backup", + "resolve_api_key_helper", + "restore_claude_settings", + "up", + "write_backup", +] diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 4de3ff5fc87..e641956b2c5 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -9,15 +9,18 @@ from litellm.proxy.client.health import HealthManagementClient from .commands.agents import agent_commands from .commands.auth import auth_group, get_stored_api_key, login, logout, whoami +from .commands.autoroute.commands import autoroute_group from .commands.chat import chat from .commands.credentials import credentials from .commands.encryption import encryption from .commands.http import http from .commands.keys import keys +from .commands.model_groups import model_groups # local imports from .commands.models import models from .commands.teams import teams +from .commands.up import down, up from .commands.users import users from .interface import interactive_shell @@ -131,6 +134,13 @@ cli.add_command(users) # Add a top-level command per coding agent (claude, codex, opencode, ...) for agent_command in agent_commands(): cli.add_command(agent_command) +# Add the up/down commands (route Claude Code through the local LiteLLM proxy) +cli.add_command(up) +cli.add_command(down) +# Add the model-groups command group (discover models your key can access) +cli.add_command(model_groups) +# Add the autoroute command group (QA auto-routing against your real proxy) +cli.add_command(autoroute_group, name="autoroute") if __name__ == "__main__": diff --git a/pyproject.toml b/pyproject.toml index d70e99c5775..2c19bf64b4f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,6 +66,7 @@ proxy = [ "litellm-enterprise==0.1.50", "RestrictedPython>=8.1,<9.0", "rich>=13.9.4,<14.0", + "InquirerPy>=0.3.4,<1.0", "polars>=1.38.1,<2.0", "soundfile>=0.12.1,<1.0", "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", @@ -74,11 +75,12 @@ proxy = [ ] # Thin client install for the `lite` CLI on developer laptops. The CLI's heavy # imports (fastapi, cryptography, ...) are all guarded, so it runs on the base -# SDK plus just these three; none of the server runtime in `proxy` is pulled in. +# SDK plus just these four; none of the server runtime in `proxy` is pulled in. cli = [ "rich>=13.9.4,<14.0", "pyyaml>=6.0.3,<7.0", "requests>=2.32.0,<3.0", + "InquirerPy>=0.3.4,<1.0", ] extra_proxy = [ "prisma>=0.11.0,<1.0", diff --git a/scripts/install-cli.sh b/scripts/install-cli.sh index d147286fcac..a39b73c2e5a 100755 --- a/scripts/install-cli.sh +++ b/scripts/install-cli.sh @@ -11,12 +11,21 @@ # Python itself (honouring litellm's requires-python), downloading a managed one # when the host has no suitable interpreter. # +# To try an unreleased branch instead of the latest PyPI release (for example, to +# QA a CLI feature before it ships), set LITELLM_CLI_REF to a branch, tag, or commit: +# curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm//scripts/install-cli.sh | \ +# LITELLM_CLI_REF= sh +# # NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian # ignores the shebang when invoked as `sh` and does not support `pipefail`). set -eu -# NOTE: before merging, this must stay as "litellm[cli]" to install from PyPI. -LITELLM_PACKAGE="litellm[cli]" +# Defaults to the PyPI release; LITELLM_CLI_REF opts into installing from source instead. +if [ -n "${LITELLM_CLI_REF:-}" ]; then + LITELLM_PACKAGE="litellm[cli] @ git+https://github.com/BerriAI/litellm.git@${LITELLM_CLI_REF}" +else + LITELLM_PACKAGE="litellm[cli]" +fi UV_VERSION="0.10.9" # ── colours ──────────────────────────────────────────────────────────────── @@ -90,7 +99,11 @@ fi # otherwise download a managed one. Either way uv honours litellm's requires-python, # so a too-old (3.9) or too-new (3.14+) system Python is skipped, not forced. echo "" -header "Installing litellm[cli]…" +if [ -n "${LITELLM_CLI_REF:-}" ]; then + header "Installing litellm[cli] from ${LITELLM_CLI_REF}…" +else + header "Installing litellm[cli]…" +fi echo "" "$UV_BIN" tool install --python-preference system --force "${LITELLM_PACKAGE}" \ diff --git a/tests/test_litellm/proxy/client/cli/autoroute/__init__.py b/tests/test_litellm/proxy/client/cli/autoroute/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py new file mode 100644 index 00000000000..e0b5d3a71af --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py @@ -0,0 +1,300 @@ +import json +import stat +from typing import Optional + +import yaml +from click.testing import CliRunner + +from litellm.proxy.client.cli.commands.autoroute import commands as commands_module +from litellm.proxy.client.cli.commands.autoroute import process as process_module +from litellm.proxy.client.cli.commands.autoroute.commands import down, up +from litellm.proxy.client.cli.commands.autoroute.process import PidRecord, ProcessLaunchError, write_pid_record +from litellm.proxy.client.cli.commands.up import BackupRecord as ClaudeBackupRecord +from litellm.proxy.client.cli.commands.up import write_backup + + +class FakeProcess: + def __init__(self, pid: int): + self.pid = pid + self.returncode: Optional[int] = None + + def poll(self) -> Optional[int]: + return self.returncode + + +def _patch_paths(monkeypatch, tmp_path): + config_path = tmp_path / "config.yaml" + log_path = tmp_path / "proxy.log" + claude_settings_path = tmp_path / "claude_settings.json" + backup_path = tmp_path / "backup.json" + pid_record_path = tmp_path / "pid.json" + + monkeypatch.setattr(commands_module, "CONFIG_PATH", config_path) + monkeypatch.setattr(commands_module, "LOG_PATH", log_path) + monkeypatch.setattr(commands_module, "CLAUDE_SETTINGS_PATH", claude_settings_path) + monkeypatch.setattr(commands_module, "AUTOROUTE_BACKUP_PATH", backup_path) + monkeypatch.setattr(process_module, "PID_RECORD_PATH", pid_record_path) + + return config_path, log_path, claude_settings_path, backup_path, pid_record_path + + +def _silence_signal_handling(monkeypatch): + monkeypatch.setattr(commands_module.signal, "signal", lambda *a, **k: None) + monkeypatch.setattr(commands_module.atexit, "register", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "stream_log", lambda *a, **k: None) + + +class TestUpCommand: + def setup_method(self): + self.runner = CliRunner() + + def test_refuses_when_never_configured(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "lite autoroute configure" in result.output + + def test_surfaces_clean_error_on_empty_config_file(self, monkeypatch, tmp_path): + """A `configure` killed between secure_create's O_TRUNC and the write completing leaves an + empty config.yaml on disk -- yaml.safe_load(empty) returns None, and validating None as the + generated-config model raises a raw pydantic.ValidationError if uncaught.""" + config_path, _log_path, _settings_path, _backup_path, _pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text("") + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert result.exception is None or isinstance(result.exception, SystemExit) + assert "lite autoroute configure" in result.output + + def test_refuses_when_pid_record_exists_and_process_still_running(self, monkeypatch, tmp_path): + config_path, _log_path, _settings_path, _backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + write_pid_record( + PidRecord(pid=123, port=4000, config_path=str(config_path), log_path="/tmp/proxy.log"), pid_record_path + ) + monkeypatch.setattr(commands_module, "is_running", lambda pid: True) + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "already running" in result.output + assert "lite autoroute down" in result.output + assert config_path.read_text() == yaml.safe_dump({"model_list": []}) + + def test_refuses_when_backup_exists_after_an_unclean_crash(self, monkeypatch, tmp_path): + """A prior `up` that was SIGKILL'd leaves no live pid but does leave a stale backup file. + + Without this guard, a fresh `up` would overwrite that backup with the currently-patched + (not original) Claude settings, so `down`/Ctrl-C would restore the wrong content forever. + """ + config_path, _log_path, claude_settings_path, backup_path, _pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + config_path.write_text(yaml.safe_dump({"model_list": []})) + claude_settings_path.write_text(json.dumps({"env": {"ANTHROPIC_AUTH_TOKEN": "stale-patched-token"}})) + write_backup(ClaudeBackupRecord(existed=True, content={"theme": "dark"}), backup_path) + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "already exists" in result.output + assert "lite autoroute down" in result.output + assert json.loads(backup_path.read_text())["content"] == {"theme": "dark"} + + def test_happy_path_patches_settings_then_restores_everything_on_stop(self, monkeypatch, tmp_path): + config_path, log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + original_settings = {"theme": "dark"} + claude_settings_path.write_text(json.dumps(original_settings)) + _silence_signal_handling(monkeypatch) + + fake_process = FakeProcess(pid=99999) + terminate_calls = [] + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "allocate_free_port", lambda: 54321) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + captured = {} + + def fake_wait(self, timeout=None): + captured["settings"] = json.loads(claude_settings_path.read_text()) + captured["backup_existed"] = backup_path.exists() + captured["settings_mode"] = stat.S_IMODE(claude_settings_path.stat().st_mode) + return True + + monkeypatch.setattr("threading.Event.wait", fake_wait) + + result = self.runner.invoke(up) + + assert result.exit_code == 0, result.output + assert captured["backup_existed"] is True + assert captured["settings"]["theme"] == "dark" + assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:54321" + assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "fixed-master-key" + assert "apiKeyHelper" not in captured["settings"] + assert captured["settings_mode"] == 0o600 + + assert terminate_calls == [99999] + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + written_config = yaml.safe_load(config_path.read_text()) + assert written_config["general_settings"]["master_key"] == "fixed-master-key" + assert stat.S_IMODE(config_path.stat().st_mode) == 0o600 + + def test_teardown_reports_clean_error_when_backup_is_corrupt(self, monkeypatch, tmp_path): + """A corrupt backup at teardown time (e.g. a concurrent process wrote garbage to it) must + not crash the whole command -- _restore_once in up.py handles the identical case in + lite up the same way, echoing the error instead of propagating it.""" + config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + claude_settings_path.write_text(json.dumps({"theme": "dark"})) + _silence_signal_handling(monkeypatch) + + fake_process = FakeProcess(pid=11111) + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "allocate_free_port", lambda: 65432) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: None) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + def fake_wait(self, timeout=None): + backup_path.write_text("not json at all {{{") + return True + + monkeypatch.setattr("threading.Event.wait", fake_wait) + + result = self.runner.invoke(up) + + assert result.exit_code == 0, result.output + assert "invalid or unexpected JSON" in result.output + assert not pid_record_path.exists() + + def test_surfaces_clean_error_and_cleans_up_when_health_check_fails(self, monkeypatch, tmp_path): + config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + original_settings = {"theme": "dark"} + claude_settings_path.write_text(json.dumps(original_settings)) + + fake_process = FakeProcess(pid=555) + terminate_calls = [] + + def _raise_launch_error(*args, **kwargs): + raise ProcessLaunchError("boom: proxy never became healthy") + + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", _raise_launch_error) + monkeypatch.setattr(commands_module, "allocate_free_port", lambda: 12345) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "boom" in result.output + assert terminate_calls == [555] + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + def test_terminates_ephemeral_proxy_when_claude_settings_is_corrupt(self, monkeypatch, tmp_path): + """The health check can pass and the proxy can come up fine, but if + ~/.claude/settings.json turns out to be corrupt, the just-started proxy must not be left + running with no pid record -- exactly the leak `lite autoroute down` exists to clean up.""" + config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + claude_settings_path.write_text("not json at all {{{") + + fake_process = FakeProcess(pid=777) + terminate_calls = [] + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "allocate_free_port", lambda: 23456) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "invalid JSON" in result.output + assert terminate_calls == [777] + assert not pid_record_path.exists() + assert not backup_path.exists() + + +class TestDownCommand: + def setup_method(self): + self.runner = CliRunner() + + def test_restores_settings_and_terminates_when_process_still_running(self, monkeypatch, tmp_path): + _config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + original_settings = {"theme": "dark"} + write_backup(ClaudeBackupRecord(existed=True, content=original_settings), backup_path) + claude_settings_path.write_text(json.dumps({"env": {"ANTHROPIC_AUTH_TOKEN": "fixed-master-key"}})) + write_pid_record(PidRecord(pid=777, port=1234, config_path="c", log_path="l"), pid_record_path) + + terminate_calls = [] + monkeypatch.setattr(commands_module, "is_running", lambda pid: True) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Stopped leftover ephemeral proxy" in result.output + assert "Restored" in result.output + assert terminate_calls == [777] + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + def test_is_a_clean_no_op_when_nothing_is_running_and_no_backup_exists(self, monkeypatch, tmp_path): + _config_path, _log_path, claude_settings_path, _backup_path, _pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Nothing to restore." in result.output + assert not claude_settings_path.exists() + + def test_clears_a_corrupt_pid_record_and_still_restores_settings(self, monkeypatch, tmp_path): + """down is specifically the crash-recovery path -- a pid file truncated by a mid-write + crash must not block it from clearing the record and restoring Claude settings anyway.""" + _config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + pid_record_path.parent.mkdir(parents=True, exist_ok=True) + pid_record_path.write_text("not json at all {{{") + original_settings = {"theme": "dark"} + write_backup(ClaudeBackupRecord(existed=True, content=original_settings), backup_path) + claude_settings_path.write_text(json.dumps({"env": {"ANTHROPIC_AUTH_TOKEN": "fixed-master-key"}})) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "invalid or unexpected JSON" in result.output + assert "Restored" in result.output + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + def test_surfaces_clean_error_when_backup_is_corrupt(self, monkeypatch, tmp_path): + _config_path, _log_path, _claude_settings_path, backup_path, _pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + backup_path.parent.mkdir(parents=True, exist_ok=True) + backup_path.write_text("not json at all {{{") + + result = self.runner.invoke(down) + + assert result.exit_code != 0 + assert "invalid or unexpected JSON" in result.output diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py new file mode 100644 index 00000000000..9fa01524ef3 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py @@ -0,0 +1,181 @@ +from typing import Any, Dict, Tuple + +import pytest + +from litellm.proxy.client.cli.commands.autoroute.config import ( + AutorouteConfig, + ConfigGenerationError, + DiscoveredModel, + HeuristicClassifier, + LLMClassifier, + NoSemanticMatching, + SemanticMatching, + build_generated_model_list, + build_generated_proxy_config, + chat_models, + embedding_models, + parse_discovered_models, + validate_config, +) + +DISCOVERED: Tuple[DiscoveredModel, ...] = ( + DiscoveredModel(name="gpt-4o-mini", mode="chat"), + DiscoveredModel(name="gpt-4o", mode="chat"), + DiscoveredModel(name="o1", mode="chat"), + DiscoveredModel(name="text-embedding-3-small", mode="embedding"), +) + + +def _base_config(**overrides: Any) -> AutorouteConfig: + defaults: Dict[str, Any] = { + "base_url": "http://real-proxy.internal:4000", + "api_key": "sk-real-key", + "tiers": { + "SIMPLE": ("gpt-4o-mini",), + "MEDIUM": ("gpt-4o",), + "COMPLEX": ("gpt-4o",), + "REASONING": ("o1",), + }, + "default_model": "gpt-4o", + } + defaults.update(overrides) + return AutorouteConfig(**defaults) + + +class TestParseDiscoveredModels: + def test_parses_valid_raw_list_into_typed_tuple(self): + raw = [ + { + "model_group": "gpt-4o", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + }, + {"model_group": "text-embedding-3-small", "mode": "embedding"}, + ] + result = parse_discovered_models(raw) + assert result == ( + DiscoveredModel(name="gpt-4o", mode="chat", input_cost_per_token=0.01, output_cost_per_token=0.02), + DiscoveredModel(name="text-embedding-3-small", mode="embedding"), + ) + + def test_ignores_unknown_extra_fields(self): + raw = [{"model_group": "gpt-4o", "mode": "chat", "totally_unknown_field": "whatever"}] + result = parse_discovered_models(raw) + assert result == (DiscoveredModel(name="gpt-4o", mode="chat"),) + + def test_missing_mode_defaults_to_chat(self): + raw = [{"model_group": "gpt-4o"}] + result = parse_discovered_models(raw) + assert result[0].mode == "chat" + + +class TestChatAndEmbeddingFiltering: + def test_filters_by_mode(self): + models = ( + DiscoveredModel(name="gpt-4o", mode="chat"), + DiscoveredModel(name="text-embedding-3-small", mode="embedding"), + DiscoveredModel(name="claude", mode="chat"), + ) + assert chat_models(models) == (models[0], models[2]) + assert embedding_models(models) == (models[1],) + + +class TestBuildGeneratedModelList: + def test_dedups_model_used_in_multiple_roles(self): + config = _base_config(classifier=LLMClassifier(model="gpt-4o")) + model_list = build_generated_model_list(config) + gpt4o_entries = [m for m in model_list if m["model_name"] == "gpt-4o"] + assert len(gpt4o_entries) == 1 + + def test_every_proxy_deployment_points_back_at_customer_proxy(self): + config = _base_config() + model_list = build_generated_model_list(config) + proxy_entries = [m for m in model_list if m["model_name"] not in ("autorouter", "*")] + names = {m["model_name"] for m in proxy_entries} + assert names == {"gpt-4o-mini", "gpt-4o", "o1"} + for entry in proxy_entries: + assert entry["litellm_params"]["model"] == f"litellm_proxy/{entry['model_name']}" + assert entry["litellm_params"]["api_base"] == config.base_url + assert entry["litellm_params"]["api_key"] == config.api_key + + def test_no_wildcard_deployment_is_generated(self): + # A bare "*" model_name looks like the obvious catch-all, but Router's auto-router + # registry is keyed by the literal requested model string with no wildcard resolution + # (litellm/router.py:10711-10717), so a "*" entry here would silently never match real + # traffic. Regression guard: don't reintroduce it. + config = _base_config() + model_list = build_generated_model_list(config) + assert not any(m["model_name"] == "*" for m in model_list) + + def test_complexity_router_config_reflects_llm_classifier(self): + config = _base_config(classifier=LLMClassifier(model="gpt-4o", timeout_ms=1234)) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert router_config["classifier_type"] == "llm" + assert router_config["classifier_llm_config"] == {"model": "gpt-4o", "timeout_ms": 1234} + assert "semantic_keyword_matching" not in router_config + assert "adaptive" not in router_config + + def test_complexity_router_config_reflects_semantic_matching(self): + config = _base_config( + semantic_matching=SemanticMatching(embedding_model="text-embedding-3-small", match_threshold=0.7) + ) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert router_config["semantic_keyword_matching"] is True + assert router_config["embedding_model"] == "text-embedding-3-small" + assert router_config["match_threshold"] == 0.7 + assert router_config["keyword_tier_rules"] + assert "classifier_type" not in router_config + + def test_complexity_router_config_reflects_adaptive(self): + config = _base_config(adaptive=True) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + assert autorouter["litellm_params"]["complexity_router_config"]["adaptive"] is True + + def test_default_classifier_and_semantic_matching_add_no_extra_keys(self): + config = _base_config(classifier=HeuristicClassifier(), semantic_matching=NoSemanticMatching()) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert set(router_config.keys()) == {"tiers", "default_model"} + + +class TestBuildGeneratedProxyConfig: + def test_embeds_master_key_under_general_settings(self): + config = _base_config() + proxy_config = build_generated_proxy_config(config, "sk-master-123") + assert proxy_config["general_settings"] == {"master_key": "sk-master-123"} + assert proxy_config["model_list"] == build_generated_model_list(config) + + +class TestValidateConfig: + def test_passes_for_fully_valid_config(self): + validate_config(_base_config(), DISCOVERED) + + def test_raises_for_tier_referencing_unknown_model(self): + config = _base_config( + tiers={ + "SIMPLE": ("unknown-model",), + "MEDIUM": ("gpt-4o",), + "COMPLEX": ("gpt-4o",), + "REASONING": ("o1",), + } + ) + with pytest.raises(ConfigGenerationError, match="unknown-model"): + validate_config(config, DISCOVERED) + + def test_raises_for_unknown_default_model(self): + config = _base_config(default_model="unknown-model") + with pytest.raises(ConfigGenerationError, match="unknown-model"): + validate_config(config, DISCOVERED) + + def test_raises_for_unknown_llm_classifier_model(self): + config = _base_config(classifier=LLMClassifier(model="unknown-model")) + with pytest.raises(ConfigGenerationError, match="unknown-model"): + validate_config(config, DISCOVERED) + + def test_raises_for_unknown_semantic_embedding_model(self): + config = _base_config(semantic_matching=SemanticMatching(embedding_model="unknown-embedding")) + with pytest.raises(ConfigGenerationError, match="unknown-embedding"): + validate_config(config, DISCOVERED) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_process.py b/tests/test_litellm/proxy/client/cli/autoroute/test_process.py new file mode 100644 index 00000000000..8e7f355adc1 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_process.py @@ -0,0 +1,139 @@ +import os +import socket +from typing import Optional +from unittest.mock import patch + +import pytest + +from litellm.proxy.client.cli.commands.autoroute import process as process_module +from litellm.proxy.client.cli.commands.autoroute.process import ( + PidRecord, + ProcessLaunchError, + UpError, + allocate_free_port, + clear_pid_record, + is_running, + launch_proxy, + poll_liveliness, + read_pid_record, + write_pid_record, +) + + +class FakeProcess: + def __init__(self, returncode: Optional[int] = None): + self.returncode = returncode + + def poll(self) -> Optional[int]: + return self.returncode + + +class FakeResponse: + def __init__(self, status_code: int): + self.status_code = status_code + + +def test_allocate_free_port_returns_a_bindable_port(): + port = allocate_free_port() + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", port)) + + +class TestLaunchProxy: + def test_binds_loopback_only_not_all_interfaces(self, tmp_path): + """proxy_cli.py's own --host default is 0.0.0.0 -- without an explicit override here, the + ephemeral proxy would be reachable from other hosts on the network despite base_url always + being built from 127.0.0.1, exposing its unauthenticated-until-master-key-lands routes.""" + config_path = tmp_path / "config.yaml" + log_path = tmp_path / "proxy.log" + + with patch.object(process_module.subprocess, "Popen") as mock_popen: + launch_proxy(config_path, 12345, log_path) + + args = mock_popen.call_args[0][0] + assert "--host" in args + assert args[args.index("--host") + 1] == "127.0.0.1" + + +class TestPidRecordRoundTrip: + def test_write_then_read_round_trips(self, tmp_path): + path = tmp_path / "pid.json" + record = PidRecord(pid=123, port=4000, config_path="/tmp/config.yaml", log_path="/tmp/proxy.log") + + write_pid_record(record, path) + + assert read_pid_record(path) == record + + def test_read_missing_file_returns_none(self, tmp_path): + assert read_pid_record(tmp_path / "missing.json") is None + + def test_read_raises_clean_error_on_corrupt_content(self, tmp_path): + path = tmp_path / "pid.json" + path.write_text("not json at all {{{") + + with pytest.raises(UpError, match="invalid or unexpected JSON"): + read_pid_record(path) + + def test_clear_removes_an_existing_record(self, tmp_path): + path = tmp_path / "pid.json" + write_pid_record(PidRecord(pid=1, port=1, config_path="a", log_path="b"), path) + assert path.exists() + + clear_pid_record(path) + + assert not path.exists() + + def test_clear_missing_file_is_a_no_op(self, tmp_path): + clear_pid_record(tmp_path / "missing.json") + + def test_write_creates_parent_directories(self, tmp_path): + path = tmp_path / "nested" / "dir" / "pid.json" + + write_pid_record(PidRecord(pid=1, port=1, config_path="a", log_path="b"), path) + + assert path.exists() + + +class TestIsRunning: + def test_current_process_is_running(self): + assert is_running(os.getpid()) is True + + def test_huge_unlikely_pid_is_not_running(self): + assert is_running(2**30) is False + + def test_permission_error_from_kill_is_treated_as_running(self, monkeypatch): + def fake_kill(pid: int, sig: int) -> None: + raise PermissionError("not permitted to signal this pid") + + monkeypatch.setattr(process_module.os, "kill", fake_kill) + + assert is_running(999) is True + + +class TestPollLiveliness: + def test_succeeds_when_health_check_returns_200_quickly(self, monkeypatch, tmp_path): + monkeypatch.setattr(process_module.requests, "get", lambda url, timeout: FakeResponse(200)) + + poll_liveliness("http://127.0.0.1:4000", tmp_path / "proxy.log", FakeProcess(), timeout=5.0) + + def test_raises_with_log_tail_when_timeout_elapses(self, monkeypatch, tmp_path): + log_path = tmp_path / "proxy.log" + log_path.write_text("line one\nline two\nline three\n") + monkeypatch.setattr(process_module.requests, "get", lambda url, timeout: FakeResponse(500)) + monkeypatch.setattr(process_module.time, "sleep", lambda seconds: None) + + with pytest.raises(ProcessLaunchError) as exc_info: + poll_liveliness("http://127.0.0.1:4000", log_path, FakeProcess(), timeout=0.05) + + assert "never became healthy" in str(exc_info.value) + assert "line three" in str(exc_info.value) + + def test_raises_immediately_when_process_already_exited(self, tmp_path): + log_path = tmp_path / "proxy.log" + log_path.write_text("crash log line") + + with pytest.raises(ProcessLaunchError) as exc_info: + poll_liveliness("http://127.0.0.1:4000", log_path, FakeProcess(returncode=1), timeout=5.0) + + assert "exited early" in str(exc_info.value) + assert "crash log line" in str(exc_info.value) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py new file mode 100644 index 00000000000..40d3e7f2aee --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py @@ -0,0 +1,56 @@ +from litellm.proxy.client.cli.commands.autoroute.settings import ( + ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, + merge_claude_settings_static_token, +) + + +def test_preserves_unrelated_top_level_keys(): + merged = merge_claude_settings_static_token({"theme": "dark"}, "http://127.0.0.1:4000", "token-abc") + assert merged["theme"] == "dark" + + +def test_preserves_unrelated_env_keys(): + settings = {"env": {"SOME_OTHER_VAR": "value"}} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert merged["env"]["SOME_OTHER_VAR"] == "value" + + +def test_sets_base_url_and_auth_token(): + merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000/", "token-abc") + assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" + assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" + + +def test_drops_stray_api_key(): + settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert "ANTHROPIC_API_KEY" not in merged["env"] + + +def test_removes_existing_api_key_helper(): + settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token"} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert "apiKeyHelper" not in merged + + +def test_does_not_mutate_input(): + settings = {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} + merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert settings == {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} + + +def test_forces_all_claude_code_default_model_tiers_to_the_autorouter(): + # A bare "*" model_name deployment looks like the obvious way to catch every request + # regardless of which model Claude Code thinks it's using, but Router's auto-router + # registry is keyed by the literal requested model string with no wildcard resolution + # (litellm/router.py:10711-10717) -- so the only reliable way to make every one of Claude + # Code's own tiers hit the auto-router is to override the env vars it reads per tier. + merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000", "token-abc") + for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: + assert merged["env"][key] == "autorouter" + + +def test_overrides_a_preexisting_default_model_env_var(): + settings = {"env": {"ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-opus-4-8"}} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert merged["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter" diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py new file mode 100644 index 00000000000..75431a2c588 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py @@ -0,0 +1,293 @@ +import asyncio +from typing import Any, Dict, List, Tuple +from unittest.mock import patch + +import click +import pytest +import yaml +from click.testing import CliRunner +from InquirerPy.base.control import Choice +from prompt_toolkit.application import create_app_session +from prompt_toolkit.input import create_pipe_input +from prompt_toolkit.output import DummyOutput + +from litellm.proxy.client.cli.commands.autoroute import wizard as wizard_module +from litellm.proxy.client.cli.commands.autoroute.config import DiscoveredModel +from litellm.proxy.client.cli.commands.autoroute.wizard import run_configure_wizard + +CHAT_AND_EMBEDDING_GROUPS: List[Dict[str, Any]] = [ + {"model_group": "gpt-4o-mini", "mode": "chat", "input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, + {"model_group": "gpt-4o", "mode": "chat", "input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, + {"model_group": "claude-opus", "mode": "chat"}, + {"model_group": "o1", "mode": "chat"}, + {"model_group": "text-embedding-3-small", "mode": "embedding"}, +] + +CHAT_ONLY_GROUPS: List[Dict[str, Any]] = [ + {"model_group": "gpt-4o-mini", "mode": "chat"}, + {"model_group": "gpt-4o", "mode": "chat"}, + {"model_group": "claude-opus", "mode": "chat"}, + {"model_group": "o1", "mode": "chat"}, +] + +EMBEDDING_ONLY_GROUPS: List[Dict[str, Any]] = [ + {"model_group": "text-embedding-3-small", "mode": "embedding"}, +] + + +@click.command() +@click.pass_context +def _invoke_wizard(ctx: click.Context) -> None: + run_configure_wizard(ctx) + + +def _run( + tmp_path, + raw_groups: List[Dict[str, Any]], + tier_picks: Dict[str, Tuple[str, ...]], + input_str: str, + classifier_pick: str = "", + embedding_pick: str = "", +): + """Drives run_configure_wizard's orchestration logic (discovery, validation, config writing, + classifier/semantic/adaptive branching) by mocking the fuzzy picker itself, since that widget + is a real prompt_toolkit application tested separately in TestFuzzyPickWidget. CliRunner's + injected input still drives the plain click.confirm() y/n prompts.""" + config_path = tmp_path / "config.yaml" + runner = CliRunner() + + def _fake_prompt_for_models(models, prompt_label): + return tier_picks[prompt_label] + + def _fake_prompt_for_model(models, prompt_label): + if prompt_label == "LLM classifier": + return classifier_pick + if prompt_label == "semantic embeddings": + return embedding_pick + raise AssertionError(f"unexpected single-pick prompt_label {prompt_label!r}") + + with ( + patch.object(wizard_module, "Client") as mock_client_cls, + patch.object(wizard_module, "CONFIG_PATH", config_path), + patch.object(wizard_module, "_is_interactive", return_value=True), + patch.object(wizard_module, "_render_and_prompt_for_models", side_effect=_fake_prompt_for_models), + patch.object(wizard_module, "_render_and_prompt_for_model", side_effect=_fake_prompt_for_model), + ): + mock_client_cls.return_value.model_groups.info.return_value = raw_groups + result = runner.invoke( + _invoke_wizard, + obj={"base_url": "http://localhost:4000", "api_key": "sk-test"}, + input=input_str, + ) + return result, config_path + + +def _router_config(config_path) -> Dict[str, Any]: + written = yaml.safe_load(config_path.read_text()) + autorouter = next(m for m in written["model_list"] if m["model_name"] == "autorouter") + return autorouter["litellm_params"]["complexity_router_config"] + + +_SIMPLE_TIER_PICKS: Dict[str, Tuple[str, ...]] = { + "SIMPLE": ("gpt-4o-mini",), + "MEDIUM": ("gpt-4o",), + "COMPLEX": ("claude-opus",), + "REASONING": ("o1",), +} + + +class TestRunConfigureWizardHappyPath: + def test_assigns_tiers_and_declines_everything(self, tmp_path): + result, config_path = _run(tmp_path, CHAT_AND_EMBEDDING_GROUPS, _SIMPLE_TIER_PICKS, input_str="n\nn\nn\n") + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["tiers"] == { + "SIMPLE": ["gpt-4o-mini"], + "MEDIUM": ["gpt-4o"], + "COMPLEX": ["claude-opus"], + "REASONING": ["o1"], + } + assert router_config["default_model"] == "gpt-4o" + assert "classifier_type" not in router_config + assert "classifier_llm_config" not in router_config + assert "semantic_keyword_matching" not in router_config + assert "adaptive" not in router_config + + def test_assigns_multiple_models_to_a_single_tier(self, tmp_path): + tier_picks = {**_SIMPLE_TIER_PICKS, "SIMPLE": ("gpt-4o-mini", "gpt-4o")} + result, config_path = _run(tmp_path, CHAT_AND_EMBEDDING_GROUPS, tier_picks, input_str="n\nn\nn\n") + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["tiers"]["SIMPLE"] == ["gpt-4o-mini", "gpt-4o"] + assert router_config["default_model"] == "gpt-4o" + + def test_writes_config_file_with_restricted_permissions(self, tmp_path): + result, config_path = _run(tmp_path, CHAT_AND_EMBEDDING_GROUPS, _SIMPLE_TIER_PICKS, input_str="n\nn\nn\n") + + assert result.exit_code == 0, result.output + assert config_path.exists() + assert oct(config_path.stat().st_mode)[-3:] == "600" + + def test_no_embedding_pool_skips_semantic_prompt_entirely(self, tmp_path): + result, config_path = _run(tmp_path, CHAT_ONLY_GROUPS, _SIMPLE_TIER_PICKS, input_str="n\nn\n") + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert "semantic_keyword_matching" not in router_config + + +class TestRunConfigureWizardLLMClassifier: + def test_accepting_llm_classifier_records_chosen_model(self, tmp_path): + result, config_path = _run( + tmp_path, CHAT_AND_EMBEDDING_GROUPS, _SIMPLE_TIER_PICKS, input_str="y\nn\nn\n", classifier_pick="gpt-4o" + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["classifier_type"] == "llm" + assert router_config["classifier_llm_config"]["model"] == "gpt-4o" + + +class TestRunConfigureWizardSemanticMatching: + def test_accepting_semantic_matching_records_embedding_model(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + _SIMPLE_TIER_PICKS, + input_str="n\ny\nn\n", + embedding_pick="text-embedding-3-small", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["semantic_keyword_matching"] is True + assert router_config["embedding_model"] == "text-embedding-3-small" + + +class TestRunConfigureWizardAdaptive: + def test_accepting_adaptive_sets_adaptive_flag(self, tmp_path): + result, config_path = _run(tmp_path, CHAT_AND_EMBEDDING_GROUPS, _SIMPLE_TIER_PICKS, input_str="n\nn\ny\n") + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["adaptive"] is True + + +class TestRunConfigureWizardNoChatModels: + def test_fails_cleanly_without_prompting_when_no_chat_models(self, tmp_path): + result, config_path = _run(tmp_path, EMBEDDING_ONLY_GROUPS, {}, input_str="") + + assert result.exit_code != 0 + assert "no chat-capable models" in result.output.lower() + assert not config_path.exists() + + def test_surfaces_clean_error_when_response_is_not_a_list(self, tmp_path): + result, config_path = _run(tmp_path, {"data": CHAT_AND_EMBEDDING_GROUPS}, {}, input_str="") + + assert result.exit_code != 0 + assert result.exception is None or not isinstance(result.exception, AssertionError) + assert "Unexpected response from /model_group/info" in result.output + assert not config_path.exists() + + +class TestRunConfigureWizardNotInteractive: + def test_fails_cleanly_when_not_a_tty(self, tmp_path): + config_path = tmp_path / "config.yaml" + runner = CliRunner() + with ( + patch.object(wizard_module, "Client") as mock_client_cls, + patch.object(wizard_module, "CONFIG_PATH", config_path), + patch.object(wizard_module, "_is_interactive", return_value=False), + ): + mock_client_cls.return_value.model_groups.info.return_value = CHAT_AND_EMBEDDING_GROUPS + result = runner.invoke(_invoke_wizard, obj={"base_url": "http://localhost:4000", "api_key": "sk-test"}) + + assert result.exit_code != 0 + assert "interactive terminal" in result.output + assert not config_path.exists() + + +def _drive_fuzzy_pick( + models: Tuple[DiscoveredModel, ...], + prompt_label: str, + multiselect: bool, + key_events: List[Tuple[str, float]], +) -> List[str]: + """Drives the real InquirerPy fuzzy prompt through prompt_toolkit's own test input/output, + exercising the actual widget (filtering, tab-to-toggle, enter-to-confirm) rather than mocking + it away. asyncio.to_thread propagates the create_app_session context into the worker thread + running _fuzzy_pick's synchronous .execute() call.""" + + async def _run() -> List[str]: + with create_pipe_input() as pipe_input: + with create_app_session(input=pipe_input, output=DummyOutput()): + task = asyncio.ensure_future( + asyncio.to_thread(wizard_module._fuzzy_pick, models, prompt_label, multiselect) + ) + await asyncio.sleep(0.05) + for text, delay in key_events: + pipe_input.send_text(text) + await asyncio.sleep(delay) + return await task + + return asyncio.run(_run()) + + +class TestFuzzyPickWidget: + def _models(self) -> Tuple[DiscoveredModel, ...]: + return tuple(DiscoveredModel(name=f"model-{i}") for i in range(20)) + + def test_single_select_filters_and_returns_highlighted_match(self): + result = _drive_fuzzy_pick( + self._models(), "test", multiselect=False, key_events=[("model-13", 0.3), ("\r", 0.1)] + ) + assert result == ["model-13"] + + def test_multiselect_requires_tab_to_toggle_before_enter(self): + result = _drive_fuzzy_pick( + self._models(), "test", multiselect=True, key_events=[("model-7", 0.3), ("\t", 0.1), ("\r", 0.1)] + ) + assert result == ["model-7"] + + def test_multiselect_can_pick_more_than_one_across_filters(self): + result = _drive_fuzzy_pick( + self._models(), + "test", + multiselect=True, + key_events=[ + ("model-3", 0.3), + ("\t", 0.1), + *[("\x7f", 0.02) for _ in range("model-3".__len__())], + ("model-15", 0.3), + ("\t", 0.1), + ("\r", 0.1), + ], + ) + assert set(result) == {"model-3", "model-15"} + + def test_choice_wraps_name_and_value_to_the_same_model_name(self): + model = DiscoveredModel(name="only-model") + choice = Choice(value=model.name, name=model.name) + assert choice.value == choice.name == "only-model" + + +class TestRenderAndPromptForModelWrappers: + def test_single_pick_wrapper_returns_bare_string(self): + with patch.object(wizard_module, "_fuzzy_pick", return_value=["model-a"]) as mock_pick: + result = wizard_module._render_and_prompt_for_model((), "tier") + assert result == "model-a" + mock_pick.assert_called_once_with((), "tier", multiselect=False) + + def test_multi_pick_wrapper_returns_tuple(self): + with patch.object(wizard_module, "_fuzzy_pick", return_value=["model-a", "model-b"]) as mock_pick: + result = wizard_module._render_and_prompt_for_models((), "tier") + assert result == ("model-a", "model-b") + mock_pick.assert_called_once_with((), "tier", multiselect=True) + + +@pytest.mark.parametrize("isatty_value", [True, False]) +def test_is_interactive_reflects_stdin_isatty(isatty_value): + with patch.object(wizard_module.sys.stdin, "isatty", return_value=isatty_value): + assert wizard_module._is_interactive() is isatty_value diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index a451889415f..e101515d4b1 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -797,14 +797,18 @@ class TestPrintTokenCommand: verbatim as the bearer token, so any diagnostic text on stdout would corrupt authentication. - apiKeyHelper is configured as a bare command (managed-settings.json sets - just `"apiKeyHelper": "lite auth print-token"`, no --base-url flag) -- - so in the common case ctx.obj has no explicit base_url at all, and the - command must resolve the server from whatever `lite login` stored in - token.json, not from a CLI default. `--base-url`/`LITELLM_PROXY_URL` - only matters when a caller explicitly overrides it (tracked via - ctx.obj["base_url_explicit"], set by the `cli` group from - click's ParameterSource). + `lite up` now writes `apiKeyHelper` with an explicit `--base-url` bound + to whatever proxy it was pointed at (resolve_api_key_helper), so + print-token enforces that the cached token was actually issued for that + server -- a token minted for a different, previously-logged-into proxy + must never be handed to whichever server the helper is invoked for. + Settings patched by an older `lite up`, or a manually-configured + apiKeyHelper, can still invoke this bare (no --base-url at all); that + case falls back to trusting whatever `lite login` stored in token.json, + since there is no explicit target to check it against. `--base-url`/ + `LITELLM_PROXY_URL` only enforces the match when a caller explicitly + passes it (tracked via ctx.obj["base_url_explicit"], set by the `cli` + group from click's ParameterSource). """ def setup_method(self): @@ -818,8 +822,9 @@ class TestPrintTokenCommand: assert "Not authenticated" in result.output def test_bare_invocation_resolves_server_from_stored_token(self): - """The apiKeyHelper's real invocation shape: no --base-url given at - all. Must use token.json's own base_url, not a hardcoded default.""" + """The legacy/manual invocation shape: no --base-url given at all + (e.g. settings patched before resolve_api_key_helper started binding + one). Must use token.json's own base_url, not a hardcoded default.""" with ( patch( "litellm.proxy.client.cli.commands.auth.load_token", @@ -839,7 +844,10 @@ class TestPrintTokenCommand: def test_explicit_base_url_mismatch_fails_cleanly(self): """When the caller *does* explicitly pass --base-url, a token issued - for a different server must never be printed.""" + for a different server must never be printed. This is the exact + scenario `lite up`'s own bound --base-url now guards against: a + token minted for proxy A must not reach a helper invocation aimed + at proxy B, even though the token itself is otherwise fresh.""" with patch( "litellm.proxy.client.cli.commands.auth.load_token", return_value={ @@ -856,6 +864,25 @@ class TestPrintTokenCommand: assert result.exit_code != 0 assert "sk-should-not-print" not in result.output + def test_explicit_base_url_match_prints_token(self): + """`lite up`'s own bound invocation shape: --base-url matching the token's origin + must succeed exactly like the bare/legacy invocation does.""" + with patch( + "litellm.proxy.client.cli.commands.auth.load_token", + return_value={ + "base_url": "http://localhost:4000", + "key": "sk-matches", + "timestamp": time.time(), + }, + ): + result = self.runner.invoke( + print_token, + obj={"base_url": "http://localhost:4000", "base_url_explicit": True}, + ) + + assert result.exit_code == 0 + assert result.output.strip() == "sk-matches" + def test_fresh_cached_key_printed_without_network_call(self): """A recently-issued key should be printed straight from cache -- no refresh call on every single invocation (apiKeyHelper gets called diff --git a/tests/test_litellm/proxy/client/cli/test_model_groups_commands.py b/tests/test_litellm/proxy/client/cli/test_model_groups_commands.py new file mode 100644 index 00000000000..c2809a90ba2 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_model_groups_commands.py @@ -0,0 +1,114 @@ +import json +import os +from typing import Any, Dict, List +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from litellm.proxy.client.cli import cli + +SAMPLE_MODEL_GROUPS: List[Dict[str, Any]] = [ + { + "model_group": "gpt-4o", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + }, + { + "model_group": "text-embedding-3-small", + "mode": "embedding", + "input_cost_per_token": 0.0001, + "output_cost_per_token": None, + }, +] + + +@pytest.fixture +def mock_client(): + with patch("litellm.proxy.client.cli.commands.model_groups.Client") as MockClient: + yield MockClient + + +@pytest.fixture +def cli_runner(): + return CliRunner() + + +@pytest.fixture(autouse=True) +def mock_env(): + with patch.dict( + os.environ, + { + "LITELLM_PROXY_URL": "http://localhost:4000", + "LITELLM_PROXY_API_KEY": "sk-test", + }, + ): + yield + + +def test_list_table_format_shows_model_names_and_modes(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = SAMPLE_MODEL_GROUPS + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code == 0, result.output + assert "gpt-4o" in result.output + assert "chat" in result.output + assert "text-embedding-3-small" in result.output + assert "embedding" in result.output + assert "0.01" in result.output + assert "0.02" in result.output + + mock_client.assert_called_once_with(base_url="http://localhost:4000", api_key="sk-test") + mock_client.return_value.model_groups.info.assert_called_once() + + +def test_list_table_format_defaults_missing_mode_to_chat(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = [{"model_group": "some-model"}] + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code == 0, result.output + assert "some-model" in result.output + assert "chat" in result.output + + +def test_list_json_format_round_trips_raw_data(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = SAMPLE_MODEL_GROUPS + + result = cli_runner.invoke(cli, ["model-groups", "list", "--format", "json"]) + + assert result.exit_code == 0, result.output + assert json.loads(result.output) == SAMPLE_MODEL_GROUPS + + +def test_list_with_custom_base_url_and_api_key(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = [] + + result = cli_runner.invoke( + cli, + ["--base-url", "http://custom.server:8000", "--api-key", "custom-key", "model-groups", "list"], + ) + + assert result.exit_code == 0, result.output + mock_client.assert_called_once_with(base_url="http://custom.server:8000", api_key="custom-key") + + +def test_list_error_handling(mock_client, cli_runner): + mock_client.return_value.model_groups.info.side_effect = Exception("API Error") + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code != 0 + assert "API Error" in str(result.exception) + + +def test_list_surfaces_clean_error_when_response_is_not_a_list(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = {"data": SAMPLE_MODEL_GROUPS} + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code != 0 + assert result.exception is None or not isinstance(result.exception, AssertionError) + assert "Unexpected response from /model_group/info" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py new file mode 100644 index 00000000000..1b182553644 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -0,0 +1,398 @@ +import json +import shutil +import stat +import sys +from unittest.mock import patch + +import click +import pytest +from click.testing import CliRunner + +from litellm.proxy.client.cli.commands import up as up_module +from litellm.proxy.client.cli.commands.agents import AgentRunError +from litellm.proxy.client.cli.commands.up import ( + BackupRecord, + UpError, + _ensure_fresh_login, + down, + load_json_or_empty, + merge_claude_settings, + read_backup, + resolve_api_key_helper, + restore_claude_settings, + up, + write_backup, +) + +UP_MODULE = "litellm.proxy.client.cli.commands.up" + + +def _patch_paths(monkeypatch, tmp_path): + settings_path = tmp_path / "claude_settings.json" + backup_path = tmp_path / "backup.json" + monkeypatch.setattr(up_module, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(up_module, "BACKUP_PATH", backup_path) + return settings_path, backup_path + + +class TestMergeClaudeSettings: + def test_preserves_unrelated_top_level_keys(self): + merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", "helper") + assert merged["theme"] == "dark" + + def test_preserves_unrelated_env_keys(self): + settings = {"env": {"SOME_OTHER_VAR": "value"}} + merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + assert merged["env"]["SOME_OTHER_VAR"] == "value" + + def test_overrides_base_url_and_helper(self): + settings = { + "env": {"ANTHROPIC_BASE_URL": "https://old.example.com"}, + "apiKeyHelper": "old-helper", + } + merged = merge_claude_settings(settings, "http://localhost:4000/", "new-helper") + assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert merged["apiKeyHelper"] == "new-helper" + + def test_drops_stray_api_key(self): + settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} + merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + assert "ANTHROPIC_API_KEY" not in merged["env"] + + def test_works_from_empty_settings(self): + merged = merge_claude_settings({}, "http://localhost:4000", "helper") + assert merged["env"] == {"ANTHROPIC_BASE_URL": "http://localhost:4000"} + assert merged["apiKeyHelper"] == "helper" + + def test_does_not_mutate_input(self): + settings = {"env": {"FOO": "bar"}} + merge_claude_settings(settings, "http://localhost:4000", "helper") + assert settings == {"env": {"FOO": "bar"}} + + +class TestLoadJsonOrEmpty: + def test_returns_empty_dict_when_file_does_not_exist(self, tmp_path): + assert load_json_or_empty(tmp_path / "missing.json") == {} + + def test_returns_empty_dict_when_file_is_empty(self, tmp_path): + path = tmp_path / "settings.json" + path.write_text("") + assert load_json_or_empty(path) == {} + + def test_returns_empty_dict_when_file_is_whitespace_only(self, tmp_path): + path = tmp_path / "settings.json" + path.write_text(" \n") + assert load_json_or_empty(path) == {} + + def test_parses_real_content(self, tmp_path): + path = tmp_path / "settings.json" + path.write_text(json.dumps({"theme": "dark"})) + assert load_json_or_empty(path) == {"theme": "dark"} + + def test_raises_clean_error_on_invalid_json(self, tmp_path): + path = tmp_path / "settings.json" + path.write_text("not json at all {{{") + with pytest.raises(UpError, match="invalid JSON"): + load_json_or_empty(path) + + def test_raises_clean_error_on_non_object_root(self, tmp_path): + path = tmp_path / "settings.json" + path.write_text(json.dumps([1, 2, 3])) + with pytest.raises(UpError, match="invalid JSON"): + load_json_or_empty(path) + + +class TestBackupRoundTrip: + def test_restores_original_content_when_file_existed(self, monkeypatch, tmp_path): + settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + original = {"apiKeyHelper": "old-helper", "theme": "dark"} + settings_path.write_text(json.dumps(original)) + + write_backup(BackupRecord(existed=True, content=original)) + settings_path.write_text(json.dumps({"apiKeyHelper": "lite-helper"})) + + restored = restore_claude_settings() + + assert restored is not None + assert restored.existed is True + assert json.loads(settings_path.read_text()) == original + assert not backup_path.exists() + + def test_deletes_settings_file_when_it_did_not_exist_before(self, monkeypatch, tmp_path): + settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + write_backup(BackupRecord(existed=False, content=None)) + settings_path.write_text(json.dumps({"apiKeyHelper": "lite-helper"})) + + restored = restore_claude_settings() + + assert restored is not None + assert restored.existed is False + assert not settings_path.exists() + assert not backup_path.exists() + + def test_no_backup_is_a_no_op_returning_none(self, monkeypatch, tmp_path): + settings_path, _backup_path = _patch_paths(monkeypatch, tmp_path) + assert restore_claude_settings() is None + assert not settings_path.exists() + + def test_recreates_claude_dir_if_it_was_deleted_while_up_was_running(self, monkeypatch, tmp_path): + """If ~/.claude/ is removed while `lite up` holds it open, restoring must recreate the + directory rather than crash with FileNotFoundError and strand the backup file, which + would otherwise permanently break every future `lite down`.""" + claude_dir = tmp_path / "claude_dir" + settings_path = claude_dir / "settings.json" + backup_path = tmp_path / "backup.json" + monkeypatch.setattr(up_module, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(up_module, "BACKUP_PATH", backup_path) + original = {"theme": "dark"} + claude_dir.mkdir(parents=True) + write_backup(BackupRecord(existed=True, content=original)) + shutil.rmtree(claude_dir) + + restored = restore_claude_settings() + + assert restored is not None + assert json.loads(settings_path.read_text()) == original + assert not backup_path.exists() + + def test_read_backup_round_trips_write_backup(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + write_backup(BackupRecord(existed=True, content={"a": 1})) + assert read_backup() == BackupRecord(existed=True, content={"a": 1}) + + def test_read_backup_missing_file_returns_none(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + assert read_backup() is None + + def test_read_backup_raises_clean_error_on_corrupt_content(self, monkeypatch, tmp_path): + _settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + backup_path.parent.mkdir(parents=True, exist_ok=True) + backup_path.write_text("not json at all {{{") + + with pytest.raises(UpError, match="invalid or unexpected JSON"): + read_backup() + + def test_write_backup_restricts_permissions_for_a_new_file(self, monkeypatch, tmp_path): + _settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + write_backup(BackupRecord(existed=True, content={"a": 1})) + assert stat.S_IMODE(backup_path.stat().st_mode) == 0o600 + + def test_write_backup_restricts_permissions_of_a_preexisting_permissive_file(self, monkeypatch, tmp_path): + _settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + backup_path.parent.mkdir(parents=True, exist_ok=True) + backup_path.write_text("{}") + backup_path.chmod(0o644) + + write_backup(BackupRecord(existed=True, content={"a": 1})) + + assert stat.S_IMODE(backup_path.stat().st_mode) == 0o600 + + def test_backup_file_always_removed_after_restore(self, monkeypatch, tmp_path): + _settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + write_backup(BackupRecord(existed=False, content=None)) + assert backup_path.exists() + + restore_claude_settings() + + assert not backup_path.exists() + + +class TestResolveApiKeyHelper: + def test_returns_helper_command_bound_to_the_selected_proxy(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") + helper = resolve_api_key_helper("http://localhost:4000") + assert helper == "/usr/local/bin/lite auth print-token --base-url http://localhost:4000" + + def test_quotes_a_base_url_containing_shell_metacharacters(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") + helper = resolve_api_key_helper("http://example.com/path; rm -rf /") + assert helper == "/usr/local/bin/lite auth print-token --base-url 'http://example.com/path; rm -rf /'" + + def test_raises_when_lite_not_on_path(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda name: None) + with pytest.raises(UpError, match="Could not find `lite`"): + resolve_api_key_helper("http://localhost:4000") + + +def _make_ctx(base_url): + return click.Context(click.Command("test"), obj={"base_url": base_url}) + + +class TestEnsureFreshLogin: + """A token that is fresh but was issued for a *different* proxy must not be trusted: without + this check, a user logged into proxy A who runs `up --base-url proxy-b` would silently get an + apiKeyHelper wired up around proxy A's real token, which print-token would then hand to proxy B.""" + + def test_reuses_a_fresh_token_issued_for_the_same_proxy(self, monkeypatch): + monkeypatch.setattr(up_module, "load_token", lambda: {"key": "sk-a", "base_url": "http://proxy-a:4000"}) + monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) + login_calls = [] + monkeypatch.setattr(up_module, "login", lambda ctx: login_calls.append(ctx)) + + _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + + assert login_calls == [] + + def test_forces_a_fresh_login_when_the_cached_token_is_for_a_different_proxy(self, monkeypatch): + monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: True) + tokens = iter( + [ + {"key": "sk-a", "base_url": "http://proxy-a:4000"}, + {"key": "sk-b", "base_url": "http://proxy-b:4000"}, + ] + ) + monkeypatch.setattr(up_module, "load_token", lambda: next(tokens)) + monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) + login_calls = [] + + @click.pass_context + def fake_login(ctx): + login_calls.append(ctx.obj["base_url"]) + + monkeypatch.setattr(up_module, "login", fake_login) + + _ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + + assert login_calls == ["http://proxy-b:4000"] + + def test_fails_cleanly_non_interactively_when_only_a_different_proxys_token_is_cached(self, monkeypatch): + monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: False) + monkeypatch.setattr(up_module, "load_token", lambda: {"key": "sk-a", "base_url": "http://proxy-a:4000"}) + monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) + + with pytest.raises(UpError, match="lite login"): + _ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + + +class TestUpCommand: + def setup_method(self): + self.runner = CliRunner() + + def test_refuses_double_start_without_touching_settings_file(self, monkeypatch, tmp_path): + settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + existing_backup = {"existed": False, "content": None} + backup_path.write_text(json.dumps(existing_backup)) + + with ( + patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}), + patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), + patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), + patch(f"{UP_MODULE}.verify_proxy_key"), + ): + result = self.runner.invoke(up, obj={"base_url": "http://localhost:4000"}) + + assert result.exit_code != 0 + assert "already" in result.output + assert "lite down" in result.output + assert not settings_path.exists() + assert json.loads(backup_path.read_text()) == existing_backup + + def test_no_fresh_login_non_interactive_fails_cleanly(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + monkeypatch.setattr(sys.stdin, "isatty", lambda: False) + + with patch(f"{UP_MODULE}.load_token", return_value=None): + result = self.runner.invoke(up, obj={"base_url": "http://localhost:4000"}) + + assert result.exit_code != 0 + assert "lite login" in result.output + + def test_unreachable_proxy_fails_cleanly(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + + with ( + patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}), + patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), + patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), + patch( + f"{UP_MODULE}.verify_proxy_key", + side_effect=AgentRunError("Could not reach the LiteLLM proxy at http://localhost:4000"), + ), + ): + result = self.runner.invoke(up, obj={"base_url": "http://localhost:4000"}) + + assert result.exit_code != 0 + assert "Could not reach the LiteLLM proxy" in result.output + + def test_happy_path_writes_settings_and_backup_then_restores_on_stop(self, monkeypatch, tmp_path): + settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + original = {"theme": "dark"} + settings_path.write_text(json.dumps(original)) + + captured = {} + + def fake_wait(self, timeout=None): + captured["settings"] = json.loads(settings_path.read_text()) + captured["backup_existed"] = backup_path.exists() + return True + + with ( + patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}), + patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), + patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), + patch(f"{UP_MODULE}.verify_proxy_key"), + patch( + f"{UP_MODULE}.resolve_api_key_helper", + return_value="/usr/local/bin/lite auth print-token", + ), + patch(f"{UP_MODULE}.signal.signal"), + patch(f"{UP_MODULE}.atexit.register"), + patch("threading.Event.wait", new=fake_wait), + ): + result = self.runner.invoke(up, obj={"base_url": "http://localhost:4000"}) + + assert result.exit_code == 0, result.output + assert captured["backup_existed"] is True + assert captured["settings"]["theme"] == "dark" + assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert captured["settings"]["apiKeyHelper"] == "/usr/local/bin/lite auth print-token" + assert json.loads(settings_path.read_text()) == original + assert not backup_path.exists() + + +class TestDownCommand: + def setup_method(self): + self.runner = CliRunner() + + def test_restores_when_backup_exists(self, monkeypatch, tmp_path): + settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + original = {"apiKeyHelper": "old-helper"} + write_backup(BackupRecord(existed=True, content=original)) + settings_path.write_text(json.dumps({"apiKeyHelper": "lite-helper"})) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Restored" in result.output + assert json.loads(settings_path.read_text()) == original + assert not backup_path.exists() + + def test_removes_settings_file_when_it_did_not_exist_before(self, monkeypatch, tmp_path): + settings_path, _backup_path = _patch_paths(monkeypatch, tmp_path) + write_backup(BackupRecord(existed=False, content=None)) + settings_path.write_text(json.dumps({"apiKeyHelper": "lite-helper"})) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Removed" in result.output + assert not settings_path.exists() + + def test_prints_nothing_to_restore_when_no_backup(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Nothing to restore." in result.output + + def test_surfaces_clean_error_on_a_corrupt_backup_file(self, monkeypatch, tmp_path): + _settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + backup_path.parent.mkdir(parents=True, exist_ok=True) + backup_path.write_text("not json at all {{{") + + result = self.runner.invoke(down) + + assert result.exit_code != 0 + assert result.exception is None or isinstance(result.exception, SystemExit) + assert "invalid or unexpected JSON" in result.output diff --git a/uv.lock b/uv.lock index 2439d1488cd..22e3a3feece 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-13T00:19:39.570486Z" +exclude-newer = "2026-07-13T03:38:04.421387Z" exclude-newer-span = "P3D" [manifest] @@ -222,9 +222,9 @@ name = "aiologic" version = "0.17.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "sniffio", marker = "python_full_version < '3.13'" }, + { name = "sniffio" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, - { name = "wrapt", marker = "python_full_version < '3.13'" }, + { name = "wrapt" }, ] sdist = { url = "https://files.pythonhosted.org/packages/53/a7/809482759f40079f4c4328c7318bf569ae25d457f5017aad30a1b9aafedc/aiologic-0.17.0.tar.gz", hash = "sha256:65aa058e858c94cd208badb188e7f00b54dcabb3ba85b34f794db98074d108b9", size = 251625, upload-time = "2026-06-14T12:24:35.367Z" } wheels = [ @@ -516,14 +516,14 @@ name = "aurelio-sdk" version = "0.0.19" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiofiles", marker = "python_full_version < '3.14'" }, - { name = "aiohttp", marker = "python_full_version < '3.14'" }, - { name = "colorlog", marker = "python_full_version < '3.14'" }, - { name = "pydantic", marker = "python_full_version < '3.14'" }, - { name = "python-dotenv", marker = "python_full_version < '3.14'" }, - { name = "requests", marker = "python_full_version < '3.14'" }, - { name = "requests-toolbelt", marker = "python_full_version < '3.14'" }, - { name = "tornado", marker = "python_full_version < '3.14'" }, + { name = "aiofiles" }, + { name = "aiohttp" }, + { name = "colorlog" }, + { name = "pydantic" }, + { name = "python-dotenv" }, + { name = "requests" }, + { name = "requests-toolbelt" }, + { name = "tornado" }, ] sdist = { url = "https://files.pythonhosted.org/packages/27/0e/c2e369ad173fb3d76448e46d10beb3dcc53388318933ddf8169a3f21a810/aurelio_sdk-0.0.19.tar.gz", hash = "sha256:14107e7440ff2efd0b4a08c52fb595e7680bd4bc973a0ddfb3b64157c6666b91", size = 15258, upload-time = "2025-03-24T14:37:32.203Z" } wheels = [ @@ -1047,7 +1047,7 @@ name = "coloredlogs" version = "15.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "humanfriendly", marker = "python_full_version < '3.14'" }, + { name = "humanfriendly" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/c7/eed8f27100517e8c0e6b923d5f0845d0cb99763da6fdee00478f91db7325/coloredlogs-15.0.1.tar.gz", hash = "sha256:7c991aa71a4577af2f82600d8f8f3a89f936baeaf9b50a9c197da014e5bf16b0", size = 278520, upload-time = "2021-06-11T10:22:45.202Z" } wheels = [ @@ -1059,7 +1059,7 @@ name = "colorlog" version = "6.10.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "python_full_version < '3.14' and sys_platform == 'win32'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a2/61/f083b5ac52e505dfc1c624eafbf8c7589a0d7f32daa398d2e7590efa5fda/colorlog-6.10.1.tar.gz", hash = "sha256:eb4ae5cb65fe7fec7773c2306061a8e63e02efc2c72eba9d27b0fa23c94f1321", size = 17162, upload-time = "2025-10-16T16:14:11.978Z" } wheels = [ @@ -1420,7 +1420,7 @@ name = "culsans" version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiologic", marker = "python_full_version < '3.13'" }, + { name = "aiologic" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/d9/e3/49afa1bc180e0d28008ec6bcdf82a4072d1c7a41032b5b759b60814ca4b0/culsans-0.11.0.tar.gz", hash = "sha256:0b43d0d05dce6106293d114c86e3fb4bfc63088cfe8ff08ed3fe36891447fe33", size = 107546, upload-time = "2025-12-31T23:15:38.196Z" } @@ -2966,7 +2966,7 @@ name = "humanfriendly" version = "10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pyreadline3", marker = "python_full_version < '3.14' and sys_platform == 'win32'" }, + { name = "pyreadline3", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/3f/2c29224acb2e2df4d2046e4c73ee2662023c58ff5b113c4c1adac0886c43/humanfriendly-10.0.tar.gz", hash = "sha256:6b0b831ce8f15f7300721aa49829fc4e83921a9a301cc7f606be6686a2288ddc", size = 360702, upload-time = "2021-09-17T21:40:43.31Z" } wheels = [ @@ -3049,6 +3049,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "inquirerpy" +version = "0.3.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pfzy" }, + { name = "prompt-toolkit" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/64/73/7570847b9da026e07053da3bbe2ac7ea6cde6bb2cbd3c7a5a950fa0ae40b/InquirerPy-0.3.4.tar.gz", hash = "sha256:89d2ada0111f337483cb41ae31073108b2ec1e618a49d7110b0d7ade89fc197e", size = 44431, upload-time = "2022-06-27T23:11:20.598Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/ff/3b59672c47c6284e8005b42e84ceba13864aa0f39f067c973d1af02f5d91/InquirerPy-0.3.4-py3-none-any.whl", hash = "sha256:c65fdfbac1fa00e3ee4fb10679f4d3ed7a012abf4833910e63c295827fe2a7d4", size = 67677, upload-time = "2022-06-27T23:11:17.723Z" }, +] + [[package]] name = "isodate" version = "0.7.2" @@ -3280,10 +3293,10 @@ name = "jsonschema-path" version = "0.3.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pathable", marker = "python_full_version < '3.14'" }, - { name = "pyyaml", marker = "python_full_version < '3.14'" }, - { name = "referencing", marker = "python_full_version < '3.14'" }, - { name = "requests", marker = "python_full_version < '3.14'" }, + { name = "pathable" }, + { name = "pyyaml" }, + { name = "referencing" }, + { name = "requests" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6e/45/41ebc679c2a4fced6a722f624c18d658dee42612b83ea24c1caf7c0eb3a8/jsonschema_path-0.3.4.tar.gz", hash = "sha256:8365356039f16cc65fddffafda5f58766e34bebab7d6d105616ab52bc4297001", size = 11159, upload-time = "2025-01-24T14:33:16.547Z" } wheels = [ @@ -3754,6 +3767,7 @@ caching = [ { name = "diskcache" }, ] cli = [ + { name = "inquirerpy" }, { name = "pyyaml" }, { name = "requests" }, { name = "rich" }, @@ -3789,6 +3803,7 @@ proxy = [ { name = "fastapi-sso" }, { name = "granian" }, { name = "gunicorn" }, + { name = "inquirerpy" }, { name = "litellm-enterprise" }, { name = "litellm-proxy-extras" }, { name = "mcp" }, @@ -3966,6 +3981,8 @@ requires-dist = [ { name = "gunicorn", marker = "extra == 'proxy'", specifier = ">=23.0.0,<24.0" }, { name = "httpx", specifier = ">=0.28.0,<1.0" }, { name = "importlib-metadata", specifier = ">=8.0.0,<9.0" }, + { name = "inquirerpy", marker = "extra == 'cli'", specifier = ">=0.3.4,<1.0" }, + { name = "inquirerpy", marker = "extra == 'proxy'", specifier = ">=0.3.4,<1.0" }, { name = "jinja2", specifier = ">=3.1.6,<4.0" }, { name = "jsonschema", specifier = ">=4.0.0,<5.0" }, { name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=2.59.7,<3.0" }, @@ -4464,7 +4481,7 @@ version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/fd/15/76f86faa0902836cc133939732f7611ace68cf54148487a99c539c272dc8/ml_dtypes-0.4.1.tar.gz", hash = "sha256:fad5f2de464fd09127e49b7fd1252b9006fb43d2edc1ff112d390c324af5ca7a", size = 692594, upload-time = "2024-09-13T19:07:11.624Z" } wheels = [ @@ -4963,14 +4980,14 @@ name = "openapi-core" version = "0.22.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "isodate", marker = "python_full_version < '3.14'" }, - { name = "jsonschema", marker = "python_full_version < '3.14'" }, - { name = "jsonschema-path", marker = "python_full_version < '3.14'" }, - { name = "more-itertools", marker = "python_full_version < '3.14'" }, - { name = "openapi-schema-validator", marker = "python_full_version < '3.14'" }, - { name = "openapi-spec-validator", marker = "python_full_version < '3.14'" }, - { name = "typing-extensions", marker = "python_full_version < '3.14'" }, - { name = "werkzeug", marker = "python_full_version < '3.14'" }, + { name = "isodate" }, + { name = "jsonschema" }, + { name = "jsonschema-path" }, + { name = "more-itertools" }, + { name = "openapi-schema-validator" }, + { name = "openapi-spec-validator" }, + { name = "typing-extensions" }, + { name = "werkzeug" }, ] sdist = { url = "https://files.pythonhosted.org/packages/fd/65/ee75f25b9459a02df6f713f8ffde5dacb57b8b4e45145cde4cab28b5abba/openapi_core-0.22.0.tar.gz", hash = "sha256:b30490dfa74e3aac2276105525590135212352f5dd7e5acf8f62f6a89ed6f2d0", size = 109242, upload-time = "2025-12-22T19:19:49.608Z" } wheels = [ @@ -4982,9 +4999,9 @@ name = "openapi-schema-validator" version = "0.6.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "jsonschema", marker = "python_full_version < '3.14'" }, - { name = "jsonschema-specifications", marker = "python_full_version < '3.14'" }, - { name = "rfc3339-validator", marker = "python_full_version < '3.14'" }, + { name = "jsonschema" }, + { name = "jsonschema-specifications" }, + { name = "rfc3339-validator" }, ] sdist = { url = "https://files.pythonhosted.org/packages/8b/f3/5507ad3325169347cd8ced61c232ff3df70e2b250c49f0fe140edb4973c6/openapi_schema_validator-0.6.3.tar.gz", hash = "sha256:f37bace4fc2a5d96692f4f8b31dc0f8d7400fd04f3a937798eaf880d425de6ee", size = 11550, upload-time = "2025-01-10T18:08:22.268Z" } wheels = [ @@ -4996,10 +5013,10 @@ name = "openapi-spec-validator" version = "0.7.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "jsonschema", marker = "python_full_version < '3.14'" }, - { name = "jsonschema-path", marker = "python_full_version < '3.14'" }, - { name = "lazy-object-proxy", marker = "python_full_version < '3.14'" }, - { name = "openapi-schema-validator", marker = "python_full_version < '3.14'" }, + { name = "jsonschema" }, + { name = "jsonschema-path" }, + { name = "lazy-object-proxy" }, + { name = "openapi-schema-validator" }, ] sdist = { url = "https://files.pythonhosted.org/packages/82/af/fe2d7618d6eae6fb3a82766a44ed87cd8d6d82b4564ed1c7cfb0f6378e91/openapi_spec_validator-0.7.2.tar.gz", hash = "sha256:cc029309b5c5dbc7859df0372d55e9d1ff43e96d678b9ba087f7c56fc586f734", size = 36855, upload-time = "2025-06-07T14:48:56.299Z" } wheels = [ @@ -5862,6 +5879,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7d/eb/b6260b31b1a96386c0a880edebe26f89669098acea8e0318bff6adb378fd/pathable-0.4.4-py3-none-any.whl", hash = "sha256:5ae9e94793b6ef5a4cbe0a7ce9dbbefc1eec38df253763fd0aeeacf2762dbbc2", size = 9592, upload-time = "2025-01-10T18:43:11.88Z" }, ] +[[package]] +name = "pfzy" +version = "0.3.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d9/5a/32b50c077c86bfccc7bed4881c5a2b823518f5450a30e639db5d3711952e/pfzy-0.3.4.tar.gz", hash = "sha256:717ea765dd10b63618e7298b2d98efd819e0b30cd5905c9707223dceeb94b3f1", size = 8396, upload-time = "2022-01-28T02:26:17.946Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/d7/8ff98376b1acc4503253b685ea09981697385ce344d4e3935c2af49e044d/pfzy-0.3.4-py3-none-any.whl", hash = "sha256:5f50d5b2b3207fa72e7ec0ef08372ef652685470974a107d0d4999fc5a903a96", size = 8537, upload-time = "2022-01-28T02:26:16.047Z" }, +] + [[package]] name = "pillow" version = "12.3.0" @@ -6075,6 +6101,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/98/745b810d822103adca2df8decd4c0bbe839ba7ad3511af3f0d09692fc0f0/prometheus_client-0.20.0-py3-none-any.whl", hash = "sha256:cde524a85bce83ca359cc837f28b8c0db5cac7aa653a588fd7e84ba061c329e7", size = 54474, upload-time = "2024-02-14T15:55:03.957Z" }, ] +[[package]] +name = "prompt-toolkit" +version = "3.0.52" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a1/96/06e01a7b38dce6fe1db213e061a4602dd6032a8a97ef6c1a862537732421/prompt_toolkit-3.0.52.tar.gz", hash = "sha256:28cde192929c8e7321de85de1ddbe736f1375148b02f2e17edd840042b1be855", size = 434198, upload-time = "2025-08-27T15:24:02.057Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/03/0d3ce49e2505ae70cf43bc5bb3033955d2fc9f932163e84dc0779cc47f48/prompt_toolkit-3.0.52-py3-none-any.whl", hash = "sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955", size = 391431, upload-time = "2025-08-27T15:23:59.498Z" }, +] + [[package]] name = "propcache" version = "0.5.2" @@ -7125,16 +7163,16 @@ name = "redisvl" version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "coloredlogs", marker = "python_full_version < '3.14'" }, - { name = "ml-dtypes", marker = "python_full_version < '3.14'" }, + { name = "coloredlogs" }, + { name = "ml-dtypes" }, { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, - { name = "pydantic", marker = "python_full_version < '3.14'" }, - { name = "python-ulid", marker = "python_full_version < '3.14'" }, - { name = "pyyaml", marker = "python_full_version < '3.14'" }, - { name = "redis", marker = "python_full_version < '3.14'" }, - { name = "tabulate", marker = "python_full_version < '3.14'" }, - { name = "tenacity", marker = "python_full_version < '3.14'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "pydantic" }, + { name = "python-ulid" }, + { name = "pyyaml" }, + { name = "redis" }, + { name = "tabulate" }, + { name = "tenacity" }, ] sdist = { url = "https://files.pythonhosted.org/packages/21/33/ab14865a0b2a31b1d003c29e7e8ea3a7a2f2c8ecb24e58e58d606e1f031b/redisvl-0.4.1.tar.gz", hash = "sha256:fd6a36426ba94792c0efca20915c31232d4ee3cc58eb23794a62c142696401e6", size = 77688, upload-time = "2025-02-21T22:51:41.389Z" } wheels = [ @@ -7368,7 +7406,7 @@ name = "rfc3339-validator" version = "0.1.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "six", marker = "python_full_version < '3.14'" }, + { name = "six" }, ] sdist = { url = "https://files.pythonhosted.org/packages/28/ea/a9387748e2d111c3c2b275ba970b735e04e15cdb1eb30693b6b5708c4dbd/rfc3339_validator-0.1.4.tar.gz", hash = "sha256:138a2abdf93304ad60530167e51d2dfb9549521a836871b88d7f4695d0022f6b", size = 5513, upload-time = "2021-05-12T16:37:54.178Z" } wheels = [ @@ -7817,20 +7855,20 @@ name = "semantic-router" version = "0.1.15" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiohttp", marker = "python_full_version < '3.14'" }, - { name = "aurelio-sdk", marker = "python_full_version < '3.14'" }, - { name = "colorama", marker = "python_full_version < '3.14'" }, - { name = "colorlog", marker = "python_full_version < '3.14'" }, - { name = "litellm", marker = "python_full_version < '3.14'" }, + { name = "aiohttp" }, + { name = "aurelio-sdk" }, + { name = "colorama" }, + { name = "colorlog" }, + { name = "litellm" }, { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, - { name = "openai", marker = "python_full_version < '3.14'" }, - { name = "pydantic", marker = "python_full_version < '3.14'" }, - { name = "pyyaml", marker = "python_full_version < '3.14'" }, - { name = "regex", marker = "python_full_version < '3.14'" }, - { name = "tiktoken", marker = "python_full_version < '3.14'" }, - { name = "tornado", marker = "python_full_version < '3.14'" }, - { name = "urllib3", marker = "python_full_version < '3.14'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "openai" }, + { name = "pydantic" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "tiktoken" }, + { name = "tornado" }, + { name = "urllib3" }, ] sdist = { url = "https://files.pythonhosted.org/packages/dc/a9/1a689e916e8b280f1fd8fb335cc059be626a22fe4533baa045d32fcd6de5/semantic_router-0.1.15.tar.gz", hash = "sha256:328256ddc3c2b713101ec69561d6585aecbf1198ea3461e1486289d8c3a35288", size = 95605, upload-time = "2026-05-23T12:58:15.444Z" } wheels = [ From bf3a05878175587c0bff6a579d93f1f97db316ad Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 21:54:15 -0700 Subject: [PATCH 38/62] feat(complexity_router): enable session_affinity by default (#33500) Pin a session's first-turn model for the rest of the session by default instead of reclassifying every turn. Keeps multi-turn sessions on a single model, preserving provider prompt caches and avoiding cross-model conversation-history errors (e.g. Anthropic rejecting a thinking block produced by a different model). Requests without a resolvable session_id are unaffected. Set session_affinity: false to opt out. Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../complexity_router/config.py | 7 ++-- .../router_strategy/test_complexity_router.py | 35 ++++++++++++++++--- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index b7ffa2866f2..1f984798970 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -363,10 +363,13 @@ class ComplexityRouterConfig(BaseModel): # Session affinity: pin the first turn's routed model for the rest of the session session_affinity: bool = Field( - default=False, + default=True, description=( "When True and a session_id is resolvable on the request, pin the model chosen on the " - "session's first turn and reuse it for every later turn, skipping re-classification." + "session's first turn and reuse it for every later turn, skipping re-classification. " + "On by default so multi-turn sessions stay on one model, preserving provider prompt " + "caches and avoiding cross-model conversation-history errors. Set False to reclassify " + "every turn." ), ) session_affinity_ttl_seconds: int = Field( diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 3404b55f0db..12b2c9abefb 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2499,7 +2499,7 @@ class TestRoutingDecisionCauseLogging: class TestSessionAffinity: - """Test the opt-in session_affinity sticky-routing behavior.""" + """Test the session_affinity sticky-routing behavior (on by default).""" REASONING_MESSAGE = [ { @@ -2513,14 +2513,19 @@ class TestSessionAffinity: def session_affinity_config(self, basic_config) -> Dict: return {**basic_config, "session_affinity": True} + @pytest.fixture + def session_affinity_disabled_config(self, basic_config) -> Dict: + return {**basic_config, "session_affinity": False} + @staticmethod def _request_kwargs(session_id: str) -> Dict: return {"metadata": {"session_id": session_id}} @pytest.mark.asyncio - async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): - """Regression: session_affinity defaults to False, so a shared session_id must - not pin the model -- each turn is still classified independently.""" + async def test_enabled_by_default_pins_model(self, mock_router_instance, basic_config): + """Regression: session_affinity defaults to True, so a shared session_id pins the + first turn's model and later turns reuse it instead of reclassifying.""" + assert "session_affinity" not in basic_config mock_router_instance.cache = DualCache() router = ComplexityRouter( model_name="test-router", @@ -2535,6 +2540,28 @@ class TestSessionAffinity: model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE ) assert first.model == "o1-preview" + assert second.model == "o1-preview" + + @pytest.mark.asyncio + async def test_can_be_disabled_reclassifies_every_turn( + self, mock_router_instance, session_affinity_disabled_config + ): + """Regression: session_affinity=False must still reclassify every turn even when a + shared session_id is present, so the opt-out keeps working.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_disabled_config, + ) + request_kwargs = self._request_kwargs("session-1") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" assert second.model == "gpt-4o-mini" @pytest.mark.asyncio From afd7917b8b39ca66e48321f4fe182914720617c6 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 23:11:37 -0700 Subject: [PATCH 39/62] fix(mcp): separate issuer identity from anchoring so carry-forward keeps endpoints Making the in-memory issuer reflect a trust-on-first-use discovered value fixed the registry/row token-identity drift, but it overloaded a single field: the carry-forward gate keyed on issuer truthiness as a proxy for "endpoints are anchored to a pinned issuer, fail-closed". A discovered issuer is truthy yet not anchored, so a resource-rooted server that had learned its issuer would drop its last-known-good endpoints on a transient discovery blip instead of carrying them forward. Anchoring is now a first-class property rather than a proxy. MCPServer carries issuer_is_anchored, set at both build paths from the single _uses_issuer_anchor definition (a pinned issuer on a discovery auth type). issuer stays the identity value used by the token-identity tuple and the serializers; issuer_is_anchored is the provenance value the carry-forward gate reads to decide fail-closed. The two properties can no longer be conflated, so a discovered issuer keeps its resource-rooted endpoints carrying forward while a pinned issuer still fails closed. Regression tests pin both directions: a discovered-but-not-anchored server restores its endpoints on a discovery blip, an anchored server does not, and the build sets issuer_is_anchored true only when the issuer is pinned --- .../mcp_server/mcp_server_manager.py | 32 +++++++++--- .../types/mcp_server/mcp_server_manager.py | 1 + .../mcp_server/test_mcp_server_manager.py | 50 ++++++++++++++++++- 3 files changed, 74 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 57cd755797b..115ff2e492c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -201,6 +201,18 @@ def _blank_to_none(value: str | None) -> str | None: return value.strip() or None +def _uses_issuer_anchor(manual_issuer: str | None, is_discovery_auth_type: bool) -> bool: + """Whether the endpoints are authoritatively anchored to an admin-pinned issuer (RFC 8414 §3.3). + + This is the trust/provenance property, distinct from whether the ``issuer`` field is merely + populated: a trust-on-first-use discovered issuer sets ``issuer`` for token identity but is NOT + anchored, so its endpoints stay resource-rooted. Anchoring holds only when the issuer was pinned + (present on the row/config) on a discovery auth type. Every consumer of "is this anchored" reads + this one definition, so the answer cannot diverge across build paths. + """ + return _blank_to_none(manual_issuer) is not None and is_discovery_auth_type + + def _endpoints_yield_to_issuer( issuer: str | None, is_discovery_auth_type: bool, @@ -292,16 +304,20 @@ def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serv consistent group) or pins the same one. An admin re-pointing ``authorization_url`` to a different server must not keep serving the old server's token endpoint or granted scopes. - When an ``issuer`` is configured the endpoints must come solely from the §3.3-validated issuer - document, so carry-forward is skipped entirely for its endpoints: a failed issuer fetch leaves - them ``None`` and must stay ``None`` (fail-closed), never resurrected from the previous registry - entry. Scopes stay resource-driven and can still carry. + When the server is issuer-anchored (``issuer_is_anchored`` -- a pinned issuer on a discovery auth + type), the endpoints come solely from the §3.3-validated issuer document, so carry-forward is + skipped entirely for its endpoints: a failed issuer fetch leaves them ``None`` and must stay + ``None`` (fail-closed), never resurrected from the previous registry entry. A merely discovered + (trust-on-first-use) issuer is NOT anchored -- ``issuer`` is set for token identity but the + endpoints are resource-rooted, so they still carry forward as last-known-good, gated by the + corroboration check below like any other resource-rooted server. Scopes stay resource-driven and + can carry either way. """ if previous_server is None: return if previous_server.url != new_server.url or previous_server.auth_type != new_server.auth_type: return - if _blank_to_none(new_server.issuer): + if new_server.issuer_is_anchored: # Endpoints come solely from the §3.3-validated issuer document; a failed fetch stays # fail-closed and must not be resurrected from the previous entry. Only the resource-driven # scopes carry as last-known-good. @@ -1185,7 +1201,7 @@ class MCPServerManager: manual_token_url = _blank_to_none(server_config.get("token_url")) manual_registration_url = _blank_to_none(server_config.get("registration_url")) is_discovery_auth_type = auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES - use_issuer_anchor = manual_issuer is not None and is_discovery_auth_type + use_issuer_anchor = _uses_issuer_anchor(manual_issuer, is_discovery_auth_type) manual_authorization_url, manual_token_url, manual_registration_url = _endpoints_yield_to_issuer( manual_issuer, is_discovery_auth_type, @@ -1291,6 +1307,7 @@ class MCPServerManager: oauth2_flow=self._explicit_oauth2_flow(config_oauth2_flow), scopes=resolved_scopes, issuer=effective_issuer, + issuer_is_anchored=use_issuer_anchor, authorization_url=resolved_authorization_url, token_url=resolved_token_url, registration_url=resolved_registration_url, @@ -1685,7 +1702,7 @@ class MCPServerManager: manual_token_url = _blank_to_none(mcp_server.token_url) manual_registration_url = _blank_to_none(mcp_server.registration_url) is_discovery_auth_type = auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES - use_issuer_anchor = manual_issuer is not None and is_discovery_auth_type + use_issuer_anchor = _uses_issuer_anchor(manual_issuer, is_discovery_auth_type) manual_authorization_url, manual_token_url, manual_registration_url = _endpoints_yield_to_issuer( manual_issuer, is_discovery_auth_type, manual_authorization_url, manual_token_url, manual_registration_url ) @@ -1732,6 +1749,7 @@ class MCPServerManager: oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)), scopes=resolved_scopes, issuer=effective_issuer, + issuer_is_anchored=use_issuer_anchor, authorization_url=manual_authorization_url or getattr(gated_oauth_metadata, "authorization_url", None), token_url=manual_token_url or getattr(gated_oauth_metadata, "token_url", None), registration_url=manual_registration_url or getattr(gated_oauth_metadata, "registration_url", None), diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index cf3eed20f5b..d0d8cc4cb28 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -66,6 +66,7 @@ class MCPServer(BaseModel): client_id: Optional[str] = None client_secret: Optional[str] = None issuer: Optional[str] = None + issuer_is_anchored: bool = False scopes: Optional[List[str]] = None authorization_url: Optional[str] = None token_url: Optional[str] = None 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 b16840741af..80bf08a5eba 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 @@ -1196,6 +1196,7 @@ class TestMCPServerManager: built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) assert built.issuer == "https://idp.example.com" + assert built.issuer_is_anchored is False assert built.authorization_url == "https://idp.example.com/authorize" @pytest.mark.asyncio @@ -1326,6 +1327,7 @@ class TestMCPServerManager: anchored.assert_awaited_once_with("https://idp.example.com", "https://up.example.com/mcp") resource_rooted.assert_not_awaited() assert built.issuer == "https://idp.example.com" + assert built.issuer_is_anchored is True assert built.authorization_url == "https://idp.example.com/authorize" assert built.token_url == "https://idp.example.com/token" assert built.registration_url == "https://idp.example.com/register" @@ -5869,11 +5871,12 @@ class TestMCPServerTimestamps: assert same_authorize.registration_url == "https://idp.example.com/register" def test_carry_forward_does_not_restore_endpoints_for_issuer_anchored_server(self): - """When an issuer is configured the endpoints come solely from the §3.3-validated issuer + """When the server is issuer-anchored the endpoints come solely from the §3.3-validated issuer document, so a failed issuer fetch (token_url None) must stay fail-closed. Carry-forward must NOT resurrect the previous registry entry's token endpoint, or the very attacker-controlled endpoint the issuer anchor distrusts would keep being served across rebuilds. Resource-driven - scopes still carry as last-known-good.""" + scopes still carry as last-known-good. Anchoring is keyed on the explicit issuer_is_anchored + flag, not on issuer truthiness, so a discovered issuer does not trip this fail-closed branch.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _carry_forward_resolved_oauth_endpoints, ) @@ -5885,6 +5888,7 @@ class TestMCPServerTimestamps: transport=MCPTransport.http, auth_type=MCPAuth.oauth2, issuer="https://idp.example.com", + issuer_is_anchored=True, authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", registration_url="https://idp.example.com/register", @@ -5897,6 +5901,7 @@ class TestMCPServerTimestamps: transport=MCPTransport.http, auth_type=MCPAuth.oauth2, issuer="https://idp.example.com", + issuer_is_anchored=True, ) _carry_forward_resolved_oauth_endpoints(new_server=failed_rebuild, previous_server=previous) @@ -5906,6 +5911,47 @@ class TestMCPServerTimestamps: assert failed_rebuild.registration_url is None assert failed_rebuild.scopes == ["read"] + def test_carry_forward_restores_endpoints_for_discovered_issuer_not_anchored(self): + """A server that merely DISCOVERED its issuer trust-on-first-use is not anchored: issuer is set + for token identity but the endpoints are resource-rooted, so on a transient discovery blip they + must still carry forward as last-known-good, the same as any resource-rooted server. This is the + regression the explicit issuer_is_anchored flag prevents: keying fail-closed on issuer truthiness + alone would drop the working endpoints the moment the server learned its issuer.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _carry_forward_resolved_oauth_endpoints, + ) + + previous = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + issuer="https://idp.example.com", + issuer_is_anchored=False, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + scopes=["read"], + ) + blipped_rebuild = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + issuer="https://idp.example.com", + issuer_is_anchored=False, + authorization_url=None, + ) + + _carry_forward_resolved_oauth_endpoints(new_server=blipped_rebuild, previous_server=previous) + + assert blipped_rebuild.authorization_url == "https://idp.example.com/authorize" + assert blipped_rebuild.token_url == "https://idp.example.com/token" + assert blipped_rebuild.registration_url == "https://idp.example.com/register" + assert blipped_rebuild.scopes == ["read"] + def test_normalized_authorize_endpoint_treats_default_port_and_slash_as_identity(self): """The corroboration check must not fail on formatting-only differences an IdP legitimately emits: default port, trailing slash, host case, and query string are not identity, but a From bbd52984b1425fd48ae8af2d2a80ed8d9d1f8dae Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 00:24:02 -0700 Subject: [PATCH 40/62] fix(anthropic): stop 500 on combined thinking+signature streaming chunk (#33505) --- .../adapters/transformation.py | 12 +--- ...al_pass_through_adapters_transformation.py | 56 +++++++++++++++---- .../test_streaming_iterator_first_delta.py | 52 +++++++++++++++++ 3 files changed, 100 insertions(+), 20 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index cd75eed2e6e..4b6617fbeac 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1403,11 +1403,6 @@ class LiteLLMAnthropicMessagesAdapter: assert isinstance(thinking, str) assert isinstance(signature, str) - if thinking and signature: - raise ValueError( - "Both `thinking` and `signature` in a single streaming chunk isn't supported." - ) - return "thinking", ChatCompletionThinkingBlock( type="thinking", thinking=thinking, signature=signature ) @@ -1463,17 +1458,14 @@ class LiteLLMAnthropicMessagesAdapter: if choice.delta.reasoning_content is not None: reasoning_content += choice.delta.reasoning_content - if reasoning_content and reasoning_signature: - raise ValueError("Both `reasoning` and `signature` in a single streaming chunk isn't supported.") - if partial_json is not None: return "input_json_delta", ContentJsonBlockDelta(type="input_json_delta", partial_json=partial_json) - elif reasoning_content: - return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content) elif reasoning_signature: return "signature_delta", ContentThinkingSignatureBlockDelta( type="signature_delta", signature=reasoning_signature ) + elif reasoning_content: + return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content) else: return "text_delta", ContentTextBlockDelta(type="text_delta", text=text) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index f9c55db72b5..dfe7e0c3a51 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -256,7 +256,14 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_signature_block( } -def test_translate_streaming_openai_chunk_to_anthropic_raises_when_thinking_and_signature_content_block(): +def test_translate_streaming_openai_chunk_to_anthropic_content_block_thinking_and_signature(): + """The content-block classifier must treat a chunk carrying both ``thinking`` + and ``signature`` as a ``thinking`` block instead of raising. + + Such a chunk is the terminal signature event of an already-open thinking block, + so classifying it as ``thinking`` keeps the stream on the same block rather than + 500'ing. Before the fix this raised ``ValueError``. + """ choices = [ StreamingChoices( finish_reason=None, @@ -289,10 +296,14 @@ def test_translate_streaming_openai_chunk_to_anthropic_raises_when_thinking_and_ ) ] - with pytest.raises(ValueError): - LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ( + block_type, + content_block_start, + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block( + choices=choices + ) + + assert block_type == "thinking" def test_translate_anthropic_messages_to_openai_thinking_blocks(): @@ -738,7 +749,17 @@ def test_translate_streaming_openai_chunk_to_anthropic_with_thinking(): assert content_block_delta["signature"] == "sigsig" -def test_translate_streaming_openai_chunk_to_anthropic_raises_when_thinking_and_signature(): +def test_translate_streaming_openai_chunk_to_anthropic_emits_signature_when_thinking_and_signature(): + """A single streaming chunk carrying both ``thinking`` and ``signature`` must + translate to a ``signature_delta``, not crash. + + litellm's Anthropic streaming handler emits the ``signature_delta`` event as an + OpenAI chunk whose ``thinking_blocks`` entry re-states the full accumulated + thinking text alongside the signature (see anthropic/chat/handler.py). That text + was already streamed as ``thinking_delta`` chunks, so the signature must win and + the duplicate thinking must not be re-emitted. Before the fix this raised + ``ValueError`` and 500'd the whole stream, breaking Claude Code through the proxy. + """ choices = [ StreamingChoices( finish_reason=None, @@ -771,10 +792,25 @@ def test_translate_streaming_openai_chunk_to_anthropic_raises_when_thinking_and_ ) ] - with pytest.raises(ValueError): - LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic( - choices=choices - ) + adapter = LiteLLMAnthropicMessagesAdapter() + + ( + type_of_content, + content_block_delta, + ) = adapter._translate_streaming_openai_chunk_to_anthropic(choices=choices) + + assert type_of_content == "signature_delta" + assert content_block_delta["type"] == "signature_delta" + assert content_block_delta["signature"] == "sigsig" + + ( + block_type, + content_block_start, + ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block( + choices=choices + ) + + assert block_type == "thinking" def test_translate_anthropic_messages_to_openai_user_message_with_base64_image(): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 6ed79763753..19ec1a04b45 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -438,6 +438,58 @@ async def test_empty_reasoning_delta_mid_thinking_block_is_suppressed_async(): _assert_empty_reasoning_delta_suppressed(await _drain_async(wrapper)) +def _full_snapshot_signature_chunks() -> List[MagicMock]: + """Mirror litellm's real Anthropic streaming: incremental ``thinking_delta`` + chunks (empty signature), then a terminal chunk whose ``thinking_blocks`` entry + re-states the *full accumulated thinking text* together with the signature + (anthropic/chat/handler.py builds the signature_delta event this way), then the + answer text. + """ + return [ + _thinking_chunk("Let me "), + _thinking_chunk("think about it."), + _thinking_chunk("Let me think about it.", signature="sig-abc"), + _make_chunk(Delta(content="42")), + _make_chunk(Delta(content=None), finish_reason="stop"), + ] + + +def _assert_full_snapshot_signature_handled(events: List[dict]) -> None: + _assert_deltas_match_their_block_type(events) + # The full-text snapshot on the signature chunk must NOT be re-emitted as an + # extra thinking_delta (it was already streamed incrementally) - otherwise the + # client renders the reasoning twice. + assert _thinking_deltas(events) == ["Let me ", "think about it."] + assert "".join(_thinking_deltas(events)) == "Let me think about it." + assert _signature_deltas(events) == ["sig-abc"] + assert _text_deltas(events) == ["42"] + + +def test_full_thinking_snapshot_with_signature_emits_signature_only_sync(): + """Regression: a terminal thinking chunk carrying both the full thinking text + and the signature used to raise ``ValueError`` (500) mid-stream, breaking every + Claude Code request routed through the proxy with an extended-thinking model. It + must instead emit a single ``signature_delta`` without duplicating the thinking. + """ + wrapper = AnthropicStreamWrapper( + completion_stream=iter(_full_snapshot_signature_chunks()), + model="claude-x", + ) + _assert_full_snapshot_signature_handled(_drain_sync(wrapper)) + + +@pytest.mark.asyncio +async def test_full_thinking_snapshot_with_signature_emits_signature_only_async(): + """Async twin - the proxy serves the async iterator, so the crash must be gone + on that path too. + """ + wrapper = AnthropicStreamWrapper( + completion_stream=_AsyncStream(_full_snapshot_signature_chunks()), + model="claude-x", + ) + _assert_full_snapshot_signature_handled(await _drain_async(wrapper)) + + def test_empty_content_chunk_mid_text_block_is_suppressed_sync(): """An empty-content chunk arriving mid-text-block (no transition) used to emit a pointless ``text_delta {"text": ""}``; it must be dropped without From 5a0e1dd1dd25a0d3e6c41b995c272dc10b01c02d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 07:43:41 +0000 Subject: [PATCH 41/62] feat(autoroute): prompt for semantic keywords per tier in configure wizard (#33508) --- .../client/cli/commands/autoroute/config.py | 32 +++++++++++----- .../client/cli/commands/autoroute/wizard.py | 24 +++++++++++- .../proxy/client/cli/autoroute/test_config.py | 27 +++++++++++++ .../proxy/client/cli/autoroute/test_wizard.py | 38 ++++++++++++++++++- 4 files changed, 109 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/client/cli/commands/autoroute/config.py b/litellm/proxy/client/cli/commands/autoroute/config.py index fe2c4e467e7..2d760ef0f8a 100644 --- a/litellm/proxy/client/cli/commands/autoroute/config.py +++ b/litellm/proxy/client/cli/commands/autoroute/config.py @@ -80,24 +80,32 @@ class NoSemanticMatching(BaseModel): kind: Literal["none"] = "none" +class KeywordTierRule(BaseModel): + model_config = ConfigDict(frozen=True) + keywords: tuple[str, ...] + tier: str + + +# Satisfies complexity_router's "semantic matching requires non-empty keyword_tier_rules" +# invariant with a sane starting point; the wizard lets the user override these per tier. +DEFAULT_KEYWORD_TIER_RULES: tuple[KeywordTierRule, ...] = ( + KeywordTierRule(keywords=("hi", "hello", "thanks"), tier="SIMPLE"), + KeywordTierRule(keywords=("explain", "how does"), tier="MEDIUM"), + KeywordTierRule(keywords=("refactor", "implement", "debug"), tier="COMPLEX"), + KeywordTierRule(keywords=("step by step", "think through", "prove"), tier="REASONING"), +) + + class SemanticMatching(BaseModel): model_config = ConfigDict(frozen=True) kind: Literal["semantic"] = "semantic" embedding_model: str match_threshold: float = 0.5 + keyword_tier_rules: tuple[KeywordTierRule, ...] = DEFAULT_KEYWORD_TIER_RULES SemanticMatchingChoice = Union[NoSemanticMatching, SemanticMatching] -# Satisfies complexity_router's "semantic matching requires non-empty keyword_tier_rules" -# invariant with a sane starting point; the generated config.yaml can be hand-edited afterward. -_DEFAULT_KEYWORD_TIER_RULES: tuple[dict[str, JsonValue], ...] = ( - {"keywords": ["hi", "hello", "thanks"], "tier": "SIMPLE"}, - {"keywords": ["explain", "how does"], "tier": "MEDIUM"}, - {"keywords": ["refactor", "implement", "debug"], "tier": "COMPLEX"}, - {"keywords": ["step by step", "think through", "prove"], "tier": "REASONING"}, -) - class AutorouteConfig(BaseModel): model_config = ConfigDict(frozen=True) @@ -181,7 +189,9 @@ def build_generated_model_list(config: AutorouteConfig) -> list[JsonValue]: complexity_router_config["semantic_keyword_matching"] = True complexity_router_config["embedding_model"] = config.semantic_matching.embedding_model complexity_router_config["match_threshold"] = config.semantic_matching.match_threshold - complexity_router_config["keyword_tier_rules"] = list(_DEFAULT_KEYWORD_TIER_RULES) + complexity_router_config["keyword_tier_rules"] = [ + {"keywords": list(rule.keywords), "tier": rule.tier} for rule in config.semantic_matching.keyword_tier_rules + ] if config.adaptive: complexity_router_config["adaptive"] = True @@ -222,8 +232,10 @@ __all__ = [ "AutorouteConfig", "ClassifierChoice", "ConfigGenerationError", + "DEFAULT_KEYWORD_TIER_RULES", "DiscoveredModel", "HeuristicClassifier", + "KeywordTierRule", "LLMClassifier", "NoSemanticMatching", "SemanticMatching", diff --git a/litellm/proxy/client/cli/commands/autoroute/wizard.py b/litellm/proxy/client/cli/commands/autoroute/wizard.py index 89b53d8f9fd..60696fb2e7e 100644 --- a/litellm/proxy/client/cli/commands/autoroute/wizard.py +++ b/litellm/proxy/client/cli/commands/autoroute/wizard.py @@ -8,11 +8,13 @@ from InquirerPy.base.control import Choice from .... import Client from .config import ( + DEFAULT_KEYWORD_TIER_RULES, TIER_NAMES, AutorouteConfig, ConfigGenerationError, DiscoveredModel, HeuristicClassifier, + KeywordTierRule, LLMClassifier, NoSemanticMatching, SemanticMatching, @@ -63,6 +65,25 @@ def _render_and_prompt_for_models(models: tuple[DiscoveredModel, ...], prompt_la return tuple(_fuzzy_pick(models, prompt_label, multiselect=True)) +def _parse_keywords(raw: str) -> tuple[str, ...]: + return tuple(keyword.strip() for keyword in raw.split(",") if keyword.strip()) + + +def _prompt_for_keyword_tier_rules() -> tuple[KeywordTierRule, ...]: + """Let the user supply the semantic-matching keywords per tier, since matching those + keywords against the request is the whole point of enabling it. Each prompt is prefilled + with the built-in default, so pressing enter keeps it.""" + click.echo("\nEnter example keywords/phrases per tier (comma-separated); press enter to keep the default:") + defaults = {rule.tier: rule.keywords for rule in DEFAULT_KEYWORD_TIER_RULES} + + def _rule_for(tier: str) -> KeywordTierRule: + default_keywords = defaults.get(tier, ()) + raw = click.prompt(f" {tier} keywords", default=", ".join(default_keywords), show_default=True) + return KeywordTierRule(keywords=_parse_keywords(raw) or default_keywords, tier=tier) + + return tuple(_rule_for(tier) for tier in TIER_NAMES) + + def run_configure_wizard(ctx: click.Context) -> Path: """Discover the caller's accessible models, walk them through tier assignment, write config.""" base_url = ctx.obj["base_url"] @@ -96,7 +117,8 @@ def run_configure_wizard(ctx: click.Context) -> Path: semantic_matching = NoSemanticMatching() if embedding_pool and click.confirm("\nEnable semantic keyword matching?", default=False): embedding_model = _render_and_prompt_for_model(embedding_pool, "semantic embeddings") - semantic_matching = SemanticMatching(embedding_model=embedding_model) + keyword_tier_rules = _prompt_for_keyword_tier_rules() + semantic_matching = SemanticMatching(embedding_model=embedding_model, keyword_tier_rules=keyword_tier_rules) adaptive = click.confirm("\nEnable adaptive (bandit) selection on top of tiering?", default=False) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py index 9fa01524ef3..f8d82476ef0 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py @@ -3,10 +3,12 @@ from typing import Any, Dict, Tuple import pytest from litellm.proxy.client.cli.commands.autoroute.config import ( + DEFAULT_KEYWORD_TIER_RULES, AutorouteConfig, ConfigGenerationError, DiscoveredModel, HeuristicClassifier, + KeywordTierRule, LLMClassifier, NoSemanticMatching, SemanticMatching, @@ -129,6 +131,31 @@ class TestBuildGeneratedModelList: assert router_config["keyword_tier_rules"] assert "classifier_type" not in router_config + def test_semantic_matching_defaults_emit_builtin_keyword_rules(self): + config = _base_config(semantic_matching=SemanticMatching(embedding_model="text-embedding-3-small")) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert router_config["keyword_tier_rules"] == [ + {"keywords": list(rule.keywords), "tier": rule.tier} for rule in DEFAULT_KEYWORD_TIER_RULES + ] + + def test_semantic_matching_serializes_custom_keyword_rules(self): + config = _base_config( + semantic_matching=SemanticMatching( + embedding_model="text-embedding-3-small", + keyword_tier_rules=( + KeywordTierRule(keywords=("yo", "sup"), tier="SIMPLE"), + KeywordTierRule(keywords=("architect", "design a system"), tier="COMPLEX"), + ), + ) + ) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert router_config["keyword_tier_rules"] == [ + {"keywords": ["yo", "sup"], "tier": "SIMPLE"}, + {"keywords": ["architect", "design a system"], "tier": "COMPLEX"}, + ] + def test_complexity_router_config_reflects_adaptive(self): config = _base_config(adaptive=True) autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py index 75431a2c588..2b9240aafc7 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py @@ -156,7 +156,7 @@ class TestRunConfigureWizardSemanticMatching: tmp_path, CHAT_AND_EMBEDDING_GROUPS, _SIMPLE_TIER_PICKS, - input_str="n\ny\nn\n", + input_str="n\ny\n\n\n\n\nn\n", embedding_pick="text-embedding-3-small", ) @@ -165,6 +165,42 @@ class TestRunConfigureWizardSemanticMatching: assert router_config["semantic_keyword_matching"] is True assert router_config["embedding_model"] == "text-embedding-3-small" + def test_blank_keyword_answers_keep_the_builtin_defaults(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + _SIMPLE_TIER_PICKS, + input_str="n\ny\n\n\n\n\nn\n", + embedding_pick="text-embedding-3-small", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["keyword_tier_rules"] == [ + {"keywords": ["hi", "hello", "thanks"], "tier": "SIMPLE"}, + {"keywords": ["explain", "how does"], "tier": "MEDIUM"}, + {"keywords": ["refactor", "implement", "debug"], "tier": "COMPLEX"}, + {"keywords": ["step by step", "think through", "prove"], "tier": "REASONING"}, + ] + + def test_custom_keyword_answers_are_recorded_per_tier(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + _SIMPLE_TIER_PICKS, + input_str="n\ny\nyo, sup\n\nbuild a service, migrate\nderive, prove rigorously\nn\n", + embedding_pick="text-embedding-3-small", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["keyword_tier_rules"] == [ + {"keywords": ["yo", "sup"], "tier": "SIMPLE"}, + {"keywords": ["explain", "how does"], "tier": "MEDIUM"}, + {"keywords": ["build a service", "migrate"], "tier": "COMPLEX"}, + {"keywords": ["derive", "prove rigorously"], "tier": "REASONING"}, + ] + class TestRunConfigureWizardAdaptive: def test_accepting_adaptive_sets_adaptive_flag(self, tmp_path): From ebc6fdb4c26c91e65001143eaf4df73ba14967bf Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 00:44:00 -0700 Subject: [PATCH 42/62] fix(cli/anthropic): unblock lite autoroute proxy deps, adaptive thinking, and thinking+signature streaming (#33507) --- litellm/proxy/client/cli/README.md | 10 ++- .../client/cli/commands/autoroute/commands.py | 11 +++ .../client/cli/commands/autoroute/process.py | 16 ++++ scripts/install.sh | 20 ++++- .../test_streaming_iterator_combined_chunk.py | 87 +++++++++++++++++++ .../test_anthropic_messages_effort.py | 31 +++++++ .../client/cli/autoroute/test_commands.py | 19 ++++ .../client/cli/autoroute/test_process.py | 19 ++++ 8 files changed, 207 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index be8905e11ca..17041751d15 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -495,19 +495,21 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi #### Install the CLI -If you don't already have the `lite` command, install it with a single curl command -- no existing Python tooling required, `uv` is bootstrapped automatically if missing: +`lite autoroute up` builds and runs a throwaway litellm proxy locally, so unlike the rest of this CLI it needs the proxy server runtime, not just the thin `litellm[cli]` client. Install `litellm[proxy]` (which ships the `lite` command too) with a single curl command -- no existing Python tooling required, `uv` is bootstrapped automatically if missing: ```bash -curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install-cli.sh | sh +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh ``` -This installs only `litellm[cli]`, the thin client (`lite`), not the full proxy server. To try an unreleased branch or commit instead of the latest PyPI release, set `LITELLM_CLI_REF`: +To QA an unreleased branch or commit instead of the latest PyPI release, set `LITELLM_CLI_REF`: ```bash -curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm//scripts/install-cli.sh | \ +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm//scripts/install.sh | \ LITELLM_CLI_REF= sh ``` +The thin `scripts/install-cli.sh` installs only `litellm[cli]`, which is enough for `lite login`, `lite claude`, and `lite up`, but not for `lite autoroute up`; running it against a `litellm[cli]` install fails fast with a message telling you to install the proxy runtime. + Point the CLI at your real proxy and key before running any `lite model-groups` or `lite autoroute` command -- like every other command in this CLI, they read `LITELLM_PROXY_URL`/`LITELLM_PROXY_API_KEY` (or `--base-url`/`--api-key`), no `lite login` required: ```bash diff --git a/litellm/proxy/client/cli/commands/autoroute/commands.py b/litellm/proxy/client/cli/commands/autoroute/commands.py index 428daeefa60..161907f5b27 100644 --- a/litellm/proxy/client/cli/commands/autoroute/commands.py +++ b/litellm/proxy/client/cli/commands/autoroute/commands.py @@ -21,6 +21,7 @@ from .process import ( clear_pid_record, is_running, launch_proxy, + missing_proxy_runtime_modules, poll_liveliness, read_pid_record, secure_create, @@ -81,6 +82,16 @@ def up() -> None: if not CONFIG_PATH.exists(): raise click.ClickException("No config found. Run `lite autoroute configure` first.") + missing = missing_proxy_runtime_modules() + if missing: + raise click.ClickException( + "lite autoroute up launches a local litellm proxy, which needs the proxy runtime that the " + f"thin `litellm[cli]` install does not include (missing: {', '.join(missing)}). Install the " + "proxy runtime with `uv tool install --force 'litellm[proxy]'`, or to QA a branch, " + "`curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm//scripts/install.sh | " + "LITELLM_CLI_REF= sh`." + ) + try: existing_pid = read_pid_record() except UpError as e: diff --git a/litellm/proxy/client/cli/commands/autoroute/process.py b/litellm/proxy/client/cli/commands/autoroute/process.py index 86fa584a223..712f2eed2da 100644 --- a/litellm/proxy/client/cli/commands/autoroute/process.py +++ b/litellm/proxy/client/cli/commands/autoroute/process.py @@ -1,4 +1,5 @@ import contextlib +import importlib.util import json import os import signal @@ -37,6 +38,20 @@ class PidRecord: _PID_RECORD_ADAPTER = TypeAdapter(PidRecord) +_PROXY_RUNTIME_MODULES: tuple[str, ...] = ("fastapi", "uvicorn", "backoff", "orjson", "websockets", "apscheduler") + + +def missing_proxy_runtime_modules() -> tuple[str, ...]: + """Proxy-server modules that ``lite autoroute up`` needs but the thin CLI install lacks. + + ``launch_proxy`` runs the full ``litellm.proxy.proxy_cli`` server, whose dependencies live in + the ``proxy`` extra, not the ``cli`` extra that installs the ``lite`` command. On a thin + ``litellm[cli]`` install the subprocess dies with a bare ``ModuleNotFoundError``; detecting the + gap here lets ``up`` fail with an actionable message instead. + """ + return tuple(name for name in _PROXY_RUNTIME_MODULES if importlib.util.find_spec(name) is None) + + def allocate_free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("127.0.0.1", 0)) @@ -165,6 +180,7 @@ __all__ = [ "clear_pid_record", "is_running", "launch_proxy", + "missing_proxy_runtime_modules", "poll_liveliness", "read_pid_record", "secure_create", diff --git a/scripts/install.sh b/scripts/install.sh index 06e6249c9ba..213f8a7b440 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -5,12 +5,24 @@ # Needs only curl: uv is bootstrapped if missing, and uv provisions a compatible # Python itself (reusing a suitable system one, else downloading a managed build). # +# To install from an unreleased branch, tag, or commit instead of the latest PyPI +# release, set LITELLM_CLI_REF: +# curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm//scripts/install.sh | \ +# LITELLM_CLI_REF= sh +# # NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian # ignores the shebang when invoked as `sh` and does not support `pipefail`). set -eu # NOTE: before merging, this must stay as "litellm[proxy]" to install from PyPI. -LITELLM_PACKAGE="litellm[proxy]" +# LITELLM_CLI_REF opts into installing from a branch, tag, or commit instead (for +# example, to QA lite autoroute against an unreleased branch, which needs this proxy +# runtime, not the thin litellm[cli] install). +if [ -n "${LITELLM_CLI_REF:-}" ]; then + LITELLM_PACKAGE="litellm[proxy] @ git+https://github.com/BerriAI/litellm.git@${LITELLM_CLI_REF}" +else + LITELLM_PACKAGE="litellm[proxy]" +fi UV_VERSION="0.10.9" # ── colours ──────────────────────────────────────────────────────────────── @@ -81,7 +93,11 @@ fi # ── install ──────────────────────────────────────────────────────────────── echo "" -header "Installing litellm[proxy]…" +if [ -n "${LITELLM_CLI_REF:-}" ]; then + header "Installing litellm[proxy] from ${LITELLM_CLI_REF}…" +else + header "Installing litellm[proxy]…" +fi echo "" # --python-preference system: reuse a compatible system Python when present, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py index d67de0dcaf8..6973340101e 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -12,6 +12,7 @@ content survives. import asyncio import json from types import SimpleNamespace +from typing import AsyncIterator from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, @@ -202,3 +203,89 @@ def test_split_clears_reasoning_and_thinking_on_finish_chunk(): assert content_chunk.choices[0].delta.thinking_blocks == [{"type": "thinking"}] assert finish_chunk.choices[0].delta.reasoning_content is None assert finish_chunk.choices[0].delta.thinking_blocks is None + + +def _thinking_delta_chunk(thinking: str) -> ModelResponseStream: + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + reasoning_content=thinking, + thinking_blocks=[{"type": "thinking", "thinking": thinking, "signature": None}], + provider_specific_fields={ + "thinking_blocks": [{"type": "thinking", "thinking": thinking, "signature": None}] + }, + ), + finish_reason=None, + ) + ], + ) + + +def _signature_chunk(recap: str, signature: str) -> ModelResponseStream: + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + reasoning_content=recap, + thinking_blocks=[{"type": "thinking", "thinking": recap, "signature": signature}], + provider_specific_fields={ + "thinking_blocks": [{"type": "thinking", "thinking": recap, "signature": signature}] + }, + ), + finish_reason=None, + ) + ], + ) + + +def test_thinking_then_signature_chunk_does_not_crash_stream(): + """Regression for the /v1/messages streaming crash reported on autoroute. + + Anthropic streams extended thinking as incremental ``thinking_delta`` chunks, then a + closing chunk that recaps the full accumulated thinking AND carries the signature. The + adapter used to raise ``ValueError`` on that closing chunk, killing the whole stream. It + must instead emit a single ``signature_delta`` for the recap chunk and never re-emit the + recap thinking, so the incremental thinking text is not duplicated. + """ + chunks = [ + _thinking_delta_chunk("First, "), + _thinking_delta_chunk("reason."), + _signature_chunk("First, reason.", "sig-abc"), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="Done"), finish_reason=None)], + ), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], + usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), + ), + ] + + async def _aiter() -> "AsyncIterator[ModelResponseStream]": + for chunk in chunks: + yield chunk + + wrapper = AnthropicStreamWrapper(completion_stream=_aiter(), model="claude-haiku-4-5") + sse = _collect_async(wrapper) + + signature_deltas = [ + json.loads(line[len("data: ") :]) + for block in sse.split("\n\n") + for line in block.splitlines() + if line.startswith("data: ") and '"signature_delta"' in line + ] + assert len(signature_deltas) == 1 + assert signature_deltas[0]["delta"]["signature"] == "sig-abc" + + thinking_text = "".join( + json.loads(line[len("data: ") :])["delta"]["thinking"] + for block in sse.split("\n\n") + for line in block.splitlines() + if line.startswith("data: ") and '"thinking_delta"' in line + ) + assert thinking_text == "First, reason." + + assert "message_stop" in sse + assert "Done" in sse diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py index 5254808e315..e9d4d625421 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py @@ -32,6 +32,37 @@ def _transform(model, params, litellm_params=None): ) +def test_adaptive_thinking_only_translated_to_legacy_for_haiku_4_5(): + """The minimal autoroute repro: Claude Code sends bare ``thinking={type: adaptive}`` + (no ``output_config``) and the complexity router picks Haiku 4.5, which does not + support adaptive thinking. Anthropic 400s with "adaptive thinking is not supported on + this model" unless the flag is dropped, so it must be translated to the legacy extended + thinking the model does support rather than forwarded raw.""" + result = _transform("claude-haiku-4-5", {"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + + assert result["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + } + assert "output_config" not in result + + +def test_adaptive_thinking_only_dropped_for_non_reasoning_model(): + """Bare adaptive thinking on a model with no reasoning support at all is silently + dropped so the request still succeeds instead of being rejected.""" + result = _transform("claude-3-5-haiku-latest", {"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + + assert "thinking" not in result + + +def test_adaptive_thinking_only_preserved_for_4_6(): + """A 4.6+ model natively supports adaptive thinking, so a bare adaptive flag must not + be rewritten even without output_config.""" + result = _transform("claude-sonnet-4-6", {"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + + assert result["thinking"] == {"type": "adaptive"} + + def test_effort_translated_to_legacy_thinking_for_haiku_4_5(): """Core regression: Claude Code sends adaptive thinking + effort to Haiku 4.5 (thinking-capable, pre-4.6). Effort must be translated to legacy extended diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py index e0b5d3a71af..9efde03e04c 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py @@ -69,6 +69,25 @@ class TestUpCommand: assert result.exception is None or isinstance(result.exception, SystemExit) assert "lite autoroute configure" in result.output + def test_refuses_with_actionable_error_when_proxy_runtime_missing(self, monkeypatch, tmp_path): + """`up` launches a real litellm proxy, which the thin `litellm[cli]` install cannot run. + It must fail fast with an actionable message pointing at the proxy install, before it ever + tries to launch the doomed subprocess (which would otherwise die with a bare ImportError).""" + config_path, _log_path, _settings_path, _backup_path, _pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + monkeypatch.setattr(commands_module, "missing_proxy_runtime_modules", lambda: ("fastapi", "websockets")) + + def _fail_if_launched(*args, **kwargs): + raise AssertionError("launch_proxy must not run when the proxy runtime is missing") + + monkeypatch.setattr(commands_module, "launch_proxy", _fail_if_launched) + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "fastapi, websockets" in result.output + assert "litellm[proxy]" in result.output + def test_refuses_when_pid_record_exists_and_process_still_running(self, monkeypatch, tmp_path): config_path, _log_path, _settings_path, _backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) config_path.write_text(yaml.safe_dump({"model_list": []})) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_process.py b/tests/test_litellm/proxy/client/cli/autoroute/test_process.py index 8e7f355adc1..a4f85ea44ff 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_process.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_process.py @@ -14,6 +14,7 @@ from litellm.proxy.client.cli.commands.autoroute.process import ( clear_pid_record, is_running, launch_proxy, + missing_proxy_runtime_modules, poll_liveliness, read_pid_record, write_pid_record, @@ -137,3 +138,21 @@ class TestPollLiveliness: assert "exited early" in str(exc_info.value) assert "crash log line" in str(exc_info.value) + + +class TestMissingProxyRuntimeModules: + def test_flags_absent_modules_only(self, monkeypatch): + """A thin litellm[cli] install lacks the proxy runtime; the missing ones must be reported + (by name, for an actionable error) while modules that are importable are not.""" + monkeypatch.setattr( + process_module, + "_PROXY_RUNTIME_MODULES", + ("os", "litellm_autoroute_definitely_absent_pkg", "socket"), + ) + + assert missing_proxy_runtime_modules() == ("litellm_autoroute_definitely_absent_pkg",) + + def test_empty_when_all_present(self, monkeypatch): + monkeypatch.setattr(process_module, "_PROXY_RUNTIME_MODULES", ("os", "socket")) + + assert missing_proxy_runtime_modules() == () From dce1beadca59c06f41692ed67327d214c4793764 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 16 Jul 2026 07:11:15 -0700 Subject: [PATCH 43/62] Merge pull request #33357 from BerriAI/litellm_/gallant-carson-880e37 refactor(ui): migrate policies, deleted keys, deleted teams, budgets, and search tools tables onto shared DataTable --- ui/litellm-dashboard/eslint-suppressions.json | 29 -- .../budgets/_components/BudgetTable.test.tsx | 87 ++++ .../budgets/_components/BudgetTable.tsx | 57 +++ .../_components/BudgetTableColumns.tsx | 124 ++++++ .../budgets/_components/budget_panel.test.tsx | 39 +- .../budgets/_components/budget_panel.tsx | 75 +--- .../src/app/(dashboard)/hooks/keys/useKeys.ts | 4 +- ...cy_table.test.tsx => PolicyTable.test.tsx} | 106 ++--- .../policies/_components/PolicyTable.tsx | 82 ++++ .../_components/PolicyTableColumns.tsx | 207 +++++++++ .../policies/_components/index.tsx | 2 +- .../policies/_components/policy_table.tsx | 326 --------------- .../_components/SearchToolColumn.tsx | 104 ----- .../_components/SearchToolTable.test.tsx | 117 ++++++ .../_components/SearchToolTable.tsx | 66 +++ .../_components/SearchToolTableColumns.tsx | 168 ++++++++ .../search-tools/_components/SearchTools.tsx | 77 ++-- .../DeletedKeysPage/DeletedKeysPage.test.tsx | 37 +- .../DeletedKeysPage/DeletedKeysPage.tsx | 12 +- .../DeletedKeysTable.test.tsx | 159 ++++--- .../DeletedKeysTable/DeletedKeysTable.tsx | 393 ++---------------- .../DeletedKeysTableColumns.tsx | 130 ++++++ .../DeletedTeamsPage.test.tsx | 14 +- .../DeletedTeamsPage/DeletedTeamsPage.tsx | 4 +- .../DeletedTeamsTable.test.tsx | 43 +- .../DeletedTeamsTable/DeletedTeamsTable.tsx | 321 ++------------ .../DeletedTeamsTableColumns.tsx | 111 +++++ 27 files changed, 1485 insertions(+), 1409 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx rename ui/litellm-dashboard/src/app/(dashboard)/policies/_components/{policy_table.test.tsx => PolicyTable.test.tsx} (58%) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_table.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolColumn.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTableColumns.tsx create mode 100644 ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx create mode 100644 ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 9b55bf6ca0d..508015dcac8 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -950,19 +950,6 @@ "count": 1 } }, - "src/app/(dashboard)/policies/_components/policy_table.test.tsx": { - "react/display-name": { - "count": 1 - } - }, - "src/app/(dashboard)/policies/_components/policy_table.tsx": { - "no-nested-ternary": { - "count": 2 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/policies/_components/policy_test_panel.tsx": { "no-restricted-imports": { "count": 1 @@ -1448,22 +1435,6 @@ "count": 1 } }, - "src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/components/EntityUsageExport/ExportSummary.tsx": { "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx new file mode 100644 index 00000000000..9a7a7bd2eb9 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx @@ -0,0 +1,87 @@ +import { screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "@/../tests/test-utils"; +import BudgetTable from "./BudgetTable"; +import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; + +const makeBudget = (overrides: Partial = {}): budgetItem => ({ + budget_id: "budget-1", + max_budget: 100, + tpm_limit: 1000, + rpm_limit: 10, + updated_at: "2024-01-01T00:00:00Z", + ...overrides, +}); + +const defaultProps = { + budgets: [makeBudget()], + isLoading: false, + canModify: true, + onEditClick: vi.fn(), + onDeleteClick: vi.fn(), +}; + +describe("BudgetTable", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should display budget information", () => { + renderWithProviders(); + expect(screen.getByText("budget-1")).toBeInTheDocument(); + expect(screen.getByText("$100.00")).toBeInTheDocument(); + expect(screen.getByText("1000")).toBeInTheDocument(); + expect(screen.getByText("10")).toBeInTheDocument(); + }); + + it("should show n/a for missing rate limits and Unlimited for a missing max budget", () => { + renderWithProviders( + , + ); + expect(screen.getAllByText("n/a")).toHaveLength(2); + expect(screen.getByText("Unlimited")).toBeInTheDocument(); + }); + + it("should sort budgets by updated_at descending", () => { + const budgets = [ + makeBudget({ budget_id: "budget-old", updated_at: "2024-01-01T00:00:00Z" }), + makeBudget({ budget_id: "budget-new", updated_at: "2024-06-01T00:00:00Z" }), + ]; + renderWithProviders(); + const rows = screen.getAllByRole("row").slice(1); + expect(within(rows[0]).getByText("budget-new")).toBeInTheDocument(); + expect(within(rows[1]).getByText("budget-old")).toBeInTheDocument(); + }); + + it("should call onEditClick from the actions menu", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByTestId("budget-actions-budget-1")); + await user.click(await screen.findByTestId("budget-action-edit")); + expect(defaultProps.onEditClick).toHaveBeenCalledWith(defaultProps.budgets[0]); + }); + + it("should call onDeleteClick from the actions menu", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByTestId("budget-actions-budget-1")); + await user.click(await screen.findByTestId("budget-action-delete")); + expect(defaultProps.onDeleteClick).toHaveBeenCalledWith(defaultProps.budgets[0]); + }); + + it("should not render the actions menu when the user cannot modify budgets", () => { + renderWithProviders(); + expect(screen.queryByTestId("budget-actions-budget-1")).not.toBeInTheDocument(); + }); + + it("should show skeleton rows when loading", () => { + renderWithProviders(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); + }); + + it("should show the empty state when there are no budgets", () => { + renderWithProviders(); + expect(screen.getByText("No budgets yet")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx new file mode 100644 index 00000000000..4bc06425f80 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx @@ -0,0 +1,57 @@ +"use client"; + +import { Inbox } from "lucide-react"; +import React, { useMemo } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; +import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; + +import { getBudgetTableColumns } from "./BudgetTableColumns"; + +interface BudgetTableProps { + budgets: budgetItem[]; + isLoading: boolean; + canModify: boolean; + onEditClick: (budget: budgetItem) => void; + onDeleteClick: (budget: budgetItem) => void; +} + +function EmptyState() { + return ( +
+
+ +
+
No budgets yet
+
+ Create a budget to set spend, TPM and RPM limits for customers. +
+
+ ); +} + +const BudgetTable: React.FC = ({ budgets, isLoading, canModify, onEditClick, onDeleteClick }) => { + const rows = useMemo( + () => [...budgets].sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()), + [budgets], + ); + + const columns = useMemo( + () => getBudgetTableColumns({ canModify, onEditClick, onDeleteClick }), + [canModify, onEditClick, onDeleteClick], + ); + + return ( + budget.budget_id || String(index)} + isLoading={isLoading} + loadingMessage="Loading budgets…" + noDataMessage={} + size="compact" + /> + ); +}; + +export default BudgetTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx new file mode 100644 index 00000000000..456ab9d6b68 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx @@ -0,0 +1,124 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { MoreHorizontal, Pencil, Trash2 } from "lucide-react"; + +import { IdCell, MoneyCell } from "@/components/shared/table_cells"; +import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; + +function RateLimitCell({ value }: { value: number | null }) { + if (value == null) { + return n/a; + } + return {value}; +} + +interface BudgetRowActionsProps { + budget: budgetItem; + onEditClick: (budget: budgetItem) => void; + onDeleteClick: (budget: budgetItem) => void; +} + +function BudgetRowActions({ budget, onEditClick, onDeleteClick }: BudgetRowActionsProps) { + return ( + + + + + + onEditClick(budget)}> + + Edit budget + + + onDeleteClick(budget)} + > + + Delete budget + + + + ); +} + +interface BudgetTableColumnsDeps { + canModify: boolean; + onEditClick: (budget: budgetItem) => void; + onDeleteClick: (budget: budgetItem) => void; +} + +export const getBudgetTableColumns = ({ + canModify, + onEditClick, + onDeleteClick, +}: BudgetTableColumnsDeps): ColumnDef[] => [ + { + id: "budget_id", + accessorKey: "budget_id", + meta: { title: "Budget ID" }, + header: "Budget ID", + size: 220, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "max_budget", + accessorKey: "max_budget", + meta: { title: "Max Budget", numeric: true }, + header: "Max Budget", + size: 120, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "tpm_limit", + accessorKey: "tpm_limit", + meta: { title: "TPM", numeric: true }, + header: "TPM", + size: 100, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "rpm_limit", + accessorKey: "rpm_limit", + meta: { title: "RPM", numeric: true }, + header: "RPM", + size: 100, + enableSorting: false, + cell: ({ row }) => , + }, + ...(canModify + ? [ + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + } satisfies ColumnDef, + ] + : []), +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx index f4d70a5e8f8..392616f1935 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx @@ -1,5 +1,6 @@ import { fireEvent, render, waitFor, screen } from "@testing-library/react"; import { act } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { afterEach, describe, expect, it, vi } from "vitest"; import BudgetPanel from "./budget_panel"; @@ -57,7 +58,8 @@ describe("Budget Panel", () => { }); }); - it("should open delete modal when clicking delete icon", async () => { + it("should open delete modal from the actions menu", async () => { + const user = userEvent.setup(); vi.mocked(useBudgets).mockReturnValue({ data: [ { @@ -77,11 +79,8 @@ describe("Budget Panel", () => { expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); }); - const deleteButton = screen.getByTestId("delete-budget-button"); - - act(() => { - fireEvent.click(deleteButton); - }); + await user.click(screen.getByTestId("budget-actions-budget-to-delete")); + await user.click(await screen.findByTestId("budget-action-delete")); await waitFor(() => { expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); @@ -89,6 +88,7 @@ describe("Budget Panel", () => { }); it("should successfully delete a budget", async () => { + const user = userEvent.setup(); const deleteMutateAsync = vi.fn().mockResolvedValue(undefined); vi.mocked(useBudgets).mockReturnValue({ data: [ @@ -113,17 +113,13 @@ describe("Budget Panel", () => { expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); }); - // Open delete modal - const deleteButton = screen.getByTestId("delete-budget-button"); - act(() => { - fireEvent.click(deleteButton); - }); + await user.click(screen.getByTestId("budget-actions-budget-to-delete")); + await user.click(await screen.findByTestId("budget-action-delete")); await waitFor(() => { expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); }); - // Confirm delete const confirmButton = screen.getByRole("button", { name: /delete/i }); act(() => { fireEvent.click(confirmButton); @@ -148,6 +144,7 @@ describe("Budget Panel", () => { }); it("should handle delete error", async () => { + const user = userEvent.setup(); const deleteMutateAsync = vi.fn().mockRejectedValue(new Error("Delete failed")); vi.mocked(useBudgets).mockReturnValue({ data: [ @@ -172,17 +169,13 @@ describe("Budget Panel", () => { expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); }); - // Open delete modal - const deleteButton = screen.getByTestId("delete-budget-button"); - act(() => { - fireEvent.click(deleteButton); - }); + await user.click(screen.getByTestId("budget-actions-budget-to-delete")); + await user.click(await screen.findByTestId("budget-action-delete")); await waitFor(() => { expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); }); - // Confirm delete const confirmButton = screen.getByRole("button", { name: /delete/i }); act(() => { fireEvent.click(confirmButton); @@ -193,7 +186,8 @@ describe("Budget Panel", () => { }); }); - it("should open edit modal when clicking edit icon", async () => { + it("should open edit modal from the actions menu", async () => { + const user = userEvent.setup(); vi.mocked(useBudgets).mockReturnValue({ data: [ { @@ -213,11 +207,8 @@ describe("Budget Panel", () => { expect(screen.getByText("budget-to-edit")).toBeInTheDocument(); }); - const editButton = screen.getByTestId("edit-budget-button"); - - act(() => { - fireEvent.click(editButton); - }); + await user.click(screen.getByTestId("budget-actions-budget-to-edit")); + await user.click(await screen.findByTestId("budget-action-edit")); await waitFor(() => { expect(screen.getByText("Edit Budget")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx index af15a99f0b4..6d5c0c7be08 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx @@ -3,30 +3,14 @@ * */ -import { - Button, - Card, - Tab, - TabGroup, - Table, - TableBody, - TableCell, - TableHead, - TableHeaderCell, - TableRow, - TabList, - TabPanel, - TabPanels, - Text, -} from "@tremor/react"; +import { Button, Tab, TabGroup, TabList, TabPanel, TabPanels, Text } from "@tremor/react"; import React, { useState } from "react"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import TableIconActionButton from "@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; import NotificationsManager from "@/components/molecules/notifications_manager"; import { useBudgets, useDeleteBudget, budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; -import { MoneyCell } from "@/components/shared/table_cells"; import BudgetModal from "./budget_modal"; +import BudgetTable from "./BudgetTable"; import EditBudgetModal from "./edit_budget_modal"; import { CREATE_END_USER_CURL_COMMAND, CHAT_COMPLETIONS_CURL_COMMAND, OPENAI_SDK_PYTHON_CODE } from "./constants"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; @@ -46,7 +30,7 @@ const BudgetPanel: React.FC = ({ accessToken }) => { // Admin Viewer follows the read-parity rule: see budgets, no writes. const canModify = isProxyAdminRole(userRole ?? ""); - const { data: budgetList = [] } = useBudgets(); + const { data: budgetList = [], isLoading } = useBudgets(); const deleteBudget = useDeleteBudget(); const handleEditCall = async (budget: budgetItem) => { @@ -109,51 +93,14 @@ const BudgetPanel: React.FC = ({ accessToken }) => { existingBudget={selectedBudget} /> )} - - Create a budget to assign to customers. - - - - Budget ID - Max Budget - TPM - RPM - - - - - {budgetList - .slice() - .sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()) - .map((value: budgetItem) => ( - - {value.budget_id} - - - - {value.tpm_limit ? value.tpm_limit : "n/a"} - {value.rpm_limit ? value.rpm_limit : "n/a"} - {canModify && ( - <> - handleEditCall(value)} - dataTestId="edit-budget-button" - /> - handleDeleteClick(value)} - dataTestId="delete-budget-button" - /> - - )} - - ))} - -
- + Create a budget to assign to customers. + => { +): UseQueryResult => { const { accessToken } = useAuthorized(); - return useQuery({ + return useQuery({ queryKey: deletedKeyKeys.list({ page, limit: pageSize, ...options }), queryFn: async () => await keyListCall(accessToken!, page, pageSize, { ...options, status: "deleted" }), enabled: Boolean(accessToken), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx similarity index 58% rename from ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_table.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx index 934114f5405..06c939aa151 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx @@ -1,47 +1,11 @@ import React from "react"; -import { screen } from "@testing-library/react"; +import { screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "@/../tests/test-utils"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import PolicyTable from "./policy_table"; +import PolicyTable from "./PolicyTable"; import { Policy } from "@/components/policies/types"; -vi.mock("@heroicons/react/outline", () => ({ - TrashIcon: function TrashIcon() { - return null; - }, - PencilIcon: function PencilIcon() { - return null; - }, - SwitchVerticalIcon: function SwitchVerticalIcon() { - return null; - }, - ChevronUpIcon: function ChevronUpIcon() { - return null; - }, - ChevronDownIcon: function ChevronDownIcon() { - return null; - }, -})); - -vi.mock("@tremor/react", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - Button: React.forwardRef(({ children, ...props }, ref) => - React.createElement("button", { ...props, ref }, children), - ), - Icon: ({ icon: IconComp, onClick, className }: any) => - React.createElement( - "button", - { type: "button", onClick, className }, - IconComp?.displayName ?? IconComp?.name ?? "icon", - ), - Tooltip: ({ children }: { children?: React.ReactNode }) => React.createElement(React.Fragment, null, children), - Badge: ({ children }: { children?: React.ReactNode }) => React.createElement("span", null, children), - }; -}); - const makePolicy = (overrides: Partial = {}): Policy => ({ policy_id: "policy-id-1", policy_name: "test-policy", @@ -71,20 +35,21 @@ describe("PolicyTable", () => { renderWithProviders(); expect(screen.getByText("Name")).toBeInTheDocument(); expect(screen.getByText("Description")).toBeInTheDocument(); - expect(screen.getByText("Actions")).toBeInTheDocument(); + expect(screen.getByText("Guardrails (Add)")).toBeInTheDocument(); + expect(screen.getByText("Created At")).toBeInTheDocument(); }); - it("should show a loading message when isLoading is true", () => { + it("should show skeleton rows when isLoading is true", () => { renderWithProviders(); - expect(screen.getByText(/loading/i)).toBeInTheDocument(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); }); - it("should show 'No policies found' when there are no policies", () => { + it("should show the empty state when there are no policies", () => { renderWithProviders(); - expect(screen.getByText(/no policies found/i)).toBeInTheDocument(); + expect(screen.getByText("No policies found")).toBeInTheDocument(); }); - it("should render a button with the policy name for each grouped policy", () => { + it("should render a clickable name cell for each grouped policy", () => { const policies = [ makePolicy({ policy_name: "alpha-policy", policy_id: "id-1" }), makePolicy({ policy_name: "beta-policy", policy_id: "id-2" }), @@ -94,7 +59,18 @@ describe("PolicyTable", () => { expect(screen.getByRole("button", { name: "beta-policy" })).toBeInTheDocument(); }); - it("should call onViewClick with the policy_id when the policy name button is clicked", async () => { + it("should sort rows by policy name ascending by default", () => { + const policies = [ + makePolicy({ policy_name: "zeta-policy", policy_id: "id-z" }), + makePolicy({ policy_name: "alpha-policy", policy_id: "id-a" }), + ]; + renderWithProviders(); + const rows = screen.getAllByRole("row").slice(1); + expect(within(rows[0]).getByText("alpha-policy")).toBeInTheDocument(); + expect(within(rows[1]).getByText("zeta-policy")).toBeInTheDocument(); + }); + + it("should call onViewClick with the policy_id when the policy name is clicked", async () => { const user = userEvent.setup(); const policy = makePolicy({ policy_name: "my-policy", policy_id: "view-id-1" }); renderWithProviders(); @@ -102,36 +78,46 @@ describe("PolicyTable", () => { expect(defaultProps.onViewClick).toHaveBeenCalledWith("view-id-1"); }); - it("should call onDeleteClick with policy_id and policy_name when the delete icon is clicked", async () => { + it("should call onDeleteClick with policy_id and policy_name from the actions menu", async () => { const user = userEvent.setup(); const policy = makePolicy({ policy_name: "del-policy", policy_id: "del-id-1" }); renderWithProviders(); - await user.click(screen.getByRole("button", { name: /TrashIcon/i })); + await user.click(screen.getByTestId("policy-actions-del-id-1")); + await user.click(await screen.findByTestId("policy-action-delete")); expect(defaultProps.onDeleteClick).toHaveBeenCalledWith("del-id-1", "del-policy"); }); - it("should call onEditClick with the policy when the edit icon is clicked", async () => { + it("should call onEditClick with the policy from the actions menu", async () => { const user = userEvent.setup(); const policy = makePolicy({ policy_name: "edit-policy", policy_id: "edit-id-1" }); renderWithProviders(); - await user.click(screen.getByRole("button", { name: /PencilIcon/i })); + await user.click(screen.getByTestId("policy-actions-edit-id-1")); + await user.click(await screen.findByTestId("policy-action-edit")); expect(defaultProps.onEditClick).toHaveBeenCalledWith(policy); }); - it("should not show admin action icons for non-admins", () => { + it("should not show the actions menu for non-admins", () => { const policy = makePolicy(); renderWithProviders(); - expect(screen.queryByRole("button", { name: /TrashIcon/i })).not.toBeInTheDocument(); - expect(screen.queryByRole("button", { name: /PencilIcon/i })).not.toBeInTheDocument(); + expect(screen.queryByTestId(`policy-actions-${policy.policy_id}`)).not.toBeInTheDocument(); }); it("should show a version badge when multiple versions of the same policy name exist", () => { - const policies = [ - makePolicy({ policy_name: "versioned", policy_id: "v1", version_status: "published", version_number: 1 }), - makePolicy({ policy_name: "versioned", policy_id: "v2", version_status: "production", version_number: 2 }), - ]; + const publishedVersion: Partial = { + policy_name: "versioned", + policy_id: "v1", + version_status: "published", + version_number: 1, + }; + const productionVersion: Partial = { + policy_name: "versioned", + policy_id: "v2", + version_status: "production", + version_number: 2, + }; + const policies = [makePolicy(publishedVersion), makePolicy(productionVersion)]; renderWithProviders(); - expect(screen.getByText(/2 version/i)).toBeInTheDocument(); + expect(screen.getByText("2 versions")).toBeInTheDocument(); }); it("should group policies with the same name into a single row", () => { @@ -140,10 +126,10 @@ describe("PolicyTable", () => { makePolicy({ policy_name: "shared", policy_id: "s2", version_status: "production" }), ]; renderWithProviders(); - expect(screen.getAllByRole("button", { name: "shared" })).toHaveLength(1); + expect(screen.getAllByText("shared")).toHaveLength(1); }); - it("should show an overflow tag when more than 2 guardrails_add exist", () => { + it("should show an overflow badge when more than 2 guardrails_add exist", () => { const policy = makePolicy({ guardrails_add: ["g1", "g2", "g3", "g4"] }); renderWithProviders(); expect(screen.getByText("+2")).toBeInTheDocument(); @@ -156,7 +142,7 @@ describe("PolicyTable", () => { makePolicy({ policy_name: "grouped", policy_id: "prod-id", version_status: "production" }), ]; renderWithProviders(); - await user.click(screen.getByRole("button", { name: "grouped" })); + await user.click(screen.getByRole("button", { name: /grouped/ })); expect(defaultProps.onViewClick).toHaveBeenCalledWith("prod-id"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx new file mode 100644 index 00000000000..d6e841c2119 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx @@ -0,0 +1,82 @@ +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; +import { Policy } from "@/components/policies/types"; + +import { getPolicyTableColumns, PolicyRow } from "./PolicyTableColumns"; + +/** One row per policy name; primaryPolicy is used for display and for Edit (FlowBuilder loads all versions) */ +function groupPoliciesByName(policies: Policy[]): PolicyRow[] { + const names = Array.from(new Set(policies.map((policy) => policy.policy_name || "(unnamed)"))); + return names.map((policyName) => { + const versions = policies.filter((policy) => (policy.policy_name || "(unnamed)") === policyName); + const primary = + versions.find((version) => version.version_status === "production") ?? + [...versions].sort((a, b) => (b.version_number ?? 0) - (a.version_number ?? 0))[0]; + return { policy_name: policyName, primaryPolicy: primary, versionCount: versions.length }; + }); +} + +interface PolicyTableProps { + policies: Policy[]; + isLoading: boolean; + onDeleteClick: (policyId: string, policyName: string) => void; + onEditClick: (policy: Policy) => void; + onViewClick: (policyId: string) => void; + isAdmin?: boolean; +} + +const DEFAULT_SORTING: SortingState = [{ id: "policy_name", desc: false }]; + +function EmptyState() { + return ( +
+
+ +
+
No policies found
+
+ Create a policy to bundle guardrails and apply them across teams. +
+
+ ); +} + +const PolicyTable: React.FC = ({ + policies, + isLoading, + onDeleteClick, + onEditClick, + onViewClick, + isAdmin = false, +}) => { + const [sorting, setSorting] = useState(DEFAULT_SORTING); + + const rows = useMemo(() => groupPoliciesByName(policies), [policies]); + + const columns = useMemo(() => { + const deps = { isAdmin, onViewClick, onEditClick, onDeleteClick }; + return getPolicyTableColumns(deps); + }, [isAdmin, onViewClick, onEditClick, onDeleteClick]); + + return ( + row.policy_name} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + isLoading={isLoading} + loadingMessage="Loading policies…" + noDataMessage={} + size="compact" + /> + ); +}; + +export default PolicyTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx new file mode 100644 index 00000000000..dd036d83283 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx @@ -0,0 +1,207 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { MoreHorizontal, Pencil, Trash2 } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell, StatusBadge } from "@/components/shared/table_cells"; +import { Policy } from "@/components/policies/types"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; + +export interface PolicyRow { + policy_name: string; + primaryPolicy: Policy; + versionCount: number; +} + +function GuardrailChips({ guardrails, tone }: { guardrails: string[]; tone: "success" | "error" }) { + if (guardrails.length === 0) { + return -; + } + return ( +
+ {guardrails.slice(0, 2).map((guardrail) => ( + + ))} + {guardrails.length > 2 && ( + + )} +
+ ); +} + +interface PolicyRowActionsProps { + policy: Policy; + onEditClick: (policy: Policy) => void; + onDeleteClick: (policyId: string, policyName: string) => void; +} + +function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActionsProps) { + return ( + + + + + + onEditClick(policy)}> + + Edit policy + + + onDeleteClick(policy.policy_id, policy.policy_name || "Unnamed Policy")} + > + + Delete policy + + + + ); +} + +interface PolicyTableColumnsDeps { + isAdmin: boolean; + onViewClick: (policyId: string) => void; + onEditClick: (policy: Policy) => void; + onDeleteClick: (policyId: string, policyName: string) => void; +} + +export const getPolicyTableColumns = ({ + isAdmin, + onViewClick, + onEditClick, + onDeleteClick, +}: PolicyTableColumnsDeps): ColumnDef[] => [ + { + id: "policy_name", + accessorKey: "policy_name", + meta: { title: "Name", skeleton: "twoLine" }, + header: ({ column }) => , + size: 220, + enableSorting: true, + cell: ({ row }) => ( + 1 ? ( + + ) : undefined + } + onClick={() => onViewClick(row.original.primaryPolicy.policy_id)} + /> + ), + }, + { + id: "description", + accessorFn: (row) => row.primaryPolicy.description ?? "", + meta: { title: "Description" }, + header: "Description", + size: 220, + enableSorting: false, + cell: ({ row }) => { + const description = row.original.primaryPolicy.description; + if (!description) { + return -; + } + return ( + + {description} + + ); + }, + }, + { + id: "inherit", + accessorFn: (row) => row.primaryPolicy.inherit ?? "", + meta: { title: "Inherits From", skeleton: "badge" }, + header: "Inherits From", + size: 150, + enableSorting: false, + cell: ({ row }) => { + const inherit = row.original.primaryPolicy.inherit; + if (!inherit) { + return -; + } + return ; + }, + }, + { + id: "guardrails_add", + meta: { title: "Guardrails (Add)", skeleton: "chips" }, + header: "Guardrails (Add)", + size: 180, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "guardrails_remove", + meta: { title: "Guardrails (Remove)", skeleton: "chips" }, + header: "Guardrails (Remove)", + size: 180, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "model_condition", + meta: { title: "Model Condition" }, + header: "Model Condition", + size: 160, + enableSorting: false, + cell: ({ row }) => { + const model = row.original.primaryPolicy.condition?.model; + if (!model) { + return -; + } + return ( + + {model} + + ); + }, + }, + { + id: "created_at", + accessorFn: (row) => row.primaryPolicy.created_at ?? "", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 150, + enableSorting: true, + cell: ({ row }) => , + }, + ...(isAdmin + ? [ + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + } satisfies ColumnDef, + ] + : []), +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/index.tsx index 0545e2f86dc..df55d2c386b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/index.tsx @@ -5,7 +5,7 @@ import { Alert } from "antd"; import MessageManager from "@/components/molecules/message_manager"; import { InfoCircleOutlined } from "@ant-design/icons"; import { isAdminRole } from "@/utils/roles"; -import PolicyTable from "./policy_table"; +import PolicyTable from "./PolicyTable"; import PolicyInfoView from "./policy_info"; import AddPolicyForm from "./add_policy_form"; import { FlowBuilderPage } from "./pipeline_flow_builder"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_table.tsx deleted file mode 100644 index 278773069e6..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_table.tsx +++ /dev/null @@ -1,326 +0,0 @@ -import React, { useMemo, useState } from "react"; -import { Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Icon, Button, Badge } from "@tremor/react"; -import { TrashIcon, PencilIcon, SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon } from "@heroicons/react/outline"; -import { Tooltip, Tag } from "antd"; -import { - ColumnDef, - flexRender, - getCoreRowModel, - getSortedRowModel, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { DateCell } from "@/components/shared/table_cells"; -import { Policy } from "@/components/policies/types"; - -/** One row per policy name; primaryPolicy is used for display and for Edit (FlowBuilder loads all versions) */ -interface PolicyRow { - policy_name: string; - primaryPolicy: Policy; - versionCount: number; -} - -function groupPoliciesByName(policies: Policy[]): PolicyRow[] { - const byName = new Map(); - for (const p of policies) { - const name = p.policy_name || "(unnamed)"; - if (!byName.has(name)) byName.set(name, []); - byName.get(name)!.push(p); - } - const rows: PolicyRow[] = []; - for (const [policyName, versions] of byName) { - // Prefer production, then highest version_number - const primary = - versions.find((v) => v.version_status === "production") ?? - [...versions].sort((a, b) => (b.version_number ?? 0) - (a.version_number ?? 0))[0] ?? - versions[0]; - rows.push({ policy_name: policyName, primaryPolicy: primary, versionCount: versions.length }); - } - return rows.sort((a, b) => a.policy_name.localeCompare(b.policy_name)); -} - -interface PolicyTableProps { - policies: Policy[]; - isLoading: boolean; - onDeleteClick: (policyId: string, policyName: string) => void; - onEditClick: (policy: Policy) => void; - onViewClick: (policyId: string) => void; - isAdmin?: boolean; -} - -const PolicyTable: React.FC = ({ - policies, - isLoading, - onDeleteClick, - onEditClick, - onViewClick, - isAdmin = false, -}) => { - const [sorting, setSorting] = useState([{ id: "policy_name", desc: false }]); - - const rows = useMemo(() => groupPoliciesByName(policies), [policies]); - - const columns: ColumnDef[] = [ - { - header: "Name", - accessorKey: "policy_name", - cell: ({ row }) => { - const { primaryPolicy, versionCount } = row.original; - return ( -
- 1 ? ` (${versionCount} versions)` : ""}`} - > - - - {versionCount > 1 && ( - - {versionCount} version{versionCount !== 1 ? "s" : ""} - - )} -
- ); - }, - }, - { - header: "Description", - accessorFn: (row) => row.primaryPolicy.description ?? "", - cell: ({ row }) => { - const policy = row.original.primaryPolicy; - return ( - - {policy.description || "-"} - - ); - }, - }, - { - header: "Inherits From", - accessorFn: (row) => row.primaryPolicy.inherit ?? "", - cell: ({ row }) => { - const policy = row.original.primaryPolicy; - return policy.inherit ? ( - - {policy.inherit} - - ) : ( - - - ); - }, - }, - { - header: "Guardrails (Add)", - accessorFn: (row) => (row.primaryPolicy.guardrails_add ?? []).join(", "), - cell: ({ row }) => { - const policy = row.original.primaryPolicy; - const guardrails = policy.guardrails_add || []; - if (guardrails.length === 0) { - return -; - } - return ( -
- {guardrails.slice(0, 2).map((g, i) => ( - - {g} - - ))} - {guardrails.length > 2 && ( - - +{guardrails.length - 2} - - )} -
- ); - }, - }, - { - header: "Guardrails (Remove)", - accessorFn: (row) => (row.primaryPolicy.guardrails_remove ?? []).join(", "), - cell: ({ row }) => { - const policy = row.original.primaryPolicy; - const guardrails = policy.guardrails_remove || []; - if (guardrails.length === 0) { - return -; - } - return ( -
- {guardrails.slice(0, 2).map((g, i) => ( - - {g} - - ))} - {guardrails.length > 2 && ( - - +{guardrails.length - 2} - - )} -
- ); - }, - }, - { - header: "Model Condition", - accessorFn: (row) => { - const m = row.primaryPolicy.condition?.model; - return typeof m === "string" ? m : JSON.stringify(m ?? ""); - }, - cell: ({ row }) => { - const policy = row.original.primaryPolicy; - const modelCondition = policy.condition?.model; - if (!modelCondition) { - return -; - } - return ( - - - {typeof modelCondition === "string" - ? modelCondition.length > 20 - ? modelCondition.slice(0, 20) + "..." - : modelCondition - : "Multiple"} - - - ); - }, - }, - { - header: "Created At", - id: "created_at", - accessorFn: (row) => row.primaryPolicy.created_at ?? "", - cell: ({ row }) => , - }, - { - id: "actions", - header: "Actions", - cell: ({ row }) => { - const { primaryPolicy } = row.original; - const policy = primaryPolicy; - return ( -
- {isAdmin && ( - <> - - onEditClick(policy)} - className="cursor-pointer hover:text-blue-500" - /> - - - - policy.policy_id && onDeleteClick(policy.policy_id, policy.policy_name || "Unnamed Policy") - } - className="cursor-pointer hover:text-red-500" - /> - - - )} -
- ); - }, - }, - ]; - - const table = useReactTable({ - data: rows, - columns, - state: { - sorting, - }, - onSortingChange: setSorting, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - enableSorting: true, - }); - - return ( -
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - -
-
- {header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())} -
- {header.id !== "actions" && ( -
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
- )} -
- - ))} - - ))} - - - {isLoading ? ( - - -
-

Loading...

-
- - - ) : rows.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No policies found

-
- - - )} - -
-
-
- ); -}; - -export default PolicyTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolColumn.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolColumn.tsx deleted file mode 100644 index 3d9fdb2866e..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolColumn.tsx +++ /dev/null @@ -1,104 +0,0 @@ -import { Tag } from "antd"; -import { ColumnsType } from "antd/es/table"; -import TableIconActionButton from "@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; -import { DateCell, IdCell } from "@/components/shared/table_cells"; -import { SearchTool } from "./types"; - -export const searchToolColumns = ( - onView: (searchToolId: string) => void, - onEdit: (searchToolId: string) => void, - onDelete: (searchToolId: string) => void, - availableProviders: Array<{ provider_name: string; ui_friendly_name: string }>, -): ColumnsType => [ - { - title: "Search Tool ID", - dataIndex: "search_tool_id", - key: "search_tool_id", - render: (_, tool) => { - const isFromConfig = tool.is_from_config; - - if (isFromConfig) { - return -; - } - - return ; - }, - }, - { - title: "Name", - dataIndex: "search_tool_name", - key: "search_tool_name", - render: (name: string) => {name}, - }, - { - title: "Provider", - key: "provider", - render: (_, tool) => { - const provider = tool.litellm_params.search_provider; - const providerInfo = availableProviders.find((p) => p.provider_name === provider); - const displayName = providerInfo?.ui_friendly_name || provider; - - return {displayName}; - }, - }, - { - title: "Created At", - dataIndex: "created_at", - key: "created_at", - render: (_, tool) => { - return ; - }, - }, - { - title: "Updated At", - dataIndex: "updated_at", - key: "updated_at", - render: (_, tool) => { - return ; - }, - }, - { - title: "Source", - key: "source", - render: (_, tool) => { - const isFromConfig = tool.is_from_config ?? false; - - return {isFromConfig ? "Config" : "DB"}; - }, - }, - { - title: "Actions", - key: "actions", - render: (_, tool) => { - const toolId = tool.search_tool_id; - const isFromConfig = tool.is_from_config ?? false; - - return ( -
- { - if (toolId && !isFromConfig) { - onEdit(toolId); - } - }} - /> - { - if (toolId && !isFromConfig) { - onDelete(toolId); - } - }} - /> -
- ); - }, - }, -]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.test.tsx new file mode 100644 index 00000000000..c4c05f97dce --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.test.tsx @@ -0,0 +1,117 @@ +import { screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "@/../tests/test-utils"; +import SearchToolTable from "./SearchToolTable"; +import { AvailableSearchProvider, SearchTool } from "./types"; + +const makeSearchTool = (overrides: Partial = {}): SearchTool => ({ + search_tool_id: "tool-1", + search_tool_name: "Perplexity Search", + litellm_params: { + search_provider: "perplexity", + }, + created_at: "2024-01-15T10:30:00Z", + updated_at: "2024-01-16T10:30:00Z", + ...overrides, +}); + +const availableProviders: AvailableSearchProvider[] = [ + { provider_name: "perplexity", ui_friendly_name: "Perplexity AI" }, +]; + +const defaultProps = { + searchTools: [makeSearchTool()], + isLoading: false, + availableProviders, + onView: vi.fn(), + onEdit: vi.fn(), + onDelete: vi.fn(), +}; + +describe("SearchToolTable", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should display search tool information with the friendly provider name", () => { + renderWithProviders(); + expect(screen.getByText("Perplexity Search")).toBeInTheDocument(); + expect(screen.getByText("tool-1")).toBeInTheDocument(); + expect(screen.getByText("Perplexity AI")).toBeInTheDocument(); + expect(screen.getByText("DB")).toBeInTheDocument(); + }); + + it("should call onView when the search tool ID is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /tool-1/ })); + expect(defaultProps.onView).toHaveBeenCalledWith("tool-1"); + }); + + it("should call onEdit and onDelete from the actions menu for a DB tool", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByTestId("search-tool-actions-tool-1")); + await user.click(await screen.findByTestId("search-tool-action-edit")); + expect(defaultProps.onEdit).toHaveBeenCalledWith("tool-1"); + + await user.click(screen.getByTestId("search-tool-actions-tool-1")); + await user.click(await screen.findByTestId("search-tool-action-delete")); + expect(defaultProps.onDelete).toHaveBeenCalledWith("tool-1"); + }); + + it("should show a dash instead of a clickable ID for config tools", () => { + const configTool = makeSearchTool({ search_tool_id: "config-tool", is_from_config: true }); + renderWithProviders(); + expect(screen.queryByRole("button", { name: /config-tool/ })).not.toBeInTheDocument(); + }); + + it("should disable Edit and Delete for config tools and suppress their callbacks", async () => { + const user = userEvent.setup(); + const configTool = makeSearchTool({ search_tool_id: "config-tool", is_from_config: true }); + renderWithProviders(); + + await user.click(screen.getByTestId("search-tool-actions-config-tool")); + + const editItem = await screen.findByTestId("search-tool-action-edit"); + const deleteItem = await screen.findByTestId("search-tool-action-delete"); + expect(editItem).toHaveAttribute("data-disabled"); + expect(deleteItem).toHaveAttribute("data-disabled"); + + await user.click(editItem); + await user.click(deleteItem); + expect(defaultProps.onEdit).not.toHaveBeenCalled(); + expect(defaultProps.onDelete).not.toHaveBeenCalled(); + }); + + it("should sort tools by created_at descending by default", () => { + const tools = [ + makeSearchTool({ + search_tool_id: "tool-old", + search_tool_name: "older-tool", + created_at: "2024-01-01T00:00:00Z", + }), + makeSearchTool({ + search_tool_id: "tool-new", + search_tool_name: "newer-tool", + created_at: "2024-06-01T00:00:00Z", + }), + ]; + renderWithProviders(); + const rows = screen.getAllByRole("row").slice(1); + expect(within(rows[0]).getByText("newer-tool")).toBeInTheDocument(); + expect(within(rows[1]).getByText("older-tool")).toBeInTheDocument(); + }); + + it("should show skeleton rows when loading", () => { + renderWithProviders(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); + }); + + it("should show the empty state when there are no search tools", () => { + renderWithProviders(); + expect(screen.getByText("No search tools configured")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.tsx new file mode 100644 index 00000000000..70fc6a376df --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.tsx @@ -0,0 +1,66 @@ +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; + +import { getSearchToolTableColumns, searchToolKey } from "./SearchToolTableColumns"; +import { AvailableSearchProvider, SearchTool } from "./types"; + +interface SearchToolTableProps { + searchTools: SearchTool[]; + isLoading: boolean; + availableProviders: AvailableSearchProvider[]; + onView: (searchToolId: string) => void; + onEdit: (searchToolId: string) => void; + onDelete: (searchToolId: string) => void; +} + +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +function EmptyState() { + return ( +
+
+ +
+
No search tools configured
+
Add a search tool to enable web search for your models.
+
+ ); +} + +const SearchToolTable: React.FC = ({ + searchTools, + isLoading, + availableProviders, + onView, + onEdit, + onDelete, +}) => { + const [sorting, setSorting] = useState(DEFAULT_SORTING); + + const columns = useMemo(() => { + const deps = { availableProviders, onView, onEdit, onDelete }; + return getSearchToolTableColumns(deps); + }, [availableProviders, onView, onEdit, onDelete]); + + return ( + searchToolKey(tool) || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + isLoading={isLoading} + loadingMessage="Loading search tools…" + noDataMessage={} + size="compact" + /> + ); +}; + +export default SearchToolTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTableColumns.tsx new file mode 100644 index 00000000000..550bf3dc7bd --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTableColumns.tsx @@ -0,0 +1,168 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { MoreHorizontal, Pencil, Trash2 } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell, StatusBadge } from "@/components/shared/table_cells"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; + +import { AvailableSearchProvider, SearchTool } from "./types"; + +const CONFIG_EDIT_HINT = "Config search tools cannot be edited on the dashboard. Please edit the config file."; +const CONFIG_DELETE_HINT = "Config search tools cannot be deleted on the dashboard. Please edit the config file."; + +export const searchToolKey = (tool: SearchTool): string => tool.search_tool_id || tool.search_tool_name; + +interface SearchToolRowActionsProps { + tool: SearchTool; + onEdit: (searchToolId: string) => void; + onDelete: (searchToolId: string) => void; +} + +function SearchToolRowActions({ tool, onEdit, onDelete }: SearchToolRowActionsProps) { + const isFromConfig = tool.is_from_config ?? false; + const toolId = tool.search_tool_id; + + return ( + + + + + + toolId && onEdit(toolId)} + > + + Edit search tool + + + toolId && onDelete(toolId)} + > + + Delete search tool + + + + ); +} + +interface SearchToolTableColumnsDeps { + availableProviders: AvailableSearchProvider[]; + onView: (searchToolId: string) => void; + onEdit: (searchToolId: string) => void; + onDelete: (searchToolId: string) => void; +} + +export const getSearchToolTableColumns = ({ + availableProviders, + onView, + onEdit, + onDelete, +}: SearchToolTableColumnsDeps): ColumnDef[] => [ + { + id: "search_tool_id", + accessorKey: "search_tool_id", + meta: { title: "Search Tool ID" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => { + const tool = row.original; + const toolId = tool.search_tool_id; + if (tool.is_from_config || !toolId) { + return -; + } + return ( + onView(toolId)} /> + ); + }, + }, + { + id: "search_tool_name", + accessorKey: "search_tool_name", + meta: { title: "Name" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => ( + + {row.original.search_tool_name || "-"} + + ), + }, + { + id: "provider", + meta: { title: "Provider" }, + header: "Provider", + size: 160, + enableSorting: false, + cell: ({ row }) => { + const provider = row.original.litellm_params.search_provider; + const providerInfo = availableProviders.find((candidate) => candidate.provider_name === provider); + return {providerInfo?.ui_friendly_name || provider}; + }, + }, + { + id: "created_at", + accessorKey: "created_at", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 130, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "updated_at", + accessorKey: "updated_at", + meta: { title: "Updated At" }, + header: ({ column }) => , + size: 130, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "source", + meta: { title: "Source", skeleton: "badge" }, + header: "Source", + size: 100, + enableSorting: false, + cell: ({ row }) => { + const isFromConfig = row.original.is_from_config ?? false; + return ; + }, + }, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchTools.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchTools.tsx index 65aee211271..c9e2a46a861 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchTools.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchTools.tsx @@ -1,8 +1,7 @@ import { isAdminRole } from "@/utils/roles"; -import { LoadingOutlined } from "@ant-design/icons"; import { useQuery } from "@tanstack/react-query"; import { Button, Text, Title } from "@tremor/react"; -import { Form, Input, Modal, Select, Spin, Table } from "antd"; +import { Form, Input, Modal, Select } from "antd"; import React, { useState } from "react"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; import NotificationsManager from "@/components/molecules/notifications_manager"; @@ -13,7 +12,7 @@ import { updateSearchTool, } from "@/components/networking"; import CreateSearchTool from "./CreateSearchTools"; -import { searchToolColumns } from "./SearchToolColumn"; +import SearchToolTable from "./SearchToolTable"; import { SearchToolView } from "./SearchToolView"; import { AvailableSearchProvider, SearchTool } from "./types"; @@ -58,34 +57,29 @@ const SearchTools: React.FC = ({ accessToken, userRole, userID const [isEditModalVisible, setEditModalVisible] = useState(false); const [form] = Form.useForm(); - const columns = React.useMemo( - () => - searchToolColumns( - (toolId: string) => { - setSelectedToolId(toolId); - setEditTool(false); - }, - (toolId: string) => { - const tool = searchTools?.find((t) => t.search_tool_id === toolId); - if (tool) { - form.setFieldsValue({ - search_tool_name: tool.search_tool_name, - search_provider: tool.litellm_params.search_provider, - api_key: tool.litellm_params.api_key, - api_base: tool.litellm_params.api_base, - timeout: tool.litellm_params.timeout, - max_retries: tool.litellm_params.max_retries, - description: tool.search_tool_info?.description, - }); - setSelectedToolId(toolId); - setEditModalVisible(true); - } - }, - handleDelete, - availableProviders, - ), - [availableProviders, searchTools, form], - ); + const handleView = (toolId: string) => { + setSelectedToolId(toolId); + setEditTool(false); + }; + + const handleEditOpen = (toolId: string) => { + const tool = searchTools?.find((t) => t.search_tool_id === toolId); + if (!tool) { + return; + } + const editFormValues = { + search_tool_name: tool.search_tool_name, + search_provider: tool.litellm_params.search_provider, + api_key: tool.litellm_params.api_key, + api_base: tool.litellm_params.api_base, + timeout: tool.litellm_params.timeout, + max_retries: tool.litellm_params.max_retries, + description: tool.search_tool_info?.description, + }; + form.setFieldsValue(editFormValues); + setSelectedToolId(toolId); + setEditModalVisible(true); + }; function handleDelete(toolId: string) { setToolToDelete(toolId); @@ -220,19 +214,14 @@ const SearchTools: React.FC = ({ accessToken, userRole, userID /> ) : (
- } size="large"> - record.search_tool_id || record.search_tool_name} - pagination={false} - locale={{ - emptyText: "No search tools configured", - }} - size="small" - /> - + ); diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx index 46df98a31de..05361f1c858 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx @@ -1,4 +1,5 @@ import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { vi, it, expect, beforeEach, MockedFunction } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import DeletedKeysPage from "./DeletedKeysPage"; @@ -13,6 +14,9 @@ const mockUseDeletedKeys = useDeletedKeys as MockedFunction { current_page: 1, total_pages: 1, }, - isPending: false, - isFetching: false, - } as any); + isLoading: false, + } as unknown as ReturnType); }); it("should render DeletedKeysPage component", () => { @@ -89,14 +92,32 @@ it("should render DeletedKeysPage component", () => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); -it("should handle loading state", () => { +it("should show skeleton rows while the initial load is pending", () => { mockUseDeletedKeys.mockReturnValue({ data: undefined, - isPending: true, - isFetching: false, - } as any); + isLoading: true, + } as unknown as ReturnType); renderWithProviders(); - expect(screen.getByText("🚅 Loading keys...")).toBeInTheDocument(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); +}); + +it("should request the next page from the hook when the pagination next button is clicked", async () => { + const user = userEvent.setup(); + mockUseDeletedKeys.mockReturnValue({ + data: { + keys: [mockDeletedKey], + total_count: 120, + current_page: 1, + total_pages: 3, + }, + isLoading: false, + } as unknown as ReturnType); + + renderWithProviders(); + + expect(mockUseDeletedKeys).toHaveBeenLastCalledWith(1, 50); + await user.click(screen.getByTestId("pagination-next")); + expect(mockUseDeletedKeys).toHaveBeenLastCalledWith(2, 50); }); diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.tsx index 8523710719e..78d500f5e64 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.tsx @@ -1,5 +1,6 @@ "use client"; import { useState } from "react"; +import { PaginationState } from "@tanstack/react-table"; import { Alert } from "antd"; import { useDeletedKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; @@ -7,10 +8,9 @@ import { DeletedKeysTable } from "./DeletedKeysTable/DeletedKeysTable"; export default function DeletedKeysPage() { const { premiumUser } = useAuthorized(); - const [pageIndex, setPageIndex] = useState(0); - const [pageSize] = useState(50); + const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: 50 }); - const { data: keysData, isPending: isLoading, isFetching } = useDeletedKeys(pageIndex + 1, pageSize); + const { data: keysData, isLoading } = useDeletedKeys(pagination.pageIndex + 1, pagination.pageSize); return (
@@ -27,10 +27,8 @@ export default function DeletedKeysPage() { keys={keysData?.keys || []} totalCount={keysData?.total_count || 0} isLoading={isLoading} - isFetching={isFetching} - pageIndex={pageIndex} - pageSize={pageSize} - onPageChange={setPageIndex} + pagination={pagination} + onPaginationChange={setPagination} />
); 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 081ae0a80b5..7e30ef2c135 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx @@ -1,101 +1,88 @@ -import { screen } from "@testing-library/react"; +import { screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { vi, it, expect, beforeEach } from "vitest"; import { renderWithProviders } from "../../../../tests/test-utils"; import { DeletedKeysTable } from "./DeletedKeysTable"; import { DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; -const mockDeletedKey: DeletedKeyResponse = { - token: "sk-1234567890abcdef", - token_id: "key-1", - key_name: "test-key", - key_alias: "Test Key Alias", - spend: 5.5, - max_budget: 100, - expires: "2024-12-31T23:59:59Z", - models: ["gpt-3.5-turbo"], - aliases: {}, - config: {}, - user_id: "user-1", - team_id: "team-1", - max_parallel_requests: 10, - metadata: {}, - tpm_limit: 1000, - rpm_limit: 100, - duration: "30d", - budget_duration: "1m", - budget_reset_at: "2024-12-01T00:00:00Z", - allowed_cache_controls: [], - allowed_routes: [], - permissions: {}, - model_spend: {}, - model_max_budget: {}, - soft_budget_cooldown: false, - blocked: false, - litellm_budget_table: {}, - organization_id: "org-1", - created_at: "2024-11-01T10:00:00Z", - updated_at: "2024-11-15T10:00:00Z", - team_spend: 5.5, - team_alias: "Test Team", - team_tpm_limit: 5000, - team_rpm_limit: 500, - team_max_budget: 500, - team_models: ["gpt-3.5-turbo"], - team_blocked: false, - soft_budget: 50, - team_model_aliases: {}, - team_member_spend: 0, - team_metadata: {}, - end_user_id: "end-user-1", - end_user_tpm_limit: 100, - end_user_rpm_limit: 10, - end_user_max_budget: 10, - last_refreshed_at: Date.now(), - api_key: "sk-1234567890abcdef", - user_role: "user", - rpm_limit_per_model: {}, - tpm_limit_per_model: {}, - user_tpm_limit: 1000, - user_rpm_limit: 100, - user_email: "user@example.com", - deleted_at: "2024-11-15T10:00:00Z", - deleted_by: "user-1", +const makeDeletedKey = (overrides: Partial = {}): DeletedKeyResponse => + ({ + token: "sk-1234567890abcdef", + token_id: "key-1", + key_name: "test-key", + key_alias: "Test Key Alias", + spend: 5.5, + max_budget: 100, + models: ["gpt-3.5-turbo"], + user_id: "user-1", + team_id: "team-1", + organization_id: "org-1", + created_at: "2024-11-01T10:00:00Z", + updated_at: "2024-11-15T10:00:00Z", + created_by: "creator-1", + team_alias: "Test Team", + user_email: "user@example.com", + deleted_at: "2024-11-15T10:00:00Z", + deleted_by: "user-1", + ...overrides, + }) as DeletedKeyResponse; + +const defaultProps = { + keys: [makeDeletedKey()], + totalCount: 1, + isLoading: false, + pagination: { pageIndex: 0, pageSize: 50 }, + onPaginationChange: vi.fn(), }; beforeEach(() => { vi.clearAllMocks(); }); -it("should render DeletedKeysTable component", () => { - renderWithProviders( - , - ); - - expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); -}); - -it("should display key information correctly", () => { - renderWithProviders( - , - ); +it("should display key information", () => { + renderWithProviders(); expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); expect(screen.getByText("sk-1234567890abcdef")).toBeInTheDocument(); - expect(screen.getByText("Showing 1 - 1 of 1 results")).toBeInTheDocument(); + expect(screen.getByText("user@example.com")).toBeInTheDocument(); +}); + +it("should show the total count in the pagination footer", () => { + renderWithProviders(); + + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-50 of 120"); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 3"); +}); + +it("should propagate pagination changes when the next page button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByTestId("pagination-next")); + + expect(defaultProps.onPaginationChange).toHaveBeenCalled(); +}); + +it("should sort the current page by deleted_at descending by default", () => { + const keys = [ + makeDeletedKey({ token: "sk-older", key_alias: "older-key", deleted_at: "2024-01-01T10:00:00Z" }), + makeDeletedKey({ token: "sk-newer", key_alias: "newer-key", deleted_at: "2024-06-01T10:00:00Z" }), + ]; + renderWithProviders(); + + const rows = screen.getAllByRole("row").slice(1); + expect(within(rows[0]).getByText("newer-key")).toBeInTheDocument(); + expect(within(rows[1]).getByText("older-key")).toBeInTheDocument(); +}); + +it("should show skeleton rows when loading", () => { + renderWithProviders(); + + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); +}); + +it("should show the empty state when there are no deleted keys", () => { + renderWithProviders(); + + expect(screen.getByText("No deleted keys found")).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.tsx index d4a120d0589..bc6941a7860 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.tsx @@ -1,364 +1,63 @@ "use client"; -import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells"; -import { ChevronDownIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; -import { - ColumnDef, - flexRender, - getCoreRowModel, - getPaginationRowModel, - getSortedRowModel, - PaginationState, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow } from "@tremor/react"; -import { Tooltip } from "antd"; -import React, { useState } from "react"; -import { KeyResponse } from "../../key_team_helpers/key_list"; + +import { OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; +import { DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; + +import { getDeletedKeysTableColumns } from "./DeletedKeysTableColumns"; interface DeletedKeysTableProps { - keys: KeyResponse[]; + keys: DeletedKeyResponse[]; totalCount: number; isLoading: boolean; - isFetching: boolean; - pageIndex: number; - pageSize: number; - onPageChange: (pageIndex: number) => void; + pagination: PaginationState; + onPaginationChange: OnChangeFn; +} + +const DEFAULT_SORTING: SortingState = [{ id: "deleted_at", desc: true }]; + +function EmptyState() { + return ( +
+
+ +
+
No deleted keys found
+
Keys deleted from this proxy will show up here.
+
+ ); } export function DeletedKeysTable({ keys, totalCount, isLoading, - isFetching, - pageIndex, - pageSize, - onPageChange, + pagination, + onPaginationChange, }: DeletedKeysTableProps) { - const [sorting, setSorting] = useState([ - { - id: "deleted_at", - desc: true, - }, - ]); + const [sorting, setSorting] = useState(DEFAULT_SORTING); - const [tablePagination, setTablePagination] = useState({ - pageIndex, - pageSize, - }); - - // Sync pagination state when prop changes - React.useEffect(() => { - setTablePagination({ pageIndex, pageSize }); - }, [pageIndex, pageSize]); - - const columns: ColumnDef[] = [ - { - id: "token", - accessorKey: "token", - header: "Key ID", - size: 150, - maxSize: 250, - cell: (info) => , - }, - { - id: "key_alias", - accessorKey: "key_alias", - header: "Key Alias", - size: 150, - maxSize: 200, - cell: (info) => { - const value = info.getValue() as string; - return ( - - {value ?? "-"} - - ); - }, - }, - { - id: "team_alias", - accessorKey: "team_alias", - header: "Team Alias", - size: 120, - maxSize: 180, - cell: (info) => { - const value = info.getValue() as string; - return {value || "-"}; - }, - }, - { - id: "spend", - accessorKey: "spend", - header: "Spend (USD)", - size: 100, - maxSize: 140, - cell: (info) => , - }, - { - id: "max_budget", - accessorKey: "max_budget", - header: "Budget (USD)", - size: 110, - maxSize: 150, - cell: (info) => ( - - ), - }, - { - id: "user_email", - accessorKey: "user_email", - header: "User Email", - size: 160, - maxSize: 250, - cell: (info) => { - const value = info.getValue() as string; - return ( - - {value ?? "-"} - - ); - }, - }, - { - id: "user_id", - accessorKey: "user_id", - header: "User ID", - size: 120, - maxSize: 200, - cell: (info) => , - }, - { - id: "created_at", - accessorKey: "created_at", - header: "Created At", - size: 120, - maxSize: 140, - cell: (info) => , - }, - { - id: "created_by", - accessorKey: "created_by", - header: "Created By", - size: 120, - maxSize: 180, - cell: (info) => { - const value = (info.row.original as any).created_by as string | null | undefined; - return ( - - {value || "-"} - - ); - }, - }, - { - id: "deleted_at", - accessorKey: "deleted_at", - header: "Deleted At", - size: 120, - maxSize: 140, - cell: (info) => ( - - ), - }, - { - id: "deleted_by", - accessorKey: "deleted_by", - header: "Deleted By", - size: 120, - maxSize: 180, - cell: (info) => { - const value = (info.row.original as any).deleted_by as string | null | undefined; - return ( - - {value || "-"} - - ); - }, - }, - ]; - - const table = useReactTable({ - data: keys, - columns, - columnResizeMode: "onChange", - columnResizeDirection: "ltr", - state: { - sorting, - pagination: tablePagination, - }, - onSortingChange: setSorting, - onPaginationChange: (updater) => { - const newPagination = typeof updater === "function" ? updater(tablePagination) : updater; - setTablePagination(newPagination); - onPageChange(newPagination.pageIndex); - }, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - getPaginationRowModel: getPaginationRowModel(), - enableSorting: true, - manualSorting: false, - manualPagination: true, - pageCount: Math.ceil(totalCount / pageSize), - }); - - const { pageIndex: currentPageIndex } = table.getState().pagination; - const start = currentPageIndex * pageSize + 1; - const end = Math.min((currentPageIndex + 1) * pageSize, totalCount); - const rangeLabel = `${start} - ${end}`; + const columns = useMemo(() => getDeletedKeysTableColumns(), []); return ( -
-
-
- {isLoading || isFetching ? ( - Loading... - ) : ( - - Showing {rangeLabel} of {totalCount} results - - )} - -
- {isLoading || isFetching ? ( - Loading... - ) : ( - - Page {currentPageIndex + 1} of {table.getPageCount()} - - )} - - - - -
-
-
-
-
-
- - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - { - const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`); - if (resizer) { - (resizer as HTMLElement).style.opacity = "0.5"; - } - }} - onMouseLeave={() => { - const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`); - if (resizer && !header.column.getIsResizing()) { - (resizer as HTMLElement).style.opacity = "0"; - } - }} - onClick={header.column.getToggleSortingHandler()} - > -
-
- {header.isPlaceholder - ? null - : flexRender(header.column.columnDef.header, header.getContext())} -
-
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
-
header.column.resetSize()} - onMouseDown={header.getResizeHandler()} - onTouchStart={header.getResizeHandler()} - className={`resizer ${table.options.columnResizeDirection} ${header.column.getIsResizing() ? "isResizing" : ""}`} - style={{ - position: "absolute", - right: 0, - top: 0, - height: "100%", - width: "5px", - background: header.column.getIsResizing() ? "#3b82f6" : "transparent", - cursor: "col-resize", - userSelect: "none", - touchAction: "none", - opacity: header.column.getIsResizing() ? 1 : 0, - }} - /> -
- - ))} - - ))} - - - {isLoading || isFetching ? ( - - -
-

🚅 Loading keys...

-
- - - ) : keys.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No deleted keys found

-
- - - )} - -
-
-
-
-
- + key.token || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + paginationMode="server" + pagination={pagination} + onPaginationChange={onPaginationChange} + rowCount={totalCount} + isLoading={isLoading} + loadingMessage="Loading deleted keys…" + noDataMessage={} + size="compact" + /> ); } diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx new file mode 100644 index 00000000000..aa7d6380cc3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx @@ -0,0 +1,130 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells"; +import { DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; + +function TruncatedTextCell({ value }: { value: string | null | undefined }) { + if (!value) { + return -; + } + return ( + + {value} + + ); +} + +export const getDeletedKeysTableColumns = (): ColumnDef[] => [ + { + id: "token", + accessorKey: "token", + meta: { title: "Key ID" }, + header: "Key ID", + size: 150, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "key_alias", + accessorKey: "key_alias", + meta: { title: "Key Alias" }, + header: "Key Alias", + size: 150, + enableSorting: false, + cell: ({ row }) => { + const value = row.original.key_alias; + if (!value) { + return -; + } + return ( + + {value} + + ); + }, + }, + { + id: "team_alias", + accessorKey: "team_alias", + meta: { title: "Team Alias" }, + header: "Team Alias", + size: 120, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "spend", + accessorKey: "spend", + meta: { title: "Spend (USD)", numeric: true }, + header: ({ column }) => , + size: 100, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "max_budget", + accessorKey: "max_budget", + meta: { title: "Budget (USD)", numeric: true }, + header: "Budget (USD)", + size: 110, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "user_email", + accessorKey: "user_email", + meta: { title: "User Email" }, + header: "User Email", + size: 160, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "user_id", + accessorKey: "user_id", + meta: { title: "User ID" }, + header: "User ID", + size: 120, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "created_at", + accessorKey: "created_at", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "created_by", + accessorKey: "created_by", + meta: { title: "Created By" }, + header: "Created By", + size: 120, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "deleted_at", + accessorKey: "deleted_at", + meta: { title: "Deleted At" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "deleted_by", + accessorKey: "deleted_by", + meta: { title: "Deleted By" }, + header: "Deleted By", + size: 120, + enableSorting: false, + cell: ({ row }) => , + }, +]; diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx index 77d2e94067e..a3eceb4b458 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx @@ -32,9 +32,8 @@ beforeEach(() => { mockUseDeletedTeams.mockReturnValue({ data: [mockDeletedTeam], - isPending: false, - isFetching: false, - } as any); + isLoading: false, + } as unknown as ReturnType); }); it("should render DeletedTeamsPage component", () => { @@ -43,14 +42,13 @@ it("should render DeletedTeamsPage component", () => { expect(screen.getByText("Test Team")).toBeInTheDocument(); }); -it("should handle loading state", () => { +it("should show skeleton rows while the initial load is pending", () => { mockUseDeletedTeams.mockReturnValue({ data: undefined, - isPending: true, - isFetching: false, - } as any); + isLoading: true, + } as unknown as ReturnType); renderWithProviders(); - expect(screen.getByText("🚅 Loading teams...")).toBeInTheDocument(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); }); diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.tsx index 30803d0a004..0265a3b623e 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.tsx @@ -6,7 +6,7 @@ import { DeletedTeamsTable } from "./DeletedTeamsTable/DeletedTeamsTable"; export default function DeletedTeamsPage() { const { premiumUser } = useAuthorized(); - const { data: teamsData, isPending: isLoading, isFetching } = useDeletedTeams(1, 100); + const { data: teamsData, isLoading } = useDeletedTeams(1, 100); return (
@@ -19,7 +19,7 @@ export default function DeletedTeamsPage() { description="Deleted team auditing is graduating from beta into our Enterprise audit & compliance suite." /> )} - +
); } 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 358a9e90fa1..c0cc5a342a8 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx @@ -1,10 +1,10 @@ -import { screen } from "@testing-library/react"; +import { screen, within } from "@testing-library/react"; import { vi, it, expect, beforeEach } from "vitest"; import { renderWithProviders } from "../../../../tests/test-utils"; import { DeletedTeamsTable } from "./DeletedTeamsTable"; import { DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; -const mockDeletedTeam: DeletedTeam = { +const makeDeletedTeam = (overrides: Partial = {}): DeletedTeam => ({ team_id: "team-1", team_alias: "Test Team", models: ["gpt-3.5-turbo", "gpt-4"], @@ -19,22 +19,41 @@ const mockDeletedTeam: DeletedTeam = { deleted_at: "2024-11-15T10:00:00Z", deleted_by: "user-1", spend: 100.5, -}; + ...overrides, +}); beforeEach(() => { vi.clearAllMocks(); }); -it("should render DeletedTeamsTable component", () => { - renderWithProviders(); - - expect(screen.getByText("Test Team")).toBeInTheDocument(); -}); - -it("should display team information correctly", () => { - renderWithProviders(); +it("should display team information", () => { + renderWithProviders(); expect(screen.getByText("Test Team")).toBeInTheDocument(); expect(screen.getByText("team-1")).toBeInTheDocument(); - expect(screen.getByText("Showing 1 team")).toBeInTheDocument(); + expect(screen.getByText("org-1")).toBeInTheDocument(); +}); + +it("should sort teams by deleted_at descending by default", () => { + const teams = [ + makeDeletedTeam({ team_id: "team-old", team_alias: "older-team", deleted_at: "2024-01-01T10:00:00Z" }), + makeDeletedTeam({ team_id: "team-new", team_alias: "newer-team", deleted_at: "2024-06-01T10:00:00Z" }), + ]; + renderWithProviders(); + + const rows = screen.getAllByRole("row").slice(1); + expect(within(rows[0]).getByText("newer-team")).toBeInTheDocument(); + expect(within(rows[1]).getByText("older-team")).toBeInTheDocument(); +}); + +it("should show skeleton rows when loading", () => { + renderWithProviders(); + + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); +}); + +it("should show the empty state when there are no deleted teams", () => { + renderWithProviders(); + + expect(screen.getByText("No deleted teams found")).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.tsx index ddfd5cf73b6..9578a52453f 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.tsx @@ -1,299 +1,50 @@ "use client"; -import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells"; -import { ChevronDownIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; -import { - ColumnDef, - flexRender, - getCoreRowModel, - getSortedRowModel, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Badge, Text } from "@tremor/react"; -import { Tooltip } from "antd"; -import React, { useState } from "react"; + +import { SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; import { DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; -import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; + +import { getDeletedTeamsTableColumns } from "./DeletedTeamsTableColumns"; interface DeletedTeamsTableProps { teams: DeletedTeam[]; isLoading: boolean; - isFetching: boolean; } -export function DeletedTeamsTable({ teams, isLoading, isFetching }: DeletedTeamsTableProps) { - const [sorting, setSorting] = useState([ - { - id: "deleted_at", - desc: true, - }, - ]); - - const columns: ColumnDef[] = [ - { - id: "team_alias", - accessorKey: "team_alias", - header: "Team Name", - size: 150, - maxSize: 200, - cell: (info) => { - const value = info.getValue() as string; - return ( - - {value || "-"} - - ); - }, - }, - { - id: "team_id", - accessorKey: "team_id", - header: "Team ID", - size: 150, - maxSize: 250, - cell: (info) => , - }, - { - id: "created_at", - accessorKey: "created_at", - header: "Created", - size: 120, - maxSize: 140, - cell: (info) => , - }, - { - id: "spend", - accessorKey: "spend", - header: "Spend (USD)", - size: 100, - maxSize: 140, - cell: (info) => , - }, - { - id: "max_budget", - accessorKey: "max_budget", - header: "Budget (USD)", - size: 110, - maxSize: 150, - cell: (info) => ( - - ), - }, - { - id: "models", - accessorKey: "models", - header: "Models", - size: 200, - maxSize: 300, - cell: (info) => { - const models = info.getValue() as string[]; - if (!Array.isArray(models) || models.length === 0) { - return ( - - All Proxy Models - - ); - } - return ( -
- {models.slice(0, 3).map((model: string, index: number) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 ? `${getModelDisplayName(model).slice(0, 30)}...` : getModelDisplayName(model)} - - - ), - )} - {models.length > 3 && ( - - - +{models.length - 3} {models.length - 3 === 1 ? "more model" : "more models"} - - - )} -
- ); - }, - }, - { - id: "organization_id", - accessorKey: "organization_id", - header: "Organization", - size: 150, - maxSize: 200, - cell: (info) => , - }, - { - id: "deleted_at", - accessorKey: "deleted_at", - header: "Deleted At", - size: 120, - maxSize: 140, - cell: (info) => , - }, - { - id: "deleted_by", - accessorKey: "deleted_by", - header: "Deleted By", - size: 120, - maxSize: 180, - cell: (info) => { - const value = (info.row.original as any).deleted_by as string | null | undefined; - return ( - - {value || "-"} - - ); - }, - }, - ]; - - const table = useReactTable({ - data: teams, - columns, - columnResizeMode: "onChange", - columnResizeDirection: "ltr", - state: { - sorting, - }, - onSortingChange: setSorting, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - enableSorting: true, - manualSorting: false, - }); +const DEFAULT_SORTING: SortingState = [{ id: "deleted_at", desc: true }]; +function EmptyState() { return ( -
-
-
- {isLoading || isFetching ? ( - Loading... - ) : ( - - Showing {teams.length} {teams.length === 1 ? "team" : "teams"} - - )} -
-
-
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - { - const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`); - if (resizer) { - (resizer as HTMLElement).style.opacity = "0.5"; - } - }} - onMouseLeave={() => { - const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`); - if (resizer && !header.column.getIsResizing()) { - (resizer as HTMLElement).style.opacity = "0"; - } - }} - onClick={header.column.getToggleSortingHandler()} - > -
-
- {header.isPlaceholder - ? null - : flexRender(header.column.columnDef.header, header.getContext())} -
-
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
-
header.column.resetSize()} - onMouseDown={header.getResizeHandler()} - onTouchStart={header.getResizeHandler()} - className={`resizer ${table.options.columnResizeDirection} ${header.column.getIsResizing() ? "isResizing" : ""}`} - style={{ - position: "absolute", - right: 0, - top: 0, - height: "100%", - width: "5px", - background: header.column.getIsResizing() ? "#3b82f6" : "transparent", - cursor: "col-resize", - userSelect: "none", - touchAction: "none", - opacity: header.column.getIsResizing() ? 1 : 0, - }} - /> -
- - ))} - - ))} - - - {isLoading || isFetching ? ( - - -
-

🚅 Loading teams...

-
- - - ) : teams.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No deleted teams found

-
- - - )} - -
-
-
-
+
+
+
+
No deleted teams found
+
Teams deleted from this proxy will show up here.
); } + +export function DeletedTeamsTable({ teams, isLoading }: DeletedTeamsTableProps) { + const [sorting, setSorting] = useState(DEFAULT_SORTING); + + const columns = useMemo(() => getDeletedTeamsTableColumns(), []); + + return ( + team.team_id || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + isLoading={isLoading} + loadingMessage="Loading deleted teams…" + noDataMessage={} + size="compact" + /> + ); +} diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx new file mode 100644 index 00000000000..e36077fd2c3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx @@ -0,0 +1,111 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdCell, ModelsCell, MoneyCell } from "@/components/shared/table_cells"; +import { DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; + +export const getDeletedTeamsTableColumns = (): ColumnDef[] => [ + { + id: "team_alias", + accessorKey: "team_alias", + meta: { title: "Team Name" }, + header: "Team Name", + size: 150, + enableSorting: false, + cell: ({ row }) => { + const value = row.original.team_alias; + if (!value) { + return -; + } + return ( + + {value} + + ); + }, + }, + { + id: "team_id", + accessorKey: "team_id", + meta: { title: "Team ID" }, + header: "Team ID", + size: 150, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "created_at", + accessorKey: "created_at", + meta: { title: "Created" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "spend", + accessorKey: "spend", + meta: { title: "Spend (USD)", numeric: true }, + header: ({ column }) => , + size: 100, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "max_budget", + accessorKey: "max_budget", + meta: { title: "Budget (USD)", numeric: true }, + header: "Budget (USD)", + size: 110, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "models", + accessorKey: "models", + meta: { title: "Models", skeleton: "chips" }, + header: "Models", + size: 200, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "organization_id", + accessorKey: "organization_id", + meta: { title: "Organization" }, + header: "Organization", + size: 150, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "deleted_at", + accessorKey: "deleted_at", + meta: { title: "Deleted At" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "deleted_by", + accessorKey: "deleted_by", + meta: { title: "Deleted By" }, + header: "Deleted By", + size: 120, + enableSorting: false, + cell: ({ row }) => { + const value = row.original.deleted_by; + if (!value) { + return -; + } + return ( + + {value} + + ); + }, + }, +]; From edc30ea515484e30356d048166e0f5d4d87aace1 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 16 Jul 2026 09:54:07 -0700 Subject: [PATCH 44/62] test(e2e): datadog log delivery for successful chat, messages, and responses (LIT-4447) (#33415) * test(e2e): datadog log delivery for successful chat, messages, and responses Covers logging.datadog.success.exports_metric on all three routes: one successful non-streaming call must reach the DataDog logs intake as exactly one log event whose StandardLoggingPayload message carries the model group, real token counts, and a response cost equal to the x-litellm-response-cost header of the same response. Delivery is judged at the intake: the compose stack gains a dd-sink service recording every batch the datadog callback ships via the DD_BASE_URL testing override, and a typed reader replays it. Writing these caught a live product bug: /v1/messages double-logs every success (two byte-identical events per call), filed as LIT-4447; the messages test tolerates byte-identical duplicates of the one event until it lands, while a second differing event still fails * test(e2e): address review findings on the datadog delivery suite Consolidates the fresh-key first_ok helper into logging_client now that the otel PR it mirrored has merged (both test files use the shared copy), moves intake batch parsing into a helper so no path can leave the batch unbound, and gives the sink's /health endpoint a truthful text/plain content type * test(e2e): tolerate same-logical-event duplicates by call id, not byte identity A clean LIT-4447 repro showed the duplicated payload is built twice and can mint a fresh synthetic completion id per emission, arriving as two separate intake POSTs with the same litellm_call_id and identical substantive fields. Byte-identity was therefore a flaky criterion; duplicates now qualify only when they share the call id, call type, model group, tokens, and cost, and a second differing event still fails * test(e2e): assert the scenario strictly; the messages test is the LIT-4447 regression pin Per review direction the tests now assert exactly what the scenario promises: exactly one DataDog log event per successful call, on every route. The /v1/messages test therefore fails on current code against the known double-log (LIT-4447) and is its regression pin; it goes green when the fix lands. The duplicate-tolerance machinery is removed * Simplify docstrings for DataDog log tests Removed redundant phrasing about cost cross-checking in docstrings. * Update test_datadog_log_e2e.py --- tests/e2e/coverage_registry/logging.yaml | 2 +- tests/e2e/docker-compose.yml | 59 +++++++- tests/e2e/e2e_config.py | 4 + tests/e2e/logging/conftest.py | 7 + tests/e2e/logging/datadog_sink.py | 103 ++++++++++++++ tests/e2e/logging/logging_client.py | 18 ++- tests/e2e/logging/test_datadog_log_e2e.py | 162 ++++++++++++++++++++++ tests/e2e/logging/test_otel_trace_e2e.py | 33 ++--- 8 files changed, 360 insertions(+), 28 deletions(-) create mode 100644 tests/e2e/logging/datadog_sink.py create mode 100644 tests/e2e/logging/test_datadog_log_e2e.py diff --git a/tests/e2e/coverage_registry/logging.yaml b/tests/e2e/coverage_registry/logging.yaml index afb6dbc964e..4348ccdcd76 100644 --- a/tests/e2e/coverage_registry/logging.yaml +++ b/tests/e2e/coverage_registry/logging.yaml @@ -5,7 +5,7 @@ - {id: logging.s3.success.writes_object, module: logging, tier: P0, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages, embeddings], source: "integrations/s3_v2.py", rationale: "Primary audit trail; batch flush no-drop"} - {id: logging.s3.failure.writes_object, module: logging, tier: P0, event: failure, assertions: [writes_object], exercised_on: [chat_completions, messages], source: "integrations/s3_v2.py", rationale: "Failed calls persisted for compliance"} - {id: logging.gcs_bucket.success.writes_object, module: logging, tier: P0, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages, embeddings], source: "integrations/gcs_bucket/gcs_bucket.py", rationale: "GCS parallel to S3"} -- {id: logging.datadog.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/datadog/datadog.py", rationale: "Powers dashboards/alerts; cardinality regressions common"} +- {id: logging.datadog.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/datadog/datadog.py", rationale: "Powers dashboards/alerts; cardinality regressions common"} - {id: logging.datadog.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions], source: "integrations/datadog/datadog.py", rationale: "Failure metrics for alerting/SLO"} - {id: logging.prometheus.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/prometheus.py", rationale: "Standard OSS metrics; per-key cardinality (existing e2e)"} - {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"} diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index b64f3d8dbfd..3e5de028114 100644 --- a/tests/e2e/docker-compose.yml +++ b/tests/e2e/docker-compose.yml @@ -1,5 +1,41 @@ # local setup to run e2e tests configs: + dd_sink_script: + content: | + # Minimal DataDog logs-intake sink for the logging suite: records every + # POST (gunzipping the compressed batches the integration sends) and + # replays them as JSON on GET /requests so tests can assert delivery. + import gzip, json + from http.server import BaseHTTPRequestHandler, HTTPServer + + REQUESTS = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + body = self.rfile.read(int(self.headers.get("Content-Length", 0))) + if self.headers.get("Content-Encoding") == "gzip": + body = gzip.decompress(body) + REQUESTS.append({"path": self.path, "body": body.decode("utf-8", "replace")}) + self.send_response(202) + self.end_headers() + self.wfile.write(b"{}") + + def do_GET(self): + self.send_response(200) + if self.path == "/health": + self.send_header("Content-Type", "text/plain") + self.end_headers() + self.wfile.write(b"ok") + return + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps({"requests": REQUESTS}).encode()) + + def log_message(self, *args): + pass + + HTTPServer(("0.0.0.0", 8080), Handler).serve_forever() + litellm_config: content: | general_settings: @@ -23,7 +59,7 @@ configs: # (PHOENIX_COLLECTOR_HTTP_ENDPOINT below points it at the jaeger service), # so gen-AI spans export through a preset-owned provider - the code path # where trace splits actually happen - with no cloud credentials needed. - callbacks: ["arize_phoenix"] + callbacks: ["arize_phoenix", "datadog"] router_settings: routing_strategy: simple-shuffle @@ -93,10 +129,15 @@ services: condition: service_healthy jaeger: condition: service_healthy + dd-sink: + condition: service_healthy env_file: .env environment: LITELLM_MASTER_KEY: sk-1234 STORE_MODEL_IN_DB: "True" + DD_API_KEY: local-sink-noauth + DD_SITE: datadoghq.com + DD_BASE_URL: http://dd-sink:8080 LITELLM_OTEL_V2: "true" PHOENIX_COLLECTOR_HTTP_ENDPOINT: http://jaeger:4318/v1/traces PHOENIX_API_KEY: local-jaeger-noauth @@ -155,3 +196,19 @@ services: interval: 3s timeout: 3s retries: 20 + +# throwaway DataDog logs-intake sink (records POSTs, replays on GET /requests; +# see E2E_DD_SINK_URL) + dd-sink: + image: python:3.12-alpine + command: ["python", "/sink.py"] + configs: + - source: dd_sink_script + target: /sink.py + ports: + - "9915:8080" + healthcheck: + test: ["CMD", "wget", "-qO-", "http://127.0.0.1:8080/health"] + interval: 3s + timeout: 3s + retries: 20 diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 6e6c30709de..e84438430fd 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -32,6 +32,10 @@ CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5") # read exported spans back through it. OTEL_QUERY_URL = os.environ.get("E2E_OTEL_QUERY_URL", "http://localhost:16686").rstrip("/") +# Query URL of the compose stack's DataDog logs-intake sink (the `dd-sink` +# service records every intake POST and replays them on GET /requests). +DD_SINK_URL = os.environ.get("E2E_DD_SINK_URL", "http://localhost:9915").rstrip("/") + # Writes on the proxy are eventually consistent (e.g. spend rows flush on # proxy_batch_write_at, ~60s). Read-backs poll to this deadline, never sleep-once. POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) diff --git a/tests/e2e/logging/conftest.py b/tests/e2e/logging/conftest.py index 43d279602ef..5ae791917fd 100644 --- a/tests/e2e/logging/conftest.py +++ b/tests/e2e/logging/conftest.py @@ -11,6 +11,7 @@ import os import pytest from logging_client import LangfuseCreds, LoggingClient, build_logging_client, load_langfuse_creds +from datadog_sink import DdSinkReader, build_dd_sink_reader from otel_client import OtelReader, build_otel_reader @@ -35,6 +36,12 @@ def otel_reader() -> OtelReader: return build_otel_reader() +@pytest.fixture(scope="session") +def dd_sink() -> DdSinkReader: + """Read-back client for the compose stack's DataDog logs-intake sink.""" + return build_dd_sink_reader() + + @pytest.fixture def datadog_creds() -> None: """Require Datadog shipping credentials. Hard-fail when absent; never skip.""" diff --git a/tests/e2e/logging/datadog_sink.py b/tests/e2e/logging/datadog_sink.py new file mode 100644 index 00000000000..5b5059d1428 --- /dev/null +++ b/tests/e2e/logging/datadog_sink.py @@ -0,0 +1,103 @@ +"""Read-back for the DataDog logging tests: typed models over the compose +stack's dd-sink service, which records every logs-intake POST the datadog +callback sends (gunzipped) and replays them as JSON. + +Delivery is judged on what the sink actually received, mirroring how the OTEL +tests read Jaeger; a failed sink query is a hard failure, never an empty +result. External reads go through ``e2e_http``. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass + +import pytest +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from e2e_config import DD_SINK_URL, POLL_INTERVAL, POLL_TIMEOUT +from e2e_http import URL, NoBody, Success, get + + +class DdSinkRequest(BaseModel): + model_config = ConfigDict(extra="ignore") + + path: str + body: str + + +class DdSinkRequests(BaseModel): + model_config = ConfigDict(extra="ignore") + + requests: list[DdSinkRequest] = [] + + +class DdLogEvent(BaseModel): + model_config = ConfigDict(extra="ignore") + + message: str + ddsource: str | None = None + service: str | None = None + status: str | None = None + + +_EVENT_BATCH: TypeAdapter[list[DdLogEvent]] = TypeAdapter(list[DdLogEvent]) + + +def _parse_batch(request: DdSinkRequest) -> list[DdLogEvent]: + """The intake accepts an array of events or a single event object.""" + try: + return _EVENT_BATCH.validate_json(request.body) + except ValidationError: + try: + return [DdLogEvent.model_validate_json(request.body)] + except ValidationError: + pytest.fail(f"dd-sink recorded a non-log body on {request.path}: {request.body[:200]}") + + +@dataclass(frozen=True, slots=True) +class DdSinkReader: + sink_url: str + + def _recorded_requests(self) -> list[DdSinkRequest]: + result = get( + URL(f"{self.sink_url}/requests"), + headers=NoBody(), + params=NoBody(), + response_type=DdSinkRequests, + timeout=30.0, + ) + match result: + case Success(data=page): + return page.requests + case failure: + pytest.fail(f"dd-sink query at {self.sink_url} failed: {failure}") + + def events_for_marker(self, marker: str) -> list[DdLogEvent]: + """Every log event across every recorded intake batch whose message + carries the marker. More than one hit for one call IS the + duplicate-delivery bug, so this never collapses to a single event.""" + events: list[DdLogEvent] = [] + for request in self._recorded_requests(): + if "/api/v2/logs" not in request.path: + continue + events.extend(event for event in _parse_batch(request) if marker in event.message) + return events + + def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]: + """Poll until at least one matching event lands (the callback flushes + in periodic batches), then re-read after one more interval so a late + duplicate cannot hide from the exactly-one assertion. At the deadline + the last result is returned as-is.""" + deadline = time.monotonic() + POLL_TIMEOUT + while time.monotonic() < deadline: + events = self.events_for_marker(marker) + if events: + time.sleep(POLL_INTERVAL) + return self.events_for_marker(marker) + time.sleep(POLL_INTERVAL) + return self.events_for_marker(marker) + + +def build_dd_sink_reader() -> DdSinkReader: + return DdSinkReader(sink_url=DD_SINK_URL) diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index e37d6175705..8be573d72a9 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -18,7 +18,7 @@ import json import os import time from dataclasses import dataclass -from typing import Literal +from typing import Callable, Literal import pytest from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError @@ -28,6 +28,7 @@ from e2e_gateway import Gateway, build_gateway from e2e_http import ( URL, AuthHeaders, + require_successful_call, NoBody, StreamingResponse, Success, @@ -617,5 +618,20 @@ class LoggingClient: return self.list_langfuse_observations(creds, trace_id=gen.trace_id) or [gen] +def first_ok(client: LoggingClient, send: Callable[[], StreamingResponse]) -> StreamingResponse: + """First successful call on a fresh key. A fresh key may briefly 401 until + the data plane's auth cache picks it up, so retry on 401 to a deadline; a + 401 is rejected before the LLM call, so it cannot contaminate delivery or + trace assertions. Any other failure is behavior under test and fails hard.""" + deadline = time.monotonic() + client.gateway.poll_timeout + while True: + outcome = send() + if outcome.ok: + return outcome + if outcome.status_code != 401 or time.monotonic() >= deadline: + require_successful_call(outcome) + time.sleep(client.gateway.poll_interval) + + def build_logging_client() -> LoggingClient: return LoggingClient(gateway=build_gateway()) diff --git a/tests/e2e/logging/test_datadog_log_e2e.py b/tests/e2e/logging/test_datadog_log_e2e.py new file mode 100644 index 00000000000..4651bbb28ba --- /dev/null +++ b/tests/e2e/logging/test_datadog_log_e2e.py @@ -0,0 +1,162 @@ +"""Live e2e: DataDog log delivery for successful non-streaming calls. + +Covers logging.datadog.success.exports_metric: one successful call on each +route must reach the DataDog logs intake as EXACTLY ONE log event whose +message (the StandardLoggingPayload) carries the model, the token counts, and +the response cost. Delivery is judged on what the intake actually received: +the compose stack's dd-sink service records every batch the datadog callback +ships (DD_BASE_URL override) and the tests read it back, so a dropped event, a +duplicated event, or a payload missing the cost all fail here. + +Both halves of the contract are asserted: the recorded state (the proxy +reports the DataDogLogger callback active via /health/readiness/details) and +the enforced behavior (the event at the intake, with the cost cross-checked +exactly against the x-litellm-response-cost header of the very response the +caller received). +""" + +from __future__ import annotations + +import pytest +from pydantic import BaseModel, ConfigDict + +from datadog_sink import DdLogEvent, DdSinkReader +from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, unique_marker +from e2e_http import NoBody, StreamingResponse +from lifecycle import ResourceManager +from logging_client import LoggingClient, first_ok + +pytestmark = pytest.mark.e2e + +#: The active DataDog callback's name in /health/readiness/details success_callbacks. +DD_LOGGER_NAME = "DataDogLogger" + + +class _DdMessagePayload(BaseModel): + """The fields of the StandardLoggingPayload the scenario pins.""" + + model_config = ConfigDict(extra="ignore") + + model_group: str + total_tokens: int + response_cost: float + status: str + call_type: str + + +def _assert_datadog_configured(client: LoggingClient) -> None: + """Recorded state: the proxy reports the DataDog callback among its active + callbacks, so a missing destination config fails here, before any + delivery-based assertion can time out confusingly.""" + result = client.gateway.probe("/health/readiness/details", params=NoBody()) + assert result.status_code == 200, ( + f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}" + ) + assert DD_LOGGER_NAME in result.body, ( + f"the proxy must report the {DD_LOGGER_NAME} callback active " + f"(callbacks + DD_* env in the compose config); got: {result.body[:400]}" + ) + + +def _assert_exactly_one_event( + events: list[DdLogEvent], *, model_group: str, call_type: str, outcome: StreamingResponse +) -> None: + """The enforced behavior: the intake holds exactly one event for the call, + sourced from litellm, whose payload names the model group and call type, + counts real tokens, and carries the same cost the response header reported.""" + assert events, "no DataDog log event for this call reached the intake within the deadline" + assert len(events) == 1, ( + f"expected exactly ONE DataDog log event for the call, got {len(events)} - " + "more than one event for one call is the duplicate-delivery bug (see LIT-4447 " + "for the currently known /v1/messages instance)" + ) + event = events[0] + assert event.ddsource == "litellm", f"event ddsource must be litellm, got {event.ddsource!r}" + assert event.status == "info", f"success events ship at status info, got {event.status!r}" + + payload = _DdMessagePayload.model_validate_json(event.message) + assert payload.status == "success", f"payload status must be success, got {payload.status!r}" + assert payload.model_group == model_group, ( + f"payload model_group must be {model_group!r}, got {payload.model_group!r}" + ) + assert payload.call_type == call_type, ( + f"payload call_type must be {call_type!r}, got {payload.call_type!r}" + ) + assert payload.total_tokens > 0, f"payload must count real tokens, got {payload.total_tokens}" + assert outcome.response_cost is not None and outcome.response_cost > 0, ( + f"the response must report x-litellm-response-cost, got {outcome.response_cost!r}" + ) + assert abs(payload.response_cost - outcome.response_cost) < 1e-12, ( + f"payload response_cost {payload.response_cost} must equal the response header " + f"cost {outcome.response_cost}" + ) + + +class TestDataDogLogDelivery: + @pytest.mark.covers("logging.datadog.success.exports_metric", exercised_on=["chat_completions"]) + def test_chat_completions_emits_one_log_event( + self, client: LoggingClient, dd_sink: DdSinkReader, resources: ResourceManager + ) -> None: + """One successful non-streaming /chat/completions call must reach the + DataDog logs intake as exactly one log event whose payload carries the + model, the token counts, and the response cost.""" + _assert_datadog_configured(client) + + key = client.key_with_alias(f"dd-chat-{unique_marker()}", models=[CHEAP_ANTHROPIC_MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = first_ok( + client, + lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16), + ) + events = dd_sink.poll_events_for_marker(marker) + _assert_exactly_one_event( + events, model_group=CHEAP_ANTHROPIC_MODEL, call_type="acompletion", outcome=outcome + ) + + @pytest.mark.covers("logging.datadog.success.exports_metric", exercised_on=["messages"]) + def test_messages_emits_one_log_event( + self, client: LoggingClient, dd_sink: DdSinkReader, resources: ResourceManager + ) -> None: + """One successful non-streaming /v1/messages call must reach the + DataDog logs intake as exactly one log event whose payload carries the + model, the token counts, and the response cost. + + This currently fails on the known /v1/messages double-log (LIT-4447); it goes green when the fix lands.""" + _assert_datadog_configured(client) + + key = client.key_with_alias(f"dd-messages-{unique_marker()}", models=[CHEAP_ANTHROPIC_MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = first_ok( + client, + lambda: client.messages_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16), + ) + events = dd_sink.poll_events_for_marker(marker) + _assert_exactly_one_event( + events, model_group=CHEAP_ANTHROPIC_MODEL, call_type="anthropic_messages", outcome=outcome + ) + + @pytest.mark.covers("logging.datadog.success.exports_metric", exercised_on=["responses"]) + def test_responses_emits_one_log_event( + self, client: LoggingClient, dd_sink: DdSinkReader, resources: ResourceManager + ) -> None: + """One successful non-streaming /v1/responses call must reach the + DataDog logs intake as exactly one log event whose payload carries the + model, the token counts, and the response cost.""" + _assert_datadog_configured(client) + + key = client.key_with_alias(f"dd-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = first_ok( + client, + lambda: client.responses_raw(key, CHEAP_OPENAI_MODEL, f"reply with one word {marker}"), + ) + events = dd_sink.poll_events_for_marker(marker) + _assert_exactly_one_event( + events, model_group=CHEAP_OPENAI_MODEL, call_type="aresponses", outcome=outcome + ) diff --git a/tests/e2e/logging/test_otel_trace_e2e.py b/tests/e2e/logging/test_otel_trace_e2e.py index b00fd91be3c..e49dd311b32 100644 --- a/tests/e2e/logging/test_otel_trace_e2e.py +++ b/tests/e2e/logging/test_otel_trace_e2e.py @@ -18,15 +18,14 @@ destination's own query API - never proxy-side "export succeeded" logs). from __future__ import annotations import time -from collections.abc import Callable import pytest from pydantic import BaseModel, ConfigDict, ValidationError from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, unique_marker -from e2e_http import NoBody, StreamingResponse, require_successful_call +from e2e_http import NoBody from lifecycle import ResourceManager -from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient +from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient, first_ok from models import LiteLLMParamsBody from otel_client import JaegerSpan, JaegerTrace, OtelReader @@ -60,22 +59,6 @@ def _assert_otel_destination_configured(client: LoggingClient) -> None: ) -def _first_ok(client: LoggingClient, send: Callable[[], StreamingResponse]) -> StreamingResponse: - """First successful call on a fresh key. A fresh key may briefly 401 until - the data plane's auth cache picks it up, so retry on 401 to a deadline; a - 401 is rejected before the LLM call so it exports no gen-AI span and cannot - contaminate the trace assertions. Any other failure is behavior under test - and fails hard.""" - deadline = time.monotonic() + client.gateway.poll_timeout - while True: - outcome = send() - if outcome.ok: - return outcome - if outcome.status_code != 401 or time.monotonic() >= deadline: - require_successful_call(outcome) - time.sleep(client.gateway.poll_interval) - - def _parent_ids(span_id: str, trace: JaegerTrace) -> list[str]: span = next(s for s in trace.spans if s.span_id == span_id) return [ref.span_id for ref in span.references if ref.ref_type == "CHILD_OF"] @@ -253,7 +236,7 @@ class TestOtelTraceCompleteness: resources.defer(lambda: client.delete_key(key)) marker = unique_marker() - outcome = _first_ok( + outcome = first_ok( client, lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16) ) assert outcome.call_id is not None, "success response must carry x-litellm-call-id" @@ -286,7 +269,7 @@ class TestOtelTraceCompleteness: resources.defer(lambda: client.delete_key(key)) marker = unique_marker() - outcome = _first_ok( + outcome = first_ok( client, lambda: client.messages_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16) ) assert outcome.call_id is not None, "success response must carry x-litellm-call-id" @@ -319,7 +302,7 @@ class TestOtelTraceCompleteness: resources.defer(lambda: client.delete_key(key)) marker = unique_marker() - outcome = _first_ok( + outcome = first_ok( client, lambda: client.responses_raw(key, CHEAP_OPENAI_MODEL, f"reply with one word {marker}"), ) @@ -360,7 +343,7 @@ class TestOtelTraceCompleteness: resources.defer(lambda: client.delete_key(key)) marker = unique_marker() - outcome = _first_ok( + outcome = first_ok( client, lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", stream=True, max_tokens=16), ) @@ -416,7 +399,7 @@ class TestOtelTraceCompleteness: resources.defer(lambda: client.delete_key(key)) marker = unique_marker() - outcome = _first_ok( + outcome = first_ok( client, lambda: client.messages_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16, stream=True), ) @@ -474,7 +457,7 @@ class TestOtelTraceCompleteness: resources.defer(lambda: client.delete_key(key)) marker = unique_marker() - outcome = _first_ok( + outcome = first_ok( client, lambda: client.responses_raw(key, CHEAP_OPENAI_MODEL, f"reply with one word {marker}", stream=True), ) From 260d1eae8e728d571ff0bbcf0fc465788ad00f2e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 10:35:06 -0700 Subject: [PATCH 45/62] fix(cli): make CLI output ASCII-only so it doesn't crash legacy Windows consoles (#33465) * fix(cli): force UTF-8 output so emoji don't crash the CLI on Windows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cli): drop dead flush calls flagged by review * fix(cli): replace non-ASCII CLI output with ASCII so legacy Windows consoles don't crash --------- Co-authored-by: ryan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/client/cli/commands/auth.py | 62 +++++++++---------- litellm/proxy/client/cli/commands/chat.py | 2 +- .../proxy/client/cli/commands/encryption.py | 2 +- litellm/proxy/client/cli/commands/keys.py | 4 +- litellm/proxy/client/cli/commands/teams.py | 20 +++--- litellm/proxy/client/cli/interface.py | 14 ++--- .../proxy/client/cli/test_auth_commands.py | 34 +++++----- .../proxy/client/cli/test_global_options.py | 15 +++++ .../proxy/client/cli/test_keys_commands.py | 6 +- 9 files changed, 86 insertions(+), 73 deletions(-) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 785d3b1e37b..61495403407 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -72,7 +72,7 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None: console = Console() if not teams: - console.print("❌ No teams found for your user.") + console.print("No teams found for your user.") return table = Table(title="Available Teams") @@ -162,7 +162,7 @@ def display_interactive_team_selection(teams: List[Dict[str, Any]], selected_ind # Clear the screen using Rich's method console.clear() - console.print("🎯 Select a Team (Use ↑↓ arrows, Enter to select, 'q' to skip):\n") + console.print("Select a Team (Use up/down arrows, Enter to select, 'q' to skip):\n") for i, team in enumerate(teams): team_alias = team.get("team_alias") or "N/A" @@ -184,7 +184,7 @@ def display_interactive_team_selection(teams: List[Dict[str, Any]], selected_ind # Highlight the selected item if i == selected_index: - console.print(f"➤ [bold cyan]{team_alias}[/bold cyan] ({team_id})") + console.print(f"> [bold cyan]{team_alias}[/bold cyan] ({team_id})") console.print(f" Models: [yellow]{models_str}[/yellow]") console.print(f" Budget: [blue]{budget_str}[/blue]\n") else: @@ -220,15 +220,13 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any # Clear screen and show selection console = Console() console.clear() - click.echo( - f"✅ Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})" - ) + click.echo(f"Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})") return selected_team elif key == "quit" or key == "escape": # Clear screen console = Console() console.clear() - click.echo("ℹ️ Team selection skipped.") + click.echo("Team selection skipped.") return None elif key is None: # If we can't get key input, fall back to simple selection @@ -237,7 +235,7 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any except KeyboardInterrupt: console = Console() console.clear() - click.echo("\n❌ Team selection cancelled.") + click.echo("\nTeam selection cancelled.") return None except Exception: # If interactive mode fails, fall back to simple selection @@ -265,15 +263,15 @@ def prompt_team_selection_fallback( if 0 <= index < len(teams): selected_team = teams[index] click.echo( - f"\n✅ Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})" + f"\nSelected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})" ) return selected_team else: - click.echo(f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}") + click.echo(f"Invalid selection. Please enter a number between 1 and {len(teams)}") except ValueError: - click.echo("❌ Invalid input. Please enter a number or 'skip'") + click.echo("Invalid input. Please enter a number or 'skip'") except KeyboardInterrupt: - click.echo("\n❌ Team selection cancelled.") + click.echo("\nTeam selection cancelled.") return None @@ -437,7 +435,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Op user_id = data.get("user_id") normalized_teams: List[Dict[str, Any]] = _normalize_teams(teams, team_details) if not normalized_teams: - click.echo("⚠️ No teams available for selection.") + click.echo("Warning: No teams available for selection.") return None # User has multiple teams - let them select @@ -457,7 +455,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Op "team_id": None, # Set by server in JWT } - click.echo("❌ Team selection cancelled or JWT generation failed.") + click.echo("Team selection cancelled or JWT generation failed.") return None # JWT is ready (single team or team already selected) @@ -468,7 +466,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Op # Show which team was assigned if team_id and len(teams) == 1: - click.echo(f"\n✅ Automatically assigned to team: {team_id}") + click.echo(f"\nAutomatically assigned to team: {team_id}") if api_key: return { @@ -494,19 +492,19 @@ def _handle_team_selection_during_polling( The JWT token with the selected team, or None if selection was skipped """ if not teams: - click.echo("ℹ️ No teams found. You can create or join teams using the web interface.") + click.echo("No teams found. You can create or join teams using the web interface.") return None click.echo("\n" + "=" * 60) - click.echo("📋 Select a team for your CLI session...") + click.echo("Select a team for your CLI session...") team_id = _render_and_prompt_for_team_selection(teams) if not team_id: - click.echo("ℹ️ No team selected.") + click.echo("No team selected.") return None - click.echo(f"\n🔄 Generating JWT for team: {team_id}") + click.echo(f"\nGenerating JWT for team: {team_id}") poll_url = f"{base_url}/sso/cli/poll/{key_id}?team_id={team_id}" data = _poll_for_ready_data( @@ -520,7 +518,7 @@ def _handle_team_selection_during_polling( return None jwt_token = data.get("key") if jwt_token: - click.echo(f"✅ Successfully generated JWT for team: {team_id}") + click.echo(f"Successfully generated JWT for team: {team_id}") return jwt_token return None @@ -568,14 +566,14 @@ def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Option selected_team = teams[index] team_id = str(selected_team.get("team_id")) team_alias = selected_team.get("team_alias") or team_id - click.echo(f"\n✅ Selected team: {team_alias} ({team_id})") + click.echo(f"\nSelected team: {team_alias} ({team_id})") return team_id - click.echo(f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}") + click.echo(f"Invalid selection. Please enter a number between 1 and {len(teams)}") except ValueError: - click.echo("❌ Invalid input. Please enter a number or 'skip'") + click.echo("Invalid input. Please enter a number or 'skip'") except KeyboardInterrupt: - click.echo("\n❌ Team selection cancelled.") + click.echo("\nTeam selection cancelled.") return None @@ -628,7 +626,7 @@ def login(ctx: click.Context): } ) - click.echo("\n✅ Login successful!") + click.echo("\nLogin successful!") click.echo(f"JWT Token: {api_key[:20]}...") click.echo("You can now use the CLI without specifying --api-key") @@ -637,7 +635,7 @@ def login(ctx: click.Context): show_commands() return else: - click.echo("❌ Authentication timed out. Please try again.") + click.echo("Authentication timed out. Please try again.") click.echo( "The proxy never reported the browser sign-in as finished. If you did complete it, " "check the proxy logs for /sso/callback errors and confirm SSO is configured on the proxy." @@ -645,10 +643,10 @@ def login(ctx: click.Context): return except KeyboardInterrupt: - click.echo("\n❌ Authentication cancelled by user.") + click.echo("\nAuthentication cancelled by user.") return except Exception as e: - click.echo(f"❌ Authentication failed: {e}") + click.echo(f"Authentication failed: {e}") return @@ -656,7 +654,7 @@ def login(ctx: click.Context): def logout(): """Logout and clear stored authentication""" clear_token() - click.echo("✅ Logged out successfully. Authentication token cleared.") + click.echo("Logged out successfully. Authentication token cleared.") @click.command(name="print-token") @@ -703,10 +701,10 @@ def whoami(): token_data = load_token() if not token_data: - click.echo("❌ Not authenticated. Run 'lite login' to authenticate.") + click.echo("Not authenticated. Run 'lite login' to authenticate.") return - click.echo("✅ Authenticated") + click.echo("Authenticated") click.echo(f"User Email: {token_data.get('user_email', 'Unknown')}") click.echo(f"User ID: {token_data.get('user_id', 'Unknown')}") click.echo(f"User Role: {token_data.get('user_role', 'Unknown')}") @@ -717,7 +715,7 @@ def whoami(): click.echo(f"Token age: {age_hours:.1f} hours") if age_hours > CLI_JWT_EXPIRATION_HOURS: - click.echo(f"⚠️ Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.") + click.echo(f"Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.") @click.group(name="auth") diff --git a/litellm/proxy/client/cli/commands/chat.py b/litellm/proxy/client/cli/commands/chat.py index 9d1a1aa7a30..d78feb84bd6 100644 --- a/litellm/proxy/client/cli/commands/chat.py +++ b/litellm/proxy/client/cli/commands/chat.py @@ -150,7 +150,7 @@ def chat( f"Max Tokens: [yellow]{max_tokens or 'unlimited'}[/yellow]\n\n" f"Type your messages and press Enter. Type '/quit' or '/exit' to end the session.\n" f"Type '/help' for more commands.", - title="🤖 Chat Session", + title="Chat Session", ) ) diff --git a/litellm/proxy/client/cli/commands/encryption.py b/litellm/proxy/client/cli/commands/encryption.py index f67c9746fa9..4b460bac19c 100644 --- a/litellm/proxy/client/cli/commands/encryption.py +++ b/litellm/proxy/client/cli/commands/encryption.py @@ -32,7 +32,7 @@ def migrate(ctx: click.Context, check_only: bool, dry_run: bool): Requires the proxy to be started with ``general_settings.encryption_algorithm: aes-256-gcm``. Idempotent and - resumable — safe to re-run after an interruption. + resumable; safe to re-run after an interruption. Examples: litellm-proxy encryption migrate --check # attestation scan, no writes diff --git a/litellm/proxy/client/cli/commands/keys.py b/litellm/proxy/client/cli/commands/keys.py index 45e27442708..afbaa3702c1 100644 --- a/litellm/proxy/client/cli/commands/keys.py +++ b/litellm/proxy/client/cli/commands/keys.py @@ -309,12 +309,12 @@ def _import_keys_to_destination( imported_count += 1 key_alias = key.get("key_alias", "N/A") - click.echo(f"✓ Imported key: {key_alias}") + click.echo(f"Imported key: {key_alias}") except Exception as e: failed_count += 1 key_alias = key.get("key_alias", "N/A") - click.echo(f"✗ Failed to import key {key_alias}: {str(e)}", err=True) + click.echo(f"Failed to import key {key_alias}: {str(e)}", err=True) return imported_count, failed_count diff --git a/litellm/proxy/client/cli/commands/teams.py b/litellm/proxy/client/cli/commands/teams.py index c0d45544b11..b0ccdc8f9bf 100644 --- a/litellm/proxy/client/cli/commands/teams.py +++ b/litellm/proxy/client/cli/commands/teams.py @@ -21,7 +21,7 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None: console = Console() if not teams: - console.print("❌ No teams found for your user.") + console.print("No teams found for your user.") return table = Table(title="Available Teams") @@ -91,10 +91,10 @@ def available(ctx: click.Context): teams = client.teams.get_available() if teams: console = Console() - console.print("\n🎯 Available Teams to Join:") + console.print("\nAvailable Teams to Join:") display_teams_table(teams) else: - click.echo("ℹ️ No available teams to join.") + click.echo("No available teams to join.") except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) error_body = e.response.json() @@ -113,7 +113,7 @@ def assign_key(ctx: click.Context, team_id: Optional[str]): api_key = ctx.obj["api_key"] if not api_key: - click.echo("❌ No API key found. Please login first using 'litellm login'") + click.echo("No API key found. Please login first using 'litellm login'") raise click.Abort() try: @@ -122,7 +122,7 @@ def assign_key(ctx: click.Context, team_id: Optional[str]): teams = client.teams.list() if not teams: - click.echo("❌ No teams found for your user.") + click.echo("No teams found for your user.") return # Use interactive selection from auth module @@ -133,14 +133,14 @@ def assign_key(ctx: click.Context, team_id: Optional[str]): if selected_team: team_id = selected_team.get("team_id") else: - click.echo("❌ Operation cancelled.") + click.echo("Operation cancelled.") return # Update the key with the selected team if team_id: - click.echo(f"\n🔄 Assigning your key to team: {team_id}") + click.echo(f"\nAssigning your key to team: {team_id}") client.keys.update(key=api_key, team_id=team_id) - click.echo(f"✅ Successfully assigned key to team: {team_id}") + click.echo(f"Successfully assigned key to team: {team_id}") # Show team details if available teams = client.teams.list() @@ -148,9 +148,9 @@ def assign_key(ctx: click.Context, team_id: Optional[str]): if team.get("team_id") == team_id: models = team.get("models", []) if models: - click.echo(f"🎯 You can now access models: {', '.join(models)}") + click.echo(f"You can now access models: {', '.join(models)}") else: - click.echo("🎯 You can now access all available models") + click.echo("You can now access all available models") break except requests.exceptions.HTTPError as e: diff --git a/litellm/proxy/client/cli/interface.py b/litellm/proxy/client/cli/interface.py index 33f1f4a4480..e953742f412 100644 --- a/litellm/proxy/client/cli/interface.py +++ b/litellm/proxy/client/cli/interface.py @@ -27,13 +27,13 @@ def styled_prompt(): verbose_logger.debug(f"Error getting terminal size: {e}") click.echo("\n" * 3) - # Unicode box drawing characters - top_left = "┌" - top_right = "┐" - bottom_left = "└" - bottom_right = "┘" - horizontal = "─" - vertical = "│" + # ASCII box drawing characters + top_left = "+" + top_right = "+" + bottom_left = "+" + bottom_right = "+" + horizontal = "-" + vertical = "|" # Create the box with increased width width = 80 diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index e101515d4b1..2fbc9c5c82f 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -78,7 +78,7 @@ class TestPollingErrorSurfacing: result = CliRunner().invoke(login, obj=mock_context.obj) assert result.exit_code == 0 - assert "❌ Authentication failed:" in result.output + assert "Authentication failed:" in result.output assert "CLI login session not found or expired." in result.output assert "Authentication timed out" not in result.output @@ -414,7 +414,7 @@ class TestLoginCommand: result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 - assert "✅ Login successful!" in result.output + assert "Login successful!" in result.output assert "Automatically assigned to team: team-1" in result.output # Verify browser was opened with correct URL @@ -456,7 +456,7 @@ class TestLoginCommand: result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 - assert "❌ Authentication timed out" in result.output + assert "Authentication timed out" in result.output def test_login_http_error(self): """Test login with HTTP error""" @@ -476,7 +476,7 @@ class TestLoginCommand: result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 - assert "❌ Authentication timed out" in result.output + assert "Authentication timed out" in result.output def test_login_request_exception(self): """Test login with request exception""" @@ -497,7 +497,7 @@ class TestLoginCommand: result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 - assert "❌ Authentication timed out" in result.output + assert "Authentication timed out" in result.output def test_login_keyboard_interrupt(self): """Test login cancelled by user""" @@ -512,7 +512,7 @@ class TestLoginCommand: result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 - assert "❌ Authentication cancelled by user" in result.output + assert "Authentication cancelled by user" in result.output def test_login_no_api_key_in_response(self): """Test login when response doesn't contain API key""" @@ -536,7 +536,7 @@ class TestLoginCommand: result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 - assert "❌ Authentication timed out" in result.output + assert "Authentication timed out" in result.output def test_login_general_exception(self): """Test login with general exception (not requests exception)""" @@ -551,7 +551,7 @@ class TestLoginCommand: result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 - assert "❌ Authentication failed: Invalid value" in result.output + assert "Authentication failed: Invalid value" in result.output class TestLogoutCommand: @@ -567,7 +567,7 @@ class TestLogoutCommand: result = self.runner.invoke(logout) assert result.exit_code == 0 - assert "✅ Logged out successfully" in result.output + assert "Logged out successfully" in result.output mock_clear.assert_called_once() @@ -591,7 +591,7 @@ class TestWhoamiCommand: result = self.runner.invoke(whoami) assert result.exit_code == 0 - assert "✅ Authenticated" in result.output + assert "Authenticated" in result.output assert "test@example.com" in result.output assert "test-user-123" in result.output assert "admin" in result.output @@ -603,7 +603,7 @@ class TestWhoamiCommand: result = self.runner.invoke(whoami) assert result.exit_code == 0 - assert "❌ Not authenticated" in result.output + assert "Not authenticated" in result.output assert "Run 'lite login'" in result.output def test_whoami_old_token(self): @@ -619,8 +619,8 @@ class TestWhoamiCommand: result = self.runner.invoke(whoami) assert result.exit_code == 0 - assert "✅ Authenticated" in result.output - assert "⚠️ Warning: Token is more than 24 hours old" in result.output + assert "Authenticated" in result.output + assert "Warning: Token is more than 24 hours old" in result.output def test_whoami_missing_fields(self): """Test whoami with token missing some fields""" @@ -633,7 +633,7 @@ class TestWhoamiCommand: result = self.runner.invoke(whoami) assert result.exit_code == 0 - assert "✅ Authenticated" in result.output + assert "Authenticated" in result.output assert "Unknown" in result.output # Should show "Unknown" for missing fields def test_whoami_no_timestamp(self): @@ -655,7 +655,7 @@ class TestWhoamiCommand: result = self.runner.invoke(whoami) assert result.exit_code == 0 - assert "✅ Authenticated" in result.output + assert "Authenticated" in result.output # Should calculate age based on timestamp=0 assert "Token age:" in result.output @@ -714,7 +714,7 @@ class TestCLIKeyRegenerationFlow: result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 - assert "✅ Login successful!" in result.output + assert "Login successful!" in result.output assert "team-beta" in result.output # Ensure we surface the human-readable team alias to the user assert "Beta Team" in result.output @@ -774,7 +774,7 @@ class TestCLIKeyRegenerationFlow: result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 - assert "✅ Login successful!" in result.output + assert "Login successful!" in result.output # Verify browser was opened mock_browser.assert_called_once() diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 53b7e4dbc29..8df763d35c2 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -1,6 +1,7 @@ # stdlib imports import os import sys +from pathlib import Path from unittest.mock import Mock, patch import pytest @@ -11,6 +12,7 @@ sys.path.insert( ) # Adds the parent directory to the system path +import litellm.proxy.client.cli from litellm._version import version as litellm_version from litellm.proxy.client.cli import cli @@ -36,6 +38,19 @@ def test_cli_version_flag(cli_runner): assert "LiteLLM Proxy Server Version: 1.2.3" in result.output +def test_cli_source_is_ascii_only(): + """Non-ASCII output (emoji, box-drawing chars) raises UnicodeEncodeError on legacy Windows + consoles (cp1252), so the whole CLI package must stay ASCII-only.""" + cli_root = Path(litellm.proxy.client.cli.__file__).parent + offenders = [ + f"{path.relative_to(cli_root)}:{line_number}: {line.strip()}" + for path in sorted(cli_root.rglob("*.py")) + for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1) + if not line.isascii() + ] + assert offenders == [] + + def test_base_url_trailing_slash_normalized(cli_runner): """A trailing slash on --base-url must not produce a double slash (e.g. '//sso/cli/start').""" with ( diff --git a/tests/test_litellm/proxy/client/cli/test_keys_commands.py b/tests/test_litellm/proxy/client/cli/test_keys_commands.py index 2c134f9defb..977aec9f5b7 100644 --- a/tests/test_litellm/proxy/client/cli/test_keys_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_keys_commands.py @@ -262,7 +262,7 @@ def test_keys_import_actual_import_success(mock_keys_client, cli_runner): assert result.exit_code == 0 assert "Found 1 keys in source instance" in result.output - assert "✓ Imported key: import-key-1" in result.output + assert "Imported key: import-key-1" in result.output assert "Successfully imported: 1" in result.output assert "Failed to import: 0" in result.output @@ -481,8 +481,8 @@ def test_keys_import_partial_failure(mock_keys_client, cli_runner): ) assert result.exit_code == 0 # Command completes even with partial failures - assert "✓ Imported key: success-key" in result.output - assert "✗ Failed to import key fail-key" in result.output + assert "Imported key: success-key" in result.output + assert "Failed to import key fail-key" in result.output assert "Successfully imported: 1" in result.output assert "Failed to import: 1" in result.output assert "Total keys processed: 2" in result.output From df51cebcd3513db14b5a8b422c9a08d60317ef6b Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 10:39:53 -0700 Subject: [PATCH 46/62] fix: remove dead user-cache lookup with None key in spend-update path (#33555) With litellm_settings.enable_redis_auth_cache enabled, user_api_key_cache is Redis-backed. _update_user_db performed a cache lookup with key=user_id where user_id can be None; the in-memory cache tolerates a None key but Redis raises redis.exceptions.DataError (Invalid input of type: NoneType) on every spend update for requests without a user_id. The looked-up value was never used by any subsequent code, so the lookup is removed along with the user_api_key_cache parameter it existed for. Spend updates for users, end users, and the global proxy budget are unchanged --- litellm/proxy/db/db_spend_update_writer.py | 12 +-- .../proxy/db/test_db_spend_update_writer.py | 86 ++++++++++++++++++- 2 files changed, 86 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ca6875ca800..54a4c2dad91 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -27,7 +27,7 @@ from typing import ( import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching import DualCache, RedisCache +from litellm.caching import RedisCache from litellm.constants import ( DB_SPEND_UPDATE_JOB_NAME, DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, @@ -44,7 +44,6 @@ from litellm.proxy._types import ( DailyUserSpendTransaction, DBSpendUpdateTransactions, Litellm_EntityType, - LiteLLM_UserTable, SpendLogsMetadata, SpendLogsPayload, SpendUpdateQueueItem, @@ -137,7 +136,6 @@ class DBSpendUpdateWriter: disable_spend_logs, litellm_proxy_budget_name, prisma_client, - user_api_key_cache, ) from litellm.proxy.utils import ProxyUpdateSpend, hash_token @@ -195,7 +193,6 @@ class DBSpendUpdateWriter: org_id=org_id, end_user_id=end_user_id, prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, litellm_proxy_budget_name=litellm_proxy_budget_name, payload=payload, ) @@ -326,7 +323,6 @@ class DBSpendUpdateWriter: org_id: Optional[str], end_user_id: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, litellm_proxy_budget_name: Optional[str], payload: SpendLogsPayload, ): @@ -345,7 +341,6 @@ class DBSpendUpdateWriter: response_cost=response_cost, user_id=user_id, prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, litellm_proxy_budget_name=litellm_proxy_budget_name, end_user_id=end_user_id, ) @@ -510,7 +505,6 @@ class DBSpendUpdateWriter: response_cost: Optional[float], user_id: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, litellm_proxy_budget_name: Optional[str], end_user_id: Optional[str] = None, ): @@ -518,10 +512,6 @@ class DBSpendUpdateWriter: - Update that user's row - Update litellm-proxy-budget row (global proxy spend) """ - ## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db - existing_user_obj = await user_api_key_cache.async_get_cache(key=user_id) - if existing_user_obj is not None and isinstance(existing_user_obj, dict): - existing_user_obj = LiteLLM_UserTable(**existing_user_obj) try: if prisma_client is not None: # update user_ids = [user_id] diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index c29cdaf4171..4c17c5d3482 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -13,7 +13,10 @@ from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, call, patch import pytest +from redis.exceptions import DataError +import litellm +from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter @@ -1418,7 +1421,6 @@ async def test_batch_database_updates_isolation_on_failure(): org_id="org1", end_user_id="eu1", prisma_client=MagicMock(), - user_api_key_cache=MagicMock(), litellm_proxy_budget_name="budget", payload={"key": "value"}, ) @@ -1818,3 +1820,85 @@ async def test_update_database_does_not_deepcopy_on_request_path(): fake_payload["nested"]["a"] = 999 assert batch_payload["model"] == "gpt-4" assert batch_payload["nested"]["a"] == 1 + + +@pytest.mark.asyncio +async def test_spend_update_path_never_queries_user_cache_with_none_user_id(): + """ + When user_id is None, the spend-update path must not perform a user-cache + lookup at all. With a Redis-backed auth cache (enable_redis_auth_cache), + a lookup with key=None raises redis.exceptions.DataError, which aborted + _update_user_db before any spend updates were enqueued. + + This test fails on the old code twice over: the cache mock records the + forbidden lookup, and the DataError it raises kills the end-user spend + update that must survive. + """ + db_writer = DBSpendUpdateWriter() + + strict_redis_backed_cache = MagicMock() + strict_redis_backed_cache.async_get_cache = AsyncMock( + side_effect=DataError("Invalid input of type: 'NoneType'") + ) + + with ( + patch.object(litellm, "max_budget", 0), + patch("litellm.proxy.proxy_server.disable_spend_logs", True), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", strict_redis_backed_cache), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "litellm-proxy-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value={ + "startTime": "2024-01-01T00:00:00", + "endTime": "2024-01-01T00:01:00", + "model": "gpt-4", + "custom_llm_provider": "openai", + "spend": 0.0, + }, + ), + ): + await db_writer.update_database( + token=None, + user_id=None, + end_user_id="end-user-1", + team_id=None, + org_id=None, + kwargs={"model": "gpt-4", "custom_llm_provider": "openai"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.1, + ) + await asyncio.sleep(0) + + strict_redis_backed_cache.async_get_cache.assert_not_called() + + queued = await db_writer.spend_update_queue.flush_all_updates_from_in_memory_queue() + end_user_updates = [u for u in queued if u["entity_type"] == Litellm_EntityType.END_USER] + assert len(end_user_updates) == 1 + assert end_user_updates[0]["entity_id"] == "end-user-1" + assert all(u["entity_type"] != Litellm_EntityType.USER for u in queued) + + +@pytest.mark.asyncio +async def test_update_user_db_enqueues_user_spend_without_cache_dependency(): + """ + _update_user_db needs no cache handle: it enqueues the user spend update + (and the end-user one) purely from the ids it is given. + """ + db_writer = DBSpendUpdateWriter() + + with patch.object(litellm, "max_budget", 0): + await db_writer._update_user_db( + response_cost=0.25, + user_id="user-123", + prisma_client=MagicMock(), + litellm_proxy_budget_name="litellm-proxy-budget", + end_user_id="end-user-9", + ) + + queued = await db_writer.spend_update_queue.flush_all_updates_from_in_memory_queue() + by_type = {u["entity_type"]: u["entity_id"] for u in queued} + assert by_type[Litellm_EntityType.USER] == "user-123" + assert by_type[Litellm_EntityType.END_USER] == "end-user-9" From 45fb9a70b9f7fdc43d434e70b671d478b9feca98 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 10:41:59 -0700 Subject: [PATCH 47/62] feat(helm): add per-component PodDisruptionBudget and topologySpreadConstraints to componentized chart (#33430) * feat(helm): add per-component PodDisruptionBudget and topologySpreadConstraints to componentized chart The componentized chart (helm/litellm) had no PodDisruptionBudget template for the gateway, backend, or ui, so voluntary disruptions (node drains, Karpenter consolidation) could evict every replica of a component at once. The legacy chart shipped one out of the box. Deployments also had no way to configure topologySpreadConstraints, blocking HA spread across AZs. Adds a shared litellm.pdb helper rendered per component, gated on .pdb.enabled with minAvailable/maxUnavailable (minAvailable wins, fallback maxUnavailable: 1), selectors matching each component's selectorLabels. Adds .topologySpreadConstraints rendered into each Deployment pod spec. PDBs default to disabled since the default hpa.minReplicas of 1 with minAvailable: 1 would block drains entirely. Resolves LIT-4452 * fix(helm): honor explicit 0 in pdb minAvailable/maxUnavailable A Go-template truthy check treated an explicit 0 (forbid all voluntary disruptions via maxUnavailable: 0) as unset and silently replaced it with the fallback maxUnavailable: 1, weakening the configured protection. Treat a value as set when it is non-nil and non-empty-string instead. --- helm/litellm/templates/NOTES.txt | 5 + helm/litellm/templates/_helpers.tpl | 46 +++++ .../litellm/templates/backend/deployment.yaml | 4 + .../backend/poddisruptionbudget.yaml | 6 + .../litellm/templates/gateway/deployment.yaml | 4 + .../gateway/poddisruptionbudget.yaml | 6 + helm/litellm/templates/ui/deployment.yaml | 4 + .../templates/ui/poddisruptionbudget.yaml | 6 + .../tests/pdb_topology_spread_tests.yaml | 188 ++++++++++++++++++ helm/litellm/values.yaml | 32 +++ 10 files changed, 301 insertions(+) create mode 100644 helm/litellm/templates/backend/poddisruptionbudget.yaml create mode 100644 helm/litellm/templates/gateway/poddisruptionbudget.yaml create mode 100644 helm/litellm/templates/ui/poddisruptionbudget.yaml create mode 100644 helm/litellm/tests/pdb_topology_spread_tests.yaml diff --git a/helm/litellm/templates/NOTES.txt b/helm/litellm/templates/NOTES.txt index 5b939fe480a..468cf621b32 100644 --- a/helm/litellm/templates/NOTES.txt +++ b/helm/litellm/templates/NOTES.txt @@ -46,4 +46,9 @@ Reminders: - gateway.config.proxy_config (rendered into a ConfigMap and mounted at /app/config/config.yaml; gateway reads it via CONFIG_FILE_PATH) + - {component}.pdb.{enabled,minAvailable,maxUnavailable} (per-component PodDisruptionBudget; disabled by + default — with hpa.minReplicas of 1, minAvailable: 1 + would block node drains) + - {component}.topologySpreadConstraints (standard k8s list, e.g. spread replicas across + topology.kubernetes.io/zone) - Enable ingress.enabled=true to dispatch / → ui, gateway data-plane prefixes → gateway, and the catch-all → backend. diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 043a0afc173..a0205c0a3a2 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -295,6 +295,52 @@ harmless no-op for the Job and authoritative for the app pods. {{- end }} {{- end -}} +{{/* +PodDisruptionBudget shared by gateway, backend, and ui. + +Invoke with a dict: + (dict "root" $ "component" .Values.gateway "componentName" "gateway" + "fullname" (include "litellm.gateway.fullname" .) + "selectorLabels" (include "litellm.gateway.selectorLabels" .)) + +Renders nothing unless both the component and its `pdb.enabled` are on. +Only one of minAvailable / maxUnavailable should be set; if both are, +minAvailable wins. If neither is set, falls back to `maxUnavailable: 1` so +an enabled-but-unconfigured PDB still permits node drains. + +"Set" means non-nil and non-empty-string, so an explicit 0 (e.g. +`maxUnavailable: 0` to forbid all voluntary disruptions) is honored rather +than silently replaced by the fallback. +*/}} +{{- define "litellm.pdb" -}} +{{- $root := .root -}} +{{- $component := .component -}} +{{- $min := $component.pdb.minAvailable -}} +{{- $max := $component.pdb.maxUnavailable -}} +{{- $minSet := not (or (kindIs "invalid" $min) (eq (printf "%v" $min) "")) -}} +{{- $maxSet := not (or (kindIs "invalid" $max) (eq (printf "%v" $max) "")) -}} +{{- if and $component.enabled $component.pdb $component.pdb.enabled }} +apiVersion: policy/v1 +kind: PodDisruptionBudget +metadata: + name: {{ .fullname }} + labels: + {{- include "litellm.commonLabels" $root | nindent 4 }} + app.kubernetes.io/component: {{ .componentName }} +spec: + selector: + matchLabels: + {{- .selectorLabels | nindent 6 }} + {{- if $minSet }} + minAvailable: {{ $min }} + {{- else if $maxSet }} + maxUnavailable: {{ $max }} + {{- else }} + maxUnavailable: 1 + {{- end }} +{{- end }} +{{- end -}} + {{/* Renders `envFrom:` block for a component's `envConfigMaps` / `envSecrets` lists. Each entry is a resource name; the chart wires the whole ConfigMap / diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 9d056167fe1..892b84ff7d5 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -98,4 +98,8 @@ spec: tolerations: {{- toYaml . | nindent 8 }} {{- end }} + {{- with .Values.backend.topologySpreadConstraints }} + topologySpreadConstraints: + {{- toYaml . | nindent 8 }} + {{- end }} {{- end }} diff --git a/helm/litellm/templates/backend/poddisruptionbudget.yaml b/helm/litellm/templates/backend/poddisruptionbudget.yaml new file mode 100644 index 00000000000..02853ac879c --- /dev/null +++ b/helm/litellm/templates/backend/poddisruptionbudget.yaml @@ -0,0 +1,6 @@ +{{- include "litellm.pdb" (dict + "root" $ + "component" .Values.backend + "componentName" "backend" + "fullname" (include "litellm.backend.fullname" .) + "selectorLabels" (include "litellm.backend.selectorLabels" .)) }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index 4c80d784156..b2e22612905 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -100,4 +100,8 @@ spec: tolerations: {{- toYaml . | nindent 8 }} {{- end }} + {{- with .Values.gateway.topologySpreadConstraints }} + topologySpreadConstraints: + {{- toYaml . | nindent 8 }} + {{- end }} {{- end }} diff --git a/helm/litellm/templates/gateway/poddisruptionbudget.yaml b/helm/litellm/templates/gateway/poddisruptionbudget.yaml new file mode 100644 index 00000000000..15e89af17d7 --- /dev/null +++ b/helm/litellm/templates/gateway/poddisruptionbudget.yaml @@ -0,0 +1,6 @@ +{{- include "litellm.pdb" (dict + "root" $ + "component" .Values.gateway + "componentName" "gateway" + "fullname" (include "litellm.gateway.fullname" .) + "selectorLabels" (include "litellm.gateway.selectorLabels" .)) }} diff --git a/helm/litellm/templates/ui/deployment.yaml b/helm/litellm/templates/ui/deployment.yaml index 79e9a3e43bb..cd1f8c08fd4 100644 --- a/helm/litellm/templates/ui/deployment.yaml +++ b/helm/litellm/templates/ui/deployment.yaml @@ -76,4 +76,8 @@ spec: tolerations: {{- toYaml . | nindent 8 }} {{- end }} + {{- with .Values.ui.topologySpreadConstraints }} + topologySpreadConstraints: + {{- toYaml . | nindent 8 }} + {{- end }} {{- end }} diff --git a/helm/litellm/templates/ui/poddisruptionbudget.yaml b/helm/litellm/templates/ui/poddisruptionbudget.yaml new file mode 100644 index 00000000000..f7a3a694e9c --- /dev/null +++ b/helm/litellm/templates/ui/poddisruptionbudget.yaml @@ -0,0 +1,6 @@ +{{- include "litellm.pdb" (dict + "root" $ + "component" .Values.ui + "componentName" "ui" + "fullname" (include "litellm.ui.fullname" .) + "selectorLabels" (include "litellm.ui.selectorLabels" .)) }} diff --git a/helm/litellm/tests/pdb_topology_spread_tests.yaml b/helm/litellm/tests/pdb_topology_spread_tests.yaml new file mode 100644 index 00000000000..8aa05f3a969 --- /dev/null +++ b/helm/litellm/tests/pdb_topology_spread_tests.yaml @@ -0,0 +1,188 @@ +suite: test pod disruption budgets and topology spread constraints +templates: + - gateway/poddisruptionbudget.yaml + - backend/poddisruptionbudget.yaml + - ui/poddisruptionbudget.yaml + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - ui/deployment.yaml +values: + - ./values/required.yaml +tests: + - it: renders no PDB by default + templates: + - gateway/poddisruptionbudget.yaml + - backend/poddisruptionbudget.yaml + - ui/poddisruptionbudget.yaml + asserts: + - hasDocuments: + count: 0 + + - it: gateway PDB uses minAvailable and matches the gateway selector labels + template: gateway/poddisruptionbudget.yaml + set: + gateway.pdb.enabled: true + gateway.pdb.minAvailable: 1 + asserts: + - isKind: + of: PodDisruptionBudget + - equal: + path: apiVersion + value: policy/v1 + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-gateway + - equal: + path: spec.minAvailable + value: 1 + - notExists: + path: spec.maxUnavailable + - equal: + path: spec.selector.matchLabels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: gateway + + - it: backend PDB uses maxUnavailable when minAvailable is unset + template: backend/poddisruptionbudget.yaml + set: + backend.pdb.enabled: true + backend.pdb.maxUnavailable: 25% + asserts: + - equal: + path: spec.maxUnavailable + value: 25% + - notExists: + path: spec.minAvailable + - equal: + path: spec.selector.matchLabels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: backend + + - it: minAvailable wins when both minAvailable and maxUnavailable are set + template: gateway/poddisruptionbudget.yaml + set: + gateway.pdb.enabled: true + gateway.pdb.minAvailable: 2 + gateway.pdb.maxUnavailable: 1 + asserts: + - equal: + path: spec.minAvailable + value: 2 + - notExists: + path: spec.maxUnavailable + + - it: an explicit maxUnavailable 0 is honored instead of the fallback + template: backend/poddisruptionbudget.yaml + set: + backend.pdb.enabled: true + backend.pdb.maxUnavailable: 0 + asserts: + - equal: + path: spec.maxUnavailable + value: 0 + - notExists: + path: spec.minAvailable + + - it: an explicit minAvailable 0 is honored and beats a set maxUnavailable + template: gateway/poddisruptionbudget.yaml + set: + gateway.pdb.enabled: true + gateway.pdb.minAvailable: 0 + gateway.pdb.maxUnavailable: 1 + asserts: + - equal: + path: spec.minAvailable + value: 0 + - notExists: + path: spec.maxUnavailable + + - it: enabled PDB with neither knob set falls back to maxUnavailable 1 + template: ui/poddisruptionbudget.yaml + set: + ui.pdb.enabled: true + asserts: + - equal: + path: spec.maxUnavailable + value: 1 + - notExists: + path: spec.minAvailable + - equal: + path: spec.selector.matchLabels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: ui + + - it: renders no PDB for a disabled component even when its pdb is enabled + template: gateway/poddisruptionbudget.yaml + set: + gateway.enabled: false + gateway.pdb.enabled: true + asserts: + - hasDocuments: + count: 0 + + - it: deployments omit topologySpreadConstraints by default + templates: + - gateway/deployment.yaml + - backend/deployment.yaml + - ui/deployment.yaml + asserts: + - notExists: + path: spec.template.spec.topologySpreadConstraints + + - it: gateway deployment renders configured topologySpreadConstraints + template: gateway/deployment.yaml + set: + gateway.topologySpreadConstraints: + - maxSkew: 1 + topologyKey: topology.kubernetes.io/zone + whenUnsatisfiable: ScheduleAnyway + labelSelector: + matchLabels: + app.kubernetes.io/component: gateway + asserts: + - equal: + path: spec.template.spec.topologySpreadConstraints + value: + - maxSkew: 1 + topologyKey: topology.kubernetes.io/zone + whenUnsatisfiable: ScheduleAnyway + labelSelector: + matchLabels: + app.kubernetes.io/component: gateway + + - it: backend deployment renders configured topologySpreadConstraints + template: backend/deployment.yaml + set: + backend.topologySpreadConstraints: + - maxSkew: 1 + topologyKey: kubernetes.io/hostname + whenUnsatisfiable: DoNotSchedule + labelSelector: + matchLabels: + app.kubernetes.io/component: backend + asserts: + - equal: + path: spec.template.spec.topologySpreadConstraints[0].topologyKey + value: kubernetes.io/hostname + - equal: + path: spec.template.spec.topologySpreadConstraints[0].whenUnsatisfiable + value: DoNotSchedule + + - it: ui deployment renders configured topologySpreadConstraints + template: ui/deployment.yaml + set: + ui.topologySpreadConstraints: + - maxSkew: 1 + topologyKey: topology.kubernetes.io/zone + whenUnsatisfiable: ScheduleAnyway + asserts: + - equal: + path: spec.template.spec.topologySpreadConstraints[0].topologyKey + value: topology.kubernetes.io/zone diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 74d02a25b7a..461935b2f50 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -190,10 +190,28 @@ gateway: maxReplicas: 10 targetCPUUtilizationPercentage: 70 targetMemoryUtilizationPercentage: 80 + # PodDisruptionBudget for the gateway pods. Set exactly one of + # `minAvailable` / `maxUnavailable` (minAvailable wins if both are set; + # enabling without either falls back to `maxUnavailable: 1`). Disabled by + # default: with the default hpa.minReplicas of 1, a `minAvailable: 1` PDB + # would block node drains entirely. + pdb: + enabled: false + minAvailable: "" + maxUnavailable: "" podAnnotations: {} nodeSelector: {} tolerations: [] affinity: {} + # Standard k8s topologySpreadConstraints for the gateway pods, e.g. to + # spread replicas across zones: + # - maxSkew: 1 + # topologyKey: topology.kubernetes.io/zone + # whenUnsatisfiable: ScheduleAnyway + # labelSelector: + # matchLabels: + # app.kubernetes.io/component: gateway + topologySpreadConstraints: [] # ---------- backend (UI / management API) ---------- backend: @@ -233,10 +251,17 @@ backend: minReplicas: 1 maxReplicas: 4 targetCPUUtilizationPercentage: 70 + # Same shape as gateway.pdb. + pdb: + enabled: false + minAvailable: "" + maxUnavailable: "" podAnnotations: {} nodeSelector: {} tolerations: [] affinity: {} + # Same shape as gateway.topologySpreadConstraints. + topologySpreadConstraints: [] # ---------- ui (Next.js static dashboard) ---------- ui: @@ -279,7 +304,14 @@ ui: minReplicas: 1 maxReplicas: 3 targetCPUUtilizationPercentage: 80 + # Same shape as gateway.pdb. + pdb: + enabled: false + minAvailable: "" + maxUnavailable: "" podAnnotations: {} nodeSelector: {} tolerations: [] affinity: {} + # Same shape as gateway.topologySpreadConstraints. + topologySpreadConstraints: [] From 3f5ed5a9c8d651896a1caf37eb9867ddfb9dff32 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 16 Jul 2026 11:05:31 -0700 Subject: [PATCH 48/62] fix(e2e/claude_code): unblock stage collection, align proxy env names, register compat models (#33433) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(e2e/claude_code): align proxy env names with the rest of tests/e2e Every claude_code compat cell used to read its own `LITELLM_PROXY_BASE_URL` and `LITELLM_PROXY_API_KEY` and duplicate the same 12-line "missing env, hard fail" block. The rest of `tests/e2e/` reads `LITELLM_PROXY_URL` and `LITELLM_MASTER_KEY` from `e2e_config.py`, so anyone standing up a live proxy for one suite had to export a second spelling for claude_code, and every cell repeated the same boilerplate. Centralize the resolution in `claude_code/_env.py`. `resolve_proxy()` prefers the suite-wide `LITELLM_PROXY_URL` / `LITELLM_MASTER_KEY` names and falls back to the legacy pair so existing CI wiring on stage keeps working during the roll-out. `require_proxy(compat_result)` is the one-liner cells call to bind `(base_url, api_key)` or hard-fail with a message that names both spellings. 55 cell files, `_basic_messaging.py`, and the driver's own unit-test fixture now go through the helper. `run_compat.sh` accepts either spelling and normalizes to the primary names before invoking pytest. `cron_vm/run_daily.sh` exports the primary names when launching pytest. `_pr_gate_unit_tests/test_env_resolution.py` pins the resolution rules so a future edit cannot silently reintroduce the drift: primary names win on tie, legacy names still resolve when primary is unset, mixed URL-primary key-legacy still resolves, empty-string exports are treated as unset, `require_proxy` names both spellings in its error message. Net diff: 71 files, +370/-1240. * fix(e2e): anchor claude_code Bash pin at parents[1] so container run collects `test_bash_tool_restrictions.py` derived `REPO_ROOT = Path(__file__).resolve().parents[4]` and then joined `tests/e2e/claude_code/`. That works locally, but the stage container mounts tests/e2e/ at /app/e2e/, so parents[4] resolves to filesystem root and the `_bash_cells()` assertion looks for `/tests/e2e/claude_code/tool_use` — a path that doesn't exist. Collection interrupts before any test runs, so the entire e2e suite appears broken. Fix: `CLAUDE_CODE_DIR = Path(__file__).resolve().parents[1]` resolves to the sibling `claude_code/` dir in either layout, and the `relative_to(REPO_ROOT)` calls become `relative_to(CLAUDE_CODE_DIR)` so test IDs and error messages read the same. Adds `test_claude_code_dir_anchor_is_layout_independent` as a regression pin: it checks the anchor lands on a directory named `claude_code` that contains this test file, which would fail under the old parents[4] anchor when run from /app/e2e/. * feat(e2e/claude_code): register compat deployments via /model/new from a session fixture Every compat cell hardcodes a virtual model name like `claude-sonnet-4-6` or `claude-sonnet-4-6-bedrock-invoke` and hits the proxy expecting it to be routable. On stage those live in the deployed model_list; locally the `docker-config.yaml` under tests/e2e/ only declares one of them, so anything past haiku 400s with `Invalid model name`. `claude_code/test_config.yaml` is the ground-truth compat matrix config the deployment already uses. `_compat_models.py` loads it, normalizes the yaml keys pydantic would silently drop (vertex_ai_* → vertex_*), and selects the subset whose provider credentials are present in the environment. An autouse session fixture in `conftest.py` POSTs each selected deployment to `/model/new`, blocks until it is servable on the data plane, and tears them all down on session exit. Skips silently when the proxy env is unset so pure-unit runs stay hermetic. `test_compat_models.py` pins the invariants that keep this safe. Every cell-referenced name must have a yaml entry (drift check catches a cell probing a name the fixture never registered); the yaml has no unused declarations; the fixture registers exactly 15 deployments (3 tiers × 5 provider surfaces); vertex_ai_* yaml keys populate the pydantic body's vertex_* fields (they got silently dropped historically); Azure needs both AZURE_FOUNDRY_* env vars; Bedrock lifts creds from the ambient AWS chain; Vertex needs both the yaml refs AND ambient GCP credentials. * refactor(e2e/claude_code): inject env + runner instead of monkeypatching `require_proxy` and `_basic_messaging.run_basic_messaging_cell` now take the env mapping (and the CLI runner) as constructor-style arguments with `os.environ` and `run_claude_models_parallel` as defaults. Tests exercise the branching by passing dicts and callables directly, so `monkeypatch.setenv` and `monkeypatch.setattr(_basic_messaging, "run_claude_models_parallel", ...)` are gone from every unit test in this refactor's blast radius. `test_env_resolution.py` drops the `monkeypatch.setenv`/`delenv` fixtures and passes `env={...}` dicts to `require_proxy`. Added a new pinned check that a successful resolution leaves `compat_result` untouched, and split the "unset env" test into three explicit shapes (empty, primary-only, legacy-only) so a regression that swaps the precedence rule can no longer hide behind a single monkeypatched fixture. `test_basic_messaging.py` (driver) replaces the `_install_fake_runner(monkeypatch, ...)` helper with `_make_fake_runner(...)` that returns a `(callable, captured_dict)` pair the test passes in via the helper's new `runner=` kwarg. Also drops the autouse `_proxy_env` fixture in favor of a module-level `_PROXY_ENV` dict each test wires through the helper's new `env=` kwarg. Added a regression pin that a missing-env call hard-fails without ever invoking the runner (so the guard order stays correct). `test_run_daily_pytest_scrubs_env.py` updates its pin to assert the new suite-wide env spellings (`LITELLM_PROXY_URL` / `LITELLM_MASTER_KEY`) instead of the legacy `LITELLM_PROXY_BASE_URL` / `LITELLM_PROXY_API_KEY` that `run_daily.sh` used to export. * handwrote rules --- tests/e2e/CLAUDE.md | 13 + tests/e2e/claude_code/_basic_messaging.py | 30 +- .../fixtures/expected_matrix.json | 2 +- .../_builder_unit_tests/fixtures/results.json | 6 +- .../test_matrix_builder.py | 2 +- .../_builder_unit_tests/test_v0_layout.py | 2 +- tests/e2e/claude_code/_compat_models.py | 86 +++ .../test_basic_messaging.py | 94 ++- .../_driver_unit_tests/test_cli_driver.py | 4 +- .../_driver_unit_tests/test_passthrough.py | 38 +- .../_driver_unit_tests/test_rate_limiter.py | 4 +- tests/e2e/claude_code/_env.py | 74 +++ tests/e2e/claude_code/_passthrough.py | 28 +- .../test_bash_tool_restrictions.py | 29 +- .../_pr_gate_unit_tests/test_compat_models.py | 166 +++++ .../test_env_resolution.py | 162 +++++ .../_publisher_unit_tests/__init__.py | 0 .../test_run_daily_pytest_scrubs_env.py | 95 --- .../test_run_daily_release_pagination.py | 289 --------- ...test_run_daily_version_probe_scrubs_env.py | 92 --- .../test_systemd_unit_credential_isolation.py | 104 --- .../test_anthropic.py | 4 +- .../test_azure.py | 4 +- .../test_bedrock_converse.py | 4 +- .../test_bedrock_invoke.py | 4 +- .../test_vertex_ai.py | 4 +- .../test_anthropic.py | 4 +- .../basic_messaging_streaming/test_azure.py | 4 +- .../test_bedrock_converse.py | 4 +- .../test_bedrock_invoke.py | 4 +- .../test_vertex_ai.py | 4 +- tests/e2e/claude_code/conftest.py | 128 +++- .../count_tokens/test_anthropic.py | 25 +- .../claude_code/count_tokens/test_azure.py | 25 +- .../count_tokens/test_bedrock_converse.py | 25 +- .../count_tokens/test_bedrock_invoke.py | 25 +- .../count_tokens/test_vertex_ai.py | 25 +- tests/e2e/claude_code/cron_vm/build_matrix.py | 50 -- .../cron_vm/litellm-compat-matrix.env.example | 59 -- .../cron_vm/litellm-compat-matrix.service | 141 ----- .../cron_vm/litellm-compat-matrix.timer | 25 - tests/e2e/claude_code/cron_vm/run_daily.sh | 590 ------------------ .../long_context_1m/test_anthropic.py | 24 +- .../claude_code/long_context_1m/test_azure.py | 24 +- .../long_context_1m/test_bedrock_converse.py | 24 +- .../long_context_1m/test_bedrock_invoke.py | 24 +- .../long_context_1m/test_vertex_ai.py | 24 +- tests/e2e/claude_code/matrix_builder.py | 2 +- .../claude_code/passthrough/test_anthropic.py | 2 +- .../e2e/claude_code/passthrough/test_azure.py | 2 +- .../passthrough/test_bedrock_invoke.py | 2 +- .../claude_code/passthrough/test_vertex_ai.py | 2 +- .../claude_code/pdf_input/test_anthropic.py | 24 +- tests/e2e/claude_code/pdf_input/test_azure.py | 24 +- .../pdf_input/test_bedrock_converse.py | 24 +- .../pdf_input/test_bedrock_invoke.py | 24 +- .../claude_code/pdf_input/test_vertex_ai.py | 24 +- .../prompt_caching_1h/test_anthropic.py | 23 +- .../prompt_caching_1h/test_azure.py | 23 +- .../test_bedrock_converse.py | 23 +- .../prompt_caching_1h/test_bedrock_invoke.py | 23 +- .../prompt_caching_1h/test_vertex_ai.py | 23 +- .../prompt_caching_5m/test_anthropic.py | 23 +- .../prompt_caching_5m/test_azure.py | 23 +- .../test_bedrock_converse.py | 23 +- .../prompt_caching_5m/test_bedrock_invoke.py | 23 +- .../prompt_caching_5m/test_vertex_ai.py | 23 +- tests/e2e/claude_code/run_compat.sh | 10 +- .../structured_outputs/test_anthropic.py | 24 +- .../structured_outputs/test_azure.py | 24 +- .../test_bedrock_converse.py | 24 +- .../structured_outputs/test_bedrock_invoke.py | 24 +- .../structured_outputs/test_vertex_ai.py | 24 +- tests/e2e/claude_code/test_config.yaml | 58 +- .../claude_code/thinking/test_anthropic.py | 23 +- tests/e2e/claude_code/thinking/test_azure.py | 23 +- .../thinking/test_bedrock_converse.py | 23 +- .../thinking/test_bedrock_invoke.py | 23 +- .../claude_code/thinking/test_vertex_ai.py | 23 +- .../thinking_with_tool_use/test_anthropic.py | 23 +- .../thinking_with_tool_use/test_azure.py | 23 +- .../test_bedrock_converse.py | 23 +- .../test_bedrock_invoke.py | 23 +- .../thinking_with_tool_use/test_vertex_ai.py | 23 +- .../claude_code/tool_search/test_anthropic.py | 25 +- .../e2e/claude_code/tool_search/test_azure.py | 25 +- .../tool_search/test_bedrock_converse.py | 25 +- .../tool_search/test_bedrock_invoke.py | 25 +- .../claude_code/tool_search/test_vertex_ai.py | 25 +- .../claude_code/tool_use/test_anthropic.py | 23 +- tests/e2e/claude_code/tool_use/test_azure.py | 23 +- .../tool_use/test_bedrock_converse.py | 23 +- .../tool_use/test_bedrock_invoke.py | 23 +- .../claude_code/tool_use/test_vertex_ai.py | 23 +- .../tool_use_streaming/test_anthropic.py | 23 +- .../tool_use_streaming/test_azure.py | 23 +- .../test_bedrock_converse.py | 23 +- .../tool_use_streaming/test_bedrock_invoke.py | 23 +- .../tool_use_streaming/test_vertex_ai.py | 23 +- .../e2e/claude_code/vision/test_anthropic.py | 24 +- tests/e2e/claude_code/vision/test_azure.py | 24 +- .../vision/test_bedrock_converse.py | 24 +- .../claude_code/vision/test_bedrock_invoke.py | 24 +- .../e2e/claude_code/vision/test_vertex_ai.py | 24 +- .../claude_code/web_search/test_anthropic.py | 23 +- .../e2e/claude_code/web_search/test_azure.py | 23 +- .../web_search/test_bedrock_converse.py | 23 +- .../web_search/test_bedrock_invoke.py | 23 +- .../claude_code/web_search/test_vertex_ai.py | 23 +- .../llm_claude_code_compat.yaml | 110 ++++ tests/e2e/coverage_registry/schema.py | 7 + tests/e2e/docker-compose.yml | 2 + tests/e2e/e2e_gateway.py | 22 +- tests/e2e/models.py | 2 + 114 files changed, 1236 insertions(+), 2873 deletions(-) create mode 100644 tests/e2e/claude_code/_compat_models.py create mode 100644 tests/e2e/claude_code/_env.py create mode 100644 tests/e2e/claude_code/_pr_gate_unit_tests/test_compat_models.py create mode 100644 tests/e2e/claude_code/_pr_gate_unit_tests/test_env_resolution.py delete mode 100644 tests/e2e/claude_code/_publisher_unit_tests/__init__.py delete mode 100644 tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_pytest_scrubs_env.py delete mode 100644 tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_release_pagination.py delete mode 100644 tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_version_probe_scrubs_env.py delete mode 100644 tests/e2e/claude_code/_publisher_unit_tests/test_systemd_unit_credential_isolation.py delete mode 100644 tests/e2e/claude_code/cron_vm/build_matrix.py delete mode 100644 tests/e2e/claude_code/cron_vm/litellm-compat-matrix.env.example delete mode 100644 tests/e2e/claude_code/cron_vm/litellm-compat-matrix.service delete mode 100644 tests/e2e/claude_code/cron_vm/litellm-compat-matrix.timer delete mode 100755 tests/e2e/claude_code/cron_vm/run_daily.sh create mode 100644 tests/e2e/coverage_registry/llm_claude_code_compat.yaml diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index f0d283629b0..0e1eafb5196 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -171,3 +171,16 @@ other... e.g. other.auth.jwt.valid_token_allows other.lifecycle.readiness.reports_db ``` + +## Hard Rules +- no monkeypatching, mock tests or unit tests of any kind. if a contributor asks you to write an end to end test, do NOT stage a unit test with it. if you find a product gap, call it out in the PR description + +- use model management endpoints to create new models for a test. this could be in a conftest / inline for each test. ask the user what they want. + +- do not overengineer a test, i need you to write readable, clean code of what would look like a natural user scenario + +- when it comes to typing an input schema for an api endpoint, have it type X = A | B | C ... where X = exhaustive union of all supported input schemas and A, B, C typically are composed by a base type. types are only pretty for a api request / response body. make sure to compose types instead of repeating the same base attributes over and over again. + +- use the docker-compose to your advantage and spin up a local proxy, make sure all tests pass. if a test fails due to an internally found issue, let users know to create a linear ticket for it. + +- do not use xfail markers, tests should be written in a form that the end user expects it to pass diff --git a/tests/e2e/claude_code/_basic_messaging.py b/tests/e2e/claude_code/_basic_messaging.py index f6b82a38f6a..7c581cc5e38 100644 --- a/tests/e2e/claude_code/_basic_messaging.py +++ b/tests/e2e/claude_code/_basic_messaging.py @@ -27,19 +27,20 @@ collecting this module as a test file. from __future__ import annotations -import os -from typing import Any, Mapping, Sequence +from typing import Any, Callable, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, + DriverResult, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" + +ClaudeRunner = Callable[..., Mapping[str, DriverResult | ClaudeCLIError]] # Floor on the number of `stream_event` records (with delta payloads) # we expect to see when the proxy actually streams. With @@ -79,6 +80,8 @@ def run_basic_messaging_cell( models: Sequence[str], prompt: str, verify_streaming: bool = False, + env: Mapping[str, str] | None = None, + runner: ClaudeRunner = run_claude_models_parallel, ) -> None: """Run the shared `basic_messaging_*` × cell body. @@ -99,28 +102,13 @@ def run_basic_messaging_cell( streamed reply to a single ``assistant`` event in ``--print --output-format stream-json`` mode). """ - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result, env=env) extra_args: Sequence[str] = ( ("--include-partial-messages",) if verify_streaming else () ) - outcomes = run_claude_models_parallel( + outcomes = runner( models=models, prompt=prompt, base_url=base_url, diff --git a/tests/e2e/claude_code/_builder_unit_tests/fixtures/expected_matrix.json b/tests/e2e/claude_code/_builder_unit_tests/fixtures/expected_matrix.json index d3aca0142dc..405a3772a90 100644 --- a/tests/e2e/claude_code/_builder_unit_tests/fixtures/expected_matrix.json +++ b/tests/e2e/claude_code/_builder_unit_tests/fixtures/expected_matrix.json @@ -26,7 +26,7 @@ "providers": { "anthropic": { "status": "fail", - "error": "[claude-sonnet-4-6] tool call dropped" + "error": "[claude-sonnet-4-5] tool call dropped" }, "bedrock_invoke": { "status": "not_applicable", diff --git a/tests/e2e/claude_code/_builder_unit_tests/fixtures/results.json b/tests/e2e/claude_code/_builder_unit_tests/fixtures/results.json index a01540c394f..f1b00385f17 100644 --- a/tests/e2e/claude_code/_builder_unit_tests/fixtures/results.json +++ b/tests/e2e/claude_code/_builder_unit_tests/fixtures/results.json @@ -10,7 +10,7 @@ { "feature_id": "basic_messaging_non_streaming", "provider": "anthropic", - "nodeid": "tests/e2e/claude_code/basic_messaging_non_streaming/test_anthropic.py::test_basic_messaging_non_streaming_anthropic[claude-sonnet-4-6]", + "nodeid": "tests/e2e/claude_code/basic_messaging_non_streaming/test_anthropic.py::test_basic_messaging_non_streaming_anthropic[claude-sonnet-4-5]", "result": {"status": "pass"} }, { @@ -28,8 +28,8 @@ { "feature_id": "tool_use", "provider": "anthropic", - "nodeid": "tests/e2e/claude_code/tool_use/test_anthropic.py::test_x[claude-sonnet-4-6]", - "result": {"status": "fail", "error": "[claude-sonnet-4-6] tool call dropped"} + "nodeid": "tests/e2e/claude_code/tool_use/test_anthropic.py::test_x[claude-sonnet-4-5]", + "result": {"status": "fail", "error": "[claude-sonnet-4-5] tool call dropped"} }, { "feature_id": "tool_use", diff --git a/tests/e2e/claude_code/_builder_unit_tests/test_matrix_builder.py b/tests/e2e/claude_code/_builder_unit_tests/test_matrix_builder.py index 9ddbdd29846..95db8acec70 100644 --- a/tests/e2e/claude_code/_builder_unit_tests/test_matrix_builder.py +++ b/tests/e2e/claude_code/_builder_unit_tests/test_matrix_builder.py @@ -393,7 +393,7 @@ def test_build_matrix_6x5_grid_matches_published_sample(): feature_ids = [feature["id"] for feature in manifest["features"]] providers = manifest["providers"] - models = ["claude-haiku-4-5", "claude-sonnet-4-6", "claude-opus-4-7"] + models = ["claude-haiku-4-5", "claude-sonnet-4-5", "claude-opus-4-7"] results = [] for feature_id in feature_ids: diff --git a/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py b/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py index a3569ebdb49..04e0facff5e 100644 --- a/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py +++ b/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py @@ -152,7 +152,7 @@ def test_per_provider_test_file_imports_and_parametrizes_three_models( `claude-opus-4-7-bedrock-invoke`), so we check for the tier substrings rather than exact alias names.""" text = (SUITE_ROOT / feature_id / f"test_{provider}.py").read_text() - for tier in ("haiku-4-5", "sonnet-4-6", "opus-4-7"): + for tier in ("haiku-4-5", "sonnet-4-5", "opus-4-7"): assert ( tier in text ), f"{feature_id}/test_{provider}.py does not reference {tier}" diff --git a/tests/e2e/claude_code/_compat_models.py b/tests/e2e/claude_code/_compat_models.py new file mode 100644 index 00000000000..23c9c82e596 --- /dev/null +++ b/tests/e2e/claude_code/_compat_models.py @@ -0,0 +1,86 @@ +"""Load the claude_code compat matrix's deployment list from +``test_config.yaml``. + +``test_config.yaml`` is the ground-truth config the stage deployment +uses; parsing it at fixture time means a change there (new tier, tier +retirement, provider swap, endpoint rename) reaches the fixture with +no extra edit. A drift-check test asserts every ``*_MODELS`` list +referenced by the compat cells is covered by the yaml, so a cell that +adds a probe for a name the yaml doesn't know about fails loudly at +collection instead of at 400-time. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Mapping + +import yaml + +from models import LiteLLMParamsBody + +CONFIG_PATH = Path(__file__).resolve().parent / "test_config.yaml" + + +@dataclass(frozen=True, slots=True) +class CompatDeployment: + model_name: str + litellm_params: LiteLLMParamsBody + + +# The yaml uses ``vertex_ai_*`` for the vertex project/location fields +# (that is the spelling the proxy config file historically standardized +# on), while ``LiteLLMParamsBody`` names them without the ``_ai`` infix +# (matching the proxy's DB column). Both spellings resolve at call time +# on the proxy side, but pydantic silently drops unknown fields, so a +# raw ``LiteLLMParamsBody(**entry)`` would produce a body with the +# vertex project stripped - the resulting deployment 400s at +# ``/v1/messages`` with "Invalid model name". Normalize the yaml keys +# to the pydantic names in one place. +_YAML_TO_PYDANTIC_ALIASES = { + "vertex_ai_project": "vertex_project", + "vertex_ai_location": "vertex_location", + "vertex_ai_credentials": "vertex_credentials", +} + + +def _normalize_params(raw: Mapping[str, object]) -> dict[str, object]: + return {_YAML_TO_PYDANTIC_ALIASES.get(k, k): v for k, v in raw.items()} + + +ConfigReader = Callable[[Path], str] + + +def _default_reader(path: Path) -> str: + return path.read_text() + + +def load_all_deployments( + config_path: Path = CONFIG_PATH, + reader: ConfigReader = _default_reader, +) -> tuple[CompatDeployment, ...]: + """Every deployment declared in the yaml, in file order.""" + doc = yaml.safe_load(reader(config_path)) + model_list = doc.get("model_list") or [] + return tuple( + CompatDeployment( + model_name=entry["model_name"], + litellm_params=LiteLLMParamsBody( + **_normalize_params(entry["litellm_params"]) + ), + ) + for entry in model_list + ) + + +def all_expected_model_names( + *, + config_path: Path = CONFIG_PATH, + reader: ConfigReader = _default_reader, +) -> frozenset[str]: + """Every virtual name the compat matrix declares - the ground truth + the cells are supposed to probe. Used by the drift-check test.""" + return frozenset( + d.model_name for d in load_all_deployments(config_path, reader) + ) diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_basic_messaging.py b/tests/e2e/claude_code/_driver_unit_tests/test_basic_messaging.py index 018121a8e5c..bb195ac50fe 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_basic_messaging.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_basic_messaging.py @@ -1,20 +1,19 @@ """Unit tests for the shared `run_basic_messaging_cell` helper. -These tests mock `run_claude_models_parallel` so they exercise the -helper's branching (env-missing guard, per-model pass/fail/empty-text, -streaming wire check) without spawning the real CLI. The streaming -check is the regression we care about: a proxy that buffers the -upstream stream must turn the cell red, not green. +These tests inject a fake ``ClaudeRunner`` and a fake env mapping so +they exercise the helper's branching (env-missing guard, per-model +pass/fail/empty-text, streaming wire check) without spawning the real +CLI or touching ``os.environ``. The streaming check is the regression +we care about: a proxy that buffers the upstream stream must turn the +cell red, not green. """ from __future__ import annotations -import os from typing import Any, Dict, List, Mapping, Optional, Sequence import pytest -from claude_code import _basic_messaging from claude_code._basic_messaging import ( MIN_STREAM_DELTA_EVENTS, _count_stream_event_deltas, @@ -23,6 +22,12 @@ from claude_code._basic_messaging import ( from claude_code.cli_driver import DriverResult +_PROXY_ENV: Mapping[str, str] = { + "LITELLM_PROXY_URL": "http://localhost:4000", + "LITELLM_MASTER_KEY": "sk-test", +} + + class _FakeResult: """Stand-in for the test's `compat_result` fixture. @@ -82,16 +87,17 @@ def _buffered_events() -> List[Dict[str, Any]]: ] -def _install_fake_runner(monkeypatch, *, outcomes_by_model): - """Patch `run_claude_models_parallel` to return canned outcomes. +def _make_fake_runner(*, outcomes_by_model): + """Build an injectable runner that returns canned outcomes and + records the kwargs the helper passed in. - Captures the kwargs the cell passed in so tests can assert on - `extra_args` (which is how the streaming variant opts into - `--include-partial-messages`). - """ + Returns a ``(runner, captured)`` pair; ``captured`` is a dict the + test can assert against without any global mutation, which is why + we prefer DI over ``monkeypatch.setattr``: the helper takes a + ``runner=`` kwarg, so tests bind their fake directly.""" captured: Dict[str, Any] = {} - def fake(*, models, prompt, base_url, api_key, extra_args=None, **_kwargs): + def runner(*, models, prompt, base_url, api_key, extra_args=None, **_kwargs): captured["models"] = list(models) captured["prompt"] = prompt captured["base_url"] = base_url @@ -99,14 +105,7 @@ def _install_fake_runner(monkeypatch, *, outcomes_by_model): captured["extra_args"] = list(extra_args) if extra_args else [] return {model: outcomes_by_model[model] for model in models} - monkeypatch.setattr(_basic_messaging, "run_claude_models_parallel", fake) - return captured - - -@pytest.fixture(autouse=True) -def _proxy_env(monkeypatch): - monkeypatch.setenv("LITELLM_PROXY_BASE_URL", "http://localhost:4000") - monkeypatch.setenv("LITELLM_PROXY_API_KEY", "sk-test") + return runner, captured def test_count_stream_event_deltas_only_counts_records_with_event_payload(): @@ -123,28 +122,30 @@ def test_count_stream_event_deltas_only_counts_records_with_event_payload(): assert _count_stream_event_deltas(events) == 2 -def test_verify_streaming_passes_when_proxy_streams(monkeypatch): +def test_verify_streaming_passes_when_proxy_streams(): fake_result = _FakeResult() model = "claude-haiku-4-5" outcome = DriverResult(text="1\n2\n3", events=_streamed_events(n_deltas=5)) - captured = _install_fake_runner(monkeypatch, outcomes_by_model={model: outcome}) + runner, captured = _make_fake_runner(outcomes_by_model={model: outcome}) run_basic_messaging_cell( compat_result=fake_result, models=[model], prompt="Count from 1 to 5, one number per line.", verify_streaming=True, + env=_PROXY_ENV, + runner=runner, ) assert captured["extra_args"] == ["--include-partial-messages"] assert fake_result.rows == [{"status": "pass"}] -def test_verify_streaming_fails_when_proxy_buffers(monkeypatch): +def test_verify_streaming_fails_when_proxy_buffers(): fake_result = _FakeResult() model = "claude-haiku-4-5" outcome = DriverResult(text="1\n2\n3", events=_buffered_events()) - _install_fake_runner(monkeypatch, outcomes_by_model={model: outcome}) + runner, _captured = _make_fake_runner(outcomes_by_model={model: outcome}) with pytest.raises(pytest.fail.Exception): run_basic_messaging_cell( @@ -152,7 +153,9 @@ def test_verify_streaming_fails_when_proxy_buffers(monkeypatch): models=[model], prompt="Count from 1 to 5, one number per line.", verify_streaming=True, - ) + env=_PROXY_ENV, + runner=runner, + ) assert len(fake_result.rows) == 1 row = fake_result.rows[0] @@ -161,33 +164,35 @@ def test_verify_streaming_fails_when_proxy_buffers(monkeypatch): assert f"< {MIN_STREAM_DELTA_EVENTS}" in row["error"] -def test_non_streaming_variant_omits_partial_messages_flag(monkeypatch): +def test_non_streaming_variant_omits_partial_messages_flag(): """Default `verify_streaming=False` keeps the non-streaming wire identical.""" fake_result = _FakeResult() model = "claude-haiku-4-5" outcome = DriverResult(text="pong", events=_buffered_events()) - captured = _install_fake_runner(monkeypatch, outcomes_by_model={model: outcome}) + runner, captured = _make_fake_runner(outcomes_by_model={model: outcome}) run_basic_messaging_cell( compat_result=fake_result, models=[model], prompt="Reply with the single word 'pong' and nothing else.", + env=_PROXY_ENV, + runner=runner, ) assert captured["extra_args"] == [] assert fake_result.rows == [{"status": "pass"}] -def test_verify_streaming_requires_all_models_to_stream(monkeypatch): +def test_verify_streaming_requires_all_models_to_stream(): """If any one tier buffers, the cell fails — same all-must-pass shape as the non-streaming check.""" fake_result = _FakeResult() outcomes = { "claude-haiku-4-5": DriverResult(text="ok", events=_streamed_events(5)), - "claude-sonnet-4-6": DriverResult(text="ok", events=_buffered_events()), + "claude-sonnet-4-5": DriverResult(text="ok", events=_buffered_events()), "claude-opus-4-7": DriverResult(text="ok", events=_streamed_events(5)), } - _install_fake_runner(monkeypatch, outcomes_by_model=outcomes) + runner, _captured = _make_fake_runner(outcomes_by_model=outcomes) with pytest.raises(pytest.fail.Exception): run_basic_messaging_cell( @@ -195,7 +200,30 @@ def test_verify_streaming_requires_all_models_to_stream(monkeypatch): models=list(outcomes.keys()), prompt="Count from 1 to 5, one number per line.", verify_streaming=True, - ) + env=_PROXY_ENV, + runner=runner, + ) statuses = [row["status"] for row in fake_result.rows] assert statuses == ["pass", "fail", "pass"] + + +def test_missing_proxy_env_hard_fails_regardless_of_runner(): + """The env guard fires before the runner is called, and takes the + env from the injected mapping (not os.environ). Passing an empty + env dict must hard-fail even if a happy runner is bound.""" + fake_result = _FakeResult() + runner, captured = _make_fake_runner(outcomes_by_model={}) + + with pytest.raises(pytest.fail.Exception): + run_basic_messaging_cell( + compat_result=fake_result, + models=["claude-haiku-4-5"], + prompt="whatever", + env={}, + runner=runner, + ) + + assert captured == {}, "runner must not be called when env resolution fails" + + diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_cli_driver.py b/tests/e2e/claude_code/_driver_unit_tests/test_cli_driver.py index f1a0534906e..ba6b9502c89 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_cli_driver.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_cli_driver.py @@ -140,7 +140,7 @@ def test_run_claude_inherits_only_allowlisted_os_environ(monkeypatch): monkeypatch.setenv("HOME", "/home/runner") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "totally-secret") monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-proxy-only") - monkeypatch.setenv("AZURE_FOUNDRY_API_KEY", "azure-secret") + monkeypatch.setenv("AZURE_AI_API_KEY", "azure-secret") monkeypatch.setenv("VERTEXAI_CREDENTIALS", '{"private_key": "leak"}') monkeypatch.setenv("GITHUB_TOKEN", "ghs_xxx") @@ -156,7 +156,7 @@ def test_run_claude_inherits_only_allowlisted_os_environ(monkeypatch): assert env["PATH"] == "/usr/bin:/usr/local/bin" assert "AWS_SECRET_ACCESS_KEY" not in env assert "ANTHROPIC_API_KEY" not in env - assert "AZURE_FOUNDRY_API_KEY" not in env + assert "AZURE_AI_API_KEY" not in env assert "VERTEXAI_CREDENTIALS" not in env assert "GITHUB_TOKEN" not in env diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_passthrough.py b/tests/e2e/claude_code/_driver_unit_tests/test_passthrough.py index 2d17a84d418..7c6c4f189ca 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_passthrough.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_passthrough.py @@ -20,6 +20,10 @@ from typing import Any, Dict, List, Mapping, Optional import pytest +from claude_code._env import ( + PRIMARY_API_KEY_ENV, + PRIMARY_BASE_URL_ENV, +) from claude_code._passthrough import ( ANTHROPIC_PASSTHROUGH_BASE_PATH, CLIENT_SIDE_AWS_REGION, @@ -33,8 +37,8 @@ from claude_code._passthrough import ( from claude_code.cli_driver import ClaudeCLIError, DriverResult PROXY_ENV = { - "LITELLM_PROXY_BASE_URL": "http://localhost:4000", - "LITELLM_PROXY_API_KEY": "sk-test", + PRIMARY_BASE_URL_ENV: "http://localhost:4000", + PRIMARY_API_KEY_ENV: "sk-test", } @@ -73,7 +77,29 @@ def test_env_missing_guard_reports_fail_and_aborts(): ) assert fake_result.single is not None assert fake_result.single["status"] == "fail" - assert "LITELLM_PROXY_BASE_URL" in fake_result.single["error"] + assert PRIMARY_BASE_URL_ENV in fake_result.single["error"] + assert PRIMARY_API_KEY_ENV in fake_result.single["error"] + + +def test_suite_wide_env_reaches_the_proxy(): + """Passthrough cells resolve via the same suite-wide env names as + every other e2e cell, so EKS wiring that only exports + LITELLM_PROXY_URL + LITELLM_MASTER_KEY reaches the ALB.""" + fake_result = _FakeResult() + captured: Dict[str, Any] = {} + outcome = DriverResult(text="pong") + + run_passthrough_cell( + compat_result=fake_result, + models=["claude-haiku-4-5"], + prompt="ping", + run_models=_fake_run_models({"claude-haiku-4-5": outcome}, captured), + env=PROXY_ENV, + ) + + assert captured["base_url"] == "http://localhost:4000" + assert captured["api_key"] == "sk-test" + assert fake_result.rows == [{"status": "pass"}] def test_anthropic_base_path_appended_to_normalized_proxy_url(): @@ -87,7 +113,7 @@ def test_anthropic_base_path_appended_to_normalized_proxy_url(): prompt="ping", passthrough_base_path=ANTHROPIC_PASSTHROUGH_BASE_PATH, run_models=_fake_run_models({"claude-haiku-4-5": outcome}, captured), - env={**PROXY_ENV, "LITELLM_PROXY_BASE_URL": "http://localhost:4000/"}, + env={**PROXY_ENV, PRIMARY_BASE_URL_ENV: "http://localhost:4000/"}, ) assert captured["base_url"] == "http://localhost:4000/anthropic" @@ -111,7 +137,7 @@ def test_extra_env_builder_receives_normalized_base_and_is_forwarded(): prompt="ping", build_extra_env=build, run_models=_fake_run_models({"claude-haiku-4-5": outcome}, captured), - env={**PROXY_ENV, "LITELLM_PROXY_BASE_URL": "http://localhost:4000/"}, + env={**PROXY_ENV, PRIMARY_BASE_URL_ENV: "http://localhost:4000/"}, ) assert seen_bases == ["http://localhost:4000"] @@ -124,7 +150,7 @@ def test_per_model_failures_reported_individually(): captured: Dict[str, Any] = {} outcomes = { "claude-haiku-4-5": DriverResult(text="pong"), - "claude-sonnet-4-6": ClaudeCLIError("claude CLI timed out after 120s"), + "claude-sonnet-4-5": ClaudeCLIError("claude CLI timed out after 120s"), "claude-opus-4-7": DriverResult(text="", exit_code=1), } diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_rate_limiter.py b/tests/e2e/claude_code/_driver_unit_tests/test_rate_limiter.py index 92907eda3c4..32a1c30af98 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_rate_limiter.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_rate_limiter.py @@ -60,10 +60,10 @@ from claude_code.rate_limiter import ( "model, expected", [ ("claude-haiku-4-5", PROVIDER_ANTHROPIC), - ("claude-sonnet-4-6", PROVIDER_ANTHROPIC), + ("claude-sonnet-4-5", PROVIDER_ANTHROPIC), ("claude-opus-4-7", PROVIDER_ANTHROPIC), ("claude-haiku-4-5-azure", PROVIDER_AZURE), - ("claude-sonnet-4-6-azure", PROVIDER_AZURE), + ("claude-sonnet-4-5-azure", PROVIDER_AZURE), ("claude-opus-4-7-vertex", PROVIDER_VERTEX_AI), ("claude-haiku-4-5-bedrock-converse", PROVIDER_BEDROCK_CONVERSE), ("claude-haiku-4-5-bedrock-invoke", PROVIDER_BEDROCK_INVOKE), diff --git a/tests/e2e/claude_code/_env.py b/tests/e2e/claude_code/_env.py new file mode 100644 index 00000000000..4f93cb57fdc --- /dev/null +++ b/tests/e2e/claude_code/_env.py @@ -0,0 +1,74 @@ +"""Proxy-env resolution for the claude_code compat cells. + +Uses the same ``LITELLM_PROXY_URL`` / ``LITELLM_MASTER_KEY`` names as +``e2e_config.py`` and the rest of ``tests/e2e/``. Everything under +``claude_code/`` goes through ``resolve_proxy`` / ``require_proxy`` here +so the naming lives in one place. +""" + +from __future__ import annotations + +import os +from typing import Mapping, NamedTuple + +import pytest + + +class ProxyConfig(NamedTuple): + base_url: str + api_key: str + + +PRIMARY_BASE_URL_ENV = "LITELLM_PROXY_URL" +PRIMARY_API_KEY_ENV = "LITELLM_MASTER_KEY" + + +def resolve_proxy_from(mapping: Mapping[str, str]) -> ProxyConfig | None: + """Pure resolver: takes an env mapping, returns a ProxyConfig if + both a base URL and an API key are present, else None. Extracted so + tests can exercise it without mutating ``os.environ``.""" + base_url = mapping.get(PRIMARY_BASE_URL_ENV) or None + api_key = mapping.get(PRIMARY_API_KEY_ENV) or None + if not base_url or not api_key: + return None + return ProxyConfig(base_url=base_url, api_key=api_key) + + +def resolve_proxy(env: Mapping[str, str] | None = None) -> ProxyConfig | None: + """Convenience wrapper that defaults to ``os.environ``. Prefer + calling ``resolve_proxy_from(env)`` from tests so nothing has to + reach into the process environment.""" + return resolve_proxy_from(os.environ if env is None else env) + + +def _fail_missing_proxy_env(compat_result) -> None: + compat_result.set( + { + "status": "fail", + "error": ( + f"missing required env: set {PRIMARY_BASE_URL_ENV} and " + f"{PRIMARY_API_KEY_ENV} to point at a running LiteLLM proxy" + ), + } + ) + pytest.fail( + f"{PRIMARY_BASE_URL_ENV} / {PRIMARY_API_KEY_ENV} not configured", + pytrace=False, + ) + + +def require_proxy( + compat_result, + *, + env: Mapping[str, str] | None = None, +) -> ProxyConfig: + """Return the proxy config (base URL + master key), or hard-fail + the test. + + ``env`` is injected for tests; production callers pass nothing and + the process env is used. This keeps tests off ``monkeypatch.setenv`` + for a check that is a pure function of its inputs.""" + cfg = resolve_proxy(env) + if cfg is None: + _fail_missing_proxy_env(compat_result) + return cfg diff --git a/tests/e2e/claude_code/_passthrough.py b/tests/e2e/claude_code/_passthrough.py index 24b6d694d3f..be7a475dff7 100644 --- a/tests/e2e/claude_code/_passthrough.py +++ b/tests/e2e/claude_code/_passthrough.py @@ -54,20 +54,17 @@ them unset. from __future__ import annotations -import os from typing import Any, Callable, Dict, Mapping, Optional, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" - ANTHROPIC_PASSTHROUGH_BASE_PATH = "/anthropic" CLIENT_SIDE_AWS_REGION = "us-east-1" @@ -140,32 +137,15 @@ def run_passthrough_cell( the trailing-slash-normalized proxy base URL and returns the provider-mode env for the CLI subprocess. """ - environ = env if env is not None else os.environ - base_url = environ.get(PROXY_BASE_URL_ENV) - api_key = environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) - - proxy_base = base_url.rstrip("/") + proxy = require_proxy(compat_result, env=env) + proxy_base = proxy.base_url.rstrip("/") extra_env = dict(build_extra_env(proxy_base)) if build_extra_env else None outcomes = run_models( models=models, prompt=prompt, base_url=proxy_base + passthrough_base_path, - api_key=api_key, + api_key=proxy.api_key, extra_env=extra_env, ) diff --git a/tests/e2e/claude_code/_pr_gate_unit_tests/test_bash_tool_restrictions.py b/tests/e2e/claude_code/_pr_gate_unit_tests/test_bash_tool_restrictions.py index d698131670a..ebd6d436d71 100644 --- a/tests/e2e/claude_code/_pr_gate_unit_tests/test_bash_tool_restrictions.py +++ b/tests/e2e/claude_code/_pr_gate_unit_tests/test_bash_tool_restrictions.py @@ -35,8 +35,7 @@ from typing import Iterable import pytest -REPO_ROOT = Path(__file__).resolve().parents[4] -CLAUDE_CODE_DIR = REPO_ROOT / "tests" / "e2e" / "claude_code" +CLAUDE_CODE_DIR = Path(__file__).resolve().parents[1] # Feature directories whose cells drive the `Bash` built-in tool. Add # new entries here when a new Bash-using feature is added; the test @@ -74,14 +73,14 @@ def _has_bare_bash_token(text: str) -> bool: @pytest.mark.parametrize( - "cell", list(_bash_cells()), ids=lambda p: str(p.relative_to(REPO_ROOT)) + "cell", list(_bash_cells()), ids=lambda p: str(p.relative_to(CLAUDE_CODE_DIR)) ) def test_bash_allow_rule_is_pinned_to_exact_echo_pong(cell: Path) -> None: """The cell must pass `Bash(echo pong)` as the allow rule, not the unrestricted `Bash` value that was originally flagged.""" text = cell.read_text() assert '"Bash(echo pong)"' in text, ( - f"{cell.relative_to(REPO_ROOT)} must restrict `--allowed-tools` to " + f"{cell.relative_to(CLAUDE_CODE_DIR)} must restrict `--allowed-tools` to " f'`Bash(echo pong)` (exact-match pattern). Unrestricted `"Bash"` ' f"grants arbitrary host command execution to model-controlled " f"tool_use blocks, which can read `docker inspect compat-proxy` " @@ -95,7 +94,7 @@ def test_bash_allow_rule_is_pinned_to_exact_echo_pong(cell: Path) -> None: # in text` short-circuits to True and lets a stray bare `"Bash"` # slip through silently. assert not _has_bare_bash_token(text), ( - f"{cell.relative_to(REPO_ROOT)} still references the unrestricted " + f"{cell.relative_to(CLAUDE_CODE_DIR)} still references the unrestricted " f'`"Bash"` value outside the `"Bash(echo pong)"` allow rule — ' f"sweep it out before merging." ) @@ -130,7 +129,7 @@ def test_has_bare_bash_token_ignores_unrelated_substrings(): @pytest.mark.parametrize( - "cell", list(_bash_cells()), ids=lambda p: str(p.relative_to(REPO_ROOT)) + "cell", list(_bash_cells()), ids=lambda p: str(p.relative_to(CLAUDE_CODE_DIR)) ) def test_bash_cell_uses_dontask_permission_mode(cell: Path) -> None: """The cell must pair the allow rule with `--permission-mode dontAsk` @@ -139,9 +138,25 @@ def test_bash_cell_uses_dontask_permission_mode(cell: Path) -> None: succeed without ever surfacing the security issue).""" text = cell.read_text() assert '"--permission-mode"' in text and '"dontAsk"' in text, ( - f"{cell.relative_to(REPO_ROOT)} must pass `--permission-mode dontAsk` " + f"{cell.relative_to(CLAUDE_CODE_DIR)} must pass `--permission-mode dontAsk` " f"alongside the `Bash(echo pong)` allow rule. Without dontAsk, " f"commands outside the allow rule fall back to the default ask-" f"mode behavior, which in `--print` (headless) mode is non-" f"interactive — defeating the explicit-allow contract." ) + + +def test_claude_code_dir_anchor_is_layout_independent() -> None: + """CLAUDE_CODE_DIR must resolve to the `claude_code/` directory that + contains this test file, regardless of how deep the repository is + mounted. The previous anchor `Path(__file__).resolve().parents[4]` + baked in the host layout (repo root sits four levels up) and broke + when the suite runs inside the stage container, where tests/e2e/ is + mounted at /app/e2e/ so `parents[4]` resolves to filesystem root and + the BASH_FEATURE_DIRS assertion looks for `/tests/e2e/claude_code/ + tool_use`. Anchoring at `parents[1]` (the sibling of this file's + parent) is the same directory in both layouts. + """ + assert CLAUDE_CODE_DIR.name == "claude_code" + assert CLAUDE_CODE_DIR.is_dir() + assert (CLAUDE_CODE_DIR / "_pr_gate_unit_tests" / Path(__file__).name).is_file() diff --git a/tests/e2e/claude_code/_pr_gate_unit_tests/test_compat_models.py b/tests/e2e/claude_code/_pr_gate_unit_tests/test_compat_models.py new file mode 100644 index 00000000000..6614e75e2f4 --- /dev/null +++ b/tests/e2e/claude_code/_pr_gate_unit_tests/test_compat_models.py @@ -0,0 +1,166 @@ +"""Regression tests for the compat-model registration loader. + +The compat cells hardcode virtual model names like ``claude-sonnet-4-5`` +and expect them to be registered on the proxy before the cell runs. The +session fixture in ``conftest.py`` reads ``test_config.yaml`` and POSTs +those deployments via ``/model/new``. These tests pin the invariants +that make that safe: + +- The yaml declares an entry for every virtual name a cell references - + otherwise a cell probes a name the fixture never registered, and the + cell hits an ``Invalid model name`` 400 that is much harder to trace. + +- The ``vertex_ai_*`` yaml keys get normalized to the ``vertex_*`` + pydantic-body names before ``LiteLLMParamsBody(**)`` sees them, so + the vertex project/location aren't silently dropped by pydantic's + ``extra="ignore"`` default. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +from claude_code._compat_models import ( + all_expected_model_names, + load_all_deployments, +) + + +CLAUDE_CODE_DIR = Path(__file__).resolve().parents[1] + + +def _cell_declared_model_names() -> frozenset[str]: + """Every ``"claude-*"`` model name a compat cell hardcodes in a + ``*_MODELS`` list. Uses a simple regex rather than importing every + cell because the cells depend on the harness which depends on env + the unit-test run does not have.""" + pattern = re.compile(r'"(claude-[a-zA-Z0-9._-]+)"') + found: set[str] = set() + for path in CLAUDE_CODE_DIR.glob("*/test_*.py"): + if path.parent.name.startswith("_"): + continue + for match in pattern.finditer(path.read_text()): + name = match.group(1) + # Skip upstream model references (they carry a version + # suffix or the ``anthropic/`` provider prefix - we only + # want proxy-side virtual names here). + if "/" in name or "@" in name: + continue + found.add(name) + return frozenset(found) + + +def test_yaml_covers_every_cell_declared_model_name() -> None: + """Every ``"claude-..."`` string a cell probes must have a + corresponding ``model_list`` entry in ``test_config.yaml``. A new + cell that adds a probe for a name the yaml doesn't know fails this + test - the alternative is a 400 at runtime that is much harder to + diagnose.""" + yaml_names = all_expected_model_names() + cell_names = _cell_declared_model_names() + missing = cell_names - yaml_names + assert not missing, ( + f"compat cells reference model names not declared in " + f"test_config.yaml: {sorted(missing)}. Add a matching " + f"model_list entry so the session fixture can register them." + ) + + +def test_yaml_has_no_unused_declarations() -> None: + """Every declaration in ``test_config.yaml`` is referenced by at + least one cell. A yaml entry no test exercises is dead + configuration and drift-prone; delete it or add the cell.""" + yaml_names = all_expected_model_names() + cell_names = _cell_declared_model_names() + unused = yaml_names - cell_names + assert not unused, ( + f"test_config.yaml declares model names no cell references: " + f"{sorted(unused)}. Delete them or add the cell." + ) + + +def test_load_returns_fifteen_deployments() -> None: + """The compat matrix is 3 tiers x 5 provider surfaces = 15. Pin the + count so a future edit to the yaml can't silently drop a tier.""" + assert len(load_all_deployments()) == 15 + + +def test_deployments_are_hashable_and_frozen() -> None: + """``CompatDeployment`` is frozen so tests cannot accidentally + mutate the shared list mid-session.""" + d = load_all_deployments()[0] + with pytest.raises((AttributeError, TypeError)): + d.model_name = "mutated" # type: ignore[misc] + + +def test_vertex_yaml_keys_populate_pydantic_body() -> None: + """The yaml spells vertex fields ``vertex_ai_project`` / + ``vertex_ai_location`` but ``LiteLLMParamsBody`` names them + ``vertex_project`` / ``vertex_location``. Without the alias + normalization the pydantic body silently drops the yaml keys, and + the deployment gets registered with no vertex project - a real + incident the drift regressed twice historically.""" + all_deployments = load_all_deployments() + vertex = [ + d for d in all_deployments if d.model_name.endswith("-vertex") + ] + assert vertex, "no vertex deployments found in yaml" + for d in vertex: + assert d.litellm_params.vertex_project, ( + f"{d.model_name} lost its vertex_project after normalization" + ) + assert d.litellm_params.vertex_location, ( + f"{d.model_name} lost its vertex_location after normalization" + ) + + +def test_vertex_deployments_keep_use_in_pass_through() -> None: + """Vertex passthrough cells need the deployment registered with + ``use_in_pass_through: true`` so the proxy wires project/location + credentials into the /vertex_ai passthrough router. ``LiteLLMParamsBody`` + defaults to ``extra="ignore"``, so a missing field on the body silently + strips the yaml flag and every vertex passthrough cell fails at runtime + with "No credentials found on proxy for project_name=...".""" + vertex = [ + d + for d in load_all_deployments() + if d.model_name.endswith("-vertex") + ] + assert vertex, "no vertex deployments found in yaml" + for d in vertex: + assert d.litellm_params.use_in_pass_through is True, ( + f"{d.model_name} lost use_in_pass_through after load; " + f"serialized body would be " + f"{d.litellm_params.model_dump(exclude_none=True)}" + ) + + +def test_yaml_litellm_params_are_all_known_body_fields() -> None: + """Every key under ``litellm_params`` in ``test_config.yaml`` must map + to a ``LiteLLMParamsBody`` field (after the vertex alias rewrite). + Without this pin, a new yaml flag can land in the fixture config and + be silently dropped by pydantic before ``/model/new`` ever sees it.""" + import yaml + from models import LiteLLMParamsBody + + from claude_code._compat_models import ( + CONFIG_PATH, + _YAML_TO_PYDANTIC_ALIASES, + ) + + known = frozenset(LiteLLMParamsBody.model_fields) + doc = yaml.safe_load(CONFIG_PATH.read_text()) + model_list = doc.get("model_list") or [] + unknown = tuple( + (entry["model_name"], key) + for entry in model_list + for key in entry["litellm_params"] + if _YAML_TO_PYDANTIC_ALIASES.get(key, key) not in known + ) + assert not unknown, ( + f"test_config.yaml litellm_params keys not on LiteLLMParamsBody " + f"(will be silently dropped at register time): {unknown}" + ) diff --git a/tests/e2e/claude_code/_pr_gate_unit_tests/test_env_resolution.py b/tests/e2e/claude_code/_pr_gate_unit_tests/test_env_resolution.py new file mode 100644 index 00000000000..f885b4baae2 --- /dev/null +++ b/tests/e2e/claude_code/_pr_gate_unit_tests/test_env_resolution.py @@ -0,0 +1,162 @@ +"""Regression tests for ``claude_code/_env.py``. + +Pin the resolution rules so a future edit cannot silently reintroduce +a private spelling that makes every ``claude_code`` cell fail with +"not configured" even when the surrounding e2e suite has a live proxy +configured under the suite-wide ``LITELLM_PROXY_URL`` / +``LITELLM_MASTER_KEY`` names. +""" + +from __future__ import annotations + +import pytest + +from claude_code._env import ( + PRIMARY_API_KEY_ENV, + PRIMARY_BASE_URL_ENV, + ProxyConfig, + require_proxy, + resolve_proxy_from, +) + + +class _CompatResultStub: + """Minimal stand-in for the compat_result fixture used by cells.""" + + def __init__(self) -> None: + self.calls: list[dict[str, str]] = [] + + def set(self, payload: dict[str, str]) -> None: + self.calls.append(payload) + + +def test_primary_env_names_match_suite_wide_config() -> None: + """The names claude_code reads must exactly match the ones + ``e2e_config.py`` reads for the rest of the suite. Anything else + silently reintroduces the drift this refactor cleaned up.""" + assert PRIMARY_BASE_URL_ENV == "LITELLM_PROXY_URL" + assert PRIMARY_API_KEY_ENV == "LITELLM_MASTER_KEY" + + +def test_returns_none_when_no_env_is_set() -> None: + assert resolve_proxy_from({}) is None + + +def test_returns_none_when_only_url_is_set() -> None: + assert ( + resolve_proxy_from({PRIMARY_BASE_URL_ENV: "http://localhost:4000"}) is None + ) + + +def test_returns_none_when_only_key_is_set() -> None: + assert resolve_proxy_from({PRIMARY_API_KEY_ENV: "sk-1234"}) is None + + +def test_primary_pair_resolves() -> None: + cfg = resolve_proxy_from( + { + PRIMARY_BASE_URL_ENV: "http://localhost:4000", + PRIMARY_API_KEY_ENV: "sk-1234", + } + ) + assert cfg == ProxyConfig("http://localhost:4000", "sk-1234") + + +def test_legacy_pair_is_ignored() -> None: + """``LITELLM_PROXY_BASE_URL`` / ``LITELLM_PROXY_API_KEY`` are not + accepted. A runner that only exports those must fail closed rather + than silently use a private spelling that the rest of the suite + does not know about.""" + assert ( + resolve_proxy_from( + { + "LITELLM_PROXY_BASE_URL": "http://legacy:4000", + "LITELLM_PROXY_API_KEY": "sk-legacy", + } + ) + is None + ) + + +def test_empty_string_env_is_treated_as_unset() -> None: + """``os.environ.get`` on an exported-but-empty var returns "" which + is falsy. The resolver must treat that as unset so a shell that + accidentally exports ``LITELLM_PROXY_URL=`` doesn't turn into a + "" base_url that hits the wrong endpoint.""" + assert ( + resolve_proxy_from( + {PRIMARY_BASE_URL_ENV: "", PRIMARY_API_KEY_ENV: "sk-1234"} + ) + is None + ) + + +def test_require_proxy_fails_with_helpful_message_when_env_empty() -> None: + """The error the user sees must name the suite-wide env vars so + they know exactly what to export.""" + compat = _CompatResultStub() + with pytest.raises(pytest.fail.Exception) as excinfo: + require_proxy(compat, env={}) + assert PRIMARY_BASE_URL_ENV in str(excinfo.value) + assert PRIMARY_API_KEY_ENV in str(excinfo.value) + assert compat.calls and compat.calls[0]["status"] == "fail" + assert PRIMARY_BASE_URL_ENV in compat.calls[0]["error"] + assert PRIMARY_API_KEY_ENV in compat.calls[0]["error"] + assert "LITELLM_PROXY_BASE_URL" not in compat.calls[0]["error"] + + +def test_require_proxy_returns_config_when_primary_env_supplied() -> None: + cfg = require_proxy( + _CompatResultStub(), + env={ + PRIMARY_BASE_URL_ENV: "http://localhost:4000", + PRIMARY_API_KEY_ENV: "sk-1234", + }, + ) + assert cfg == ProxyConfig("http://localhost:4000", "sk-1234") + + +class TestControlGatewayFollowsResolvedProxy: + """The session fixture that registers the compat deployments must talk + to the *same* proxy the cells do. + + Building its Gateway off ``e2e_config``'s own env read instead of the + resolved ``ProxyConfig`` would send ``/model/new`` to whatever + ``e2e_config`` defaults to when the process env is empty, while the + cells drive a different host — so registration silently lands + somewhere else and every cell 400s with "Invalid model name". + """ + + RESOLVED = ProxyConfig("http://eks-alb.internal:4000", "sk-eks") + + def _gateway(self): + from claude_code.conftest import _build_control_gateway + + return _build_control_gateway(self.RESOLVED) + + def test_management_calls_go_to_the_resolved_host_and_key(self) -> None: + control = self._gateway().transport.control + assert control.base_url == self.RESOLVED.base_url + assert control.master_key == self.RESOLVED.api_key + + def test_both_planes_share_the_one_address_the_cells_use(self) -> None: + """The deployment is fronted by a single address that routes + management and LLM paths itself, so a resolved proxy pins both.""" + transport = self._gateway().transport + assert transport.data.base_url == self.RESOLVED.base_url + assert transport.data.master_key == self.RESOLVED.api_key + assert transport.control.base_url == transport.data.base_url + + +def test_require_proxy_leaves_compat_result_untouched_on_success() -> None: + """A successful resolution must NOT append a spurious fail entry. + Would have silently poisoned every compat cell's result rows.""" + compat = _CompatResultStub() + require_proxy( + compat, + env={ + PRIMARY_BASE_URL_ENV: "http://localhost:4000", + PRIMARY_API_KEY_ENV: "sk-1234", + }, + ) + assert compat.calls == [] diff --git a/tests/e2e/claude_code/_publisher_unit_tests/__init__.py b/tests/e2e/claude_code/_publisher_unit_tests/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_pytest_scrubs_env.py b/tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_pytest_scrubs_env.py deleted file mode 100644 index 418d308674b..00000000000 --- a/tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_pytest_scrubs_env.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Pin: the cron `pytest` invocation must run under `env -i`. - -The systemd service `litellm-compat-matrix.service` loads provider -credentials (`ANTHROPIC_API_KEY`, `AWS_BEARER_TOKEN_BEDROCK`, -`AZURE_FOUNDRY_API_KEY`, `VERTEXAI_*`) and the agent-shin GitHub token -(`AGENT_SHIN_GITHUB_TOKEN`) into `run_daily.sh`'s environment from -`/etc/litellm-compat-matrix.env`. Pytest only needs to talk to the -loopback proxy at `127.0.0.1:${PROXY_PORT}` and has no legitimate reason -to see provider creds in its own `os.environ`. Leaving them in would -let a test under `tests/e2e/claude_code/` read them via `os.environ` and -exfiltrate them, and would also let a model-directed `Read` tool call -during a PDF/vision cell reach `/proc//environ`. The -PR-gate's pytest step in `.circleci/config.yml` already runs under -`env -i`; this pin enforces the same scrub on the cron path. -""" - -from __future__ import annotations - -from pathlib import Path - -REPO_ROOT = Path(__file__).resolve().parents[4] -RUN_DAILY = REPO_ROOT / "tests" / "e2e" / "claude_code" / "cron_vm" / "run_daily.sh" - - -def _pytest_invocation_block() -> str: - """Return only the executable lines around the pytest invocation. - - Comment text in run_daily.sh explains *why* certain credential - names must not appear, so a naïve substring scan over the whole - region would false-positive on the rationale itself. Strip lines - whose first non-space character is `#`. - """ - body = RUN_DAILY.read_text() - start = body.index('log "running pytest"') - end = body.index("PYTEST_EXIT=$?", start) - return "\n".join( - line for line in body[start:end].splitlines() - if line.lstrip()[:1] != "#" - ) - - -def test_pytest_invocation_wraps_in_env_i() -> None: - block = _pytest_invocation_block() - assert "env -i" in block, ( - "run_daily.sh: the pytest invocation must run under `env -i` so " - "PR-controlled test code under tests/e2e/claude_code/ cannot read " - "provider/agent-shin credentials out of the systemd service " - "environment, and so a model-directed `Read` tool call cannot " - "reach /proc//environ to pull them out." - ) - assert block.index("env -i") < block.index('"${WORKTREE_UV}" run pytest'), ( - "run_daily.sh: `env -i` must precede the pytest invocation; " - "otherwise pytest inherits the full credential-bearing env." - ) - - -def test_pytest_invocation_env_i_excludes_provider_secrets() -> None: - block = _pytest_invocation_block() - for forbidden in ( - "ANTHROPIC_API_KEY", - "AWS_BEARER_TOKEN_BEDROCK", - "AWS_ACCESS_KEY_ID", - "AWS_SECRET_ACCESS_KEY", - "VERTEXAI_CREDENTIALS", - "VERTEXAI_PROJECT", - "VERTEXAI_LOCATION", - "AZURE_FOUNDRY_API_KEY", - "AZURE_FOUNDRY_API_BASE", - "GITHUB_TOKEN", - "AGENT_SHIN_GITHUB_TOKEN", - ): - assert forbidden not in block, ( - f"run_daily.sh: the pytest-step `env -i` allowlist must not " - f"pass {forbidden} through. Found it inside the pytest " - f"invocation block." - ) - - -def test_pytest_invocation_passes_proxy_url_and_key_explicitly() -> None: - block = _pytest_invocation_block() - assert "LITELLM_PROXY_BASE_URL=" in block, ( - "run_daily.sh: the pytest `env -i` block must still pass " - "LITELLM_PROXY_BASE_URL so the test suite knows where to find " - "the loopback proxy." - ) - assert "LITELLM_PROXY_API_KEY=" in block, ( - "run_daily.sh: the pytest `env -i` block must still pass " - "LITELLM_PROXY_API_KEY so the test suite can authenticate to " - "the loopback proxy." - ) - assert "COMPAT_RESULTS_PATH=" in block, ( - "run_daily.sh: the pytest `env -i` block must still pass " - "COMPAT_RESULTS_PATH so the conftest writes the per-cell " - "tagged-union artifact to the script-managed path." - ) diff --git a/tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_release_pagination.py b/tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_release_pagination.py deleted file mode 100644 index fc733845ba3..00000000000 --- a/tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_release_pagination.py +++ /dev/null @@ -1,289 +0,0 @@ -"""Regression tests for the GitHub release pagination in `run_daily.sh`. - -The cron job resolves "newest LiteLLM v*-stable" via the GitHub Releases -API. A previous version of the loop broke as soon as the current page -contained ANY v*-stable tag. The Releases endpoint orders by -`created_at`, NOT by semver, so a backport on an older series cut today -(e.g. v1.80.1-stable) can land on an earlier page than a higher-version -release cut two weeks ago (e.g. v1.83.0-stable). The early-break would -silently pin the cron to a stale tag because the higher-version release -on a later page never made it into the merged set the final `sort_by` -consumed. - -These tests pin two things: - - 1. The buggy early-break-on-first-stable pattern must not return. - 2. The loop still terminates early on the standard "empty page" guard - so a quiet release feed doesn't burn API quota. - -The shell loop itself is exercised end-to-end with a fake `curl` that -serves canned page JSON, demonstrating that the resolved tag is the -highest-semver stable across all pages even when the highest tag lives -on page 2+. -""" - -from __future__ import annotations - -import os -import shutil -import subprocess -import textwrap -from pathlib import Path - -import pytest - -REPO_ROOT = Path(__file__).resolve().parents[4] -RUN_DAILY = REPO_ROOT / "tests" / "e2e" / "claude_code" / "cron_vm" / "run_daily.sh" - -# The extracted snippet starts AFTER `log`/`die` are defined in run_daily.sh, -# so the test harness has to provide its own stubs. Without them, a failure -# inside the snippet (e.g. jq returning an empty LITELLM_VERSION) would crash -# with `bash: die: command not found` (exit 127) instead of the intended -# diagnostic, making test failures unnecessarily hard to debug. -_PREAMBLE = ( - "set -Eeuo pipefail\n" - "log() { printf '==> %s\\n' \"$*\" >&2; }\n" - "die() { printf 'ERROR: %s\\n' \"$*\" >&2; exit 1; }\n" -) - - -def test_run_daily_does_not_early_break_on_first_stable_page() -> None: - """The regex pattern `select(test("...stable$"))] | length > 0` followed - by `break` is exactly the buggy early-stop. If it ever returns the - cron will silently start testing against a stale stable tag. - """ - body = RUN_DAILY.read_text() - assert ( - "length > 0" not in body - or "break" not in body - or ( - # If both substrings exist, make sure they aren't both inside the - # same release-pagination loop. The current loop only contains - # a `break` for the empty-page guard, not for any "length > 0" - # condition. - not _shares_loop_body(body, "length > 0", "break") - ) - ), ( - "run_daily.sh contains the old early-break-on-stable pattern. The " - "Releases endpoint orders by created_at, not semver, so breaking " - "on first-stable-seen can miss higher-versioned releases sitting " - "on later pages." - ) - - -def _shares_loop_body(body: str, needle_a: str, needle_b: str) -> bool: - """Heuristic: do both needles live inside a `for page in ...; do ... done` - block? Used as a defensive guard for the static check above.""" - in_loop = False - saw_a = False - saw_b = False - for line in body.splitlines(): - stripped = line.strip() - if stripped.startswith("for page in"): - in_loop = True - saw_a = False - saw_b = False - continue - if in_loop and stripped == "done": - if saw_a and saw_b: - return True - in_loop = False - continue - if in_loop: - if needle_a in line: - saw_a = True - if needle_b in line: - saw_b = True - return False - - -def test_run_daily_keeps_empty_page_break_guard() -> None: - """The empty-page break is the only break that should remain in the - pagination loop — without it a quiet release feed wastes API quota - walking past the last real page.""" - body = RUN_DAILY.read_text() - assert "jq 'length' \"${PAGE_JSON}\"" in body, ( - "run_daily.sh must still detect empty pages via `jq 'length' " - "${PAGE_JSON}`; without this the loop walks the full 5-page cap " - "even when there are no more releases." - ) - assert ( - '== "0"' in body - ), 'The empty-page guard must compare jq\'s length output to "0".' - - -def _make_fake_curl(scratch: Path, pages: dict[int, str]) -> Path: - """Build a fake `curl` shim that serves the canned page JSON for - each `page=N` request and an empty array for any page past the - last canned one. - - The shim mimics just enough of curl's CLI surface for the cron - script: it accepts the headers + URL we pass, ignores everything - we don't need, and writes the canned body to either stdout or the - --output target if one is given. - """ - pages_dir = scratch / "pages" - pages_dir.mkdir() - for page_num, body in pages.items(): - (pages_dir / f"page{page_num}.json").write_text(body) - - curl_path = scratch / "curl" - curl_path.write_text( - textwrap.dedent( - f"""\ - #!/usr/bin/env bash - # Fake curl for run_daily.sh release pagination tests. Serves - # page JSON from {pages_dir} keyed by the `page=` query value, - # and returns "[]" for pages past the last canned one (which - # is exactly how the real GitHub API behaves past the end). - url="" - output="" - while [[ $# -gt 0 ]]; do - case "$1" in - -fsS|-fsSL|-H|-o|--output) - if [[ "$1" == "-o" || "$1" == "--output" ]]; then - output="$2"; shift 2 - elif [[ "$1" == "-H" ]]; then - shift 2 - else - shift - fi - ;; - http*) - url="$1"; shift - ;; - *) - shift - ;; - esac - done - page="$(printf '%s' "$url" | sed -n 's/.*[?&]page=\\([0-9]*\\).*/\\1/p')" - [[ -z "$page" ]] && page=1 - file="{pages_dir}/page${{page}}.json" - if [[ -f "$file" ]]; then - if [[ -n "$output" ]]; then cp "$file" "$output"; else cat "$file"; fi - else - if [[ -n "$output" ]]; then printf '[]' > "$output"; else printf '[]'; fi - fi - """ - ) - ) - curl_path.chmod(0o755) - return curl_path - - -def _extract_resolution_snippet() -> str: - """Pull the pagination + sort_by + assignment block out of run_daily.sh - so the test exercises the actual production code path (not a copy). - - The block is everything from the GH_AUTH_HEADER setup down through - the LITELLM_VERSION emission. - """ - body = RUN_DAILY.read_text() - start = body.index("GH_AUTH_HEADER=()") - end = body.index('log "resolved litellm:') - return body[start:end] - - -@pytest.mark.skipif(shutil.which("jq") is None, reason="jq not available") -def test_run_daily_resolves_highest_semver_across_pages(tmp_path: Path) -> None: - """End-to-end: drive the actual run_daily.sh pagination loop with a - fake curl whose page 1 contains a freshly-cut LOW-version backport - (v1.80.1-stable) and page 2 contains a two-weeks-old HIGH-version - release (v1.83.0-stable). The correct behavior is to resolve - v1.83.0-stable. The pre-fix behavior would resolve v1.80.1-stable - because the early-break consumed only page 1. - """ - pages = { - # Page 1: most-recently-created releases. The order here matches - # what /releases?page=1 returns: created-at descending. The - # freshly-cut v1.80.1-stable backport sits at the top, plus a - # bunch of non-stable releases. - 1: """[ - {"tag_name": "v1.84.0-nightly.1"}, - {"tag_name": "v1.80.1-stable"}, - {"tag_name": "v1.84.0-nightly.0"} - ]""", - # Page 2: older releases. The HIGHER-version stable lives here - # because it was cut two weeks ago, before the v1.80.1 backport. - 2: """[ - {"tag_name": "v1.83.0-rc.5"}, - {"tag_name": "v1.83.0-stable"}, - {"tag_name": "v1.82.4-stable"} - ]""", - # Page 3+: empty -> the loop's empty-page guard fires here. - } - fake_curl_dir = tmp_path / "shim" - fake_curl_dir.mkdir() - _make_fake_curl(fake_curl_dir, pages) - - workdir = tmp_path / "work" - workdir.mkdir() - - snippet = _extract_resolution_snippet() - script = ( - _PREAMBLE - + f"WORKDIR={workdir!s}\n" - + snippet - + 'printf "%s" "${LITELLM_VERSION}"\n' - ) - - env = { - **os.environ, - "PATH": f"{fake_curl_dir}:{os.environ.get('PATH', '')}", - } - # Make sure the loop hits the fake curl, not the system one. - env.pop("GITHUB_TOKEN", None) - result = subprocess.run( - ["bash", "-c", script], - capture_output=True, - text=True, - env=env, - check=True, - ) - assert result.stdout == "v1.83.0-stable", ( - f"Expected the highest-semver stable across pages 1-2, got " - f"{result.stdout!r}. stderr={result.stderr!r}" - ) - - -@pytest.mark.skipif(shutil.which("jq") is None, reason="jq not available") -def test_run_daily_terminates_on_empty_page(tmp_path: Path) -> None: - """The empty-page guard must fire so we don't always walk all 5 - pages. With a single populated page and an empty page 2 we should - stop after fetching page 2 (the first empty response).""" - pages = {1: '[{"tag_name": "v1.50.0-stable"}]'} - fake_curl_dir = tmp_path / "shim" - fake_curl_dir.mkdir() - _make_fake_curl(fake_curl_dir, pages) - - workdir = tmp_path / "work" - workdir.mkdir() - - snippet = _extract_resolution_snippet() - script = ( - _PREAMBLE - + f"WORKDIR={workdir!s}\n" - + snippet - + 'printf "%s" "${LITELLM_VERSION}"\n' - ) - - env = { - **os.environ, - "PATH": f"{tmp_path}/shim:{os.environ.get('PATH', '')}", - } - env.pop("GITHUB_TOKEN", None) - result = subprocess.run( - ["bash", "-c", script], - capture_output=True, - text=True, - env=env, - check=True, - ) - assert result.stdout == "v1.50.0-stable" - # Only pages 1 and 2 should have been fetched (2 is empty -> break). - assert (workdir / "releases.page2.json").exists() - assert not (workdir / "releases.page3.json").exists(), ( - "Empty-page guard didn't fire — the loop kept walking past the " - "first empty response." - ) diff --git a/tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_version_probe_scrubs_env.py b/tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_version_probe_scrubs_env.py deleted file mode 100644 index 1c3959764f3..00000000000 --- a/tests/e2e/claude_code/_publisher_unit_tests/test_run_daily_version_probe_scrubs_env.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Pin: the cron `claude --version` probe must run under `env -i`. - -The systemd service `litellm-compat-matrix.service` loads provider -credentials (`ANTHROPIC_API_KEY`, `AWS_BEARER_TOKEN_BEDROCK`, -`AZURE_FOUNDRY_API_KEY`) and the agent-shin GitHub token -(`AGENT_SHIN_GITHUB_TOKEN`) into `run_daily.sh`'s environment from -`/etc/litellm-compat-matrix.env`. Running the npm-installed `claude` -binary directly there would hand that full env to package code, so a -compromised `@anthropic-ai/claude-code` release could read those -secrets out of `os.environ` before the proxy or test harness ever -starts. The version probe must be wrapped in `env -i` with a minimal -PATH/HOME/USER/TERM/LANG/LC_ALL/TMPDIR allowlist — matching the -PR-gate's resolver/npm-install/pytest scrubs. -""" - -from __future__ import annotations - -from pathlib import Path - -REPO_ROOT = Path(__file__).resolve().parents[4] -RUN_DAILY = REPO_ROOT / "tests" / "e2e" / "claude_code" / "cron_vm" / "run_daily.sh" - - -def _version_probe_block() -> str: - body = RUN_DAILY.read_text() - start = body.index("CLAUDE_CODE_VERSION=") - end = body.index('[[ -n "${CLAUDE_CODE_VERSION}" ]]', start) - return body[start:end] - - -def test_version_probe_wraps_claude_in_env_i() -> None: - block = _version_probe_block() - assert "env -i" in block, ( - "run_daily.sh: the `claude --version` probe must run under " - "`env -i` so a compromised @anthropic-ai/claude-code package " - "cannot read provider/GitHub credentials out of the systemd " - "service environment." - ) - assert block.index("env -i") < block.index("claude --version"), ( - "run_daily.sh: `env -i` must precede `claude --version`; " - "otherwise the binary inherits the full credential-bearing env." - ) - - -def test_version_probe_env_i_excludes_provider_secrets() -> None: - block = _version_probe_block() - for forbidden in ( - "ANTHROPIC_API_KEY", - "AWS_BEARER_TOKEN_BEDROCK", - "AWS_ACCESS_KEY_ID", - "AWS_SECRET_ACCESS_KEY", - "VERTEXAI_CREDENTIALS", - "AZURE_FOUNDRY_API_KEY", - "GITHUB_TOKEN", - "AGENT_SHIN_GITHUB_TOKEN", - ): - assert forbidden not in block, ( - f"run_daily.sh: the version-probe `env -i` allowlist must " - f"not pass {forbidden} through. Found it inside the probe " - f"block." - ) - - -def test_version_probe_uses_isolated_home_not_runtime_user_home() -> None: - """Pin: the `claude --version` probe runs under a fresh empty HOME. - - `ProtectHome=read-only` in the systemd unit allows reads of the - runtime user's real home directory. If the probe's `env -i` - block forwards `HOME=${HOME}`, a compromised `claude` package - can `os.path.expanduser("~/.config/gh/hosts.yml")` or - `os.path.expanduser("~/.ssh/...")` and exfiltrate the contents - before the proxy or test harness ever starts. The probe must - point HOME at a per-run tmpdir under `${WORKDIR}` so the CLI - sees an empty HOME instead. - """ - block = _version_probe_block() - body = RUN_DAILY.read_text() - - assert "CLAUDE_PROBE_HOME=" in body, ( - "run_daily.sh: must define a `CLAUDE_PROBE_HOME` per-run tmpdir " - "for the `claude --version` probe so the CLI never sees the " - "runtime user's real $HOME." - ) - assert 'HOME="${CLAUDE_PROBE_HOME}"' in block, ( - "run_daily.sh: the probe's `env -i` block must set HOME to " - "the per-run isolated tmpdir, not to the runtime user's $HOME." - ) - assert 'HOME="${HOME}"' not in block, ( - "run_daily.sh: the probe's `env -i` block must not forward the " - "runtime user's $HOME to `claude --version`. Use the isolated " - "$CLAUDE_PROBE_HOME tmpdir instead." - ) diff --git a/tests/e2e/claude_code/_publisher_unit_tests/test_systemd_unit_credential_isolation.py b/tests/e2e/claude_code/_publisher_unit_tests/test_systemd_unit_credential_isolation.py deleted file mode 100644 index 12edce3cb50..00000000000 --- a/tests/e2e/claude_code/_publisher_unit_tests/test_systemd_unit_credential_isolation.py +++ /dev/null @@ -1,104 +0,0 @@ -"""Pin: the cron systemd unit hides credential-bearing dotdirs. - -`ProtectHome=read-only` blocks writes to /home/mateo but still allows -reads. A model-directed `Read` tool call (the PDF cells pass -`--allowed-tools Read` to the `claude` CLI) or a compromised -`@anthropic-ai/claude-code` package can read absolute paths under -the runtime user's home and exfiltrate the contents — even with the -per-`claude`-invocation HOME isolation in place, because absolute -paths bypass `~`-expansion. - -This file pins the second line of defense: the systemd unit lists -the credential-bearing dotdirs (`~/.config/gh`, `~/.ssh`, `~/.aws`, -`~/.docker`, `~/.kube`, `~/.gnupg`) under `InaccessiblePaths=` so -the kernel hides them from every process in the unit's mount -namespace, including any child of `claude --version` or the pytest -run. It also pins that `~/.config/gh` is *not* in `ReadWritePaths=` -— we pass `GH_TOKEN` inline to every `gh` invocation in -`run_daily.sh`, so the host gh-cli config is unused. -""" - -from __future__ import annotations - -import re -from pathlib import Path - -REPO_ROOT = Path(__file__).resolve().parents[4] -SERVICE = ( - REPO_ROOT / "tests" / "e2e" / "claude_code" / "cron_vm" / "litellm-compat-matrix.service" -) - - -def _service_text() -> str: - return SERVICE.read_text() - - -def _directive(name: str) -> str: - """Return the value of a single-line systemd directive (or empty).""" - text = _service_text() - match = re.search(rf"^\s*{re.escape(name)}\s*=\s*(.*)$", text, re.MULTILINE) - return match.group(1).strip() if match else "" - - -def test_inaccessible_paths_hides_credential_dotdirs() -> None: - """Every credential-bearing dotdir must be under `InaccessiblePaths=`.""" - inaccessible = _directive("InaccessiblePaths") - assert inaccessible, ( - "litellm-compat-matrix.service: must declare `InaccessiblePaths=` " - "to hide credential dotdirs from the `claude` subprocess and the " - "model-directed Read tool. Without this, an absolute-path read " - "like `Read('/home/mateo/.config/gh/hosts.yml')` exfiltrates " - "the gh-cli token despite the per-invocation HOME isolation." - ) - for path in ( - "/home/mateo/.config/gh", - "/home/mateo/.ssh", - "/home/mateo/.aws", - "/home/mateo/.docker", - "/home/mateo/.kube", - "/home/mateo/.gnupg", - ): - # Tolerated `-` prefix means "ignore if missing on host". - assert path in inaccessible, ( - f"litellm-compat-matrix.service: `{path}` must appear in " - f"`InaccessiblePaths=` so the cron `claude` subprocess can " - f"never read it (even via an absolute path that bypasses " - f"the per-invocation HOME override)." - ) - - -def test_gh_config_is_not_writeable() -> None: - """`~/.config/gh` is not whitelisted under `ReadWritePaths=`. - - We pass `GH_TOKEN` inline to every `gh` invocation in - `run_daily.sh` (`gh repo clone`, `gh pr create`, `gh pr edit`). - The host `~/.config/gh/hosts.yml` is therefore never consulted - or written to. Keeping it out of `ReadWritePaths=` is the second - line of defense: a future regression that drops the inline-token - convention will fail loudly (gh writes a new login config and - hits a read-only filesystem) rather than silently re-introduce - the credential exfiltration surface that - `InaccessiblePaths=/home/mateo/.config/gh` is closing. - """ - rw = _directive("ReadWritePaths") - assert ".config/gh" not in rw, ( - "litellm-compat-matrix.service: `/home/mateo/.config/gh` must " - "*not* appear in `ReadWritePaths=`. We pass `GH_TOKEN` inline " - "to every `gh` invocation in run_daily.sh, so the host gh-cli " - "config is never consulted or written to. Keeping the path out " - "of ReadWritePaths means a future regression that drops the " - "inline-token convention will fail loudly instead of silently " - "re-opening the credential exfiltration surface that " - "`InaccessiblePaths=` is closing." - ) - - -def test_protect_home_is_read_only_or_stricter() -> None: - """`ProtectHome=` must be at least `read-only`.""" - value = _directive("ProtectHome") - assert value in ("read-only", "tmpfs", "yes", "true"), ( - f"litellm-compat-matrix.service: `ProtectHome=` must be `read-only`, " - f"`tmpfs`, or `yes`. Got: {value!r}. Without this, the unit can " - f"write anywhere under /home/mateo, including overwriting " - f"~/.config/gh/hosts.yml." - ) diff --git a/tests/e2e/claude_code/basic_messaging_non_streaming/test_anthropic.py b/tests/e2e/claude_code/basic_messaging_non_streaming/test_anthropic.py index c06fff28d2d..21383b85da5 100644 --- a/tests/e2e/claude_code/basic_messaging_non_streaming/test_anthropic.py +++ b/tests/e2e/claude_code/basic_messaging_non_streaming/test_anthropic.py @@ -20,6 +20,7 @@ the matrix builder still sees three rows for this (feature, provider). from __future__ import annotations +import pytest from claude_code._basic_messaging import run_basic_messaging_cell # Per the PRD: each cell is exercised against three Claude tiers via the @@ -27,11 +28,12 @@ from claude_code._basic_messaging import run_basic_messaging_cell # routing config; the driver only sends the alias. ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] +@pytest.mark.covers("llm.messages.anthropic.basic.nonstream.works") def test_basic_messaging_non_streaming_anthropic(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a reply. diff --git a/tests/e2e/claude_code/basic_messaging_non_streaming/test_azure.py b/tests/e2e/claude_code/basic_messaging_non_streaming/test_azure.py index 2a962b244a8..19e88dbe3cb 100644 --- a/tests/e2e/claude_code/basic_messaging_non_streaming/test_azure.py +++ b/tests/e2e/claude_code/basic_messaging_non_streaming/test_azure.py @@ -25,6 +25,7 @@ the matrix builder still sees three rows for this (feature, provider). from __future__ import annotations +import pytest from claude_code._basic_messaging import run_basic_messaging_cell # Per-model aliases registered in the LiteLLM proxy's routing config to @@ -33,11 +34,12 @@ from claude_code._basic_messaging import run_basic_messaging_cell # resource URL and API key. AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] +@pytest.mark.covers("llm.messages.azure_foundry.basic.nonstream.works") def test_basic_messaging_non_streaming_azure(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a reply. diff --git a/tests/e2e/claude_code/basic_messaging_non_streaming/test_bedrock_converse.py b/tests/e2e/claude_code/basic_messaging_non_streaming/test_bedrock_converse.py index 2245ed7417a..2b0f49bc205 100644 --- a/tests/e2e/claude_code/basic_messaging_non_streaming/test_bedrock_converse.py +++ b/tests/e2e/claude_code/basic_messaging_non_streaming/test_bedrock_converse.py @@ -20,6 +20,7 @@ the matrix builder still sees three rows for this (feature, provider). from __future__ import annotations +import pytest from claude_code._basic_messaging import run_basic_messaging_cell # Per-model aliases registered in the LiteLLM proxy's routing config to @@ -28,11 +29,12 @@ from claude_code._basic_messaging import run_basic_messaging_cell # strategy. BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] +@pytest.mark.covers("llm.messages.bedrock_converse.basic.nonstream.works") def test_basic_messaging_non_streaming_bedrock_converse(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a reply.""" run_basic_messaging_cell( diff --git a/tests/e2e/claude_code/basic_messaging_non_streaming/test_bedrock_invoke.py b/tests/e2e/claude_code/basic_messaging_non_streaming/test_bedrock_invoke.py index e0a6e77f3c1..937ea5ee27e 100644 --- a/tests/e2e/claude_code/basic_messaging_non_streaming/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/basic_messaging_non_streaming/test_bedrock_invoke.py @@ -20,6 +20,7 @@ the matrix builder still sees three rows for this (feature, provider). from __future__ import annotations +import pytest from claude_code._basic_messaging import run_basic_messaging_cell # Per-model aliases registered in the LiteLLM proxy's routing config to @@ -28,11 +29,12 @@ from claude_code._basic_messaging import run_basic_messaging_cell # routing strategy. BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] +@pytest.mark.covers("llm.messages.bedrock_invoke.basic.nonstream.works") def test_basic_messaging_non_streaming_bedrock_invoke(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a reply.""" run_basic_messaging_cell( diff --git a/tests/e2e/claude_code/basic_messaging_non_streaming/test_vertex_ai.py b/tests/e2e/claude_code/basic_messaging_non_streaming/test_vertex_ai.py index e4e2a39e6cd..c46e5a8f762 100644 --- a/tests/e2e/claude_code/basic_messaging_non_streaming/test_vertex_ai.py +++ b/tests/e2e/claude_code/basic_messaging_non_streaming/test_vertex_ai.py @@ -20,6 +20,7 @@ the matrix builder still sees three rows for this (feature, provider). from __future__ import annotations +import pytest from claude_code._basic_messaging import run_basic_messaging_cell # Per-model aliases registered in the LiteLLM proxy's routing config to @@ -28,11 +29,12 @@ from claude_code._basic_messaging import run_basic_messaging_cell # model id and the GCP region. VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] +@pytest.mark.covers("llm.messages.vertex.basic.nonstream.works") def test_basic_messaging_non_streaming_vertex_ai(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a reply.""" run_basic_messaging_cell( diff --git a/tests/e2e/claude_code/basic_messaging_streaming/test_anthropic.py b/tests/e2e/claude_code/basic_messaging_streaming/test_anthropic.py index 56e3fb6c181..ce453f3e523 100644 --- a/tests/e2e/claude_code/basic_messaging_streaming/test_anthropic.py +++ b/tests/e2e/claude_code/basic_messaging_streaming/test_anthropic.py @@ -25,15 +25,17 @@ sees three rows for this (feature, provider). from __future__ import annotations +import pytest from claude_code._basic_messaging import run_basic_messaging_cell ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] +@pytest.mark.covers("llm.messages.anthropic.basic.stream.works") def test_basic_messaging_streaming_anthropic(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a non-empty streamed reply (one row per Claude tier). diff --git a/tests/e2e/claude_code/basic_messaging_streaming/test_azure.py b/tests/e2e/claude_code/basic_messaging_streaming/test_azure.py index b6c002d0b27..3307194e862 100644 --- a/tests/e2e/claude_code/basic_messaging_streaming/test_azure.py +++ b/tests/e2e/claude_code/basic_messaging_streaming/test_azure.py @@ -19,15 +19,17 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations +import pytest from claude_code._basic_messaging import run_basic_messaging_cell AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] +@pytest.mark.covers("llm.messages.azure_foundry.basic.stream.works") def test_basic_messaging_streaming_azure(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a non-empty streamed reply (one row per Claude tier). diff --git a/tests/e2e/claude_code/basic_messaging_streaming/test_bedrock_converse.py b/tests/e2e/claude_code/basic_messaging_streaming/test_bedrock_converse.py index 44ac54515f0..a8bc0b77a5d 100644 --- a/tests/e2e/claude_code/basic_messaging_streaming/test_bedrock_converse.py +++ b/tests/e2e/claude_code/basic_messaging_streaming/test_bedrock_converse.py @@ -15,15 +15,17 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations +import pytest from claude_code._basic_messaging import run_basic_messaging_cell BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] +@pytest.mark.covers("llm.messages.bedrock_converse.basic.stream.works") def test_basic_messaging_streaming_bedrock_converse(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a non-empty streamed reply (one row per Claude tier). diff --git a/tests/e2e/claude_code/basic_messaging_streaming/test_bedrock_invoke.py b/tests/e2e/claude_code/basic_messaging_streaming/test_bedrock_invoke.py index 1d59d16cdc1..c0ece0e0721 100644 --- a/tests/e2e/claude_code/basic_messaging_streaming/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/basic_messaging_streaming/test_bedrock_invoke.py @@ -15,15 +15,17 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations +import pytest from claude_code._basic_messaging import run_basic_messaging_cell BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] +@pytest.mark.covers("llm.messages.bedrock_invoke.basic.stream.works") def test_basic_messaging_streaming_bedrock_invoke(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a non-empty streamed reply (one row per Claude tier). diff --git a/tests/e2e/claude_code/basic_messaging_streaming/test_vertex_ai.py b/tests/e2e/claude_code/basic_messaging_streaming/test_vertex_ai.py index 014a31160a8..13f1a0abf40 100644 --- a/tests/e2e/claude_code/basic_messaging_streaming/test_vertex_ai.py +++ b/tests/e2e/claude_code/basic_messaging_streaming/test_vertex_ai.py @@ -15,15 +15,17 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations +import pytest from claude_code._basic_messaging import run_basic_messaging_cell VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] +@pytest.mark.covers("llm.messages.vertex.basic.stream.works") def test_basic_messaging_streaming_vertex_ai(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a non-empty streamed reply (one row per Claude tier). diff --git a/tests/e2e/claude_code/conftest.py b/tests/e2e/claude_code/conftest.py index 6ee8b940648..8f5c09fa4e2 100644 --- a/tests/e2e/claude_code/conftest.py +++ b/tests/e2e/claude_code/conftest.py @@ -169,8 +169,9 @@ def _manifest_feature_ids() -> FrozenSet[str]: Used as a positive filter so only directories that correspond to a real matrix row contribute results — utility/support directories - (e.g. `cron_vm`, `_driver_unit_tests`) are dropped regardless of - naming convention, and the rate-limit summary stays clean. + (e.g. `_driver_unit_tests`, `_builder_unit_tests`) are dropped + regardless of naming convention, and the rate-limit summary stays + clean. Returns an empty set if the manifest is missing or malformed; the caller treats that as "no path is a feature path", which is the @@ -199,11 +200,10 @@ def _infer_feature_and_provider(node_path: Path) -> Optional[tuple]: Path shape: tests/e2e/claude_code//test_.py Returns None if the file is not a per-feature test (e.g. unit tests - under `_driver_unit_tests/` or support code under `cron_vm/`), so - those don't pollute the matrix artifact. We positively filter the - parent directory against `manifest.yaml` rather than relying on - naming conventions, because non-feature siblings don't all share - an underscore prefix. + under `_driver_unit_tests/`), so those don't pollute the matrix + artifact. We positively filter the parent directory against + `manifest.yaml` rather than relying on naming conventions, because + non-feature siblings don't all share an underscore prefix. """ name = node_path.name if not name.startswith("test_") or not name.endswith(".py"): @@ -548,3 +548,117 @@ def pytest_sessionfinish(session, exitstatus): ) summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True)) _print_rate_limit_summary(summary) + + +# --------------------------------------------------------------------------- +# Session-scoped compat model registration. +# +# The compat cells probe hardcoded virtual names like ``claude-sonnet-4-5`` +# and ``claude-sonnet-4-5-bedrock-invoke``. On stage those live in the +# gateway's model_list at deploy time; locally the docker-config.yaml +# under tests/e2e/ only declares one of them, so every non-haiku cell +# 400s with ``Invalid model name``. The fixture here reconciles the two: +# it reads ``test_config.yaml`` (the ground-truth compat matrix config) +# and POSTs ``/model/new`` for the subset whose provider credentials are +# actually set in the current environment, then tears them all down at +# session end. +# +# Kept below the rest of the conftest so the compat-artifact hooks stay +# grouped up top. The fixture is opt-in via autouse=True on the session +# scope, so a cell that hits the proxy sees the deployment ready without +# any per-cell wiring, and pure unit tests that never reach the proxy +# pay only one skipped-liveness check. +# --------------------------------------------------------------------------- + +from claude_code._env import ProxyConfig, resolve_proxy # noqa: E402 +from claude_code._compat_models import ( # noqa: E402 + CompatDeployment, + load_all_deployments, +) + + +def _build_control_gateway(proxy: ProxyConfig): + """Local import of the shared harness so the pure-unit-test tree + under ``_driver_unit_tests/`` etc. never has to pull it in. The + control plane transport is what /model/new lives on; SplitTransport + routes it correctly for both monolithic and split deployments. + + The endpoints come from the *resolved* proxy, not from a second + independent env read, so registration and the cells always hit the + same host and key. Both planes get the one URL the cells use; the + deployment is fronted by a single address that routes management + and LLM paths itself.""" + from e2e_gateway import build_gateway + + return build_gateway( + base_url=proxy.base_url, + master_key=proxy.api_key, + control_plane_base_url=proxy.base_url, + ) + + +def _register_deployment(gateway, deployment: CompatDeployment) -> str: + """Register one deployment and return its proxy-assigned model_id + once it is servable on the data plane.""" + return gateway.create_model( + deployment.model_name, + deployment.litellm_params, + ) + + +@pytest.fixture(scope="session", autouse=True) +def _compat_models_registered() -> Any: + """Register every compat deployment against the running proxy, then + tear them all down on session exit. + + Skips silently if the proxy env is not configured (no + ``LITELLM_PROXY_URL``/``LITELLM_MASTER_KEY``) so unit-test runs + stay hermetic. + + Design note: we always attempt to register all 15 deployments, + regardless of what credentials are exported in the test-runner's + shell. The credentials live in the proxy container's environment + (via docker-compose ``env_file``), not the shell running pytest - + so gating on shell env would filter out deployments the proxy can + actually serve. Per-deployment ``/model/new`` failures are printed + but do not abort the session: the cells that need that specific + deployment will 400 with "Invalid model name" and fail loudly, + which is the right signal (missing cred on the proxy side).""" + proxy = resolve_proxy() + if proxy is None: + yield + return + + from requests import RequestException + + gateway = _build_control_gateway(proxy) + registered_ids: list[str] = [] + failures: list[tuple[str, str]] = [] + try: + for deployment in load_all_deployments(): + try: + model_id = _register_deployment(gateway, deployment) + registered_ids.append(model_id) + except (AssertionError, RequestException) as exc: + failures.append((deployment.model_name, str(exc))) + if failures: + summary = "\n".join( + f" - {name}: {reason}" for name, reason in failures + ) + print( + f"[compat fixture] {len(failures)} of " + f"{len(failures) + len(registered_ids)} deployments " + f"failed to register (proxy likely missing that provider's " + f"credentials); cells that target them will fail loudly:\n" + f"{summary}" + ) + yield + finally: + for model_id in registered_ids: + try: + gateway.delete_model(model_id) + except (AssertionError, RequestException): + # Best-effort — teardown surfaces via warnings inside + # ``delete_model`` already; swallowing here so one flaky + # delete does not mask real test failures. + pass diff --git a/tests/e2e/claude_code/count_tokens/test_anthropic.py b/tests/e2e/claude_code/count_tokens/test_anthropic.py index 3508063459c..2fbdd4212c4 100644 --- a/tests/e2e/claude_code/count_tokens/test_anthropic.py +++ b/tests/e2e/claude_code/count_tokens/test_anthropic.py @@ -37,44 +37,27 @@ CLI rows isn't useful here. from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] +@pytest.mark.covers("llm.messages.anthropic.count_tokens.nonstream.works") def test_count_tokens_anthropic(compat_result): """Probe `/v1/messages/count_tokens` for each Anthropic tier and assert the response shape.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) failures = [] for model in ANTHROPIC_MODELS: diff --git a/tests/e2e/claude_code/count_tokens/test_azure.py b/tests/e2e/claude_code/count_tokens/test_azure.py index 2b8707b50b0..a9aa168ccea 100644 --- a/tests/e2e/claude_code/count_tokens/test_azure.py +++ b/tests/e2e/claude_code/count_tokens/test_azure.py @@ -37,44 +37,27 @@ CLI rows isn't useful here. from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] +@pytest.mark.covers("llm.messages.azure_foundry.count_tokens.nonstream.works") def test_count_tokens_azure(compat_result): """Probe `/v1/messages/count_tokens` for each Azure (Microsoft Foundry) tier and assert the response shape.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) failures = [] for model in AZURE_MODELS: diff --git a/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py b/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py index 4221773ead2..6dcff3ecae3 100644 --- a/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py +++ b/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py @@ -37,44 +37,27 @@ CLI rows isn't useful here. from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] +@pytest.mark.covers("llm.messages.bedrock_converse.count_tokens.nonstream.works") def test_count_tokens_bedrock_converse(compat_result): """Probe `/v1/messages/count_tokens` for each Bedrock (Converse) tier and assert the response shape.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) failures = [] for model in BEDROCK_CONVERSE_MODELS: diff --git a/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py b/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py index cc70bf12392..ae89067dc00 100644 --- a/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py @@ -37,44 +37,27 @@ CLI rows isn't useful here. from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] +@pytest.mark.covers("llm.messages.bedrock_invoke.count_tokens.nonstream.works") def test_count_tokens_bedrock_invoke(compat_result): """Probe `/v1/messages/count_tokens` for each Bedrock (Invoke) tier and assert the response shape.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) failures = [] for model in BEDROCK_INVOKE_MODELS: diff --git a/tests/e2e/claude_code/count_tokens/test_vertex_ai.py b/tests/e2e/claude_code/count_tokens/test_vertex_ai.py index 8c2678f7010..0f952496566 100644 --- a/tests/e2e/claude_code/count_tokens/test_vertex_ai.py +++ b/tests/e2e/claude_code/count_tokens/test_vertex_ai.py @@ -37,44 +37,27 @@ CLI rows isn't useful here. from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] +@pytest.mark.covers("llm.messages.vertex.count_tokens.nonstream.works") def test_count_tokens_vertex_ai(compat_result): """Probe `/v1/messages/count_tokens` for each Vertex AI tier and assert the response shape.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) failures = [] for model in VERTEX_AI_MODELS: diff --git a/tests/e2e/claude_code/cron_vm/build_matrix.py b/tests/e2e/claude_code/cron_vm/build_matrix.py deleted file mode 100644 index 128f041cced..00000000000 --- a/tests/e2e/claude_code/cron_vm/build_matrix.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Tiny CLI wrapper around `claude_code.matrix_builder.build_from_paths`. - -Exists only so `run_daily.sh` can hand the version metadata + paths into -the matrix builder without re-implementing it in bash. All real logic -lives in `matrix_builder.py`, which has its own unit tests under -`_builder_unit_tests/`. - -Invoked from the cron worktree (where `uv sync` has installed pyyaml), -not the dev checkout — the bash script `cd`s into the worktree before -`uv run python`-ing this file. -""" - -from __future__ import annotations - -import argparse -import datetime -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - -from claude_code.matrix_builder import build_from_paths # noqa: E402 # import needs the sys.path bootstrap above - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--manifest", type=Path, required=True) - parser.add_argument("--results", type=Path, required=True) - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--litellm-version", required=True) - parser.add_argument("--claude-code-version", required=True) - args = parser.parse_args() - - generated_at = datetime.datetime.now(datetime.timezone.utc).strftime( - "%Y-%m-%dT%H:%M:%SZ" - ) - build_from_paths( - manifest_path=args.manifest, - results_path=args.results, - litellm_version=args.litellm_version, - claude_code_version=args.claude_code_version, - generated_at=generated_at, - output_path=args.output, - ) - print(f"wrote {args.output}") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.env.example b/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.env.example deleted file mode 100644 index 5ca7937a426..00000000000 --- a/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.env.example +++ /dev/null @@ -1,59 +0,0 @@ -# Environment file consumed by `litellm-compat-matrix.service`. -# -# Install at `/etc/litellm-compat-matrix.env` and chmod 0600. -# `EnvironmentFile=-` in the unit means the service is allowed to start -# even if this file is missing, but the populator will fail at the -# first provider request without these credentials. - -# Anthropic -ANTHROPIC_API_KEY= - -# Bedrock (invoke + converse columns). -# Use Anthropic's Bedrock API-key passthrough (long-lived bearer token). -# No AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY required for the matrix -- -# both the LiteLLM invoke and converse routes pick up -# AWS_BEARER_TOKEN_BEDROCK when present. -AWS_BEARER_TOKEN_BEDROCK= -AWS_REGION_NAME=us-east-1 - -# Vertex AI. -# On the GCP VM, the default service-account ADC from the metadata server -# is used -- no JSON key file is needed. If you ever need to run outside -# GCP, also export GOOGLE_APPLICATION_CREDENTIALS=/path/to/sa.json. -VERTEXAI_PROJECT= -VERTEXAI_LOCATION=global - -# Microsoft Foundry (Azure column) -AZURE_FOUNDRY_API_KEY= -AZURE_FOUNDRY_API_BASE= - -# Azure cell of the `passthrough` row. Foundry-mode Claude Code sends -# the model in the request body, so the proxy's /azure passthrough -# cannot resolve a router alias and falls back to these env vars. -# AZURE_API_BASE is the Foundry resource's Anthropic surface, i.e. -# https://.services.ai.azure.com/anthropic ; AZURE_API_KEY -# is the same key as AZURE_FOUNDRY_API_KEY. -AZURE_API_BASE= -AZURE_API_KEY= - -# REQUIRED for publishing: PAT for the `agent-shin` user, used to push -# the daily compat-matrix branch to its fork (agent-shin/litellm-docs) -# and open the cross-repo PR against BerriAI/litellm-docs. Scopes: -# classic `repo` + `workflow`, or fine-grained on agent-shin/litellm-docs -# with Contents:RW + Pull requests:RW + Workflows:RW. -# Skip by setting SKIP_PUBLISH=1 (publishes nothing; only writes the -# matrix JSON locally). -AGENT_SHIN_GITHUB_TOKEN= - -# Optional: lifts the unauthenticated rate limit on the GitHub Releases -# API used by `resolver.py`. Any token works (read-only). Not required. -# GITHUB_TOKEN= - -# Optional overrides; defaults are sensible for the cron VM. -# PROXY_PORT=4100 -# LITELLM_WORKTREE=/home/mateo/litellm-cron-worktree -# DOCS_REPO=BerriAI/litellm-docs -# DOCS_BRANCH=main -# DOCS_TARGET_PATH=src/data/compatibility-matrix.json -# FORK_OWNER=agent-shin -# FORK_REPO=agent-shin/litellm-docs diff --git a/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.service b/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.service deleted file mode 100644 index c05ece90f50..00000000000 --- a/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.service +++ /dev/null @@ -1,141 +0,0 @@ -# systemd service for the Claude Code compatibility-matrix populator. -# -# Triggered by `litellm-compat-matrix.timer`; not started directly. The -# unit is a `Type=oneshot` so the timer's `OnCalendar=` semantics -# describe "run once per day" cleanly — there's no long-lived daemon to -# supervise; each invocation runs the populator end-to-end and exits. -# -# Install -# ------- -# -# sudo cp tests/e2e/claude_code/cron_vm/litellm-compat-matrix.service /etc/systemd/system/ -# sudo cp tests/e2e/claude_code/cron_vm/litellm-compat-matrix.timer /etc/systemd/system/ -# sudo systemctl daemon-reload -# sudo systemctl enable --now litellm-compat-matrix.timer -# -# Paths are hard-coded to /home/mateo rather than using systemd's %h -# specifier. Why: in *system* units (this one), %h is expanded at -# parse time against the *manager's* home -- which is /root for PID 1 -# -- and *not* against the User= directive. That mismatch makes -# ReadWritePaths point at /root/.cache (which doesn't exist), causing -# the namespace setup to fail with status=226/NAMESPACE before the -# script ever runs. The runtime user (`User=mateo`) must: -# -# * have a checkout of `BerriAI/litellm` at `~/litellm/litellm` so the -# publisher module is importable; -# * have a uv venv at `~/litellm/litellm/.venv` (created by -# `uv sync --frozen` inside that checkout once); -# * have `gh` already authenticated against an account with -# `pull-requests: write` on `BerriAI/litellm-docs`; -# * have provider credentials exported in `/etc/litellm-compat-matrix.env` -# (see `litellm-compat-matrix.env.example` in this directory). - -[Unit] -Description=Claude Code compatibility-matrix populator (oneshot) -Wants=network-online.target -After=network-online.target - -[Service] -Type=oneshot -User=mateo -Group=mateo - -# Provider credentials + any gh/PROXY_PORT overrides live here. Format -# is the standard `KEY=value` one line per env var. -EnvironmentFile=-/etc/litellm-compat-matrix.env - -# systemd starts with a minimal PATH (~/usr/local/bin:/usr/bin:/bin). -# `uv` and `claude` are installed under the runtime user's `~/.local/bin` -# so we have to prepend it explicitly; otherwise run_daily.sh fails at -# the up-front command-presence check. -Environment=PATH=/home/mateo/.local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin - -# `HOME` is auto-set to /home/mateo when User=mateo is honored, but be -# explicit so anything that reads $HOME (e.g. uv's cache lookup, the -# claude CLI's per-session dir) sees the right value even if a future -# refactor flips DynamicUser= or PrivateUsers= on. -Environment=HOME=/home/mateo - -WorkingDirectory=/home/mateo/litellm/litellm - -ExecStart=/home/mateo/litellm/litellm/tests/e2e/claude_code/cron_vm/run_daily.sh - -# 90 minutes is generous: cold runs do `git clone` + `uv sync` of a new -# tag's lockfile, which can take a couple of minutes on a 2-vCPU VM, -# plus 30 cells of pytest hitting four cloud providers. -TimeoutStartSec=90min - -# A failed run shouldn't restart automatically — the next timer fire is -# the right retry. Reruns of the same day's matrix are idempotent. -Restart=no - -# Security hardening: the populator only reads the litellm checkout and -# the env-file; everything else it writes lives in either the worktree -# (managed) or `/tmp` (cleaned up by tempfile). -# -# `ProtectHome=read-only` blocks writes to /home/mateo but still -# allows reads. That's safe for the trusted run_daily.sh script -# itself, but unsafe for any subprocess we don't control: a -# compromised npm-installed `claude` package, or a model-directed -# `Read` tool call during a PDF/vision cell, could read sensitive -# host files like `~/.config/gh/hosts.yml` (gh-host token), -# `~/.ssh/`, or `~/.bash_history`. We mitigate that at the call -# boundary: every `claude` subprocess (the up-front `claude --version` -# probe in run_daily.sh, plus every CLI invocation routed through -# tests/e2e/claude_code/cli_driver.py) runs with `HOME` pointed at a -# fresh empty per-invocation tmpdir, not at /home/mateo. The CLI -# never sees the runtime user's real dotfiles. `gh` invocations in -# run_daily.sh pass `GH_TOKEN` inline, so they never need to read -# ~/.config/gh either; that path is intentionally NOT in the -# whitelist below — keeping it out is the second line of defense if -# the inline-token convention is ever accidentally regressed. -# -# ReadWritePaths whitelist: -# * litellm-cron-worktree - the long-lived stable-tag checkout + -# its `.venv` (`uv sync` rewrites every -# run) + `.uv-bin` (pinned `uv` binary -# cache). -# * .cache - uv's wheel cache (~/.cache/uv) so we -# don't redownload pinned deps each -# run. Used only by the trusted `uv` -# process; not exposed to `claude`. -# * /tmp - mktemp -d workdir, proxy logs, and -# the per-`claude`-invocation isolated -# HOME tmpdirs. PrivateTmp=true below -# gives the service its own tmpfs view -# so these don't escape to the host. -NoNewPrivileges=true -ProtectSystem=strict -ProtectHome=read-only -ReadWritePaths=/home/mateo/litellm-cron-worktree /home/mateo/.cache /tmp -PrivateTmp=true - -# Filesystem-level hiding for credential-bearing dotdirs/files. Even -# though `ProtectHome=read-only` prevents writes, a model-directed -# `Read` tool call (the PDF cells pass `--allowed-tools Read`) or a -# compromised `claude` package can read absolute paths under -# /home/mateo and exfiltrate the contents. `InaccessiblePaths=` makes -# the listed paths look like empty/missing to every process in the -# unit's mount namespace -- including the trusted populator script, -# which is fine because it doesn't need any of these: -# -# * .config/gh - gh CLI host token; we pass GH_TOKEN inline to -# every `gh` invocation (clone/PR/reviewer) so the -# host config is never consulted. -# * .ssh - never used by the populator. -# * .aws - upstream AWS credentials are passed to the proxy -# via the EnvironmentFile (provider env vars), not -# via shared SDK config files. -# * .docker - the populator never talks to a docker socket. -# * .kube - the populator never talks to a k8s API. -# * .gnupg - no GPG signing on the bot's commits. -# -# Leading `-` makes systemd tolerant if a path doesn't exist on the -# host (the unit is portable across VMs that may not have all of -# them set up). Anything else under /home/mateo (the litellm -# checkout, the cron worktree, the uv cache, .local/bin for the -# claude/uv/gh binaries on PATH) stays read-accessible. -InaccessiblePaths=-/home/mateo/.config/gh -/home/mateo/.ssh -/home/mateo/.aws -/home/mateo/.docker -/home/mateo/.kube -/home/mateo/.gnupg - -[Install] -WantedBy=multi-user.target diff --git a/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.timer b/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.timer deleted file mode 100644 index ee22538c6ed..00000000000 --- a/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.timer +++ /dev/null @@ -1,25 +0,0 @@ -# Daily timer for the compatibility-matrix populator. -# -# 06:00 UTC matches the original GitHub Actions cron schedule; chosen so -# operators in US/EU timezones see fresh PRs at the start of their work -# day. -# -# `Persistent=true` causes a missed run (VM was off / suspended) to -# fire the next time the timer is started, which is the property we -# want for a once-a-day job: the matrix should refresh as soon as the -# VM is reachable again, not wait another 24h. -# -# `RandomizedDelaySec=10min` smears load if multiple matrix-style -# pipelines are ever colocated on the same VM in the future. - -[Unit] -Description=Run the Claude Code compatibility-matrix populator daily - -[Timer] -OnCalendar=*-*-* 06:00:00 UTC -Persistent=true -RandomizedDelaySec=10min -Unit=litellm-compat-matrix.service - -[Install] -WantedBy=timers.target diff --git a/tests/e2e/claude_code/cron_vm/run_daily.sh b/tests/e2e/claude_code/cron_vm/run_daily.sh deleted file mode 100755 index ae2d67c070c..00000000000 --- a/tests/e2e/claude_code/cron_vm/run_daily.sh +++ /dev/null @@ -1,590 +0,0 @@ -#!/usr/bin/env bash -# Daily Claude Code compatibility-matrix populator. -# -# Runs from the GCP VM `litellm-compatibility-matrix-populator` via the -# systemd timer in this directory. The flow is: -# -# 1. Resolve the latest LiteLLM v*-stable tag from the GitHub Releases API. -# 2. Update a long-lived worktree at $WORKTREE to that tag and `uv sync` it. -# 3. Boot the proxy as a background subprocess on $PROXY_PORT (default -# 4100; a separate port from the human-tended :4000 proxy). -# 4. Run `pytest tests/e2e/claude_code/` against the proxy. Test failures -# become `fail` cells in the JSON, not script errors. -# 5. Hand the per-test results artifact + manifest to a small Python -# CLI (`build_matrix.py`) that wraps the existing -# `matrix_builder.build_from_paths` to produce the published -# compatibility-matrix.json. -# 6. `gh repo clone` litellm-docs, write the JSON to a deterministic -# branch (`compat-matrix/--`), commit, -# `git push --force`, and `gh pr create`. -# -# Same-day reruns land on the same branch so they update the existing PR -# rather than spawning a new one. If the JSON is byte-identical to the -# docs branch, we skip the push entirely. -# -# Required commands on $PATH: git, uv, gh, jq, curl, claude. -# Required state: ~/litellm/litellm checked out (this file lives in it), -# $WORKTREE is created on first run, gh is already authenticated. -# -# Override any default by setting the matching env var; see the systemd -# unit for the production wiring. - -set -Eeuo pipefail - -LITELLM_REPO="${LITELLM_REPO:-${HOME}/litellm/litellm}" -WORKTREE="${LITELLM_WORKTREE:-${HOME}/litellm-cron-worktree}" -PROXY_PORT="${PROXY_PORT:-4100}" -PROXY_API_KEY="${PROXY_API_KEY:-sk-cron-matrix}" -DOCS_REPO="${DOCS_REPO:-BerriAI/litellm-docs}" -DOCS_BRANCH="${DOCS_BRANCH:-main}" -DOCS_TARGET_PATH="${DOCS_TARGET_PATH:-src/data/compatibility-matrix.json}" -SKIP_PUBLISH="${SKIP_PUBLISH:-0}" -PYTEST_K="${PYTEST_K:-}" -# Comma-separated GitHub usernames to request a review from on every PR. -# Reviewers must have at least read access to ${DOCS_REPO}. PR-author -# (agent-shin) has implicit rights to request reviews from anyone with -# read access, so no extra token scope is needed. Set to empty to skip. -PR_REVIEWERS="${PR_REVIEWERS:-mateo-berri}" - -POPULATOR_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -WORKDIR="$(mktemp -d -t litellm-compat-matrix.XXXXXX)" -PROXY_PID_FILE="${WORKDIR}/proxy.pid" - -# Cleanup is intentionally aggressive: it can run on normal exit, on a -# signal received by the script, or after a partial failure where the -# proxy is up but ${PROXY_PID_FILE} is stale. We try four things in -# order and stop as soon as the proxy port is free: -# -# 1. SIGTERM the pid recorded in proxy.pid. -# 2. SIGKILL anything from `pgrep -f "litellm.*--port ${PROXY_PORT}"` -# that survived. This catches the common case where the recorded -# pid was the sh wrapper, not the long-lived python child. -# 3. ss -K on the port (kernel kills sockets but not processes; -# mostly useful for catching lingering CLOSE_WAITs). -# 4. wipe ${WORKDIR}. -cleanup() { - local rc=$? - set +e - local proxy_pid - if [[ -f "${PROXY_PID_FILE}" ]]; then - proxy_pid="$(cat "${PROXY_PID_FILE}")" - if [[ -n "${proxy_pid}" ]]; then - kill -TERM "-${proxy_pid}" 2>/dev/null || kill -TERM "${proxy_pid}" 2>/dev/null || true - for _ in 1 2 3 4 5; do - kill -0 "${proxy_pid}" 2>/dev/null || break - sleep 1 - done - fi - fi - # Belt-and-braces: any python or uv talking to ${PROXY_PORT} that - # survived the SIGTERM gets SIGKILL'd by name. - pgrep -f "litellm.*--port[ =]?${PROXY_PORT}([^0-9]|$)" 2>/dev/null \ - | xargs -r kill -KILL 2>/dev/null || true - pgrep -f "${WORKTREE}/.uv-bin/uv.*run litellm" 2>/dev/null \ - | xargs -r kill -KILL 2>/dev/null || true - rm -rf "${WORKDIR}" - exit "${rc}" -} -trap cleanup EXIT INT TERM - -log() { printf '==> %s\n' "$*" >&2; } -die() { printf 'ERROR: %s\n' "$*" >&2; exit 1; } - -for cmd in git uv gh jq curl claude; do - command -v "${cmd}" >/dev/null 2>&1 || die "missing required command: ${cmd}" -done - -# Publishing is from a fork (agent-shin/litellm-docs) so neither the cron -# host nor the bot identity needs write access to BerriAI/litellm-docs. We -# require the fork token up front -- failing 30 minutes into a run because -# the env file is missing one line is a waste of CI quota. -if [[ "${SKIP_PUBLISH}" != "1" ]]; then - [[ -n "${AGENT_SHIN_GITHUB_TOKEN:-}" ]] \ - || die "AGENT_SHIN_GITHUB_TOKEN required to open PRs from agent-shin/litellm-docs (or set SKIP_PUBLISH=1)" -fi - -# --------------------------------------------------------------------------- -# 1. Resolve versions -# --------------------------------------------------------------------------- - -# Newest v*-stable release on BerriAI/litellm. The `select(...)` filter -# drops drafts/non-stable, the version_key sort handles 1.10 > 1.9. -# -# Paginate through the releases endpoint instead of grabbing only page 1 -# (default page_size=30). LiteLLM ships multiple non-stable releases per -# day, so it's common to need to walk past 30+ entries before hitting -# the most recent v*-stable. We cap at 5 pages (500 releases) which is -# conservatively beyond the worst observed gap. -# -# We deliberately do NOT short-circuit on the first page that contains a -# v*-stable tag. The /releases endpoint orders by `created_at`, not by -# semver, so a backport on an older series (e.g. v1.80.1-stable cut -# today) can show up on an earlier page than a higher-versioned release -# (v1.83.0-stable cut two weeks ago). Breaking early on first-stable-seen -# would silently pin the cron to the stale tag because the -# higher-versioned release still on a later page would never make it -# into the merged set the `sort_by` below consumes. The only break we -# keep is the empty-page guard, which means a quiet period in the -# release feed doesn't waste API quota — we just always walk far enough -# to be confident we've seen the highest stable tag. -GH_AUTH_HEADER=() -if [[ -n "${GITHUB_TOKEN:-}" ]]; then - GH_AUTH_HEADER=(-H "Authorization: Bearer ${GITHUB_TOKEN}") -fi -RELEASES_JSON="${WORKDIR}/releases.json" -echo "[]" >"${RELEASES_JSON}" -for page in 1 2 3 4 5; do - PAGE_JSON="${WORKDIR}/releases.page${page}.json" - curl -fsS \ - -H 'Accept: application/vnd.github+json' \ - -H 'User-Agent: litellm-compat-matrix' \ - "${GH_AUTH_HEADER[@]}" \ - "https://api.github.com/repos/BerriAI/litellm/releases?per_page=100&page=${page}" \ - >"${PAGE_JSON}" - jq -s '.[0] + .[1]' "${RELEASES_JSON}" "${PAGE_JSON}" >"${RELEASES_JSON}.merged" - mv "${RELEASES_JSON}.merged" "${RELEASES_JSON}" - # No more pages? GitHub returns an empty array past the last page. - if [[ "$(jq 'length' "${PAGE_JSON}")" == "0" ]]; then - break - fi -done -LITELLM_VERSION="$( - jq -r ' - [ .[] | .tag_name // empty - | select(test("^v[0-9]+\\.[0-9]+\\.[0-9]+-stable$")) - ] - | sort_by( - capture("^v(?[0-9]+)\\.(?[0-9]+)\\.(?[0-9]+)-stable$") - | [(.a|tonumber), (.b|tonumber), (.c|tonumber)] - ) - | last // empty - ' "${RELEASES_JSON}" -)" -[[ -n "${LITELLM_VERSION}" ]] || die "could not resolve latest v*-stable tag in 5 pages of releases" -log "resolved litellm: ${LITELLM_VERSION}" - -# The systemd unit loads provider credentials and the agent-shin GitHub -# token from /etc/litellm-compat-matrix.env into this script's -# environment. Running the npm-installed `claude` binary directly here -# would hand that full env to package code -- a compromised -# @anthropic-ai/claude-code release could read ANTHROPIC_API_KEY / -# AWS_BEARER_TOKEN_BEDROCK / AZURE_FOUNDRY_API_KEY / -# AGENT_SHIN_GITHUB_TOKEN from os.environ and exfiltrate them before -# the proxy or test harness ever starts. Probe under `env -i` with the -# same minimal allowlist the PR-gate uses (the matrix run itself goes -# through cli_driver.py, which already scrubs the CLI env). -# -# The probe also runs under a fresh empty HOME instead of the runtime -# user's real $HOME. `ProtectHome=read-only` in the systemd unit -# blocks *writes* to /home/mateo but still allows reads, so a -# compromised claude package invoked here with HOME=/home/mateo could -# read ~/.config/gh/hosts.yml (the gh-host token), ~/.bash_history, -# or ~/.ssh/. Pointing HOME at a per-run dir under ${WORKDIR} hides -# those entirely from the subprocess; ${WORKDIR} is rm -rf'd by the -# script-wide cleanup() trap regardless of probe outcome. -CLAUDE_PROBE_HOME="${WORKDIR}/claude-probe-home" -mkdir -p "${CLAUDE_PROBE_HOME}" -CLAUDE_CODE_VERSION="$(env -i \ - PATH="${PATH}" \ - HOME="${CLAUDE_PROBE_HOME}" \ - USER="${USER:-mateo}" \ - TERM="${TERM:-dumb}" \ - LANG="${LANG:-C.UTF-8}" \ - LC_ALL="${LC_ALL:-}" \ - TMPDIR="${TMPDIR:-/tmp}" \ - claude --version 2>/dev/null \ - | grep -oE '[0-9]+\.[0-9]+\.[0-9]+([.-][A-Za-z0-9.-]+)?' \ - | head -n1 || true)" -# `|| true` above keeps `set -Eeuo pipefail` from aborting silently when -# `grep` finds no match (exit 1) — without it the assignment inherits the -# pipeline's non-zero exit, `set -e` kills the script, and the operator -# never sees the helpful diagnostic below. -[[ -n "${CLAUDE_CODE_VERSION}" ]] || die "could not parse semver from 'claude --version'" -log "local claude code: ${CLAUDE_CODE_VERSION}" - -# --------------------------------------------------------------------------- -# 2. Update the worktree to that tag -# --------------------------------------------------------------------------- - -if [[ ! -d "${WORKTREE}/.git" ]]; then - log "first run: cloning litellm into ${WORKTREE}" - mkdir -p "$(dirname "${WORKTREE}")" - git clone https://github.com/BerriAI/litellm.git "${WORKTREE}" -fi - -log "updating worktree to ${LITELLM_VERSION}" -git -C "${WORKTREE}" fetch --tags --force -git -C "${WORKTREE}" reset --hard -# Keep the venv and the .uv-bin cache around — uv sync will reconcile -# the venv on every run, and we don't want to re-download the pinned -# uv binary each time. Drop everything else (including any prior -# tests/e2e/claude_code/ shim) so each run starts clean before the shim -# below rewrites it from the dev checkout. -git -C "${WORKTREE}" clean -fdx -e .venv -e .uv-bin -git -C "${WORKTREE}" checkout --force "${LITELLM_VERSION}" - -# Always overwrite tests/e2e/claude_code/ in the worktree with the copy -# from the dev checkout, regardless of whether the resolved -# ${LITELLM_VERSION} tag already ships a tests/e2e/claude_code/ tree of -# its own. Rationale: the matrix populator's job is to exercise -# today's tests against the latest stable proxy. The dev checkout -# carries the most recent test fixes (e.g. the stream-json vision -# rewrite, the --effort thinking knob, the WebSearch tool_use -# assertion) that haven't yet rolled into a v*-stable, and we want -# every cron run to pick those up the moment they land on -# ${LITELLM_REPO}, not whenever the next stable release happens. -# -# Concretely this means a fresh `rm -rf` + `cp -r` every run so the -# tree is byte-identical to ${LITELLM_REPO}/tests/e2e/claude_code (no -# stale files left over from the tag's own checkout, no drift across -# runs). -if [[ ! -d "${LITELLM_REPO}/tests/e2e/claude_code" ]]; then - die "no shim source at ${LITELLM_REPO}/tests/e2e/claude_code" -fi -log "shimming tests/e2e/claude_code/ from ${LITELLM_REPO} (always-overwrite)" -rm -rf "${WORKTREE}/tests/e2e/claude_code" -mkdir -p "${WORKTREE}/tests/e2e" -cp -r "${LITELLM_REPO}/tests/e2e/claude_code" "${WORKTREE}/tests/e2e/" - -# litellm pins an exact uv version in pyproject.toml's [tool.uv] -# `required-version` field, so a system uv that's newer or older -# refuses to sync. We pin our own local copy at the version the -# checked-out tag asks for, cached under .uv-bin/ inside the worktree -# so subsequent runs skip the download. -PINNED_UV_VERSION="$( - awk -F'"' ' - /^required-version[[:space:]]*=/ { - # Field 2 is the value between the quotes, e.g. ">=0.10.9" or - # "0.10.9". Strip any leading specifier prefix so we end up with - # the bare version string, which is what /releases/download// - # expects. - v = $2 - sub(/^[[:space:]=<>!~]+/, "", v) - if (v != "") { print v; exit } - } - ' "${WORKTREE}/pyproject.toml" -)" -if [[ -z "${PINNED_UV_VERSION}" ]]; then - log "no uv version pin in pyproject.toml; using system uv" - WORKTREE_UV="$(command -v uv)" -else - WORKTREE_UV="${WORKTREE}/.uv-bin/uv-${PINNED_UV_VERSION}" - if [[ ! -x "${WORKTREE_UV}" ]]; then - log "downloading uv ${PINNED_UV_VERSION} for the worktree" - mkdir -p "${WORKTREE}/.uv-bin" - # Detect host arch so the same script works on x86_64 GCP VMs and on - # aarch64 hosts (Astral publishes both `uv-x86_64-unknown-linux-gnu` - # and `uv-aarch64-unknown-linux-gnu` tarballs under the same release - # tag, and `uname -m` already returns the exact token uv uses). - UV_ARCH="$(uname -m)" - UV_TRIPLE="uv-${UV_ARCH}-unknown-linux-gnu" - UV_TARBALL_NAME="${UV_TRIPLE}.tar.gz" - UV_DOWNLOAD_URL="https://github.com/astral-sh/uv/releases/download/${PINNED_UV_VERSION}/${UV_TARBALL_NAME}" - UV_TMPDIR="$(mktemp -d -t uv-download.XXXXXX)" - # Download the tarball and Astral's official .sha256 sidecar to disk - # and verify the digest before extracting/executing anything. This - # closes the supply-chain trust gap of piping a remote binary - # straight into `tar -xzO ... > file ; chmod +x` (see CLAUDE.md - # "CI Supply-Chain Safety"). - curl -fsSL --output "${UV_TMPDIR}/${UV_TARBALL_NAME}" "${UV_DOWNLOAD_URL}" - curl -fsSL --output "${UV_TMPDIR}/${UV_TARBALL_NAME}.sha256" "${UV_DOWNLOAD_URL}.sha256" - (cd "${UV_TMPDIR}" && sha256sum -c "${UV_TARBALL_NAME}.sha256") \ - || { rm -rf "${UV_TMPDIR}"; die "uv ${PINNED_UV_VERSION} sha256 mismatch — refusing to install"; } - tar -xzf "${UV_TMPDIR}/${UV_TARBALL_NAME}" -C "${UV_TMPDIR}" "${UV_TRIPLE}/uv" - mv "${UV_TMPDIR}/${UV_TRIPLE}/uv" "${WORKTREE_UV}.tmp" - chmod +x "${WORKTREE_UV}.tmp" - mv "${WORKTREE_UV}.tmp" "${WORKTREE_UV}" - rm -rf "${UV_TMPDIR}" - fi -fi -# `--extra proxy` pulls fastapi/uvicorn/etc. so `uv run litellm` can -# actually serve. `--group proxy-dev` brings in pytest and the rest of -# what tests/e2e/claude_code/ needs. -log "uv sync --frozen --group proxy-dev --extra proxy (uv ${PINNED_UV_VERSION:-system})" -(cd "${WORKTREE}" && "${WORKTREE_UV}" sync --frozen --group proxy-dev --extra proxy) - -PROXY_CONFIG="${WORKTREE}/tests/e2e/claude_code/test_config.yaml" -[[ -f "${PROXY_CONFIG}" ]] || die "proxy config not found at ${PROXY_CONFIG} (does ${LITELLM_VERSION} predate the compat matrix work?)" - -# --------------------------------------------------------------------------- -# 3. Boot the proxy -# --------------------------------------------------------------------------- - -log "starting proxy on 127.0.0.1:${PROXY_PORT}" -# Bind the proxy to loopback only. The populator proxy is talked to -# exclusively by the pytest run on the same host (the health check and -# the test env set `LITELLM_PROXY_BASE_URL=http://127.0.0.1:...`), -# so there's no reason to expose it on the VM's external interfaces. -# Without `--host`, `litellm` defaults to 0.0.0.0, which combined with -# the predictable default `LITELLM_MASTER_KEY=sk-cron-matrix` would -# allow anything that can reach :${PROXY_PORT} on the VM to authenticate -# and burn upstream provider credentials. -# -# `setsid` puts the proxy in its own session+pgroup so cleanup() can -# SIGTERM the whole tree by passing the pgid as a negative pid. We -# write that pid to a file so cleanup() doesn't need to remember a -# variable that might be stale by the time the trap fires. -# -# Pass the master key as a shell-prefix assignment on `setsid` (inherited -# via the environment) rather than as `env KEY=VAL ...` argv. The argv -# form would land the literal key in /proc//cmdline, where -# any local reader (a model-directed `Read` tool call, another user on -# the VM, a crash dump) could pick it up before the process execs into -# the litellm child. The shell-prefix form keeps the key out of argv at -# every layer (setsid → bash → uv → litellm). -LITELLM_MASTER_KEY="${PROXY_API_KEY}" setsid bash -c ' - echo "$$" > "$0" - cd "$1" - exec "$2" run litellm --config "$3" --host 127.0.0.1 --port "$4" -' "${PROXY_PID_FILE}" "${WORKTREE}" "${WORKTREE_UV}" "${PROXY_CONFIG}" "${PROXY_PORT}" \ - >"${WORKDIR}/proxy.log" 2>&1 & -disown - -HEALTH_URL="http://127.0.0.1:${PROXY_PORT}/health/liveliness" -for _ in $(seq 1 45); do - if curl -fsS "${HEALTH_URL}" >/dev/null 2>&1; then - break - fi - sleep 2 -done -curl -fsS "${HEALTH_URL}" >/dev/null \ - || { tail -50 "${WORKDIR}/proxy.log" >&2; die "proxy did not become healthy"; } - -# --------------------------------------------------------------------------- -# 4. Run pytest -# --------------------------------------------------------------------------- - -RESULTS_JSON="${WORKDIR}/compat-results.json" -PYTEST_ARGS=( - tests/e2e/claude_code/ - --ignore=tests/e2e/claude_code/_driver_unit_tests - --ignore=tests/e2e/claude_code/_builder_unit_tests - --ignore=tests/e2e/claude_code/_publisher_unit_tests - --ignore=tests/e2e/claude_code/_pr_gate_unit_tests -) -if [[ -n "${PYTEST_K}" ]]; then - log "PYTEST_K set; narrowing to: ${PYTEST_K}" - PYTEST_ARGS+=(-k "${PYTEST_K}") -fi - -log "running pytest" -set +e -# Pytest only needs to talk to the loopback proxy at 127.0.0.1:${PROXY_PORT} -# — it has no legitimate reason to see ANTHROPIC_API_KEY / -# AWS_BEARER_TOKEN_BEDROCK / VERTEXAI_* / AZURE_FOUNDRY_* / -# AGENT_SHIN_GITHUB_TOKEN / GITHUB_TOKEN in its own env. The systemd -# unit's EnvironmentFile injects all of those into this script for the -# proxy to consume, and pytest inherits them by default. Wrap the -# invocation in `env -i` so: -# -# 1. test code under tests/e2e/claude_code/ (or anything it imports) -# cannot read provider/agent-shin creds out of `os.environ` and -# exfiltrate them via an outbound call from inside a conftest hook -# or a fixture (a sibling vector to the model-controlled Bash/Read -# concern handled by `cli_driver.py`'s own env scrub); -# 2. a model-directed `Read` tool call during a PDF/vision cell -# cannot reach /proc//environ and pull the creds out -# of the parent process the way it can today; -# 3. this matches the PR-gate pytest step in `.circleci/config.yml`, -# which already runs under `env -i` with the same minimal -# allowlist. -# -# `cli_driver.py` re-allowlists its own subset (PATH/USER/LOGNAME/etc.) -# when spawning the `claude` binary, so the CLI still finds Node + the -# claude shim on PATH and gets a fresh isolated HOME per invocation. -( - cd "${WORKTREE}" \ - && env -i \ - PATH="${PATH}" \ - HOME="${HOME}" \ - USER="${USER:-mateo}" \ - TERM="${TERM:-dumb}" \ - LANG="${LANG:-C.UTF-8}" \ - LC_ALL="${LC_ALL:-}" \ - TMPDIR="${TMPDIR:-/tmp}" \ - LITELLM_PROXY_BASE_URL="http://127.0.0.1:${PROXY_PORT}" \ - LITELLM_PROXY_API_KEY="${PROXY_API_KEY}" \ - COMPAT_RESULTS_PATH="${RESULTS_JSON}" \ - "${WORKTREE_UV}" run pytest "${PYTEST_ARGS[@]}" -) -PYTEST_EXIT=$? -set -e -log "pytest exit code: ${PYTEST_EXIT} (failures become 'fail' cells, not script errors)" -[[ -f "${RESULTS_JSON}" ]] || die "pytest did not produce ${RESULTS_JSON}" - -# --------------------------------------------------------------------------- -# 5. Build the matrix JSON -# --------------------------------------------------------------------------- - -MATRIX_JSON="${WORKDIR}/compatibility-matrix.json" -log "building ${MATRIX_JSON}" -( - cd "${WORKTREE}" \ - && "${WORKTREE_UV}" run python "${POPULATOR_DIR}/build_matrix.py" \ - --manifest "${WORKTREE}/tests/e2e/claude_code/manifest.yaml" \ - --results "${RESULTS_JSON}" \ - --output "${MATRIX_JSON}" \ - --litellm-version "${LITELLM_VERSION}" \ - --claude-code-version "${CLAUDE_CODE_VERSION}" -) - -# --------------------------------------------------------------------------- -# 6. Open a docs-repo PR -# --------------------------------------------------------------------------- - -if [[ "${SKIP_PUBLISH}" == "1" ]]; then - cp "${MATRIX_JSON}" "${LITELLM_REPO}/compatibility-matrix.json" - log "SKIP_PUBLISH=1; matrix written to ${LITELLM_REPO}/compatibility-matrix.json" - exit 0 -fi - -DATE_UTC="$(date -u +%Y-%m-%d)" -BRANCH_NAME="compat-matrix/${LITELLM_VERSION}-${CLAUDE_CODE_VERSION}-${DATE_UTC}" -DOCS_CLONE="${WORKDIR}/litellm-docs" -FORK_OWNER="${FORK_OWNER:-agent-shin}" -FORK_REPO="${FORK_REPO:-${FORK_OWNER}/litellm-docs}" - -log "cloning ${DOCS_REPO}@${DOCS_BRANCH}" -# Use the agent-shin token inline rather than the host gh-cli config. -# `BerriAI/litellm-docs` is a public repo so unauthenticated clone -# would also work, but passing the token explicitly means the systemd -# unit can hide `~/.config/gh` (`InaccessiblePaths=`) without breaking -# this clone — closing the model-directed `Read("/home/mateo/.config/gh/...")` -# exfiltration path on the cron VM. -GH_TOKEN="${AGENT_SHIN_GITHUB_TOKEN}" \ - gh repo clone "${DOCS_REPO}" "${DOCS_CLONE}" -- --depth 1 --branch "${DOCS_BRANCH}" - -cd "${DOCS_CLONE}" -git config user.email "litellm-bot@berri.ai" -git config user.name "litellm-compat-matrix-bot" -git checkout -b "${BRANCH_NAME}" - -mkdir -p "$(dirname "${DOCS_TARGET_PATH}")" -cp "${MATRIX_JSON}" "${DOCS_TARGET_PATH}" -git add "${DOCS_TARGET_PATH}" - -if git diff --cached --quiet; then - log "matrix JSON unchanged from ${DOCS_BRANCH}; skipping PR" - exit 0 -fi - -GENERATED_AT="$(jq -r '.generated_at' "${MATRIX_JSON}")" -COMMIT_MSG="$(cat </dev/null || true -git remote add fork "${FORK_PUSH_URL}" -git push --force --set-upstream fork "${BRANCH_NAME}" -git remote remove fork -unset FORK_PUSH_URL - -# Per-feature status table for the PR body. Reviewers triage from this. -PR_FEATURE_TABLE="$(jq -r ' - .features[] as $f - | "- **\($f.name)**: " + - ([ .providers[] as $p - | "\($p)=\($f.providers[$p].status // "not_tested")" - ] | join(", ")) -' "${MATRIX_JSON}")" - -PR_TITLE="chore(compat-matrix): refresh for ${LITELLM_VERSION} + claude-code ${CLAUDE_CODE_VERSION}" -PR_BODY="$(cat < ${DOCS_REPO}:${DOCS_BRANCH}" -# GH_TOKEN here is scoped to this single subshell so we don't bleed the -# fork token into the rest of the script (release-listing earlier uses -# ${GITHUB_TOKEN}, which may be a different identity). gh's --head accepts -# `OWNER:BRANCH` for cross-repo PRs from a fork. -# -# Reviewer assignment is done in a *separate* call below: as the PR -# author from a fork, agent-shin has no write/triage access on -# ${DOCS_REPO} and the `RequestReviewsByLogin` GraphQL mutation -# (which backs `gh pr create --reviewer` and `gh pr edit --add-reviewer`) -# rejects with "does not have the correct permissions". We use the -# collaborator-scoped ${GITHUB_TOKEN} for that instead. Don't fold -# --reviewer into `gh pr create` here -- it would fail the whole -# create on the very first cron run. -set +e -PR_OUT="$( - GH_TOKEN="${AGENT_SHIN_GITHUB_TOKEN}" gh pr create \ - --repo "${DOCS_REPO}" \ - --base "${DOCS_BRANCH}" \ - --head "${FORK_OWNER}:${BRANCH_NAME}" \ - --title "${PR_TITLE}" \ - --body "${PR_BODY}" 2>&1 -)" -PR_EXIT=$? -set -e -echo "${PR_OUT}" - -if [[ ${PR_EXIT} -ne 0 ]]; then - if grep -q "a pull request for branch.*already exists" <<<"${PR_OUT}"; then - log "PR already exists for ${FORK_OWNER}:${BRANCH_NAME}; updated branch in place" - else - die "gh pr create failed (exit ${PR_EXIT})" - fi -fi - -# Request reviews from PR_REVIEWERS using the collaborator-scoped -# ${GITHUB_TOKEN} (mateo-berri's token, already provisioned for release -# listing). This is idempotent: `gh pr edit --add-reviewer` is a no-op -# on a user who's already in reviewRequests, and silently re-adds -# anyone whose prior review was dismissed -- so same-day reruns stay -# clean. Reviewer-add failures are non-fatal: the matrix JSON has -# already landed on the PR; the worst case is a manual ping. -if [[ -n "${PR_REVIEWERS}" ]]; then - if [[ -z "${GITHUB_TOKEN:-}" ]]; then - log "WARN: PR_REVIEWERS set but GITHUB_TOKEN missing -- cannot request reviews; skipping" - else - log "requesting reviews from: ${PR_REVIEWERS}" - set +e - GH_TOKEN="${GITHUB_TOKEN}" gh pr edit \ - "${FORK_OWNER}:${BRANCH_NAME}" \ - --repo "${DOCS_REPO}" \ - --add-reviewer "${PR_REVIEWERS}" 2>&1 | sed 's/^/ /' - REVIEWER_EXIT=${PIPESTATUS[0]} - set -e - if [[ ${REVIEWER_EXIT} -ne 0 ]]; then - log "WARN: gh pr edit --add-reviewer exited ${REVIEWER_EXIT} (non-fatal)" - fi - fi -fi - -log "done" diff --git a/tests/e2e/claude_code/long_context_1m/test_anthropic.py b/tests/e2e/claude_code/long_context_1m/test_anthropic.py index fb74d5fd40d..b9bbd1c2fe7 100644 --- a/tests/e2e/claude_code/long_context_1m/test_anthropic.py +++ b/tests/e2e/claude_code/long_context_1m/test_anthropic.py @@ -54,25 +54,23 @@ $10/day on this row. from __future__ import annotations -import os from typing import Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" # Haiku 4.5 is excluded -- only Sonnet 4.6 and Opus 4.7 support the # 1M-context beta. See module docstring for the per-cell-aggregator # rationale. ANTHROPIC_MODELS: Sequence[str] = ( - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ) @@ -155,26 +153,12 @@ def _build_long_prompt(target_tokens: int = TARGET_INPUT_TOKENS) -> str: return preamble + "".join(pad_lines) + closing +@pytest.mark.covers("llm.messages.anthropic.long_context_1m.nonstream.works") def test_long_context_1m_anthropic(compat_result): """Drive the `claude` CLI with a ~210k-token prompt and the `context-1m-2025-08-07` beta header; assert no 400 / 413 and a non-empty reply for Sonnet + Opus.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) long_prompt = _build_long_prompt() diff --git a/tests/e2e/claude_code/long_context_1m/test_azure.py b/tests/e2e/claude_code/long_context_1m/test_azure.py index 5800fdadbfc..d62214d2758 100644 --- a/tests/e2e/claude_code/long_context_1m/test_azure.py +++ b/tests/e2e/claude_code/long_context_1m/test_azure.py @@ -54,25 +54,23 @@ $10/day on this row. from __future__ import annotations -import os from typing import Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" # Haiku 4.5 is excluded -- only Sonnet 4.6 and Opus 4.7 support the # 1M-context beta. See module docstring for the per-cell-aggregator # rationale. AZURE_MODELS: Sequence[str] = ( - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ) @@ -155,26 +153,12 @@ def _build_long_prompt(target_tokens: int = TARGET_INPUT_TOKENS) -> str: return preamble + "".join(pad_lines) + closing +@pytest.mark.covers("llm.messages.azure_foundry.long_context_1m.nonstream.works") def test_long_context_1m_azure(compat_result): """Drive the `claude` CLI (Azure (Microsoft Foundry)) with a ~210k-token prompt and the `context-1m-2025-08-07` beta header; assert no 400 / 413 and a non-empty reply for Sonnet + Opus.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) long_prompt = _build_long_prompt() diff --git a/tests/e2e/claude_code/long_context_1m/test_bedrock_converse.py b/tests/e2e/claude_code/long_context_1m/test_bedrock_converse.py index 18587f7c2d6..3c2fd4f02cc 100644 --- a/tests/e2e/claude_code/long_context_1m/test_bedrock_converse.py +++ b/tests/e2e/claude_code/long_context_1m/test_bedrock_converse.py @@ -54,25 +54,23 @@ $10/day on this row. from __future__ import annotations -import os from typing import Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" # Haiku 4.5 is excluded -- only Sonnet 4.6 and Opus 4.7 support the # 1M-context beta. See module docstring for the per-cell-aggregator # rationale. BEDROCK_CONVERSE_MODELS: Sequence[str] = ( - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ) @@ -155,26 +153,12 @@ def _build_long_prompt(target_tokens: int = TARGET_INPUT_TOKENS) -> str: return preamble + "".join(pad_lines) + closing +@pytest.mark.covers("llm.messages.bedrock_converse.long_context_1m.nonstream.works") def test_long_context_1m_bedrock_converse(compat_result): """Drive the `claude` CLI (Bedrock (Converse)) with a ~210k-token prompt and the `context-1m-2025-08-07` beta header; assert no 400 / 413 and a non-empty reply for Sonnet + Opus.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) long_prompt = _build_long_prompt() diff --git a/tests/e2e/claude_code/long_context_1m/test_bedrock_invoke.py b/tests/e2e/claude_code/long_context_1m/test_bedrock_invoke.py index 0270197ce2a..4801d405760 100644 --- a/tests/e2e/claude_code/long_context_1m/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/long_context_1m/test_bedrock_invoke.py @@ -54,25 +54,23 @@ $10/day on this row. from __future__ import annotations -import os from typing import Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" # Haiku 4.5 is excluded -- only Sonnet 4.6 and Opus 4.7 support the # 1M-context beta. See module docstring for the per-cell-aggregator # rationale. BEDROCK_INVOKE_MODELS: Sequence[str] = ( - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ) @@ -155,26 +153,12 @@ def _build_long_prompt(target_tokens: int = TARGET_INPUT_TOKENS) -> str: return preamble + "".join(pad_lines) + closing +@pytest.mark.covers("llm.messages.bedrock_invoke.long_context_1m.nonstream.works") def test_long_context_1m_bedrock_invoke(compat_result): """Drive the `claude` CLI (Bedrock (Invoke)) with a ~210k-token prompt and the `context-1m-2025-08-07` beta header; assert no 400 / 413 and a non-empty reply for Sonnet + Opus.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) long_prompt = _build_long_prompt() diff --git a/tests/e2e/claude_code/long_context_1m/test_vertex_ai.py b/tests/e2e/claude_code/long_context_1m/test_vertex_ai.py index d2db4a1b4ee..efa96bf076d 100644 --- a/tests/e2e/claude_code/long_context_1m/test_vertex_ai.py +++ b/tests/e2e/claude_code/long_context_1m/test_vertex_ai.py @@ -54,25 +54,23 @@ $10/day on this row. from __future__ import annotations -import os from typing import Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" # Haiku 4.5 is excluded -- only Sonnet 4.6 and Opus 4.7 support the # 1M-context beta. See module docstring for the per-cell-aggregator # rationale. VERTEX_AI_MODELS: Sequence[str] = ( - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ) @@ -155,26 +153,12 @@ def _build_long_prompt(target_tokens: int = TARGET_INPUT_TOKENS) -> str: return preamble + "".join(pad_lines) + closing +@pytest.mark.covers("llm.messages.vertex.long_context_1m.nonstream.works") def test_long_context_1m_vertex_ai(compat_result): """Drive the `claude` CLI (Vertex AI) with a ~210k-token prompt and the `context-1m-2025-08-07` beta header; assert no 400 / 413 and a non-empty reply for Sonnet + Opus.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) long_prompt = _build_long_prompt() diff --git a/tests/e2e/claude_code/matrix_builder.py b/tests/e2e/claude_code/matrix_builder.py index 5641e488da2..d9a13d17ea4 100644 --- a/tests/e2e/claude_code/matrix_builder.py +++ b/tests/e2e/claude_code/matrix_builder.py @@ -183,7 +183,7 @@ def build_from_paths( generated_at: str, output_path: Optional[Path] = None, ) -> Dict[str, Any]: - """I/O wrapper around build_matrix used by the publisher script.""" + """I/O wrapper around ``build_matrix``: reads the manifest and per-test results from disk, calls ``build_matrix``, and (optionally) writes the compat-matrix JSON to ``output_path``. Whatever orchestrator publishes the matrix (currently the ECR image) invokes this.""" manifest = load_manifest(manifest_path) results = load_results(results_path) matrix = build_matrix( diff --git a/tests/e2e/claude_code/passthrough/test_anthropic.py b/tests/e2e/claude_code/passthrough/test_anthropic.py index aa0443e0625..8382342ae12 100644 --- a/tests/e2e/claude_code/passthrough/test_anthropic.py +++ b/tests/e2e/claude_code/passthrough/test_anthropic.py @@ -29,7 +29,7 @@ from claude_code._passthrough import ( ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] diff --git a/tests/e2e/claude_code/passthrough/test_azure.py b/tests/e2e/claude_code/passthrough/test_azure.py index 09b0824047a..21100a49c16 100644 --- a/tests/e2e/claude_code/passthrough/test_azure.py +++ b/tests/e2e/claude_code/passthrough/test_azure.py @@ -44,7 +44,7 @@ from claude_code._passthrough import foundry_extra_env, run_passthrough_cell AZURE_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] diff --git a/tests/e2e/claude_code/passthrough/test_bedrock_invoke.py b/tests/e2e/claude_code/passthrough/test_bedrock_invoke.py index 6e84dea6779..f1f28ab5b4c 100644 --- a/tests/e2e/claude_code/passthrough/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/passthrough/test_bedrock_invoke.py @@ -27,7 +27,7 @@ from claude_code._passthrough import bedrock_extra_env, run_passthrough_cell BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] diff --git a/tests/e2e/claude_code/passthrough/test_vertex_ai.py b/tests/e2e/claude_code/passthrough/test_vertex_ai.py index 5e3c6bce419..790f8b60c8f 100644 --- a/tests/e2e/claude_code/passthrough/test_vertex_ai.py +++ b/tests/e2e/claude_code/passthrough/test_vertex_ai.py @@ -30,7 +30,7 @@ from claude_code._passthrough import run_passthrough_cell, vertex_extra_env VERTEX_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] diff --git a/tests/e2e/claude_code/pdf_input/test_anthropic.py b/tests/e2e/claude_code/pdf_input/test_anthropic.py index 36fb69a1db6..21c8028ef1c 100644 --- a/tests/e2e/claude_code/pdf_input/test_anthropic.py +++ b/tests/e2e/claude_code/pdf_input/test_anthropic.py @@ -22,22 +22,19 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] @@ -108,24 +105,11 @@ def _build_minimal_pdf(marker: str) -> bytes: return bytes(out) +@pytest.mark.covers("llm.messages.anthropic.pdf_input.nonstream.works") def test_pdf_input_anthropic(compat_result, tmp_path): """Drive the `claude` CLI against the LiteLLM proxy with a PDF attached via the Read tool and assert the reply references it.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) pdf_path = tmp_path / "marker.pdf" pdf_path.write_bytes(_build_minimal_pdf(PDF_MARKER)) diff --git a/tests/e2e/claude_code/pdf_input/test_azure.py b/tests/e2e/claude_code/pdf_input/test_azure.py index 810c857e407..34ae3732b99 100644 --- a/tests/e2e/claude_code/pdf_input/test_azure.py +++ b/tests/e2e/claude_code/pdf_input/test_azure.py @@ -15,22 +15,19 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] @@ -85,22 +82,9 @@ def _build_minimal_pdf(marker: str) -> bytes: return bytes(out) +@pytest.mark.covers("llm.messages.azure_foundry.pdf_input.nonstream.works") def test_pdf_input_azure(compat_result, tmp_path): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) pdf_path = tmp_path / "marker.pdf" pdf_path.write_bytes(_build_minimal_pdf(PDF_MARKER)) diff --git a/tests/e2e/claude_code/pdf_input/test_bedrock_converse.py b/tests/e2e/claude_code/pdf_input/test_bedrock_converse.py index 191a27c6d46..76aa84f0f47 100644 --- a/tests/e2e/claude_code/pdf_input/test_bedrock_converse.py +++ b/tests/e2e/claude_code/pdf_input/test_bedrock_converse.py @@ -21,22 +21,19 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] @@ -91,22 +88,9 @@ def _build_minimal_pdf(marker: str) -> bytes: return bytes(out) +@pytest.mark.covers("llm.messages.bedrock_converse.pdf_input.nonstream.works") def test_pdf_input_bedrock_converse(compat_result, tmp_path): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) pdf_path = tmp_path / "marker.pdf" pdf_path.write_bytes(_build_minimal_pdf(PDF_MARKER)) diff --git a/tests/e2e/claude_code/pdf_input/test_bedrock_invoke.py b/tests/e2e/claude_code/pdf_input/test_bedrock_invoke.py index 163cabb45a0..4450266bb6b 100644 --- a/tests/e2e/claude_code/pdf_input/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/pdf_input/test_bedrock_invoke.py @@ -20,22 +20,19 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] @@ -90,22 +87,9 @@ def _build_minimal_pdf(marker: str) -> bytes: return bytes(out) +@pytest.mark.covers("llm.messages.bedrock_invoke.pdf_input.nonstream.works") def test_pdf_input_bedrock_invoke(compat_result, tmp_path): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) pdf_path = tmp_path / "marker.pdf" pdf_path.write_bytes(_build_minimal_pdf(PDF_MARKER)) diff --git a/tests/e2e/claude_code/pdf_input/test_vertex_ai.py b/tests/e2e/claude_code/pdf_input/test_vertex_ai.py index 0d0573d05b3..b78f58cfda1 100644 --- a/tests/e2e/claude_code/pdf_input/test_vertex_ai.py +++ b/tests/e2e/claude_code/pdf_input/test_vertex_ai.py @@ -15,22 +15,19 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] @@ -85,22 +82,9 @@ def _build_minimal_pdf(marker: str) -> bytes: return bytes(out) +@pytest.mark.covers("llm.messages.vertex.pdf_input.nonstream.works") def test_pdf_input_vertex_ai(compat_result, tmp_path): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) pdf_path = tmp_path / "marker.pdf" pdf_path.write_bytes(_build_minimal_pdf(PDF_MARKER)) diff --git a/tests/e2e/claude_code/prompt_caching_1h/test_anthropic.py b/tests/e2e/claude_code/prompt_caching_1h/test_anthropic.py index d81887231d8..637be1c551d 100644 --- a/tests/e2e/claude_code/prompt_caching_1h/test_anthropic.py +++ b/tests/e2e/claude_code/prompt_caching_1h/test_anthropic.py @@ -24,23 +24,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Optional import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] @@ -65,25 +63,12 @@ def _cache_tokens(usage: Optional[Mapping[str, Any]]) -> int: return 0 +@pytest.mark.covers("llm.messages.anthropic.prompt_cache_1h.nonstream.works") def test_prompt_caching_1h_anthropic(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with the 1h TTL opt-in env var set, and assert the upstream usage block surfaces a non-zero cache token count.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=ANTHROPIC_MODELS, diff --git a/tests/e2e/claude_code/prompt_caching_1h/test_azure.py b/tests/e2e/claude_code/prompt_caching_1h/test_azure.py index 416757f8691..f34557b3c5f 100644 --- a/tests/e2e/claude_code/prompt_caching_1h/test_azure.py +++ b/tests/e2e/claude_code/prompt_caching_1h/test_azure.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Optional import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] @@ -49,22 +47,9 @@ def _cache_tokens(usage: Optional[Mapping[str, Any]]) -> int: return 0 +@pytest.mark.covers("llm.messages.azure_foundry.prompt_cache_1h.nonstream.works") def test_prompt_caching_1h_azure(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=AZURE_MODELS, diff --git a/tests/e2e/claude_code/prompt_caching_1h/test_bedrock_converse.py b/tests/e2e/claude_code/prompt_caching_1h/test_bedrock_converse.py index 5bc632c6f1b..bf62a49444c 100644 --- a/tests/e2e/claude_code/prompt_caching_1h/test_bedrock_converse.py +++ b/tests/e2e/claude_code/prompt_caching_1h/test_bedrock_converse.py @@ -20,23 +20,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Optional import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] @@ -57,22 +55,9 @@ def _cache_tokens(usage: Optional[Mapping[str, Any]]) -> int: return 0 +@pytest.mark.covers("llm.messages.bedrock_converse.prompt_cache_1h.nonstream.works") def test_prompt_caching_1h_bedrock_converse(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_CONVERSE_MODELS, diff --git a/tests/e2e/claude_code/prompt_caching_1h/test_bedrock_invoke.py b/tests/e2e/claude_code/prompt_caching_1h/test_bedrock_invoke.py index 4501834956b..dc3468702d4 100644 --- a/tests/e2e/claude_code/prompt_caching_1h/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/prompt_caching_1h/test_bedrock_invoke.py @@ -22,23 +22,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Optional import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] @@ -61,22 +59,9 @@ def _cache_tokens(usage: Optional[Mapping[str, Any]]) -> int: return 0 +@pytest.mark.covers("llm.messages.bedrock_invoke.prompt_cache_1h.nonstream.works") def test_prompt_caching_1h_bedrock_invoke(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_INVOKE_MODELS, diff --git a/tests/e2e/claude_code/prompt_caching_1h/test_vertex_ai.py b/tests/e2e/claude_code/prompt_caching_1h/test_vertex_ai.py index 09ded634b45..66cf961fcfc 100644 --- a/tests/e2e/claude_code/prompt_caching_1h/test_vertex_ai.py +++ b/tests/e2e/claude_code/prompt_caching_1h/test_vertex_ai.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Optional import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] @@ -49,22 +47,9 @@ def _cache_tokens(usage: Optional[Mapping[str, Any]]) -> int: return 0 +@pytest.mark.covers("llm.messages.vertex.prompt_cache_1h.nonstream.works") def test_prompt_caching_1h_vertex_ai(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=VERTEX_AI_MODELS, diff --git a/tests/e2e/claude_code/prompt_caching_5m/test_anthropic.py b/tests/e2e/claude_code/prompt_caching_5m/test_anthropic.py index 4b20a65f31b..ef551beb45c 100644 --- a/tests/e2e/claude_code/prompt_caching_5m/test_anthropic.py +++ b/tests/e2e/claude_code/prompt_caching_5m/test_anthropic.py @@ -22,23 +22,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Optional import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] @@ -56,24 +54,11 @@ def _cache_tokens(usage: Optional[Mapping[str, Any]]) -> int: return 0 +@pytest.mark.covers("llm.messages.anthropic.prompt_cache_5m.nonstream.works") def test_prompt_caching_5m_anthropic(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the upstream usage block surfaces a non-zero cache token count.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=ANTHROPIC_MODELS, diff --git a/tests/e2e/claude_code/prompt_caching_5m/test_azure.py b/tests/e2e/claude_code/prompt_caching_5m/test_azure.py index 22bd5aa7048..9d4137e0726 100644 --- a/tests/e2e/claude_code/prompt_caching_5m/test_azure.py +++ b/tests/e2e/claude_code/prompt_caching_5m/test_azure.py @@ -22,23 +22,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Optional import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] @@ -54,24 +52,11 @@ def _cache_tokens(usage: Optional[Mapping[str, Any]]) -> int: return 0 +@pytest.mark.covers("llm.messages.azure_foundry.prompt_cache_5m.nonstream.works") def test_prompt_caching_5m_azure(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the upstream usage block surfaces a non-zero cache token count.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=AZURE_MODELS, diff --git a/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_converse.py b/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_converse.py index 681a6ecce10..c9b34c010b0 100644 --- a/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_converse.py +++ b/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_converse.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Optional import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] @@ -47,24 +45,11 @@ def _cache_tokens(usage: Optional[Mapping[str, Any]]) -> int: return 0 +@pytest.mark.covers("llm.messages.bedrock_converse.prompt_cache_5m.nonstream.works") def test_prompt_caching_5m_bedrock_converse(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the upstream usage block surfaces a non-zero cache token count.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_CONVERSE_MODELS, diff --git a/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_invoke.py b/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_invoke.py index f1a3109b3a1..b95c509ba3c 100644 --- a/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_invoke.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Optional import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] @@ -47,24 +45,11 @@ def _cache_tokens(usage: Optional[Mapping[str, Any]]) -> int: return 0 +@pytest.mark.covers("llm.messages.bedrock_invoke.prompt_cache_5m.nonstream.works") def test_prompt_caching_5m_bedrock_invoke(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the upstream usage block surfaces a non-zero cache token count.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_INVOKE_MODELS, diff --git a/tests/e2e/claude_code/prompt_caching_5m/test_vertex_ai.py b/tests/e2e/claude_code/prompt_caching_5m/test_vertex_ai.py index cc5d337dfbe..f79377b7372 100644 --- a/tests/e2e/claude_code/prompt_caching_5m/test_vertex_ai.py +++ b/tests/e2e/claude_code/prompt_caching_5m/test_vertex_ai.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Optional import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] @@ -47,24 +45,11 @@ def _cache_tokens(usage: Optional[Mapping[str, Any]]) -> int: return 0 +@pytest.mark.covers("llm.messages.vertex.prompt_cache_5m.nonstream.works") def test_prompt_caching_5m_vertex_ai(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the upstream usage block surfaces a non-zero cache token count.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=VERTEX_AI_MODELS, diff --git a/tests/e2e/claude_code/run_compat.sh b/tests/e2e/claude_code/run_compat.sh index 4d8d0b6d7b2..e383792f45e 100755 --- a/tests/e2e/claude_code/run_compat.sh +++ b/tests/e2e/claude_code/run_compat.sh @@ -11,9 +11,9 @@ # 4. If a provider has `rate_limited > 0`, halve its rate; else, double it. # 5. Repeat until the highest no-429 rate is found. # -# Required env (proxy connection): -# LITELLM_PROXY_BASE_URL e.g. http://localhost:4000 -# LITELLM_PROXY_API_KEY e.g. sk-1234 +# Required env (proxy connection), same names as the rest of tests/e2e: +# LITELLM_PROXY_URL e.g. http://localhost:4000 +# LITELLM_MASTER_KEY e.g. sk-1234 # # Optional env (rate limits, all default to 5 req/s; 0 disables a column): # LITELLM_COMPAT_RATE_ANTHROPIC @@ -32,8 +32,8 @@ set -euo pipefail -if [[ -z "${LITELLM_PROXY_BASE_URL:-}" || -z "${LITELLM_PROXY_API_KEY:-}" ]]; then - echo "error: LITELLM_PROXY_BASE_URL and LITELLM_PROXY_API_KEY must be set" >&2 +if [[ -z "${LITELLM_PROXY_URL:-}" || -z "${LITELLM_MASTER_KEY:-}" ]]; then + echo "error: LITELLM_PROXY_URL and LITELLM_MASTER_KEY must be set" >&2 exit 64 fi diff --git a/tests/e2e/claude_code/structured_outputs/test_anthropic.py b/tests/e2e/claude_code/structured_outputs/test_anthropic.py index 610d8433b72..3dc4c7ab8f2 100644 --- a/tests/e2e/claude_code/structured_outputs/test_anthropic.py +++ b/tests/e2e/claude_code/structured_outputs/test_anthropic.py @@ -48,24 +48,22 @@ tier so the matrix's "all three must pass" rule applies. from __future__ import annotations import json -import os import re from typing import Any, Mapping, Optional, Sequence, Tuple import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] @@ -152,26 +150,12 @@ def _validate_against_schema( return None +@pytest.mark.covers("llm.messages.anthropic.structured_output.nonstream.works") def test_structured_outputs_anthropic(compat_result): """Drive `claude --json-schema ...` against the LiteLLM proxy and assert the trailing `result` event contains a schema-conforming `structured_output`.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=ANTHROPIC_MODELS, diff --git a/tests/e2e/claude_code/structured_outputs/test_azure.py b/tests/e2e/claude_code/structured_outputs/test_azure.py index 290f9156910..7a776ed55ad 100644 --- a/tests/e2e/claude_code/structured_outputs/test_azure.py +++ b/tests/e2e/claude_code/structured_outputs/test_azure.py @@ -48,24 +48,22 @@ tier so the matrix's "all three must pass" rule applies. from __future__ import annotations import json -import os import re from typing import Any, Mapping, Optional, Sequence, Tuple import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] @@ -152,26 +150,12 @@ def _validate_against_schema( return None +@pytest.mark.covers("llm.messages.azure_foundry.structured_output.nonstream.works") def test_structured_outputs_azure(compat_result): """Drive `claude --json-schema ...` against the LiteLLM proxy and assert the trailing `result` event contains a schema-conforming `structured_output`.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=AZURE_MODELS, diff --git a/tests/e2e/claude_code/structured_outputs/test_bedrock_converse.py b/tests/e2e/claude_code/structured_outputs/test_bedrock_converse.py index 5179014773c..345d7c327cf 100644 --- a/tests/e2e/claude_code/structured_outputs/test_bedrock_converse.py +++ b/tests/e2e/claude_code/structured_outputs/test_bedrock_converse.py @@ -48,24 +48,22 @@ tier so the matrix's "all three must pass" rule applies. from __future__ import annotations import json -import os import re from typing import Any, Mapping, Optional, Sequence, Tuple import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] @@ -152,26 +150,12 @@ def _validate_against_schema( return None +@pytest.mark.covers("llm.messages.bedrock_converse.structured_output.nonstream.works") def test_structured_outputs_bedrock_converse(compat_result): """Drive `claude --json-schema ...` against the LiteLLM proxy and assert the trailing `result` event contains a schema-conforming `structured_output`.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_CONVERSE_MODELS, diff --git a/tests/e2e/claude_code/structured_outputs/test_bedrock_invoke.py b/tests/e2e/claude_code/structured_outputs/test_bedrock_invoke.py index 313a714be34..0cf48c72d4f 100644 --- a/tests/e2e/claude_code/structured_outputs/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/structured_outputs/test_bedrock_invoke.py @@ -48,24 +48,22 @@ tier so the matrix's "all three must pass" rule applies. from __future__ import annotations import json -import os import re from typing import Any, Mapping, Optional, Sequence, Tuple import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] @@ -152,26 +150,12 @@ def _validate_against_schema( return None +@pytest.mark.covers("llm.messages.bedrock_invoke.structured_output.nonstream.works") def test_structured_outputs_bedrock_invoke(compat_result): """Drive `claude --json-schema ...` against the LiteLLM proxy and assert the trailing `result` event contains a schema-conforming `structured_output`.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_INVOKE_MODELS, diff --git a/tests/e2e/claude_code/structured_outputs/test_vertex_ai.py b/tests/e2e/claude_code/structured_outputs/test_vertex_ai.py index ec04c724193..24f5a0c35d4 100644 --- a/tests/e2e/claude_code/structured_outputs/test_vertex_ai.py +++ b/tests/e2e/claude_code/structured_outputs/test_vertex_ai.py @@ -48,24 +48,22 @@ tier so the matrix's "all three must pass" rule applies. from __future__ import annotations import json -import os import re from typing import Any, Mapping, Optional, Sequence, Tuple import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] @@ -152,26 +150,12 @@ def _validate_against_schema( return None +@pytest.mark.covers("llm.messages.vertex.structured_output.nonstream.works") def test_structured_outputs_vertex_ai(compat_result): """Drive `claude --json-schema ...` against the LiteLLM proxy and assert the trailing `result` event contains a schema-conforming `structured_output`.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=VERTEX_AI_MODELS, diff --git a/tests/e2e/claude_code/test_config.yaml b/tests/e2e/claude_code/test_config.yaml index e9253da2b3c..ba57ca4ccb7 100644 --- a/tests/e2e/claude_code/test_config.yaml +++ b/tests/e2e/claude_code/test_config.yaml @@ -21,42 +21,54 @@ model_list: litellm_params: model: anthropic/claude-haiku-4-5 api_key: os.environ/ANTHROPIC_API_KEY - - model_name: claude-sonnet-4-6 + - model_name: claude-sonnet-4-5 litellm_params: - model: anthropic/claude-sonnet-4-6 + model: anthropic/claude-sonnet-4-5 api_key: os.environ/ANTHROPIC_API_KEY + extra_headers: + anthropic-beta: "context-1m-2025-08-07" - model_name: claude-opus-4-7 litellm_params: model: anthropic/claude-opus-4-7 api_key: os.environ/ANTHROPIC_API_KEY + extra_headers: + anthropic-beta: "context-1m-2025-08-07" # ---- Bedrock (InvokeModel) ---- - model_name: claude-haiku-4-5-bedrock-invoke litellm_params: model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0 aws_region_name: us-east-1 - - model_name: claude-sonnet-4-6-bedrock-invoke + - model_name: claude-sonnet-4-5-bedrock-invoke litellm_params: - model: bedrock/us.anthropic.claude-sonnet-4-6 + model: bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0 aws_region_name: us-east-1 + extra_headers: + anthropic-beta: "context-1m-2025-08-07" - model_name: claude-opus-4-7-bedrock-invoke litellm_params: model: bedrock/us.anthropic.claude-opus-4-7 aws_region_name: us-east-1 + extra_headers: + anthropic-beta: "context-1m-2025-08-07" # ---- Bedrock (Converse) ---- - model_name: claude-haiku-4-5-bedrock-converse litellm_params: model: bedrock/converse/us.anthropic.claude-haiku-4-5-20251001-v1:0 aws_region_name: us-east-1 - - model_name: claude-sonnet-4-6-bedrock-converse + - model_name: claude-sonnet-4-5-bedrock-converse litellm_params: - model: bedrock/converse/us.anthropic.claude-sonnet-4-6 + model: bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0 aws_region_name: us-east-1 + extra_headers: + anthropic-beta: "context-1m-2025-08-07" - model_name: claude-opus-4-7-bedrock-converse litellm_params: model: bedrock/converse/us.anthropic.claude-opus-4-7 aws_region_name: us-east-1 + extra_headers: + anthropic-beta: "context-1m-2025-08-07" # ---- Vertex AI ---- # `use_in_pass_through: true` registers each deployment's @@ -70,37 +82,45 @@ model_list: litellm_params: model: vertex_ai/claude-haiku-4-5 vertex_project: os.environ/VERTEXAI_PROJECT - vertex_location: os.environ/VERTEXAI_LOCATION + vertex_location: global use_in_pass_through: true - - model_name: claude-sonnet-4-6-vertex + - model_name: claude-sonnet-4-5-vertex litellm_params: - model: vertex_ai/claude-sonnet-4-6 + model: vertex_ai/claude-sonnet-4-5 vertex_project: os.environ/VERTEXAI_PROJECT - vertex_location: os.environ/VERTEXAI_LOCATION + vertex_location: global use_in_pass_through: true + extra_headers: + anthropic-beta: "context-1m-2025-08-07" - model_name: claude-opus-4-7-vertex litellm_params: model: vertex_ai/claude-opus-4-7 vertex_project: os.environ/VERTEXAI_PROJECT - vertex_location: os.environ/VERTEXAI_LOCATION + vertex_location: global use_in_pass_through: true + extra_headers: + anthropic-beta: "context-1m-2025-08-07" # ---- Microsoft Foundry (Anthropic deployments on Azure) ---- - model_name: claude-haiku-4-5-azure litellm_params: model: azure_ai/claude-haiku-4-5 - api_base: os.environ/AZURE_FOUNDRY_API_BASE - api_key: os.environ/AZURE_FOUNDRY_API_KEY - - model_name: claude-sonnet-4-6-azure + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY + - model_name: claude-sonnet-4-5-azure litellm_params: - model: azure_ai/claude-sonnet-4-6 - api_base: os.environ/AZURE_FOUNDRY_API_BASE - api_key: os.environ/AZURE_FOUNDRY_API_KEY + model: azure_ai/claude-sonnet-4-5 + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY + extra_headers: + anthropic-beta: "context-1m-2025-08-07" - model_name: claude-opus-4-7-azure litellm_params: model: azure_ai/claude-opus-4-7 - api_base: os.environ/AZURE_FOUNDRY_API_BASE - api_key: os.environ/AZURE_FOUNDRY_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY + extra_headers: + anthropic-beta: "context-1m-2025-08-07" general_settings: # Claude Code sends provider-specific headers (e.g. anthropic-beta) we diff --git a/tests/e2e/claude_code/thinking/test_anthropic.py b/tests/e2e/claude_code/thinking/test_anthropic.py index 1090d1b384e..ebb2445fb6d 100644 --- a/tests/e2e/claude_code/thinking/test_anthropic.py +++ b/tests/e2e/claude_code/thinking/test_anthropic.py @@ -20,23 +20,21 @@ still sees three rows for this (feature, provider). from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] @@ -76,24 +74,11 @@ def _has_thinking_block(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.anthropic.thinking.nonstream.works") def test_thinking_anthropic(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with thinking enabled and assert a `thinking` content block was emitted.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=ANTHROPIC_MODELS, diff --git a/tests/e2e/claude_code/thinking/test_azure.py b/tests/e2e/claude_code/thinking/test_azure.py index 1fd5138d574..ffd5ca92df0 100644 --- a/tests/e2e/claude_code/thinking/test_azure.py +++ b/tests/e2e/claude_code/thinking/test_azure.py @@ -23,23 +23,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] @@ -64,24 +62,11 @@ def _has_thinking_block(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.azure_foundry.thinking.nonstream.works") def test_thinking_azure(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with thinking enabled and assert a `thinking` content block was emitted.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=AZURE_MODELS, diff --git a/tests/e2e/claude_code/thinking/test_bedrock_converse.py b/tests/e2e/claude_code/thinking/test_bedrock_converse.py index 793ce8542da..0b409f18ea7 100644 --- a/tests/e2e/claude_code/thinking/test_bedrock_converse.py +++ b/tests/e2e/claude_code/thinking/test_bedrock_converse.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] @@ -56,24 +54,11 @@ def _has_thinking_block(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.bedrock_converse.thinking.nonstream.works") def test_thinking_bedrock_converse(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with thinking enabled and assert a `thinking` content block was emitted.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_CONVERSE_MODELS, diff --git a/tests/e2e/claude_code/thinking/test_bedrock_invoke.py b/tests/e2e/claude_code/thinking/test_bedrock_invoke.py index e31b60eb004..a2c97eae321 100644 --- a/tests/e2e/claude_code/thinking/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/thinking/test_bedrock_invoke.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] @@ -56,24 +54,11 @@ def _has_thinking_block(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.bedrock_invoke.thinking.nonstream.works") def test_thinking_bedrock_invoke(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with thinking enabled and assert a `thinking` content block was emitted.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_INVOKE_MODELS, diff --git a/tests/e2e/claude_code/thinking/test_vertex_ai.py b/tests/e2e/claude_code/thinking/test_vertex_ai.py index c5c7df1f9b8..f1a1c5b6cee 100644 --- a/tests/e2e/claude_code/thinking/test_vertex_ai.py +++ b/tests/e2e/claude_code/thinking/test_vertex_ai.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] @@ -56,24 +54,11 @@ def _has_thinking_block(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.vertex.thinking.nonstream.works") def test_thinking_vertex_ai(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with thinking enabled and assert a `thinking` content block was emitted.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=VERTEX_AI_MODELS, diff --git a/tests/e2e/claude_code/thinking_with_tool_use/test_anthropic.py b/tests/e2e/claude_code/thinking_with_tool_use/test_anthropic.py index 2c573ea039e..7e39ea26d42 100644 --- a/tests/e2e/claude_code/thinking_with_tool_use/test_anthropic.py +++ b/tests/e2e/claude_code/thinking_with_tool_use/test_anthropic.py @@ -24,23 +24,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] @@ -90,25 +88,12 @@ def _has_block_type( return False +@pytest.mark.covers("llm.messages.anthropic.thinking_with_tool_use.nonstream.works") def test_thinking_with_tool_use_anthropic(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with thinking enabled and tool use, and assert both `thinking` and `tool_use` content blocks landed in the same turn.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=ANTHROPIC_MODELS, diff --git a/tests/e2e/claude_code/thinking_with_tool_use/test_azure.py b/tests/e2e/claude_code/thinking_with_tool_use/test_azure.py index 3d65e82cdec..0371a10f8a6 100644 --- a/tests/e2e/claude_code/thinking_with_tool_use/test_azure.py +++ b/tests/e2e/claude_code/thinking_with_tool_use/test_azure.py @@ -18,23 +18,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] @@ -71,22 +69,9 @@ def _has_block_type( return False +@pytest.mark.covers("llm.messages.azure_foundry.thinking_with_tool_use.nonstream.works") def test_thinking_with_tool_use_azure(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=AZURE_MODELS, diff --git a/tests/e2e/claude_code/thinking_with_tool_use/test_bedrock_converse.py b/tests/e2e/claude_code/thinking_with_tool_use/test_bedrock_converse.py index eb916323546..026d2a3707f 100644 --- a/tests/e2e/claude_code/thinking_with_tool_use/test_bedrock_converse.py +++ b/tests/e2e/claude_code/thinking_with_tool_use/test_bedrock_converse.py @@ -23,23 +23,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] @@ -76,22 +74,9 @@ def _has_block_type( return False +@pytest.mark.covers("llm.messages.bedrock_converse.thinking_with_tool_use.nonstream.works") def test_thinking_with_tool_use_bedrock_converse(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_CONVERSE_MODELS, diff --git a/tests/e2e/claude_code/thinking_with_tool_use/test_bedrock_invoke.py b/tests/e2e/claude_code/thinking_with_tool_use/test_bedrock_invoke.py index d1a61a59772..1dd4cf0a73c 100644 --- a/tests/e2e/claude_code/thinking_with_tool_use/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/thinking_with_tool_use/test_bedrock_invoke.py @@ -25,23 +25,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] @@ -78,22 +76,9 @@ def _has_block_type( return False +@pytest.mark.covers("llm.messages.bedrock_invoke.thinking_with_tool_use.nonstream.works") def test_thinking_with_tool_use_bedrock_invoke(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_INVOKE_MODELS, diff --git a/tests/e2e/claude_code/thinking_with_tool_use/test_vertex_ai.py b/tests/e2e/claude_code/thinking_with_tool_use/test_vertex_ai.py index 285419c67f7..b25228edb55 100644 --- a/tests/e2e/claude_code/thinking_with_tool_use/test_vertex_ai.py +++ b/tests/e2e/claude_code/thinking_with_tool_use/test_vertex_ai.py @@ -23,23 +23,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] @@ -76,22 +74,9 @@ def _has_block_type( return False +@pytest.mark.covers("llm.messages.vertex.thinking_with_tool_use.nonstream.works") def test_thinking_with_tool_use_vertex_ai(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=VERTEX_AI_MODELS, diff --git a/tests/e2e/claude_code/tool_search/test_anthropic.py b/tests/e2e/claude_code/tool_search/test_anthropic.py index 3495c882e06..7b8ea07aa07 100644 --- a/tests/e2e/claude_code/tool_search/test_anthropic.py +++ b/tests/e2e/claude_code/tool_search/test_anthropic.py @@ -43,45 +43,28 @@ per-cell aggregator. from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.http_probe import ( assert_tool_search_shape, probe_tool_search, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] +@pytest.mark.covers("llm.messages.anthropic.tool_search.nonstream.works") def test_tool_search_anthropic(compat_result): """Probe `/v1/messages` with a `tool_search_tool_regex_20251119` tool and assert the proxy + upstream accept it for every Anthropic tier.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) failures = [] for model in ANTHROPIC_MODELS: diff --git a/tests/e2e/claude_code/tool_search/test_azure.py b/tests/e2e/claude_code/tool_search/test_azure.py index 1d9cb5673c5..4eee13e4ecc 100644 --- a/tests/e2e/claude_code/tool_search/test_azure.py +++ b/tests/e2e/claude_code/tool_search/test_azure.py @@ -43,45 +43,28 @@ per-cell aggregator. from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.http_probe import ( assert_tool_search_shape, probe_tool_search, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] +@pytest.mark.covers("llm.messages.azure_foundry.tool_search.nonstream.works") def test_tool_search_azure(compat_result): """Probe `/v1/messages` with a `tool_search_tool_regex_20251119` tool and assert the proxy + upstream accept it for every Azure (Microsoft Foundry) tier.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) failures = [] for model in AZURE_MODELS: diff --git a/tests/e2e/claude_code/tool_search/test_bedrock_converse.py b/tests/e2e/claude_code/tool_search/test_bedrock_converse.py index 5ca0792529a..7951f8ecdb4 100644 --- a/tests/e2e/claude_code/tool_search/test_bedrock_converse.py +++ b/tests/e2e/claude_code/tool_search/test_bedrock_converse.py @@ -43,45 +43,28 @@ per-cell aggregator. from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.http_probe import ( assert_tool_search_shape, probe_tool_search, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] +@pytest.mark.covers("llm.messages.bedrock_converse.tool_search.nonstream.works") def test_tool_search_bedrock_converse(compat_result): """Probe `/v1/messages` with a `tool_search_tool_regex_20251119` tool and assert the proxy + upstream accept it for every Bedrock (Converse) tier.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) failures = [] for model in BEDROCK_CONVERSE_MODELS: diff --git a/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py b/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py index 21bb33e34bd..654c2aa18d1 100644 --- a/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py @@ -43,45 +43,28 @@ per-cell aggregator. from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.http_probe import ( assert_tool_search_shape, probe_tool_search, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] +@pytest.mark.covers("llm.messages.bedrock_invoke.tool_search.nonstream.works") def test_tool_search_bedrock_invoke(compat_result): """Probe `/v1/messages` with a `tool_search_tool_regex_20251119` tool and assert the proxy + upstream accept it for every Bedrock (Invoke) tier.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) failures = [] for model in BEDROCK_INVOKE_MODELS: diff --git a/tests/e2e/claude_code/tool_search/test_vertex_ai.py b/tests/e2e/claude_code/tool_search/test_vertex_ai.py index f91400b1817..f6ff855fa78 100644 --- a/tests/e2e/claude_code/tool_search/test_vertex_ai.py +++ b/tests/e2e/claude_code/tool_search/test_vertex_ai.py @@ -43,45 +43,28 @@ per-cell aggregator. from __future__ import annotations -import os - import pytest +from claude_code._env import require_proxy from claude_code.http_probe import ( assert_tool_search_shape, probe_tool_search, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] +@pytest.mark.covers("llm.messages.vertex.tool_search.nonstream.works") def test_tool_search_vertex_ai(compat_result): """Probe `/v1/messages` with a `tool_search_tool_regex_20251119` tool and assert the proxy + upstream accept it for every Vertex AI tier.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", - pytrace=False, - ) + base_url, api_key = require_proxy(compat_result) failures = [] for model in VERTEX_AI_MODELS: diff --git a/tests/e2e/claude_code/tool_use/test_anthropic.py b/tests/e2e/claude_code/tool_use/test_anthropic.py index 7d2aa4be683..9ff4c58907f 100644 --- a/tests/e2e/claude_code/tool_use/test_anthropic.py +++ b/tests/e2e/claude_code/tool_use/test_anthropic.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] @@ -73,24 +71,11 @@ def _has_tool_use_event(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.anthropic.tool_use.nonstream.works") def test_tool_use_anthropic(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a tool call was emitted on the wire.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=ANTHROPIC_MODELS, diff --git a/tests/e2e/claude_code/tool_use/test_azure.py b/tests/e2e/claude_code/tool_use/test_azure.py index 484f50a5508..9e7398267c4 100644 --- a/tests/e2e/claude_code/tool_use/test_azure.py +++ b/tests/e2e/claude_code/tool_use/test_azure.py @@ -19,23 +19,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] @@ -67,24 +65,11 @@ def _has_tool_use_event(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.azure_foundry.tool_use.nonstream.works") def test_tool_use_azure(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a tool call was emitted on the wire.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=AZURE_MODELS, diff --git a/tests/e2e/claude_code/tool_use/test_bedrock_converse.py b/tests/e2e/claude_code/tool_use/test_bedrock_converse.py index 7d1b58fce90..33d4d3820d2 100644 --- a/tests/e2e/claude_code/tool_use/test_bedrock_converse.py +++ b/tests/e2e/claude_code/tool_use/test_bedrock_converse.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] @@ -63,24 +61,11 @@ def _has_tool_use_event(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.bedrock_converse.tool_use.nonstream.works") def test_tool_use_bedrock_converse(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a tool call was emitted on the wire.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_CONVERSE_MODELS, diff --git a/tests/e2e/claude_code/tool_use/test_bedrock_invoke.py b/tests/e2e/claude_code/tool_use/test_bedrock_invoke.py index 7d2b72b951d..47ae3aef1da 100644 --- a/tests/e2e/claude_code/tool_use/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/tool_use/test_bedrock_invoke.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] @@ -63,24 +61,11 @@ def _has_tool_use_event(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.bedrock_invoke.tool_use.nonstream.works") def test_tool_use_bedrock_invoke(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a tool call was emitted on the wire.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_INVOKE_MODELS, diff --git a/tests/e2e/claude_code/tool_use/test_vertex_ai.py b/tests/e2e/claude_code/tool_use/test_vertex_ai.py index 0a8ecc9f7a7..79a3016345c 100644 --- a/tests/e2e/claude_code/tool_use/test_vertex_ai.py +++ b/tests/e2e/claude_code/tool_use/test_vertex_ai.py @@ -15,23 +15,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] @@ -63,24 +61,11 @@ def _has_tool_use_event(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.vertex.tool_use.nonstream.works") def test_tool_use_vertex_ai(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a tool call was emitted on the wire.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=VERTEX_AI_MODELS, diff --git a/tests/e2e/claude_code/tool_use_streaming/test_anthropic.py b/tests/e2e/claude_code/tool_use_streaming/test_anthropic.py index 9aa94c89241..152652dcf3c 100644 --- a/tests/e2e/claude_code/tool_use_streaming/test_anthropic.py +++ b/tests/e2e/claude_code/tool_use_streaming/test_anthropic.py @@ -25,23 +25,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] @@ -99,24 +97,11 @@ def _count_input_json_deltas(events: Sequence[Mapping[str, Any]]) -> int: ) +@pytest.mark.covers("llm.messages.anthropic.tool_use.stream.works") def test_tool_use_streaming_anthropic(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the proxy preserves fine-grained tool streaming end-to-end.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=ANTHROPIC_MODELS, diff --git a/tests/e2e/claude_code/tool_use_streaming/test_azure.py b/tests/e2e/claude_code/tool_use_streaming/test_azure.py index c73062b72cd..8a1cc1852dd 100644 --- a/tests/e2e/claude_code/tool_use_streaming/test_azure.py +++ b/tests/e2e/claude_code/tool_use_streaming/test_azure.py @@ -17,23 +17,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] @@ -84,22 +82,9 @@ def _count_input_json_deltas(events: Sequence[Mapping[str, Any]]) -> int: ) +@pytest.mark.covers("llm.messages.azure_foundry.tool_use.stream.works") def test_tool_use_streaming_azure(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=AZURE_MODELS, diff --git a/tests/e2e/claude_code/tool_use_streaming/test_bedrock_converse.py b/tests/e2e/claude_code/tool_use_streaming/test_bedrock_converse.py index 3642551c7c3..3b04ed5962f 100644 --- a/tests/e2e/claude_code/tool_use_streaming/test_bedrock_converse.py +++ b/tests/e2e/claude_code/tool_use_streaming/test_bedrock_converse.py @@ -23,23 +23,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] @@ -90,22 +88,9 @@ def _count_input_json_deltas(events: Sequence[Mapping[str, Any]]) -> int: ) +@pytest.mark.covers("llm.messages.bedrock_converse.tool_use.stream.works") def test_tool_use_streaming_bedrock_converse(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_CONVERSE_MODELS, diff --git a/tests/e2e/claude_code/tool_use_streaming/test_bedrock_invoke.py b/tests/e2e/claude_code/tool_use_streaming/test_bedrock_invoke.py index af4689b2847..c7b61129782 100644 --- a/tests/e2e/claude_code/tool_use_streaming/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/tool_use_streaming/test_bedrock_invoke.py @@ -21,23 +21,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] @@ -88,22 +86,9 @@ def _count_input_json_deltas(events: Sequence[Mapping[str, Any]]) -> int: ) +@pytest.mark.covers("llm.messages.bedrock_invoke.tool_use.stream.works") def test_tool_use_streaming_bedrock_invoke(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_INVOKE_MODELS, diff --git a/tests/e2e/claude_code/tool_use_streaming/test_vertex_ai.py b/tests/e2e/claude_code/tool_use_streaming/test_vertex_ai.py index 19ef9a4e90e..2912e3aae3d 100644 --- a/tests/e2e/claude_code/tool_use_streaming/test_vertex_ai.py +++ b/tests/e2e/claude_code/tool_use_streaming/test_vertex_ai.py @@ -20,23 +20,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] @@ -87,22 +85,9 @@ def _count_input_json_deltas(events: Sequence[Mapping[str, Any]]) -> int: ) +@pytest.mark.covers("llm.messages.vertex.tool_use.stream.works") def test_tool_use_streaming_vertex_ai(compat_result): - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=VERTEX_AI_MODELS, diff --git a/tests/e2e/claude_code/vision/test_anthropic.py b/tests/e2e/claude_code/vision/test_anthropic.py index 650940248ea..f681b2be5ae 100644 --- a/tests/e2e/claude_code/vision/test_anthropic.py +++ b/tests/e2e/claude_code/vision/test_anthropic.py @@ -25,22 +25,19 @@ the proxy must preserve. from __future__ import annotations import json -import os - import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] @@ -85,24 +82,11 @@ def _build_stdin_input() -> str: return json.dumps(user_event) + "\n" +@pytest.mark.covers("llm.messages.anthropic.vision.nonstream.works") def test_vision_anthropic(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with an image attached via stream-json input and assert a non-empty reply.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=ANTHROPIC_MODELS, diff --git a/tests/e2e/claude_code/vision/test_azure.py b/tests/e2e/claude_code/vision/test_azure.py index 3b03c0f2b35..f0eaaad84a2 100644 --- a/tests/e2e/claude_code/vision/test_azure.py +++ b/tests/e2e/claude_code/vision/test_azure.py @@ -25,22 +25,19 @@ the proxy must preserve. from __future__ import annotations import json -import os - import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] @@ -85,24 +82,11 @@ def _build_stdin_input() -> str: return json.dumps(user_event) + "\n" +@pytest.mark.covers("llm.messages.azure_foundry.vision.nonstream.works") def test_vision_azure(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with an image attached via stream-json input and assert a non-empty reply.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=AZURE_MODELS, diff --git a/tests/e2e/claude_code/vision/test_bedrock_converse.py b/tests/e2e/claude_code/vision/test_bedrock_converse.py index 4201f9e64fc..2a5aba5a393 100644 --- a/tests/e2e/claude_code/vision/test_bedrock_converse.py +++ b/tests/e2e/claude_code/vision/test_bedrock_converse.py @@ -25,22 +25,19 @@ the proxy must preserve. from __future__ import annotations import json -import os - import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] @@ -85,24 +82,11 @@ def _build_stdin_input() -> str: return json.dumps(user_event) + "\n" +@pytest.mark.covers("llm.messages.bedrock_converse.vision.nonstream.works") def test_vision_bedrock_converse(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with an image attached via stream-json input and assert a non-empty reply.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_CONVERSE_MODELS, diff --git a/tests/e2e/claude_code/vision/test_bedrock_invoke.py b/tests/e2e/claude_code/vision/test_bedrock_invoke.py index d2e641f1462..5c995cd479e 100644 --- a/tests/e2e/claude_code/vision/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/vision/test_bedrock_invoke.py @@ -25,22 +25,19 @@ the proxy must preserve. from __future__ import annotations import json -import os - import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] @@ -85,24 +82,11 @@ def _build_stdin_input() -> str: return json.dumps(user_event) + "\n" +@pytest.mark.covers("llm.messages.bedrock_invoke.vision.nonstream.works") def test_vision_bedrock_invoke(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with an image attached via stream-json input and assert a non-empty reply.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_INVOKE_MODELS, diff --git a/tests/e2e/claude_code/vision/test_vertex_ai.py b/tests/e2e/claude_code/vision/test_vertex_ai.py index a39ef1a34b7..8d385e295d0 100644 --- a/tests/e2e/claude_code/vision/test_vertex_ai.py +++ b/tests/e2e/claude_code/vision/test_vertex_ai.py @@ -25,22 +25,19 @@ the proxy must preserve. from __future__ import annotations import json -import os - import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] @@ -85,24 +82,11 @@ def _build_stdin_input() -> str: return json.dumps(user_event) + "\n" +@pytest.mark.covers("llm.messages.vertex.vision.nonstream.works") def test_vision_vertex_ai(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with an image attached via stream-json input and assert a non-empty reply.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=VERTEX_AI_MODELS, diff --git a/tests/e2e/claude_code/web_search/test_anthropic.py b/tests/e2e/claude_code/web_search/test_anthropic.py index b8fa806f923..a20a2133dc9 100644 --- a/tests/e2e/claude_code/web_search/test_anthropic.py +++ b/tests/e2e/claude_code/web_search/test_anthropic.py @@ -27,23 +27,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" ANTHROPIC_MODELS = [ "claude-haiku-4-5", - "claude-sonnet-4-6", + "claude-sonnet-4-5", "claude-opus-4-7", ] @@ -86,26 +84,13 @@ def _has_web_search_tool_use(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.anthropic.web_search.nonstream.works") def test_web_search_anthropic(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the upstream emitted a `tool_use` block calling `WebSearch`, proving the proxy preserved both the request-side tool definition and the response-side tool_use block.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=ANTHROPIC_MODELS, diff --git a/tests/e2e/claude_code/web_search/test_azure.py b/tests/e2e/claude_code/web_search/test_azure.py index e70dc848dcf..8f9f638fbee 100644 --- a/tests/e2e/claude_code/web_search/test_azure.py +++ b/tests/e2e/claude_code/web_search/test_azure.py @@ -27,23 +27,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" AZURE_MODELS = [ "claude-haiku-4-5-azure", - "claude-sonnet-4-6-azure", + "claude-sonnet-4-5-azure", "claude-opus-4-7-azure", ] @@ -86,26 +84,13 @@ def _has_web_search_tool_use(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.azure_foundry.web_search.nonstream.works") def test_web_search_azure(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the upstream emitted a `tool_use` block calling `WebSearch`, proving the proxy preserved both the request-side tool definition and the response-side tool_use block.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=AZURE_MODELS, diff --git a/tests/e2e/claude_code/web_search/test_bedrock_converse.py b/tests/e2e/claude_code/web_search/test_bedrock_converse.py index cbeea03df40..32f37b2be79 100644 --- a/tests/e2e/claude_code/web_search/test_bedrock_converse.py +++ b/tests/e2e/claude_code/web_search/test_bedrock_converse.py @@ -27,23 +27,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_CONVERSE_MODELS = [ "claude-haiku-4-5-bedrock-converse", - "claude-sonnet-4-6-bedrock-converse", + "claude-sonnet-4-5-bedrock-converse", "claude-opus-4-7-bedrock-converse", ] @@ -86,26 +84,13 @@ def _has_web_search_tool_use(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.bedrock_converse.web_search.nonstream.works") def test_web_search_bedrock_converse(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the upstream emitted a `tool_use` block calling `WebSearch`, proving the proxy preserved both the request-side tool definition and the response-side tool_use block.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_CONVERSE_MODELS, diff --git a/tests/e2e/claude_code/web_search/test_bedrock_invoke.py b/tests/e2e/claude_code/web_search/test_bedrock_invoke.py index 86068e1e22b..68d1b30e83f 100644 --- a/tests/e2e/claude_code/web_search/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/web_search/test_bedrock_invoke.py @@ -27,23 +27,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" BEDROCK_INVOKE_MODELS = [ "claude-haiku-4-5-bedrock-invoke", - "claude-sonnet-4-6-bedrock-invoke", + "claude-sonnet-4-5-bedrock-invoke", "claude-opus-4-7-bedrock-invoke", ] @@ -86,26 +84,13 @@ def _has_web_search_tool_use(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.bedrock_invoke.web_search.nonstream.works") def test_web_search_bedrock_invoke(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the upstream emitted a `tool_use` block calling `WebSearch`, proving the proxy preserved both the request-side tool definition and the response-side tool_use block.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=BEDROCK_INVOKE_MODELS, diff --git a/tests/e2e/claude_code/web_search/test_vertex_ai.py b/tests/e2e/claude_code/web_search/test_vertex_ai.py index a33515771f3..540a8396c98 100644 --- a/tests/e2e/claude_code/web_search/test_vertex_ai.py +++ b/tests/e2e/claude_code/web_search/test_vertex_ai.py @@ -27,23 +27,21 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -import os from typing import Any, Mapping, Sequence import pytest +from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, run_claude_models_parallel, ) -PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" -PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" VERTEX_AI_MODELS = [ "claude-haiku-4-5-vertex", - "claude-sonnet-4-6-vertex", + "claude-sonnet-4-5-vertex", "claude-opus-4-7-vertex", ] @@ -86,26 +84,13 @@ def _has_web_search_tool_use(events: Sequence[Mapping[str, Any]]) -> bool: return False +@pytest.mark.covers("llm.messages.vertex.web_search.nonstream.works") def test_web_search_vertex_ai(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert the upstream emitted a `tool_use` block calling `WebSearch`, proving the proxy preserved both the request-side tool definition and the response-side tool_use block.""" - base_url = os.environ.get(PROXY_BASE_URL_ENV) - api_key = os.environ.get(PROXY_API_KEY_ENV) - if not base_url or not api_key: - compat_result.set( - { - "status": "fail", - "error": ( - f"missing required env: set {PROXY_BASE_URL_ENV} and " - f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" - ), - } - ) - pytest.fail( - f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", pytrace=False - ) + base_url, api_key = require_proxy(compat_result) outcomes = run_claude_models_parallel( models=VERTEX_AI_MODELS, diff --git a/tests/e2e/coverage_registry/llm_claude_code_compat.yaml b/tests/e2e/coverage_registry/llm_claude_code_compat.yaml new file mode 100644 index 00000000000..6edf890f7ec --- /dev/null +++ b/tests/e2e/coverage_registry/llm_claude_code_compat.yaml @@ -0,0 +1,110 @@ +# Claude Code compatibility matrix: /v1/messages coverage across the five provider surfaces +# claude-code drives (anthropic direct, azure ai foundry, bedrock invoke, bedrock converse, +# vertex ai). Each row is one (feature x provider) cell in the matrix. The seven anthropic-direct +# rows already declared in llm_conversational.yaml are NOT duplicated here; the four other +# provider surfaces plus every feature not already listed for anthropic direct are declared below. +# +# Grammar: llm.messages....works +# route : anthropic | azure_foundry | bedrock_converse | bedrock_invoke | vertex +# capability : basic | tool_use | vision | thinking | prompt_cache_5m | prompt_cache_1h +# | structured_output | pdf_input | long_context_1m +# | thinking_with_tool_use | tool_search | count_tokens | web_search +# streaming : stream | nonstream + +# ---- basic / non-streaming ---- +- {id: llm.messages.azure_foundry.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: basic, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Basic messaging over Azure AI Foundry Anthropic deployments"} +- {id: llm.messages.bedrock_converse.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Basic messaging over Bedrock Converse Anthropic"} +- {id: llm.messages.bedrock_invoke.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: basic, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Basic messaging over Bedrock Invoke Anthropic"} +- {id: llm.messages.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Basic messaging over Vertex AI Anthropic"} + +# ---- basic / streaming ---- +- {id: llm.messages.azure_foundry.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: basic, streaming: stream, assertions: [works], source: "claude_code compat matrix", rationale: "Basic streaming over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "claude_code compat matrix", rationale: "Basic streaming over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: basic, streaming: stream, assertions: [works], source: "claude_code compat matrix", rationale: "Basic streaming over Bedrock Invoke"} +- {id: llm.messages.vertex.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: basic, streaming: stream, assertions: [works], source: "claude_code compat matrix", rationale: "Basic streaming over Vertex AI"} + +# ---- tool_use / non-streaming ---- +- {id: llm.messages.azure_foundry.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: tool_use, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Tool use over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: tool_use, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Tool use over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: tool_use, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Tool use over Bedrock Invoke"} +- {id: llm.messages.vertex.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: tool_use, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Tool use over Vertex AI"} + +# ---- tool_use / streaming ---- +- {id: llm.messages.azure_foundry.tool_use.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: tool_use, streaming: stream, assertions: [works], source: "claude_code compat matrix", rationale: "Streaming tool use over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.tool_use.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: tool_use, streaming: stream, assertions: [works], source: "claude_code compat matrix", rationale: "Streaming tool use over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.tool_use.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: tool_use, streaming: stream, assertions: [works], source: "claude_code compat matrix", rationale: "Streaming tool use over Bedrock Invoke"} +- {id: llm.messages.vertex.tool_use.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: tool_use, streaming: stream, assertions: [works], source: "claude_code compat matrix", rationale: "Streaming tool use over Vertex AI"} + +# ---- vision ---- +- {id: llm.messages.azure_foundry.vision.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: vision, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Vision over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.vision.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: vision, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Vision over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.vision.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: vision, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Vision over Bedrock Invoke"} +- {id: llm.messages.vertex.vision.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: vision, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Vision over Vertex AI"} + +# ---- thinking ---- +- {id: llm.messages.azure_foundry.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: thinking, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Extended thinking over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: thinking, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Extended thinking over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: thinking, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Extended thinking over Bedrock Invoke"} +- {id: llm.messages.vertex.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: thinking, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Extended thinking over Vertex AI"} + +# ---- prompt_cache_5m ---- +- {id: llm.messages.azure_foundry.prompt_cache_5m.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "5m prompt cache over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.prompt_cache_5m.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "5m prompt cache over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.prompt_cache_5m.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "5m prompt cache over Bedrock Invoke"} +- {id: llm.messages.vertex.prompt_cache_5m.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "5m prompt cache over Vertex AI"} + +# ---- prompt_cache_1h ---- +- {id: llm.messages.anthropic.prompt_cache_1h.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: prompt_cache_1h, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "1h prompt cache over Anthropic direct"} +- {id: llm.messages.azure_foundry.prompt_cache_1h.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: prompt_cache_1h, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "1h prompt cache over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.prompt_cache_1h.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: prompt_cache_1h, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "1h prompt cache over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.prompt_cache_1h.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: prompt_cache_1h, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "1h prompt cache over Bedrock Invoke"} +- {id: llm.messages.vertex.prompt_cache_1h.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: prompt_cache_1h, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "1h prompt cache over Vertex AI"} + +# ---- structured_output ---- +- {id: llm.messages.anthropic.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: structured_output, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Structured outputs (--json-schema) over Anthropic direct"} +- {id: llm.messages.azure_foundry.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: structured_output, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Structured outputs over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: structured_output, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Structured outputs over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: structured_output, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Structured outputs over Bedrock Invoke"} +- {id: llm.messages.vertex.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: structured_output, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Structured outputs over Vertex AI"} + +# ---- pdf_input ---- +- {id: llm.messages.anthropic.pdf_input.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: pdf_input, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "PDF document input over Anthropic direct"} +- {id: llm.messages.azure_foundry.pdf_input.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: pdf_input, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "PDF document input over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.pdf_input.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: pdf_input, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "PDF document input over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.pdf_input.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: pdf_input, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "PDF document input over Bedrock Invoke"} +- {id: llm.messages.vertex.pdf_input.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: pdf_input, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "PDF document input over Vertex AI"} + +# ---- long_context_1m ---- +- {id: llm.messages.anthropic.long_context_1m.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: long_context_1m, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "1M context beta over Anthropic direct"} +- {id: llm.messages.azure_foundry.long_context_1m.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: long_context_1m, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "1M context beta over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.long_context_1m.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: long_context_1m, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "1M context beta over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.long_context_1m.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: long_context_1m, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "1M context beta over Bedrock Invoke"} +- {id: llm.messages.vertex.long_context_1m.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: long_context_1m, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "1M context beta over Vertex AI"} + +# ---- thinking_with_tool_use ---- +- {id: llm.messages.anthropic.thinking_with_tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: thinking_with_tool_use, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Thinking + tool_use interleaved over Anthropic direct"} +- {id: llm.messages.azure_foundry.thinking_with_tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: thinking_with_tool_use, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Thinking + tool_use interleaved over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.thinking_with_tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: thinking_with_tool_use, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Thinking + tool_use interleaved over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.thinking_with_tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: thinking_with_tool_use, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Thinking + tool_use interleaved over Bedrock Invoke"} +- {id: llm.messages.vertex.thinking_with_tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: thinking_with_tool_use, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Thinking + tool_use interleaved over Vertex AI"} + +# ---- tool_search ---- +- {id: llm.messages.anthropic.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search_tool_regex_20251119 discovery tool over Anthropic direct"} +- {id: llm.messages.azure_foundry.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search discovery tool over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search discovery tool over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search discovery tool over Bedrock Invoke"} +- {id: llm.messages.vertex.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search discovery tool over Vertex AI"} + +# ---- count_tokens ---- +- {id: llm.messages.anthropic.count_tokens.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: count_tokens, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "/v1/messages/count_tokens over Anthropic direct"} +- {id: llm.messages.azure_foundry.count_tokens.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: count_tokens, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "/v1/messages/count_tokens over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.count_tokens.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: count_tokens, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "/v1/messages/count_tokens over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.count_tokens.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: count_tokens, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "/v1/messages/count_tokens over Bedrock Invoke"} +- {id: llm.messages.vertex.count_tokens.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: count_tokens, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "/v1/messages/count_tokens over Vertex AI"} + +# ---- web_search ---- +- {id: llm.messages.anthropic.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Web search server tool over Anthropic direct"} +- {id: llm.messages.azure_foundry.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Web search server tool over Azure AI Foundry"} +- {id: llm.messages.bedrock_converse.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Web search server tool over Bedrock Converse"} +- {id: llm.messages.bedrock_invoke.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Web search server tool over Bedrock Invoke"} +- {id: llm.messages.vertex.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Web search server tool over Vertex AI"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 7482088d93f..a76774c3bde 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -54,13 +54,20 @@ LlmRoute = Literal[ LlmCapability = Literal[ "basic", + "count_tokens", + "long_context_1m", "mid_conversation_system", + "pdf_input", + "prompt_cache_1h", "prompt_cache_5m", "service_tier", "structured_output", "thinking", + "thinking_with_tool_use", + "tool_search", "tool_use", "vision", + "web_search", ] diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index 3e5de028114..5f75409f025 100644 --- a/tests/e2e/docker-compose.yml +++ b/tests/e2e/docker-compose.yml @@ -157,6 +157,8 @@ services: MISTRAL_API_KEY: ${MISTRAL_API_KEY:-} AZURE_API_BASE: ${AZURE_API_BASE:-} AZURE_API_KEY: ${AZURE_API_KEY:-} + AZURE_AI_API_BASE: ${AZURE_AI_API_BASE:-} + AZURE_AI_API_KEY: ${AZURE_AI_API_KEY:-} ports: - "4000:4000" configs: diff --git a/tests/e2e/e2e_gateway.py b/tests/e2e/e2e_gateway.py index d40b96d60fa..ad8b2e833a8 100644 --- a/tests/e2e/e2e_gateway.py +++ b/tests/e2e/e2e_gateway.py @@ -319,21 +319,31 @@ class Gateway: return self.transport.probe(path, params=params) -def build_gateway() -> Gateway: +def build_gateway( + *, + base_url: str = PROXY_BASE_URL, + master_key: str = MASTER_KEY, + control_plane_base_url: str = CONTROL_PLANE_BASE_URL, +) -> Gateway: """The Gateway every suite's client is built from: a SplitTransport that routes LLM calls to the data plane (PROXY_BASE_URL) and management/admin calls to the control plane (CONTROL_PLANE_BASE_URL), with the shared poll budget. The two - base URLs are the same for a monolithic proxy, so routing is then a no-op.""" + base URLs are the same for a monolithic proxy, so routing is then a no-op. + + The endpoints are injectable for callers that resolve the proxy some other + way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must + pass all three together, since a caller that overrides only the data plane + would leave management calls pointed at the env default.""" return Gateway( transport=SplitTransport( data=HttpTransport( - base_url=PROXY_BASE_URL, - master_key=MASTER_KEY, + base_url=base_url, + master_key=master_key, request_timeout=REQUEST_TIMEOUT, ), control=HttpTransport( - base_url=CONTROL_PLANE_BASE_URL, - master_key=MASTER_KEY, + base_url=control_plane_base_url, + master_key=master_key, request_timeout=REQUEST_TIMEOUT, ), ), diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 4140967f3e0..c4ecd0cfa63 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -440,6 +440,8 @@ class LiteLLMParamsBody(BaseModel): aws_batch_role_arn: str | None = None input_cost_per_token: float | None = None output_cost_per_token: float | None = None + extra_headers: dict[str, str] | None = None + use_in_pass_through: bool | None = None ModelMode = Literal["batch", "realtime", "image_generation"] From 9d88f9a8946543c2620555a05167d9eeb0df0362 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 16 Jul 2026 11:44:35 -0700 Subject: [PATCH 49/62] ci: run zizmor and proxy-db unit tests on PRs targeting litellm_ branches --- .github/workflows/test-unit-proxy-db.yml | 2 ++ .github/workflows/zizmor.yml | 6 +++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 2ac9a3b7c1c..b0ee56f5a5c 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -5,6 +5,8 @@ on: branches: - main - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" permissions: contents: read diff --git a/.github/workflows/zizmor.yml b/.github/workflows/zizmor.yml index db79fe43038..df242e5a3b6 100644 --- a/.github/workflows/zizmor.yml +++ b/.github/workflows/zizmor.yml @@ -4,7 +4,11 @@ on: push: branches: [main, litellm_internal_staging] pull_request: - branches: [main, litellm_internal_staging] + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} From 899ddef21934bf309fc3916a9bd1028c878cbd86 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 16 Jul 2026 12:01:11 -0700 Subject: [PATCH 50/62] feat(logging): add user and team level spend and budget to StandardLoggingPayload metadata (#33459) * feat(logging): add user and team level spend and budget to StandardLoggingPayload metadata * fix(logging): include user and team budget fields in dummy standard logging payload --- .../pagerduty/pagerduty.py | 8 +++ litellm/litellm_core_utils/litellm_logging.py | 12 ++++ litellm/proxy/litellm_pre_call_utils.py | 4 ++ litellm/types/utils.py | 4 ++ .../proxy/test_litellm_pre_call_utils.py | 64 +++++++++++++++++++ 5 files changed, 92 insertions(+) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py index 12fdaeb6a81..f920aa7ac13 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py @@ -113,6 +113,10 @@ class PagerDutyAlerting(SlackAlerting): user_api_key_spend=_meta.get("user_api_key_spend"), user_api_key_max_budget=_meta.get("user_api_key_max_budget"), user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"), + user_api_key_user_spend=_meta.get("user_api_key_user_spend"), + user_api_key_user_max_budget=_meta.get("user_api_key_user_max_budget"), + user_api_key_team_spend=_meta.get("user_api_key_team_spend"), + user_api_key_team_max_budget=_meta.get("user_api_key_team_max_budget"), user_api_key_org_id=_meta.get("user_api_key_org_id"), user_api_key_org_alias=_meta.get("user_api_key_org_alias"), user_api_key_team_id=_meta.get("user_api_key_team_id"), @@ -196,6 +200,10 @@ class PagerDutyAlerting(SlackAlerting): if user_api_key_dict.budget_reset_at else None ), + user_api_key_user_spend=user_api_key_dict.user_spend, + user_api_key_user_max_budget=user_api_key_dict.user_max_budget, + user_api_key_team_spend=user_api_key_dict.team_spend, + user_api_key_team_max_budget=user_api_key_dict.team_max_budget, user_api_key_org_id=user_api_key_dict.org_id, user_api_key_org_alias=user_api_key_dict.organization_alias, user_api_key_team_id=user_api_key_dict.team_id, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 268dcc3df78..34912c81277 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4597,6 +4597,10 @@ class StandardLoggingPayloadSetup: user_api_key_spend=None, user_api_key_max_budget=None, user_api_key_budget_reset_at=None, + user_api_key_user_spend=None, + user_api_key_user_max_budget=None, + user_api_key_team_spend=None, + user_api_key_team_max_budget=None, user_api_key_team_id=None, user_api_key_org_id=None, user_api_key_org_alias=None, @@ -5428,6 +5432,10 @@ def get_standard_logging_metadata( user_api_key_spend=None, user_api_key_max_budget=None, user_api_key_budget_reset_at=None, + user_api_key_user_spend=None, + user_api_key_user_max_budget=None, + user_api_key_team_spend=None, + user_api_key_team_max_budget=None, user_api_key_team_id=None, user_api_key_org_id=None, user_api_key_org_alias=None, @@ -5527,6 +5535,10 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: user_api_key_team_id=str("test_team"), user_api_key_user_id=str("test_user"), user_api_key_team_alias=str("test_team_alias"), + user_api_key_user_spend=None, + user_api_key_user_max_budget=None, + user_api_key_team_spend=None, + user_api_key_team_max_budget=None, user_api_key_org_id=None, spend_logs_metadata=None, requester_ip_address=str("127.0.0.1"), diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index afd8a437cc1..15d1876e5a2 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -949,6 +949,10 @@ class LiteLLMProxyRequestSetup: user_api_key_alias=user_api_key_dict.key_alias, user_api_key_spend=user_api_key_dict.spend, user_api_key_max_budget=user_api_key_dict.max_budget, + user_api_key_user_spend=user_api_key_dict.user_spend, + user_api_key_user_max_budget=user_api_key_dict.user_max_budget, + user_api_key_team_spend=user_api_key_dict.team_spend, + user_api_key_team_max_budget=user_api_key_dict.team_max_budget, user_api_key_team_id=user_api_key_dict.team_id, user_api_key_project_id=user_api_key_dict.project_id, user_api_key_project_alias=user_api_key_dict.project_alias, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4fb3284c05b..1633f1293fe 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2475,6 +2475,10 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict): user_api_key_spend: Optional[float] user_api_key_max_budget: Optional[float] user_api_key_budget_reset_at: Optional[str] + user_api_key_user_spend: Optional[float] + user_api_key_user_max_budget: Optional[float] + user_api_key_team_spend: Optional[float] + user_api_key_team_max_budget: Optional[float] user_api_key_org_id: Optional[str] user_api_key_org_alias: Optional[str] user_api_key_team_id: Optional[str] diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 913ac116866..d2b8b7ec23d 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2688,6 +2688,70 @@ def test_get_sanitized_user_information_from_key_includes_guardrails_metadata(): assert result["user_api_key_auth_metadata"]["other_field"] == "value" +def test_user_and_team_spend_and_budget_flow_to_standard_logging_metadata(): + """ + Full flow: UserAPIKeyAuth -> get_sanitized_user_information_from_key -> + get_standard_logging_metadata. User-level and team-level spend + max budget + must reach the StandardLoggingPayload metadata that custom loggers receive, + alongside the key-level values + """ + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key-hash", + spend=1.5, + max_budget=10.0, + user_id="test-user", + user_spend=25.5, + user_max_budget=100.0, + team_id="test-team", + team_spend=250.75, + team_max_budget=1000.0, + ) + + sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=user_api_key_dict + ) + + assert sanitized["user_api_key_spend"] == 1.5 + assert sanitized["user_api_key_max_budget"] == 10.0 + assert sanitized["user_api_key_user_spend"] == 25.5 + assert sanitized["user_api_key_user_max_budget"] == 100.0 + assert sanitized["user_api_key_team_spend"] == 250.75 + assert sanitized["user_api_key_team_max_budget"] == 1000.0 + + logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( + dict(sanitized) + ) + + assert logging_metadata["user_api_key_user_spend"] == 25.5 + assert logging_metadata["user_api_key_user_max_budget"] == 100.0 + assert logging_metadata["user_api_key_team_spend"] == 250.75 + assert logging_metadata["user_api_key_team_max_budget"] == 1000.0 + + +def test_user_and_team_spend_and_budget_default_to_none_in_standard_logging_metadata(): + """ + Keys with no user or team level budgets report None for the new fields in the + StandardLoggingPayload metadata instead of raising + """ + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + user_api_key_dict = UserAPIKeyAuth(api_key="test-key-hash") + + sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=user_api_key_dict + ) + logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( + dict(sanitized) + ) + + assert logging_metadata["user_api_key_user_spend"] is None + assert logging_metadata["user_api_key_user_max_budget"] is None + assert logging_metadata["user_api_key_team_spend"] is None + assert logging_metadata["user_api_key_team_max_budget"] is None + + @pytest.mark.asyncio async def test_team_guardrails_append_to_key_guardrails(): """ From c012373e1c8b0bee79e1339194263d34eaeb8247 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 12:05:36 -0700 Subject: [PATCH 51/62] fix(router): cast model_info cost values to float in _set_model_group_info (#33556) Cost values read from deployment model_info can be strings when the config YAML contains scientific notation with an integer mantissa (e.g. 1e-05), which YAML 1.2 parsers such as PyYAML 6.x treat as a string. Comparing that string against the running float aggregate in _set_model_group_info raised TypeError and broke /model_group/info, the prometheus remaining-usage callback, and the x-litellm-response-cost header. Coerce input/output cost values to float before comparing and storing them. --- litellm/router.py | 27 +++++-- tests/test_litellm/test_router.py | 126 ++++++++++++++++++++++++++++++ 2 files changed, 145 insertions(+), 8 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index f408d030b8b..809aef695a7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -251,6 +251,15 @@ else: PreRoutingHookResponse = Any +def _cost_value_as_float(value: Union[str, int, float, None]) -> Optional[float]: + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + class RoutingArgs(enum.Enum): ttl = 60 # 1min (RPM/TPM expire key) @@ -8750,8 +8759,8 @@ class Router: # Get mode from database model_info if available, otherwise default to "chat" db_model_info = model.get("model_info", {}) mode = db_model_info.get("mode", "chat") - input_cost_per_token = db_model_info.get("input_cost_per_token") - output_cost_per_token = db_model_info.get("output_cost_per_token") + input_cost_per_token = _cost_value_as_float(db_model_info.get("input_cost_per_token")) + output_cost_per_token = _cost_value_as_float(db_model_info.get("output_cost_per_token")) model_info = ModelMapInfo( key=model_group, @@ -8802,16 +8811,18 @@ class Router: ) ): model_group_info.max_output_tokens = model_info["max_output_tokens"] - if model_info.get("input_cost_per_token", None) is not None and ( + _input_cost_per_token = _cost_value_as_float(model_info.get("input_cost_per_token")) + if _input_cost_per_token is not None and ( model_group_info.input_cost_per_token is None - or (model_info["input_cost_per_token"] or 0.0) > (model_group_info.input_cost_per_token or 0.0) + or _input_cost_per_token > (model_group_info.input_cost_per_token or 0.0) ): - model_group_info.input_cost_per_token = model_info["input_cost_per_token"] - if model_info.get("output_cost_per_token", None) is not None and ( + model_group_info.input_cost_per_token = _input_cost_per_token + _output_cost_per_token = _cost_value_as_float(model_info.get("output_cost_per_token")) + if _output_cost_per_token is not None and ( model_group_info.output_cost_per_token is None - or (model_info["output_cost_per_token"] or 0.0) > (model_group_info.output_cost_per_token or 0.0) + or _output_cost_per_token > (model_group_info.output_cost_per_token or 0.0) ): - model_group_info.output_cost_per_token = model_info["output_cost_per_token"] + model_group_info.output_cost_per_token = _output_cost_per_token if ( model_info.get("supports_parallel_function_calling", None) is not None and model_info["supports_parallel_function_calling"] is True # type: ignore diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 9c4d83ff7ea..c2c98c8869c 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1455,6 +1455,132 @@ def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): assert result.output_cost_per_token is None +@pytest.mark.parametrize( + "value,expected", + [ + ("1e-05", 1e-05), + ("0.00001", 1e-05), + (1e-05, 1e-05), + (5, 5.0), + (None, None), + ("not-a-number", None), + ], +) +def test_cost_value_as_float(value, expected): + from litellm.router import _cost_value_as_float + + assert _cost_value_as_float(value) == expected + + +def test_model_group_info_with_stringified_cost_values(): + """ + YAML 1.2 parsers emit '1e-05' (integer mantissa) as a string, so cost + values in deployment model_info can arrive as str. Aggregating the model + group must not raise TypeError('>' between str and float) and must return + float costs. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "my-custom-model", + "litellm_params": { + "model": "openai/my-custom-backend-1", + "api_key": "fake", + }, + "model_info": { + "input_cost_per_token": "1e-05", + "output_cost_per_token": "1e-05", + }, + }, + { + "model_name": "my-custom-model", + "litellm_params": { + "model": "openai/my-custom-backend-2", + "api_key": "fake", + }, + "model_info": { + "input_cost_per_token": "2e-05", + "output_cost_per_token": "2e-05", + }, + }, + ] + ) + + def _model_info_with_str_costs(model_id: str, model_name: str): + for model in router.model_list: + if model["model_info"]["id"] == model_id: + return { + "key": model_name, + "input_cost_per_token": model["model_info"]["input_cost_per_token"], + "output_cost_per_token": model["model_info"]["output_cost_per_token"], + "litellm_provider": "openai", + "mode": "chat", + } + return None + + with patch.object( + router, "get_deployment_model_info", side_effect=_model_info_with_str_costs + ): + result = router._set_model_group_info( + model_group="my-custom-model", + user_facing_model_group_name="my-custom-model", + ) + + assert result is not None + assert result.input_cost_per_token == 2e-05 + assert result.output_cost_per_token == 2e-05 + assert isinstance(result.input_cost_per_token, float) + assert isinstance(result.output_cost_per_token, float) + + +def test_model_group_info_db_fallback_with_stringified_cost_values(): + """ + Fallback path: when get_deployment_model_info returns nothing, costs are + read straight from the deployment's model_info dict, which can hold + stringified floats parsed from YAML. They must be coerced to float. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "my-custom-model", + "litellm_params": { + "model": "openai/my-custom-backend-1", + "api_key": "fake", + }, + "model_info": { + "input_cost_per_token": "1e-05", + "output_cost_per_token": "3e-05", + }, + }, + { + "model_name": "my-custom-model", + "litellm_params": { + "model": "openai/my-custom-backend-2", + "api_key": "fake", + }, + "model_info": { + "input_cost_per_token": "2e-05", + "output_cost_per_token": "2e-05", + }, + }, + ] + ) + + with patch.object( + router, "get_deployment_model_info", side_effect=Exception("not found") + ): + result = router._set_model_group_info( + model_group="my-custom-model", + user_facing_model_group_name="my-custom-model", + ) + + assert result is not None + assert result.input_cost_per_token == 2e-05 + assert result.output_cost_per_token == 3e-05 + assert isinstance(result.input_cost_per_token, float) + assert isinstance(result.output_cost_per_token, float) + + def test_get_model_access_groups_caching(): """ Test that get_model_access_groups caches the no-args result From ebdf0bbfd73e5e1e602d5b84039b6705c8494c68 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 16 Jul 2026 12:09:24 -0700 Subject: [PATCH 52/62] chore(e2e): establish litellm_e2e_staging integration line (#33502) * chore(e2e): establish litellm_e2e_staging integration line Long-lived berri branch for e2e suite recovery work (LIT-4479 through LIT-4486) before merge to litellm_internal_staging * test(e2e): remove langfuse_otel logging e2e suite (#33558) * test(e2e): remove langfuse_otel logging e2e suite Removes the LIT-4483 dynamic per-team/key/org langfuse_otel logging e2e tests (tests/e2e/logging/test_langfuse_e2e.py, added in #32857). The shared logging_client harness and the langfuse coverage-registry cells are left in place; only the test module is removed. The otel and prometheus logging e2e suites are unaffected. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): drop orphaned langfuse coverage-registry cells The three logging.langfuse.*.logs_spend P0 cells were only exercised by the deleted langfuse_otel e2e suite. Remove them so the coverage registry has no orphaned rows. 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> * fix(e2e): log into the react admin ui in the management browser fixture (#33562) The management ui_page fixture drove the old server-rendered login form: it clicked input[type="submit"] and treated wait_for_url("**/ui/**") as the done signal. /ui/ now serves the react (antd) dashboard whose submit is a +
); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx index 3a22a55298e..dfe307c7edf 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx @@ -1,5 +1,5 @@ import * as networking from "@/components/networking"; -import { afterEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; import ModelHubTable from "./ModelHubTable"; @@ -7,6 +7,7 @@ const mockUseUISettings = vi.hoisted(() => vi.fn()); const mockGetCookie = vi.hoisted(() => vi.fn()); const mockCheckTokenValidity = vi.hoisted(() => vi.fn()); const mockRouterReplace = vi.hoisted(() => vi.fn()); +const mockLocationReplace = vi.hoisted(() => vi.fn()); vi.mock("@/components/networking", () => ({ getUiConfig: vi.fn(), @@ -43,7 +44,28 @@ vi.mock("@/utils/jwtUtils", () => ({ })); describe("ModelHubTable", () => { + const originalLocation = window.location; + + beforeEach(() => { + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:4000/ui/model_hub_table", + origin: "http://localhost:4000", + hostname: "localhost", + pathname: "/ui/model_hub_table", + search: "", + protocol: "http:", + replace: mockLocationReplace, + }, + writable: true, + }); + }); + afterEach(() => { + Object.defineProperty(window, "location", { + value: originalLocation, + writable: true, + }); vi.clearAllMocks(); }); @@ -60,6 +82,7 @@ describe("ModelHubTable", () => { mockGetCookie.mockReturnValue(tokenValue); mockCheckTokenValidity.mockReturnValue(isTokenValid); mockRouterReplace.mockClear(); + mockLocationReplace.mockClear(); // Setup other required mocks vi.mocked(networking.getUiConfig).mockResolvedValue({ @@ -92,9 +115,10 @@ describe("ModelHubTable", () => { await waitFor(() => { if (shouldRedirect) { - expect(mockRouterReplace).toHaveBeenCalledWith("http://localhost:4000/ui/login"); - } else { + expect(mockLocationReplace).toHaveBeenCalledWith("http://localhost:4000/ui/login/"); expect(mockRouterReplace).not.toHaveBeenCalled(); + } else { + expect(mockLocationReplace).not.toHaveBeenCalled(); } }); }); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index eb47f775d90..1f64d175052 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -33,6 +33,7 @@ import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import { checkTokenValidity } from "@/utils/jwtUtils"; import { getCookie } from "@/utils/cookieUtils"; +import { getLoginUrl } from "@/utils/returnUrlUtils"; interface ModelHubTableProps { accessToken: string | null; @@ -108,12 +109,12 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, // If token is invalid, redirect to login if (!isTokenValid) { - router.replace(`${getProxyBaseUrl()}/ui/login`); + window.location.replace(getLoginUrl(getProxyBaseUrl())); return; } } // If require_auth_for_public_ai_hub is false, allow public access (no change) - }, [isUISettingsLoading, publicPage, uiSettings, router]); + }, [isUISettingsLoading, publicPage, uiSettings]); useEffect(() => { const fetchData = async (accessToken: string) => { diff --git a/ui/litellm-dashboard/src/components/DashboardHeader.tsx b/ui/litellm-dashboard/src/components/DashboardHeader.tsx index 8b928f8e6fe..babf007d475 100644 --- a/ui/litellm-dashboard/src/components/DashboardHeader.tsx +++ b/ui/litellm-dashboard/src/components/DashboardHeader.tsx @@ -18,7 +18,7 @@ import WorkerDropdown from "@/components/Navbar/WorkerDropdown/WorkerDropdown"; import { useWorker } from "@/hooks/useWorker"; import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPrompts"; import { clearTokenCookies } from "@/utils/cookieUtils"; -import { clearStoredReturnUrl } from "@/utils/returnUrlUtils"; +import { clearStoredReturnUrl, getLoginUrl } from "@/utils/returnUrlUtils"; interface DashboardHeaderProps { page: string; @@ -37,7 +37,7 @@ export function DashboardHeader({ page }: DashboardHeaderProps) { clearStoredReturnUrl(); localStorage.removeItem("litellm_selected_worker_id"); localStorage.removeItem("litellm_worker_url"); - window.location.href = `/ui/login?worker=${encodeURIComponent(workerId)}`; + window.location.href = `${getLoginUrl()}?worker=${encodeURIComponent(workerId)}`; }; return ( diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index 1a0b1c38520..40638e7b8ba 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -5,7 +5,7 @@ import { useWorker } from "@/hooks/useWorker"; import { getProxyBaseUrl } from "@/components/networking"; import { useTheme } from "@/contexts/ThemeContext"; import { clearTokenCookies } from "@/utils/cookieUtils"; -import { clearStoredReturnUrl } from "@/utils/returnUrlUtils"; +import { clearStoredReturnUrl, getLoginUrl } from "@/utils/returnUrlUtils"; import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; import { DownOutlined, MenuFoldOutlined, MenuUnfoldOutlined } from "@ant-design/icons"; import { Tag } from "antd"; @@ -56,7 +56,7 @@ const Navbar: React.FC = ({ clearStoredReturnUrl(); localStorage.removeItem("litellm_selected_worker_id"); localStorage.removeItem("litellm_worker_url"); - window.location.href = `/ui/login?worker=${encodeURIComponent(workerId)}`; + window.location.href = `${getLoginUrl()}?worker=${encodeURIComponent(workerId)}`; }; return ( diff --git a/ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts b/ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts index 56de299049c..0f3a8b4bf69 100644 --- a/ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts @@ -3,6 +3,7 @@ import { clearStoredReturnUrl, consumeReturnUrl, getCurrentUrl, + getLoginUrl, getReturnUrl, getReturnUrlFromParams, getStoredReturnUrl, @@ -98,6 +99,29 @@ describe("returnUrlUtils", () => { }); }); + describe("getLoginUrl", () => { + it("should build a relative login URL with a trailing slash", () => { + expect(getLoginUrl()).toBe("/ui/login/"); + }); + + it("should prepend the given base URL and keep the trailing slash", () => { + expect(getLoginUrl("http://proxy.example")).toBe("http://proxy.example/ui/login/"); + }); + + it("should keep the trailing slash before the query when composed with buildLoginUrlWithReturn", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + href: "http://localhost:3000/ui?page=api-keys", + }, + writable: true, + }); + + const loginUrl = buildLoginUrlWithReturn(getLoginUrl()); + expect(loginUrl).toBe("/ui/login/?redirect_to=http%3A%2F%2Flocalhost%3A3000%2Fui%3Fpage%3Dapi-keys"); + }); + }); + describe("buildLoginUrlWithReturn", () => { it("should build login URL with return URL parameter", () => { Object.defineProperty(window, "location", { diff --git a/ui/litellm-dashboard/src/utils/returnUrlUtils.ts b/ui/litellm-dashboard/src/utils/returnUrlUtils.ts index 76562a0122a..b3cc5345bb1 100644 --- a/ui/litellm-dashboard/src/utils/returnUrlUtils.ts +++ b/ui/litellm-dashboard/src/utils/returnUrlUtils.ts @@ -13,6 +13,10 @@ const RETURN_URL_COOKIE_NAME = "litellm_return_url"; const RETURN_URL_PARAM = "redirect_to"; +export function getLoginUrl(baseUrl: string = ""): string { + return `${baseUrl}/ui/login/`; +} + /** * Gets the current URL with all query parameters. * Returns null if running on server-side. diff --git a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx index 62346eff057..211a754dad4 100644 --- a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx +++ b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx @@ -232,7 +232,7 @@ describe("CreateKeyPage auth behavior", () => { // Assert: we eventually redirect to SSO login with return URL (single replace, not assign/href) await waitFor(() => { expect(window.location.replace).toHaveBeenCalledWith( - expect.stringContaining("https://example.com/ui/login?redirect_to="), + expect.stringContaining("https://example.com/ui/login/?redirect_to="), ); }); From 7fe3dd86a4ed42cdf67fe21f1f70b796332954e6 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 16 Jul 2026 12:39:04 -0700 Subject: [PATCH 54/62] feat(logging): add structured budget fields to budget rejection failure logs (#33460) --- litellm/exceptions.py | 4 ++ litellm/litellm_core_utils/litellm_logging.py | 6 +++ litellm/proxy/auth/auth_checks.py | 26 +++++++++- litellm/proxy/auth/user_api_key_auth.py | 2 + .../proxy/hooks/model_max_budget_limiter.py | 6 ++- .../spend_tracking/budget_reservation.py | 16 ++++++- litellm/types/utils.py | 4 ++ .../test_litellm_logging.py | 47 +++++++++++++++++++ .../proxy/auth/test_auth_checks.py | 10 ++++ .../proxy/test_budget_reservation.py | 35 +++++++++++++- 10 files changed, 152 insertions(+), 4 deletions(-) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index aca3fb551cc..fd0a2afb3e8 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -966,11 +966,15 @@ class BudgetExceededError(Exception): max_budget: float, message: Optional[str] = None, llm_provider: Optional[str] = None, + entity_type: Optional[str] = None, + entity_id: Optional[str] = None, ): self.current_cost = current_cost self.max_budget = max_budget self.status_code = 429 self.llm_provider = llm_provider or "" + self.entity_type = entity_type + self.entity_id = entity_id # Surface unified rate-limit fields without joining the RateLimitError # hierarchy so existing `except BudgetExceededError:` handlers keep # working; custom callbacks reading StandardLoggingPayload pick these diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 34912c81277..9a0b4937fdb 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -38,6 +38,7 @@ from litellm import ( ) from litellm._logging import _is_debugging_on, _redact_string, verbose_logger from litellm.exceptions import ( + BudgetExceededError, validate_rate_limit_category, validate_rate_limit_type, ) @@ -4947,6 +4948,7 @@ class StandardLoggingPayloadSetup: rate_limit_category = validate_rate_limit_category(getattr(original_exception, "category", None)) rate_limit_type = validate_rate_limit_type(getattr(original_exception, "rate_limit_type", None)) + budget_error = original_exception if isinstance(original_exception, BudgetExceededError) else None return StandardLoggingPayloadErrorInformation( error_code=error_status, @@ -4956,6 +4958,10 @@ class StandardLoggingPayloadSetup: error_message=error_message, error_rate_limit_category=rate_limit_category, error_rate_limit_type=rate_limit_type, + error_budget_entity_type=budget_error.entity_type if budget_error else None, + error_budget_entity_id=budget_error.entity_id if budget_error else None, + error_budget_limit=budget_error.max_budget if budget_error else None, + error_budget_spend=budget_error.current_cost if budget_error else None, ) @staticmethod diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 00f6d44e25a..fa354b8cccb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -361,7 +361,11 @@ def _global_proxy_budget_check(global_proxy_spend: Optional[float], skip_budget_ and route != "/models" ): if math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget: - raise litellm.BudgetExceededError(current_cost=global_proxy_spend, max_budget=litellm.max_budget) + raise litellm.BudgetExceededError( + current_cost=global_proxy_spend, + max_budget=litellm.max_budget, + entity_type=Litellm_EntityType.PROXY.value, + ) _GUARDRAIL_MODIFICATION_KEYS: tuple = ( @@ -648,6 +652,8 @@ async def common_checks( current_cost=user_spend, max_budget=user_budget, message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", + entity_type=Litellm_EntityType.USER.value, + entity_id=user_object.user_id, ) # Each scope reads a distinct counter key with no cross-scope ordering @@ -1093,6 +1099,8 @@ async def _check_end_user_budget( current_cost=end_user_spend, max_budget=end_user_budget, message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}", + entity_type=Litellm_EntityType.END_USER.value, + entity_id=end_user_obj.user_id, ) @@ -3552,6 +3560,8 @@ async def _virtual_key_max_budget_check( current_cost=spend, max_budget=valid_token.max_budget, message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}", + entity_type=Litellm_EntityType.KEY.value, + entity_id=valid_token.token, ) @@ -3593,6 +3603,8 @@ async def _virtual_key_multi_budget_check( f"ExceededBudget: Key over {w['budget_duration']} budget. " f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" ), + entity_type=Litellm_EntityType.KEY.value, + entity_id=valid_token.token, ) @@ -3824,6 +3836,8 @@ async def _check_team_member_budget( current_cost=team_member_spend, max_budget=team_member_budget, message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}", + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=f"{valid_token.user_id}:{team_object.team_id}", ) @@ -3923,6 +3937,8 @@ async def _team_max_budget_check( current_cost=spend, max_budget=team_object.max_budget, message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}", + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_object.team_id, ) @@ -3960,6 +3976,8 @@ async def _team_multi_budget_check( f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. " f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" ), + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_object.team_id, ) @@ -4081,6 +4099,8 @@ async def _project_max_budget_check( current_cost=project_object.spend, max_budget=max_budget, message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_object.spend}, Max budget: {max_budget}", + entity_type=Litellm_EntityType.PROJECT.value, + entity_id=project_object.project_id, ) @@ -4269,6 +4289,8 @@ async def _organization_max_budget_check( current_cost=org_spend, max_budget=org_max_budget, message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}", + entity_type=Litellm_EntityType.ORGANIZATION.value, + entity_id=org_id, ) @@ -4326,6 +4348,8 @@ async def _tag_max_budget_check( current_cost=tag_spend, max_budget=tag_object.litellm_budget_table.max_budget, message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}", + entity_type=Litellm_EntityType.TAG.value, + entity_id=tag_name, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b402212fb2e..0519b5ef0b6 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1797,6 +1797,8 @@ async def _user_api_key_auth_builder( raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=f"{valid_token.user_id}:{valid_token.team_id}", ) # Check 3. If token is expired diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 803fe64c193..e8cf5fbc718 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -5,7 +5,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( @@ -76,6 +76,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}", current_cost=_current_spend, max_budget=_current_model_budget_info.max_budget, + entity_type=Litellm_EntityType.KEY.value, + entity_id=user_api_key_dict.token, ) return True @@ -140,6 +142,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", current_cost=_current_spend, max_budget=_current_model_budget_info.max_budget, + entity_type=Litellm_EntityType.END_USER.value, + entity_id=end_user_id, ) return True diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index e1a093a4a48..80fd8a1594e 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -4,7 +4,7 @@ import asyncio import json from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Optional, Sequence, cast +from typing import Any, Dict, List, Mapping, Optional, Sequence, cast import litellm from litellm._logging import verbose_proxy_logger @@ -12,6 +12,7 @@ from litellm.caching import DualCache from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.proxy._types import ( + Litellm_EntityType, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, @@ -36,6 +37,17 @@ class _BudgetCounter: window_start: Optional[datetime] = None +_COUNTER_ENTITY_TYPES: Mapping[str, str] = { + "Key": Litellm_EntityType.KEY.value, + "Team": Litellm_EntityType.TEAM.value, + "TeamMember": Litellm_EntityType.TEAM_MEMBER.value, + "User": Litellm_EntityType.USER.value, + "EndUser": Litellm_EntityType.END_USER.value, + "Tag": Litellm_EntityType.TAG.value, + "Organization": Litellm_EntityType.ORGANIZATION.value, +} + + class _CounterReservationUnavailable(Exception): def __init__( self, @@ -108,6 +120,8 @@ async def _apply_over_budget_reservation_policy( f"Current cost: {current_spend}, " f"Max budget: {counter.max_budget}" ), + entity_type=_COUNTER_ENTITY_TYPES.get(counter.entity_type), + entity_id=counter.spend_log_entity_id or counter.entity_id, ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 1633f1293fe..db39d58a925 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2697,6 +2697,10 @@ class StandardLoggingPayloadErrorInformation(TypedDict, total=False): # hints). Lets dashboards split rate-limit failures by cause without # parsing free-text error messages. error_rate_limit_type: Optional[str] + error_budget_entity_type: Optional[str] + error_budget_entity_id: Optional[str] + error_budget_limit: Optional[float] + error_budget_spend: Optional[float] class GuardrailMode(TypedDict, total=False): diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f99cfb953ba..6875894c1bf 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2247,6 +2247,53 @@ def test_get_error_information_prefers_message_attribute_over_str(): assert result["error_class"] == "ProxyExceptionLike" +def test_get_error_information_budget_exceeded_structured_fields(): + """ + Regression for LIT-4458: a budget-rejected request's failure + StandardLoggingPayload must identify WHICH budget blocked the call + as structured fields, not only inside the free-text error_str + ("ExceededBudget: User=... over budget. Spend=..., Budget=..."). + + Asserts get_error_information copies entity_type / entity_id / + max_budget / current_cost off BudgetExceededError into + error_budget_entity_type / error_budget_entity_id / + error_budget_limit / error_budget_spend, and leaves all four None + for non-budget exceptions. + """ + from litellm.exceptions import BudgetExceededError + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + exc = BudgetExceededError( + current_cost=3.4e-05, + max_budget=1e-06, + message="ExceededBudget: User=repro-user over budget. Spend=3.4e-05, Budget=1e-06", + entity_type="user", + entity_id="repro-user", + ) + + result = StandardLoggingPayloadSetup.get_error_information(exc) + assert result["error_budget_entity_type"] == "user" + assert result["error_budget_entity_id"] == "repro-user" + assert result["error_budget_limit"] == 1e-06 + assert result["error_budget_spend"] == 3.4e-05 + assert result["error_code"] == "429" + assert result["error_class"] == "BudgetExceededError" + assert result["error_rate_limit_type"] == "budget" + + legacy_exc = BudgetExceededError(current_cost=2.0, max_budget=1.0) + legacy_result = StandardLoggingPayloadSetup.get_error_information(legacy_exc) + assert legacy_result["error_budget_entity_type"] is None + assert legacy_result["error_budget_entity_id"] is None + assert legacy_result["error_budget_limit"] == 1.0 + assert legacy_result["error_budget_spend"] == 2.0 + + non_budget_result = StandardLoggingPayloadSetup.get_error_information(ValueError("boom")) + assert non_budget_result["error_budget_entity_type"] is None + assert non_budget_result["error_budget_entity_id"] is None + assert non_budget_result["error_budget_limit"] is None + assert non_budget_result["error_budget_spend"] is None + + def test_get_error_information_preserves_explicit_empty_message(): """ An exception that deliberately sets `.message = ""` must surface diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 27f43c4948f..8365909f314 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2634,6 +2634,8 @@ async def test_virtual_key_budget_check_reads_from_spend_counter(): ) assert exc_info.value.current_cost == 1.5 assert exc_info.value.max_budget == 1.0 + assert exc_info.value.entity_type == "key" + assert exc_info.value.entity_id == "test-hashed-token" @pytest.mark.asyncio @@ -2861,6 +2863,8 @@ async def test_team_budget_check_reads_from_spend_counter(): proxy_logging_obj=proxy_logging_obj, ) assert exc_info.value.current_cost == 1.5 + assert exc_info.value.entity_type == "team" + assert exc_info.value.entity_id == "test-team" @pytest.mark.asyncio @@ -2888,6 +2892,8 @@ async def test_end_user_budget_check_reads_from_spend_counter(): ) assert exc_info.value.current_cost == 1.5 assert exc_info.value.max_budget == 1.0 + assert exc_info.value.entity_type == "end_user" + assert exc_info.value.entity_id == "customer-1" @pytest.mark.asyncio @@ -2926,6 +2932,8 @@ async def test_tag_budget_check_reads_from_spend_counter(): ) assert exc_info.value.current_cost == 1.5 assert exc_info.value.max_budget == 1.0 + assert exc_info.value.entity_type == "tag" + assert exc_info.value.entity_id == "paid-tag" @pytest.mark.asyncio @@ -2976,6 +2984,8 @@ async def test_team_member_budget_check_reads_from_spend_counter(): proxy_logging_obj=proxy_logging_obj, ) assert exc_info.value.current_cost == 1.5 + assert exc_info.value.entity_type == "team_member" + assert exc_info.value.entity_id == "test-user:test-team" class TestGuardrailModificationCheck: diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 540f017ee88..0b304f2fec7 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -150,8 +150,41 @@ async def test_reservation_blocks_over_budget_non_throttled_key( await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) # counter -> 1.0 - with pytest.raises(litellm.BudgetExceededError): + with pytest.raises(litellm.BudgetExceededError) as exc_info: await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert exc_info.value.entity_type == "key" + assert exc_info.value.entity_id == "key-no-optin-over" + + +@pytest.mark.asyncio +async def test_over_budget_window_counter_tags_clean_entity_id(): + from litellm.proxy.spend_tracking.budget_reservation import ( + _apply_over_budget_reservation_policy, + _BudgetCounter, + ) + + counter = _BudgetCounter( + counter_key="spend:key:test-token:window:1d", + max_budget=1.0, + fallback_spend=0.0, + entity_type="Key", + entity_id="test-token:1d", + spend_log_entity_id="test-token", + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _apply_over_budget_reservation_policy( + counter=counter, + valid_token=None, + entry={"counter_key": counter.counter_key}, + applied_entries=[], + reservation_cost=0.5, + current_spend=2.0, + ) + assert exc_info.value.entity_type == "key" + assert exc_info.value.entity_id == "test-token" + assert exc_info.value.max_budget == 1.0 + assert exc_info.value.current_cost == 2.0 def test_should_not_serialize_budget_reservation_on_user_api_key_auth(): From a8ae515bee19ef4724c4af293aee4e216e935426 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 13:16:41 -0700 Subject: [PATCH 55/62] fix(proxy_cli): reap orphaned prisma query-engine processes when a worker dies (#33424) * fix(proxy_cli): reap orphaned prisma query-engine processes when a worker dies When the proxy runs multi-worker (uvicorn multiprocess supervisor or the gunicorn arbiter), a worker that crashes or is force-killed never runs its in-process atexit cleanup, so its prisma query-engine subprocess reparents to PID 1 and keeps its database connection pool established forever while the replacement worker opens a fresh pool. Active DB connections then grow past database_connection_pool_limit with every worker death. Run a reaper thread in the supervisor process that marks itself a child subreaper on Linux, scans for adopted query-engine children whose worker is gone, and terminates them with SIGTERM escalating to SIGKILL after a bounded grace period. Engines owned by live workers are children of those workers, never of the supervisor, so they are structurally out of reach. Resolves LIT-4449 Fixes https://github.com/BerriAI/litellm/issues/33023 * fix(proxy_cli): address review findings on the query-engine reaper Make start_query_engine_reaper idempotent, reap simultaneous orphans under one shared grace period instead of serially, and log when a PID survives SIGKILL. Also regenerate schema.d.ts for the update_team docstring line that documents the existing mcp_rpm_limit param (fixes the walk-order-dependent documentation CI failure) and avoid a cast in the prctl wrapper * test(proxy): fix reaper idempotency-test isolation and widen coverage The daemon-thread startup test now stubs threading.enumerate so a reaper thread left running by an earlier test in the same xdist worker cannot satisfy the idempotency guard and skip the code under test. Add coverage for stat-file truncation, non-numeric ppid, non-child reap, signal-to-dead-pid, subreaper capability, and reaper-loop resilience --- litellm/proxy/db/query_engine_reaper.py | 212 +++++++++++++++ litellm/proxy/proxy_cli.py | 4 + .../proxy/db/test_query_engine_reaper.py | 244 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_cli.py | 79 ++++++ 4 files changed, 539 insertions(+) create mode 100644 litellm/proxy/db/query_engine_reaper.py create mode 100644 tests/test_litellm/proxy/db/test_query_engine_reaper.py diff --git a/litellm/proxy/db/query_engine_reaper.py b/litellm/proxy/db/query_engine_reaper.py new file mode 100644 index 00000000000..0e5f0e68910 --- /dev/null +++ b/litellm/proxy/db/query_engine_reaper.py @@ -0,0 +1,212 @@ +"""Supervisor-side reaper for orphaned Prisma query-engine processes. + +Each proxy worker owns a Prisma query-engine subprocess whose only cleanup +hook is an in-process ``atexit`` handler. When a multi-worker supervisor +(uvicorn's multiprocess manager, the gunicorn arbiter) force-kills a hung or +crashed worker, that handler never runs: the engine reparents to the nearest +subreaper (PID 1 in a container, which is the supervisor itself under the +standard docker entrypoint) and keeps its database connection pool +established forever, while the replacement worker opens a fresh pool. Over +repeated worker deaths the active database connections grow without bound. + +The reaper runs only in the supervisor process, where a query-engine process +can never be a legitimate direct child: workers own their engines, and the +supervisor never starts one. Any direct child whose command name begins with +``query-engine`` is therefore an adopted orphan and is terminated +(SIGTERM, bounded grace, SIGKILL) and reaped. On Linux the supervisor also +marks itself a child subreaper so orphans reparent to it even when it is not +PID 1. + +Linux-only by construction (``/proc`` scan, ``prctl``); a no-op elsewhere. +""" + +import ctypes +import os +import signal +import sys +import threading +import time +from typing import Optional + +from litellm._logging import verbose_proxy_logger + +QUERY_ENGINE_COMM_PREFIX = "query-engine" +REAPER_SCAN_INTERVAL_SECONDS = 5.0 +SIGTERM_GRACE_SECONDS = 10.0 +PR_SET_CHILD_SUBREAPER = 36 + + +def set_child_subreaper() -> bool: + """Mark this process as a child subreaper so orphaned descendants + reparent to it instead of PID 1. Best-effort: when it fails (or on + non-Linux) the reaper still covers the containerized case where the + supervisor already is PID 1.""" + if not sys.platform.startswith("linux"): + return False + try: + libc = ctypes.CDLL(None, use_errno=True) + result: int = libc.prctl( # pyright: ignore[reportAny] # ctypes types foreign calls as Any; default restype is c_int + PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0 + ) + return result == 0 + except (OSError, AttributeError): + return False + + +def _read_comm_and_ppid(pid: int, proc_root: str) -> Optional[tuple[str, int]]: + try: + with open(f"{proc_root}/{pid}/stat", encoding="ascii", errors="replace") as stat_file: + data = stat_file.read() + except (FileNotFoundError, ProcessLookupError, PermissionError, OSError): + return None + lparen = data.find("(") + rparen = data.rfind(")") + if lparen == -1 or rparen == -1 or rparen < lparen: + return None + comm = data[lparen + 1 : rparen] + fields = data[rparen + 2 :].split() + if len(fields) < 2: + return None + try: + ppid = int(fields[1]) + except ValueError: + return None + return comm, ppid + + +def list_orphaned_engine_pids(parent_pid: int, proc_root: str = "/proc") -> tuple[int, ...]: + """PIDs of direct children of ``parent_pid`` whose command name marks + them as Prisma query engines. In the supervisor these are always + adopted orphans: live engines are children of workers, not of the + supervisor.""" + try: + entries = os.listdir(proc_root) + except (FileNotFoundError, OSError): + return () + candidate_pids = (int(entry) for entry in entries if entry.isdigit()) + return tuple( + pid + for pid in candidate_pids + if (info := _read_comm_and_ppid(pid, proc_root)) is not None + and info[1] == parent_pid + and info[0].startswith(QUERY_ENGINE_COMM_PREFIX) + ) + + +def _try_reap(pid: int) -> bool: + try: + reaped_pid, _ = os.waitpid(pid, os.WNOHANG) + except ChildProcessError: + return True + except OSError: + return True + return reaped_pid == pid + + +def _send_signal(pid: int, signum: int) -> None: + try: + os.kill(pid, signum) + except (ProcessLookupError, PermissionError, OSError): + pass + + +def _await_reaped(pids: tuple[int, ...], timeout_seconds: float) -> tuple[int, ...]: + """Poll until every PID is reaped or the shared deadline passes. + Returns the PIDs still alive at the deadline.""" + deadline = time.monotonic() + timeout_seconds + remaining = pids + while remaining and time.monotonic() < deadline: + remaining = tuple(pid for pid in remaining if not _try_reap(pid)) + if remaining: + time.sleep(0.2) + return remaining + + +def terminate_and_reap(pid: int, grace_seconds: float = SIGTERM_GRACE_SECONDS) -> None: + """SIGTERM the orphaned engine, escalate to SIGKILL after the grace + period, and reap it so it does not linger as a zombie.""" + terminate_and_reap_all((pid,), grace_seconds=grace_seconds) + + +def terminate_and_reap_all( + pids: tuple[int, ...], + grace_seconds: float = SIGTERM_GRACE_SECONDS, +) -> None: + """Terminate a batch of orphaned engines concurrently: SIGTERM all of + them, share one grace period, SIGKILL the stragglers, and reap. The + shared deadline keeps cleanup time bounded when several workers die + at once instead of paying the grace period once per orphan.""" + for pid in pids: + verbose_proxy_logger.warning( + "Reaping orphaned prisma query-engine PID %s (its worker process exited without cleanup).", + pid, + ) + _send_signal(pid, signal.SIGTERM) + survivors = _await_reaped(pids, grace_seconds) + if not survivors: + return + for pid in survivors: + verbose_proxy_logger.warning( + "Orphaned prisma query-engine PID %s did not exit within %.1fs of SIGTERM; sending SIGKILL.", + pid, + grace_seconds, + ) + _send_signal(pid, signal.SIGKILL) + unkillable = _await_reaped(survivors, 5.0) + for pid in unkillable: + verbose_proxy_logger.error( + "Orphaned prisma query-engine PID %s survived SIGKILL; will retry on the next scan.", + pid, + ) + + +def reap_orphaned_engines(parent_pid: int, proc_root: str = "/proc") -> tuple[int, ...]: + """One scan-and-reap pass. Returns the PIDs it acted on.""" + orphaned_pids = list_orphaned_engine_pids(parent_pid, proc_root=proc_root) + if orphaned_pids: + terminate_and_reap_all(orphaned_pids) + return orphaned_pids + + +def _reaper_loop(parent_pid: int) -> None: + while True: + try: + reap_orphaned_engines(parent_pid) + except Exception as scan_error: # noqa: BLE001 # reaper thread must survive any scan failure + verbose_proxy_logger.debug("Orphaned query-engine scan failed: %s", scan_error) + time.sleep(REAPER_SCAN_INTERVAL_SECONDS) + + +REAPER_THREAD_NAME = "litellm-orphan-query-engine-reaper" + + +def start_query_engine_reaper() -> Optional[threading.Thread]: + """Start the reaper daemon thread in the supervisor process. + + Must only be called from a process that never hosts the proxy app + itself (uvicorn with ``workers > 1``, the gunicorn arbiter): with a + single in-process uvicorn worker the query engine is a legitimate + direct child and must not be touched. Idempotent: a reaper already + running in this process is returned instead of starting a second one. + """ + if not sys.platform.startswith("linux"): + return None + existing = next( + (thread for thread in threading.enumerate() if thread.name == REAPER_THREAD_NAME), + None, + ) + if existing is not None: + return existing + set_child_subreaper() + reaper_thread = threading.Thread( + target=_reaper_loop, + args=(os.getpid(),), + daemon=True, + name=REAPER_THREAD_NAME, + ) + reaper_thread.start() + verbose_proxy_logger.info( + "Started orphaned prisma query-engine reaper in supervisor process %s.", + os.getpid(), + ) + return reaper_thread diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 74ec0cc8700..9bed3657b20 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -15,6 +15,7 @@ from dotenv import load_dotenv import litellm from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY +from litellm.proxy.db.query_engine_reaper import start_query_engine_reaper from litellm.secret_managers.main import get_secret_bool if TYPE_CHECKING: @@ -495,6 +496,7 @@ class ProxyInitializationHelpers: gunicorn_options["certfile"] = ssl_certfile_path gunicorn_options["keyfile"] = ssl_keyfile_path + start_query_engine_reaper() StandaloneApplication(app=app, options=gunicorn_options).run() # Run gunicorn @staticmethod @@ -1261,6 +1263,8 @@ def run_server( if reload: ProxyInitializationHelpers._configure_dev_reload(uvicorn_args, config) + if num_workers > 1: + start_query_engine_reaper() uvicorn.run( **uvicorn_args, workers=num_workers, diff --git a/tests/test_litellm/proxy/db/test_query_engine_reaper.py b/tests/test_litellm/proxy/db/test_query_engine_reaper.py new file mode 100644 index 00000000000..efcecb4bc08 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_query_engine_reaper.py @@ -0,0 +1,244 @@ +import os +import signal +import subprocess +import sys +import time +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.proxy.db.query_engine_reaper import ( + REAPER_THREAD_NAME, + _read_comm_and_ppid, + _reaper_loop, + _send_signal, + _try_reap, + list_orphaned_engine_pids, + reap_orphaned_engines, + set_child_subreaper, + start_query_engine_reaper, + terminate_and_reap, + terminate_and_reap_all, +) + + +def _write_stat(proc_root, pid, comm, ppid): + pid_dir = proc_root / str(pid) + pid_dir.mkdir() + (pid_dir / "stat").write_text(f"{pid} ({comm}) S {ppid} {pid} {pid} 0 -1 4194304 100 0 0 0") + + +class TestReadCommAndPpid: + def test_parses_comm_and_ppid(self, tmp_path): + _write_stat(tmp_path, 137, "query-engine-de", 81) + assert _read_comm_and_ppid(137, str(tmp_path)) == ("query-engine-de", 81) + + def test_comm_containing_parens_and_spaces(self, tmp_path): + _write_stat(tmp_path, 42, "weird) (name", 1) + assert _read_comm_and_ppid(42, str(tmp_path)) == ("weird) (name", 1) + + def test_missing_pid_returns_none(self, tmp_path): + assert _read_comm_and_ppid(999, str(tmp_path)) is None + + def test_malformed_stat_returns_none(self, tmp_path): + pid_dir = tmp_path / "55" + pid_dir.mkdir() + (pid_dir / "stat").write_text("garbage with no parens") + assert _read_comm_and_ppid(55, str(tmp_path)) is None + + def test_truncated_fields_after_comm_returns_none(self, tmp_path): + pid_dir = tmp_path / "56" + pid_dir.mkdir() + (pid_dir / "stat").write_text("56 (proc) S") + assert _read_comm_and_ppid(56, str(tmp_path)) is None + + def test_non_numeric_ppid_returns_none(self, tmp_path): + pid_dir = tmp_path / "57" + pid_dir.mkdir() + (pid_dir / "stat").write_text("57 (proc) S notanint 57") + assert _read_comm_and_ppid(57, str(tmp_path)) is None + + +class TestListOrphanedEnginePids: + def test_finds_only_engine_children_of_parent(self, tmp_path): + _write_stat(tmp_path, 137, "query-engine-de", 1) + _write_stat(tmp_path, 138, "query-engine-de", 1) + _write_stat(tmp_path, 260, "python", 1) + _write_stat(tmp_path, 285, "query-engine-de", 260) + (tmp_path / "not-a-pid").mkdir() + + assert sorted(list_orphaned_engine_pids(1, proc_root=str(tmp_path))) == [137, 138] + + def test_no_matches_returns_empty(self, tmp_path): + _write_stat(tmp_path, 260, "python", 1) + assert list_orphaned_engine_pids(1, proc_root=str(tmp_path)) == () + + def test_missing_proc_root_returns_empty(self, tmp_path): + assert list_orphaned_engine_pids(1, proc_root=str(tmp_path / "absent")) == () + + +class TestSetChildSubreaper: + def test_matches_platform_capability(self): + result = set_child_subreaper() + if sys.platform.startswith("linux"): + assert result is True + else: + assert result is False + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX signals and waitpid") +class TestSignalHelpers: + def test_try_reap_true_for_non_child_pid(self): + assert _try_reap(1) is True + + def test_send_signal_swallows_missing_pid(self): + child = subprocess.Popen([sys.executable, "-c", "pass"]) + child.wait() + _send_signal(child.pid, signal.SIGTERM) + + +class TestReaperLoop: + def test_survives_scan_failure_and_continues(self): + calls = [] + + def flaky_scan(parent_pid, proc_root="/proc"): + calls.append(parent_pid) + if len(calls) == 1: + raise RuntimeError("scan blew up") + raise KeyboardInterrupt + + with ( + patch( + "litellm.proxy.db.query_engine_reaper.reap_orphaned_engines", + side_effect=flaky_scan, + ), + patch("litellm.proxy.db.query_engine_reaper.time.sleep"), + pytest.raises(KeyboardInterrupt), + ): + _reaper_loop(1234) + + assert calls == [1234, 1234] + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX signals and waitpid") +class TestTerminateAndReap: + def test_sigterm_terminates_and_reaps_child(self): + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(300)"]) + terminate_and_reap(child.pid, grace_seconds=10.0) + + with pytest.raises((ChildProcessError, OSError)): + os.waitpid(child.pid, os.WNOHANG) + child.returncode = -signal.SIGTERM + + def test_escalates_to_sigkill_when_sigterm_ignored(self): + child = subprocess.Popen( + [ + sys.executable, + "-c", + "import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(300)", + ] + ) + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + probe = subprocess.run( + [sys.executable, "-c", f"import os, signal; os.kill({child.pid}, 0)"], + capture_output=True, + ) + if probe.returncode == 0: + break + time.sleep(0.05) + time.sleep(0.3) + + terminate_and_reap(child.pid, grace_seconds=0.5) + + with pytest.raises((ChildProcessError, OSError)): + os.waitpid(child.pid, os.WNOHANG) + child.returncode = -signal.SIGKILL + + +class TestReapOrphanedEngines: + def test_terminates_each_orphan(self, tmp_path): + _write_stat(tmp_path, 137, "query-engine-de", 1) + _write_stat(tmp_path, 138, "query-engine-de", 1) + _write_stat(tmp_path, 285, "query-engine-de", 260) + + with patch("litellm.proxy.db.query_engine_reaper.terminate_and_reap_all") as mock_terminate: + acted_on = reap_orphaned_engines(1, proc_root=str(tmp_path)) + + assert sorted(acted_on) == [137, 138] + assert sorted(mock_terminate.call_args.args[0]) == [137, 138] + + def test_no_orphans_no_kills(self, tmp_path): + _write_stat(tmp_path, 285, "query-engine-de", 260) + + with patch("litellm.proxy.db.query_engine_reaper.terminate_and_reap_all") as mock_terminate: + assert reap_orphaned_engines(1, proc_root=str(tmp_path)) == () + + mock_terminate.assert_not_called() + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX signals and waitpid") +class TestTerminateAndReapAll: + def test_batch_shares_one_grace_period(self): + children = [ + subprocess.Popen( + [ + sys.executable, + "-c", + "import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(300)", + ] + ) + for _ in range(3) + ] + time.sleep(0.5) + + start = time.monotonic() + terminate_and_reap_all(tuple(child.pid for child in children), grace_seconds=1.0) + elapsed = time.monotonic() - start + + assert elapsed < 3.0 + for child in children: + with pytest.raises((ChildProcessError, OSError)): + os.waitpid(child.pid, os.WNOHANG) + child.returncode = -signal.SIGKILL + + +class TestStartQueryEngineReaper: + def test_noop_on_non_linux(self): + with patch("litellm.proxy.db.query_engine_reaper.sys.platform", "darwin"): + assert start_query_engine_reaper() is None + + def test_starts_daemon_thread_on_linux(self): + with ( + patch("litellm.proxy.db.query_engine_reaper.sys.platform", "linux"), + patch( + "litellm.proxy.db.query_engine_reaper.threading.enumerate", + return_value=[], + ), + patch("litellm.proxy.db.query_engine_reaper.set_child_subreaper") as mock_subreaper, + patch("litellm.proxy.db.query_engine_reaper.threading.Thread") as mock_thread_cls, + ): + thread = start_query_engine_reaper() + + mock_subreaper.assert_called_once() + mock_thread_cls.assert_called_once() + assert mock_thread_cls.call_args.kwargs["daemon"] is True + assert mock_thread_cls.call_args.kwargs["args"] == (os.getpid(),) + mock_thread_cls.return_value.start.assert_called_once() + assert thread is mock_thread_cls.return_value + + def test_second_call_returns_existing_thread(self): + existing = MagicMock() + existing.name = REAPER_THREAD_NAME + with ( + patch("litellm.proxy.db.query_engine_reaper.sys.platform", "linux"), + patch( + "litellm.proxy.db.query_engine_reaper.threading.enumerate", + return_value=[existing], + ), + patch("litellm.proxy.db.query_engine_reaper.threading.Thread") as mock_thread_cls, + ): + thread = start_query_engine_reaper() + + assert thread is existing + mock_thread_cls.assert_not_called() diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 88dbec4020f..c029ca18307 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1557,6 +1557,85 @@ class TestProxyInitializationHelpers: mock_uvicorn_run.assert_called_once() +class TestQueryEngineReaperWiring: + def _invoke_run_server(self, args): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + }, + ), + patch("uvicorn.run") as mock_uvicorn_run, + patch( + "litellm.proxy.proxy_cli.start_query_engine_reaper" + ) as mock_start_reaper, + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + result = runner.invoke(run_server, args) + return result, mock_uvicorn_run, mock_start_reaper + + def test_multi_worker_uvicorn_starts_reaper(self): + result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server( + ["--local", "--num_workers", "2"] + ) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" + mock_uvicorn_run.assert_called_once() + mock_start_reaper.assert_called_once() + + def test_single_worker_uvicorn_does_not_start_reaper(self): + result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server( + ["--local", "--num_workers", "1"] + ) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" + mock_uvicorn_run.assert_called_once() + mock_start_reaper.assert_not_called() + + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") + def test_gunicorn_arbiter_starts_reaper(self): + pytest.importorskip("gunicorn") + + with ( + patch("gunicorn.app.base.BaseApplication.run"), + patch( + "litellm.proxy.proxy_cli.start_query_engine_reaper" + ) as mock_start_reaper, + ): + ProxyInitializationHelpers._run_gunicorn_server( + host="127.0.0.1", + port=4010, + app=MagicMock(), + num_workers=1, + ssl_certfile_path=None, + ssl_keyfile_path=None, + ) + + mock_start_reaper.assert_called_once() + + class TestRunServerDbSetup: """Tests for run_server's prisma setup_database behavior.""" From 03e7dc4ac5d8da3732a659a0f2ab6049481d41f9 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 13:21:29 -0700 Subject: [PATCH 56/62] build(deps): bump uvicorn lock to 0.51.0 so worker health-check and jitter flags take effect (#33574) --- tests/test_litellm/proxy/test_proxy_cli.py | 13 +++ uv.lock | 122 ++++++++++----------- 2 files changed, 74 insertions(+), 61 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index c029ca18307..6b0c0dba40f 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1,3 +1,4 @@ +import inspect import os import sys from pathlib import Path @@ -15,6 +16,8 @@ sys.path.insert( import builtins import types +import uvicorn + from litellm.proxy.proxy_cli import ProxyInitializationHelpers @@ -135,6 +138,16 @@ class TestProxyInitializationHelpers: ) assert args["timeout_worker_healthcheck"] == 15 + def test_installed_uvicorn_supports_worker_flags(self): + params = inspect.signature(uvicorn.Config.__init__).parameters + assert "timeout_worker_healthcheck" in params + assert "limit_max_requests_jitter" in params + + args = ProxyInitializationHelpers._get_default_unvicorn_init_args( + "localhost", 8000, timeout_worker_healthcheck=30 + ) + assert args["timeout_worker_healthcheck"] == 30 + def test_get_reload_options_no_config_still_watches_env(self): opts = ProxyInitializationHelpers._get_reload_options(None) assert opts["reload"] is True diff --git a/uv.lock b/uv.lock index 22e3a3feece..5a76b6a4531 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-13T03:38:04.421387Z" +exclude-newer = "2026-07-13T19:39:19.414377Z" exclude-newer-span = "P3D" [manifest] @@ -222,9 +222,9 @@ name = "aiologic" version = "0.17.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "sniffio" }, + { name = "sniffio", marker = "python_full_version < '3.13'" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, - { name = "wrapt" }, + { name = "wrapt", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/53/a7/809482759f40079f4c4328c7318bf569ae25d457f5017aad30a1b9aafedc/aiologic-0.17.0.tar.gz", hash = "sha256:65aa058e858c94cd208badb188e7f00b54dcabb3ba85b34f794db98074d108b9", size = 251625, upload-time = "2026-06-14T12:24:35.367Z" } wheels = [ @@ -516,14 +516,14 @@ name = "aurelio-sdk" version = "0.0.19" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiofiles" }, - { name = "aiohttp" }, - { name = "colorlog" }, - { name = "pydantic" }, - { name = "python-dotenv" }, - { name = "requests" }, - { name = "requests-toolbelt" }, - { name = "tornado" }, + { name = "aiofiles", marker = "python_full_version < '3.14'" }, + { name = "aiohttp", marker = "python_full_version < '3.14'" }, + { name = "colorlog", marker = "python_full_version < '3.14'" }, + { name = "pydantic", marker = "python_full_version < '3.14'" }, + { name = "python-dotenv", marker = "python_full_version < '3.14'" }, + { name = "requests", marker = "python_full_version < '3.14'" }, + { name = "requests-toolbelt", marker = "python_full_version < '3.14'" }, + { name = "tornado", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/27/0e/c2e369ad173fb3d76448e46d10beb3dcc53388318933ddf8169a3f21a810/aurelio_sdk-0.0.19.tar.gz", hash = "sha256:14107e7440ff2efd0b4a08c52fb595e7680bd4bc973a0ddfb3b64157c6666b91", size = 15258, upload-time = "2025-03-24T14:37:32.203Z" } wheels = [ @@ -1047,7 +1047,7 @@ name = "coloredlogs" version = "15.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "humanfriendly" }, + { name = "humanfriendly", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/c7/eed8f27100517e8c0e6b923d5f0845d0cb99763da6fdee00478f91db7325/coloredlogs-15.0.1.tar.gz", hash = "sha256:7c991aa71a4577af2f82600d8f8f3a89f936baeaf9b50a9c197da014e5bf16b0", size = 278520, upload-time = "2021-06-11T10:22:45.202Z" } wheels = [ @@ -1059,7 +1059,7 @@ name = "colorlog" version = "6.10.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "python_full_version < '3.14' and sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a2/61/f083b5ac52e505dfc1c624eafbf8c7589a0d7f32daa398d2e7590efa5fda/colorlog-6.10.1.tar.gz", hash = "sha256:eb4ae5cb65fe7fec7773c2306061a8e63e02efc2c72eba9d27b0fa23c94f1321", size = 17162, upload-time = "2025-10-16T16:14:11.978Z" } wheels = [ @@ -1420,7 +1420,7 @@ name = "culsans" version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiologic" }, + { name = "aiologic", marker = "python_full_version < '3.13'" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/d9/e3/49afa1bc180e0d28008ec6bcdf82a4072d1c7a41032b5b759b60814ca4b0/culsans-0.11.0.tar.gz", hash = "sha256:0b43d0d05dce6106293d114c86e3fb4bfc63088cfe8ff08ed3fe36891447fe33", size = 107546, upload-time = "2025-12-31T23:15:38.196Z" } @@ -2966,7 +2966,7 @@ name = "humanfriendly" version = "10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pyreadline3", marker = "sys_platform == 'win32'" }, + { name = "pyreadline3", marker = "python_full_version < '3.14' and sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/3f/2c29224acb2e2df4d2046e4c73ee2662023c58ff5b113c4c1adac0886c43/humanfriendly-10.0.tar.gz", hash = "sha256:6b0b831ce8f15f7300721aa49829fc4e83921a9a301cc7f606be6686a2288ddc", size = 360702, upload-time = "2021-09-17T21:40:43.31Z" } wheels = [ @@ -3293,10 +3293,10 @@ name = "jsonschema-path" version = "0.3.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pathable" }, - { name = "pyyaml" }, - { name = "referencing" }, - { name = "requests" }, + { name = "pathable", marker = "python_full_version < '3.14'" }, + { name = "pyyaml", marker = "python_full_version < '3.14'" }, + { name = "referencing", marker = "python_full_version < '3.14'" }, + { name = "requests", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6e/45/41ebc679c2a4fced6a722f624c18d658dee42612b83ea24c1caf7c0eb3a8/jsonschema_path-0.3.4.tar.gz", hash = "sha256:8365356039f16cc65fddffafda5f58766e34bebab7d6d105616ab52bc4297001", size = 11159, upload-time = "2025-01-24T14:33:16.547Z" } wheels = [ @@ -4481,7 +4481,7 @@ version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/fd/15/76f86faa0902836cc133939732f7611ace68cf54148487a99c539c272dc8/ml_dtypes-0.4.1.tar.gz", hash = "sha256:fad5f2de464fd09127e49b7fd1252b9006fb43d2edc1ff112d390c324af5ca7a", size = 692594, upload-time = "2024-09-13T19:07:11.624Z" } wheels = [ @@ -4980,14 +4980,14 @@ name = "openapi-core" version = "0.22.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "isodate" }, - { name = "jsonschema" }, - { name = "jsonschema-path" }, - { name = "more-itertools" }, - { name = "openapi-schema-validator" }, - { name = "openapi-spec-validator" }, - { name = "typing-extensions" }, - { name = "werkzeug" }, + { name = "isodate", marker = "python_full_version < '3.14'" }, + { name = "jsonschema", marker = "python_full_version < '3.14'" }, + { name = "jsonschema-path", marker = "python_full_version < '3.14'" }, + { name = "more-itertools", marker = "python_full_version < '3.14'" }, + { name = "openapi-schema-validator", marker = "python_full_version < '3.14'" }, + { name = "openapi-spec-validator", marker = "python_full_version < '3.14'" }, + { name = "typing-extensions", marker = "python_full_version < '3.14'" }, + { name = "werkzeug", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/fd/65/ee75f25b9459a02df6f713f8ffde5dacb57b8b4e45145cde4cab28b5abba/openapi_core-0.22.0.tar.gz", hash = "sha256:b30490dfa74e3aac2276105525590135212352f5dd7e5acf8f62f6a89ed6f2d0", size = 109242, upload-time = "2025-12-22T19:19:49.608Z" } wheels = [ @@ -4999,9 +4999,9 @@ name = "openapi-schema-validator" version = "0.6.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "jsonschema" }, - { name = "jsonschema-specifications" }, - { name = "rfc3339-validator" }, + { name = "jsonschema", marker = "python_full_version < '3.14'" }, + { name = "jsonschema-specifications", marker = "python_full_version < '3.14'" }, + { name = "rfc3339-validator", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/8b/f3/5507ad3325169347cd8ced61c232ff3df70e2b250c49f0fe140edb4973c6/openapi_schema_validator-0.6.3.tar.gz", hash = "sha256:f37bace4fc2a5d96692f4f8b31dc0f8d7400fd04f3a937798eaf880d425de6ee", size = 11550, upload-time = "2025-01-10T18:08:22.268Z" } wheels = [ @@ -5013,10 +5013,10 @@ name = "openapi-spec-validator" version = "0.7.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "jsonschema" }, - { name = "jsonschema-path" }, - { name = "lazy-object-proxy" }, - { name = "openapi-schema-validator" }, + { name = "jsonschema", marker = "python_full_version < '3.14'" }, + { name = "jsonschema-path", marker = "python_full_version < '3.14'" }, + { name = "lazy-object-proxy", marker = "python_full_version < '3.14'" }, + { name = "openapi-schema-validator", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/82/af/fe2d7618d6eae6fb3a82766a44ed87cd8d6d82b4564ed1c7cfb0f6378e91/openapi_spec_validator-0.7.2.tar.gz", hash = "sha256:cc029309b5c5dbc7859df0372d55e9d1ff43e96d678b9ba087f7c56fc586f734", size = 36855, upload-time = "2025-06-07T14:48:56.299Z" } wheels = [ @@ -7163,16 +7163,16 @@ name = "redisvl" version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "coloredlogs" }, - { name = "ml-dtypes" }, + { name = "coloredlogs", marker = "python_full_version < '3.14'" }, + { name = "ml-dtypes", marker = "python_full_version < '3.14'" }, { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "pydantic" }, - { name = "python-ulid" }, - { name = "pyyaml" }, - { name = "redis" }, - { name = "tabulate" }, - { name = "tenacity" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, + { name = "pydantic", marker = "python_full_version < '3.14'" }, + { name = "python-ulid", marker = "python_full_version < '3.14'" }, + { name = "pyyaml", marker = "python_full_version < '3.14'" }, + { name = "redis", marker = "python_full_version < '3.14'" }, + { name = "tabulate", marker = "python_full_version < '3.14'" }, + { name = "tenacity", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/21/33/ab14865a0b2a31b1d003c29e7e8ea3a7a2f2c8ecb24e58e58d606e1f031b/redisvl-0.4.1.tar.gz", hash = "sha256:fd6a36426ba94792c0efca20915c31232d4ee3cc58eb23794a62c142696401e6", size = 77688, upload-time = "2025-02-21T22:51:41.389Z" } wheels = [ @@ -7406,7 +7406,7 @@ name = "rfc3339-validator" version = "0.1.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "six" }, + { name = "six", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/28/ea/a9387748e2d111c3c2b275ba970b735e04e15cdb1eb30693b6b5708c4dbd/rfc3339_validator-0.1.4.tar.gz", hash = "sha256:138a2abdf93304ad60530167e51d2dfb9549521a836871b88d7f4695d0022f6b", size = 5513, upload-time = "2021-05-12T16:37:54.178Z" } wheels = [ @@ -7855,20 +7855,20 @@ name = "semantic-router" version = "0.1.15" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiohttp" }, - { name = "aurelio-sdk" }, - { name = "colorama" }, - { name = "colorlog" }, - { name = "litellm" }, + { name = "aiohttp", marker = "python_full_version < '3.14'" }, + { name = "aurelio-sdk", marker = "python_full_version < '3.14'" }, + { name = "colorama", marker = "python_full_version < '3.14'" }, + { name = "colorlog", marker = "python_full_version < '3.14'" }, + { name = "litellm", marker = "python_full_version < '3.14'" }, { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "openai" }, - { name = "pydantic" }, - { name = "pyyaml" }, - { name = "regex" }, - { name = "tiktoken" }, - { name = "tornado" }, - { name = "urllib3" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, + { name = "openai", marker = "python_full_version < '3.14'" }, + { name = "pydantic", marker = "python_full_version < '3.14'" }, + { name = "pyyaml", marker = "python_full_version < '3.14'" }, + { name = "regex", marker = "python_full_version < '3.14'" }, + { name = "tiktoken", marker = "python_full_version < '3.14'" }, + { name = "tornado", marker = "python_full_version < '3.14'" }, + { name = "urllib3", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/dc/a9/1a689e916e8b280f1fd8fb335cc059be626a22fe4533baa045d32fcd6de5/semantic_router-0.1.15.tar.gz", hash = "sha256:328256ddc3c2b713101ec69561d6585aecbf1198ea3461e1486289d8c3a35288", size = 95605, upload-time = "2026-05-23T12:58:15.444Z" } wheels = [ @@ -8958,16 +8958,16 @@ wheels = [ [[package]] name = "uvicorn" -version = "0.33.0" +version = "0.51.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, { name = "h11" }, { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/cb/81/a083ae41716b00df56d45d4b5f6ca8e90fc233a62e6c04ab3ad3c476b6c4/uvicorn-0.33.0.tar.gz", hash = "sha256:3577119f82b7091cf4d3d4177bfda0bae4723ed92ab1439e8d779de880c9cc59", size = 76590, upload-time = "2024-12-14T11:14:46.526Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/65/b7c6c443ccc58678c91e1e973bbe2a878591538655d6e1d47f24ba1c51f3/uvicorn-0.51.0.tar.gz", hash = "sha256:f6f4b69b657c312f516dd2d268ab9ae6f254b11e4bac504f37b2ab58b24dd0b0", size = 94412, upload-time = "2026-07-08T10:59:05.962Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/98/79/2e2620337ef1e4ef7a058b351603b765f59ac28e6e3ac7c5e7cdee9ea1ab/uvicorn-0.33.0-py3-none-any.whl", hash = "sha256:2c30de4aeea83661a520abab179b24084a0019c0c1bbe137e5409f741cbde5f8", size = 62297, upload-time = "2024-12-14T11:14:43.408Z" }, + { url = "https://files.pythonhosted.org/packages/45/ec/dbb7e5a6b91f86bfb9eb7d2988a2730907b6a729875b949c7f022e8b88fa/uvicorn-0.51.0-py3-none-any.whl", hash = "sha256:5d38af6cd620f2ae3849fb44fd4879e0890aa1febe8d47eb355fb45d93fe6a5b", size = 73219, upload-time = "2026-07-08T10:59:04.44Z" }, ] [[package]] From 51305536bff4e7b6445f76a5b880bf02afafdfc5 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 13:34:43 -0700 Subject: [PATCH 57/62] fix(proxy): coerce default_internal_user_params.max_budget to float on config load (#32434) * fix(proxy): coerce default_internal_user_params.max_budget to float on config load * fix(proxy): log coerced default_internal_user_params and cover absent max_budget in tests --- litellm/proxy/proxy_server.py | 9 +++ tests/test_litellm/proxy/test_proxy_server.py | 78 +++++++++++++++++++ 2 files changed, 87 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6661474d215..7e6ca5a0108 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4425,6 +4425,15 @@ class ProxyConfig: litellm.default_max_internal_user_budget = float(value) if litellm.max_internal_user_budget is None: litellm.max_internal_user_budget = litellm.default_max_internal_user_budget + elif key == "default_internal_user_params" and isinstance(value, dict): + litellm.default_internal_user_params = ( + {**value, "max_budget": float(value["max_budget"])} + if value.get("max_budget") is not None + else value + ) + verbose_proxy_logger.debug( + f"{blue_color_code} setting litellm.{key}={_redact_general_setting_value(key, litellm.default_internal_user_params, is_full_admin=False)}{reset_color_code}" + ) elif key == "custom_provider_map": from litellm.utils import custom_llm_setup diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 2f0924e9192..079b844638c 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2588,6 +2588,84 @@ async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeyp litellm.max_budget = original_max_budget +@pytest.mark.asyncio +async def test_load_config_default_internal_user_params_max_budget_scientific_notation(tmp_path): + """ + Helm's toYaml renders large floats in scientific notation without a + decimal mantissa (e.g. 1e+09), which PyYAML parses as a string. + load_config must coerce default_internal_user_params.max_budget to + float, otherwise every consumer of the raw dict (/user/new, SSO, + SCIM user creation) passes the string to Prisma, which rejects it + since max_budget must be Float or Null. Keys outside the coercion + (including ones not on DefaultInternalUserParams, like + auto_create_key) must pass through unchanged. + """ + from litellm.proxy.proxy_server import ProxyConfig + + config_file = tmp_path / "config.yaml" + config_file.write_text( + "model_list: []\n" + "litellm_settings:\n" + " default_internal_user_params:\n" + " user_role: internal_user\n" + " max_budget: 1e+09\n" + " budget_duration: 30d\n" + " auto_create_key: false\n" + ) + + original_params = litellm.default_internal_user_params + try: + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + assert litellm.default_internal_user_params == { + "user_role": "internal_user", + "max_budget": 1000000000.0, + "budget_duration": "30d", + "auto_create_key": False, + } + assert isinstance(litellm.default_internal_user_params["max_budget"], float) + finally: + litellm.default_internal_user_params = original_params + + +@pytest.mark.asyncio +async def test_load_config_default_internal_user_params_without_max_budget(tmp_path): + """ + default_internal_user_params without max_budget (or with an explicit + null) must be stored as-is and not gain a max_budget key. + """ + from litellm.proxy.proxy_server import ProxyConfig + + absent_config_file = tmp_path / "absent_config.yaml" + absent_config_file.write_text( + "model_list: []\n" + "litellm_settings:\n" + " default_internal_user_params:\n" + " user_role: internal_user\n" + ) + + null_config_file = tmp_path / "null_config.yaml" + null_config_file.write_text( + "model_list: []\n" + "litellm_settings:\n" + " default_internal_user_params:\n" + " user_role: internal_user\n" + " max_budget: null\n" + ) + + original_params = litellm.default_internal_user_params + try: + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(absent_config_file)) + assert litellm.default_internal_user_params == {"user_role": "internal_user"} + + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(null_config_file)) + assert litellm.default_internal_user_params == { + "user_role": "internal_user", + "max_budget": None, + } + finally: + litellm.default_internal_user_params = original_params + + @pytest.mark.asyncio async def test_load_config_user_url_validation_handles_null_and_string_false(tmp_path, monkeypatch): from litellm.proxy.proxy_server import ProxyConfig From c6778b79c3919ee4b947c5194ac9262e454f6995 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 13:36:03 -0700 Subject: [PATCH 58/62] fix(router): honor per-request routing_strategy from key/team router_settings (#33429) * fix(router): honor per-request routing_strategy from key/team router_settings Key and team router_settings.routing_strategy was stored and shown in the UI but never forwarded to the shared Router, so the global strategy always won. Forward it through router_settings_override and resolve it in _get_routing_context: a validated per-request strategy takes precedence over routing groups and the top-level strategy, with lazily built cached selectors for strategies that need one. Unknown or unsupported strategy values are ignored with a warning instead of failing the request, and routing_strategy is registered in all_litellm_params so it is stripped before the provider call. * fix(router): sweep override selectors on strategy re-init and cover coverage-gate helpers routing_strategy_init now unregisters cached per-request override selectors so a later update_settings strategy change cannot leave a zombie selector receiving callback events. Adds direct tests for the two new helpers so the router code coverage gate passes. * docs(team): document mcp_rpm_limit in update_team docstring The documentation CI job walks management_endpoints and requires every UpdateTeamRequest field to appear in the update_team docstring; mcp_rpm_limit was added to the model without a docstring line, failing the job on unrelated PRs depending on walk order. Regenerates schema.d.ts since the docstring feeds the OpenAPI spec. --- litellm/proxy/route_llm_request.py | 1 + litellm/router.py | 80 +++++++++++++-- litellm/types/utils.py | 1 + .../proxy/test_route_llm_request.py | 4 +- .../test_router_routing_groups.py | 97 +++++++++++++++++++ 5 files changed, 173 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 94871ff072d..0ca2a75990b 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -409,6 +409,7 @@ async def route_request( "num_retries", "timeout", "model_group_retry_policy", + "routing_strategy", ] # Merge override settings into data (only if not already set in request) diff --git a/litellm/router.py b/litellm/router.py index 809aef695a7..9b869d741ba 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -631,6 +631,8 @@ class Router: routing_strategy_args=routing_strategy_args, ) self._init_routing_groups(self._routing_groups_input) + self._override_selectors: dict[str, Any] = {} + self._override_selectors_lock = threading.Lock() self.access_groups = None ## USAGE TRACKING ## if isinstance(litellm._async_success_callback, list): @@ -902,7 +904,9 @@ class Router: self._unregister_router_selectors( [getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()] + + list(getattr(self, "_override_selectors", {}).values()) ) + self._override_selectors = {} self.leastbusy_logger: Optional[LeastBusyLoggingHandler] = None self.lowesttpm_logger: Optional[LowestTPMLoggingHandler] = None @@ -992,12 +996,67 @@ class Router: {strategy_value: group_selector} if group_selector is not None else {} ) - def _get_routing_context(self, model: str) -> Tuple[Optional[str], Optional[Any]]: + _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) + + def _get_request_routing_strategy_override(self, request_kwargs: Optional[dict]) -> Optional[str]: + """ + Reads a per-request `routing_strategy` override (forwarded by the proxy + from key/team `router_settings`) out of the request kwargs. + + Only strategies with a per-request-capable selector are honored; + anything else (unknown strings, `lar1`, `provider-budget-routing`) is + ignored with a warning so a bad value stored on a key or team can + never take down that caller's traffic. + """ + if not request_kwargs: + return None + raw_strategy = request_kwargs.get("routing_strategy") + if raw_strategy is None: + return None + strategy = self._normalize_strategy(raw_strategy) if isinstance(raw_strategy, (str, RoutingStrategy)) else None + if not isinstance(strategy, str) or strategy not in self._OVERRIDABLE_ROUTING_STRATEGIES: + verbose_router_logger.warning( + "Ignoring per-request routing_strategy override '%s'; supported overrides: %s.", + raw_strategy, + sorted(self._OVERRIDABLE_ROUTING_STRATEGIES), + ) + return None + return strategy + + def _get_override_strategy_selector(self, strategy: str) -> Optional[Any]: + """ + Returns the selector for a per-request strategy override. + + Reuses the default group's selector when the override matches the + router's configured strategy (so shared state keeps accumulating in + one place); otherwise lazily builds one selector per strategy and + caches it for the router's lifetime so its usage/latency state + persists across requests. + """ + if strategy == self._normalize_strategy(self.routing_strategy): + attr = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy) + return getattr(self, attr, None) if attr is not None else None + with self._override_selectors_lock: + if strategy not in self._override_selectors: + self._override_selectors[strategy] = self._build_strategy_selector( + strategy=strategy, + routing_strategy_args={}, + ) + return self._override_selectors[strategy] + + def _get_routing_context( + self, model: str, request_kwargs: Optional[dict] = None + ) -> tuple[Optional[str], Optional[Any]]: """ Resolves the routing strategy and selector to use for the given model. - Every model belongs to exactly one group: an explicit entry from - `routing_groups`, or the implicit `"default"` group driven by the + A per-request `routing_strategy` in `request_kwargs` (forwarded by the + proxy from key/team `router_settings`) takes precedence over both the + model's routing group and the router's top-level strategy, since it is + the most specific expression of caller intent. + + Otherwise every model belongs to exactly one group: an explicit entry + from `routing_groups`, or the implicit `"default"` group driven by the router's top-level `routing_strategy` / `routing_strategy_args`. `self.routing_strategy` may be either a string or a `RoutingStrategy` @@ -1005,6 +1064,11 @@ class Router: string here. Downstream call sites and `_select_deployment_*` arms compare against string literals. """ + override = self._get_request_routing_strategy_override(request_kwargs) + if override is not None: + verbose_router_logger.debug("routing_group=request-override model=%s strategy=%s", model, override) + return override, self._get_override_strategy_selector(override) + group_name = self._model_to_group.get(model) if group_name is None: strategy = self._normalize_strategy(self.routing_strategy) @@ -5945,7 +6009,7 @@ class Router: input_kwargs: dict, ) -> Optional[Any]: """Same-model-group retry after a failed deployment; returns None if not applicable.""" - strategy, _ = self._get_routing_context(original_model_group) + strategy, _ = self._get_routing_context(original_model_group, kwargs) if strategy != "simple-shuffle": return None @@ -10473,7 +10537,7 @@ class Router: # Resolve the strategy and logger AFTER the pre-routing hook, since # the hook can replace `model` and routing-group lookup must key # off the final model name. - strategy, strategy_selector = self._get_routing_context(model) + strategy, strategy_selector = self._get_routing_context(model, request_kwargs) healthy_deployments = await self.async_get_healthy_deployments( model=model, @@ -10617,7 +10681,7 @@ class Router: # 5. Apply load balancing strategy start_time = time.perf_counter() - strategy, strategy_selector = self._get_routing_context(model) + strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if strategy == "simple-shuffle": return simple_shuffle( llm_router_instance=self, @@ -10904,7 +10968,7 @@ class Router: cooldown_list=_cooldown_list, ) - strategy, strategy_selector = self._get_routing_context(model) + strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if strategy == "simple-shuffle": # if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm ############## Check 'weight' param set for weighted pick ################# @@ -11044,7 +11108,7 @@ class Router: ) # 6. Apply load balancing strategy - strategy, strategy_selector = self._get_routing_context(model) + strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if strategy == "simple-shuffle": return simple_shuffle( llm_router_instance=self, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index db39d58a925..8c4739fd417 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3150,6 +3150,7 @@ all_litellm_params = ( "use_client", "id", "fallbacks", + "routing_strategy", "azure", "headers", "model_list", diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 303871e3981..c96b86f84b0 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -382,8 +382,8 @@ async def test_route_request_with_router_settings_override(): "num_retries": 5, "timeout": 30, "model_group_retry_policy": {"gpt-3.5-turbo": {"RateLimitErrorRetries": 3}}, - # These settings should be ignored (not in per_request_settings list) "routing_strategy": "least-busy", + # This setting should be ignored (not in per_request_settings list) "model_group_alias": {"alias": "real_model"}, }, } @@ -400,8 +400,8 @@ async def test_route_request_with_router_settings_override(): assert call_kwargs["num_retries"] == 5 assert call_kwargs["timeout"] == 30 assert call_kwargs["model_group_retry_policy"] == {"gpt-3.5-turbo": {"RateLimitErrorRetries": 3}} + assert call_kwargs["routing_strategy"] == "least-busy" # Verify unsupported settings were NOT merged - assert "routing_strategy" not in call_kwargs assert "model_group_alias" not in call_kwargs # Verify router_settings_override was removed from data assert "router_settings_override" not in call_kwargs diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 7368c72836d..b8dcdacd8a3 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -629,3 +629,100 @@ async def test_async_dispatch_falls_back_to_sync_for_usage_based_routing_v1(): ) assert v1_spy.called, "async dispatch must route v1 strategy through sync method" + + +def test_request_routing_strategy_override_beats_top_level(): + router = _build_router(routing_strategy="least-busy") + strategy, selector = router._get_routing_context( + "other-model", {"routing_strategy": "simple-shuffle"} + ) + assert strategy == "simple-shuffle" + assert selector is None + + +def test_request_routing_strategy_override_beats_explicit_group(): + router = _build_router( + routing_strategy="simple-shuffle", + routing_groups=[ + { + "group_name": "fast", + "models": ["filtered-model"], + "routing_strategy": "latency-based-routing", + } + ], + ) + strategy, _ = router._get_routing_context( + "filtered-model", {"routing_strategy": "least-busy"} + ) + assert strategy == "least-busy" + + +def test_request_routing_strategy_override_builds_and_caches_selector(): + router = _build_router(routing_strategy="simple-shuffle") + strategy, selector = router._get_routing_context( + "other-model", {"routing_strategy": "latency-based-routing"} + ) + assert strategy == "latency-based-routing" + assert selector is not None + _, selector_again = router._get_routing_context( + "other-model", {"routing_strategy": "latency-based-routing"} + ) + assert selector_again is selector + + +def test_request_routing_strategy_override_matching_global_reuses_default_selector(): + router = _build_router(routing_strategy="least-busy") + _, selector = router._get_routing_context( + "other-model", {"routing_strategy": "least-busy"} + ) + assert selector is router.leastbusy_logger + assert router._override_selectors == {} + + +def test_invalid_request_routing_strategy_override_falls_back(): + router = _build_router(routing_strategy="least-busy") + strategy, selector = router._get_routing_context( + "other-model", {"routing_strategy": "not-a-real-strategy"} + ) + assert strategy == "least-busy" + assert selector is router.leastbusy_logger + + +def test_no_override_key_keeps_existing_behavior(): + router = _build_router(routing_strategy="least-busy") + strategy, _ = router._get_routing_context("other-model", {"messages": []}) + assert strategy == "least-busy" + strategy_none_kwargs, _ = router._get_routing_context("other-model", None) + assert strategy_none_kwargs == "least-busy" + + +def test_request_routing_strategy_override_helper_validates_directly(): + router = _build_router(routing_strategy="least-busy") + assert router._get_request_routing_strategy_override({"routing_strategy": "simple-shuffle"}) == "simple-shuffle" + assert router._get_request_routing_strategy_override({"routing_strategy": RoutingStrategy.LEAST_BUSY}) == "least-busy" + assert router._get_request_routing_strategy_override({"routing_strategy": "lar1"}) is None + assert router._get_request_routing_strategy_override({"routing_strategy": {"bad": "type"}}) is None + assert router._get_request_routing_strategy_override({}) is None + assert router._get_request_routing_strategy_override(None) is None + + +def test_override_strategy_selector_helper_builds_per_strategy(): + router = _build_router(routing_strategy="least-busy") + latency_selector = router._get_override_strategy_selector("latency-based-routing") + assert latency_selector is not None + assert router._get_override_strategy_selector("latency-based-routing") is latency_selector + assert router._get_override_strategy_selector("least-busy") is router.leastbusy_logger + assert router._get_override_strategy_selector("simple-shuffle") is None + + +def test_strategy_reinit_unregisters_override_selectors(): + router = _build_router(routing_strategy="least-busy") + override_selector = router._get_override_strategy_selector("latency-based-routing") + assert override_selector is not None + assert any(id(cb) == id(override_selector) for cb in litellm.callbacks) + + router.update_settings(routing_strategy="latency-based-routing") + + assert router._override_selectors == {} + assert not any(id(cb) == id(override_selector) for cb in litellm.callbacks) + assert router._get_override_strategy_selector("latency-based-routing") is router.lowestlatency_logger From 903219a8b10e805ea2108ae022fcf3c70a01bd9f Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 13:38:01 -0700 Subject: [PATCH 59/62] fix(redis): honor ssl value instead of key presence when building async connection pool (#32590) * fix(redis): honor ssl value instead of key presence when building async connection pool * ci: rerun codspeed after cross-runtime-environment flake --- litellm/_redis.py | 6 ++---- tests/test_litellm/test_redis.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/litellm/_redis.py b/litellm/_redis.py index 0b91cdabffc..fe5c5cdabe9 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -688,10 +688,8 @@ def get_redis_connection_pool( elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"): redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account) - connection_class = async_redis.Connection - if redis_kwargs.pop("ssl", False): - connection_class = async_redis.SSLConnection - redis_kwargs["connection_class"] = connection_class + if redis_kwargs.pop("ssl", None): + redis_kwargs["connection_class"] = async_redis.SSLConnection return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs) diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 4b9f13c340b..0818237655d 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -724,3 +724,16 @@ def test_connection_pool_without_ssl_kwarg_uses_plain_connection(monkeypatch): call_kwargs = mock_pool.call_args.kwargs assert call_kwargs.get("connection_class") is not async_redis.SSLConnection assert "ssl" not in call_kwargs + + +def test_connection_pool_env_redis_ssl_false_uses_plain_connection(monkeypatch): + """REDIS_SSL=false from the environment must not select SSLConnection.""" + monkeypatch.delenv("REDIS_URL", raising=False) + monkeypatch.delenv("REDIS_CLUSTER_NODES", raising=False) + monkeypatch.setenv("REDIS_SSL", "false") + + pool = get_redis_connection_pool(host="plain-host", port=6379) + + assert pool is not None + assert pool.connection_class is async_redis.Connection + assert "ssl" not in pool.connection_kwargs From 2162da501527e6bfa3ba0d162711240eaaefebb1 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 13:39:10 -0700 Subject: [PATCH 60/62] fix(langfuse_otel): build per-request OTLP exporter from key and team dynamic Langfuse credentials (#32437) * fix(langfuse_otel): build per-request OTLP exporter from key/team dynamic Langfuse credentials Key-scoped langfuse_otel callbacks only injected Authorization headers into the init-time exporter, so a proxy without global LANGFUSE_* env vars kept its fallback exporter and never exported traces to Langfuse. Dynamic params now build a full per-request OTLP config (endpoint from the key's langfuse_host, otlp_http, basic auth from the key's credentials). Resolves LIT-3976 * fix(otel): log dynamic config endpoint in span processor debug output * fix(otel): redact authorization headers in exporter debug logs --- .../integrations/langfuse/langfuse_otel.py | 79 +++++---- litellm/integrations/opentelemetry.py | 111 +++++++++--- .../integrations/test_langfuse_otel.py | 165 ++++++++++++++++++ 3 files changed, 293 insertions(+), 62 deletions(-) diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index fc7c1b211c0..449457bd123 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -267,29 +267,10 @@ class LangfuseOtelLogger(OpenTelemetry): # If no keys, return default from env (likely logging to console or something else) return OpenTelemetryConfig.from_env() - # Determine endpoint - default to US cloud - langfuse_host = LangfuseOtelLogger._get_langfuse_otel_host() - - if langfuse_host: - # If LANGFUSE_HOST is provided, construct OTEL endpoint from it - if not langfuse_host.startswith("http"): - langfuse_host = "https://" + langfuse_host - endpoint = f"{langfuse_host.rstrip('/')}/api/public/otel" - verbose_logger.debug(f"Using Langfuse OTEL endpoint from host: {endpoint}") - else: - # Default to US cloud endpoint - endpoint = LANGFUSE_CLOUD_US_ENDPOINT - verbose_logger.debug(f"Using Langfuse US cloud endpoint: {endpoint}") - - auth_header = LangfuseOtelLogger._get_langfuse_authorization_header( - public_key=public_key, secret_key=secret_key - ) - otlp_auth_headers = f"Authorization={auth_header}" - - return OpenTelemetryConfig( - exporter="otlp_http", - endpoint=endpoint, - headers=otlp_auth_headers, + return LangfuseOtelLogger._build_langfuse_otel_config( + public_key=public_key, + secret_key=secret_key, + langfuse_host=LangfuseOtelLogger._get_langfuse_otel_host(), ) @staticmethod @@ -316,33 +297,36 @@ class LangfuseOtelLogger(OpenTelemetry): "LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY must be set for Langfuse OpenTelemetry integration." ) - # Determine endpoint - default to US cloud - langfuse_host = LangfuseOtelLogger._get_langfuse_otel_host() + return LangfuseOtelLogger._build_langfuse_otel_config( + public_key=public_key, + secret_key=secret_key, + langfuse_host=LangfuseOtelLogger._get_langfuse_otel_host(), + ) + @staticmethod + def _build_langfuse_otel_config( + public_key: str, secret_key: str, langfuse_host: Optional[str] + ) -> "OpenTelemetryConfig": + """ + Builds an OTLP HTTP config pointing at the Langfuse OTEL endpoint for the + given host (US cloud when no host is provided), authorized with the given keys. + """ if langfuse_host: - # If LANGFUSE_HOST is provided, construct OTEL endpoint from it - if not langfuse_host.startswith("http"): - langfuse_host = "https://" + langfuse_host - endpoint = f"{langfuse_host.rstrip('/')}/api/public/otel" + normalized_host = langfuse_host if langfuse_host.startswith("http") else f"https://{langfuse_host}" + endpoint = f"{normalized_host.rstrip('/')}/api/public/otel" verbose_logger.debug(f"Using Langfuse OTEL endpoint from host: {endpoint}") else: - # Default to US cloud endpoint endpoint = LANGFUSE_CLOUD_US_ENDPOINT verbose_logger.debug(f"Using Langfuse US cloud endpoint: {endpoint}") auth_header = LangfuseOtelLogger._get_langfuse_authorization_header( public_key=public_key, secret_key=secret_key ) - otlp_auth_headers = f"Authorization={auth_header}" - - # Prevent modification of global env vars which causes leakage - # os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint - # os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = otlp_auth_headers return OpenTelemetryConfig( exporter="otlp_http", endpoint=endpoint, - headers=otlp_auth_headers, + headers=f"Authorization={auth_header}", ) @staticmethod @@ -378,6 +362,29 @@ class LangfuseOtelLogger(OpenTelemetry): return dynamic_headers + def construct_dynamic_otel_config( + self, standard_callback_dynamic_params: StandardCallbackDynamicParams + ) -> Optional["OpenTelemetryConfig"]: + """ + Build a full per-request OTLP config from team/key dynamic Langfuse credentials. + + Key-scoped credentials must define the export target, not just the auth + headers: without this, a proxy with no global LANGFUSE_* env vars keeps its + init-time fallback exporter (console), so key-level langfuse_otel silently + never reaches Langfuse. + """ + public_key = standard_callback_dynamic_params.get("langfuse_public_key") + secret_key = standard_callback_dynamic_params.get("langfuse_secret_key") + if not public_key or not secret_key: + return None + + langfuse_host = standard_callback_dynamic_params.get("langfuse_host") or self._get_langfuse_otel_host() + return LangfuseOtelLogger._build_langfuse_otel_config( + public_key=public_key, + secret_key=secret_key, + langfuse_host=langfuse_host, + ) + def create_litellm_proxy_request_started_span( self, start_time: datetime, diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index de543fa042b..fea55cd1db4 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -28,6 +28,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import ( parse_semconv_opt_in, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.secret_managers.main import get_secret_bool, str_to_bool from litellm.types.services import ServiceLoggerPayload from litellm.types.utils import ( @@ -948,12 +949,22 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): Returns: Tracer: The tracer to use for this request """ + dynamic_config = self._get_dynamic_otel_config_from_kwargs(kwargs) + if dynamic_config is not None: + verbose_logger.debug( + "[OTEL DEBUG] Using DYNAMIC config tracer with endpoint: %s", + dynamic_config.endpoint, + ) + return self._get_tracer_with_dynamic_config(dynamic_config) + dynamic_headers = self._get_dynamic_otel_headers_from_kwargs(kwargs) if dynamic_headers is not None: # Create spans using a temporary tracer with dynamic headers tracer_to_use = self._get_tracer_with_dynamic_headers(dynamic_headers) - verbose_logger.debug("[OTEL DEBUG] Using DYNAMIC tracer with headers: %s", dynamic_headers) + verbose_logger.debug( + "[OTEL DEBUG] Using DYNAMIC tracer with headers: %s", redact_string(str(dynamic_headers)) + ) else: # For langfuse_otel without dynamic headers, create a provider with env var credentials if hasattr(self, "callback_name") and self.callback_name == "langfuse_otel": @@ -989,6 +1000,32 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): return dynamic_headers if dynamic_headers else None + def _get_dynamic_otel_config_from_kwargs(self, kwargs: dict) -> Optional[OpenTelemetryConfig]: + """Extract a full dynamic exporter config from kwargs if available.""" + standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = kwargs.get( + "standard_callback_dynamic_params" + ) + + if not standard_callback_dynamic_params: + return None + + return self.construct_dynamic_otel_config(standard_callback_dynamic_params=standard_callback_dynamic_params) + + def _get_tracer_with_dynamic_config(self, dynamic_config: OpenTelemetryConfig): + """Create (or reuse) a tracer whose exporter target comes from a per-request config.""" + from opentelemetry.sdk.trace import TracerProvider + + cache_key = f"dynamic_config:{dynamic_config.exporter}:{dynamic_config.endpoint}:{dynamic_config.headers}" + if cache_key in self._tracer_provider_cache: + return self._tracer_provider_cache[cache_key].get_tracer(LITELLM_TRACER_NAME) + + temp_provider = TracerProvider(resource=self._get_litellm_resource(self.config)) + temp_provider.add_span_processor(self._get_span_processor(config_override=dynamic_config)) + + self._tracer_provider_cache[cache_key] = temp_provider + + return temp_provider.get_tracer(LITELLM_TRACER_NAME) + def _get_tracer_with_dynamic_headers(self, dynamic_headers: dict): """Create a temporary tracer with dynamic headers for this request only.""" from opentelemetry.sdk.trace import TracerProvider @@ -1020,6 +1057,19 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): """ return None + def construct_dynamic_otel_config( + self, standard_callback_dynamic_params: StandardCallbackDynamicParams + ) -> Optional[OpenTelemetryConfig]: + """ + Construct a full exporter config from standard callback dynamic params. + + Override this when team/key dynamic params must control the export + target (exporter kind + endpoint), not just the request headers. When + this returns a config, it takes precedence over + construct_dynamic_otel_headers for the request. + """ + return None + ######################################################### # End of Team/Key Based Logging Control Flow ######################################################### @@ -2747,7 +2797,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): verbose_logger.debug("OpenTelemetry: No parent context found, creating root span") return None, None - def _get_span_processor(self, dynamic_headers: Optional[dict] = None): + def _get_span_processor( + self, + dynamic_headers: Optional[dict] = None, + config_override: Optional[OpenTelemetryConfig] = None, + ): from opentelemetry.sdk.trace.export import ( BatchSpanProcessor, ConsoleSpanExporter, @@ -2755,40 +2809,45 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): SpanExporter, ) + otel_exporter = config_override.exporter if config_override else self.OTEL_EXPORTER + otel_endpoint = config_override.endpoint if config_override else self.OTEL_ENDPOINT + otel_headers = config_override.headers if config_override else self.OTEL_HEADERS + verbose_logger.debug( - "OpenTelemetry Logger, initializing span processor \nself.OTEL_EXPORTER: %s\nself.OTEL_ENDPOINT: %s\nself.OTEL_HEADERS: %s", - self.OTEL_EXPORTER, - self.OTEL_ENDPOINT, - self.OTEL_HEADERS, + "OpenTelemetry Logger, initializing span processor \nexporter: %s\nendpoint: %s\nheaders: %s", + otel_exporter, + otel_endpoint, + redact_string(str(otel_headers)), ) - _split_otel_headers = OpenTelemetry._get_headers_dictionary(headers=dynamic_headers or self.OTEL_HEADERS) + _split_otel_headers = OpenTelemetry._get_headers_dictionary(headers=dynamic_headers or otel_headers) if dynamic_headers: verbose_logger.debug( "[OTEL DEBUG] Creating span processor with DYNAMIC headers: %s", - {k: v[:20] + "..." if len(str(v)) > 20 else v for k, v in _split_otel_headers.items()}, + redact_string(str(_split_otel_headers)), + ) + elif config_override: + verbose_logger.debug( + "[OTEL DEBUG] Creating span processor with DYNAMIC config, endpoint: %s", + otel_endpoint, ) else: verbose_logger.debug("[OTEL DEBUG] Creating span processor with GLOBAL headers") - if hasattr(self.OTEL_EXPORTER, "export"): # Check if it has the export method that SpanExporter requires + if hasattr(otel_exporter, "export"): # Check if it has the export method that SpanExporter requires verbose_logger.debug( "OpenTelemetry: intiializing SpanExporter. Value of OTEL_EXPORTER: %s", - self.OTEL_EXPORTER, + otel_exporter, ) - return SimpleSpanProcessor(cast(SpanExporter, self.OTEL_EXPORTER)) + return SimpleSpanProcessor(cast(SpanExporter, otel_exporter)) - if self.OTEL_EXPORTER == "console": + if otel_exporter == "console": verbose_logger.debug( "OpenTelemetry: intiializing console exporter. Value of OTEL_EXPORTER: %s", - self.OTEL_EXPORTER, + otel_exporter, ) return BatchSpanProcessor(ConsoleSpanExporter()) - elif ( - self.OTEL_EXPORTER == "otlp_http" - or self.OTEL_EXPORTER == "http/protobuf" - or self.OTEL_EXPORTER == "http/json" - ): + elif otel_exporter == "otlp_http" or otel_exporter == "http/protobuf" or otel_exporter == "http/json": try: from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter as OTLPSpanExporterHTTP, @@ -2801,13 +2860,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): verbose_logger.debug( "OpenTelemetry: intiializing http exporter. Value of OTEL_EXPORTER: %s", - self.OTEL_EXPORTER, + otel_exporter, ) - normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "traces") + normalized_endpoint = self._normalize_otel_endpoint(otel_endpoint, "traces") return BatchSpanProcessor( OTLPSpanExporterHTTP(endpoint=normalized_endpoint, headers=_split_otel_headers), ) - elif self.OTEL_EXPORTER == "otlp_grpc" or self.OTEL_EXPORTER == "grpc": + elif otel_exporter == "otlp_grpc" or otel_exporter == "grpc": try: from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( OTLPSpanExporter as OTLPSpanExporterGRPC, @@ -2820,16 +2879,16 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): verbose_logger.debug( "OpenTelemetry: intiializing grpc exporter. Value of OTEL_EXPORTER: %s", - self.OTEL_EXPORTER, + otel_exporter, ) - normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "traces") + normalized_endpoint = self._normalize_otel_endpoint(otel_endpoint, "traces") return BatchSpanProcessor( OTLPSpanExporterGRPC(endpoint=normalized_endpoint, headers=_split_otel_headers), ) else: verbose_logger.debug( "OpenTelemetry: intiializing console exporter. Value of OTEL_EXPORTER: %s", - self.OTEL_EXPORTER, + otel_exporter, ) return BatchSpanProcessor(ConsoleSpanExporter()) @@ -2841,7 +2900,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): "OpenTelemetry Logger, initializing log exporter \nself.OTEL_EXPORTER: %s\nself.OTEL_ENDPOINT: %s\nself.OTEL_HEADERS: %s", self.OTEL_EXPORTER, self.OTEL_ENDPOINT, - self.OTEL_HEADERS, + redact_string(str(self.OTEL_HEADERS)), ) _split_otel_headers = OpenTelemetry._get_headers_dictionary(self.OTEL_HEADERS) @@ -2928,7 +2987,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): "OpenTelemetry Logger, initializing metric reader\nself.OTEL_EXPORTER: %s\nself.OTEL_ENDPOINT: %s\nself.OTEL_HEADERS: %s", self.OTEL_EXPORTER, self.OTEL_ENDPOINT, - self.OTEL_HEADERS, + redact_string(str(self.OTEL_HEADERS)), ) _split_otel_headers = OpenTelemetry._get_headers_dictionary(self.OTEL_HEADERS) diff --git a/tests/test_litellm/integrations/test_langfuse_otel.py b/tests/test_litellm/integrations/test_langfuse_otel.py index 3aade7514e4..2f2675ca790 100644 --- a/tests/test_litellm/integrations/test_langfuse_otel.py +++ b/tests/test_litellm/integrations/test_langfuse_otel.py @@ -412,6 +412,171 @@ class TestLangfuseOtelIntegration: # Endpoint assertion removed as side effect is gone +class TestLangfuseOtelKeyDynamicConfig: + """Key/team-scoped Langfuse credentials must define the full export target + (OTLP endpoint + auth), not just auth headers on the init-time exporter.""" + + CLEAN_ENV_VARS = [ + "LANGFUSE_PUBLIC_KEY", + "LANGFUSE_SECRET_KEY", + "LANGFUSE_HOST", + "LANGFUSE_OTEL_HOST", + "OTEL_EXPORTER", + "OTEL_EXPORTER_OTLP_PROTOCOL", + "OTEL_ENDPOINT", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_HEADERS", + "OTEL_EXPORTER_OTLP_HEADERS", + ] + + def _clean_env(self): + cleaned = {k: v for k, v in os.environ.items() if k not in self.CLEAN_ENV_VARS} + return patch.dict(os.environ, cleaned, clear=True) + + def _dynamic_params(self, **overrides): + from litellm.types.utils import StandardCallbackDynamicParams + + params = { + "langfuse_public_key": "key_public", + "langfuse_secret_key": "key_secret", + "langfuse_host": "https://langfuse.example.com", + } + params.update(overrides) + return StandardCallbackDynamicParams(**{k: v for k, v in params.items() if v is not None}) + + def test_construct_dynamic_otel_config_with_key_credentials(self): + with self._clean_env(): + logger = LangfuseOtelLogger() + config = logger.construct_dynamic_otel_config(self._dynamic_params()) + + assert config is not None + assert config.exporter == "otlp_http" + assert config.endpoint == "https://langfuse.example.com/api/public/otel" + + import base64 + + expected_auth = base64.b64encode(b"key_public:key_secret").decode() + assert config.headers == f"Authorization=Basic {expected_auth}" + + def test_construct_dynamic_otel_config_host_without_protocol(self): + with self._clean_env(): + logger = LangfuseOtelLogger() + config = logger.construct_dynamic_otel_config(self._dynamic_params(langfuse_host="langfuse.example.com")) + + assert config is not None + assert config.endpoint == "https://langfuse.example.com/api/public/otel" + + def test_construct_dynamic_otel_config_defaults_to_us_cloud(self): + with self._clean_env(): + logger = LangfuseOtelLogger() + config = logger.construct_dynamic_otel_config(self._dynamic_params(langfuse_host=None)) + + assert config is not None + assert config.endpoint == "https://us.cloud.langfuse.com/api/public/otel" + + def test_construct_dynamic_otel_config_falls_back_to_env_host(self): + with self._clean_env(): + with patch.dict(os.environ, {"LANGFUSE_HOST": "https://env-host.example.com"}): + logger = LangfuseOtelLogger() + config = logger.construct_dynamic_otel_config(self._dynamic_params(langfuse_host=None)) + + assert config is not None + assert config.endpoint == "https://env-host.example.com/api/public/otel" + + def test_construct_dynamic_otel_config_requires_both_keys(self): + with self._clean_env(): + logger = LangfuseOtelLogger() + + assert logger.construct_dynamic_otel_config(self._dynamic_params(langfuse_secret_key=None)) is None + assert logger.construct_dynamic_otel_config(self._dynamic_params(langfuse_public_key=None)) is None + + def test_key_dynamic_params_create_otlp_exporter_without_global_env(self): + """Without global LANGFUSE_* env vars, a request carrying key-scoped Langfuse + credentials must get a tracer exporting via OTLP HTTP to that key's host.""" + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter, + ) + from opentelemetry.sdk.trace.export import BatchSpanProcessor + + with self._clean_env(): + logger = LangfuseOtelLogger() + assert logger.OTEL_EXPORTER == "console" + + tracer = logger.get_tracer_to_use_for_request( + {"standard_callback_dynamic_params": self._dynamic_params()} + ) + + assert tracer is not logger.tracer + assert len(logger._tracer_provider_cache) == 1 + + provider = next(iter(logger._tracer_provider_cache.values())) + span_processors = provider._active_span_processor._span_processors + assert len(span_processors) == 1 + assert isinstance(span_processors[0], BatchSpanProcessor) + + exporter = span_processors[0].span_exporter + assert isinstance(exporter, OTLPSpanExporter) + assert exporter._endpoint == "https://langfuse.example.com/api/public/otel/v1/traces" + + import base64 + + expected_auth = base64.b64encode(b"key_public:key_secret").decode() + assert exporter._headers == {"Authorization": f"Basic {expected_auth}"} + + def test_key_dynamic_params_reuse_cached_provider(self): + with self._clean_env(): + logger = LangfuseOtelLogger() + kwargs = {"standard_callback_dynamic_params": self._dynamic_params()} + logger.get_tracer_to_use_for_request(kwargs) + logger.get_tracer_to_use_for_request(kwargs) + + assert len(logger._tracer_provider_cache) == 1 + + def test_no_dynamic_params_keeps_default_tracer(self): + with self._clean_env(): + logger = LangfuseOtelLogger() + tracer = logger.get_tracer_to_use_for_request({}) + + assert tracer is logger.tracer + assert logger._tracer_provider_cache == {} + + def test_key_credentials_never_passed_to_debug_logger(self): + """The span-processor debug logs must receive a redacted header value, so the + key-scoped Langfuse secret never enters a log record regardless of downstream + handler configuration, while the exporter still gets the real header.""" + import base64 + + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter, + ) + + from litellm.integrations import opentelemetry as otel_module + + secret = base64.b64encode(b"key_public:key_secret").decode() + + recorded_arguments = [] + + def _spy(message, *args, **kwargs): + recorded_arguments.append(" ".join(str(part) for part in (message, *args))) + + with self._clean_env(): + logger = LangfuseOtelLogger() + with patch.object(otel_module.verbose_logger, "debug", side_effect=_spy): + logger.get_tracer_to_use_for_request( + {"standard_callback_dynamic_params": self._dynamic_params()} + ) + + logged = "\n".join(recorded_arguments) + assert "initializing span processor" in logged + assert secret not in logged + assert f"Basic {secret}" not in logged + + provider = next(iter(logger._tracer_provider_cache.values())) + exporter = provider._active_span_processor._span_processors[0].span_exporter + assert isinstance(exporter, OTLPSpanExporter) + assert exporter._headers == {"Authorization": f"Basic {secret}"} + + class TestLangfuseOtelResponsesAPI: """Test suite for Langfuse OTEL integration with ResponsesAPI""" From 21ba9692c3720ffd278f4035c4e733c0f34f492d Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 14:41:24 -0700 Subject: [PATCH 61/62] fix(router): apply team/key enable_tag_filtering to tag routing (#33436) Team/key router_settings.enable_tag_filtering was stored and echoed by /team/info but never applied at request time: the per-request override whitelist in route_llm_request.py dropped it, tag filtering only read the router-level flag, and UpdateRouterConfig silently discarded the field on /key/generate and /config/update. Requests from teams with the toggle on were load balanced across all deployments instead of tag-matched ones. - add enable_tag_filtering to the router_settings_override whitelist and strip any client-supplied copy from the request body first, so only the key/team value reaches the router - run tag filtering when the request carries enable_tag_filtering=True; a request-level False cannot disable a router-level True, so per-request settings can only scope down, never escape the global policy - add the field to UpdateRouterConfig so key and config update paths stop dropping it, and to all_litellm_params so it never leaks into provider request bodies - allow it through Router.update_settings/get_settings so the global UI toggle persists across DB config reloads Resolves LIT-4390 --- litellm/proxy/route_llm_request.py | 3 + litellm/router.py | 2 + litellm/router_strategy/tag_based_routing.py | 8 +- litellm/types/router.py | 1 + litellm/types/utils.py | 1 + tests/test_litellm/proxy/test_proxy_types.py | 15 +++ .../proxy/test_route_llm_request.py | 68 +++++++++++++ .../test_router_tag_routing.py | 96 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 9 files changed, 195 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 0ca2a75990b..25fa0819930 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -362,6 +362,8 @@ async def route_request( for _key in _MOCK_TESTING_KWARG_NAMES: data.pop(_key, None) + data.pop("enable_tag_filtering", None) + team_id = get_team_id_from_data(data) router_model_names = llm_router.model_names if llm_router is not None else [] is_proxy_admin_without_team = team_id is None and _is_proxy_admin_request(data) @@ -410,6 +412,7 @@ async def route_request( "timeout", "model_group_retry_policy", "routing_strategy", + "enable_tag_filtering", ] # Merge override settings into data (only if not already set in request) diff --git a/litellm/router.py b/litellm/router.py index 9b869d741ba..78e156801f8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9778,6 +9778,7 @@ class Router: "retry_policy", "model_group_alias", "enable_weighted_failover", + "enable_tag_filtering", ] for var in vars_to_include: @@ -9814,6 +9815,7 @@ class Router: "model_group_retry_policy", "model_group_alias", "enable_weighted_failover", + "enable_tag_filtering", ] _int_settings = [ diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index 6ca4e1de322..710c2199107 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -160,8 +160,14 @@ async def get_deployments_for_tag( Returns a list of deployments that match the requested model and tags in the request. Executes tag based filtering based on the tags in request metadata and the tags on the deployments + + Runs when the router-level `enable_tag_filtering` is True or the request carries + `enable_tag_filtering=True` (set from key/team router_settings by the proxy). + A request-level False never disables a router-level True, so per-request settings + cannot escape an operator's global tag-routing policy. """ - if llm_router_instance.enable_tag_filtering is not True: + request_enable_tag_filtering = request_kwargs.get("enable_tag_filtering") if request_kwargs else None + if request_enable_tag_filtering is not True and llm_router_instance.enable_tag_filtering is not True: return healthy_deployments if request_kwargs is None: diff --git a/litellm/types/router.py b/litellm/types/router.py index d62c613bf57..69a8ca9f19e 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -117,6 +117,7 @@ class UpdateRouterConfig(BaseModel): fallbacks: Optional[List[dict]] = None context_window_fallbacks: Optional[List[dict]] = None model_group_alias: Optional[Dict[str, Union[str, Dict]]] = {} + enable_tag_filtering: Optional[bool] = None model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8c4739fd417..88b3a39844f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3222,6 +3222,7 @@ all_litellm_params = ( "shared_session", "search_tool_name", "order", + "enable_tag_filtering", "enable_json_schema_validation", "use_xai_oauth", "_litellm_rate_limit_descriptors", diff --git a/tests/test_litellm/proxy/test_proxy_types.py b/tests/test_litellm/proxy/test_proxy_types.py index bc77a9ba3c0..c8e0b3a730a 100644 --- a/tests/test_litellm/proxy/test_proxy_types.py +++ b/tests/test_litellm/proxy/test_proxy_types.py @@ -124,3 +124,18 @@ def test_proxy_exception_str_returns_message(): "param": "key", "code": "401", } + + +def test_key_request_router_settings_keeps_enable_tag_filtering(): + """``router_settings`` on key requests validates through + ``UpdateRouterConfig``; a field missing from that model is silently + dropped at parse time, so a key's "Enable Tag Filtering" toggle would + never reach the DB even though the team path (plain dict) kept it.""" + from litellm.proxy._types import GenerateKeyRequest + + req = GenerateKeyRequest(router_settings={"enable_tag_filtering": True, "num_retries": 2}) + + assert req.router_settings is not None + dumped = req.router_settings.model_dump(exclude_none=True) + assert dumped["enable_tag_filtering"] is True + assert dumped["num_retries"] == 2 diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index c96b86f84b0..f506b9665a6 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -819,3 +819,71 @@ async def test_route_request_realtime_transcription_session_resolves_credentials ) assert mock_handler.call_args.kwargs["api_key"] == "transcription-key" + + +@pytest.mark.asyncio +async def test_route_request_merges_enable_tag_filtering_from_override(): + """Key/team router_settings carry enable_tag_filtering; the override + whitelist must forward it to the router call or the team's tag-routing + toggle saved in the UI is silently ignored at request time.""" + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + "router_settings_override": { + "enable_tag_filtering": True, + }, + } + + llm_router = MagicMock() + llm_router.acompletion.return_value = "success" + + response = await route_request(data, llm_router, None, "acompletion") + + assert response == "success" + call_kwargs = llm_router.acompletion.call_args[1] + assert call_kwargs["enable_tag_filtering"] is True + + +@pytest.mark.asyncio +async def test_route_request_strips_client_supplied_enable_tag_filtering(): + """enable_tag_filtering influences deployment selection and is only + trusted when it comes from key/team router_settings via + router_settings_override. A caller putting it in the request body must + not reach the router with it.""" + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + "enable_tag_filtering": True, + } + + llm_router = MagicMock() + llm_router.acompletion.return_value = "ok" + + await route_request(data, llm_router, None, "acompletion") + + call_kwargs = llm_router.acompletion.call_args[1] + assert "enable_tag_filtering" not in call_kwargs + assert "enable_tag_filtering" not in data + + +@pytest.mark.asyncio +async def test_route_request_override_enable_tag_filtering_beats_body_value(): + """A client-sent enable_tag_filtering must not shadow the key/team + setting: the body copy is stripped first, so the override value is the + one the router sees.""" + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + "enable_tag_filtering": False, + "router_settings_override": { + "enable_tag_filtering": True, + }, + } + + llm_router = MagicMock() + llm_router.acompletion.return_value = "ok" + + await route_request(data, llm_router, None, "acompletion") + + call_kwargs = llm_router.acompletion.call_args[1] + assert call_kwargs["enable_tag_filtering"] is True diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index eb289095c51..98506aad594 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -1019,3 +1019,99 @@ async def test_negation_removes_tag_regex_deployment_falls_to_ban_only(): mock_response="hi", ) assert response._hidden_params["model_id"] == "openai-deployment" + + +@pytest.mark.asyncio() +async def test_request_level_enable_tag_filtering_applies_when_global_off(): + """ + A request carrying enable_tag_filtering=True (set by the proxy from key/team + router_settings) must activate tag filtering even when the router-level flag + is off. Without this, a team's "Enable Tag Filtering" toggle saved in the UI + is silently ignored at request time. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["teamA"], + }, + "model_info": {"id": "team-a-deployment"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["teamB"], + }, + "model_info": {"id": "team-b-deployment"}, + }, + ], + enable_tag_filtering=False, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["teamA"]}, + enable_tag_filtering=True, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "team-a-deployment" + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["teamB"]}, + enable_tag_filtering=True, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "team-b-deployment" + + +@pytest.mark.asyncio() +async def test_request_level_enable_tag_filtering_false_cannot_disable_global(): + """ + A request-level enable_tag_filtering=False must not bypass a router-level + True: tag filtering can be an operator-level restriction on which + deployments a caller may reach, so per-request settings may only scope + down, never escape the global policy. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["teamA"], + }, + "model_info": {"id": "team-a-deployment"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["teamB"], + }, + "model_info": {"id": "team-b-deployment"}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["teamA"]}, + enable_tag_filtering=False, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "team-a-deployment" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 87c760257d1..96e2c450e00 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -32245,6 +32245,8 @@ export interface components { }[] | null; /** Cooldown Time */ cooldown_time?: number | null; + /** Enable Tag Filtering */ + enable_tag_filtering?: boolean | null; /** Fallbacks */ fallbacks?: { [key: string]: unknown; From 5ab160113f9f5045031eea619398847efb53d2af Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 16 Jul 2026 14:50:40 -0700 Subject: [PATCH 62/62] feat(proxy): add disable_auto_add_proxy_admin_to_teams flag (#33563) --- litellm/proxy/_types.py | 4 + .../management_endpoints/team_endpoints.py | 26 ++++-- litellm/proxy/proxy_server.py | 8 ++ .../test_team_endpoints.py | 79 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 26 ++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 ++ 6 files changed, 140 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 053f78a3698..d102c1d1e37 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2369,6 +2369,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="If True, stores request messages and responses in spend logs. Default is False.", ) + disable_auto_add_proxy_admin_to_teams: bool | None = Field( + None, + description="By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False.", + ) maximum_spend_logs_retention_period: Optional[str] = Field( None, description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d267a3cac69..70c002d2d2d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -14,7 +14,7 @@ import json import math import traceback from datetime import datetime, timezone -from typing import Annotated, Any, Dict, List, Optional, Tuple, Union, cast +from typing import Annotated, Any, Dict, List, Mapping, Optional, Tuple, Union, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -894,6 +894,17 @@ def _check_team_budget_update_authority( ) +def _should_auto_add_team_creator( + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], +) -> bool: + if user_api_key_dict.user_id is None: + return False + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + return True + return general_settings.get("disable_auto_add_proxy_admin_to_teams") is not True + + #### TEAM MANAGEMENT #### @router.post( "/team/new", @@ -997,6 +1008,7 @@ async def new_team( from litellm.proxy.proxy_server import ( _license_check, create_audit_log_for_update, + general_settings, litellm_proxy_admin_name, prisma_client, user_api_key_cache, @@ -1123,13 +1135,11 @@ async def new_team( user_api_key_cache=user_api_key_cache, ) - if user_api_key_dict.user_id is not None: - creating_user_in_list = False - for member in data.members_with_roles: - if member.user_id == user_api_key_dict.user_id: - creating_user_in_list = True - - if creating_user_in_list is False: + if _should_auto_add_team_creator(user_api_key_dict, general_settings): + creating_user_in_list = any( + member.user_id == user_api_key_dict.user_id for member in data.members_with_roles + ) + if not creating_user_in_list: data.members_with_roles.append(Member(role="admin", user_id=user_api_key_dict.user_id)) _check_passthrough_routes_caller_permission(data, user_api_key_dict, entity="team") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7e6ca5a0108..b0eb42b266e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5794,6 +5794,13 @@ class ProxyConfig: # For other types, convert to bool general_settings["store_prompts_in_spend_logs"] = bool(value) + if "disable_auto_add_proxy_admin_to_teams" in _general_settings: + value = _general_settings["disable_auto_add_proxy_admin_to_teams"] + if isinstance(value, str): + general_settings["disable_auto_add_proxy_admin_to_teams"] = value.lower() == "true" + else: + general_settings["disable_auto_add_proxy_admin_to_teams"] = value if value is None else bool(value) + ## STORE MODEL IN DB ## if "store_model_in_db" in _general_settings: value = _general_settings["store_model_in_db"] @@ -14907,6 +14914,7 @@ async def get_config_list( "mcp_required_fields": {"type": "List"}, "cancel_on_disconnect": {"type": "Boolean"}, "skip_user_budget_on_team_key": {"type": "Boolean"}, + "disable_auto_add_proxy_admin_to_teams": {"type": "Boolean"}, } return_val = [] 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 f7b2df45a85..4936191c344 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -555,6 +555,85 @@ async def test_new_team_with_mcp_tool_permissions(mock_db_client, mock_admin_aut assert created_permission_data["mcp_servers"] == ["server_a", "server_b"] +@pytest.mark.parametrize( + "user_role,user_id,flag_value,expected", + [ + (LitellmUserRoles.PROXY_ADMIN, "admin-1", True, False), + (LitellmUserRoles.PROXY_ADMIN, "admin-1", False, True), + (LitellmUserRoles.PROXY_ADMIN, "admin-1", None, True), + (LitellmUserRoles.INTERNAL_USER, "user-1", True, True), + (LitellmUserRoles.ORG_ADMIN, "org-admin-1", True, True), + (LitellmUserRoles.PROXY_ADMIN, None, False, False), + ], +) +def test_should_auto_add_team_creator(user_role, user_id, flag_value, expected): + from litellm.proxy.management_endpoints.team_endpoints import ( + _should_auto_add_team_creator, + ) + + general_settings = ( + {} if flag_value is None else {"disable_auto_add_proxy_admin_to_teams": flag_value} + ) + auth = UserAPIKeyAuth(user_role=user_role, user_id=user_id) + assert _should_auto_add_team_creator(auth, general_settings) is expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "disable_flag,expect_creator_added", [(True, False), (False, True)] +) +async def test_new_team_disable_auto_add_proxy_admin_flag( + mock_db_client, disable_flag, expect_creator_added +): + """ + When general_settings.disable_auto_add_proxy_admin_to_teams is True, a proxy + admin calling /team/new must NOT be auto-added to the team's members. When + the flag is off, the creator is auto-added as a team admin (default + behavior, regression guard for LIT-3739). + """ + mock_db_client.jsonify_team_object = lambda db_data: db_data + mock_db_client.get_data = AsyncMock(return_value=None) + mock_db_client.update_data = AsyncMock(return_value=MagicMock()) + mock_db_client.db = MagicMock() + + team_create_result = MagicMock(team_id="team-789") + team_create_result.model_dump.return_value = {"team_id": "team-789"} + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = AsyncMock( + return_value=team_create_result + ) + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + admin_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user-1" + ) + + with patch( + "litellm.proxy.proxy_server.general_settings", + {"disable_auto_add_proxy_admin_to_teams": disable_flag}, + ), patch( + "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", + new_callable=AsyncMock, + ) as mock_add_members: + await new_team( + data=NewTeamRequest(team_alias="flag-test-team"), + http_request=MagicMock(spec=Request), + user_api_key_dict=admin_auth, + ) + + mock_add_members.assert_called_once() + member_add_request = mock_add_members.call_args.kwargs["data"] + member_user_ids = [m.user_id for m in member_add_request.member] + assert ("admin-user-1" in member_user_ids) is expect_creator_added + + @pytest.mark.asyncio async def test_team_update_object_permissions_existing_permission(monkeypatch): """ diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 079b844638c..2f0e71c4205 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -6265,6 +6265,32 @@ async def test_update_general_settings_store_model_in_db_false(): assert ps.general_settings["store_model_in_db"] is False +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_value,expected", + [(True, True), (False, False), ("true", True), ("false", False), (None, None)], +) +async def test_update_general_settings_disable_auto_add_proxy_admin_to_teams(db_value, expected): + """ + Verify _update_general_settings propagates disable_auto_add_proxy_admin_to_teams + from the DB config into the live general_settings dict, so a UI toggle via + /config/field/update takes effect on the next config poll instead of + requiring a proxy restart. + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + with patch("litellm.proxy.proxy_server.general_settings", {}): + await proxy_config._update_general_settings( + db_general_settings={"disable_auto_add_proxy_admin_to_teams": db_value} + ) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["disable_auto_add_proxy_admin_to_teams"] is expected + + @pytest.mark.asyncio async def test_update_general_settings_store_model_in_db_string_normalization(): """ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 96e2c450e00..11fc38c7c34 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22504,6 +22504,11 @@ export interface components { * @description connect to a postgres db - needed for generating temporary keys + tracking spend / key */ database_url?: string | null; + /** + * Disable Auto Add Proxy Admin To Teams + * @description By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False. + */ + disable_auto_add_proxy_admin_to_teams?: boolean | null; /** * Disable Budget Reservation * @description If True, disables the optimistic per-request budget reservation introduced in v1.84.0. WARNING: This weakens hard budget enforcement. Without the reservation, a burst of concurrent requests from a single key can each pass the read-time spend check before any of them is charged, allowing a configured budget to be exceeded under high concurrency. Budgets are still evaluated on every request at read time, so an already-exhausted budget is still rejected. Enable only if your deployment is experiencing phantom BudgetExceededError responses caused by leaked reservations (see GitHub issue #27639). A proxy-level WARNING is logged on every request while this flag is active as a reminder that hard enforcement is relaxed.