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"