fix(litellm): keep every tool call in a finish-bearing stream chunk

[elbourne12345]

fix(litellm): keep every tool call in a finish-bearing stream chunk

• Closes: LiteLlm streaming keeps only the last tool call when one chunk carries N calls plus a finish_reason (latent) #7004
• Related: LiteLLM streaming silently drops tool call responses when finish_reason is "length" (max output tokens reached) #4482, LiteLlm generate_content_async doesn't handle multiple tool calls in the same model response in SSE mode #484, LiteLLM parallel tool calls does not have correct args in StreamingMode.SSE #1038

Problem:

_model_response_to_chunk yields one FunctionChunk per tool call in a delta, each paired
with the same choice-level finish_reason, and the tool-call finalizer ran inside that
per-chunk loop:

for chunk, finish_reason in _model_response_to_chunk(part):
    ...
    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 a delta carrying N complete tool calls plus finish_reason="tool_calls" (or "length")
rebuilt the aggregated response N times, each rebuild replacing the previous one after
_reset_stream_buffers() had cleared the accumulator. Only the last call survived, silently,
in a well-formed response — the end-of-stream fallback cannot rescue the earlier ones because
the buffers are empty and the aggregate is already set. Any content or reasoning_content
in the same delta was folded into the first, overwritten response and lost with it.

Scope, stated plainly: this shape does not arrive through litellm.acompletion. Its
CustomStreamWrapper pops finish_reason off every non-empty chunk
(litellm_core_utils/streaming_handler.py) and re-emits it on a trailing empty-delta chunk,
for fake-streamed and natively-streamed providers alike — I verified that end to end for
openai, azure, bedrock, vertex_ai, gemini, ollama_chat, anthropic, groq, together_ai,
hosted_vllm and custom providers. The shape does arrive via a custom LiteLlm.llm_client
(a public field), and ADK's own unit tests drive this code with raw ModelResponseStream
chunks, i.e. through exactly the vulnerable path. So this is a latent robustness fix plus the
missing coverage, not a fix for a live provider regression. It is worth doing because the
behaviour contradicts BaseLlm.generate_content_async's documented contract that the final
partial=False chunk equals the stream=False output, and because it fails silently.

The placement predates multi-call support: the check has been inside the inner loop since the
initial commit 982782014 (2025-04-08), when the aggregator tracked a single function_id,
and was inherited unchanged by 05f48347 (#759, index-keyed dict), e8019b1b (#4225, the
stop-only guard), 4c6096baa (#4482, the "length" arm), 36fd2c8e and eaed0aa8. No commit
in that history considers more than one tool call per chunk.

Solution:

Decide the finalization once per streamed part. The per-chunk loop now records what the part
asks for; the finalization happens after the loop:

finalize_tool_calls = False
finalize_text = False
finalize_finish_reason: str | None = None
for chunk, finish_reason in _model_response_to_chunk(part):
    ...
    if function_calls and (...same conditions...):
        finalize_tool_calls = True
        finalize_finish_reason = finish_reason
    elif (text_parts or reasoning_parts) and (...same conditions...):
        finalize_text = True
        finalize_finish_reason = finish_reason

if finalize_tool_calls:
    aggregated_llm_response_with_tool_call = _finalize_tool_call_response(...)
    _reset_stream_buffers()
elif finalize_text:
    aggregated_llm_response = _finalize_text_response(...)
    _reset_stream_buffers()

The trigger conditions are copied unchanged, so the chunk is None guard that protects
against LiteLLM 1.81+ setting finish_reason="stop" on partial chunks still applies exactly
as before. Tool calls take precedence over text for a part that asks for both, which loses
nothing: _finalize_tool_call_response already carries the buffered text and reasoning into
the tool-call response.

I deliberately did not implement this as "merge instead of replace". The "length" arm
routes through _parse_tool_call_arguments and can produce an error LlmResponse, so merging
would need error/normal reconciliation for no benefit over deciding once.

Testing Plan

Unit Tests:

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

Added _same_chunk_multiple_function_calls_stream(...) plus three cases in
tests/unittests/models/test_litellm.py:

• test_generate_content_async_same_chunk_multiple_function_calls[tool_calls]
• test_generate_content_async_same_chunk_multiple_function_calls[length]
• test_generate_content_async_same_chunk_calls_keep_text_and_reasoning

No existing fixture has more than one tool call per delta, or a tool-call delta sharing a chunk
with a finish reason, which is why MULTIPLE_FUNCTION_CALLS_STREAM (calls in separate chunks,
finish reason on its own empty chunk) passes both before and after.

All three fail on main and pass with the change. Reverting only
src/google/adk/models/lite_llm.py and rerunning the new tests:

FAILED tests/unittests/models/test_litellm.py::test_generate_content_async_same_chunk_multiple_function_calls[tool_calls]
FAILED tests/unittests/models/test_litellm.py::test_generate_content_async_same_chunk_multiple_function_calls[length]
FAILED tests/unittests/models/test_litellm.py::test_generate_content_async_same_chunk_calls_keep_text_and_reasoning
E     AssertionError: assert ['function_2'] == ['function_1', 'function_2']
3 failed, 417 deselected in 7.16s

With the change applied:

$ pytest tests/unittests/models/test_litellm.py -q
420 passed, 1 warning in 7.19s

$ pytest tests/unittests/models/ -q
1235 passed, 35 warnings in 138.58s

pre-commit run --files src/google/adk/models/lite_llm.py tests/unittests/models/test_litellm.py
passes (ruff, isort, pyink, addlicense, codespell, end-of-file, trailing-whitespace); the two
repo-local hooks were run directly because they shell out to /bin/bash and a relative script
path that don't resolve on Windows — scripts/compliance_checks.py and
scripts/check_new_py_files.py both exit 0, and the change adds no new Python files.

Manual End-to-End (E2E) Tests:

Self-contained, no network and no credentials. Feeds LiteLlm real
litellm.types.utils.ModelResponseStream chunks through a custom llm_client across six
stream shapes:

import asyncio
from google.adk.models.lite_llm import LiteLlm, LiteLLMClient
from google.adk.models.llm_request import LlmRequest
from google.genai import types
from litellm.types.utils import (
    ChatCompletionDeltaToolCall, Delta, Function, ModelResponseStream, StreamingChoices,
)

def tc(id_, name, args, index):
    return ChatCompletionDeltaToolCall(type="function", id=id_, function=Function(name=name, arguments=args), index=index)

def chunk(tool_calls, finish_reason):
    return ModelResponseStream(model="openai/gpt-4o", choices=[
        StreamingChoices(finish_reason=finish_reason, delta=Delta(role="assistant", tool_calls=tool_calls or None))])

F1 = ("call_1", "f1", '{"a": 1}', 0)
F2 = ("call_2", "f2", '{"b": 2}', 1)
F3 = ("call_3", "f3", '{"c": 3}', 2)
CASES = {
    "A_two_calls_one_chunk_finish_tool_calls":    [chunk([tc(*F1), tc(*F2)], "tool_calls")],
    "A3_three_calls_one_chunk_finish_tool_calls": [chunk([tc(*F1), tc(*F2), tc(*F3)], "tool_calls")],
    "L_two_calls_one_chunk_finish_length":        [chunk([tc(*F1), tc(*F2)], "length")],
    "B_separate_chunks_then_empty_finish":        [chunk([tc(*F1)], None), chunk([tc(*F2)], None), chunk([], "tool_calls")],
    "C_two_calls_one_chunk_finish_stop":          [chunk([tc(*F1), tc(*F2)], "stop")],
    "D_two_calls_one_chunk_then_empty_finish":    [chunk([tc(*F1), tc(*F2)], None), chunk([], "tool_calls")],
}

class FakeClient(LiteLLMClient):
    def __init__(self, chunks):
        self._chunks = chunks
    async def acompletion(self, model, messages, tools, **kwargs):
        async def gen():
            for c in self._chunks:
                yield c
        return gen()
    def completion(self, *a, **k):
        raise NotImplementedError

REQ = LlmRequest(
    contents=[types.Content(role="user", parts=[types.Part.from_text(text="go")])],
    config=types.GenerateContentConfig(tools=[types.Tool(function_declarations=[
        types.FunctionDeclaration(name=n, description=n, parameters=types.Schema(
            type=types.Type.OBJECT, properties={k: types.Schema(type=types.Type.INTEGER)}))
        for n, k in (("f1", "a"), ("f2", "b"), ("f3", "c"))])]),
)

async def run(name, chunks):
    expected = ["f1", "f2", "f3"] if "three" in name else ["f1", "f2"]
    llm = LiteLlm(model="openai/gpt-4o", llm_client=FakeClient(chunks))
    finals = [r async for r in llm.generate_content_async(REQ, stream=True) if not r.partial]
    survivors = [p.function_call.name for r in finals for p in (r.content.parts if r.content else []) if p.function_call]
    print(f"{name:44s} -> survivors={survivors}  {'OK' if sorted(survivors) == expected else 'BUG'}")

async def main():
    for name, chunks in CASES.items():
        await run(name, chunks)

asyncio.run(main())

Before (google-adk 2.8.0 and main @ d637d1b):

A_two_calls_one_chunk_finish_tool_calls      -> survivors=['f2']  BUG
A3_three_calls_one_chunk_finish_tool_calls   -> survivors=['f3']  BUG
L_two_calls_one_chunk_finish_length          -> survivors=['f2']  BUG
B_separate_chunks_then_empty_finish          -> survivors=['f1', 'f2']  OK
C_two_calls_one_chunk_finish_stop            -> survivors=['f1', 'f2']  OK
D_two_calls_one_chunk_then_empty_finish      -> survivors=['f1', 'f2']  OK

After this change:

A_two_calls_one_chunk_finish_tool_calls      -> survivors=['f1', 'f2']  OK
A3_three_calls_one_chunk_finish_tool_calls   -> survivors=['f1', 'f2', 'f3']  OK
L_two_calls_one_chunk_finish_length          -> survivors=['f1', 'f2']  OK
B_separate_chunks_then_empty_finish          -> survivors=['f1', 'f2']  OK
C_two_calls_one_chunk_finish_stop            -> survivors=['f1', 'f2']  OK
D_two_calls_one_chunk_then_empty_finish      -> survivors=['f1', 'f2']  OK

The three control shapes are unchanged, which is the point: the fix only affects a delta that
carries tool calls and a finish reason together.

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. (N/A — no dependent changes.)

Additional context

Related but distinct, so this is not a duplicate: #4482 (closed 2026-03-10) reported a tool
call being dropped entirely when finish_reason == "length", because that value was missing
from the yield condition; its fix added the "length" arm to the same if this PR restructures.
That fixed "nothing is yielded"; this fixes "only the last of N is yielded". #484 / #1038 (fixed
by PR #759, which created this aggregation loop and the index-keyed dict) cover the
separate-chunk shape, where the finish reason arrives on its own empty chunk — the shape the
existing tests exercise and which was already correct.

Not addressed here, to keep this to one concern: text arriving after a mid-stream text
finalization is single-slot-overwritten in the same way, but
test_streaming_text_buffer_is_reset_between_aggregated_responses (from 36fd2c8e) pins the
current last-segment-wins behaviour there, and moving the tool-call check does not affect it.
Happy to file that separately if it is worth changing.