@@ -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+
87108class 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
0 commit comments