fix: preserve function call IDs across SSE streaming partial/final events#4611
fix: preserve function call IDs across SSE streaming partial/final events#4611stakeswky wants to merge 2 commits intogoogle:mainfrom
Conversation
…ents
In SSE streaming mode, _finalize_model_response_event creates a brand-new
Event object for every LlmResponse chunk (partial and final). Because each
new Event's function calls start with an empty .id, populate_client_function_call_id
generates a fresh adk-{uuid} every time. This means the partial event and the
final event for the same function call end up with different IDs, breaking
LongRunningFunctionTool / HITL workflows that match responses by ID.
Fix: add an optional function_call_ids dict parameter to
_finalize_model_response_event. The dict maps (function_name, index) to the
ID that was assigned the first time that function call was seen. Before
populate_client_function_call_id runs, any previously stored ID is restored
onto the function call so the guard 'if not function_call.id' keeps it.
After population, newly generated IDs are written back into the dict.
_run_one_step_async creates one such dict per LLM call and threads it
through _postprocess_async for the lifetime of the streaming sequence, so
all partial and final events share the same stable IDs.
Fixes google#4609
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
Summary of ChangesHello @stakeswky, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a critical bug where function call IDs were not consistently maintained between partial and final events in SSE streaming. Previously, each streaming chunk generated a new Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
adk-bot
added
the
live
|
Response from ADK Triaging Agent Hello @stakeswky, thank you for your contribution! To help reviewers evaluate your PR, could you please add a Also, it looks like the Contributor License Agreement (CLA) check is failing. Please make sure you have signed the CLA to proceed. Thank you! |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request addresses an issue where function call IDs were not preserved across partial and final events during SSE streaming, which is a critical fix for workflows relying on stable IDs. The solution, which involves passing a dictionary to track and persist these IDs throughout the streaming sequence, is well-implemented and effectively solves the problem. The addition of comprehensive regression tests is also a great inclusion, ensuring the fix is robust and prevents future regressions. I have one minor suggestion to refactor a small part of the logic for improved conciseness using a more idiomatic Python dictionary method.
| if function_call_ids is not None: | ||
| for i, fc in enumerate(function_calls): | ||
| key = (fc.name, i) | ||
| if fc.id and key not in function_call_ids: | ||
| function_call_ids[key] = fc.id |
There was a problem hiding this comment.
For conciseness, you can use dict.setdefault() to persist the newly generated IDs. This method is more idiomatic for this pattern and achieves the same result as the current if condition, but in a more compact way.
| if function_call_ids is not None: | |
| for i, fc in enumerate(function_calls): | |
| key = (fc.name, i) | |
| if fc.id and key not in function_call_ids: | |
| function_call_ids[key] = fc.id | |
| if function_call_ids is not None: | |
| for i, fc in enumerate(function_calls): | |
| key = (fc.name, i) | |
| if fc.id: | |
| function_call_ids.setdefault(key, fc.id) |
There was a problem hiding this comment.
Good suggestion! Applied in cae3567 — switched to dict.setdefault() for cleaner ID persistence.
|
The PR description already includes a Testing Plan section with 4 unit tests and the test command. I'll sign the CLA to unblock the review. |
|
I have signed the Google CLA. Please re-check. |
|
@googlebot I signed it! |
stakeswky
force-pushed
the
fix/streaming-function-call-id-mismatch
branch
from
bc6aed5 to
94ba594
Compare
|
Closing in favor of #4619 which has a cleaner implementation. |
2 participants
Summary
Fixes #4609
When SSE streaming is enabled (default since ADK 1.22+), _finalize_model_response_event creates a fresh Event object for the final (non-partial) response. This caused populate_client_function_call_id to generate a new UUID for the same logical function call, breaking HITL workflows with LongRunningFunctionTool.
Root Cause
The
if not function_call.idguard inpopulate_client_function_call_idonly prevents re-assignment on the same Event object. It does not prevent assigning a different ID to the same logical function call across different Event objects (partial vs final).Fix
Introduces a
function_call_idsdict that tracks assigned IDs across partial and final events in the same streaming sequence:populate_client_function_call_idruns, newly generated IDs are persisted back into the dict for reuseThe dict is keyed by
(function_name, index)to handle multi-function-call responses correctly.Changes
base_llm_flow.py: Addedfunction_call_idsparameter threading through_finalize_model_response_event,_postprocess_after_model_call, and the streaming looptest_streaming_function_call_id.py: 4 regression testsTesting Plan
Unit tests: 4 new tests in
test_streaming_function_call_id.py:test_partial_and_final_share_same_id: Core regression testtest_without_function_call_ids_dict_generates_different_ids: Demonstrates old buggy behaviortest_multiple_function_calls_preserve_ids: Multi-call response stabilitytest_server_provided_id_is_preserved: Server IDs not overwrittenRegression testing: All 352 existing tests in
tests/unittests/flows/llm_flows/passTest command:
python -m pytest tests/unittests/flows/llm_flows/ -v— 356 passed (352 existing + 4 new)