From fa5a4f06edec9355d91b864fdc63341b708db8f0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:24:30 +0000 Subject: [PATCH] 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. --- .../chat_completion_agentic_loop.py | 49 +++++-- .../guardrail_hooks/headroom/headroom.py | 22 ++- litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/types/integrations/custom_logger.py | 8 +- litellm/types/utils.py | 1 + .../guardrail_hooks/test_headroom.py | 136 +++++++++++++++++- 6 files changed, 202 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index b7262a42324..39d36f25143 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 2735acd7787..1189d9841e4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/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, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d94fed0ee5b..bd451cb4a18 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/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", ) diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 04e490f79ee..8dddeb40e09 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/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( ( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e4dfac48141..5edd2806f7d 100644 --- a/litellm/types/utils.py +++ b/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 = ( 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 248893ed153..9f556a103da 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/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])