diff --git a/src/google/adk/flows/llm_flows/core/_resume.py b/src/google/adk/flows/llm_flows/core/_resume.py index ffcb9cb78c..9723bffe07 100644 --- a/src/google/adk/flows/llm_flows/core/_resume.py +++ b/src/google/adk/flows/llm_flows/core/_resume.py @@ -110,9 +110,8 @@ def _pause_left_calls_unanswered( fr.id for ev in events for fr in ev.get_function_responses() if fr.id } # `issubset`, not `&`: this asks whether *any* awaited id is still open, so a - # partially answered pause keeps waiting. `decide_resume` asks the opposite - # question of its own ids -- whether *none* are answered -- and drops - # `issubset` for that reason. The two are not interchangeable. + # partially answered pause keeps waiting. Intersection would treat one + # sibling answer as coverage for the rest. return bool(awaited) and not awaited.issubset(answered) @@ -191,20 +190,28 @@ def _needs_call_replay( call_names: set[str | None], answers: list[types.FunctionResponse], from_sub_branch: bool, + call_ids: set[str | None] | None = None, + answered_ids: set[str] | None = None, ) -> bool: """Whether the calls named by `call_names` still have to be run. - `call_names` holds every name on the call event, not just the first: one - event can carry parallel calls, and an answer to the second is not evidence - the first never ran. + Coverage is all-or-nothing: an answer to one parallel call is not + evidence the others ran. Names and ids both have to be covered -- + two calls can share a name, so names alone cannot see a missing twin. + `answers` is every function response after the call, not only the last + matching event: parallel answers often arrive separately. """ if not call_names: return False - return ( - not answers - or any(fr.name not in call_names for fr in answers) - or from_sub_branch - ) + answered_names = {fr.name for fr in answers} + names_uncovered = not call_names.issubset(answered_names) + ids_uncovered = False + if call_ids is not None and answered_ids is not None: + concrete_ids = {i for i in call_ids if i is not None} + ids_uncovered = bool(concrete_ids) and not concrete_ids.issubset( + answered_ids + ) + return names_uncovered or ids_uncovered or from_sub_branch def decide_resume( @@ -257,18 +264,32 @@ def decide_resume( # short-circuits both unanswered tests rather than being repeated in each. from_sub_branch = _is_sub_branch_answer(answer_event, call_event) answers = answer_event.get_function_responses() - # `ids & answered` alone decides these: a set that is a subset of the - # answered ids necessarily intersects it, so testing `issubset` as well - # never changes the outcome. - lro_unanswered = bool(lro_ids) and not lro_ids & answered_ids + # Coverage looks at every response after the call, not only the last + # matching event: parallel answers often arrive as separate events. + all_answers = [ + fr + for ev in events[call_idx + 1 :] + for fr in ev.get_function_responses() + ] + concrete_call_ids = {i for i in call_ids if i is not None} + # `issubset`, not `&`: one answered id does not cover a sibling that + # never ran. Pause only when nothing matched (no id and no name); + # leftover ids after a sibling answer are a replay, not a pause. + lro_unanswered = bool(lro_ids) and not lro_ids.issubset(answered_ids) call_unanswered = ( - bool(call_ids) - and not call_ids & answered_ids + bool(concrete_call_ids) + and not concrete_call_ids.issubset(answered_ids) and not any(fr.name in call_names for fr in answers) ) if not from_sub_branch and (lro_unanswered or call_unanswered): pause = True - elif _needs_call_replay(call_names, answers, from_sub_branch): + elif _needs_call_replay( + call_names, + all_answers, + from_sub_branch, + call_ids=concrete_call_ids, + answered_ids=answered_ids, + ): return ResumeDecision(ResumeAction.REPLAY_CALLS, call_event) return ResumeDecision(ResumeAction.PAUSE if pause else ResumeAction.CONTINUE) diff --git a/tests/unittests/flows/llm_flows/core/test_resume.py b/tests/unittests/flows/llm_flows/core/test_resume.py index afebde3132..fec288c96a 100644 --- a/tests/unittests/flows/llm_flows/core/test_resume.py +++ b/tests/unittests/flows/llm_flows/core/test_resume.py @@ -52,6 +52,22 @@ def _call_event(name: str, call_id: str, *, lro: bool = False) -> Event: ) +def _parallel_call_event(pairs: list[tuple[str, str]]) -> Event: + return Event( + author='agent', + invocation_id='inv-1', + content=types.Content( + role='model', + parts=[ + types.Part( + function_call=types.FunctionCall(id=i, name=n, args={}) + ) + for i, n in pairs + ], + ), + ) + + def _response_event( name: str, response_id: str | None, @@ -292,25 +308,7 @@ def test_parallel_calls_all_answered_continue(self): # One event can carry parallel calls. Matching answers against only the # first call's name reads the second answer as a foreign name, so a fully # answered event is replayed and both tools run a second time. - call = Event( - author='agent', - invocation_id='inv-1', - content=types.Content( - role='model', - parts=[ - types.Part( - function_call=types.FunctionCall( - id='c1', name='ask', args={} - ) - ), - types.Part( - function_call=types.FunctionCall( - id='c2', name='fetch', args={} - ) - ), - ], - ), - ) + call = _parallel_call_event([('c1', 'ask'), ('c2', 'fetch')]) events = [ call, _response_event('ask', 'c1'), @@ -321,6 +319,29 @@ def test_parallel_calls_all_answered_continue(self): ) assert decision.action is ResumeAction.CONTINUE + def test_parallel_calls_partially_answered_replay(self): + # A sibling answer is not coverage for a call that never ran. Any-answered + # (name match or id intersection) would CONTINUE here and drop fetch. + call = _parallel_call_event([('c1', 'ask'), ('c2', 'fetch')]) + events = [call, _response_event('ask', 'c1')] + + decision = decide_resume( + self._ctx(), events, {'ask': object(), 'fetch': object()} + ) + + assert decision.action is ResumeAction.REPLAY_CALLS + assert decision.event is call + + def test_parallel_same_name_partially_answered_replay(self): + # Two calls can share a name, so names alone cannot see the missing twin. + call = _parallel_call_event([('c1', 'ask'), ('c2', 'ask')]) + events = [call, _response_event('ask', 'c1')] + + decision = decide_resume(self._ctx(), events, {'ask': object()}) + + assert decision.action is ResumeAction.REPLAY_CALLS + assert decision.event is call + def test_sub_branch_answer_replays_instead_of_pausing(self): # A HITL answer returned against the branch the call opened resolves it, # even though it carries none of the call's ids. @@ -434,3 +455,15 @@ def test_the_agents_own_trailing_call_is_still_replayed(self): ) assert decision.action is ResumeAction.REPLAY_CALLS assert decision.replay_event() is call + + def test_a_partially_answered_parallel_step_is_replayed(self): + # The entry point must not CONTINUE just because a sibling was answered. + call = _parallel_call_event([('c1', 'ask'), ('c2', 'fetch')]) + events = [call, _response_event('ask', 'c1')] + + decision = decide_step_resume( + self._ctx(events), {'ask': object(), 'fetch': object()} + ) + + assert decision.action is ResumeAction.REPLAY_CALLS + assert decision.replay_event() is call