diff --git a/src/google/adk/flows/llm_flows/tools/_rearranger.py b/src/google/adk/flows/llm_flows/tools/_rearranger.py index 5e3e1ff209..17732301e3 100644 --- a/src/google/adk/flows/llm_flows/tools/_rearranger.py +++ b/src/google/adk/flows/llm_flows/tools/_rearranger.py @@ -350,6 +350,7 @@ def rearrange_events_for_latest_function_response( event = events[idx] function_calls = event.get_function_calls() if function_calls: + matched = False for function_call in function_calls: if function_call.id in function_responses_ids: function_call_event_idx = idx @@ -368,7 +369,14 @@ def rearrange_events_for_latest_function_response( # collect all function responses from the function call event to # the last response event function_responses_ids = function_call_ids + matched = True break + # A call id can be reused by a later call in the same session (some + # model providers do this). Stop at the nearest preceding match rather + # than continuing to walk past it, or a reused id attributes the + # response to a stale, already-answered call further back in history. + if matched: + break if function_call_event_idx == -1: logger.debug( diff --git a/tests/unittests/flows/llm_flows/tools/test_rearranger.py b/tests/unittests/flows/llm_flows/tools/test_rearranger.py index d9399b3427..104b661800 100644 --- a/tests/unittests/flows/llm_flows/tools/test_rearranger.py +++ b/tests/unittests/flows/llm_flows/tools/test_rearranger.py @@ -382,6 +382,90 @@ def test_rearrange_latest_response_missing_matching_call_raises_value_error(): rearrange_events_for_latest_function_response(events) +def test_rearrange_latest_response_reused_call_id_pairs_with_nearest_call(): + """A reused call id pairs the latest response with the nearest preceding call. + + Some model providers reuse function-call ids across turns. The latest + response should attribute to the closest matching call, not the oldest + one, so earlier, already-answered turns are left untouched. + """ + call1 = _call_event("call_1", "lookup") + resp1 = _resp_event("call_1", "lookup", "looked up") + intervening_msg = Event(author="user", content=types.UserContent("q2")) + call2 = _call_event("call_1", "update") + placeholder = _resp_event("call_1", "update", {"placeholder": True}) + final = _resp_event("call_1", "update", {"applied": True}) + events = [call1, resp1, intervening_msg, call2, placeholder, final] + + result = rearrange_events_for_latest_function_response(events) + + assert result[:4] == [call1, resp1, intervening_msg, call2] + assert len(result) == 5 + merged_responses = result[-1].get_function_responses() + assert len(merged_responses) == 1 + assert merged_responses[0].response == {"applied": True} + + +def test_rearrange_latest_response_reused_id_in_parallel_batch_pairs_with_nearest_call(): + """A reused id beside another call in a batch pairs with the nearest call. + + Without the fix, the search walks past the nearest matching call event to + an older call event that does not carry every id in the batch, and the + subset check raises instead of truncating. + """ + call2 = Event( + author="test_agent", + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id="call_1", name="update", args={} + ) + ), + types.Part( + function_call=types.FunctionCall( + id="call_2", name="list", args={} + ) + ), + ], + ), + ) + paused = Event( + author="user", + content=types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id="call_1", name="update", response={"placeholder": True} + ) + ), + types.Part( + function_response=types.FunctionResponse( + id="call_2", name="list", response={"rows": 3} + ) + ), + ], + ), + ) + events = [ + _call_event("call_1", "lookup"), + _resp_event("call_1", "lookup", "looked up"), + call2, + paused, + _resp_event("call_1", "update", {"applied": True}), + ] + + result = rearrange_events_for_latest_function_response(events) + + assert result[:3] == events[:3] + assert [(r.id, r.response) for r in result[-1].get_function_responses()] == [ + ("call_1", {"applied": True}), + ("call_2", {"rows": 3}), + ] + + def test_rearrange_history_reused_id_across_tools_pairs_correctly(): """Reused call IDs across different tools pair each tool with its own response.""" events = [