fix(flows): stop after final preprocessing responses

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -1158,12 +1158,15 @@ async def _run_one_step_async(
     run_config = _require_run_config(invocation_context)
 
     # Preprocess before calling the LLM.
+    preprocess_yielded_final_response = False
     async with Aclosing(
         self._preprocess_async(invocation_context, llm_request)
     ) as agen:
       async for event in agen:
+        if event.is_final_response():
+          preprocess_yielded_final_response = True
         yield event
-    if invocation_context.end_invocation:
+    if invocation_context.end_invocation or preprocess_yielded_final_response:
       return
 
     # Resume the LLM agent based on the last event from the current branch.

tests/unittests/flows/llm_flows/test_base_llm_flow.py

@@ -2902,6 +2902,48 @@ async def _drive_one_llm_call(flow, invocation_context):
       pass
 
 
+@pytest.mark.asyncio
+async def test_preprocess_final_response_skips_llm_call():
+  """A final response from preprocessing must finish the current step."""
+  agent = Agent(
+      name='root_agent', model=testing_utils.MockModel.create(responses=[])
+  )
+  flow = BaseLlmFlowForTesting()
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='resume'
+  )
+  function_response_event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=types.Content(
+          role='user',
+          parts=[
+              types.Part.from_function_response(
+                  name='resumed_tool', response={'result': 'done'}
+              )
+          ],
+      ),
+  )
+  function_response_event.actions.skip_summarization = True
+
+  async def mock_preprocess(_ctx, _request):
+    yield function_response_event
+
+  async def fail_if_llm_called(*_args, **_kwargs):
+    raise AssertionError('LLM should not be called after a final response')
+    yield  # pylint: disable=unreachable
+
+  with (
+      mock.patch.object(flow, '_preprocess_async', side_effect=mock_preprocess),
+      mock.patch.object(
+          flow, '_call_llm_async', side_effect=fail_if_llm_called
+      ),
+  ):
+    events = [event async for event in flow.run_async(invocation_context)]
+
+  assert events == [function_response_event]
+
+
 @pytest.mark.asyncio
 async def test_cfc_llm_calls_are_counted_against_max_llm_calls():
   """support_cfc must not exempt a run from the max_llm_calls spend cap."""