From fa6b20916572821e58618b0e927b2764fceed768 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 09:56:59 -0700 Subject: [PATCH 1/3] feat(guardrails): add only_scan_new_messages for per-session incremental scanning (#33278) * feat(guardrails): add only_scan_new_messages for per-session incremental scanning Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(guardrails): use fixed TTL constant and revert unrelated test formatting Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(guardrails): run only_scan_new_messages in the unified apply_guardrail path The initial wiring lived in BedrockGuardrail.async_pre_call_hook, but the proxy routes Bedrock through the unified apply_guardrail interface, so the flag had no effect live. Move incremental selection into apply_guardrail: filter the flat texts list against per-session scanned hashes, skip the Bedrock call when nothing is new, and mark hashes only after a successful (non-blocked) scan. Full-context fallback is preserved when there is no session id, the cache is unavailable, or a masking guardrail is configured. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover session-id fallbacks and mark_texts_scanned guards Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(guardrails): fall back to full scan when incremental guardrail masks content, use shared cache Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover generic agent multi-turn incremental scan Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover incremental scan cache resolver fallbacks Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover flag interactions and /v1/messages incremental scan semantics * feat(guardrails): make GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS env configurable * test(guardrails): prove skip_system/skip_tool are enforced upstream of incremental scan --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> Co-authored-by: Yucheng Zhu --- litellm/constants.py | 3 + litellm/integrations/custom_guardrail.py | 102 +++++- .../guardrail_hooks/bedrock_guardrails.py | 93 +++++ .../guardrails/guardrail_initializers.py | 1 + litellm/types/guardrails.py | 12 + .../integrations/test_custom_guardrail.py | 198 ++++++++++ .../test_anthropic_guardrail_handler.py | 129 +++++++ .../test_openai_guardrail_handler.py | 92 +++++ .../test_bedrock_guardrails.py | 340 ++++++++++++++++++ 9 files changed, 969 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 2af84c139a1..84ac9e29729 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -264,6 +264,9 @@ MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) +GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS = int( + os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60) +) # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 856556f7c56..cf9dafcb222 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,3 +1,4 @@ +import hashlib import os import secrets from datetime import datetime @@ -46,7 +47,10 @@ if TYPE_CHECKING: dc = DualCache() -from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.constants import ( + GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + PRE_CALL_EXECUTED_GUARDRAILS_KEY, +) from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, @@ -113,6 +117,7 @@ class CustomGuardrail(CustomLogger): on_sensitive_data: Optional[str] = None, sensitive_data_route_to_model: Optional[str] = None, sticky_session_routing: bool = True, + only_scan_new_messages: bool = False, **kwargs, ): """ @@ -145,6 +150,7 @@ class CustomGuardrail(CustomLogger): self.on_sensitive_data: Optional[str] = on_sensitive_data self.sensitive_data_route_to_model: Optional[str] = sensitive_data_route_to_model self.sticky_session_routing: bool = sticky_session_routing + self.only_scan_new_messages: bool = only_scan_new_messages if supported_event_hooks: ## validate event_hook is in supported_event_hooks @@ -269,6 +275,100 @@ class CustomGuardrail(CustomLogger): """Extract session_id from request data.""" return get_session_id_from_request_data(request_data) + @staticmethod + def _scanned_text_hash(text: str) -> str: + """Stable content hash for a single scannable text segment. + + Hashing the exact text the provider would receive means an edited earlier + segment produces a different hash and gets re-scanned, while an unchanged + segment repeated on a later turn is skipped. + """ + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + def _scanned_texts_cache_key(self, session_id: str) -> str: + return f"guardrail_scanned_texts:{self.guardrail_name}:{session_id}" + + async def filter_new_texts_for_session( + self, + texts: list[str] | None, + request_data: dict[str, object], + cache: DualCache, + ) -> list[str] | None: + """Return only the text segments not already scanned earlier in this session. + + Returns ``None`` when incremental scanning is inactive (feature off, no + session id, masking enabled, or the cache read failed). ``None`` signals + the caller to fall back to a full scan; a returned list (possibly empty) + signals the caller to scan only that subset and skip masking write-back. + """ + if not self.only_scan_new_messages or not texts: + return None + + if self.mask_request_content or self.mask_response_content: + verbose_logger.warning( + "Guardrail %s: only_scan_new_messages is not supported with masking; scanning full context.", + self.guardrail_name, + ) + return None + + session_id = get_session_id_from_request_data(request_data) + if not session_id: + verbose_logger.debug( + "Guardrail %s: only_scan_new_messages enabled but request has no session id; scanning full context.", + self.guardrail_name, + ) + return None + + try: + cached: object = await cache.async_get_cache(key=self._scanned_texts_cache_key(session_id)) + except Exception as e: # noqa: BLE001 # cache is best-effort; any failure must fall back to a full scan + verbose_logger.warning( + "Guardrail %s: failed to read scanned-message cache (%s); scanning full context.", + self.guardrail_name, + e, + ) + return None + + seen: set[str] = {str(h) for h in cached} if isinstance(cached, list) else set() + return [text for text in texts if self._scanned_text_hash(text) not in seen] + + async def mark_texts_scanned( + self, + texts: list[str] | None, + request_data: dict[str, object], + cache: DualCache, + ) -> None: + """Record the hashes of all text segments present on a successful (non-blocked) scan. + + Called only after the guardrail allows the request, so a blocked segment is + never marked scanned and will be re-checked if the client retries. + """ + if not self.only_scan_new_messages or not texts: + return + if self.mask_request_content or self.mask_response_content: + return + session_id = get_session_id_from_request_data(request_data) + if not session_id: + return + + cache_key = self._scanned_texts_cache_key(session_id) + current_hashes = [self._scanned_text_hash(text) for text in texts] + try: + existing: object = await cache.async_get_cache(key=cache_key) + existing_hashes: list[str] = [str(h) for h in existing] if isinstance(existing, list) else [] + merged: list[str] = list(dict.fromkeys(existing_hashes + current_hashes)) + await cache.async_set_cache( + key=cache_key, + value=merged, + ttl=GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + ) + except Exception as e: # noqa: BLE001 # cache is best-effort; any failure must not block the request + verbose_logger.warning( + "Guardrail %s: failed to persist scanned-message cache (%s); next call will re-scan.", + self.guardrail_name, + e, + ) + def should_route_on_sensitive_data(self) -> bool: """ Returns True if this guardrail is configured to route requests diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 54156715da8..cec682d772a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -2046,6 +2046,90 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): masking_index += 1 verbose_proxy_logger.debug("Applied masking to choice text content") + @staticmethod + def _incremental_scan_cache() -> DualCache: + """Resolve the cache used to remember which segments a session already scanned. + + Prefers the proxy's shared cache (``internal_usage_cache.dual_cache``), which is + backed by Redis when the deployment configures it, so incremental state is shared + across proxy instances. Falls back to a process-local ``DualCache`` singleton when + the proxy is not running (e.g. unit tests), where sharing does not apply. + """ + from litellm.integrations.custom_guardrail import dc as fallback_cache + + try: + from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging + except Exception: # noqa: BLE001 # proxy not importable outside the server; use local fallback + return fallback_cache + if _proxy_logging is not None: + return _proxy_logging.internal_usage_cache.dual_cache + return fallback_cache + + def _bedrock_response_has_masked_output(self, response: BedrockGuardrailResponse) -> bool: + """Return True if the guardrail rewrote (masked/anonymized) any scanned text. + + Bedrock returns non-empty ``output``/``outputs`` text only when it changed the + content; an ``action == "NONE"`` response leaves both empty. + """ + for field in ("output", "outputs"): + items = response.get(field) or [] + if any(isinstance(item, dict) and item.get("text") for item in items): + return True + return False + + async def _apply_incremental_request_scan( + self, + texts: list[str], + inputs: "GenericGuardrailAPIInputs", + request_data: dict, + ) -> Optional["GenericGuardrailAPIInputs"]: + """Scan only the text segments not already seen earlier in this session. + + Returns ``None`` when incremental scanning is inactive (feature off, no + session id, masking enabled, or cache unavailable) or when the guardrail + turns out to mask content, telling the caller to run the normal full scan. + Otherwise scans only the new segments and skips the Bedrock call entirely + when nothing is new. Incremental mode is for blocking/detection guardrails + only: if the guardrail returns masked output it cannot be applied to the + skipped context, so the scan falls back to the full path and no session + state is recorded. + """ + cache = self._incremental_scan_cache() + + new_texts = await self.filter_new_texts_for_session( + texts=texts, + request_data=request_data, + cache=cache, + ) + if new_texts is None: + return None + + if not new_texts: + verbose_proxy_logger.debug("Bedrock Guardrail: no new messages to scan for this session, skipping API call") + return inputs + + bedrock_response = await self.make_bedrock_api_request( + source="INPUT", + messages=[ChatCompletionUserMessage(role="user", content=text) for text in new_texts], + request_data=request_data, + logging_event_type=GuardrailEventHooks.pre_call, + ) + + if self._bedrock_response_has_masked_output(bedrock_response): + verbose_proxy_logger.warning( + "Bedrock Guardrail %s: guardrail returned masked/anonymized content; " + "only_scan_new_messages cannot apply masking to skipped context, falling back to a full-context scan", + self.guardrail_name, + ) + return None + + await self.mark_texts_scanned( + texts=texts, + request_data=request_data, + cache=cache, + ) + return inputs + async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", @@ -2077,6 +2161,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): try: verbose_proxy_logger.debug(f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)") + if input_type == "request": + incremental_result = await self._apply_incremental_request_scan( + texts=texts, + inputs=inputs, + request_data=request_data, + ) + if incremental_result is not None: + return incremental_result + masked_texts = [] selection = self._select_messages_for_apply_guardrail( diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 14e76a21093..e909c15382b 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -35,6 +35,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): aws_sts_endpoint=litellm_params.aws_sts_endpoint, aws_bedrock_runtime_endpoint=litellm_params.aws_bedrock_runtime_endpoint, experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only, + only_scan_new_messages=litellm_params.only_scan_new_messages or False, ) litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback) return _bedrock_callback diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 47d93fc2d7a..c86794b90f8 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -725,6 +725,18 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up description="When True, guardrails only receive the latest message for the relevant role (e.g., newest user input pre-call, newest assistant output post-call)", ) + only_scan_new_messages: Optional[bool] = Field( + default=False, + description=( + "When True, the guardrail only scans messages that have not already been scanned " + "earlier in the same session (identified by litellm_session_id / session_id). " + "Message content is hashed per session and cached; only the diff (new or edited " + "messages) is sent to the guardrail provider on follow-up calls. Falls back to a " + "full scan when the request has no session id or the cache is unavailable. Intended " + "for blocking/detection guardrails; not applied when mask_request_content is set." + ), + ) + skip_system_message_in_guardrail: Optional[bool] = Field( default=None, description=( diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 9289dece83f..64813c1eda7 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1716,3 +1716,201 @@ class TestApplyGuardrailStyleDeploymentDispatch: await guardrail.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) assert guardrail.apply_called is False + + +class TestOnlyScanNewMessages: + """Incremental guardrail scanning: only send text segments not already scanned this session.""" + + def _guardrail(self, **overrides): + params = dict(guardrail_name="test-guard", only_scan_new_messages=True) + params.update(overrides) + return CustomGuardrail(**params) + + def _cache(self): + from litellm.caching import DualCache + + return DualCache() + + @pytest.mark.asyncio + async def test_disabled_returns_none(self): + guardrail = self._guardrail(only_scan_new_messages=False) + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"litellm_session_id": "s1"}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_no_session_id_fails_safe_to_full_scan(self): + guardrail = self._guardrail() + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"metadata": {}}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_masking_guardrail_not_supported(self): + guardrail = self._guardrail(mask_request_content=True) + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"litellm_session_id": "s1"}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_cache_read_failure_fails_safe_to_full_scan(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail() + cache = self._cache() + cache.async_get_cache = AsyncMock(side_effect=RuntimeError("redis down")) + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"litellm_session_id": "s1"}, + cache=cache, + ) + assert result is None + + @pytest.mark.asyncio + async def test_dedupes_previously_scanned_texts(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-dedupe"} + turn1 = ["you are helpful", "first question"] + + first = await guardrail.filter_new_texts_for_session(texts=turn1, request_data=request, cache=cache) + assert first == turn1 + await guardrail.mark_texts_scanned(texts=turn1, request_data=request, cache=cache) + + turn2 = turn1 + ["an answer", "second question"] + second = await guardrail.filter_new_texts_for_session(texts=turn2, request_data=request, cache=cache) + assert second == ["an answer", "second question"] + + @pytest.mark.asyncio + async def test_no_new_texts_returns_empty(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-empty"} + texts = ["only message"] + + await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=texts, request_data=request, cache=cache) + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == [] + + @pytest.mark.asyncio + async def test_modified_earlier_text_is_rescanned(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-edit"} + original = ["original"] + + await guardrail.filter_new_texts_for_session(texts=original, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=original, request_data=request, cache=cache) + + edited = ["original EDITED"] + result = await guardrail.filter_new_texts_for_session(texts=edited, request_data=request, cache=cache) + assert result == edited + + @pytest.mark.asyncio + async def test_blocked_scan_does_not_persist_hashes(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-blocked"} + texts = ["please block me"] + + filtered = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert filtered == texts + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == texts + + @pytest.mark.asyncio + async def test_scanned_hashes_written_with_fixed_ttl(self): + from unittest.mock import AsyncMock + + from litellm.constants import GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS + + guardrail = self._guardrail() + cache = self._cache() + cache.async_set_cache = AsyncMock() + request = {"litellm_session_id": "sess-ttl"} + + await guardrail.mark_texts_scanned(texts=["a", "b"], request_data=request, cache=cache) + + cache.async_set_cache.assert_awaited_once() + assert cache.async_set_cache.await_args.kwargs["ttl"] == GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS + + @pytest.mark.asyncio + async def test_session_id_from_metadata_is_used_for_dedupe(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"metadata": {"session_id": "sess-meta"}} + texts = ["shared message"] + + await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=texts, request_data=request, cache=cache) + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == [] + + @pytest.mark.asyncio + async def test_session_id_from_litellm_metadata_is_used_for_dedupe(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_metadata": {"session_id": "sess-lmeta"}} + texts = ["shared message"] + + await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=texts, request_data=request, cache=cache) + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == [] + + @pytest.mark.asyncio + async def test_mark_texts_scanned_disabled_does_not_persist(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail(only_scan_new_messages=False) + cache = self._cache() + cache.async_set_cache = AsyncMock() + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache) + cache.async_set_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mark_texts_scanned_masking_does_not_persist(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail(mask_request_content=True) + cache = self._cache() + cache.async_set_cache = AsyncMock() + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache) + cache.async_set_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mark_texts_scanned_without_session_does_not_persist(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail() + cache = self._cache() + cache.async_set_cache = AsyncMock() + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"metadata": {}}, cache=cache) + cache.async_set_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mark_texts_scanned_survives_cache_write_failure(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail() + cache = self._cache() + cache.async_set_cache = AsyncMock(side_effect=RuntimeError("redis down")) + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index c5422e0d70f..9cd1fbb59a6 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -373,3 +373,132 @@ class TestAnthropicMessagesHandlerToolInjection: if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +class TestAnthropicMessagesIncrementalScan: + """PR #33278: only_scan_new_messages through the real /v1/messages translation + handler (the path Claude Code uses). Encodes the wire payloads observed in the + live validation against a real Bedrock guardrail. + """ + + def _bedrock_guardrail(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + return BedrockGuardrail( + guardrail_name="bedrock-incremental-anthropic", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + + def _data(self, messages, session_id): + return { + "model": "claude-sonnet-4-5", + "messages": messages, + "system": "You are a helpful geography assistant.", + "litellm_session_id": session_id, + } + + @pytest.mark.asyncio + async def test_first_turn_scans_all_eligible_then_second_turn_scans_only_diff(self): + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-diff" + turn1 = [{"role": "user", "content": "What is the capital of France?"}] + turn2 = turn1 + [ + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "What is the capital of Germany?"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages( + data=self._data(turn1, sid), guardrail_to_apply=guardrail + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "What is the capital of France?" + ] + mock_api.reset_mock() + await handler.process_input_messages( + data=self._data(turn2, sid), guardrail_to_apply=guardrail + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "Paris.", + "What is the capital of Germany?", + ] + + @pytest.mark.asyncio + async def test_identical_resend_makes_no_guardrail_call(self): + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-resend" + msgs = [ + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "What is the capital of Germany?"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + mock_api.reset_mock() + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + mock_api.assert_not_called() + + @pytest.mark.asyncio + async def test_edited_history_message_is_rescanned(self): + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-edit" + msgs = [{"role": "user", "content": "What is the capital of France?"}] + edited = [{"role": "user", "content": "What is the capital and population of France?"}] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + mock_api.reset_mock() + await handler.process_input_messages(data=self._data(edited, sid), guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "What is the capital and population of France?" + ] + + @pytest.mark.asyncio + async def test_mixed_text_and_tool_use_keeps_text_segments(self): + """A message carrying both text and a tool_use block must not lose its text. + (tool_use inputs and tool_result content are dropped from texts on the + anthropic input path today; that is pre-existing baseline behavior.)""" + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-tools" + msgs = [ + {"role": "user", "content": "Search for the weather in Paris"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me look that up for you."}, + {"type": "tool_use", "id": "toolu_1", "name": "search", "input": {"query": "canary-args"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "canary-result"}], + }, + {"role": "user", "content": "Thanks, summarize the result."}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert "Let me look that up for you." in scanned, "text beside a tool_use must be scanned" + assert "Search for the weather in Paris" in scanned + assert "Thanks, summarize the result." in scanned diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 4c268d9dfc9..7730b664c5e 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1137,3 +1137,95 @@ class TestGetStructuredMessages: if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +class TestIncrementalScanRespectsSkipFlags: + """PR #33278: skip_system_message_in_guardrail and skip_tool_message_in_guardrail + are enforced while this handler builds inputs["texts"] (_extract_inputs early + returns for system/tool roles), upstream of BedrockGuardrail's incremental path. + Bypassing _select_messages_for_apply_guardrail therefore cannot resurrect skipped + content on any turn, including a session's first turn where every segment is new. + Verified live against a real Bedrock ApplyGuardrail before being encoded here. + The flags are set as instance attributes, mirroring how guardrail_registry + applies litellm_params to the callback (they are not constructor kwargs). + """ + + def _bedrock_guardrail(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + guardrail = BedrockGuardrail( + guardrail_name="bedrock-incremental-skip-flags", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + guardrail.skip_system_message_in_guardrail = True + guardrail.skip_tool_message_in_guardrail = True + return guardrail + + def _messages(self, followup=None): + base = [ + {"role": "system", "content": "SYSTEM-PROMPT-must-not-be-scanned"}, + {"role": "user", "content": "Search for the weather in Paris"}, + { + "role": "assistant", + "content": "Let me look that up.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "search", "arguments": '{"query": "weather"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-must-not-be-scanned"}, + {"role": "user", "content": "Thanks, summarize."}, + ] + return base + (followup or []) + + @pytest.mark.asyncio + async def test_first_turn_scans_no_system_or_tool_content(self): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + data = {"messages": self._messages(), "litellm_session_id": "skip-flags-turn1"} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == [ + "Search for the weather in Paris", + "Let me look that up.", + "Thanks, summarize.", + ] + assert not any("SYSTEM-PROMPT" in text for text in scanned) + assert not any("TOOL-RESULT" in text for text in scanned) + + @pytest.mark.asyncio + async def test_second_turn_scans_only_new_eligible_content(self): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + session = "skip-flags-turn2" + followup = [ + {"role": "assistant", "content": "It is sunny in Paris."}, + {"role": "user", "content": "And tomorrow?"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages( + data={"messages": self._messages(), "litellm_session_id": session}, + guardrail_to_apply=guardrail, + ) + mock_api.reset_mock() + await handler.process_input_messages( + data={"messages": self._messages(followup), "litellm_session_id": session}, + guardrail_to_apply=guardrail, + ) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["It is sunny in Paris.", "And tomorrow?"] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 15827b80bcf..53f32fd96fb 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -3274,3 +3274,343 @@ async def test_chat_completion_modify_response_exception_streaming_logging_obj_n # CustomStreamWrapper would raise AttributeError inside __init__ and this # call would never reach here. assert response is not None + + +class TestBedrockOnlyScanNewMessages: + """Bedrock apply_guardrail honors only_scan_new_messages: scans only the per-session diff. + + apply_guardrail is the path the proxy actually runs for Bedrock (via the unified + guardrail interface), so these tests exercise it directly rather than the legacy + async_pre_call_hook. Each test uses a unique session id to isolate the process-wide + incremental cache. + """ + + def _guardrail(self): + return BedrockGuardrail( + guardrail_name="bedrock-incremental", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + + @pytest.mark.asyncio + async def test_second_turn_scans_only_new_messages(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-diff"} + bedrock_none = {"action": "NONE", "output": [], "outputs": []} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = bedrock_none + + await guardrail.apply_guardrail( + inputs={"texts": ["be helpful", "first question"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count == 1 + first_scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in first_scanned] == ["be helpful", "first question"] + + mock_api.reset_mock() + + await guardrail.apply_guardrail( + inputs={"texts": ["be helpful", "first question", "first answer", "second question"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count == 1 + second_scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in second_scanned] == ["first answer", "second question"] + + @pytest.mark.asyncio + async def test_identical_resend_skips_api_call(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-resend"} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + + await guardrail.apply_guardrail( + inputs={"texts": ["only question"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + + mock_api.reset_mock() + result = await guardrail.apply_guardrail( + inputs={"texts": ["only question"]}, request_data=session, input_type="request" + ) + mock_api.assert_not_called() + assert result["texts"] == ["only question"] + + @pytest.mark.asyncio + async def test_no_session_id_scans_full_context(self): + guardrail = self._guardrail() + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + + await guardrail.apply_guardrail( + inputs={"texts": ["q1", "a1", "q2"]}, + request_data={"metadata": {}}, + input_type="request", + ) + assert mock_api.call_count == 1 + scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in scanned] == ["q1", "a1", "q2"] + + @pytest.mark.asyncio + async def test_masking_guardrail_falls_back_and_does_not_persist(self): + """A guardrail that anonymizes content must not be short-circuited. + + Regression: the incremental fast path used to ignore the guardrail response, + so masked/anonymized output was dropped, the raw text reached the model, and + the segment was marked scanned so it was never re-checked. Detecting masked + output must force a full-context scan (which applies the masking) and must not + persist session state, so an identical resend is scanned again. + """ + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-mask"} + masked = { + "action": "GUARDRAIL_INTERVENED", + "output": [], + "outputs": [{"text": "my ssn is [REDACTED]"}], + } + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = masked + + result = await guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123-45-6789"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count == 2 + assert result["texts"] == ["my ssn is [REDACTED]"] + + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123-45-6789"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count >= 1 + first_scanned = mock_api.call_args_list[0].kwargs.get("messages") + assert first_scanned is not None + assert [m["content"] for m in first_scanned] == ["my ssn is 123-45-6789"] + + @pytest.mark.asyncio + async def test_generic_agent_multi_turn_scans_only_new_each_turn(self): + """A generic agent (not Claude Code) opts in by propagating a session id. + + Agent frameworks on the OpenAI SDK carry the session through the request + body (metadata.session_id here), not the x-claude-code-session-id header. + Across a growing multi-turn conversation every turn after the first must + send Bedrock only the newly appended segments, never the whole context. + """ + guardrail = self._guardrail() + session = {"metadata": {"session_id": "agent-multi-turn"}} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + + await guardrail.apply_guardrail( + inputs={"texts": ["system prompt", "turn 1 question"]}, + request_data=session, + input_type="request", + ) + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "system prompt", + "turn 1 question", + ] + + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["system prompt", "turn 1 question", "turn 1 answer", "turn 2 question"]}, + request_data=session, + input_type="request", + ) + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "turn 1 answer", + "turn 2 question", + ] + + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={ + "texts": [ + "system prompt", + "turn 1 question", + "turn 1 answer", + "turn 2 question", + "turn 2 answer", + "turn 3 question", + ] + }, + request_data=session, + input_type="request", + ) + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "turn 2 answer", + "turn 3 question", + ] + + def test_incremental_scan_cache_prefers_proxy_shared_cache(self): + guardrail = self._guardrail() + shared = DualCache() + proxy_logging = MagicMock() + proxy_logging.internal_usage_cache.dual_cache = shared + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging): + assert guardrail._incremental_scan_cache() is shared + + def test_incremental_scan_cache_falls_back_when_proxy_logging_missing(self): + from litellm.integrations.custom_guardrail import dc as fallback_cache + + guardrail = self._guardrail() + with patch("litellm.proxy.proxy_server.proxy_logging_obj", None): + assert guardrail._incremental_scan_cache() is fallback_cache + + def test_incremental_scan_cache_falls_back_when_proxy_not_importable(self): + from litellm.integrations.custom_guardrail import dc as fallback_cache + + guardrail = self._guardrail() + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": None}): + assert guardrail._incremental_scan_cache() is fallback_cache + + @pytest.mark.asyncio + async def test_blocked_turn_is_rescanned_on_retry(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-blocked"} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.side_effect = HTTPException(status_code=400, detail="blocked") + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs={"texts": ["blocked prompt"]}, request_data=session, input_type="request" + ) + + mock_api.reset_mock() + mock_api.side_effect = None + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["blocked prompt"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in scanned] == ["blocked prompt"] + + +class TestBedrockIncrementalFlagInteractions: + """Regression coverage for only_scan_new_messages combined with the other + Bedrock guardrail flags, from the PR #33278 live validation. Live evidence: + each of these was reproduced against a real Bedrock ApplyGuardrail first; + the mocks here encode the wire payloads observed there. + """ + + def _guardrail(self, **overrides): + params = dict( + guardrail_name="bedrock-incremental-flags", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + params.update(overrides) + return BedrockGuardrail(**params) + + @pytest.mark.asyncio + async def test_edited_history_segment_rescans_only_that_segment(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-flags-edit"} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["q1", "a1", "q2"]}, request_data=session, input_type="request" + ) + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["q1 EDITED", "a1", "q2"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == ["q1 EDITED"] + + @pytest.mark.asyncio + async def test_same_content_different_session_rescans_everything(self): + guardrail = self._guardrail() + texts = ["shared question", "shared answer"] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": list(texts)}, request_data={"litellm_session_id": "sess-x1"}, input_type="request" + ) + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": list(texts)}, request_data={"litellm_session_id": "sess-x2"}, input_type="request" + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == texts + + @pytest.mark.asyncio + async def test_litellm_masking_flag_disables_incremental_single_full_scan(self): + """mask_request_content must fall back to exactly ONE full scan per turn + and never persist hashes (verified live: 1 call/turn, no cache writes).""" + guardrail = self._guardrail(mask_request_content=True) + session = {"litellm_session_id": "sess-flags-mask"} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1, "masking mode must re-scan every turn, exactly once" + + @pytest.mark.asyncio + async def test_server_side_anonymize_falls_back_full_scan_and_never_persists(self): + """A guardrail that rewrites content (Bedrock-side ANONYMIZE) must fall back + to the full scan so masking applies, and record no session state. Live + validation showed this costs 2 provider calls per turn; the count is + asserted here as documentation of that intended-tradeoff behavior.""" + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-flags-anon"} + masked = {"action": "NONE", "output": [{"text": "MASKED q1"}], "outputs": [{"text": "MASKED q1"}]} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = masked + result = await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 2, "incremental attempt + full-scan fallback" + assert result["texts"] == ["MASKED q1"], "masked content must be applied" + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 2, "no hashes persisted, so the double scan repeats" + + @pytest.mark.asyncio + @pytest.mark.xfail( + reason="PR #33278 known gap: incremental path bypasses _select_messages_for_apply_guardrail, " + "so experimental_use_latest_role_message_only is silently ignored. Intended semantics " + "(pending DRI decision): incremental mode defers to the latest-role selection.", + strict=False, + ) + async def test_latest_role_only_is_respected_with_incremental(self): + guardrail = self._guardrail(experimental_use_latest_role_message_only=True) + session = {"litellm_session_id": "sess-flags-latestrole"} + structured = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "q1"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["sys", "q1"], "structured_messages": structured}, + request_data=session, + input_type="request", + ) + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["q1"], "latest-role selection must exclude the system prompt" From 17a83aa89665ee5e640c9a032dd6d51b5b127cb5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 10:29:34 -0700 Subject: [PATCH 2/3] fix(proxy): share CLI SSO login sessions across workers without enable_redis_auth_cache (#33261) * fix(proxy): share CLI SSO login sessions across workers without enable_redis_auth_cache * fix(proxy): make CLI SSO flow state redis-authoritative across workers The CLI SSO flow is stored in a DualCache whose get_cache is memory-first, so the worker that served /sso/cli/start keeps serving its stale in-memory flow and never observes the sso_complete/session_data update another worker writes during the OAuth callback. Attaching Redis alone is not enough; poll on the original worker returns pending forever. Read and write the flow directly through the attached Redis backend when present so every worker sees the same authoritative state, falling back to the in-memory DualCache only when no Redis is configured. * fix(proxy): serialize CLI SSO flow as JSON for the redis round trip RedisCache stores values via str(value) and parses reads with json.loads then ast.literal_eval. The completed flow contains a LitellmUserRoles enum in session_data.user_role, whose repr is not a parseable literal, so any worker reading the completed flow from redis raised SyntaxError and returned 400 "CLI login session not found". Writing the flow as json.dumps makes the round trip lossless (the enum is a str subclass) and fails loudly at write time if a non-serializable value is ever added to the flow. * fix(proxy): point CLI SSO session-not-found hint at configuring Redis The error message and warning still told users to set enable_redis_auth_cache, but the CLI SSO session cache now gets Redis unconditionally whenever one is configured, so that flag no longer affects CLI login --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: ryan-crabbe-berri --- litellm/proxy/management_endpoints/ui_sso.py | 56 +++-- litellm/proxy/proxy_server.py | 15 +- .../proxy/management_endpoints/test_ui_sso.py | 197 ++++++++++++++---- .../proxy/test_redis_auth_cache_flag.py | 40 +++- 4 files changed, 240 insertions(+), 68 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 3c8444ecf26..de988a0140f 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -12,6 +12,7 @@ import asyncio import base64 import hashlib import inspect +import json import os import re import secrets @@ -258,11 +259,20 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic raise HTTPException(status_code=400, detail="Invalid CLI login session id") cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id)) - flow = cache.get_cache(key=cache_key) + redis_cache = cache.redis_cache + if redis_cache is not None: + flow = redis_cache.get_cache(key=cache_key) + else: + flow = cache.get_cache(key=cache_key) + if isinstance(flow, str): + try: + flow = json.loads(flow) + except ValueError: + flow = None if not isinstance(flow, dict) or "poll_secret_hash" not in flow: verbose_proxy_logger.warning( "CLI SSO login session not found in cache for login_id=%s. If the proxy runs multiple replicas, " - "a shared Redis cache (enable_redis_auth_cache: true) is required for CLI login to work.", + "a shared Redis cache is required for CLI login to work.", login_id, ) raise HTTPException( @@ -270,7 +280,7 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic detail=( "CLI login session not found or expired. Run `litellm-proxy login` again. " "If this happens immediately after starting a login, the proxy is likely running multiple " - "replicas without a shared cache; configure Redis with `enable_redis_auth_cache: true` " + "replicas without a shared cache; configure a Redis cache " "so every replica can see the login session." ), ) @@ -278,11 +288,12 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None: - cache.set_cache( - key=_get_cli_sso_flow_cache_key(login_id), - value=flow, - ttl=CLI_SSO_SESSION_TTL_SECONDS, - ) + cache_key = _get_cli_sso_flow_cache_key(login_id) + redis_cache = cache.redis_cache + if redis_cache is not None: + redis_cache.set_cache(key=cache_key, value=json.dumps(flow), ttl=CLI_SSO_SESSION_TTL_SECONDS) + else: + cache.set_cache(key=cache_key, value=flow, ttl=CLI_SSO_SESSION_TTL_SECONDS) def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool: @@ -593,11 +604,11 @@ def _render_cli_sso_verification_page( @router.post("/sso/cli/start", tags=["experimental"], include_in_schema=False) async def cli_sso_start(request: Request): - from litellm.proxy.proxy_server import general_settings, user_api_key_cache + from litellm.proxy.proxy_server import cli_sso_session_cache, general_settings _check_cli_sso_start_rate_limit( request=request, - cache=user_api_key_cache, + cache=cli_sso_session_cache, use_x_forwarded_for=bool((general_settings or {}).get("use_x_forwarded_for", False)), ) @@ -612,7 +623,7 @@ async def cli_sso_start(request: Request): "user_code_verified": False, "session_data": None, } - _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) + _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow) verification_uri_complete: str | None = ( ( @@ -644,9 +655,9 @@ async def cli_sso_complete(request: Request, login_id: str): from litellm.proxy.common_utils.html_forms.cli_sso_success import ( render_cli_sso_success_page, ) - from litellm.proxy.proxy_server import user_api_key_cache + from litellm.proxy.proxy_server import cli_sso_session_cache - flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=user_api_key_cache) + flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cli_sso_session_cache) if not flow.get("sso_complete") or not flow.get("session_data"): raise HTTPException(status_code=400, detail="CLI login is not ready") @@ -670,7 +681,7 @@ async def cli_sso_complete(request: Request, login_id: str): raise HTTPException(status_code=400, detail="Invalid verification code") flow["user_code_verified"] = True - _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) + _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow) html_content = render_cli_sso_success_page() return HTMLResponse(content=html_content, status_code=200) @@ -861,10 +872,10 @@ async def google_login( Example: """ from litellm.proxy.proxy_server import ( + cli_sso_session_cache, general_settings, premium_user, prisma_client, - user_api_key_cache, user_custom_ui_sso_sign_in_handler, ) @@ -912,7 +923,7 @@ async def google_login( ) if source == LITELLM_CLI_SOURCE_IDENTIFIER: - _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) + _get_cli_sso_flow_or_raise(login_id=key, cache=cli_sso_session_cache) # Store CLI login handle in state for OAuth flow cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( @@ -1957,6 +1968,7 @@ async def _complete_cli_sso_callback_session( user_defined_values: Optional[SSOUserDefinedValues], prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, + cli_sso_session_cache: DualCache, proxy_logging_obj: ProxyLogging, prefill_user_code: str | None = None, sso_assertion: SSOIdentityAssertion | None = None, @@ -2006,7 +2018,7 @@ async def _complete_cli_sso_callback_session( flow["sso_complete"] = True browser_complete_token = secrets.token_urlsafe(32) flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token) - _set_cli_sso_flow(login_id=key, cache=user_api_key_cache, flow=flow) + _set_cli_sso_flow(login_id=key, cache=cli_sso_session_cache, flow=flow) verbose_proxy_logger.info( f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}" @@ -2037,13 +2049,14 @@ async def cli_sso_callback( verbose_proxy_logger.info("CLI SSO callback") from litellm.proxy.proxy_server import ( + cli_sso_session_cache, general_settings, prisma_client, proxy_logging_obj, user_api_key_cache, ) - flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) + flow = _get_cli_sso_flow_or_raise(login_id=key, cache=cli_sso_session_cache) if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) @@ -2083,6 +2096,7 @@ async def cli_sso_callback( user_defined_values=user_defined_values, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + cli_sso_session_cache=cli_sso_session_cache, proxy_logging_obj=proxy_logging_obj, prefill_user_code=prefill_user_code, sso_assertion=sso_assertion, @@ -2114,10 +2128,10 @@ async def cli_poll_key( team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams. """ from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken - from litellm.proxy.proxy_server import user_api_key_cache + from litellm.proxy.proxy_server import cli_sso_session_cache try: - flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache) + flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=cli_sso_session_cache) if not _verify_cli_sso_poll_secret(flow=flow, poll_secret=x_litellm_cli_poll_secret): raise HTTPException(status_code=403, detail="Invalid CLI polling secret") @@ -2192,7 +2206,7 @@ async def cli_poll_key( ) # Delete cache entry (single-use) - user_api_key_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id)) + cli_sso_session_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id)) verbose_proxy_logger.info(f"CLI JWT generated for user: {user_id}, team: {team_id}") poll_response = { diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6de3e43fc1a..d4bff81ea6a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -226,6 +226,7 @@ from litellm.constants import ( APSCHEDULER_MAX_INSTANCES, APSCHEDULER_MISFIRE_GRACE_TIME, APSCHEDULER_REPLACE_EXISTING, + CLI_SSO_SESSION_TTL_SECONDS, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, @@ -1970,6 +1971,7 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) spend_counter_cache = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value) +cli_sso_session_cache = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS) model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=user_api_key_cache) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) redis_usage_cache: Optional[RedisCache] = None # redis cache used for tracking spend, tpm/rpm limits @@ -3696,13 +3698,22 @@ def _build_redis_usage_cache_from_environment() -> RedisCache | None: def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: bool) -> None: """ Wires an established coordination Redis into the proxy-level caches that - consume it directly: the spend counter cache, the cluster-wide config - cache, and (only when opted in) the virtual-key auth cache. + consume it directly: the spend counter cache, the CLI SSO login-session + cache, the cluster-wide config cache, and (only when opted in) the + virtual-key auth cache. + + The CLI SSO login-session cache is always backed by Redis when available so + that the browser SSO flow behind `lite login` survives landing on different + workers; it must not be gated behind enable_redis_auth_cache. """ spend_counter_cache.attach_redis_cache( redis_cache, default_redis_ttl=litellm.default_redis_ttl, ) + cli_sso_session_cache.attach_redis_cache( + redis_cache, + default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS, + ) if enable_redis_auth_cache is True: user_api_key_cache.attach_redis_cache( redis_cache, diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index e1856860c8a..47ceb0c05fa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2214,7 +2214,95 @@ class TestCLIKeyRegenerationFlow: _get_cli_sso_flow_or_raise(login_id="cli-test_1234567890", cache=mock_cache) assert expired_exc.value.status_code == 400 assert "session not found or expired" in expired_exc.value.detail - assert "enable_redis_auth_cache" in expired_exc.value.detail + assert "configure a Redis cache" in expired_exc.value.detail + assert "enable_redis_auth_cache" not in expired_exc.value.detail + + def test_cli_sso_flow_is_redis_authoritative_when_redis_attached(self): + """ + When Redis is attached, the CLI SSO flow must be read from and written to + Redis directly, never the in-memory layer. Otherwise the worker that served + /sso/cli/start keeps serving its stale in-memory flow and never sees the + sso_complete/session_data update another worker wrote, which is exactly the + multi-worker failure this fix targets. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + CLI_SSO_SESSION_TTL_SECONDS, + _get_cli_sso_flow_cache_key, + _get_cli_sso_flow_or_raise, + _set_cli_sso_flow, + ) + + login_id = "cli-redis_authoritative_1234567890" + cache_key = _get_cli_sso_flow_cache_key(login_id) + fresh_flow = {"poll_secret_hash": "fresh", "sso_complete": True} + stale_flow = {"poll_secret_hash": "stale", "sso_complete": False} + + redis_cache = MagicMock() + redis_cache.get_cache.return_value = fresh_flow + cache = MagicMock() + cache.redis_cache = redis_cache + cache.get_cache.return_value = stale_flow + + result = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cache) + + assert result == fresh_flow + redis_cache.get_cache.assert_called_once_with(key=cache_key) + cache.get_cache.assert_not_called() + + _set_cli_sso_flow(login_id=login_id, cache=cache, flow=fresh_flow) + + redis_cache.set_cache.assert_called_once_with( + key=cache_key, value=json.dumps(fresh_flow), ttl=CLI_SSO_SESSION_TTL_SECONDS + ) + cache.set_cache.assert_not_called() + + def test_cli_sso_flow_with_enum_survives_redis_round_trip(self): + """ + RedisCache stores values via str(value) and reads them back through + json.loads/ast.literal_eval. A raw flow dict containing a Python enum + (session_data.user_role after the SSO callback) produces an unparseable + repr, so every worker reading the completed flow from Redis got a + SyntaxError and returned 400 "session not found". The flow must survive + a real Redis serialization round trip. + """ + from litellm.caching.redis_cache import RedisCache + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_or_raise, + _set_cli_sso_flow, + ) + + login_id = "cli-enum_round_trip_1234567890" + completed_flow = { + "poll_secret_hash": "hash", + "sso_complete": True, + "user_code_verified": False, + "session_data": { + "user_id": "user-1", + "user_role": LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + "models": [], + "teams": ["team-1"], + "team_details": [{"team_id": "team-1", "team_alias": "alias"}], + }, + } + + redis_store: dict = {} + redis_cache = MagicMock() + redis_cache.set_cache.side_effect = lambda key, value, ttl: redis_store.__setitem__( + key, str(value).encode("utf-8") + ) + redis_cache.get_cache.side_effect = lambda key: RedisCache._get_cache_logic( + MagicMock(), redis_store.get(key) + ) + cache = MagicMock() + cache.redis_cache = redis_cache + + _set_cli_sso_flow(login_id=login_id, cache=cache, flow=completed_flow) + flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cache) + + assert flow["sso_complete"] is True + assert flow["session_data"]["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value + assert flow["session_data"]["team_details"] == [{"team_id": "team-1", "team_alias": "alias"}] @pytest.mark.asyncio async def test_cli_sso_start_creates_bound_flow(self): @@ -2228,10 +2316,13 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 1 - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): result = await cli_sso_start(request=mock_request) assert result["login_id"].startswith("cli-") @@ -2259,10 +2350,13 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 31 - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_sso_start(request=mock_request) @@ -2281,7 +2375,7 @@ class TestCLIKeyRegenerationFlow: mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} mock_request.base_url = "https://proxy.example.com/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 1 with ( @@ -2315,7 +2409,7 @@ class TestCLIKeyRegenerationFlow: mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} mock_request.base_url = "https://proxy.example.com/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 1 with ( @@ -2349,7 +2443,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.base_url = "https://proxy.example.com/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = {"poll_secret_hash": "h"} async def drive(enabled: bool): @@ -2358,6 +2452,7 @@ class TestCLIKeyRegenerationFlow: patch("litellm.proxy.proxy_server.premium_user", True), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler", None, @@ -2525,7 +2620,7 @@ class TestCLIKeyRegenerationFlow: ) mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"} - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -2544,6 +2639,7 @@ class TestCLIKeyRegenerationFlow: ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), ): result = await cli_sso_callback( request=mock_request, @@ -2568,7 +2664,7 @@ class TestCLIKeyRegenerationFlow: mock_request.body = AsyncMock( return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( @@ -2582,6 +2678,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", return_value="Success", @@ -2606,7 +2703,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.body = AsyncMock(return_value=b"user_code=ABCD-EFGH") - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( @@ -2618,7 +2715,10 @@ class TestCLIKeyRegenerationFlow: "session_data": {"user_id": "test-user-123"}, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_sso_complete( request=mock_request, login_id="cli-session-4567890" @@ -2640,7 +2740,7 @@ class TestCLIKeyRegenerationFlow: mock_request.body = AsyncMock( return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( @@ -2651,7 +2751,10 @@ class TestCLIKeyRegenerationFlow: "session_data": None, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_sso_complete( request=mock_request, login_id="cli-session-4567890" @@ -2687,7 +2790,7 @@ class TestCLIKeyRegenerationFlow: mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"} # Mock cache - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -2709,6 +2812,7 @@ class TestCLIKeyRegenerationFlow: ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", return_value="Success", @@ -2769,7 +2873,7 @@ class TestCLIKeyRegenerationFlow: } # Mock cache - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -2777,7 +2881,10 @@ class TestCLIKeyRegenerationFlow: "session_data": session_data, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): # Act - First poll without team_id result = await cli_poll_key( key_id=session_key, @@ -2803,7 +2910,7 @@ class TestCLIKeyRegenerationFlow: cli_poll_key, ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -2816,7 +2923,10 @@ class TestCLIKeyRegenerationFlow: }, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_poll_key(key_id="cli-session-789123", team_id=None) @@ -2830,7 +2940,7 @@ class TestCLIKeyRegenerationFlow: cli_poll_key, ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -2843,7 +2953,10 @@ class TestCLIKeyRegenerationFlow: }, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): result = await cli_poll_key( key_id="cli-session-789123", team_id=None, @@ -3011,7 +3124,7 @@ class TestCLIKeyRegenerationFlow: ) # Mock cache - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -3023,6 +3136,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", @@ -3086,7 +3200,7 @@ class TestCLIKeyRegenerationFlow: models=["gpt-4"], max_budget=100.0, ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -3097,6 +3211,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", @@ -3142,7 +3257,7 @@ class TestCLIKeyRegenerationFlow: "models": ["gpt-4"], "user_email": "unbudgeted@example.com", } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -3153,6 +3268,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", return_value=mock_jwt_token, @@ -4082,7 +4198,7 @@ class TestPKCEFunctionality: mock_request.query_params = {"state": test_state} # Mock cache with async methods — use dict format (primary path) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) test_code_verifier = "test_code_verifier_abc123xyz" mock_cache.async_get_cache = AsyncMock( return_value={"code_verifier": test_code_verifier} @@ -4133,7 +4249,7 @@ class TestPKCEFunctionality: mock_sso.__exit__ = MagicMock(return_value=False) test_state = "test456" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_set_cache = AsyncMock() @@ -4657,7 +4773,7 @@ class TestPKCEFunctionality: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=None) # verifier not found mock_request = MagicMock(spec=Request) @@ -4783,7 +4899,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Cache returns an integer — unexpected format - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=12345) mock_cache.async_delete_cache = AsyncMock() @@ -4825,7 +4941,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=None) # verifier not found mock_request = MagicMock(spec=Request) @@ -4913,7 +5029,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Cache returns an integer — unexpected format - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=12345) mock_cache.async_delete_cache = AsyncMock() @@ -4965,7 +5081,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler legacy_verifier = "legacy_plain_string_verifier_abc123" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=legacy_verifier) mock_request = MagicMock(spec=Request) @@ -6249,7 +6365,7 @@ class TestCliSsoAttributionMetadata: provider="generic", team_ids=[], ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -6266,6 +6382,7 @@ class TestCliSsoAttributionMetadata: ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.user_custom_sso", None), ): await ui_sso.cli_sso_callback( @@ -6290,7 +6407,7 @@ class TestCliSsoAttributionMetadata: mock_request = MagicMock(spec=Request) mock_request.base_url = "http://internal-proxy.local/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -6313,6 +6430,7 @@ class TestCliSsoAttributionMetadata: ) as get_user_info_mock, patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.user_custom_sso", None), patch( "litellm.proxy.proxy_server.general_settings", @@ -6359,7 +6477,7 @@ class TestCliSsoAttributionMetadata: "user_id": "test-user-123", "employment_type": "contractor", } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -6387,6 +6505,7 @@ class TestCliSsoAttributionMetadata: ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.user_custom_sso", None), patch( "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", @@ -6428,7 +6547,7 @@ class TestCliSsoAttributionMetadata: "org": {"cost_center": "CC-42"}, }, } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -6436,7 +6555,10 @@ class TestCliSsoAttributionMetadata: "session_data": session_data, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): result = await cli_poll_key( key_id=session_key, team_id=None, @@ -7287,7 +7409,7 @@ async def test_cli_poll_key_tolerates_missing_user_row(): "models": ["gpt-4"], } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -7299,6 +7421,7 @@ async def test_cli_poll_key_tolerates_missing_user_row(): with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py index d0cb5ec5465..849d5494c6e 100644 --- a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -54,8 +54,8 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): _FakeRedisCache (passes the isinstance guard in _init_cache). 3. Extracts enable_redis_auth_cache from litellm_settings and passes it as the second argument to _init_cache (matching production behaviour). - 4. Yields (user_api_key_cache, spend_counter_cache) after calling - _init_cache, then restores everything. + 4. Yields (user_api_key_cache, spend_counter_cache, cli_sso_session_cache) + after calling _init_cache, then restores everything. """ fake_redis = _FakeRedisCache() @@ -64,19 +64,21 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): fresh_user_cache = DualCache() fresh_spend_cache = DualCache() + fresh_cli_sso_cache = DualCache() enable_redis_auth_cache = litellm_settings.get("enable_redis_auth_cache", False) with ( patch.object(ps, "user_api_key_cache", fresh_user_cache), patch.object(ps, "spend_counter_cache", fresh_spend_cache), + patch.object(ps, "cli_sso_session_cache", fresh_cli_sso_cache), patch.object(ps, "llm_router", None), # Cache is locally imported inside _init_cache: patch it at source. patch("litellm.Cache", return_value=mock_litellm_cache), ): litellm.cache = None ps.ProxyConfig()._init_cache(cache_params, enable_redis_auth_cache) - yield fresh_user_cache, fresh_spend_cache + yield fresh_user_cache, fresh_spend_cache, fresh_cli_sso_cache # --------------------------------------------------------------------------- @@ -90,7 +92,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={"enable_redis_auth_cache": True}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, _): + ) as (user_cache, _, _cli_sso_cache): assert user_cache.redis_cache is not None, ( "Redis should be attached to user_api_key_cache when " "enable_redis_auth_cache=True" @@ -101,7 +103,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={"enable_redis_auth_cache": False}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, _): + ) as (user_cache, _, _cli_sso_cache): assert user_cache.redis_cache is None, ( "user_api_key_cache must remain in-memory-only when " "enable_redis_auth_cache=False" @@ -112,7 +114,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, _): + ) as (user_cache, _, _cli_sso_cache): assert user_cache.redis_cache is None, ( "user_api_key_cache must remain in-memory-only when " "enable_redis_auth_cache is absent from litellm_settings" @@ -129,7 +131,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings=ls, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (_, spend_cache): + ) as (_, spend_cache, _cli_sso_cache): assert spend_cache.redis_cache is not None, ( f"spend_counter_cache must always get Redis " f"(enable_redis_auth_cache={flag_value!r})" @@ -140,6 +142,28 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={"enable_redis_auth_cache": False}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, spend_cache): + ) as (user_cache, spend_cache, _cli_sso_cache): assert spend_cache.redis_cache is not None assert user_cache.redis_cache is None + + def test_cli_sso_session_cache_always_gets_redis_regardless_of_flag(self): + """ + cli_sso_session_cache must receive Redis regardless of the auth-cache + flag so that `lite login` works on multi-worker deployments without + enable_redis_auth_cache (regression for the CLI SSO "Invalid CLI login + session" bug) + """ + for flag_value in (True, False, None): + ls = ( + {"enable_redis_auth_cache": flag_value} + if flag_value is not None + else {} + ) + with _patched_init_cache( + litellm_settings=ls, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (_, _, cli_sso_cache): + assert cli_sso_cache.redis_cache is not None, ( + f"cli_sso_session_cache must always get Redis " + f"(enable_redis_auth_cache={flag_value!r})" + ) From 0fcaadf11ca1676f5f3e041808caf837fdb71ee3 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 22 Jul 2026 12:43:10 -0700 Subject: [PATCH 3/3] test(e2e): move Admin UI Playwright suite to tests/e2e/ui (#34196) Relocates ui/litellm-dashboard/e2e_tests to tests/e2e/ui so all end to end suites live under tests/e2e. The suite stays in TypeScript and becomes a self-contained npm package with its own package.json, lockfile and tsconfig instead of leaning on the dashboard's toolchain; the dashboard drops its @playwright/test dependency, e2e scripts and knip/vitest/tsconfig carve-outs. CI paths follow the move: both CircleCI jobs (main e2e and the SERVER_ROOT_PATH migration smoke) and the test_server_root_path workflow now install and run Playwright from tests/e2e/ui, with the node cache keyed on both lockfiles. classify_changes.sh treats tests/e2e/ui as client so spec edits keep skipping backend jobs. The suite's mock LLM fixture is excluded from the e2e basedpyright zero-error gate in pyrightconfig.json since it belongs to the TS suite, not the typed Python harness. --- .circleci/config.yml | 42 ++++--- .circleci/scripts/classify_changes.sh | 2 +- .github/workflows/test_server_root_path.yml | 10 +- pyrightconfig.json | 2 +- tests/e2e/CLAUDE.md | 1 + tests/e2e/load/test_session_anomaly.py | 4 +- .../e2e_tests => tests/e2e/ui}/constants.ts | 0 .../e2e/ui}/fixtures/config.yml | 0 .../e2e/ui}/fixtures/menuMappings.ts | 0 .../e2e/ui}/fixtures/migratedPages.ts | 0 .../ui}/fixtures/mock_llm_server/server.py | 0 .../e2e/ui}/fixtures/pages.ts | 0 .../e2e/ui}/fixtures/roles.ts | 0 .../e2e/ui}/fixtures/seed.sql | 0 .../e2e/ui}/fixtures/users.ts | 0 .../e2e_tests => tests/e2e/ui}/globalSetup.ts | 0 .../e2e/ui}/helpers/navigation.ts | 0 .../ui}/migration.serverRootPath.config.ts | 0 .../migration.serverRootPath.globalSetup.ts | 0 tests/e2e/ui/package-lock.json | 111 ++++++++++++++++++ tests/e2e/ui/package.json | 16 +++ .../e2e/ui}/playwright.config.ts | 0 .../e2e_tests => tests/e2e/ui}/run_e2e.sh | 6 +- .../e2e/ui}/serverRootPath.config.ts | 0 .../e2e/ui}/tests/auth/logout.spec.ts | 0 .../e2e/ui}/tests/auth/proxyLogoutUrl.spec.ts | 0 .../auth/unauthenticatedRedirect.spec.ts | 0 .../tests/internal-user/internalUser.spec.ts | 0 .../internal-user/internalUserNoTeam.spec.ts | 0 .../internalUserWithTeams.spec.ts | 0 .../internal-viewer/internalViewer.spec.ts | 0 .../tests/login/internalUserIdentity.spec.ts | 0 .../e2e/ui}/tests/login/login.spec.ts | 0 .../login/serverRootPathRedirect.spec.ts | 0 .../e2e/ui}/tests/mcp/mcpServers.spec.ts | 0 .../e2e/ui}/tests/migration/README.md | 5 +- .../ui}/tests/migration/migratedPages.spec.ts | 0 .../e2e/ui}/tests/modelHub/modelHub.spec.ts | 0 .../e2e/ui}/tests/modelsPage/addModel.spec.ts | 0 .../modelsPage/clearCustomPricing.spec.ts | 0 .../ui}/tests/modelsPage/credentials.spec.ts | 0 .../e2e/ui}/tests/navigation/sidebar.spec.ts | 0 .../e2e/ui}/tests/proxy-admin/keys.spec.ts | 0 .../e2e/ui}/tests/proxy-admin/license.spec.ts | 0 .../e2e/ui}/tests/proxy-admin/teams.spec.ts | 0 .../ui}/tests/settings/adminSettings.spec.ts | 0 .../ui}/tests/settings/routerSettings.spec.ts | 2 +- .../ui}/tests/team-admin/teamAdmin.spec.ts | 0 .../e2e/ui}/tests/users/searchUsers.spec.ts | 0 .../ui}/tests/users/viewInternalUsers.spec.ts | 0 tests/e2e/ui/tsconfig.json | 16 +++ .../proxy/management_endpoints/test_ui_sso.py | 1 + ui/litellm-dashboard/knip.json | 10 +- ui/litellm-dashboard/package-lock.json | 11 +- ui/litellm-dashboard/package.json | 5 - ui/litellm-dashboard/tsconfig.json | 2 +- ui/litellm-dashboard/vitest.config.ts | 3 +- 57 files changed, 194 insertions(+), 55 deletions(-) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/constants.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/config.yml (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/menuMappings.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/migratedPages.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/mock_llm_server/server.py (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/pages.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/roles.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/seed.sql (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/users.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/globalSetup.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/helpers/navigation.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/migration.serverRootPath.config.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/migration.serverRootPath.globalSetup.ts (100%) create mode 100644 tests/e2e/ui/package-lock.json create mode 100644 tests/e2e/ui/package.json rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/playwright.config.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/run_e2e.sh (98%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/serverRootPath.config.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/auth/logout.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/auth/proxyLogoutUrl.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/auth/unauthenticatedRedirect.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/internal-user/internalUser.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/internal-user/internalUserNoTeam.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/internal-user/internalUserWithTeams.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/internal-viewer/internalViewer.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/login/internalUserIdentity.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/login/login.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/login/serverRootPathRedirect.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/mcp/mcpServers.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/migration/README.md (84%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/migration/migratedPages.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/modelHub/modelHub.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/modelsPage/addModel.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/modelsPage/clearCustomPricing.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/modelsPage/credentials.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/navigation/sidebar.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/proxy-admin/keys.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/proxy-admin/license.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/proxy-admin/teams.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/settings/adminSettings.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/settings/routerSettings.spec.ts (98%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/team-admin/teamAdmin.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/users/searchUsers.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/users/viewInternalUsers.spec.ts (100%) create mode 100644 tests/e2e/ui/tsconfig.json diff --git a/.circleci/config.yml b/.circleci/config.yml index b0a705966a2..2f01b6de4f3 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2731,7 +2731,7 @@ jobs: - ~/.cache/uv - restore_cache: keys: - - ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + - ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} - run: name: Install Node dependencies and Playwright # The cimg/python:3.12-browsers image already ships the Chromium system @@ -2742,11 +2742,14 @@ jobs: command: | cd ui/litellm-dashboard npm ci + cd ../../tests/e2e/ui + npm ci npx playwright install chromium - save_cache: - key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} paths: - ui/litellm-dashboard/node_modules + - tests/e2e/ui/node_modules - ~/.cache/ms-playwright - run: name: Build UI from source @@ -2777,10 +2780,10 @@ jobs: name: Seed database command: | PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \ - -f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql + -f tests/e2e/ui/fixtures/seed.sql - run: name: Start mock LLM server - command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py + command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py background: true - run: name: Start LiteLLM proxy @@ -2798,7 +2801,7 @@ jobs: command: | LITELLM_LICENSE="$LITELLM_LICENSE" \ uv run --no-sync python -m litellm.proxy.proxy_cli \ - --config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \ + --config tests/e2e/ui/fixtures/config.yml \ --port 4000 background: true - run: @@ -2819,15 +2822,15 @@ jobs: # Forward LITELLM_LICENSE so license.spec.ts can detect that the # proxy was launched with a license and assert premium_user=true. command: | - cd ui/litellm-dashboard + cd tests/e2e/ui LITELLM_LICENSE="$LITELLM_LICENSE" \ - npx playwright test --config e2e_tests/playwright.config.ts + npx playwright test --config playwright.config.ts no_output_timeout: 10m - store_artifacts: - path: ui/litellm-dashboard/test-results + path: tests/e2e/ui/test-results destination: e2e-test-results - store_artifacts: - path: ui/litellm-dashboard/playwright-report + path: tests/e2e/ui/playwright-report destination: e2e-playwright-report e2e_ui_testing_server_root_path: @@ -2870,17 +2873,20 @@ jobs: - ~/.cache/uv - restore_cache: keys: - - ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + - ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} - run: name: Install Node dependencies and Playwright command: | cd ui/litellm-dashboard npm ci + cd ../../tests/e2e/ui + npm ci npx playwright install chromium - save_cache: - key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} paths: - ui/litellm-dashboard/node_modules + - tests/e2e/ui/node_modules - ~/.cache/ms-playwright - run: name: Build UI from source @@ -2902,10 +2908,10 @@ jobs: name: Seed database command: | PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \ - -f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql + -f tests/e2e/ui/fixtures/seed.sql - run: name: Start mock LLM server - command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py + command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py background: true - run: name: Start LiteLLM proxy under a server root path @@ -2918,7 +2924,7 @@ jobs: command: | LITELLM_LICENSE="$LITELLM_LICENSE" \ uv run --no-sync python -m litellm.proxy.proxy_cli \ - --config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \ + --config tests/e2e/ui/fixtures/config.yml \ --port 4000 background: true - run: @@ -2937,15 +2943,15 @@ jobs: - run: name: Run migration smoke under SERVER_ROOT_PATH command: | - cd ui/litellm-dashboard + cd tests/e2e/ui LITELLM_LICENSE="$LITELLM_LICENSE" \ - npx playwright test --config e2e_tests/migration.serverRootPath.config.ts + npx playwright test --config migration.serverRootPath.config.ts no_output_timeout: 10m - store_artifacts: - path: ui/litellm-dashboard/test-results + path: tests/e2e/ui/test-results destination: e2e-server-root-path-test-results - store_artifacts: - path: ui/litellm-dashboard/playwright-report + path: tests/e2e/ui/playwright-report destination: e2e-server-root-path-playwright-report build_docker_database_image: diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index 2c15428be6a..2ca2654a207 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -8,7 +8,7 @@ has_backend=false while IFS= read -r file || [ -n "$file" ]; do [ -n "$file" ] || continue case "$file" in - ui/*) has_client=true ;; + ui/* | tests/e2e/ui/*) has_client=true ;; docs/* | *.md | *.mdx) : ;; *) has_backend=true ;; esac diff --git a/.github/workflows/test_server_root_path.yml b/.github/workflows/test_server_root_path.yml index f59cee29893..01f70511e79 100644 --- a/.github/workflows/test_server_root_path.yml +++ b/.github/workflows/test_server_root_path.yml @@ -106,8 +106,8 @@ jobs: with: node-version: "20" - - name: Install UI deps and Chromium - working-directory: ui/litellm-dashboard + - name: Install e2e deps and Chromium + working-directory: tests/e2e/ui run: | retry() { local attempt=1 @@ -131,17 +131,17 @@ jobs: retry npx playwright install --with-deps chromium - name: Run SERVER_ROOT_PATH redirect e2e - working-directory: ui/litellm-dashboard + working-directory: tests/e2e/ui env: SERVER_ROOT_PATH: ${{ matrix.root_path }} - run: npx playwright test --config=e2e_tests/serverRootPath.config.ts + run: npx playwright test --config=serverRootPath.config.ts - name: Upload Playwright artifacts on failure if: failure() uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 with: name: playwright-trace-${{ strategy.job-index }} - path: ui/litellm-dashboard/test-results/ + path: tests/e2e/ui/test-results/ retention-days: 7 - name: Cleanup diff --git a/pyrightconfig.json b/pyrightconfig.json index eabfbf515c4..2686ccd73d9 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -1,7 +1,7 @@ { "include": ["litellm"], "ignore": [], - "exclude": ["**/node_modules", "**/__pycache__", "tests/e2e/claude_code", "litellm/types/utils.py", "litellm/proxy/_types.py"], + "exclude": ["**/node_modules", "**/__pycache__", "tests/e2e/claude_code", "tests/e2e/ui", "litellm/types/utils.py", "litellm/proxy/_types.py"], "pythonVersion": "3.12", "typeCheckingMode": "strict", "enableTypeIgnoreComments": false, diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 186517290ea..0e39664e358 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -22,6 +22,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke +- `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` ## MCP suite: real Datadog only diff --git a/tests/e2e/load/test_session_anomaly.py b/tests/e2e/load/test_session_anomaly.py index 80f8aff3ba4..7062587352b 100644 --- a/tests/e2e/load/test_session_anomaly.py +++ b/tests/e2e/load/test_session_anomaly.py @@ -62,7 +62,7 @@ class TestSummarizePlannedTurns: class TestRetried: def test_transient_failures_then_success_returns_the_success(self) -> None: - outcome = Success(data=SessionMessagesResponse()) + outcome = Success[SessionMessagesResponse](status_code=200, data=SessionMessagesResponse()) calls = iter( (NetworkError(message="overloaded"), NetworkError(message="overloaded"), outcome) ) @@ -88,7 +88,7 @@ class TestRetried: raise AssertionError("slept after a successful attempt") result = retried( - lambda: Success(data=SessionMessagesResponse()), + lambda: Success[SessionMessagesResponse](status_code=200, data=SessionMessagesResponse()), attempts=3, sleep=sleep_means_retry, ) diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/tests/e2e/ui/constants.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/constants.ts rename to tests/e2e/ui/constants.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/config.yml b/tests/e2e/ui/fixtures/config.yml similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/config.yml rename to tests/e2e/ui/fixtures/config.yml diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts b/tests/e2e/ui/fixtures/menuMappings.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts rename to tests/e2e/ui/fixtures/menuMappings.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts rename to tests/e2e/ui/fixtures/migratedPages.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py b/tests/e2e/ui/fixtures/mock_llm_server/server.py similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py rename to tests/e2e/ui/fixtures/mock_llm_server/server.py diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/pages.ts b/tests/e2e/ui/fixtures/pages.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/pages.ts rename to tests/e2e/ui/fixtures/pages.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts b/tests/e2e/ui/fixtures/roles.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/roles.ts rename to tests/e2e/ui/fixtures/roles.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/seed.sql b/tests/e2e/ui/fixtures/seed.sql similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/seed.sql rename to tests/e2e/ui/fixtures/seed.sql diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/users.ts b/tests/e2e/ui/fixtures/users.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/users.ts rename to tests/e2e/ui/fixtures/users.ts diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/tests/e2e/ui/globalSetup.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/globalSetup.ts rename to tests/e2e/ui/globalSetup.ts diff --git a/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/helpers/navigation.ts rename to tests/e2e/ui/helpers/navigation.ts diff --git a/ui/litellm-dashboard/e2e_tests/migration.serverRootPath.config.ts b/tests/e2e/ui/migration.serverRootPath.config.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/migration.serverRootPath.config.ts rename to tests/e2e/ui/migration.serverRootPath.config.ts diff --git a/ui/litellm-dashboard/e2e_tests/migration.serverRootPath.globalSetup.ts b/tests/e2e/ui/migration.serverRootPath.globalSetup.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/migration.serverRootPath.globalSetup.ts rename to tests/e2e/ui/migration.serverRootPath.globalSetup.ts diff --git a/tests/e2e/ui/package-lock.json b/tests/e2e/ui/package-lock.json new file mode 100644 index 00000000000..b22673a3535 --- /dev/null +++ b/tests/e2e/ui/package-lock.json @@ -0,0 +1,111 @@ +{ + "name": "litellm-ui-e2e", + "version": "0.0.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "litellm-ui-e2e", + "version": "0.0.0", + "devDependencies": { + "@playwright/test": "1.58.1", + "@types/node": "20.19.37", + "typescript": "5.9.3" + } + }, + "node_modules/@playwright/test": { + "version": "1.58.1", + "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.58.1.tgz", + "integrity": "sha512-6LdVIUERWxQMmUSSQi0I53GgCBYgM2RpGngCPY7hSeju+VrKjq3lvs7HpJoPbDiY5QM5EYRtRX5fvrinnMAz3w==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "playwright": "1.58.1" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/@types/node": { + "version": "20.19.37", + "resolved": "https://registry.npmjs.org/@types/node/-/node-20.19.37.tgz", + "integrity": "sha512-8kzdPJ3FsNsVIurqBs7oodNnCEVbni9yUEkaHbgptDACOPW04jimGagZ51E6+lXUwJjgnBw+hyko/lkFWCldqw==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~6.21.0" + } + }, + "node_modules/fsevents": { + "version": "2.3.2", + "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", + "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^8.16.0 || ^10.6.0 || >=11.0.0" + } + }, + "node_modules/playwright": { + "version": "1.58.1", + "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.58.1.tgz", + "integrity": "sha512-+2uTZHxSCcxjvGc5C891LrS1/NlxglGxzrC4seZiVjcYVQfUa87wBL6rTDqzGjuoWNjnBzRqKmF6zRYGMvQUaQ==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "playwright-core": "1.58.1" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + }, + "optionalDependencies": { + "fsevents": "2.3.2" + } + }, + "node_modules/playwright-core": { + "version": "1.58.1", + "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.58.1.tgz", + "integrity": "sha512-bcWzOaTxcW+VOOGBCQgnaKToLJ65d6AqfLVKEWvexyS3AS6rbXl+xdpYRMGSRBClPvyj44njOWoxjNdL/H9UNg==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "playwright-core": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/typescript": { + "version": "5.9.3", + "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", + "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "tsc": "bin/tsc", + "tsserver": "bin/tsserver" + }, + "engines": { + "node": ">=14.17" + } + }, + "node_modules/undici-types": { + "version": "6.21.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz", + "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", + "dev": true, + "license": "MIT" + } + } +} diff --git a/tests/e2e/ui/package.json b/tests/e2e/ui/package.json new file mode 100644 index 00000000000..ede759d97cb --- /dev/null +++ b/tests/e2e/ui/package.json @@ -0,0 +1,16 @@ +{ + "name": "litellm-ui-e2e", + "version": "0.0.0", + "private": true, + "scripts": { + "e2e": "playwright test --config playwright.config.ts", + "e2e:ui": "playwright test --ui --config playwright.config.ts", + "e2e:migration": "playwright test tests/migration/migratedPages.spec.ts --config playwright.config.ts", + "e2e:migration:root": "playwright test --config migration.serverRootPath.config.ts" + }, + "devDependencies": { + "@playwright/test": "1.58.1", + "@types/node": "20.19.37", + "typescript": "5.9.3" + } +} diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/tests/e2e/ui/playwright.config.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/playwright.config.ts rename to tests/e2e/ui/playwright.config.ts diff --git a/ui/litellm-dashboard/e2e_tests/run_e2e.sh b/tests/e2e/ui/run_e2e.sh similarity index 98% rename from ui/litellm-dashboard/e2e_tests/run_e2e.sh rename to tests/e2e/ui/run_e2e.sh index ea95f18890c..858eb401c8e 100755 --- a/ui/litellm-dashboard/e2e_tests/run_e2e.sh +++ b/tests/e2e/ui/run_e2e.sh @@ -20,8 +20,8 @@ set -euo pipefail # ================================================================ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -DASHBOARD_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)" +DASHBOARD_DIR="$REPO_ROOT/ui/litellm-dashboard" IS_CI="${CI:-false}" CONTAINER_NAME="litellm-e2e-postgres-$$" MOCK_PID="" @@ -187,12 +187,12 @@ PGPASSWORD="$DB_PASS" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAM # --- Playwright --- echo "=== Installing Playwright dependencies ===" -cd "$DASHBOARD_DIR" +cd "$SCRIPT_DIR" npm install --silent 2>/dev/null || true npx playwright install chromium --with-deps 2>/dev/null || npx playwright install chromium echo "=== Running Playwright tests ===" -npx playwright test --config e2e_tests/playwright.config.ts "$@" +npx playwright test --config playwright.config.ts "$@" EXIT_CODE=$? exit $EXIT_CODE diff --git a/ui/litellm-dashboard/e2e_tests/serverRootPath.config.ts b/tests/e2e/ui/serverRootPath.config.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/serverRootPath.config.ts rename to tests/e2e/ui/serverRootPath.config.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts b/tests/e2e/ui/tests/auth/logout.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts rename to tests/e2e/ui/tests/auth/logout.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts b/tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts rename to tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts b/tests/e2e/ui/tests/auth/unauthenticatedRedirect.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts rename to tests/e2e/ui/tests/auth/unauthenticatedRedirect.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUser.spec.ts rename to tests/e2e/ui/tests/internal-user/internalUser.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts rename to tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts rename to tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-viewer/internalViewer.spec.ts b/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-viewer/internalViewer.spec.ts rename to tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts b/tests/e2e/ui/tests/login/internalUserIdentity.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts rename to tests/e2e/ui/tests/login/internalUserIdentity.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts b/tests/e2e/ui/tests/login/login.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts rename to tests/e2e/ui/tests/login/login.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/serverRootPathRedirect.spec.ts b/tests/e2e/ui/tests/login/serverRootPathRedirect.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/login/serverRootPathRedirect.spec.ts rename to tests/e2e/ui/tests/login/serverRootPathRedirect.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts b/tests/e2e/ui/tests/mcp/mcpServers.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts rename to tests/e2e/ui/tests/mcp/mcpServers.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/migration/README.md b/tests/e2e/ui/tests/migration/README.md similarity index 84% rename from ui/litellm-dashboard/e2e_tests/tests/migration/README.md rename to tests/e2e/ui/tests/migration/README.md index 4b3a391d421..d6b33598ec4 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/migration/README.md +++ b/tests/e2e/ui/tests/migration/README.md @@ -9,8 +9,9 @@ the default mount and a non-root `SERVER_ROOT_PATH` mount. ## Adding a page When a page's migration merges, add its route segment to -`e2e_tests/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES` -in `src/utils/migratedPages.ts`). Both suites pick it up automatically. +`tests/e2e/ui/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES` +in `ui/litellm-dashboard/src/utils/migratedPages.ts`). Both suites pick it up +automatically. ## Running diff --git a/ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts rename to tests/e2e/ui/tests/migration/migratedPages.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts rename to tests/e2e/ui/tests/modelHub/modelHub.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts rename to tests/e2e/ui/tests/modelsPage/addModel.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts rename to tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts b/tests/e2e/ui/tests/modelsPage/credentials.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts rename to tests/e2e/ui/tests/modelsPage/credentials.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/tests/e2e/ui/tests/navigation/sidebar.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts rename to tests/e2e/ui/tests/navigation/sidebar.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts rename to tests/e2e/ui/tests/proxy-admin/keys.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts b/tests/e2e/ui/tests/proxy-admin/license.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts rename to tests/e2e/ui/tests/proxy-admin/license.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts rename to tests/e2e/ui/tests/proxy-admin/teams.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/settings/adminSettings.spec.ts b/tests/e2e/ui/tests/settings/adminSettings.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/settings/adminSettings.spec.ts rename to tests/e2e/ui/tests/settings/adminSettings.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts b/tests/e2e/ui/tests/settings/routerSettings.spec.ts similarity index 98% rename from ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts rename to tests/e2e/ui/tests/settings/routerSettings.spec.ts index 3e140b9ab56..ffa5f2c2ae2 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts +++ b/tests/e2e/ui/tests/settings/routerSettings.spec.ts @@ -6,7 +6,7 @@ import { Role, users } from "../../fixtures/users"; // Type-only import of the OpenAPI-generated backend schema, erased at runtime by // esbuild. It types the round-trips below so mistakes surface in the editor; the live // test against the real proxy is what actually enforces the contract. -import type { components } from "../../../src/lib/http/schema"; +import type { components } from "../../../../../ui/litellm-dashboard/src/lib/http/schema"; // These tests mutate the proxy's shared router_settings, and the Loadbalancing save // echoes the whole settings object, so they must not run concurrently. diff --git a/ui/litellm-dashboard/e2e_tests/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/team-admin/teamAdmin.spec.ts rename to tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts rename to tests/e2e/ui/tests/users/searchUsers.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts b/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts rename to tests/e2e/ui/tests/users/viewInternalUsers.spec.ts diff --git a/tests/e2e/ui/tsconfig.json b/tests/e2e/ui/tsconfig.json new file mode 100644 index 00000000000..f9290fe7b49 --- /dev/null +++ b/tests/e2e/ui/tsconfig.json @@ -0,0 +1,16 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "commonjs", + "moduleResolution": "node", + "lib": ["ES2022", "DOM"], + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "resolveJsonModule": true, + "noEmit": true, + "types": ["node"] + }, + "include": ["**/*.ts"], + "exclude": ["node_modules"] +} diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 47ceb0c05fa..c693017e134 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -7756,6 +7756,7 @@ async def test_cli_completion_persists_assertion_under_db_user_id(): user_defined_values=None, prisma_client=MagicMock(), user_api_key_cache=MagicMock(), + cli_sso_session_cache=MagicMock(), proxy_logging_obj=MagicMock(), sso_assertion=assertion, ) diff --git a/ui/litellm-dashboard/knip.json b/ui/litellm-dashboard/knip.json index afed6b0f90e..48b39e8122d 100644 --- a/ui/litellm-dashboard/knip.json +++ b/ui/litellm-dashboard/knip.json @@ -1,7 +1,7 @@ { "$schema": "https://unpkg.com/knip@5/schema.json", "entry": ["scripts/**/*.{ts,mjs}", "src/components/ui/**/*.{ts,tsx}"], - "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.{ts,mjs}", "e2e_tests/**/*.ts"], + "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.{ts,mjs}"], "ignore": ["src/lib/http/schema.d.ts"], "ignoreDependencies": [ "openapi-typescript", @@ -10,14 +10,6 @@ "tailwindcss", "tw-animate-css" ], - "playwright": { - "config": [ - "e2e_tests/playwright.config.ts", - "e2e_tests/serverRootPath.config.ts", - "e2e_tests/migration.serverRootPath.config.ts" - ], - "entry": ["e2e_tests/**/*.spec.ts", "e2e_tests/**/*.setup.ts", "e2e_tests/globalSetup.ts"] - }, "vitest": { "config": ["vitest.config.ts"] }, diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 742c1e4a63f..14c8f0fdc18 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -47,7 +47,6 @@ }, "devDependencies": { "@eslint/js": "9.39.2", - "@playwright/test": "1.58.1", "@tailwindcss/forms": "0.5.11", "@tailwindcss/postcss": "4.3.2", "@testing-library/dom": "10.4.1", @@ -2723,8 +2722,9 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.58.1.tgz", "integrity": "sha512-6LdVIUERWxQMmUSSQi0I53GgCBYgM2RpGngCPY7hSeju+VrKjq3lvs7HpJoPbDiY5QM5EYRtRX5fvrinnMAz3w==", - "devOptional": true, "license": "Apache-2.0", + "optional": true, + "peer": true, "dependencies": { "playwright": "1.58.1" }, @@ -7361,7 +7361,6 @@ "version": "2.3.2", "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", - "dev": true, "hasInstallScript": true, "license": "MIT", "optional": true, @@ -10948,8 +10947,9 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.58.1.tgz", "integrity": "sha512-+2uTZHxSCcxjvGc5C891LrS1/NlxglGxzrC4seZiVjcYVQfUa87wBL6rTDqzGjuoWNjnBzRqKmF6zRYGMvQUaQ==", - "devOptional": true, "license": "Apache-2.0", + "optional": true, + "peer": true, "dependencies": { "playwright-core": "1.58.1" }, @@ -10967,8 +10967,9 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.58.1.tgz", "integrity": "sha512-bcWzOaTxcW+VOOGBCQgnaKToLJ65d6AqfLVKEWvexyS3AS6rbXl+xdpYRMGSRBClPvyj44njOWoxjNdL/H9UNg==", - "devOptional": true, "license": "Apache-2.0", + "optional": true, + "peer": true, "bin": { "playwright-core": "cli.js" }, diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 0f54c536297..5add2ad4e9e 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -14,10 +14,6 @@ "test:coverage": "vitest run --coverage", "format": "prettier --write .", "format:check": "prettier --check .", - "e2e": "playwright test --config e2e_tests/playwright.config.ts", - "e2e:ui": "playwright test --ui --config e2e_tests/playwright.config.ts", - "e2e:migration": "playwright test e2e_tests/tests/migration/migratedPages.spec.ts --config e2e_tests/playwright.config.ts", - "e2e:migration:root": "playwright test --config e2e_tests/migration.serverRootPath.config.ts", "knip": "knip", "knip:ci": "knip --exclude exports,nsExports,types,nsTypes,enumMembers,classMembers,duplicates", "knip:fix": "knip --fix", @@ -63,7 +59,6 @@ }, "devDependencies": { "@eslint/js": "9.39.2", - "@playwright/test": "1.58.1", "@tailwindcss/forms": "0.5.11", "@tailwindcss/postcss": "4.3.2", "@testing-library/dom": "10.4.1", diff --git a/ui/litellm-dashboard/tsconfig.json b/ui/litellm-dashboard/tsconfig.json index 8ca1752013a..5ca97e3e9db 100644 --- a/ui/litellm-dashboard/tsconfig.json +++ b/ui/litellm-dashboard/tsconfig.json @@ -23,5 +23,5 @@ "target": "ES2017" }, "include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts", ".next/dev/types/**/*.ts"], - "exclude": ["node_modules", "e2e_tests", "scripts"] + "exclude": ["node_modules", "scripts"] } diff --git a/ui/litellm-dashboard/vitest.config.ts b/ui/litellm-dashboard/vitest.config.ts index e1c58a0c2b5..da4734eeaf2 100644 --- a/ui/litellm-dashboard/vitest.config.ts +++ b/ui/litellm-dashboard/vitest.config.ts @@ -32,7 +32,6 @@ const config: ViteUserConfig = { "**/*.spec.*", "tests/**", - "e2e_tests/**", "node_modules/**", ".next/**", @@ -44,7 +43,7 @@ const config: ViteUserConfig = { "next.config.*", ], }, - exclude: ["e2e_tests/**", "node_modules/**"], + exclude: ["node_modules/**"], include: ["src/**/*.test.ts", "src/**/*.test.tsx", "tests/**/*.test.ts", "tests/**/*.test.tsx"], }, resolve: {