diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 67a0b2a2d15..ea757edb215 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -11273,6 +11273,12 @@ ], "description": "Threshold configuration for Lakera guardrail categories" }, + "ccr_retrieval": { + "default": true, + "description": "Inject the Headroom retrieval tool for hashes declared by the compression service.", + "title": "Ccr Retrieval", + "type": "boolean" + }, "checks": { "anyOf": [ { diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py index d569802ce89..b7e4d9275fa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py @@ -36,6 +36,7 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> default_on=litellm_params.default_on or False, unreachable_fallback=litellm_params.unreachable_fallback, timeout=litellm_params.timeout, + ccr_retrieval=litellm_params.ccr_retrieval, ) litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] _callback diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 685b90f1754..fa113aa4d33 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -6,6 +6,7 @@ import re import time import uuid from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypeGuard import httpx @@ -60,7 +61,7 @@ _STREAM_CONVERTIBLE_CALL_TYPES: Final = frozenset( # stalled service holds the caller's request and a pooled connection for 600s or more. _COMPRESS_TIMEOUT_SECONDS: Final = 60.0 HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" -_HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})") +_HASH_PATTERN: Final = re.compile(r"[a-f0-9]{12,24}") _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 # Narrows the base class's bare-dict ``request_data`` at the boundary so its # untranslated messages can be read with concrete types (values pass through by @@ -239,15 +240,25 @@ def _protected_indices( tool exchanges the way ``compress()`` expands it, so a protected assistant tool call cannot end up answered by a marker standing in for the result the model just asked for. + + Every assistant row is then withheld without expanding its tool exchange: + the service protects assistant text blocks but has no gate for assistant + strings, and the Anthropic adapter hands assistant blocks over as strings, + so the model's own earlier tables came back rewritten and it imitated the + shape. The tool results those turns asked for stay compressible. """ protected: Final = frozenset(get_protected_indices(messages)) | _retrieval_result_indices( messages, extra_retrieve_call_ids ) - return protected | frozenset( - index - for group in group_tool_exchanges(messages) - if any(member in protected for member in group) - for index in group + return ( + protected + | frozenset( + index + for group in group_tool_exchanges(messages) + if any(member in protected for member in group) + for index in group + ) + | frozenset(index for index, message in enumerate(messages) if message.get("role") == "assistant") ) @@ -290,19 +301,23 @@ def _build_compress_failure_detail(status_code: int, body: str) -> dict[str, obj return {"status_code": status_code, "body": body} -def extract_hashes_from_messages(messages: list[dict[str, object]]) -> list[str]: - hashes: Final[list[str]] = [] - for msg in messages: - content = msg.get("content") - if isinstance(content, str): - hashes.extend(_HASH_PATTERN.findall(content)) - elif isinstance(content, list): - for block in content: - if isinstance(block, dict): - text = block.get("text") - if isinstance(text, str): - hashes.extend(_HASH_PATTERN.findall(text)) - return hashes +def _read_ccr_hashes(body: Mapping[str, object]) -> frozenset[str]: + ccr_hashes: Final = body.get("ccr_hashes") + if not isinstance(ccr_hashes, list): + return frozenset() + return frozenset( + hash_value.lower() + for hash_value in ccr_hashes + if isinstance(hash_value, str) and _HASH_PATTERN.fullmatch(hash_value.lower()) + ) + + +@dataclass(frozen=True, slots=True) +class _CompressResult: + messages: list[dict[str, object]] + succeeded: bool + stats: dict[str, object] + ccr_hashes: frozenset[str] = frozenset() def _build_headroom_retrieve_tool() -> dict[str, object]: @@ -319,7 +334,7 @@ def _build_headroom_retrieve_tool() -> dict[str, object]: "properties": { "hash": { "type": "string", - "description": "The 24-character hex hash from the compression marker.", + "description": "The hex hash from the compression marker.", }, "query": { "type": "string", @@ -479,6 +494,7 @@ class HeadroomGuardrail(CustomGuardrail): default_on: bool = False, unreachable_fallback: str | None = None, timeout: float | None = None, + ccr_retrieval: bool = True, ): self.headroom_api_base = (api_base or get_secret_str("HEADROOM_API_BASE") or "").rstrip("/") if not self.headroom_api_base: @@ -492,6 +508,7 @@ class HeadroomGuardrail(CustomGuardrail): "fail_open" if unreachable_fallback == "fail_open" else "fail_closed" ) self.timeout: httpx.Timeout = self._resolve_timeout(timeout) + self.ccr_retrieval = ccr_retrieval self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, ) @@ -569,7 +586,7 @@ class HeadroomGuardrail(CustomGuardrail): self, messages: list[dict[str, object]], model: str | None, - ) -> tuple[list[dict[str, object]], bool, dict[str, object]]: + ) -> _CompressResult: payload: Final[dict[str, object]] = {"messages": messages} if model: payload["model"] = model @@ -582,7 +599,7 @@ class HeadroomGuardrail(CustomGuardrail): timeout=self.timeout, ) except httpx.HTTPStatusError as e: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned an error", @@ -592,7 +609,7 @@ class HeadroomGuardrail(CustomGuardrail): {}, ) except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service unreachable", @@ -604,7 +621,7 @@ class HeadroomGuardrail(CustomGuardrail): response: Final[HttpxResponse] = raw_response if response.status_code != 200: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned an error", @@ -617,7 +634,7 @@ class HeadroomGuardrail(CustomGuardrail): try: body: Final[object] = response.json() except ValueError: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned non-JSON response", @@ -627,7 +644,7 @@ class HeadroomGuardrail(CustomGuardrail): {}, ) if not _is_str_object_dict(body): - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned unexpected response shape", @@ -639,7 +656,7 @@ class HeadroomGuardrail(CustomGuardrail): compressed_messages: Final = body.get("messages") if not _is_object_list(compressed_messages): - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service response missing 'messages'", @@ -651,7 +668,7 @@ class HeadroomGuardrail(CustomGuardrail): filtered: Final = [item for item in compressed_messages if _is_str_object_dict(item)] if not filtered: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned empty message list", @@ -664,7 +681,7 @@ class HeadroomGuardrail(CustomGuardrail): if len(filtered) != len(messages): # Rows are matched positionally when the never-compressed messages # are put back, so a reshaped conversation cannot be applied at all. - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service changed the message count", @@ -705,7 +722,7 @@ class HeadroomGuardrail(CustomGuardrail): # tokens_saved, which the live compression service omits; derive it # so savings are counted, but let a service-sent value win. stats["tokens_saved"] = tokens_before - tokens_after - return filtered, True, stats + return _CompressResult(filtered, True, stats, _read_ccr_hashes(body)) async def _call_retrieve(self, hash_value: str, query: str | None = None) -> str: params: Final[dict[str, str]] = {} @@ -793,7 +810,7 @@ class HeadroomGuardrail(CustomGuardrail): model: Final = self.headroom_model or request_data.get("model") start_time: Final = time.time() - returned, compression_succeeded, stats = await self._call_compress( + result: Final = await self._call_compress( messages=_flatten_messages_for_compression(compressible), model=model if isinstance(model, str) else None, ) @@ -803,7 +820,7 @@ class HeadroomGuardrail(CustomGuardrail): add_guardrail_to_applied_guardrails_header, ) - if not compression_succeeded: + if not result.succeeded: self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response={"error": "headroom compression unavailable; request forwarded uncompressed"}, request_data=request_data, @@ -822,12 +839,12 @@ class HeadroomGuardrail(CustomGuardrail): compressed: Final = _restore_protected_messages( messages=messages, - compressed=_restore_content_shapes(originals=compressible, returned=returned), + compressed=_restore_content_shapes(originals=compressible, returned=result.messages), protected_indices=protected_indices, ) self.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response=stats, + guardrail_json_response=result.stats, request_data=request_data, guardrail_status="success", guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER, @@ -837,7 +854,7 @@ class HeadroomGuardrail(CustomGuardrail): ) add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) - hashes: Final = extract_hashes_from_messages(compressed) + hashes: Final = result.ccr_hashes if self.ccr_retrieval else frozenset() if not hashes: return {**inputs, "structured_messages": compressed} # pyright: ignore[reportReturnType] @@ -918,14 +935,15 @@ class HeadroomGuardrail(CustomGuardrail): retrieved: Final[list[tuple[dict[str, object], str]]] = [] for tc in tool_calls: arguments = tc.get("arguments", {}) - hash_value = arguments.get("hash", "") if isinstance(arguments, dict) else "" + raw_hash = arguments.get("hash", "") if isinstance(arguments, dict) else "" + hash_value = str(raw_hash).lower() query = arguments.get("query") if isinstance(arguments, dict) else None # A hash is only honored if it was issued by *this request's own* # Headroom /v1/compress call, scoped by litellm_call_id. Scoping by # message text alone is forgeable -- an attacker can plant a # hash-shaped string in their own prompt, and a hash issued for one # request would validate for any other request that echoes it back. - if str(hash_value) not in valid_hashes: + if hash_value not in valid_hashes: verbose_proxy_logger.warning( "Headroom CCR: rejecting hash=%s not produced by current request compression", hash_value, @@ -933,7 +951,7 @@ class HeadroomGuardrail(CustomGuardrail): content = f"[Headroom: hash={hash_value} was not produced by the current request]" else: content = await self._call_retrieve( - hash_value=str(hash_value), + hash_value=hash_value, query=str(query) if query else None, ) verbose_proxy_logger.debug("Headroom CCR: retrieved hash=%s (%d chars)", hash_value, len(content)) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py index d517e508596..0df0f0c80ef 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py @@ -26,6 +26,10 @@ class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]): "forwards the request uncompressed instead of blocking it." ), ) + ccr_retrieval: bool = Field( + default=True, + description="Inject the Headroom retrieval tool for hashes declared by the compression service.", + ) @staticmethod def ui_friendly_name() -> str: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index a49d7723bcc..d8eeb8d2b8a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -35,7 +35,6 @@ import litellm from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( HeadroomGuardrail, - extract_hashes_from_messages, has_headroom_retrieve_tool, HEADROOM_RETRIEVE_TOOL_NAME, ) @@ -93,7 +92,11 @@ def _make_guardrail(**kwargs) -> HeadroomGuardrail: return HeadroomGuardrail(**defaults) -def _make_compress_response(messages: list, status: int = 200) -> MagicMock: +def _make_compress_response( + messages: list, + status: int = 200, + ccr_hashes: list[str] | None = None, +) -> MagicMock: mock = MagicMock() mock.status_code = status mock.json.return_value = { @@ -102,6 +105,7 @@ def _make_compress_response(messages: list, status: int = 200) -> MagicMock: "tokens_after": 100, "compression_ratio": 0.1, "transforms_applied": ["router:smart_crusher:0.35"], + **({} if ccr_hashes is None else {"ccr_hashes": ccr_hashes}), } mock.text = "" return mock @@ -335,7 +339,7 @@ async def test_apply_guardrail_injects_retrieve_tool_when_hashes_present( texts=["A" * 5000], structured_messages=ORIGINAL_MESSAGES, ) - mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH) + mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]) with patch.object( guardrail.async_handler, @@ -390,7 +394,7 @@ async def test_apply_guardrail_preserves_existing_tools_when_injecting( structured_messages=ORIGINAL_MESSAGES, tools=[existing_tool], ) - mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH) + mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]) with patch.object( guardrail.async_handler, @@ -489,17 +493,18 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( original_content = "This is the full compressed content." mock_retrieve = _make_retrieve_response(original_content) + # Registered hashes are lowercase; a model may echo the marker's hex in uppercase. tool_calls = [ { "id": "call_abc123", "type": "function", "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": {"hash": "b573993006976af767214fac"}, + "arguments": {"hash": "B573993006976AF767214FAC"}, } ] response = _make_openai_response_with_tool_call( tool_name=HEADROOM_RETRIEVE_TOOL_NAME, - arguments={"hash": "b573993006976af767214fac"}, + arguments={"hash": "B573993006976AF767214FAC"}, tool_id="call_abc123", ) messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] @@ -843,33 +848,110 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag assert tool_result_block["content"] == original_content -def test_extract_hashes_from_messages_finds_hashes(): +HASH_SHAPED_HISTORY = [ + {"role": "system", "content": "You are Claude Code."}, + {"role": "user", "content": [{"type": "text", "text": "Run git log."}]}, + { + "role": "assistant", + "content": [{"type": "text", "text": "Done."}], + "tool_calls": [{"id": "tu_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "tu_1", "content": "hash=3f2a9c1d7e5b4a6f8c0d1e2f9a8b7c6d5e4f3a2b"}, + {"role": "user", "content": "Please fetch hash=deadbeef000000000000dead for me."}, +] + + +async def _apply(guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None) -> dict: + request_data = {"model": "claude-sonnet-5"} + + def _echo(**kwargs): + return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"])), ccr_hashes=ccr_hashes) + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): + return await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))), + request_data=request_data, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_hash_shaped_text_in_history_never_injects_retrieve_tool(guardrail: HeadroomGuardrail): + """Regression for LIT-7086: a git SHA in a tool result and a hash= string the + caller typed both look like markers, but the service stored nothing, so the + tool must not appear and no hash may be registered as issued. Covers a + service that omits ccr_hashes, returns it empty, or returns a non-list.""" + for ccr_hashes in (None, [], "b573993006976af767214fac"): + result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes) + + assert not has_headroom_retrieve_tool(result.get("tools") or []) + assert not guardrail._issued_hashes_by_call_id + + +@pytest.mark.asyncio +async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardrail: HeadroomGuardrail): + """Only the hashes the service reports in ccr_hashes are honored, in the + service's own 12 to 24 hex grammar; anything else is dropped because each + entry is interpolated into the /v1/retrieve URL.""" + result = await _apply( + guardrail, + HASH_SHAPED_HISTORY, + ccr_hashes=["98CA69107318", "b573993006976af767214fac", "../../etc/passwd", "tooshort", 42], + ) + + assert has_headroom_retrieve_tool(result.get("tools") or []) + (issued, _expiry), = guardrail._issued_hashes_by_call_id.values() + assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"}) + + +@pytest.mark.asyncio +async def test_ccr_retrieval_disabled_ignores_service_declared_hashes(monkeypatch: pytest.MonkeyPatch): + """`ccr_retrieval: false` in config.yaml compresses without any retrieval + round trip, so it has to reach the instance through the initializer.""" + from litellm.proxy.guardrails.guardrail_hooks.headroom import initialize_guardrail + from litellm.types.guardrails import LitellmParams + + monkeypatch.setattr(litellm.logging_callback_manager, "add_litellm_callback", lambda callback: None) + params = LitellmParams(guardrail="headroom", mode="pre_call", api_base=FAKE_API_BASE, ccr_retrieval=False) + guardrail = initialize_guardrail(params, {"guardrail_name": "headroom", "litellm_params": params}) + + result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=["b573993006976af767214fac"]) + + assert result["structured_messages"][-1] == HASH_SHAPED_HISTORY[-1] + assert not has_headroom_retrieve_tool(result.get("tools") or []) + assert not guardrail._issued_hashes_by_call_id + + +@pytest.mark.asyncio +async def test_anthropic_assistant_history_never_reaches_compression_service(guardrail: HeadroomGuardrail): + """The public Anthropic handler translates assistant content blocks to a + string before Headroom sees them, so model-authored rows must be excluded + from the compression payload rather than protected by their content shape.""" + from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler + + table = "| Guardrail | Model |\n|---|---|\n" + "\n".join(f"| gr-{i} | model-{i} |" for i in range(40)) messages = [ - {"role": "user", "content": "Retrieve more: hash=b573993006976af767214fac"}, - {"role": "assistant", "content": "Also: hash=aabbccdd001122334455aabb"}, + {"role": "user", "content": [{"type": "text", "text": "List the guardrails."}]}, + {"role": "assistant", "content": [{"type": "text", "text": table}]}, + {"role": "user", "content": [{"type": "text", "text": "Earlier follow-up. " + "B" * 5000}]}, + {"role": "assistant", "content": [{"type": "text", "text": "Noted."}]}, + {"role": "user", "content": "Re-print the table."}, ] - hashes = extract_hashes_from_messages(messages) - assert "b573993006976af767214fac" in hashes - assert "aabbccdd001122334455aabb" in hashes + sent: dict = {} + def _echo(**kwargs): + sent["messages"] = kwargs["json"]["messages"] + return _make_compress_response(json.loads(json.dumps(sent["messages"]))) -def test_extract_hashes_from_messages_ignores_short_hashes(): - messages = [{"role": "user", "content": "hash=tooshort"}] - hashes = extract_hashes_from_messages(messages) - assert not hashes + data = {"model": "claude-sonnet-5", "messages": json.loads(json.dumps(messages))} + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): + result = await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + assert [row["role"] for row in sent["messages"]] == ["user", "user"] + assert sent["messages"][1]["content"] == "Earlier follow-up. " + "B" * 5000 + assert table not in json.dumps(sent["messages"]) + assert result["messages"][1]["content"] == [{"type": "text", "text": table}] -def test_extract_hashes_from_list_content_blocks(): - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "hash=b573993006976af767214fac found here"}, - ], - } - ] - hashes = extract_hashes_from_messages(messages) - assert "b573993006976af767214fac" in hashes def test_has_headroom_retrieve_tool_recognizes_anthropic_native_shape(): @@ -1001,7 +1083,7 @@ async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstre guardrail.async_handler, "post", new_callable=AsyncMock, - return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH), + return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]), ): result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) @@ -1793,7 +1875,7 @@ async def test_apply_guardrail_restores_rewritten_all_text_row( ) compressed = _echo_wire_view() compressed[0]["content"] = "compressed history. Retrieve more: hash=b573993006976af767214fac" - mock_response = _make_compress_response(compressed) + mock_response = _make_compress_response(compressed, ccr_hashes=["b573993006976af767214fac"]) with patch.object( guardrail.async_handler, @@ -1819,7 +1901,7 @@ async def test_apply_guardrail_restores_rewritten_all_text_row( assert history_content[0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} # Mixed row passes through byte-identical. assert messages[2]["content"] == PARTS_MESSAGES[2]["content"] - # Hashes inside restored parts still drive retrieve-tool injection. + # The service-declared hash still drives retrieve-tool injection on a restored row. assert has_headroom_retrieve_tool(result.get("tools") or []) @@ -2159,7 +2241,7 @@ async def test_pre_call_deployment_hook_converts_stream_after_deployment_level_c guardrail.async_handler, "post", new_callable=AsyncMock, - return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH), + return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]), ): result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion) @@ -2345,7 +2427,11 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( AGENTIC_MESSAGES = [ {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, {"role": "user", "content": "H" * 5000}, - {"role": "assistant", "content": "Older answer. " + "O" * 5000}, + { + "role": "assistant", + "content": "Older answer. " + "O" * 5000, + "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], + }, {"role": "tool", "tool_call_id": "old_1", "content": "older tool output " + "T" * 5000}, { "role": "assistant", @@ -2420,23 +2506,21 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): """Negative control: protection must not turn compression into a no-op.""" compressed_history = [ {"role": "user", "content": "hist. hash=b573993006976af767214fac"}, - {"role": "assistant", "content": "older. hash=a73993006976af767214fac1"}, {"role": "tool", "tool_call_id": "old_1", "content": "older tool. hash=c73993006976af767214fac2"}, ] wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES, returned=compressed_history) - # Exactly the three history rows go to the service, in order. - assert [row["role"] for row in wire] == ["user", "assistant", "tool"] + # Older user and tool rows go to the service, in order. Every assistant row + # stays out, but the tool results those turns asked for remain compressible. + assert [row["role"] for row in wire] == ["user", "tool"] assert wire[0]["content"] == "H" * 5000 - assert wire[2]["tool_call_id"] == "old_1" + assert wire[1]["tool_call_id"] == "old_1" messages = result["structured_messages"] assert len(messages) == len(AGENTIC_MESSAGES) assert messages[1] == compressed_history[0] - assert messages[2] == compressed_history[1] - assert messages[3] == compressed_history[2] - # Hashes in the compressed history still drive retrieve-tool injection. - assert has_headroom_retrieve_tool(result.get("tools") or []) + assert messages[2] == AGENTIC_MESSAGES[2] + assert messages[3] == compressed_history[1] # --------------------------------------------------------------------------- diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 680da00e602..8a850598259 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30499,6 +30499,12 @@ export interface components { categories?: components["schemas"]["ContentFilterCategoryConfig"][] | null; /** @description Threshold configuration for Lakera guardrail categories */ category_thresholds?: components["schemas"]["LakeraCategoryThresholds"] | null; + /** + * Ccr Retrieval + * @description Inject the Headroom retrieval tool for hashes declared by the compression service. + * @default true + */ + ccr_retrieval: boolean; /** @description Inline safeguards for the resource-less InvokeGuardrailChecks API (contentFilter / promptAttack / sensitiveInformation). When set, the guardrail calls InvokeGuardrailChecks instead of ApplyGuardrail and no guardrailIdentifier is required. Mutually exclusive with guardrailIdentifier. */ checks?: components["schemas"]["BedrockChecksConfigModel"] | null; /**