Skip to content

Commit d99b23f

Browse files
DeanChensjcopybara-github
authored andcommitted
feat(context): add abort_signal support to InvocationContext and Context
Introduces thread-safe abort signal primitives onto InvocationContext, Context, and ReadonlyContext: - Add _abort_state: _AbortState as a PrivateAttr container to preserve signal and loop references across Pydantic model_copy() derivations. - Capture event loop at creation for cross-thread dispatch via call_soon_threadsafe. - Expose read-only is_aborted property on ReadonlyContext (inherited by Context). This PR introduces context-level signal primitives; a follow-up PR wires them into Runner and Workflow execution loops. Co-authored-by: Shangjie Chen <[email protected]> PiperOrigin-RevId: 983478644
1 parent 4ff3cdc commit d99b23f

4 files changed

Lines changed: 454 additions & 0 deletions

File tree

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

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,27 @@ def increment_and_enforce_llm_calls_limit(
8484
)
8585

8686

87+
class _AbortState:
88+
"""Shared mutable state container for invocation abort signals.
89+
90+
Because Pydantic model_copy() shallow-copies __pydantic_private__, all
91+
derived contexts within the same Runner share the exact same _AbortState
92+
instance reference. Updates to loop, signal, or aborted propagate across all
93+
model_copy() clones in the tree. Cross-Runner sub-runs (such as AgentTool or
94+
nested Workflow node runners) propagate cancellation by passing
95+
``_abort_signal`` to the child Runner's ``run_async``.
96+
"""
97+
98+
def __init__(
99+
self,
100+
signal: asyncio.Event | None = None,
101+
loop: asyncio.AbstractEventLoop | None = None,
102+
) -> None:
103+
self.signal = signal if signal is not None else asyncio.Event()
104+
self.loop = loop
105+
self.aborted = False
106+
107+
87108
class InvocationContext(BaseModel):
88109
"""An invocation context represents the data of a single invocation of an agent.
89110
@@ -258,11 +279,56 @@ class InvocationContext(BaseModel):
258279
of this invocation.
259280
"""
260281

282+
_abort_state: _AbortState = PrivateAttr(default_factory=_AbortState)
283+
"""Captured abort state (signal, loop, and aborted flag) shared across copies."""
284+
261285
@override
262286
def model_post_init(self, __context: Any) -> None:
263287
super().model_post_init(__context)
264288
if self.run_config and self.run_config.custom_metadata:
265289
self._custom_metadata.update(self.run_config.custom_metadata)
290+
try:
291+
self._abort_state.loop = asyncio.get_running_loop()
292+
except RuntimeError:
293+
pass
294+
295+
@property
296+
def _abort_signal(self) -> asyncio.Event:
297+
"""The internal abort signal event for this invocation.
298+
299+
Private on purpose so callers use ``is_aborted``, while internal runners use
300+
``_abort_signal`` to await cancellation. Sub-contexts are shallow
301+
``model_copy`` clones that share this exact Event instance via
302+
``_abort_state``.
303+
"""
304+
if self._abort_state.loop is None:
305+
try:
306+
self._abort_state.loop = asyncio.get_running_loop()
307+
except RuntimeError:
308+
pass
309+
return self._abort_state.signal
310+
311+
def _attach_abort_signal(self, abort_signal: asyncio.Event) -> None:
312+
"""Replaces this context's abort signal with a caller-owned event.
313+
314+
Only for the runner to call on a freshly built root context, before any
315+
sub-context is derived from it. The runner cannot pass the signal to the
316+
constructor because the context is produced by ``_new_invocation_context``,
317+
an overridable factory whose subclass overrides do not accept the argument.
318+
319+
Because sub-contexts share ``_abort_state``, replacing the signal here
320+
propagates to all derived contexts. Rebinding after execution begins is
321+
unsafe as active tasks may already be awaiting the previous event.
322+
323+
Args:
324+
abort_signal: The caller-owned event to abort this invocation with.
325+
"""
326+
self._abort_state.signal = abort_signal
327+
if self._abort_state.loop is None:
328+
try:
329+
self._abort_state.loop = asyncio.get_running_loop()
330+
except RuntimeError:
331+
pass
266332

267333
@property
268334
def is_resumable(self) -> bool:
@@ -272,6 +338,44 @@ def is_resumable(self) -> bool:
272338
and self.resumability_config.is_resumable
273339
)
274340

341+
@property
342+
def is_aborted(self) -> bool:
343+
"""Returns whether the current invocation has been requested to abort."""
344+
return self._abort_state.aborted or self._abort_state.signal.is_set()
345+
346+
def abort(self) -> None:
347+
"""Trip the abort signal in a thread-safe manner.
348+
349+
Can be safely called from either the event loop thread or an external
350+
worker thread. When called from within the running event loop, trips the
351+
signal immediately on the same tick. When called from a foreign thread,
352+
schedules the trip thread-safely onto the captured event loop.
353+
"""
354+
self._abort_state.aborted = True
355+
signal = self._abort_state.signal
356+
if signal.is_set():
357+
return
358+
359+
try:
360+
running_loop = asyncio.get_running_loop()
361+
except RuntimeError:
362+
running_loop = None
363+
364+
target_loop = self._abort_state.loop
365+
if target_loop is not None and running_loop is not target_loop:
366+
try:
367+
target_loop.call_soon_threadsafe(signal.set)
368+
except RuntimeError:
369+
try:
370+
signal.set()
371+
except RuntimeError:
372+
pass
373+
else:
374+
try:
375+
signal.set()
376+
except RuntimeError:
377+
pass
378+
275379
async def _enqueue_event(self, event: Event) -> None:
276380
"""Enqueue an event for the Runner main loop to process.
277381

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,11 @@ def custom_metadata(self) -> Mapping[str, Any]:
7979
# pylint: disable=protected-access
8080
return MappingProxyType(self._invocation_context._custom_metadata)
8181

82+
@property
83+
def is_aborted(self) -> bool:
84+
"""Returns whether the current invocation has been requested to abort."""
85+
return self._invocation_context.is_aborted
86+
8287
def get_credential(self, key: str) -> AuthCredential | None:
8388
"""Gets a resolved credential by key for this invocation."""
8489
return self._invocation_context.credential_by_key.get(key)

‎tests/unittests/agents/test_context.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import asyncio
1516
from unittest.mock import AsyncMock
1617
from unittest.mock import MagicMock
1718
from unittest.mock import patch
@@ -44,6 +45,7 @@ def mock_invocation_context():
4445
mock_context.artifact_service = None
4546
mock_context.credential_service = None
4647
mock_context.memory_service = None
48+
mock_context.is_aborted = False
4749
return mock_context
4850

4951

@@ -55,6 +57,43 @@ def test_context_branch_returns_invocation_branch(mock_invocation_context):
5557
assert context.branch == "test-branch"
5658

5759

60+
def test_context_is_aborted(mock_invocation_context):
61+
"""Context.is_aborted delegates to invocation context."""
62+
mock_invocation_context.is_aborted = False
63+
context = Context(invocation_context=mock_invocation_context)
64+
assert context.is_aborted is False
65+
66+
mock_invocation_context.is_aborted = True
67+
assert context.is_aborted is True
68+
69+
70+
@pytest.mark.asyncio
71+
async def test_context_is_aborted_with_real_invocation_context():
72+
"""Context reflects underlying InvocationContext cancellation."""
73+
from google.adk.agents.base_agent import BaseAgent
74+
from google.adk.agents.invocation_context import InvocationContext
75+
from google.adk.sessions.base_session_service import BaseSessionService
76+
from google.adk.sessions.session import Session
77+
78+
abort_signal = asyncio.Event()
79+
inv_ctx = InvocationContext(
80+
session_service=MagicMock(spec=BaseSessionService),
81+
agent=MagicMock(spec=BaseAgent),
82+
invocation_id="inv_1",
83+
session=Session(id="s1", app_name="test_app", user_id="test_user"),
84+
)
85+
inv_ctx._attach_abort_signal(abort_signal)
86+
context = Context(invocation_context=inv_ctx)
87+
assert context.is_aborted is False
88+
assert inv_ctx.is_aborted is False
89+
90+
inv_ctx.abort()
91+
92+
assert context.is_aborted is True
93+
assert inv_ctx.is_aborted is True
94+
assert abort_signal.is_set() is True
95+
96+
5897
@pytest.fixture
5998
def mock_artifact_service():
6099
"""Create a mock artifact service for testing."""

0 commit comments

Comments
 (0)