fix(models): keep every parallel tool call streamed with the finish reason

[AtulJoshi1206]

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

No existing issue; described below following the bug-report structure.

Describe the Bug:

When a provider streams several parallel tool calls on the chunk that also carries finish_reason, LiteLlm.generate_content_async keeps only the last one. The others are dropped silently: no error, no log, no partial response.

_model_response_to_chunk yields one FunctionChunk per tool call in a chunk and stamps that chunk's finish_reason on every one of those yields. The finalizer sits inside that same per-chunk loop:

if function_calls and (
    finish_reason == "tool_calls"
    or finish_reason == "length"
    or (finish_reason == "stop" and chunk is None)
):
  aggregated_llm_response_with_tool_call = _finalize_tool_call_response(...)
  _reset_stream_buffers()

So on a chunk carrying N calls, the first call satisfies the condition, the segment is finalized, and _reset_stream_buffers() clears function_calls. The second call then enters an empty accumulator and finalizes again, overwriting aggregated_llm_response_with_tool_call. Only the final call survives.

Impact:

The agent executes one of N tools and the model is never told the rest were requested, so a parallel-tool turn silently does part of its work. The non-streaming path on the same payload returns every call, so stream=True and stream=False disagree.

This is the normal shape on LiteLLM's Gemini/Vertex route rather than an exotic one. In the installed litellm, VertexGeminiConfig._check_finish_reason returns "tool_calls" for any message carrying tool calls, _create_streaming_choice stamps it on the very chunk that carries them, and _apply_stream_candidates rewrites "stop" to "tool_calls" on tool-call chunks. Gemini emits all functionCall parts of a candidate in a single chunk.

Steps to Reproduce:

On current main (665ec9835), stream one ModelResponseStream with finish_reason="tool_calls" and two tool calls:

ModelResponseStream(
    model="test_model",
    choices=[StreamingChoices(
        finish_reason="tool_calls",
        delta=Delta(role="assistant", tool_calls=[
            ChatCompletionDeltaToolCall(type="function", id="call_1", index=0,
                function=Function(name="get_weather", arguments='{"city":"SF"}')),
            ChatCompletionDeltaToolCall(type="function", id="call_2", index=1,
                function=Function(name="get_time", arguments='{"tz":"UTC"}')),
        ]),
    )],
)

Observed Behavior:

assert [('get_time', {'tz': 'UTC'})] == [('get_weather', {'city': 'SF'}), ('get_time', {'tz': 'UTC'})]
  At index 0 diff: ('get_time', {'tz': 'UTC'}) != ('get_weather', {'city': 'SF'})
  Right contains one more item: ('get_time', {'tz': 'UTC'})

get_weather is gone.

Expected Behavior:

Both calls present on the final response, matching what the non-streaming path returns for the same payload.

Solution:

A finish reason only ends the tool-call segment when it arrives on a chunk that carries no further tool-call delta, that is when chunk is None. Every other terminal reason falls through to the end-of-stream finalizer that already exists below, which replays last_finish_reason, so the reported finish_reason is unchanged.

Two details worth flagging for review:

• The simpler-looking alternative, never finalizing in the loop, breaks test_generate_content_async_stream_reason_does_not_carry_over. The chunk is None guard is what preserves that test's intent.
• Narrowing the tool-call branch means a "length" reason arriving on a tool-call chunk no longer lands there, so the text elif below could have claimed it and wiped the buffered calls via _reset_stream_buffers(). That branch therefore gains an explicit not function_calls, which keeps the original precedence of tool calls over buffered text.

Testing Plan

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

Added test_generate_content_async_stream_keeps_parallel_tool_calls, which asserts both calls and their arguments survive. Verified it fails on an unmodified lite_llm.py and passes with the change.

$ pytest tests/unittests/models -q
1303 passed

$ pytest tests/unittests -q
15226 passed, 86 skipped, 27 xfailed, 2 xpassed in 342.58s (0:05:42)
# 1 unrelated failure, test_import_loading.py::...[agent], is an artifact of the
# git-worktree checkout I built this on; it fails there on an unmodified main
# and passes in a normal clone.

$ pre-commit run --files src/google/adk/models/lite_llm.py tests/unittests/models/test_litellm.py
# all hooks pass

Manual End-to-End (E2E) Tests:

Driven through LiteLlm.generate_content_async with a stubbed acompletion, since a live parallel-tool turn needs provider credentials. Three shapes were checked, all failing before and passing after:

stream shape	before	after
both calls on one chunk carrying finish_reason	last call only	both calls
calls on separate chunks, each carrying the reason	last call only	both calls
text, then both calls on one chunk	last call only	text plus both calls

Checklist

• I have read the CONTRIBUTING.md document.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules. (none)

Additional context

Streaming branch only. Partial yields, usage and grounding attachment, the MAX_TOKENS and MALFORMED_FUNCTION_CALL branches, and the text path are untouched. A single-call stream produces an identical response, finalized one loop iteration later by the end-of-stream finalizer.

🤖 Generated with Claude Code