fix: stop after final preprocessing responses

Merge #6869

Fixes #4472

PiperOrigin-RevId: 974149791

Author: sylvesterkaczmarek

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

@@ -1304,12 +1304,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.get_function_responses() and 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

@@ -2991,6 +2991,100 @@ 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_preprocess_non_function_response_does_not_skip_llm_call():
+  """Non-function-response events in preprocessing must not skip the LLM call."""
+  mock_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role='model',
+                  parts=[types.Part.from_text(text='Analysis done.')],
+              ),
+              finish_reason='STOP',
+          )
+      ]
+  )
+  agent = Agent(
+      name='root_agent',
+      model=testing_utils.MockModel.create(responses=[mock_response]),
+  )
+  flow = BaseLlmFlowForTesting()
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test'
+  )
+  processing_file_event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=types.Content(
+          role='model',
+          parts=[
+              types.Part(text='Processing input file: `data.csv`'),
+              types.Part(
+                  executable_code=types.ExecutableCode(
+                      code='import pandas as pd', language='PYTHON'
+                  )
+              ),
+          ],
+      ),
+  )
+  assert processing_file_event.is_final_response()
+
+  async def mock_preprocess(_ctx, _request):
+    yield processing_file_event
+
+  with mock.patch.object(
+      flow, '_preprocess_async', side_effect=mock_preprocess
+  ):
+    events = [event async for event in flow.run_async(invocation_context)]
+
+  assert len(events) == 2
+  assert events[0] == processing_file_event
+  assert events[1].content.parts[0].text == 'Analysis done.'
+
+
 @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."""