Skip to content

Commit da65851

Browse files
DeanChensjcopybara-github
DeanChensj
authored and
copybara-github
committed
feat: synthesize abort event and function response on cancellation
When an agent or workflow invocation is aborted, answer every function call it left without a response with a synthetic error response, carrying the author, branch and isolation scope of the call. If no call is pending, record a single abort event instead. Request building already drops orphaned function calls, so this is not needed to pass provider validation. What it adds is that the model sees the call was aborted instead of having it silently dropped, and that task-scope resolution, resume and compaction read consistent history. Behavior changes: - Resumable apps: abort followed by resume used to replay the dangling tool. The call is now answered, so the flow continues to the model with the abort error and a side-effecting tool is not re-run after the user cancelled. - Long-running calls (including confirmation and credential requests) are sealed too, instead of being kept pending. - The sealing events are persisted and visible to callers. They carry error_code="INVOCATION_ABORTED". A synthetic response is authored by the agent that issued the call; the event recorded when nothing is pending is authored by the root agent. Sealing after an early close or an error is best-effort: a failure is logged instead of masking the in-flight exception or skipping after_run. Part of #4796 Co-authored-by: Shangjie Chen <[email protected]> PiperOrigin-RevId: 991216123
1 parent 7d51295 commit da65851

9 files changed

Lines changed: 1060 additions & 34 deletions

File tree

‎src/google/adk/agents/_agent_router.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from typing import Optional
2222
from typing import TYPE_CHECKING
2323

24+
from ..events._abort_events import _is_abort_event
2425
from ..events._branch_path import _BranchPath
2526
from ..events._node_path_builder import _NodePathBuilder
2627
from ..events._rewind_events import _apply_rewinds
@@ -116,7 +117,11 @@ def find_agent_to_run(
116117
# type of the agent. e.g. a remote a2a agent may surface a credential
117118
# request as a special long-running function tool call.
118119
filtered_events = _apply_rewinds(session.events)
119-
event = find_matching_function_call(filtered_events)
120+
event = (
121+
find_matching_function_call(filtered_events)
122+
if filtered_events and not _is_abort_event(filtered_events[-1])
123+
else None
124+
)
120125
is_resumable = resumability_config and resumability_config.is_resumable
121126
# Only route based on a past function response if resumability is enabled.
122127
# In non-resumable scenarios, a turn ending with function call response
@@ -134,8 +139,8 @@ def find_agent_to_run(
134139
return resumed_agent
135140

136141
def _event_filter(event: Event) -> bool:
137-
"""Filters out user-authored events and agent state change events."""
138-
if event.author == "user":
142+
"""Filters out user, abort-sealing and agent state change events."""
143+
if event.author == "user" or _is_abort_event(event):
139144
return False
140145
if event.actions.agent_state is not None or event.actions.end_of_agent:
141146
return False

‎src/google/adk/agents/invocation_context.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ def __init__(
104104
self.signal = signal if signal is not None else asyncio.Event()
105105
self.loop = loop
106106
self.aborted = False
107+
self.event_synthesized = False
107108

108109
def __deepcopy__(self, memo: dict[int, Any] | None) -> _AbortState:
109110
# Preserve single-instance sharing across deepcopies and avoid traversing

‎src/google/adk/events/_abort_events.py‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Helpers for sealing an aborted invocation in session history.
16+
17+
When an invocation is cancelled mid-flight, the function calls it issued
18+
without a matching response are answered with synthetic error responses built
19+
here, and a plain abort event is recorded when there is nothing to answer.
20+
Request building would drop such orphaned calls anyway; sealing them instead
21+
lets the model see that the call was aborted, and gives task-scope resolution,
22+
resume and compaction a consistent history. On resume, a sealed call is not
23+
replayed: the flow continues to the model with the abort error, so a
24+
side-effecting tool is not re-run after the user cancelled.
25+
26+
The sealing events are persisted and visible to callers: they carry
27+
``error_code='INVOCATION_ABORTED'``, and every one is authored by an agent,
28+
either the agent that issued the sealed call or, when nothing is dangling, the
29+
root agent. Readers of session history (task-scope resolution, agent routing)
30+
use ``_is_abort_event`` to recognize them rather than their author.
31+
"""
32+
33+
from __future__ import annotations
34+
35+
import collections
36+
37+
from google.genai import types
38+
39+
from .event import Event
40+
41+
_INVOCATION_ABORTED = 'INVOCATION_ABORTED'
42+
_ABORT_MESSAGE = 'Invocation was aborted by client.'
43+
44+
45+
def _is_abort_event(event: Event) -> bool:
46+
"""Returns whether ``event`` was synthesized to seal an aborted invocation."""
47+
return event.error_code == _INVOCATION_ABORTED
48+
49+
50+
def _is_paused_task_reply(event: Event) -> bool:
51+
"""Whether a task agent already paused on this event to wait for the user."""
52+
return (
53+
bool(event.isolation_scope)
54+
and event.author != 'user'
55+
and event.is_final_response()
56+
)
57+
58+
59+
def _build_abort_events(
60+
events: list[Event],
61+
*,
62+
invocation_id: str,
63+
root_agent_name: str,
64+
branch: str | None,
65+
) -> list[Event]:
66+
"""Returns the events that seal an aborted invocation.
67+
68+
Each dangling function call gets a synthetic error response carrying the
69+
author, branch and isolation scope of the event that issued it, so the
70+
response pairs with its call in the issuing agent's own view. Calls sharing
71+
those three values are grouped into one event. If nothing is dangling, a
72+
single content-less abort event authored by the root agent is returned
73+
instead.
74+
75+
Long-running calls (including confirmation and credential requests) are
76+
sealed too: the invocation that issued them is gone, so they are not left
77+
pending for a later turn.
78+
79+
Args:
80+
events: The session history, in chronological order.
81+
invocation_id: The aborted invocation; only its events are considered.
82+
root_agent_name: Author of the abort event returned when nothing is
83+
dangling.
84+
branch: Branch for calls whose event has no branch, and for the abort
85+
event returned when nothing is dangling.
86+
87+
Returns:
88+
The synthetic events, in the order they should be appended.
89+
"""
90+
invocation_events = [e for e in events if e.invocation_id == invocation_id]
91+
if invocation_events and _is_paused_task_reply(invocation_events[-1]):
92+
return []
93+
answered_call_ids = {
94+
fr.id
95+
for e in invocation_events
96+
for fr in e.get_function_responses()
97+
if fr.id
98+
}
99+
100+
grouped_calls: dict[
101+
tuple[str, str | None, str | None], list[types.FunctionCall]
102+
] = collections.defaultdict(list)
103+
for event in invocation_events:
104+
for fc in event.get_function_calls():
105+
if fc.id and fc.id not in answered_call_ids:
106+
key = (event.author, event.branch or branch, event.isolation_scope)
107+
grouped_calls[key].append(fc)
108+
109+
if not grouped_calls:
110+
return [
111+
Event(
112+
invocation_id=invocation_id,
113+
author=root_agent_name,
114+
branch=branch,
115+
error_code=_INVOCATION_ABORTED,
116+
error_message=_ABORT_MESSAGE,
117+
)
118+
]
119+
120+
return [
121+
Event(
122+
invocation_id=invocation_id,
123+
author=author,
124+
branch=call_branch,
125+
isolation_scope=isolation_scope,
126+
content=types.Content(
127+
role='user',
128+
parts=[
129+
types.Part(
130+
function_response=types.FunctionResponse(
131+
id=fc.id,
132+
name=fc.name,
133+
response={'error': _ABORT_MESSAGE},
134+
)
135+
)
136+
for fc in calls
137+
],
138+
),
139+
error_code=_INVOCATION_ABORTED,
140+
error_message=_ABORT_MESSAGE,
141+
)
142+
for (author, call_branch, isolation_scope), calls in grouped_calls.items()
143+
]

‎src/google/adk/runners.py‎

Lines changed: 86 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@
4949
from .auth.credential_service.base_credential_service import BaseCredentialService
5050
from .errors._stale_session_error import StaleSessionError
5151
from .errors.session_not_found_error import SessionNotFoundError
52+
from .events._abort_events import _build_abort_events
53+
from .events._abort_events import _is_abort_event
5254
from .events._internal_metadata import internal_metadata
5355
from .events._internal_metadata import without_internal_metadata
5456
from .events._rewind_events import _apply_rewinds
@@ -107,10 +109,11 @@ def _find_active_task_scope(session: Session) -> Optional[tuple[str, str]]:
107109
'result' key matching ``FINISH_TASK_SUCCESS_RESULT`` or
108110
``FINISH_TASK_ERROR_RESULT``. A FunctionResponse containing an 'error' key
109111
(indicating a tool validation failure) does NOT close the scope: the task
110-
agent is still active, will see the validation error, and retry. Walking
111-
backward, the first non-empty scope we encounter that hasn't been closed by a
112-
later successful or failed terminal ``finish_task`` is the paused task
113-
awaiting the user's next reply.
112+
agent is still active, will see the validation error, and retry. Scopes whose
113+
invocation was aborted are also closed, since the task agent was cancelled
114+
without calling ``finish_task``. Walking backward, the first non-empty scope
115+
we encounter that hasn't been closed is the paused task awaiting the user's
116+
next reply.
114117
115118
Used by ``Runner._append_user_event`` to scope the new user message
116119
to that task agent's view.
@@ -125,7 +128,10 @@ def _find_active_task_scope(session: Session) -> Optional[tuple[str, str]]:
125128
# the older success FR, falsely indicating the scope is still active.
126129
live_events = _apply_rewinds(session.events)
127130
finished_scopes: set[str] = set()
131+
aborted_invocations: set[str] = set()
128132
for event in live_events:
133+
if _is_abort_event(event):
134+
aborted_invocations.add(event.invocation_id)
129135
scope = event.isolation_scope
130136
if not scope:
131137
continue
@@ -146,7 +152,10 @@ def _find_active_task_scope(session: Session) -> Optional[tuple[str, str]]:
146152
scope = event.isolation_scope
147153
if not scope:
148154
continue
149-
if scope not in finished_scopes:
155+
if (
156+
scope not in finished_scopes
157+
and event.invocation_id not in aborted_invocations
158+
):
150159
return scope, event.invocation_id
151160
return None
152161

@@ -879,6 +888,27 @@ async def _cleanup_root_task(
879888
logger.error('Root node %s failed.', node_name, exc_info=True)
880889
raise
881890

891+
async def _synthesize_abort_events_if_needed(
892+
self, ic: InvocationContext
893+
) -> list[Event]:
894+
"""Seals an aborted invocation in session history, at most once.
895+
896+
Returns the synthetic events after plugin processing and persistence.
897+
"""
898+
if ic._abort_state.event_synthesized or ic.session is None: # pylint: disable=protected-access
899+
return []
900+
ic._abort_state.event_synthesized = True # pylint: disable=protected-access
901+
abort_events = _build_abort_events(
902+
ic.session.events,
903+
invocation_id=ic.invocation_id,
904+
root_agent_name=self.agent.name,
905+
branch=ic.branch,
906+
)
907+
return [
908+
await self._process_and_append_event(event=e, invocation_context=ic)
909+
for e in abort_events
910+
]
911+
882912
async def _run_post_invocation_compaction(
883913
self,
884914
*,
@@ -1339,12 +1369,10 @@ async def execute(
13391369
) as agen:
13401370
async for event in agen:
13411371
yield event
1342-
if invocation_context.is_aborted:
1343-
break
1344-
# Run compaction after all events are yielded from the agent.
1345-
# (We don't compact in the middle of an invocation, we only compact at
1346-
# the end of an invocation.)
13471372
if not invocation_context.is_aborted:
1373+
# Run compaction after all events are yielded from the agent.
1374+
# (We don't compact in the middle of an invocation, we only compact at
1375+
# the end of an invocation.)
13481376
await self._run_post_invocation_compaction(
13491377
session=invocation_context.session,
13501378
skip_token_compaction=(
@@ -1489,6 +1517,34 @@ async def _process_event_with_plugin_callbacks(
14891517
run_config=invocation_context.run_config,
14901518
)
14911519

1520+
async def _process_and_append_event(
1521+
self,
1522+
event: Event,
1523+
invocation_context: InvocationContext,
1524+
*,
1525+
is_live_call: bool = False,
1526+
) -> Event:
1527+
"""Applies custom metadata, runs on_event callbacks, and appends to session."""
1528+
output_event = await self._process_event_with_plugin_callbacks(
1529+
invocation_context=invocation_context,
1530+
event=event,
1531+
)
1532+
1533+
if is_live_call:
1534+
if output_event.partial is not True and self._should_append_event(
1535+
output_event, is_live_call
1536+
):
1537+
logger.debug('Appending live event: %s', output_event)
1538+
await self.session_service.append_event(
1539+
session=invocation_context.session, event=output_event
1540+
)
1541+
else:
1542+
if output_event.partial is not True:
1543+
await self.session_service.append_event(
1544+
session=invocation_context.session, event=output_event
1545+
)
1546+
return output_event
1547+
14921548
async def _exec_with_plugin(
14931549
self,
14941550
invocation_context: InvocationContext,
@@ -1538,30 +1594,18 @@ async def _exec_with_plugin(
15381594
# Step 2: Otherwise continue with normal execution
15391595
async with aclosing(execute_fn(invocation_context)) as agen:
15401596
async for event in agen:
1541-
# Step 3: Run the on_event callbacks before persisting so callback
1542-
# changes are stored in the session and match the streamed event.
1543-
output_event = await self._process_event_with_plugin_callbacks(
1544-
invocation_context=invocation_context,
1597+
output_event = await self._process_and_append_event(
15451598
event=event,
1599+
invocation_context=invocation_context,
1600+
is_live_call=is_live_call,
15461601
)
1547-
1548-
if is_live_call:
1549-
# Skip partial transcriptions for Live
1550-
if (
1551-
output_event.partial is not True
1552-
and self._should_append_event(output_event, is_live_call)
1553-
):
1554-
logger.debug('Appending live event: %s', output_event)
1555-
await self.session_service.append_event(
1556-
session=invocation_context.session, event=output_event
1557-
)
1558-
else:
1559-
if output_event.partial is not True:
1560-
await self.session_service.append_event(
1561-
session=invocation_context.session, event=output_event
1562-
)
1563-
15641602
yield output_event
1603+
if not is_live_call and invocation_context.is_aborted:
1604+
abort_events = await self._synthesize_abort_events_if_needed(
1605+
invocation_context
1606+
)
1607+
for abort_event in abort_events:
1608+
yield abort_event
15651609
except GeneratorExit:
15661610
# Early generator close is treated as a clean completion.
15671611
closing_early = True
@@ -1584,6 +1628,17 @@ async def _exec_with_plugin(
15841628
run_error = e
15851629
raise
15861630
finally:
1631+
if not is_live_call and invocation_context.is_aborted:
1632+
# Best-effort: only reached on early close or error, where raising
1633+
# would mask the in-flight exception and skip after_run.
1634+
try:
1635+
await self._synthesize_abort_events_if_needed(invocation_context)
1636+
except Exception: # pylint: disable=broad-exception-caught
1637+
logger.error(
1638+
'Failed to seal aborted invocation %s.',
1639+
invocation_context.invocation_id,
1640+
exc_info=True,
1641+
)
15871642
# Step 4: Run after_run callbacks on successful completion or early exit.
15881643
if run_error is None:
15891644
try:

‎src/google/adk/workflow/_node_runner_utils.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,12 @@ async def _drive_root_node() -> None:
283283
await runner._cleanup_root_task( # pylint: disable=protected-access
284284
task, runner.agent.name
285285
)
286+
if ic.is_aborted:
287+
abort_events = await runner._synthesize_abort_events_if_needed( # pylint: disable=protected-access
288+
ic
289+
)
290+
for abort_event in abort_events:
291+
yield abort_event
286292
except GeneratorExit:
287293
# Caller cancelled or broke out of the generator early (e.g. via aclosing).
288294
# GeneratorExit is a BaseException, but early generator close is considered
@@ -314,6 +320,17 @@ async def _drive_root_node() -> None:
314320
# after_run plugin raising, which PluginManager surfaces as a
315321
# RuntimeError) is itself an unhandled runner error, so notify
316322
# on_run_error_callback once and re-raise (unless closing_early).
323+
if ic.is_aborted:
324+
# Best-effort: only reached on early close or error, where raising
325+
# would mask the in-flight exception and skip after_run.
326+
try:
327+
await runner._synthesize_abort_events_if_needed(ic) # pylint: disable=protected-access
328+
except Exception: # pylint: disable=broad-exception-caught
329+
logger.error(
330+
"Failed to seal aborted invocation %s.",
331+
ic.invocation_id,
332+
exc_info=True,
333+
)
317334
if run_error is None:
318335
try:
319336
await ic.plugin_manager.run_after_run_callback(

0 commit comments

Comments
 (0)