Skip to content

fix: preserve pending function calls during event compaction #4743

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
atian8179 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: preserve pending function calls during event compaction #4743

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 33 additions & 2 deletions src/google/adk/apps/compaction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -273,6 +273,22 @@ def _event_function_response_ids(event: Event) -> set[str]:
return function_response_ids


def _pending_function_call_ids(events: list[Event]) -> set[str]:
"""Returns function call IDs that have no matching response in the event list.

These represent pending/unanswered function calls (e.g., from
adk_request_confirmation) that must not be compacted away, otherwise
a later FunctionResponse will fail to find its matching FunctionCall.
See #4740.
"""
all_call_ids: set[str] = set()
all_response_ids: set[str] = set()
for event in events:
all_call_ids.update(_event_function_call_ids(event))
all_response_ids.update(_event_function_response_ids(event))
return all_call_ids - all_response_ids


def _safe_token_compaction_split_index(
*,
candidate_events: list[Event],
Expand All @@ -285,6 +301,10 @@ def _safe_token_compaction_split_index(
assembly can fail. This method shifts the split earlier so matching function
call events are retained together with their responses.

Also preserves pending (unanswered) function call events to prevent
compacting away function calls that have not yet received a response
(e.g., adk_request_confirmation). See #4740.

Iterates backwards through candidate_events once, maintaining a running set
of unmatched response IDs. The latest valid split point where no unmatched
responses remain is returned.
Expand All @@ -293,6 +313,9 @@ def _safe_token_compaction_split_index(
if initial_split <= 0:
return 0

# Identify pending function calls across ALL candidate events
pending_call_ids = _pending_function_call_ids(candidate_events)

unmatched_response_ids: set[str] = set()
best_split = 0

Expand All @@ -302,9 +325,17 @@ def _safe_token_compaction_split_index(
call_ids = _event_function_call_ids(event)
unmatched_response_ids -= call_ids

# Also check: would compacting events [0..i) remove a pending call?
if not unmatched_response_ids and i <= initial_split:
best_split = i
break
# Verify no pending function calls exist in the compacted prefix
pending_in_prefix = False
for j in range(i):
if _event_function_call_ids(candidate_events[j]) & pending_call_ids:
pending_in_prefix = True
break
if not pending_in_prefix:
best_split = i
break
Comment on lines +332 to +338
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

security-medium medium

The _safe_token_compaction_split_index function introduces a nested loop, leading to $O(N^2)$ time complexity. This can be exploited by an attacker to cause a Denial of Service (DoS) by generating a large number of events, consuming excessive CPU. To remediate this, optimize the check for pending function calls in the prefix to achieve $O(N)$ complexity, for example, by pre-calculating the index of the first event that contains a pending function call.

   # Identify pending function calls across ALL candidate events
  pending_call_ids = _pending_function_call_ids(candidate_events)

  # Pre-calculate the index of the first event containing a pending call
  first_pending_call_idx = -1
  for idx, event in enumerate(candidate_events):
    if _event_function_call_ids(event) & pending_call_ids:
      first_pending_call_idx = idx
      break

  unmatched_response_ids: set[str] = set()
  best_split = 0

  for i in range(len(candidate_events) - 1, -1, -1):
    event = candidate_events[i]
    unmatched_response_ids.update(_event_function_response_ids(event))
    call_ids = _event_function_call_ids(event)
    unmatched_response_ids -= call_ids

    # Also check: would compacting events [0..i) remove a pending call?
    if not unmatched_response_ids and i <= initial_split:
      # Verify no pending function calls exist in the compacted prefix
      if first_pending_call_idx == -1 or first_pending_call_idx >= i:
        best_split = i
        break


return best_split

Expand Down