diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index be81e1df46..c6e17b56d0 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -3026,12 +3026,16 @@ def _finalize_tool_call_response( for index, func_data in function_calls.items(): if func_data["id"]: args = "".join(func_data["args_parts"]) - if finish_reason == "length": - try: - _parse_tool_call_arguments(args) - except json.JSONDecodeError: - has_incomplete_tool_call_args = True - continue + # Validate regardless of finish_reason: a stream that aborts + # mid-tool-call (e.g. transport cut, provider incident) ends + # with no finish_reason at all, so finalization falls back to a + # hardcoded "tool_calls" below and would otherwise let partial + # arguments through to an uncaught JSONDecodeError downstream. + try: + _parse_tool_call_arguments(args) + except json.JSONDecodeError: + has_incomplete_tool_call_args = True + continue tool_calls.append( ChatCompletionMessageToolCall( type="function", diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 519a5aa64f..2db655a3cc 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -4706,6 +4706,53 @@ async def test_streaming_tool_call_truncated_by_max_tokens( assert "max_output_tokens" in error_response.error_message +@pytest.mark.asyncio +async def test_streaming_tool_call_aborted_mid_stream( + mock_completion, lite_llm_instance +): + """Tests that a stream aborting mid-tool-call (no finish_reason at all) + + yields a graceful error LlmResponse instead of raising JSONDecodeError. + """ + stream_chunks = [ + ModelResponseStream( + choices=[ + StreamingChoices( + finish_reason=None, + delta=Delta( + role="assistant", + tool_calls=[ + ChatCompletionDeltaToolCall( + type="function", + id="call_789", + function=Function( + name="test_function", + arguments='{"test_arg":', + ), + index=0, + ) + ], + ), + ) + ] + ), + ] + mock_completion.return_value = iter(stream_chunks) + + responses = [ + response + async for response in lite_llm_instance.generate_content_async( + LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True + ) + ] + + assert len(responses) == 1 + error_response = responses[0] + assert error_response.error_code == types.FinishReason.MAX_TOKENS + assert error_response.finish_reason == types.FinishReason.MAX_TOKENS + assert "truncated" in error_response.error_message + + @pytest.mark.asyncio async def test_streaming_tool_call_complete_with_length_finish_reason( mock_completion, lite_llm_instance