Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions litellm/proxy/guardrails/guardrail_hooks/tool_permission.py
Original file line number Diff line number Diff line change
Expand Up @@ -799,6 +799,11 @@ async def async_post_call_streaming_iterator_hook(
verbose_proxy_logger.debug(
"Tool Permission Guardrail: No tool uses found"
)
mock_response = MockResponseIterator(
model_response=assembled_model_response
)
async for chunk in mock_response:
yield chunk
return

verbose_proxy_logger.debug(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
ChatCompletionMessageToolCall,
Choices,
ModelResponse,
ModelResponseStream,
)


Expand Down Expand Up @@ -676,6 +677,49 @@ async def test_async_pre_call_hook_rewrite_mode_filters_legacy_functions(self):
assert new_data["function_call"] == "none"
assert new_data["tool_choice"] == "none"

@pytest.mark.asyncio
async def test_async_post_call_streaming_iterator_hook_plain_text_yields_chunks(
self,
):
"""Regression test: hook must re-emit chunks when LLM replies with plain text.

Before the fix, the `if not tool_calls:` branch did a bare `return` inside
the async generator, which yielded nothing. Clients received only
`data: [DONE]` with no content.
"""
text_chunk = ModelResponseStream(
id="chatcmpl-plain-text",
created=1700000000,
model="gpt-4",
object="chat.completion.chunk",
choices=[],
)

async def _fake_stream():
yield text_chunk

assembled = ModelResponse(
choices=[Choices(message={"content": "Hello, world!"})]
)

with patch("litellm.main.stream_chunk_builder", return_value=assembled):
chunks = []
async for chunk in self.guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(),
response=_fake_stream(),
request_data={},
):
chunks.append(chunk)

assert len(chunks) >= 1, (
"Hook must yield at least one chunk for plain-text responses; "
"got none — bare return bug"
)
Comment on lines +714 to +717

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Loose regression assertion

len(chunks) >= 1 only proves that something was yielded; it doesn't verify that the yielded chunk carries the expected content ("Hello, world!"). A stronger assertion would also confirm that the content attribute in at least one chunk matches the assembled response, making the test a true guard against silent content corruption in addition to the dropped-stream bug.

Suggested change
assert len(chunks) >= 1, (
"Hook must yield at least one chunk for plain-text responses; "
"got none — bare return bug"
)
assert len(chunks) >= 1, (
"Hook must yield at least one chunk for plain-text responses; "
"got none — bare return bug"
)
# Verify the content of the yielded chunks matches the assembled response.
content = "".join(
getattr(c.choices[0].delta, "content", "") or ""
for c in chunks
if c.choices
)
assert "Hello, world!" in content, (
f"Expected plain-text content to be re-emitted; got: {content!r}"
)

assert chunks[0].choices[0].delta.content == "Hello, world!", (
"Hook must preserve the original response content; "
f"got: {chunks[0].choices[0].delta.content!r}"
)

def test_modify_response_with_permission_errors(self):
# Setup a response with one tool_call
tool_call = ChatCompletionMessageToolCall(
Expand Down
Loading