Merge fc113ad742 into 85fdea0fd1

litellm/litellm_core_utils/streaming_handler.py

@@ -2594,18 +2594,38 @@ class CustomStreamWrapper:
                         if processed_chunk is None:
                             continue
 
+                        if self.logging_obj.completion_start_time is None:
+                            self.logging_obj.update_completion_start_time(completion_start_time=datetime.datetime.now())
+
                         if processed_chunk.choices:
                             choice = processed_chunk.choices[0]
                             if isinstance(choice, StreamingChoices):
                                 self._accumulate_streamed_delta(choice.delta)
                         self.rules.post_call_rules(input=self.response_uptil_now, model=self.model)
+                        # Add mcp_list_tools to first chunk if present
+                        if not self.sent_first_chunk and processed_chunk.choices:
+                            processed_chunk = self._add_mcp_list_tools_to_first_chunk(processed_chunk)
+                            self.sent_first_chunk = True
+
                         if not self.send_stream_usage or getattr(processed_chunk, "usage", None) is None:
                             self.chunks.append(processed_chunk)
-                            return processed_chunk
-                        usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk)
-                        if usage_free_chunk is None:
-                            continue
-                        return usage_free_chunk
+                        else:
+                            usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk)
+                            if usage_free_chunk is None:
+                                continue
+                            processed_chunk = usage_free_chunk
+
+                        # add usage as hidden param
+                        if self.sent_last_chunk is True and self.stream_options is None:
+                            self._set_final_chunk_usage(processed_chunk)
+
+                        # Call post-call streaming deployment hook for final chunk
+                        if self.sent_last_chunk is True:
+                            processed_chunk = await self._call_post_streaming_deployment_hook(processed_chunk)
+                            # Add MCP metadata to final chunk if present (after hooks)
+                            processed_chunk = self._add_mcp_metadata_to_final_chunk(processed_chunk)
+
+                        return processed_chunk
         except (StopAsyncIteration, StopIteration):
             if self._should_fail_on_clean_eof():
                 await self.aclose()

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -9571,3 +9571,167 @@ def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details(
     assert usage.prompt_tokens_details.cached_tokens == 80
     assert usage.completion_tokens_details is not None
     assert usage.completion_tokens_details.reasoning_tokens == 5
+
+
+class _SyncListIterator:
+    """Sync iterator over pre-built chunks.
+
+    Drives ``__anext__``'s non-async-iterable branch (the ``asyncio.to_thread``
+    path taken by boto3 bedrock streams)."""
+
+    def __init__(self, chunks: list[dict[str, object]]):
+        self._chunks = list(chunks)
+        self._index = 0
+
+    def __iter__(self) -> "_SyncListIterator":
+        return self
+
+    def __next__(self) -> dict[str, object]:
+        if self._index >= len(self._chunks):
+            raise StopIteration
+        chunk = self._chunks[self._index]
+        self._index += 1
+        return chunk
+
+
+def _generic_text_chunk(text: str, is_finished: bool, finish_reason: str | None, usage: dict | None) -> dict:
+    chunk: dict[str, object] = {
+        "text": text,
+        "is_finished": is_finished,
+        "finish_reason": finish_reason,
+        "usage": usage,
+    }
+    return chunk
+
+
+@pytest.mark.asyncio
+async def test_async_noniterable_branch_sets_completion_start_time(logging_obj: Logging):
+    """The non-async-iterable branch of __anext__ (boto3 bedrock streams) must
+    set logging_obj.completion_start_time on the first processed chunk, exactly
+    like the sync path and the main async-iterable branch. Without this, TTFT
+    logging falls back to end_time and reports the whole stream duration."""
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=_SyncListIterator(
+            [_generic_text_chunk("hello", is_finished=False, finish_reason=None, usage=None)]
+        ),
+        model="gpt-4o",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+    )
+
+    out = await wrapper.__anext__()
+
+    assert out.choices[0].delta.content == "hello"
+    assert logging_obj.completion_start_time is not None, (
+        "non-async-iterable branch never set completion_start_time - TTFT logging degrades to full stream duration"
+    )
+
+
+@pytest.mark.asyncio
+async def test_async_noniterable_branch_sets_hidden_usage_on_final_chunk(logging_obj: Logging):
+    """A provider that puts usage on the finish_reason chunk (e.g. bedrock's
+    amazon-bedrock-invocationMetrics on message_stop) keeps top-level usage on
+    the chunk when include_usage was not requested (behavior pinned by
+    test_plain_iterator_stream_keeps_usage_on_its_chunk_when_usage_was_not_requested),
+    and must ALSO carry usage in _hidden_params via _set_final_chunk_usage,
+    like the main async path with stream_options=None."""
+    usage: Final = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=_SyncListIterator(
+            [
+                _generic_text_chunk("hello", is_finished=False, finish_reason=None, usage=None),
+                _generic_text_chunk("", is_finished=True, finish_reason="stop", usage=usage),
+            ]
+        ),
+        model="gpt-4o",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+    )
+
+    first = await wrapper.__anext__()
+    assert first.choices[0].delta.content == "hello"
+
+    final_chunk = await wrapper.__anext__()
+
+    assert final_chunk.choices[0].finish_reason == "stop"
+    assert final_chunk.usage.total_tokens == 15, (
+        "top-level usage must stay on the chunk when include_usage was not requested"
+    )
+    hidden_usage: Final = getattr(final_chunk, "_hidden_params", {}).get("usage")
+    assert hidden_usage is not None, "_set_final_chunk_usage never ran - hidden usage missing on final chunk"
+    assert hidden_usage.prompt_tokens == 10
+    assert hidden_usage.completion_tokens == 5
+
+
+@pytest.mark.asyncio
+async def test_async_noniterable_branch_include_usage_final_chunk_still_carries_usage(logging_obj: Logging):
+    """Regression guard: with stream_options.include_usage, the non-iterable
+    branch must still strip per-chunk usage and deliver the final usage chunk
+    at end-of-stream, matching the main async-iterable branch."""
+    usage: Final = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=_SyncListIterator(
+            [
+                _generic_text_chunk("hello", is_finished=False, finish_reason=None, usage=None),
+                _generic_text_chunk("", is_finished=True, finish_reason="stop", usage=usage),
+            ]
+        ),
+        model="gpt-4o",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+        stream_options={"include_usage": True},  # mutable-ok: constructor payload
+    )
+
+    first = await wrapper.__anext__()
+    assert first.choices[0].delta.content == "hello"
+
+    final_content_chunk = await wrapper.__anext__()
+    assert final_content_chunk.choices[0].finish_reason == "stop"
+    assert getattr(final_content_chunk, "usage", None) is None, "usage was not stripped from the finish_reason chunk"
+
+    usage_chunk = await wrapper.__anext__()
+    final_usage: Final = getattr(usage_chunk, "usage", None)
+    assert final_usage is not None
+    assert final_usage.prompt_tokens == 10
+    assert final_usage.completion_tokens == 5
+
+
+@pytest.mark.asyncio
+async def test_async_noniterable_branch_final_chunk_hook_and_mcp_metadata(logging_obj: Logging):
+    """The final chunk from the non-iterable branch must flow through
+    async_post_call_streaming_deployment_hook and gain MCP metadata, like the
+    main async-iterable branch."""
+    from litellm.types.utils import ModelResponseStream as MRS
+
+    recorded_chunks: list[MRS] = []
+
+    class _RecordingLogger(CustomLogger):
+        async def async_post_call_streaming_deployment_hook(self, request_data, response_chunk, call_type=None):
+            recorded_chunks.append(response_chunk)
+            return response_chunk
+
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=_SyncListIterator(
+            [
+                _generic_text_chunk("hello", is_finished=False, finish_reason=None, usage=None),
+                _generic_text_chunk("", is_finished=True, finish_reason="stop", usage=None),
+            ]
+        ),
+        model="gpt-4o",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+    )
+    wrapper._hidden_params["mcp_metadata"] = {"mcp_list_tools": [{"name": "get_weather"}]}
+
+    await wrapper.__anext__()
+
+    with patch.object(litellm, "callbacks", [_RecordingLogger()]):
+        final_chunk = await wrapper.__anext__()
+
+    assert final_chunk.choices[0].finish_reason == "stop"
+    provider_fields = final_chunk.choices[0].delta.provider_specific_fields
+    assert provider_fields is not None
+    assert provider_fields.get("mcp_list_tools") == [{"name": "get_weather"}], (
+        "MCP metadata was not added to the final chunk of the non-iterable branch"
+    )
+    assert len(recorded_chunks) == 1, "post-call streaming deployment hook did not run for the final chunk"