fix(guardrails/headroom): stop compressing the turn the model must act on

.github/workflows/test-unit-misc.yml

@@ -27,6 +27,7 @@ jobs:
         tests/test_litellm/a2a_protocol
         tests/test_litellm/anthropic_interface
         tests/test_litellm/completion_extras
+        tests/test_litellm/compression
         tests/test_litellm/containers
         tests/test_litellm/experimental_mcp_client
         tests/test_litellm/models

litellm/compression/compress.py

@@ -3,6 +3,7 @@
 scoring, message stubbing, and retrieval tool injection.
 """
 
+from collections.abc import Mapping, Sequence
 from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
 
 from litellm.caching.dual_cache import DualCache
@@ -204,33 +205,21 @@ def _extract_anthropic_tool_exchange_spans(
     return spans, None
 
 
-def _get_protected_indices(messages: List[dict]) -> List[int]:
+def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
     """
     Return indices of messages that must never be compressed:
     - All system messages
     - The last user message
     - The last assistant message
-    """
-    protected: List[int] = []
-
-    last_user_idx = None
-    last_assistant_idx = None
-
-    for i, msg in enumerate(messages):
-        role = msg.get("role", "")
-        if role == "system":
-            protected.append(i)
-        elif role == "user":
-            last_user_idx = i
-        elif role == "assistant":
-            last_assistant_idx = i
 
-    if last_user_idx is not None:
-        protected.append(last_user_idx)
-    if last_assistant_idx is not None:
-        protected.append(last_assistant_idx)
-
-    return protected
+    The last user message is what the model is being asked to act on right now,
+    so compressing it replaces the live instruction with a marker. Compression
+    guardrails share this policy; see the Headroom guardrail.
+    """
+    system_indices = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system")
+    last_user = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:]
+    last_assistant = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:]
+    return system_indices + last_user + last_assistant
 
 
 def _combine_scores(
@@ -432,7 +421,7 @@ def compress(
         combined_scores = bm25_scores
 
     # Protected messages are never compressed
-    protected_indices = _get_protected_indices(normalized_messages)
+    protected_indices = get_protected_indices(normalized_messages)
     kept_indices: Set[int] = set(protected_indices)
 
     tool_exchange_spans: List[Set[int]] = []

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -6,7 +6,7 @@
 import re
 import xml.etree.ElementTree as ET
 from enum import Enum
-from collections.abc import Mapping
+from collections.abc import Iterator, Mapping, Sequence
 from typing import Any, Dict, List, Optional, Set, Tuple, TypedDict, Union, cast, overload
 
 from jinja2.sandbox import ImmutableSandboxedEnvironment
@@ -2210,6 +2210,49 @@ def _is_orphaned_tool_result(
     return False
 
 
+def _declared_tool_call_ids(message: Mapping[str, Any]) -> frozenset[str]:
+    tool_calls = message.get("tool_calls")
+    if not isinstance(tool_calls, list):
+        return frozenset()
+    return frozenset(
+        str(tool_call["id"]) for tool_call in tool_calls if isinstance(tool_call, Mapping) and tool_call.get("id")
+    )
+
+
+def group_tool_exchanges(messages: Sequence[Mapping[str, Any]]) -> tuple[tuple[int, ...], ...]:
+    """Group message indices into tool exchanges: an assistant row that made
+    tool calls, together with the tool rows answering the ids it declared.
+
+    Membership is by ``tool_call_id`` ownership rather than adjacency, so a tool
+    row belonging to some other call opens its own group instead of being swept
+    into the exchange it happens to sit next to. Every other row is its own
+    group. Groups stay contiguous and in order, so a caller can convert or
+    protect them without reordering the conversation.
+
+    Callers need this because an assistant row and the tool rows answering it
+    are only well-formed together: ``sanitize_messages_for_tool_calling`` reads
+    an assistant row whose results are missing as an orphaned tool call, and
+    a tool row whose call is missing as an orphaned result.
+    """
+    return tuple(_iter_tool_exchange_groups(messages))
+
+
+def _iter_tool_exchange_groups(messages: Sequence[Mapping[str, Any]]) -> Iterator[tuple[int, ...]]:
+    index = 0
+    while index < len(messages):
+        declared = _declared_tool_call_ids(messages[index])
+        end = index + 1
+        while (
+            declared
+            and end < len(messages)
+            and messages[end].get("role") in ("tool", "function")
+            and str(messages[end].get("tool_call_id")) in declared
+        ):
+            end += 1
+        yield tuple(range(index, end))
+        index = end
+
+
 def sanitize_messages_for_tool_calling(
     messages: List[AllMessageValues],
 ) -> List[AllMessageValues]:

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -361,14 +361,34 @@ async def process_input_messages(
 
     @staticmethod
     def _write_back_structured_messages(data: dict, structured_messages: list) -> None:
-        """Convert compressed structured_messages back to Anthropic format and write to data."""
+        """Convert compressed structured_messages back to Anthropic format and write to data.
+
+        ``anthropic_messages_pt`` merges every run of consecutive user/tool rows
+        into a single message, so a turn carrying only tool results and the user
+        turn that follows it come back fused, and the request the model sees no
+        longer has the boundaries the client sent. Converting a row at a time
+        would keep them apart but breaks tool pairing: an assistant row whose
+        tool results sit outside its own call reads as an orphaned tool call,
+        and under ``modify_params`` the sanitizer answers it with a synthetic
+        "tool execution skipped" result and drops the real one. Converting each
+        assistant row together with the tool rows that answer it, and every
+        other row on its own, satisfies both.
+        """
         from litellm.litellm_core_utils.prompt_templates.factory import (
             anthropic_messages_pt,
+            group_tool_exchanges,
         )
 
         model = str(data.get("model") or "")
         non_system = [m for m in structured_messages if m.get("role") != "system"]
-        converted = anthropic_messages_pt(messages=non_system, model=model, llm_provider="anthropic")
+        groups = tuple([non_system[index] for index in group] for group in group_tool_exchanges(non_system)) or (
+            non_system,
+        )
+        converted = [
+            message
+            for group in groups
+            for message in anthropic_messages_pt(messages=group, model=model, llm_provider="anthropic")
+        ]
         for msg in converted:
             content = msg.get("content")
             if isinstance(content, list):

litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py

@@ -48,6 +48,7 @@
 )
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.content_text import (
+    assistant_text_from_response,
     content_to_text,
     is_all_text_parts,
     merge_rewritten_text_parts,
@@ -391,47 +392,6 @@ def _is_anthropic_messages_response(response: object) -> bool:
     return isinstance(get_attribute_or_key(response, "content", None), list)
 
 
-def _assistant_text_from_response(response: object) -> str | None:
-    """The assistant's natural-language text from a model response, across chat,
-    Anthropic, and Responses shapes. Preserved when the turn is rebuilt for the
-    retrieval follow-up so the model's reasoning is not lost."""
-    choices = get_attribute_or_key(response, "choices", None)
-    if isinstance(choices, list) and choices:
-        message = get_attribute_or_key(choices[0], "message", None)
-        if message is not None:
-            text = content_to_text(get_attribute_or_key(message, "content", None))
-            if text:
-                return text
-    content = get_attribute_or_key(response, "content", None)
-    if isinstance(content, list):
-        parts = [
-            text
-            for block in content
-            if get_attribute_or_key(block, "type", None) == "text"
-            for text in (get_attribute_or_key(block, "text", None),)
-            if isinstance(text, str) and text
-        ]
-        if parts:
-            return "".join(parts)
-    output = get_attribute_or_key(response, "output", None)
-    if isinstance(output, list):
-        parts = []
-        for item in output:
-            if get_attribute_or_key(item, "type", None) != "message":
-                continue
-            item_content = get_attribute_or_key(item, "content", None)
-            if not isinstance(item_content, list):
-                continue
-            for chunk in item_content:
-                if get_attribute_or_key(chunk, "type", None) == "output_text":
-                    text = get_attribute_or_key(chunk, "text", None)
-                    if isinstance(text, str) and text:
-                        parts.append(text)
-        if parts:
-            return "".join(parts)
-    return None
-
-
 def _build_assistant_message_from_response(
     response: object,
     retrieved: list[tuple[dict[str, object], str]],
@@ -446,7 +406,7 @@ def _build_assistant_message_from_response(
     """
     return {
         "role": "assistant",
-        "content": _assistant_text_from_response(response),
+        "content": assistant_text_from_response(response),
         "tool_calls": [
             {
                 "id": tool_call.get("id"),
@@ -470,7 +430,7 @@ def _build_anthropic_followup_messages(
     assistant text is preserved; non-retrieve tool calls are re-planned by the
     follow-up (see _build_assistant_message_from_response)."""
     assistant_content: list[dict[str, object]] = []
-    text = _assistant_text_from_response(response)
+    text = assistant_text_from_response(response)
     if text:
         assistant_content.append({"type": "text", "text": text})
     assistant_content.extend(
@@ -501,7 +461,7 @@ def _build_responses_followup_items(
     with a function_call_output keyed by the same call_id. The assistant text is
     preserved; non-retrieve tool calls are re-planned by the follow-up."""
     items: list[dict[str, object]] = []
-    text = _assistant_text_from_response(response)
+    text = assistant_text_from_response(response)
     if text:
         items.append({"role": "assistant", "content": text})
     for tool_call, content in retrieved:

litellm/proxy/guardrails/guardrail_hooks/content_text.py

@@ -14,6 +14,8 @@
 
 from collections.abc import Sequence
 
+from litellm.litellm_core_utils.prompt_templates.factory import get_attribute_or_key
+
 
 def content_to_text(content: object) -> str:
     """Collapse a message ``content`` (str or list-of-parts) to plain text.
@@ -53,3 +55,41 @@ def merge_rewritten_text_parts(parts: Sequence[object], new_text: str) -> list[o
     breakpoints = tuple(part["cache_control"] for part in dict_parts if part.get("cache_control") is not None)
     base = {**dict_parts[0], "text": new_text} if dict_parts else {"type": "text", "text": new_text}
     return [{**base, "cache_control": breakpoints[-1]} if breakpoints else base]
+
+
+def assistant_text_from_response(response: object) -> str | None:
+    """The assistant's natural-language text from a model response, across chat,
+    Anthropic, and Responses shapes. Preserved when the turn is rebuilt for the
+    retrieval follow-up so the model's reasoning is not lost."""
+    choices = get_attribute_or_key(response, "choices", None)
+    if isinstance(choices, list) and choices:
+        message = get_attribute_or_key(choices[0], "message", None)
+        if message is not None:
+            text = content_to_text(get_attribute_or_key(message, "content", None))
+            if text:
+                return text
+    content = get_attribute_or_key(response, "content", None)
+    if isinstance(content, list):
+        parts = [
+            text
+            for block in content
+            if get_attribute_or_key(block, "type", None) == "text"
+            for text in (get_attribute_or_key(block, "text", None),)
+            if isinstance(text, str) and text
+        ]
+        if parts:
+            return "".join(parts)
+    output = get_attribute_or_key(response, "output", None)
+    if isinstance(output, list):
+        output_parts = [
+            text
+            for item in output
+            if get_attribute_or_key(item, "type", None) == "message"
+            for chunk in (get_attribute_or_key(item, "content", None) or ())
+            if get_attribute_or_key(chunk, "type", None) == "output_text"
+            for text in (get_attribute_or_key(chunk, "text", None),)
+            if isinstance(text, str) and text
+        ]
+        if output_parts:
+            return "".join(output_parts)
+    return None