fix(headroom): resolve CCR retrieval on streaming /chat/completions

Streaming chat completions returned a CustomStreamWrapper, so the agentic loop dispatch in main.py (guarded on ModelResponse) never ran and the headroom_retrieve tool call was streamed straight to a client that never declared the tool. The guardrail now converts a CCR stream request to a non-streaming call in its deployment hook and the loop fake-streams the resolved answer back.

litellm/litellm_core_utils/chat_completion_agentic_loop.py

@@ -1,12 +1,14 @@
 # this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api
 
 import json
+from collections.abc import Mapping
 from typing import cast
 
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.types.integrations.custom_logger import (
     CHAT_COMPLETION_AGENTIC_SURFACE,
+    HEADROOM_CONVERTED_STREAM_KEY,
     NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
     AgenticLoopPlan,
     AgenticLoopRequestPatch,
@@ -46,6 +48,12 @@ def _post_hook_overridden(callback: CustomLogger) -> bool:
     return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
 
 
+def _converted_stream_requested(kwargs: Mapping[str, object]) -> bool:
+    return bool(
+        kwargs.get("_code_interpreter_interception_converted_stream") or kwargs.get(HEADROOM_CONVERTED_STREAM_KEY)
+    )
+
+
 def _coerce_int(value: object, default: int) -> int:
     return int(value) if isinstance(value, (int, str)) else default
 
@@ -80,16 +88,25 @@ def _check_agentic_loop_safety(
     return fingerprint
 
 
-def _wrap_response_as_fake_stream(response: object) -> object:
-    if getattr(response, "object", None) == "chat.completion.chunk":
+def _wrap_response_as_fake_stream(
+    response: object,
+    *,
+    model: str,
+    custom_llm_provider: str,
+    logging_obj: object,
+) -> object:
+    if isinstance(response, CustomStreamWrapper):
         return response
-    if not hasattr(response, "choices"):
+    if not isinstance(response, ModelResponse):
         return response
-    from litellm.llms.base_llm.base_model_iterator import (
-        convert_model_response_to_streaming,
-    )
+    from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
 
-    return convert_model_response_to_streaming(cast(ModelResponse, response))
+    return CustomStreamWrapper(
+        completion_stream=MockResponseIterator(model_response=response),
+        model=model,
+        custom_llm_provider=custom_llm_provider,
+        logging_obj=logging_obj,
+    )
 
 
 def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None:
@@ -170,8 +187,13 @@ async def _execute_chat_completion_agentic_plan(
                     model,
                     str(e),
                 )
-        if kwargs.get("_code_interpreter_interception_converted_stream") and not depth:
-            return _wrap_response_as_fake_stream(response_followup)
+        if _converted_stream_requested(kwargs) and not depth:
+            return _wrap_response_as_fake_stream(
+                response_followup,
+                model=model,
+                custom_llm_provider=custom_llm_provider,
+                logging_obj=logging_obj,
+            )
         return response_followup
     finally:
         try:
@@ -295,9 +317,14 @@ async def maybe_run_chat_completion_agentic_loop(
                 str(e),
             )
 
-    if kwargs.get("_code_interpreter_interception_converted_stream") and not depth and hasattr(response, "choices"):
+    if _converted_stream_requested(kwargs) and not depth:
         return cast(
             "ModelResponse | CustomStreamWrapper",
-            _wrap_response_as_fake_stream(response),
+            _wrap_response_as_fake_stream(
+                response,
+                model=model,
+                custom_llm_provider=custom_llm_provider,
+                logging_obj=logging_obj,
+            ),
         )
     return None

litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py

@@ -4,6 +4,7 @@ import json
 import re
 import time
 import uuid
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, ClassVar, List, Literal, Optional
 
 import httpx
@@ -35,8 +36,12 @@ from litellm.proxy.guardrails.guardrail_hooks.content_text import (
 )
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.guardrails import GuardrailEventHooks, Mode
-from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch
-from litellm.types.utils import GenericGuardrailAPIInputs
+from litellm.types.integrations.custom_logger import (
+    HEADROOM_CONVERTED_STREAM_KEY,
+    AgenticLoopPlan,
+    AgenticLoopRequestPatch,
+)
+from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -606,6 +611,19 @@ class HeadroomGuardrail(CustomGuardrail):
 
         return {**inputs, "structured_messages": compressed, "tools": merged_tools}  # pyright: ignore[reportReturnType]
 
+    async def async_pre_call_deployment_hook(
+        self,
+        kwargs: Mapping[str, Any],
+        call_type: CallTypes | None,
+    ) -> dict[str, Any] | None:  # mutable-ok: overrides CustomLogger hook whose contract is a plain dict
+        if call_type not in (CallTypes.completion, CallTypes.acompletion):
+            return None
+        if not kwargs.get("stream"):
+            return None
+        if not has_headroom_retrieve_tool(kwargs.get("tools")):
+            return None
+        return {**kwargs, "stream": False, HEADROOM_CONVERTED_STREAM_KEY: True}
+
     async def async_should_run_agentic_loop(
         self,
         response: Any,

litellm/proxy/litellm_pre_call_utils.py

@@ -180,6 +180,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS = (
     "_code_interpreter_interception_converted_stream",
     "_code_interpreter_interception_sandbox_key",
     "_code_interpreter_interception_session_scoped",
+    "_headroom_interception_converted_stream",
     "max_agentic_loops",
 )
 

litellm/types/integrations/custom_logger.py

@@ -5,8 +5,14 @@ from pydantic import BaseModel, Field
 CHAT_COMPLETION_AGENTIC_SURFACE = "chat_completions"
 RESPONSES_AGENTIC_SURFACE = "responses"
 CODE_INTERPRETER_INTERCEPTION_PREFIX = "_code_interpreter_interception"
+HEADROOM_INTERCEPTION_PREFIX = "_headroom_interception"
+HEADROOM_CONVERTED_STREAM_KEY = f"{HEADROOM_INTERCEPTION_PREFIX}_converted_stream"
 NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES = frozenset(
-    ("_websearch_interception", "_compression_interception")
+    (
+        "_websearch_interception",
+        "_compression_interception",
+        HEADROOM_INTERCEPTION_PREFIX,
+    )
 )
 INTERCEPTION_INTERNAL_PREFIXES = frozenset(
     (

litellm/types/utils.py

@@ -3192,6 +3192,7 @@ agentic_loop_internal_litellm_params = [
     "_code_interpreter_interception_sandbox_key",
     "_code_interpreter_interception_session_scoped",
     "_code_interpreter_interception_converted_stream",
+    "_headroom_interception_converted_stream",
 ]
 
 all_litellm_params = (

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py

@@ -17,10 +17,13 @@ Tests cover:
 - CCR: headroom_retrieve tool injected when compressed messages contain hashes
 - CCR: async_should_run_agentic_loop returns True when response has headroom_retrieve tool calls
 - CCR: async_build_agentic_loop_plan calls retrieve endpoint and builds follow-up messages
+- CCR: streaming /chat/completions is converted to a non-streaming call so the agentic
+  loop resolves the retrieve tool call, then fake-streamed back to the client
 """
 
 import json
 import time
+from typing import Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
@@ -38,7 +41,16 @@ from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import (
 from litellm.proxy.spend_tracking.compression_savings import (
     extract_compression_saved_tokens,
 )
-from litellm.types.utils import GenericGuardrailAPIInputs
+from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY
+from litellm.types.utils import (
+    CallTypes,
+    ChatCompletionMessageToolCall,
+    Choices,
+    Function,
+    GenericGuardrailAPIInputs,
+    Message,
+    ModelResponse,
+)
 
 FAKE_API_BASE = "https://headroom.example.com"
 FAKE_API_KEY = "test-key"
@@ -1782,3 +1794,125 @@ async def test_fail_open_returns_original_parts_shapes():
 
     messages = result["structured_messages"]
     assert [m["content"] for m in messages] == [m["content"] for m in PARTS_MESSAGES]
+
+
+CCR_HASH = "b573993006976af767214fac"
+
+
+def _retrieve_tool_definition() -> dict:
+    return {
+        "type": "function",
+        "function": {
+            "name": HEADROOM_RETRIEVE_TOOL_NAME,
+            "description": "retrieve compressed content",
+            "parameters": {"type": "object", "properties": {"hash": {"type": "string"}}},
+        },
+    }
+
+
+def _model_response_with_retrieve_call() -> ModelResponse:
+    return ModelResponse(
+        choices=[
+            Choices(
+                finish_reason="tool_calls",
+                message=Message(
+                    role="assistant",
+                    content=None,
+                    tool_calls=[
+                        ChatCompletionMessageToolCall(
+                            id="call_ccr",
+                            type="function",
+                            function=Function(
+                                name=HEADROOM_RETRIEVE_TOOL_NAME,
+                                arguments=json.dumps({"hash": CCR_HASH}),
+                            ),
+                        )
+                    ],
+                ),
+            )
+        ]
+    )
+
+
+@pytest.mark.parametrize(
+    "call_type, stream, tools, expect_conversion",
+    [
+        (CallTypes.acompletion, True, [_retrieve_tool_definition()], True),
+        (CallTypes.completion, True, [_retrieve_tool_definition()], True),
+        (CallTypes.acompletion, False, [_retrieve_tool_definition()], False),
+        (CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False),
+        (CallTypes.acompletion, True, None, False),
+        (CallTypes.aresponses, True, [_retrieve_tool_definition()], False),
+        (CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False),
+    ],
+)
+@pytest.mark.asyncio
+async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions(
+    guardrail: HeadroomGuardrail,
+    call_type: CallTypes,
+    stream: bool,
+    tools: Optional[list],
+    expect_conversion: bool,
+):
+    kwargs = {"model": "gpt-4o", "stream": stream, "tools": tools}
+
+    result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type)
+
+    if not expect_conversion:
+        assert result is None
+        assert kwargs["stream"] is stream
+        return
+
+    assert result is not None
+    assert result["stream"] is False
+    assert result[HEADROOM_CONVERTED_STREAM_KEY] is True
+    assert kwargs["stream"] is True
+
+
+@pytest.mark.asyncio
+async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end(
+    guardrail: HeadroomGuardrail,
+):
+    """Regression test for streaming /chat/completions: the retrieve tool call the
+    model emits must be resolved by the agentic loop instead of being streamed back
+    to a client that never declared the tool."""
+    original_content = "the full uncompressed document"
+    final_answer = "the document says hello"
+    guardrail._issued_hashes_by_call_id["ccr-call-id"] = (
+        frozenset({CCR_HASH}),
+        time.monotonic() + 999,
+    )
+
+    real_acompletion = litellm.acompletion
+
+    async def acompletion_with_followup_answer(*args, **kwargs):
+        if kwargs.get("_agentic_loop_depth"):
+            kwargs["mock_response"] = final_answer
+        return await real_acompletion(*args, **kwargs)
+
+    saved_callbacks = list(litellm.callbacks)
+    litellm.callbacks = [guardrail]
+    try:
+        with patch.object(
+            guardrail.async_handler,
+            "get",
+            new_callable=AsyncMock,
+            return_value=_make_retrieve_response(original_content),
+        ) as mock_get, patch.object(litellm, "acompletion", new=acompletion_with_followup_answer):
+            response = await litellm.acompletion(
+                model="openai/gpt-4o",
+                messages=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}],
+                tools=[_retrieve_tool_definition()],
+                stream=True,
+                litellm_call_id="ccr-call-id",
+                mock_response=_model_response_with_retrieve_call(),
+            )
+            chunks = [chunk async for chunk in response]
+    finally:
+        litellm.callbacks = saved_callbacks
+
+    streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
+    assert streamed_text == final_answer
+    assert not any(chunk.choices and chunk.choices[0].delta.tool_calls for chunk in chunks)
+    mock_get.assert_called_once()
+    assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0])