Merge 339639b0e7 into 5d207d85ee

litellm/proxy/common_request_processing.py

@@ -16,6 +16,7 @@ from typing import (
     Protocol,
     TypeAlias,
     TypeVar,
+    cast,  # noqa: TID251  # preserve the mutable JSON payload at the existing untyped boundary
     overload,
     runtime_checkable,
 )
@@ -116,13 +117,16 @@ from litellm.proxy.common_utils.sse_keepalive import (
 from litellm.proxy.dd_span_tagger import DDSpanTagger
 from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression
 from litellm.proxy.native_compaction import with_proxy_compaction_executor
-from litellm.proxy.route_llm_request import (
-    route_request,
-)
+from litellm.proxy.route_llm_request import route_request
 from litellm.proxy.utils import (  # noqa: F401  # legacy module exports
     ProxyLogging,
+    StreamingHookResponse,
+    StreamingReasoningState,
+    StreamingToolCallState,
     _check_and_merge_model_level_guardrails,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     check_and_merge_model_level_guardrails,
+    streaming_reasoning_with_response,
+    streaming_tool_calls_with_response,
 )
 from litellm.router import Router
 from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
@@ -3894,6 +3898,44 @@ class ProxyBaseLLMRequestProcessing:
 
     _finalize_streaming_generator_cleanup = finalize_streaming_generator_cleanup
 
+    @staticmethod
+    def _streaming_chunk_response_text(chunk: StreamingHookResponse | Mapping[str, object]) -> str:
+        if isinstance(chunk, (ModelResponse, ModelResponseStream)):
+            return litellm.get_response_string(response_obj=chunk)
+        if isinstance(chunk, Mapping):
+            return str(chunk.get("content", ""))
+        if isinstance(chunk, str):
+            return ""
+        try:
+            data: Final = chunk.model_dump(mode="json", exclude_none=True)
+            return str(data.get("content", ""))
+        except Exception:
+            return ""
+
+    @staticmethod
+    async def _apply_streaming_chunk_hook(
+        *,
+        chunk: Any,  # noqa: ANN401  # streaming callbacks may replace chunks with provider-specific response types
+        proxy_logging_obj: ProxyLogging,
+        user_api_key_dict: UserAPIKeyAuth,
+        request_data: dict[str, object],  # mutable-ok: ProxyLogging requires the existing mutable request payload
+        str_so_far: str,
+        streaming_tool_calls_so_far: StreamingToolCallState,
+        streaming_reasoning_so_far: StreamingReasoningState = (),
+    ) -> tuple[StreamingHookResponse, StreamingToolCallState, StreamingReasoningState]:
+        return (
+            await proxy_logging_obj.async_post_call_streaming_hook(
+                user_api_key_dict=user_api_key_dict,
+                response=chunk,
+                data=request_data,
+                str_so_far=str_so_far,
+                streaming_tool_calls_so_far=streaming_tool_calls_so_far,
+                streaming_reasoning_so_far=streaming_reasoning_so_far,
+            ),
+            streaming_tool_calls_with_response(streaming_tool_calls_so_far, chunk),
+            streaming_reasoning_with_response(streaming_reasoning_so_far, chunk),
+        )
+
     @staticmethod
     async def async_streaming_data_generator(
         response: object,
@@ -3945,6 +3987,8 @@ class ProxyBaseLLMRequestProcessing:
         )
         try:
             str_so_far = ""
+            streaming_tool_calls_so_far: StreamingToolCallState = ()  # rebind-ok: accumulated streaming state advances after each chunk
+            streaming_reasoning_so_far: StreamingReasoningState = ()  # rebind-ok: accumulated streaming state advances after each chunk
             async for chunk in guarded_stream:
                 # ``.format(chunk)`` was previously evaluated for every chunk
                 # regardless of log level; gate it behind the level check.
@@ -3952,25 +3996,21 @@ class ProxyBaseLLMRequestProcessing:
                     verbose_proxy_logger.debug("async_data_generator: received streaming chunk - %s", chunk)
 
                 if not fast_path:
-                    chunk = await proxy_logging_obj.async_post_call_streaming_hook(
+                    (
+                        chunk,
+                        streaming_tool_calls_so_far,
+                        streaming_reasoning_so_far,
+                    ) = await ProxyBaseLLMRequestProcessing._apply_streaming_chunk_hook(
+                        chunk=chunk,
+                        proxy_logging_obj=proxy_logging_obj,
                         user_api_key_dict=user_api_key_dict,
-                        response=chunk,
-                        data=request_data,
+                        request_data=cast(dict[str, object], request_data),  # cast-ok: request JSON has string keys
                         str_so_far=str_so_far,
+                        streaming_tool_calls_so_far=streaming_tool_calls_so_far,
+                        streaming_reasoning_so_far=streaming_reasoning_so_far,
                     )
 
-                    if isinstance(chunk, (ModelResponse, ModelResponseStream)):
-                        response_str = litellm.get_response_string(response_obj=chunk)
-                        str_so_far += response_str
-                    elif hasattr(chunk, "model_dump"):
-                        try:
-                            d = chunk.model_dump(mode="json", exclude_none=True)
-                            if isinstance(d, dict):
-                                str_so_far += str(d.get("content", ""))
-                        except Exception:
-                            pass
-                    elif isinstance(chunk, dict):
-                        str_so_far += str(chunk.get("content", ""))
+                    str_so_far += ProxyBaseLLMRequestProcessing._streaming_chunk_response_text(chunk)
 
                     model_name = request_data.get("model", "")
                     chunk = ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection(

litellm/proxy/proxy_server.py

@@ -866,6 +866,8 @@ from litellm.proxy.utils import (  # noqa: F401, RUF100  # legacy module exports
     PrismaClient,
     ProxyLogging,
     ProxyUpdateSpend,
+    StreamingReasoningState,
+    StreamingToolCallState,
     _cache_user_row,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     _get_docs_url,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     _get_openapi_url,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
@@ -893,6 +895,8 @@ from litellm.proxy.utils import (  # noqa: F401, RUF100  # legacy module exports
     migrate_passwords_to_scrypt_async,
     model_dump_with_preserved_fields,
     prefetch_config_params,
+    streaming_reasoning_with_response,
+    streaming_tool_calls_with_response,
     update_spend,
 )
 from litellm.proxy.video_endpoints.endpoints import router as video_router
@@ -9848,19 +9852,26 @@ async def _apply_streaming_chunk_hooks(
     user_api_key_dict: UserAPIKeyAuth,
     request_data: dict,
     str_so_far: str,
-) -> tuple[Any, str]:
+    streaming_tool_calls_so_far: StreamingToolCallState = (),
+    streaming_reasoning_so_far: StreamingReasoningState = (),
+) -> tuple[Any, str, StreamingToolCallState, StreamingReasoningState]:
+    stream_chunk: Final = chunk
     chunk = await proxy_logging_obj.async_post_call_streaming_hook(
         user_api_key_dict=user_api_key_dict,
         response=chunk,
         data=request_data,
         str_so_far=str_so_far if str_so_far else None,
+        streaming_tool_calls_so_far=streaming_tool_calls_so_far,
+        streaming_reasoning_so_far=streaming_reasoning_so_far,
     )
 
     if isinstance(chunk, (ModelResponse, ModelResponseStream)):
         response_str: Final = litellm.get_response_string(response_obj=chunk)
         str_so_far += response_str
 
-    return chunk, str_so_far
+    updated_tool_calls: Final = streaming_tool_calls_with_response(streaming_tool_calls_so_far, stream_chunk)
+    updated_reasoning: Final = streaming_reasoning_with_response(streaming_reasoning_so_far, stream_chunk)
+    return chunk, str_so_far, updated_tool_calls, updated_reasoning
 
 
 def _format_streaming_sse_chunk(chunk: str | bytes) -> str | bytes:
@@ -10128,6 +10139,8 @@ async def async_data_generator(
         # Previously "".join(str_so_far_parts) was called every chunk, re-joining
         # the entire accumulated response. String += is O(n) amortized total.
         _str_so_far: str = ""
+        _streaming_tool_calls_so_far: StreamingToolCallState = ()  # rebind-ok: accumulated streaming state advances after each chunk
+        _streaming_reasoning_so_far: StreamingReasoningState = ()  # rebind-ok: accumulated streaming state advances after each chunk
         # Separate iterator-level vs per-chunk hook decisions. The iterator
         # wrap is needed when any callback overrides
         # ``async_post_call_streaming_iterator_hook`` or has
@@ -10167,11 +10180,18 @@ async def async_data_generator(
             chunk = cast(Any, item)  # cast-ok: sentinel already handled above, item is a real chunk here
             if needs_per_chunk_hook:
                 ### CALL HOOKS ### - modify outgoing data
-                chunk, _str_so_far = await _apply_streaming_chunk_hooks(
+                (
+                    chunk,
+                    _str_so_far,
+                    _streaming_tool_calls_so_far,
+                    _streaming_reasoning_so_far,
+                ) = await _apply_streaming_chunk_hooks(
                     chunk=chunk,
                     user_api_key_dict=user_api_key_dict,
                     request_data=request_data,
                     str_so_far=_str_so_far,
+                    streaming_tool_calls_so_far=_streaming_tool_calls_so_far,
+                    streaming_reasoning_so_far=_streaming_reasoning_so_far,
                 )
 
             # Mid-stream fallbacks surface metadata on individual chunks rather than

litellm/proxy/utils.py

@@ -19,6 +19,7 @@ from collections.abc import (
     Awaitable,
     Callable,
     Coroutine,
+    Iterator,
     Mapping,
     Sequence,
 )
@@ -255,7 +256,7 @@ from litellm.repositories.verification_token_repository import (
 from litellm.router_utils.common_utils import resolve_model_group_alias
 from litellm.secret_managers.main import str_to_bool
 from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
-from litellm.types.llms.openai import ResponsesAPIResponse
+from litellm.types.llms.openai import ChatCompletionReasoningSummaryTextBlock, ResponsesAPIResponse
 from litellm.types.mcp import (
     MCPDuringCallResponseObject,
     MCPPreCallRequestObject,
@@ -263,7 +264,12 @@ from litellm.types.mcp import (
 )
 from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
 from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult
-from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
+from litellm.types.utils import (
+    ChatCompletionDeltaToolCall,
+    Delta,
+    LLMResponseTypes,
+    LoggedLiteLLMParams,
+)
 from litellm.utils import (
     add_custom_logger_callback_to_specific_event,
 )
@@ -1102,6 +1108,198 @@ def _call_type_for_route(route: str | None) -> str | None:
     return call_types[0].value if len(operations) == 1 else None
 
 
+StreamingHookResponse: TypeAlias = str | ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream
+
+
+class _StreamingHookResponseText(str):
+    """Marks the exact text object passed to a per-chunk streaming hook."""
+
+
+class _StructuredStreamingGuardrailText(_StreamingHookResponseText):
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class _StreamingToolCallFragment:
+    choice_index: int
+    tool_index: int
+    name: str
+    arguments: str
+
+    @property
+    def key(self) -> tuple[int, int]:
+        return self.choice_index, self.tool_index
+
+
+StreamingToolCallState: TypeAlias = tuple[_StreamingToolCallFragment, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class _StreamingReasoningFragment:
+    choice_index: int
+    text: str
+
+
+StreamingReasoningState: TypeAlias = tuple[_StreamingReasoningFragment, ...]
+
+
+def _streaming_reasoning_text(delta: Delta) -> str:
+    reasoning: Final = delta.reasoning_content if hasattr(delta, "reasoning_content") else None
+    thinking_blocks: Final = delta.thinking_blocks if hasattr(delta, "thinking_blocks") else None
+    reasoning_items: Final = delta.reasoning_items if hasattr(delta, "reasoning_items") else None
+    thinking: Final = tuple(
+        _thinking_block_text(block) for block in thinking_blocks or () if block["type"] == "thinking"
+    )
+    summaries: Final = tuple(_reasoning_summary_text(item.get("summary") or ()) for item in reasoning_items or ())
+    return "".join((reasoning or "", *thinking, *summaries))
+
+
+def _thinking_block_text(block: Mapping[str, object]) -> str:
+    text: Final = block.get("thinking")
+    return text if isinstance(text, str) else ""
+
+
+def _streaming_reasoning_fragments(response: ModelResponseStream) -> Iterator[_StreamingReasoningFragment]:
+    for choice in response.choices:
+        if text := _streaming_reasoning_text(choice.delta):
+            yield _StreamingReasoningFragment(choice_index=choice.index, text=text)
+
+
+def _reasoning_summary_text(summary: Sequence[ChatCompletionReasoningSummaryTextBlock]) -> str:
+    return "".join(block.get("text", "") for block in summary)
+
+
+def streaming_reasoning_with_response(
+    reasoning: Sequence[_StreamingReasoningFragment], response: object
+) -> StreamingReasoningState:
+    if not isinstance(response, ModelResponseStream):
+        return tuple(reasoning)
+    fragments: Final = (*reasoning, *_streaming_reasoning_fragments(response))
+    indices: Final = tuple(dict.fromkeys(fragment.choice_index for fragment in fragments))
+    return tuple(
+        _StreamingReasoningFragment(
+            choice_index=index,
+            text="".join(fragment.text for fragment in fragments if fragment.choice_index == index),
+        )
+        for index in indices
+    )
+
+
+def _streaming_hook_response_text(*, response_str: str, str_so_far: str | None, response: object) -> str:
+    complete_response: Final = str_so_far + response_str if str_so_far is not None else response_str
+    if complete_response == "" and isinstance(response, (ModelResponse, ModelResponseStream)):
+        return _StreamingHookResponseText(complete_response)
+    return complete_response
+
+
+def _streaming_tool_call_fragments(response: ModelResponseStream) -> Iterator[_StreamingToolCallFragment]:
+    for choice in response.choices:
+        for tool_call in choice.delta.tool_calls or ():
+            if isinstance(tool_call, ChatCompletionDeltaToolCall):
+                yield _StreamingToolCallFragment(
+                    choice_index=choice.index,
+                    tool_index=tool_call.index,
+                    name=tool_call.function.name or "",
+                    arguments=tool_call.function.arguments or "",
+                )
+            else:
+                yield _StreamingToolCallFragment(
+                    choice_index=choice.index,
+                    tool_index=tool_call.index,
+                    name=tool_call.custom.name or "",
+                    arguments=tool_call.custom.input or "",
+                )
+        if choice.delta.function_call is not None:
+            yield _StreamingToolCallFragment(
+                choice_index=choice.index,
+                tool_index=-1,
+                name=choice.delta.function_call.name or "",
+                arguments=choice.delta.function_call.arguments or "",
+            )
+
+
+def _assembled_streaming_tool_calls(
+    fragments: Sequence[_StreamingToolCallFragment],
+) -> StreamingToolCallState:
+    keys: Final = tuple(
+        fragment.key
+        for position, fragment in enumerate(fragments)
+        if fragment.key not in tuple(previous.key for previous in fragments[:position])
+    )
+    return tuple(
+        _StreamingToolCallFragment(
+            choice_index=key[0],
+            tool_index=key[1],
+            name="".join(fragment.name for fragment in fragments if fragment.key == key),
+            arguments="".join(fragment.arguments for fragment in fragments if fragment.key == key),
+        )
+        for key in keys
+    )
+
+
+def streaming_tool_calls_with_response(
+    tool_calls: Sequence[_StreamingToolCallFragment], response: object
+) -> StreamingToolCallState:
+    if not isinstance(response, ModelResponseStream):
+        return tuple(tool_calls)
+    fragments: Final = _streaming_tool_call_fragments(response)
+    return _assembled_streaming_tool_calls((*tool_calls, *fragments))
+
+
+def _streaming_guardrail_response_text(
+    *,
+    complete_response: str,
+    response: object,
+    streaming_tool_calls_so_far: Sequence[_StreamingToolCallFragment],
+    streaming_reasoning_so_far: Sequence[_StreamingReasoningFragment],
+) -> str:
+    if not isinstance(response, ModelResponseStream):
+        return complete_response
+    tool_calls: Final = streaming_tool_calls_with_response(streaming_tool_calls_so_far, response)
+    tool_fields: Final = tuple(
+        "tool_call:"
+        + json.dumps(
+            (tool_call.choice_index, tool_call.tool_index, tool_call.name, tool_call.arguments),
+            ensure_ascii=False,
+            separators=(",", ":"),
+        )
+        for tool_call in tool_calls
+    )
+    reasoning: Final = streaming_reasoning_with_response(streaming_reasoning_so_far, response)
+    reasoning_fields: Final = tuple(
+        "reasoning:" + json.dumps((fragment.choice_index, fragment.text), ensure_ascii=False, separators=(",", ":"))
+        for fragment in reasoning
+    )
+    structured_fields: Final = (*tool_fields, *reasoning_fields)
+    if not structured_fields:
+        return complete_response
+    return _StructuredStreamingGuardrailText("\n".join((str(complete_response), *structured_fields)))
+
+
+def _is_unchanged_structured_streaming_hook_response(
+    *, callback_response: object, complete_response: str, response_str: str, response: object
+) -> bool:
+    if not isinstance(response, (ModelResponse, ModelResponseStream)):
+        return False
+    if isinstance(complete_response, _StructuredStreamingGuardrailText):
+        return callback_response == complete_response
+    if isinstance(complete_response, _StreamingHookResponseText):
+        return callback_response is complete_response
+    if response_str != "" and not (
+        isinstance(response, ModelResponseStream)
+        and any(
+            choice.delta.tool_calls
+            or choice.delta.function_call is not None
+            or getattr(choice.delta, "reasoning_content", None)
+            or getattr(choice.delta, "thinking_blocks", None)
+            or getattr(choice.delta, "reasoning_items", None)
+            for choice in response.choices
+        )
+    ):
+        return False
+    return callback_response == complete_response
+
+
 _PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException)
 _LOG_DB_METRICS_CALL_TYPES: Final = frozenset(("get_data", "insert_data", "update_data", "delete_data"))
 
@@ -3885,6 +4083,8 @@ class ProxyLogging:
         response: ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream,
         user_api_key_dict: UserAPIKeyAuth,
         str_so_far: str | None = None,
+        streaming_tool_calls_so_far: Sequence[_StreamingToolCallFragment] = (),
+        streaming_reasoning_so_far: Sequence[_StreamingReasoningFragment] = (),
     ):
         """
         Allow user to modify outgoing streaming data -> per chunk
@@ -3904,6 +4104,7 @@ class ProxyLogging:
 
         from litellm.proxy.proxy_server import llm_router
 
+        hook_response: StreamingHookResponse = response  # rebind-ok: each callback may replace the chunk
         response_str: str | None = None
         if isinstance(response, (ModelResponse, ModelResponseStream)):
             response_str = litellm.get_response_string(response_obj=response)
@@ -3951,22 +4152,35 @@ class ProxyLogging:
                     else:
                         _callback = callback
                     if _callback is not None and isinstance(_callback, CustomLogger):
-                        if str_so_far is not None:
-                            complete_response = str_so_far + response_str
-                        else:
-                            complete_response = response_str
+                        complete_response = _streaming_hook_response_text(
+                            response_str=response_str,
+                            str_so_far=str_so_far,
+                            response=hook_response,
+                        )
+                        if isinstance(_callback, CustomGuardrail):
+                            complete_response = _streaming_guardrail_response_text(
+                                complete_response=complete_response,
+                                response=hook_response,
+                                streaming_tool_calls_so_far=streaming_tool_calls_so_far,
+                                streaming_reasoning_so_far=streaming_reasoning_so_far,
+                            )
                         callback_response: (
-                            ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream | None
+                            str | ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream | None
                         )
                         callback_response = await _callback.async_post_call_streaming_hook(
                             user_api_key_dict=user_api_key_dict,
                             response=complete_response,
                         )
-                        if callback_response is not None:
-                            response = callback_response
+                        if callback_response is not None and not _is_unchanged_structured_streaming_hook_response(
+                            callback_response=callback_response,
+                            complete_response=complete_response,
+                            response_str=response_str,
+                            response=hook_response,
+                        ):
+                            hook_response = callback_response
                 except Exception as e:
                     raise e
-        return response
+        return hook_response
 
     async def async_post_call_streaming_iterator_hook(
         self,

tests/unit/proxy/proxy_server/test_streaming_helpers.py

@@ -31,6 +31,7 @@ from pydantic import BaseModel
 import litellm
 import litellm.proxy.proxy_server as ps
 from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY
+from litellm.integrations.custom_logger import CustomLogger
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.proxy_server import (
     _apply_streaming_chunk_hooks,
@@ -78,6 +79,30 @@ def _simple_chunk(model: str = "gpt-4", content: str = "hi") -> ModelResponseStr
     )
 
 
+def _tool_call_chunk(arguments: str) -> ModelResponseStream:
+    return ModelResponseStream(
+        id="chatcmpl-tool-call",
+        choices=[
+            {
+                "index": 0,
+                "delta": {
+                    "tool_calls": [
+                        {
+                            "index": 0,
+                            "type": "function",
+                            "function": {"name": "run_command", "arguments": arguments},
+                        }
+                    ]
+                },
+                "finish_reason": None,
+            }
+        ],
+        created=0,
+        model="gpt-4",
+        object="chat.completion.chunk",
+    )
+
+
 async def _async_iter(items):
     for it in items:
         yield it
@@ -589,12 +614,15 @@ def test_serialize_streaming_chunk_invalid_input_raises_attribute_error():
 async def test_apply_streaming_chunk_hooks_appends_to_str_so_far(monkeypatch):
     chunk = _simple_chunk(content="abc")
 
-    async def _passthrough(*, user_api_key_dict, response, data, str_so_far=None):
+    async def _passthrough(
+        *, user_api_key_dict, response, data, str_so_far=None, streaming_tool_calls_so_far=(),
+        streaming_reasoning_so_far=(),
+    ):
         return response
 
     monkeypatch.setattr(ps.proxy_logging_obj, "async_post_call_streaming_hook", _passthrough)
 
-    new_chunk, new_str = await _apply_streaming_chunk_hooks(
+    new_chunk, new_str, tool_calls, _ = await _apply_streaming_chunk_hooks(
         chunk=chunk,
         user_api_key_dict=_user_auth(),
         request_data={},
@@ -605,14 +633,40 @@ async def test_apply_streaming_chunk_hooks_appends_to_str_so_far(monkeypatch):
         "chunk_is_basemodel": isinstance(new_chunk, ModelResponseStream),
         "str_so_far": new_str,
         "grew": len(new_str) > len("prior:"),
+        "tool_calls": tool_calls,
     }
     assert observed == {
         "chunk_is_basemodel": True,
         "str_so_far": "prior:abc",
         "grew": True,
+        "tool_calls": (),
     }
 
 
+@pytest.mark.asyncio
+async def test_apply_streaming_chunk_hooks_compacts_tool_call_history(monkeypatch):
+    async def _passthrough(
+        *, user_api_key_dict, response, data, str_so_far=None, streaming_tool_calls_so_far=(),
+        streaming_reasoning_so_far=(),
+    ):
+        return response
+
+    monkeypatch.setattr(ps.proxy_logging_obj, "async_post_call_streaming_hook", _passthrough)
+    stream_state = ()
+
+    for arguments in ('{"command":"', "blocked-", "command", '"}'):
+        _, _, stream_state, _ = await _apply_streaming_chunk_hooks(
+            chunk=_tool_call_chunk(arguments),
+            user_api_key_dict=_user_auth(),
+            request_data={},
+            str_so_far="",
+            streaming_tool_calls_so_far=stream_state,
+        )
+
+    assert len(stream_state) == 1
+    assert stream_state[0].arguments == '{"command":"blocked-command"}'
+
+
 @pytest.mark.asyncio
 async def test_apply_streaming_chunk_hooks_hook_raises_exception(monkeypatch):
     async def _boom(*args, **kwargs):
@@ -715,6 +769,80 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch):
     }
 
 
+@pytest.mark.asyncio
+async def test_async_data_generator_preserves_tool_calls_through_per_chunk_hook(
+    monkeypatch,
+):
+    class _PassThrough(CustomLogger):
+        async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+            return response
+
+    callback = _PassThrough()
+    monkeypatch.setattr(litellm, "callbacks", [callback])
+    _patch_logging_flags(monkeypatch, needs_per_chunk=True)
+
+    chunk = ModelResponseStream(
+        id="chatcmpl-tools",
+        choices=[
+            {
+                "index": 0,
+                "delta": {
+                    "tool_calls": [
+                        {
+                            "index": 0,
+                            "id": "call-weather",
+                            "type": "function",
+                            "function": {
+                                "name": "get_weather",
+                                "arguments": "",
+                            },
+                        },
+                        {
+                            "index": 1,
+                            "id": "call-time",
+                            "type": "function",
+                            "function": {"name": "get_time", "arguments": ""},
+                        },
+                    ]
+                },
+                "finish_reason": None,
+            }
+        ],
+        created=0,
+        model="gpt-4o-mini",
+        object="chat.completion.chunk",
+    )
+
+    out = []
+    async for line in async_data_generator(
+        response=_async_iter([chunk]),
+        user_api_key_dict=_user_auth(),
+        request_data={"model": "gpt-4o-mini"},
+    ):
+        out.append(line)
+
+    first = out[0]
+    assert isinstance(first, (str, bytes))
+    first_text = first.decode() if isinstance(first, bytes) else first
+    assert first_text.startswith("data: {")
+    payload = json.loads(first_text.removeprefix("data: ").removesuffix("\n\n"))
+    assert payload["choices"][0]["delta"]["tool_calls"] == [
+        {
+            "index": 0,
+            "id": "call-weather",
+            "type": "function",
+            "function": {"name": "get_weather", "arguments": ""},
+        },
+        {
+            "index": 1,
+            "id": "call-time",
+            "type": "function",
+            "function": {"name": "get_time", "arguments": ""},
+        },
+    ]
+    assert out[-1] == "data: [DONE]\n\n"
+
+
 @pytest.mark.asyncio
 async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch):
     _patch_logging_flags(monkeypatch)

tests/unit/proxy/test_common_request_processing.py

@@ -4824,6 +4824,48 @@ class TestAsyncStreamingDataGeneratorFastPath:
 
         ProxyLogging._callback_capabilities_cache.clear()
 
+    @pytest.mark.asyncio
+    async def test_apply_streaming_chunk_hook_compacts_tool_call_history(self):
+        proxy_logging_obj = MagicMock(spec=ProxyLogging)
+        proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"])
+
+        def _chunk(arguments: str):
+            return litellm.ModelResponseStream(
+                id="chatcmpl-tool-call",
+                choices=[
+                    {
+                        "index": 0,
+                        "delta": {
+                            "tool_calls": [
+                                {
+                                    "index": 0,
+                                    "type": "function",
+                                    "function": {"name": "run_command", "arguments": arguments},
+                                }
+                            ]
+                        },
+                        "finish_reason": None,
+                    }
+                ],
+                created=0,
+                model="gpt-4",
+                object="chat.completion.chunk",
+            )
+
+        stream_state = ()
+        for arguments in ("blocked-", "command"):
+            _, stream_state, _ = await ProxyBaseLLMRequestProcessing._apply_streaming_chunk_hook(
+                chunk=_chunk(arguments),
+                proxy_logging_obj=proxy_logging_obj,
+                user_api_key_dict=MagicMock(spec=ProxyUserAPIKeyAuth),
+                request_data={},
+                str_so_far="",
+                streaming_tool_calls_so_far=stream_state,
+            )
+
+        assert len(stream_state) == 1
+        assert stream_state[0].arguments == "blocked-command"
+
 
 class TestDisconnectGatherCleanup:
     def _disconnect_request(self) -> Request:

tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py

@@ -10,11 +10,13 @@ Covers ``_wrap_streaming_iterator_with_enrichment``,
 from __future__ import annotations
 
 import asyncio
+import json
 from collections.abc import AsyncGenerator, AsyncIterator
 from datetime import datetime
 from typing import Any, Dict, Final, List
 from unittest.mock import AsyncMock, MagicMock
 
+import httpx
 import pytest
 from fastapi import HTTPException
 
@@ -27,7 +29,7 @@ from litellm.llms.anthropic.pass_through.messages.streaming_iterator import (
     BaseAnthropicMessagesStreamingIterator,
 )
 from litellm.proxy._types import UserAPIKeyAuth
-from litellm.proxy.utils import ProxyLogging
+from litellm.proxy.utils import ProxyLogging, streaming_tool_calls_with_response
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import Usage
 
@@ -384,6 +386,241 @@ async def test_async_post_call_streaming_hook_invokes_per_chunk_callback(proxy_l
     assert out.startswith("modified-")
 
 
+@pytest.mark.parametrize("str_so_far", [None, "I will check that. "])
+@pytest.mark.parametrize("content", [None, "Checking those now. "])
+@pytest.mark.asyncio
+async def test_async_post_call_streaming_hook_preserves_tool_calls_when_callback_returns_unmodified_text(
+    proxy_logging, make_user_api_key_auth, monkeypatch, str_so_far, content
+):
+    class _PassThrough(CustomLogger):
+        async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+            return response
+
+    monkeypatch.setattr(litellm, "callbacks", [_PassThrough()])
+
+    response = litellm.ModelResponseStream(
+        id="chatcmpl-tools",
+        choices=[
+            {
+                "index": 0,
+                "delta": {
+                    "content": content,
+                    "tool_calls": [
+                        {
+                            "index": 0,
+                            "id": "call-weather",
+                            "type": "function",
+                            "function": {"name": "get_weather", "arguments": ""},
+                        },
+                        {
+                            "index": 1,
+                            "id": "call-time",
+                            "type": "function",
+                            "function": {"name": "get_time", "arguments": ""},
+                        },
+                    ]
+                },
+                "finish_reason": None,
+            }
+        ],
+        created=0,
+        model="gpt-4o-mini",
+        object="chat.completion.chunk",
+    )
+
+    out = await proxy_logging.async_post_call_streaming_hook(
+        data={},
+        response=response,
+        user_api_key_dict=make_user_api_key_auth(),
+        str_so_far=str_so_far,
+    )
+
+    assert out is response
+    assert [
+        tool_call.model_dump(exclude_none=True)
+        for tool_call in out.choices[0].delta.tool_calls
+    ] == [
+        {
+            "id": "call-weather",
+            "function": {"arguments": "", "name": "get_weather"},
+            "type": "function",
+            "index": 0,
+        },
+        {
+            "id": "call-time",
+            "function": {"arguments": "", "name": "get_time"},
+            "type": "function",
+            "index": 1,
+        },
+    ]
+
+    class _Replace(CustomLogger):
+        async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+            return "replacement"
+
+    monkeypatch.setattr(litellm, "callbacks", [_PassThrough(), _Replace()])
+    replaced = await proxy_logging.async_post_call_streaming_hook(
+        data={},
+        response=response,
+        user_api_key_dict=make_user_api_key_auth(),
+        str_so_far=str_so_far,
+    )
+    assert replaced == "replacement"
+
+    if str_so_far is not None:
+        class _EquivalentCopy(CustomLogger):
+            async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+                return response.encode().decode()
+
+        monkeypatch.setattr(litellm, "callbacks", [_EquivalentCopy()])
+        equivalent = await proxy_logging.async_post_call_streaming_hook(
+            data={},
+            response=response,
+            user_api_key_dict=make_user_api_key_auth(),
+            str_so_far=str_so_far,
+        )
+        assert equivalent is response
+
+    class _Suppress(CustomLogger):
+        async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+            return ""
+
+    monkeypatch.setattr(litellm, "callbacks", [_Suppress()])
+    suppressed = await proxy_logging.async_post_call_streaming_hook(
+        data={},
+        response=response,
+        user_api_key_dict=make_user_api_key_auth(),
+        str_so_far=str_so_far,
+    )
+    assert suppressed == ""
+
+
+@pytest.mark.asyncio
+async def test_async_post_call_streaming_hook_exposes_tool_calls_to_guardrails(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    class _BlockingGuardrail(CustomGuardrail):
+        def should_run_guardrail(self, data, event_type):
+            return True
+
+        async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+            if "blocked-command" in response:
+                return "data: blocked\n\n"
+            return response
+
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [_BlockingGuardrail(guardrail_name="tool-call-scanner")],
+    )
+
+    first_chunk = litellm.ModelResponseStream(
+        id="chatcmpl-guardrail-tools",
+        choices=[
+            {
+                "index": 0,
+                "delta": {
+                    "tool_calls": [
+                        {
+                            "index": 0,
+                            "id": "call-shell",
+                            "type": "function",
+                            "function": {
+                                "name": "run_command",
+                                "arguments": '{"command":"blocked-',
+                            },
+                        }
+                    ]
+                },
+                "finish_reason": None,
+            }
+        ],
+        created=0,
+        model="gpt-4o-mini",
+        object="chat.completion.chunk",
+    )
+    second_chunk = litellm.ModelResponseStream(
+        id="chatcmpl-guardrail-tools",
+        choices=[
+            {
+                "index": 0,
+                "delta": {
+                    "tool_calls": [
+                        {
+                            "index": 0,
+                            "function": {"arguments": 'command"}'},
+                        }
+                    ]
+                },
+                "finish_reason": None,
+            }
+        ],
+        created=0,
+        model="gpt-4o-mini",
+        object="chat.completion.chunk",
+    )
+
+    out = await proxy_logging.async_post_call_streaming_hook(
+        data={},
+        response=second_chunk,
+        user_api_key_dict=make_user_api_key_auth(),
+        streaming_tool_calls_so_far=streaming_tool_calls_with_response((), first_chunk),
+    )
+
+    assert out == "data: blocked\n\n"
+
+
+@pytest.mark.asyncio
+async def test_async_post_call_streaming_hook_preserves_equal_guardrail_projection(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    class _CopyingGuardrail(CustomGuardrail):
+        def should_run_guardrail(self, data, event_type):
+            return True
+
+        async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+            return response.encode().decode()
+
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [_CopyingGuardrail(guardrail_name="tool-call-scanner")],
+    )
+    response = litellm.ModelResponseStream(
+        id="chatcmpl-guardrail-copy",
+        choices=[
+            {
+                "index": 0,
+                "delta": {
+                    "tool_calls": [
+                        {
+                            "index": 0,
+                            "id": "call-weather",
+                            "type": "function",
+                            "function": {
+                                "name": "get_weather",
+                                "arguments": '{"city":"Paris"}',
+                            },
+                        }
+                    ]
+                },
+                "finish_reason": None,
+            }
+        ],
+        created=0,
+        model="gpt-4o-mini",
+        object="chat.completion.chunk",
+    )
+
+    out = await proxy_logging.async_post_call_streaming_hook(
+        data={},
+        response=response,
+        user_api_key_dict=make_user_api_key_auth(),
+    )
+
+    assert out is response
+
+
 @pytest.mark.asyncio
 async def test_async_post_call_streaming_hook_callback_error_raises(proxy_logging, make_user_api_key_auth, monkeypatch):
     class _Per(CustomLogger):
@@ -901,3 +1138,173 @@ async def test_stream_provider_failure_is_not_attributed_to_any_guardrail(
     with pytest.raises(RuntimeError, match="provider connection dropped"):
         await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _failing_provider_stream(), request_data)
     assert request_data["metadata"] == {}
+
+
+@pytest.fixture
+def reasoning_guardrail(proxy_logging, monkeypatch):
+    from litellm.proxy.guardrails.guardrail_hooks.azure.text_moderation import (
+        AzureContentSafetyTextModerationGuardrail,
+    )
+
+    guardrail: Final = AzureContentSafetyTextModerationGuardrail(
+        guardrail_name="reasoning-scanner", api_key="test-key", api_base="https://moderation.test",
+        event_hook=GuardrailEventHooks.post_call, default_on=True,
+    )
+    monkeypatch.setattr(litellm, "callbacks", [guardrail])
+    return guardrail
+
+
+def _moderate_reasoning(request: httpx.Request) -> httpx.Response:
+    text: Final = json.loads(request.content)["text"]
+    assert "opaque-marker" not in text
+    return httpx.Response(200, json={
+        "blocklistsMatch": [],
+        "categoriesAnalysis": [{"category": "Violence", "severity": 6 if "blocked-command" in text else 0}],
+    })
+
+
+def _reasoning_chunk(field: str, text: str, choice_index: int = 0) -> litellm.ModelResponseStream:
+    delta: Final = {
+        "reasoning_content": {"reasoning_content": text},
+        "thinking_blocks": {"thinking_blocks": [
+            {"type": "thinking", "thinking": text, "signature": "opaque-marker"},
+            {"type": "redacted_thinking", "data": "opaque-marker"},
+        ]},
+        "reasoning_items": {"reasoning_items": [{
+            "type": "reasoning", "id": "r1", "encrypted_content": "opaque-marker",
+            "summary": [{"type": "summary_text", "text": text}],
+        }]},
+    }[field]
+    return litellm.ModelResponseStream(choices=[{"index": choice_index, "delta": delta}])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("field", ["reasoning_content", "thinking_blocks", "reasoning_items"])
+@pytest.mark.parametrize("text", ["blocked-command", "safe reasoning", ""])
+async def test_native_azure_moderates_plaintext_reasoning(
+    field, text, reasoning_guardrail, proxy_logging, make_user_api_key_auth,
+):
+    chunk: Final = _reasoning_chunk(field, text)
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_moderate_reasoning)) as handler:
+        reasoning_guardrail.async_handler = handler
+        result: Final = await proxy_logging.async_post_call_streaming_hook(
+            data={}, response=chunk, user_api_key_dict=make_user_api_key_auth(),
+        )
+    if text == "blocked-command":
+        assert isinstance(result, str)
+        assert "Azure Content Safety Guardrail" in result
+        assert "blocked-command" not in result
+    else:
+        assert result is chunk
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("field", ["reasoning_content", "thinking_blocks", "reasoning_items"])
+@pytest.mark.parametrize("path", ["proxy_server", "common_request_processing"])
+@pytest.mark.parametrize("same_choice", [True, False])
+async def test_native_azure_moderates_split_reasoning_in_both_streams(
+    field, path, same_choice, reasoning_guardrail, proxy_logging, make_user_api_key_auth,
+):
+    from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+    from litellm.proxy.proxy_server import async_data_generator
+
+    chunks: Final = (_reasoning_chunk(field, "blocked-"), _reasoning_chunk(field, "command", 0 if same_choice else 1))
+
+    async def source():
+        for chunk in chunks:
+            yield chunk
+
+    stream: Final = (
+        async_data_generator(response=source(), user_api_key_dict=make_user_api_key_auth(), request_data={})
+        if path == "proxy_server"
+        else ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
+            response=source(), user_api_key_dict=make_user_api_key_auth(), request_data={},
+            proxy_logging_obj=proxy_logging,
+            serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk,
+            serialize_error=lambda error: f"data: {error}\n\n",
+        )
+    )
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_moderate_reasoning)) as handler:
+        reasoning_guardrail.async_handler = handler
+        output: Final = tuple([part async for part in stream])
+    assert len(output) >= 2
+    if same_choice:
+        assert isinstance(output[1], str)
+        assert "Azure Content Safety Guardrail" in output[1]
+        assert '"command"' not in output[1]
+    else:
+        if path == "common_request_processing":
+            assert output[1] is chunks[1]
+        else:
+            assert isinstance(output[1], str)
+            assert json.loads(output[1].removeprefix("data: "))["choices"] == chunks[1].model_dump(exclude_none=True)["choices"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("callback_kind", ["ordinary", "disabled_guardrail"])
+@pytest.mark.parametrize("content", [None, "answer"])
+async def test_reasoning_chunks_preserve_ordinary_and_disabled_callback_behavior(
+    callback_kind, content, proxy_logging, make_user_api_key_auth, monkeypatch,
+):
+    class OrdinaryCallback(CustomLogger):
+        async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+            assert "blocked-command" not in response
+            return response
+
+    class DisabledGuardrail(CustomGuardrail):
+        def should_run_guardrail(self, data, event_type):
+            return False
+
+        async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+            pytest.fail("Disabled guardrail ran")
+
+    callback: Final = OrdinaryCallback() if callback_kind == "ordinary" else DisabledGuardrail(guardrail_name="disabled")
+    monkeypatch.setattr(litellm, "callbacks", [callback])
+    chunk: Final = litellm.ModelResponseStream(choices=[{
+        "index": 0, "delta": {"content": content, "reasoning_content": "blocked-command"},
+    }])
+    result: Final = await proxy_logging.async_post_call_streaming_hook(
+        data={}, response=chunk, user_api_key_dict=make_user_api_key_auth(),
+    )
+    assert result is chunk
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("path", ["proxy_server", "common_request_processing"])
+async def test_reasoning_stream_preserves_opaque_items_without_optional_plaintext(
+    path, reasoning_guardrail, proxy_logging, make_user_api_key_auth,
+):
+    from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+    from litellm.proxy.proxy_server import async_data_generator
+
+    chunks: Final = tuple(litellm.ModelResponseStream(choices=[{
+        "index": 0, "delta": {"reasoning_items": [item]},
+    }]) for item in (
+        {"type": "reasoning", "id": "encrypted", "encrypted_content": "opaque-marker"},
+        {"type": "reasoning", "id": "empty-summary", "summary": [{"type": "summary_text"}]},
+    ))
+
+    async def source():
+        for chunk in chunks:
+            yield chunk
+
+    stream: Final = (
+        async_data_generator(response=source(), user_api_key_dict=make_user_api_key_auth(), request_data={})
+        if path == "proxy_server"
+        else ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
+            response=source(), user_api_key_dict=make_user_api_key_auth(), request_data={},
+            proxy_logging_obj=proxy_logging,
+            serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk,
+            serialize_error=lambda error: f"data: {error}\n\n",
+        )
+    )
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_moderate_reasoning)) as handler:
+        reasoning_guardrail.async_handler = handler
+        output: Final = tuple([part async for part in stream])
+    assert len(output) >= 2
+    if path == "common_request_processing":
+        assert output[:2] == chunks
+    else:
+        for chunk, frame in zip(chunks, output[:2]):
+            assert isinstance(frame, str)
+            assert json.loads(frame.removeprefix("data: "))["choices"] == chunk.model_dump(exclude_none=True)["choices"]