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 1/5] 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]) From bff46fd782c709dc5ccc37d10082d52ae59e9441 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 27 Aug 2026 22:15:41 +0000 Subject: [PATCH 2/5] chore: reformat headroom stream constant and regenerate stale openapi snapshot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 44 +++++++++++++++++++ litellm/types/integrations/custom_logger.py | 4 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 3 files changed, 47 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index e30750ef565..1963c7799a2 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -15038,6 +15038,17 @@ } ], "title": "Upstream Resource" + }, + "upstream_token_header": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Token Header" } }, "title": "MCPCredentials", @@ -17518,6 +17529,17 @@ } ], "title": "Upstream Resource" + }, + "upstream_token_header": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Token Header" } }, "title": "MCPCredentials", @@ -20352,6 +20374,17 @@ } ], "title": "Upstream Resource" + }, + "upstream_token_header": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Token Header" } }, "title": "MCPCredentials", @@ -23699,6 +23732,17 @@ } ], "title": "Upstream Resource" + }, + "upstream_token_header": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Token Header" } }, "title": "MCPCredentials", diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index fd4648bff14..9a714e1724e 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -6,9 +6,7 @@ CHAT_COMPLETION_AGENTIC_SURFACE: Final = "chat_completions" RESPONSES_AGENTIC_SURFACE: Final = "responses" CODE_INTERPRETER_INTERCEPTION_PREFIX: Final = "_code_interpreter_interception" HEADROOM_INTERCEPTION_PREFIX: Final = "_headroom_interception" -HEADROOM_CONVERTED_STREAM_KEY: Final = ( - f"{HEADROOM_INTERCEPTION_PREFIX}_converted_stream" -) +HEADROOM_CONVERTED_STREAM_KEY: Final = f"{HEADROOM_INTERCEPTION_PREFIX}_converted_stream" NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES: Final = frozenset( ( "_websearch_interception", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index c124cc2e9c8..03a4d1142f8 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30208,6 +30208,8 @@ export interface components { token_exchange_profile?: string | null; /** Upstream Resource */ upstream_resource?: string | null; + /** Upstream Token Header */ + upstream_token_header?: string | null; }; /** * MCPEnvVar From 1cf1a97b731f99b9276d299bf5c7fc0487802b52 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 27 Aug 2026 22:37:08 +0000 Subject: [PATCH 3/5] chore: annotate headroom pre-call hook dict for the type-discipline budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/guardrails/guardrail_hooks/headroom/headroom.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index e4fe392b97f..2b34c0eeccb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -728,7 +728,11 @@ class HeadroomGuardrail(CustomGuardrail): return None if not has_headroom_retrieve_tool(kwargs.get("tools")): return None - return {**kwargs, "stream": False, HEADROOM_CONVERTED_STREAM_KEY: True} + return { # mutable-ok: the hook contract is a plain dict the router merges into the request kwargs + **kwargs, + "stream": False, + HEADROOM_CONVERTED_STREAM_KEY: True, + } async def async_should_run_agentic_loop( self, From d556a042866921af4f0ec583167159be660d551b Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 27 Aug 2026 22:49:23 +0000 Subject: [PATCH 4/5] test(headroom): fake the HTTP boundary in the streaming CCR regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 35 ++++++ .../guardrail_hooks/test_headroom.py | 116 ++++++++++-------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 3 files changed, 104 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 1963c7799a2..040d258f97a 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -10238,6 +10238,18 @@ "description": "AWS Bedrock runtime endpoint URL", "title": "Aws Bedrock Runtime Endpoint" }, + "aws_external_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "External ID required by the target role's trust policy on sts:AssumeRole", + "title": "Aws External Id" + }, "aws_profile_name": { "anyOf": [ { @@ -25237,6 +25249,9 @@ }, { "$ref": "#/components/schemas/ChatCompletionImageObject" + }, + { + "$ref": "#/components/schemas/ChatCompletionToolReferenceObject" } ] }, @@ -25324,6 +25339,26 @@ "title": "ChatCompletionToolParamFunctionChunk", "type": "object" }, + "ChatCompletionToolReferenceObject": { + "description": "Anthropic tool-search result block, carried through untouched so it survives a round trip.", + "properties": { + "tool_name": { + "title": "Tool Name", + "type": "string" + }, + "type": { + "const": "tool_reference", + "title": "Type", + "type": "string" + } + }, + "required": [ + "type", + "tool_name" + ], + "title": "ChatCompletionToolReferenceObject", + "type": "object" + }, "ChatCompletionUserMessage": { "properties": { "cache_control": { 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 420b91c983b..52df875e714 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -28,6 +28,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx from fastapi import HTTPException import litellm @@ -44,12 +45,7 @@ from litellm.proxy.spend_tracking.compression_savings import ( 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" @@ -1919,30 +1915,41 @@ def _retrieve_tool_definition() -> dict: } -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}), - ), - ) - ], - ), - ) - ] +def _openai_completion_payload(message: dict, finish_reason: str) -> dict: + return { + "id": "chatcmpl-ccr", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o", + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +def _openai_tool_call_payload() -> dict: + return _openai_completion_payload( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_ccr", + "type": "function", + "function": { + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": CCR_HASH}), + }, + } + ], + }, + "tool_calls", ) +def _openai_text_payload(content: str) -> dict: + return _openai_completion_payload({"role": "assistant", "content": content}, "stop") + + @pytest.mark.parametrize( "call_type, stream, tools, expect_conversion", [ @@ -1981,6 +1988,8 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple @pytest.mark.asyncio async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ): """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 @@ -1992,33 +2001,30 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( time.monotonic() + 999, ) - real_acompletion = litellm.acompletion + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=[ + httpx.Response(200, json=_openai_tool_call_payload()), + httpx.Response(200, json=_openai_text_payload(final_answer)), + ] + ) - 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 + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + 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", + ) + chunks = [chunk async for chunk in response] streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) assert streamed_text == final_answer @@ -2026,6 +2032,12 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( mock_get.assert_called_once() assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + assert len(upstream.calls) == 2 + followup_body = json.loads(upstream.calls[1].request.content) + assert not followup_body.get("stream") + assert original_content in json.dumps(followup_body["messages"]) + assert not any(key.startswith("_headroom_interception") for key in followup_body) + # --------------------------------------------------------------------------- # LIT-5018: the turn the model is being asked to act on is never compressed. diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 25fbd53018a..8945710cba2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29466,6 +29466,11 @@ export interface components { * @description AWS Bedrock runtime endpoint URL */ aws_bedrock_runtime_endpoint?: string | null; + /** + * Aws External Id + * @description External ID required by the target role's trust policy on sts:AssumeRole + */ + aws_external_id?: string | null; /** * Aws Profile Name * @description AWS profile name for credential retrieval From a14daf18d6ebbe399ec271207a9ea7f9124155b8 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 27 Aug 2026 23:19:17 +0000 Subject: [PATCH 5/5] refactor(agentic-loop): narrow logging_obj before building the fake stream wrapper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/chat_completion_agentic_loop.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index a5f5c6ac68f..c8e9e2583ba 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -10,6 +10,8 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, ) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, HEADROOM_CONVERTED_STREAM_KEY, @@ -104,9 +106,8 @@ def _wrap_response_as_fake_stream( ) -> object: if isinstance(response, CustomStreamWrapper): return response - if not isinstance(response, ModelResponse): + if not isinstance(response, ModelResponse) or not isinstance(logging_obj, LiteLLMLoggingObject): return response - from litellm.llms.base_llm.base_model_iterator import MockResponseIterator return CustomStreamWrapper( completion_stream=MockResponseIterator(model_response=response),