fix: recover compacted function calls during prompt assembly

Close #5602

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 944697651

Author: GWeale

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

@@ -539,6 +539,96 @@ def _is_timestamp_compacted(ts: float) -> bool:
   return [event for _, _, event in processed_items]
 
 
+def _recover_compacted_function_calls(
+    events: list[Event],
+    source_events: list[Event],
+) -> list[Event]:
+  """Re-injects function-call events that compaction removed.
+
+  Compaction can summarize away a function_call while a matching
+  function_response survives outside the compacted range. The clearest case
+  is a long-running tool call: the call is compacted along with its
+  intermediate placeholder response, then the real result arrives on resume
+  (a later event not covered by the summary). That surviving response would
+  be orphaned, which breaks call/response pairing during prompt assembly (it
+  raises in `_rearrange_events_for_latest_function_response`).
+
+  For each response whose call is no longer present, this restores the
+  original call event from `source_events` (the pre-compaction list),
+  inserting it immediately before the first surviving response that
+  references it. The whole call event is re-injected verbatim (rather than
+  trimmed to the resumed call) so parallel-call thought signatures, which only
+  the first part carries, are preserved. Any sibling responses that compaction
+  removed are re-injected too, so a sibling is not surfaced as a phantom
+  pending call.
+
+  Args:
+    events: The post-compaction events being assembled into request contents.
+    source_events: The pre-compaction events to recover missing calls from.
+
+  Returns:
+    `events` with any recoverable missing function-call events (and their
+    compacted sibling responses) re-injected; the original list is returned
+    unchanged when nothing needs recovery.
+  """
+  call_ids_present: set[str] = set()
+  response_ids_present: set[str] = set()
+  for event in events:
+    for function_call in event.get_function_calls():
+      if function_call.id:
+        call_ids_present.add(function_call.id)
+    for function_response in event.get_function_responses():
+      if function_response.id:
+        response_ids_present.add(function_response.id)
+
+  orphaned_ids = {
+      response_id
+      for response_id in response_ids_present
+      if response_id not in call_ids_present
+  }
+  if not orphaned_ids:
+    return events
+
+  call_event_by_id: dict[str, Event] = {}
+  for event in source_events:
+    for function_call in event.get_function_calls():
+      if function_call.id in orphaned_ids:
+        call_event_by_id.setdefault(function_call.id, event)
+
+  if not call_event_by_id:
+    return events
+
+  response_event_by_id: dict[str, Event] = {}
+  for event in source_events:
+    for function_response in event.get_function_responses():
+      if function_response.id:
+        response_event_by_id.setdefault(function_response.id, event)
+
+  result: list[Event] = []
+  reinjected_ids: set[str] = set()
+  for event in events:
+    for function_response in event.get_function_responses():
+      call_event = call_event_by_id.get(function_response.id)
+      if call_event is None or function_response.id in reinjected_ids:
+        continue
+      result.append(call_event)
+      sibling_ids = [
+          function_call.id
+          for function_call in call_event.get_function_calls()
+          if function_call.id
+      ]
+      reinjected_ids.update(sibling_ids)
+      # Recover sibling responses that compaction removed so a parallel sibling
+      # is not left looking like a pending call.
+      for sibling_id in sibling_ids:
+        if sibling_id not in response_ids_present:
+          sibling_response = response_event_by_id.get(sibling_id)
+          if sibling_response is not None:
+            result.append(sibling_response)
+    result.append(event)
+  return result
+
+
 def _copy_content_for_request(
     content: types.Content,
     *,
@@ -660,6 +750,12 @@ def _get_contents(
 
   if has_compaction_events:
     events_to_process = _process_compaction_events(raw_filtered_events)
+    # Compaction may have removed a function_call whose response survives
+    # (e.g. a long-running call resumed after it was compacted); restore it so
+    # the call/response pairing is intact.
+    events_to_process = _recover_compacted_function_calls(
+        events_to_process, raw_filtered_events
+    )
   else:
     events_to_process = raw_filtered_events
 

tests/unittests/flows/llm_flows/test_contents.py

@@ -15,6 +15,7 @@
 from google.adk.agents.llm_agent import Agent
 from google.adk.events.event import Event
 from google.adk.events.event_actions import EventActions
+from google.adk.events.event_actions import EventCompaction
 from google.adk.flows.llm_flows import _nl_planning
 from google.adk.flows.llm_flows import contents
 from google.adk.flows.llm_flows.contents import request_processor
@@ -1691,3 +1692,205 @@ def test_rearrange_async_function_responses_early_returns_when_no_responses():
       events
   )
   assert result is events
+
+
+def _long_running_call_event() -> Event:
+  return Event(
+      invocation_id="inv2",
+      author="model",
+      timestamp=2.0,
+      long_running_tool_ids={"lr-1"},
+      content=types.Content(
+          role="model",
+          parts=[
+              types.Part(
+                  function_call=types.FunctionCall(
+                      id="lr-1", name="lr_tool", args={}
+                  )
+              )
+          ],
+      ),
+  )
+
+
+def _long_running_response_event(response: dict[str, str]) -> Event:
+  return Event(
+      invocation_id="inv2",
+      author="user",
+      timestamp=4.0,
+      content=types.Content(
+          role="user",
+          parts=[
+              types.Part(
+                  function_response=types.FunctionResponse(
+                      id="lr-1", name="lr_tool", response=response
+                  )
+              )
+          ],
+      ),
+  )
+
+
+def test_recover_compacted_function_calls_reinjects_missing_call():
+  """A response whose call was compacted gets its call re-injected before it."""
+  summary_event = Event(
+      invocation_id="compacted",
+      author="model",
+      timestamp=3.0,
+      content=types.Content(role="model", parts=[types.Part(text="summary")]),
+  )
+  call_event = _long_running_call_event()
+  resume_response = _long_running_response_event({"result": "done"})
+
+  # After compaction the call is gone from the effective list but survives in
+  # the source (pre-compaction) list.
+  effective = [summary_event, resume_response]
+  source = [call_event, resume_response]
+
+  result = contents._recover_compacted_function_calls(effective, source)  # pylint: disable=protected-access
+
+  assert result == [summary_event, call_event, resume_response]
+
+
+def test_recover_compacted_function_calls_noop_when_call_present():
+  """No change when every response already has its call in the list."""
+  call_event = _long_running_call_event()
+  resume_response = _long_running_response_event({"result": "done"})
+  effective = [call_event, resume_response]
+
+  result = contents._recover_compacted_function_calls(effective, effective)  # pylint: disable=protected-access
+
+  assert result is effective
+
+
+def test_get_contents_recovers_compacted_long_running_call_on_resume():
+  """A long-running call compacted before resume is restored during assembly.
+
+  Reproduces issue #5602: the call and its intermediate placeholder response are
+  summarized away, then the real result arrives on resume. Without recovery,
+  assembly raises because the resumed response has no matching call.
+  """
+  compaction = EventCompaction(
+      start_timestamp=1.0,
+      end_timestamp=3.0,
+      compacted_content=types.Content(
+          role="model", parts=[types.Part(text="summary of earlier turns")]
+      ),
+  )
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          timestamp=1.0,
+          content=types.UserContent("start the long job"),
+      ),
+      _long_running_call_event(),
+      # Intermediate placeholder response (same invocation as the call).
+      _long_running_response_event({"status": "pending"}).model_copy(
+          update={"timestamp": 3.0}
+      ),
+      Event(
+          invocation_id="compacted",
+          author="model",
+          timestamp=3.0,
+          content=compaction.compacted_content,
+          actions=EventActions(compaction=compaction),
+      ),
+      # Real result delivered on resume; the runner stamps it with the call's
+      # invocation id, and its timestamp is outside the compacted range.
+      _long_running_response_event({"result": "done"}),
+  ]
+
+  result = contents._get_contents(None, events, agent_name="model")  # pylint: disable=protected-access
+
+  assert len(result) == 3
+  assert result[0].parts[0].text == "summary of earlier turns"
+  assert result[1].parts[0].function_call.id == "lr-1"
+  assert result[2].parts[0].function_response.id == "lr-1"
+  assert result[2].parts[0].function_response.response == {"result": "done"}
+
+
+def test_recover_compacted_parallel_call_reinjects_sibling_response():
+  """A recovered parallel call keeps every part and restores sibling responses.
+
+  The call event issues a long-running call (lr-1) and a regular call (reg-1)
+  together. Both, plus reg-1's response, are compacted; only lr-1 is resumed.
+  The whole call event is re-injected (so parallel-call thought signatures on
+  the first part survive), and reg-1's compacted response is restored so reg-1
+  is not surfaced as a phantom pending call.
+  """
+  parallel_call = Event(
+      invocation_id="inv2",
+      author="model",
+      timestamp=2.0,
+      long_running_tool_ids={"lr-1"},
+      content=types.Content(
+          role="model",
+          parts=[
+              types.Part(
+                  function_call=types.FunctionCall(
+                      id="lr-1", name="lr_tool", args={}
+                  )
+              ),
+              types.Part(
+                  function_call=types.FunctionCall(
+                      id="reg-1", name="reg_tool", args={}
+                  )
+              ),
+          ],
+      ),
+  )
+  compaction = EventCompaction(
+      start_timestamp=1.0,
+      end_timestamp=3.5,
+      compacted_content=types.Content(
+          role="model", parts=[types.Part(text="summary")]
+      ),
+  )
+  events = [
+      parallel_call,
+      # reg-1's response and lr-1's placeholder, both compacted.
+      _long_running_response_event({"status": "pending"}).model_copy(
+          update={"timestamp": 3.0}
+      ),
+      Event(
+          invocation_id="inv2",
+          author="user",
+          timestamp=3.5,
+          content=types.Content(
+              role="user",
+              parts=[
+                  types.Part(
+                      function_response=types.FunctionResponse(
+                          id="reg-1", name="reg_tool", response={"result": "ok"}
+                      )
+                  )
+              ],
+          ),
+      ),
+      Event(
+          invocation_id="compacted",
+          author="model",
+          timestamp=3.5,
+          content=compaction.compacted_content,
+          actions=EventActions(compaction=compaction),
+      ),
+      _long_running_response_event({"result": "done"}),
+  ]
+
+  result = contents._get_contents(None, events, agent_name="model")  # pylint: disable=protected-access
+
+  assert len(result) == 3
+  # The whole parallel call event is preserved (both parts, so a thought
+  # signature on the first part is not stripped).
+  call_ids = {
+      part.function_call.id for part in result[1].parts if part.function_call
+  }
+  assert call_ids == {"lr-1", "reg-1"}
+  # Both responses are present, so neither call looks pending.
+  response_ids = {
+      part.function_response.id
+      for part in result[2].parts
+      if part.function_response
+  }
+  assert response_ids == {"lr-1", "reg-1"}