88import random
99import threading
1010import time
11+ from collections import deque
1112from dataclasses import dataclass , field , replace
12- from typing import Any , Dict , Iterable , List , Optional , Set
13+ from typing import Any , Deque , Dict , Iterable , List , Optional , Set
1314
1415from .tool_result_store import FileToolResultStore
1516from .tools import Tool , ToolContext , ToolResult , ToolSet , error_result
2728 EVENT_TYPE_SESSION_STATUS_RUNNING ,
2829 EVENT_TYPE_SESSION_STATUS_TERMINATED ,
2930 EVENT_TYPE_USER_CUSTOM_TOOL_RESULT ,
31+ EVENT_TYPE_USER_INTERRUPT ,
3032 EVENT_TYPE_USER_TOOL_CONFIRMATION ,
3133 EVENT_TYPE_USER_TOOL_RESULT ,
3234 PERMISSION_ALLOW ,
@@ -98,6 +100,25 @@ class SessionToolRunnerOptions:
98100 on_tool_error : Any = None
99101
100102
103+ @dataclass
104+ class _PendingToolEvent :
105+ event : Event
106+ custom : bool
107+ confirmation : str = ""
108+
109+
110+ @dataclass
111+ class _ActiveToolExecution :
112+ pending : _PendingToolEvent
113+ cancel_event : _ToolCancelEvent
114+
115+
116+ @dataclass
117+ class _ToolExecutionResult :
118+ pending : _PendingToolEvent
119+ result : ToolResult
120+
121+
101122class SessionToolRunner :
102123 def __init__ (self , api : Any , session_id : str , options : SessionToolRunnerOptions ) -> None :
103124 if not session_id :
@@ -128,26 +149,29 @@ def _is_stopped(self) -> bool:
128149 return self ._stop .is_set () or bool (self .options .stop_event and self .options .stop_event .is_set ())
129150
130151 def run (self ) -> List [ToolCallResult ]:
131- if self .options .result_store is not None :
132- pending , processed = self .options .result_store .recover ()
133- self ._state .pending_results .update (pending )
134- self ._state .recovered_results .update (pending )
135- self ._state .processed .update (processed )
136- self ._state .answered .update (processed )
137- if self .options .prefer_stream and hasattr (self .api , "stream_events" ):
138- try :
139- self ._consume_stream_loop ()
140- return self ._results
141- except EventStreamUnsupported :
142- pass
143- except (IdleTimeout , SessionTerminated ):
144- raise
145- except Exception :
146- if self ._is_stopped ():
152+ try :
153+ if self .options .result_store is not None :
154+ pending , processed = self .options .result_store .recover ()
155+ self ._state .pending_results .update (pending )
156+ self ._state .recovered_results .update (pending )
157+ self ._state .processed .update (processed )
158+ self ._state .answered .update (processed )
159+ if self .options .prefer_stream and hasattr (self .api , "stream_events" ):
160+ try :
161+ self ._consume_stream_loop ()
147162 return self ._results
148- raise
149- self ._consume_list ()
150- return self ._results
163+ except EventStreamUnsupported :
164+ pass
165+ except (IdleTimeout , SessionTerminated ):
166+ raise
167+ except Exception :
168+ if self ._is_stopped ():
169+ return self ._results
170+ raise
171+ self ._consume_list ()
172+ return self ._results
173+ finally :
174+ self ._state .cancel_active_execution ()
151175
152176 def _consume_stream_loop (self ) -> None :
153177 backoff = STREAM_BACKOFF_START
@@ -158,6 +182,7 @@ def _consume_stream_loop(self) -> None:
158182 pump .start ()
159183 self ._state .reconcile ()
160184 while not self ._is_stopped () and (pump .is_alive () or not event_queue .empty ()):
185+ self ._state .drain_execution_done ()
161186 self ._state .flush_results ()
162187 self ._raise_if_idle_expired ()
163188 try :
@@ -204,6 +229,7 @@ def _put_stream_item(self, event_queue: "queue.Queue[object]", item: object) ->
204229
205230 def _consume_list (self ) -> None :
206231 while not self ._is_stopped ():
232+ self ._state .drain_execution_done ()
207233 self ._state .reconcile (reconcile = False )
208234 self ._state .flush_results ()
209235 self ._raise_if_idle_expired ()
@@ -212,6 +238,7 @@ def _consume_list(self) -> None:
212238 def _sleep_or_idle (self , seconds : float ) -> None :
213239 deadline = time .monotonic () + max (seconds , 0 )
214240 while not self ._is_stopped ():
241+ self ._state .drain_execution_done ()
215242 self ._raise_if_idle_expired ()
216243 remaining = deadline - time .monotonic ()
217244 if remaining <= 0 :
@@ -239,6 +266,11 @@ def __init__(self, runner: SessionToolRunner) -> None:
239266 self .pending_ask : Dict [str , Event ] = {}
240267 self .confirmations : Dict [str , Event ] = {}
241268 self .external_tools : Dict [str , Event ] = {}
269+ self .tool_use_events : Dict [str , Event ] = {}
270+ self .scheduled : Set [str ] = set ()
271+ self .execution_queue : Deque [_PendingToolEvent ] = deque ()
272+ self .active_execution : Optional [_ActiveToolExecution ] = None
273+ self .execution_done : "queue.Queue[_ToolExecutionResult]" = queue .Queue ()
242274 self .session_tool_uses : Set [str ] = set ()
243275 self .tool_uses_since_status : Set [str ] = set ()
244276 self .blocking_event_ids : Set [str ] = set ()
@@ -284,6 +316,7 @@ def _reconcile_once(self, *, reconcile: bool) -> None:
284316 def process_listed_events (self , events : Iterable [Event ], reconcile : bool = False ) -> None :
285317 pending : List [Event ] = []
286318 pending_ids : Dict [str , bool ] = {}
319+ replayed_tool_uses : Set [str ] = set ()
287320 touched_idle = False
288321 last_was_end_turn = False
289322 for event in events :
@@ -299,10 +332,13 @@ def process_listed_events(self, events: Iterable[Event], reconcile: bool = False
299332 )
300333 if event .type == EVENT_TYPE_USER_TOOL_CONFIRMATION :
301334 self .record_confirmation (event )
335+ elif event .type == EVENT_TYPE_USER_INTERRUPT :
336+ self .handle_interrupt (event , replayed_tool_uses if reconcile else None )
302337 elif event .type in (EVENT_TYPE_USER_TOOL_RESULT , EVENT_TYPE_USER_CUSTOM_TOOL_RESULT ):
303338 self .mark_answered (tool_result_call_id (event ))
304339 elif event .type in (EVENT_TYPE_AGENT_TOOL_USE , EVENT_TYPE_AGENT_CUSTOM_TOOL_USE ):
305340 call_id = tool_use_call_id (event )
341+ replayed_tool_uses .add (call_id )
306342 if call_id and not pending_ids .get (call_id ):
307343 pending .append (event )
308344 pending_ids [call_id ] = True
@@ -318,10 +354,7 @@ def process_listed_events(self, events: Iterable[Event], reconcile: bool = False
318354 self .handle_tool_use (event , event .type == EVENT_TYPE_AGENT_CUSTOM_TOOL_USE )
319355 self .release_confirmed_tool_uses ()
320356 if touched_idle and last_was_end_turn :
321- if self .has_unblocked_outstanding_tool (pending ):
322- self .disarm_idle ()
323- else :
324- self .arm_idle ()
357+ self .arm_idle ()
325358
326359 def note_idle_event (self , event : Event ) -> None :
327360 if event .type == EVENT_TYPE_USER_TOOL_CONFIRMATION :
@@ -343,6 +376,8 @@ def handle_event(self, event: Event) -> None:
343376 if event .type == EVENT_TYPE_USER_TOOL_CONFIRMATION :
344377 self .record_confirmation (event )
345378 self .release_confirmed_tool_uses ()
379+ elif event .type == EVENT_TYPE_USER_INTERRUPT :
380+ self .handle_interrupt (event )
346381 elif event .type in (EVENT_TYPE_USER_TOOL_RESULT , EVENT_TYPE_USER_CUSTOM_TOOL_RESULT ):
347382 self .mark_answered (tool_result_call_id (event ))
348383 elif event .type in (EVENT_TYPE_AGENT_TOOL_USE , EVENT_TYPE_AGENT_CUSTOM_TOOL_USE ):
@@ -352,6 +387,8 @@ def handle_event(self, event: Event) -> None:
352387
353388 def mark_event_seen (self , event : Event ) -> bool :
354389 key = event .id or tool_use_call_id (event )
390+ if not key and event .type == EVENT_TYPE_USER_INTERRUPT :
391+ key = f"interrupt:{ event .processed_at } :{ event .session_thread_id } "
355392 if not key :
356393 return True
357394 if self .seen .get (key ):
@@ -368,6 +405,7 @@ def mark_answered(self, call_id: str) -> None:
368405 self .recovered_results .discard (call_id )
369406 self .pending_ask .pop (call_id , None )
370407 self .external_tools .pop (call_id , None )
408+ self .tool_use_events .pop (call_id , None )
371409 self .maybe_arm_pending_idle ()
372410
373411 def is_answered (self , call_id : str ) -> bool :
@@ -384,19 +422,9 @@ def release_confirmed_tool_uses(self) -> None:
384422 self .pending_ask .pop (tool_use_call_id (event ), None )
385423 self .handle_tool_use (event , event .type == EVENT_TYPE_AGENT_CUSTOM_TOOL_USE )
386424
387- def has_unblocked_outstanding_tool (self , pending : Iterable [Event ]) -> bool :
388- for event in pending :
389- call_id = tool_use_call_id (event )
390- if not call_id or self .is_answered (call_id ) or not self .should_handle_tool_use (call_id ):
391- continue
392- if call_id in self .pending_ask or call_id in self .pending_results :
393- continue
394- return True
395- return False
396-
397425 def handle_tool_use (self , event : Event , custom : bool ) -> None :
398426 call_id = tool_use_call_id (event )
399- if not call_id or self .is_answered (call_id ):
427+ if not call_id or self .is_answered (call_id ) or call_id in self . scheduled :
400428 return
401429 pending = self .pending_results .get (call_id )
402430 if pending is not None :
@@ -424,8 +452,97 @@ def handle_tool_use(self, event: Event, custom: bool) -> None:
424452 self .pending_results [call_id ] = decision .result
425453 self .send_result (call_id , event , custom , "" , decision .result )
426454 return
427- result = self .execute_tool (event , custom )
428- self .post_result (event , custom , call_id , result , confirmation )
455+ self .scheduled .add (call_id )
456+ self .execution_queue .append (_PendingToolEvent (event , custom , confirmation ))
457+ self .start_next_tool_execution ()
458+
459+ def start_next_tool_execution (self ) -> None :
460+ if self .active_execution is not None :
461+ return
462+ while self .execution_queue :
463+ pending = self .execution_queue .popleft ()
464+ call_id = tool_use_call_id (pending .event )
465+ if self .is_answered (call_id ):
466+ self .scheduled .discard (call_id )
467+ continue
468+ cancel_event = _ToolCancelEvent (self .runner .options .tool_context .cancel_event )
469+ self .active_execution = _ActiveToolExecution (pending , cancel_event )
470+
471+ def execute (item : _PendingToolEvent = pending , cancel : _ToolCancelEvent = cancel_event ) -> None :
472+ result = self .execute_tool (item .event , item .custom , cancel )
473+ self .execution_done .put (_ToolExecutionResult (item , result ))
474+
475+ threading .Thread (target = execute , name = "ma-self-host-tool-dispatch" , daemon = True ).start ()
476+ return
477+
478+ def drain_execution_done (self ) -> None :
479+ while True :
480+ try :
481+ completed = self .execution_done .get_nowait ()
482+ except queue .Empty :
483+ return
484+ self .finish_tool_execution (completed )
485+
486+ def finish_tool_execution (self , completed : _ToolExecutionResult ) -> None :
487+ call_id = tool_use_call_id (completed .pending .event )
488+ active = self .active_execution
489+ if active is not None and tool_use_call_id (active .pending .event ) == call_id :
490+ active .cancel_event .set ()
491+ self .active_execution = None
492+ self .scheduled .discard (call_id )
493+ if not self .is_answered (call_id ):
494+ self .post_result (
495+ completed .pending .event ,
496+ completed .pending .custom ,
497+ call_id ,
498+ completed .result ,
499+ completed .pending .confirmation ,
500+ )
501+ self .start_next_tool_execution ()
502+
503+ def cancel_active_execution (self ) -> None :
504+ if self .active_execution is not None :
505+ self .active_execution .cancel_event .set ()
506+
507+ def handle_interrupt (self , event : Event , eligible : Optional [Set [str ]] = None ) -> None :
508+ thread_id = event .session_thread_id
509+
510+ def matches (call_id : str , tool_event : Event ) -> bool :
511+ if eligible is not None and call_id not in eligible :
512+ return False
513+ return not thread_id or tool_event .session_thread_id == thread_id
514+
515+ active = self .active_execution
516+ if active is not None :
517+ call_id = tool_use_call_id (active .pending .event )
518+ if matches (call_id , active .pending .event ):
519+ active .cancel_event .set ()
520+ self .settle_interrupted_tool_use (call_id )
521+ retained : Deque [_PendingToolEvent ] = deque ()
522+ for pending in self .execution_queue :
523+ call_id = tool_use_call_id (pending .event )
524+ if matches (call_id , pending .event ):
525+ self .settle_interrupted_tool_use (call_id )
526+ else :
527+ retained .append (pending )
528+ self .execution_queue = retained
529+ for call_id , tool_event in list (self .tool_use_events .items ()):
530+ if not self .is_answered (call_id ) and matches (call_id , tool_event ):
531+ self .settle_interrupted_tool_use (call_id )
532+
533+ def settle_interrupted_tool_use (self , call_id : str ) -> None :
534+ if not call_id or self .is_answered (call_id ):
535+ return
536+ self .scheduled .discard (call_id )
537+ self .mark_answered (call_id )
538+ if self .runner .options .result_store is None :
539+ return
540+ try :
541+ self .runner .options .result_store .discard (call_id )
542+ except Exception as exc : # noqa: BLE001 - interruption must not stop event consumption.
543+ self .runner .options .logger .warning (
544+ "discard interrupted tool result failed tool_use_id=%s err=%s" , call_id , exc
545+ )
429546
430547 def owns_tool (self , event : Event , custom : bool ) -> bool :
431548 if custom :
@@ -453,13 +570,15 @@ def permission_allows(self, event: Event, custom: bool, call_id: str) -> tuple:
453570 self .pending_ask [call_id ] = event
454571 return "" , False
455572
456- def execute_tool (self , event : Event , custom : bool ) -> ToolResult :
573+ def execute_tool (
574+ self , event : Event , custom : bool , execution_cancel : Optional [_ToolCancelEvent ] = None
575+ ) -> ToolResult :
457576 context = replace (self .runner .options .tool_context )
458577 if self .runner .options .tool_timeout_seconds is not None and self .runner .options .tool_timeout_seconds > 0 :
459578 context .tool_timeout_seconds = self .runner .options .tool_timeout_seconds
460579 if context .tool_timeout_seconds <= 0 :
461580 context .tool_timeout_seconds = DEFAULT_TOOL_TIMEOUT_SECONDS
462- cancel_event = _ToolCancelEvent (context .cancel_event )
581+ cancel_event = execution_cancel or _ToolCancelEvent (context .cancel_event )
463582 context .cancel_event = cancel_event
464583 results : "queue.Queue[ToolResult]" = queue .Queue (maxsize = 1 )
465584
@@ -579,6 +698,8 @@ def observe_session_state(self, event: Event) -> None:
579698 call_id = tool_use_call_id (event )
580699 if call_id :
581700 self .session_tool_uses .add (call_id )
701+ if not self .is_answered (call_id ):
702+ self .tool_use_events [call_id ] = event
582703 self .tool_uses_since_status .add (call_id )
583704 return
584705 if event .type == EVENT_TYPE_SESSION_STATUS_IDLE :
@@ -642,7 +763,7 @@ def maybe_arm_pending_idle(self) -> None:
642763 self .idle_armed_at = time .monotonic ()
643764
644765 def has_idle_blockers (self ) -> bool :
645- return bool (self .pending_ask or self .pending_results or self .external_tools )
766+ return bool (self .pending_ask or self .pending_results or self .external_tools or self . scheduled )
646767
647768 def idle_expired (self ) -> bool :
648769 max_idle = self .max_idle_seconds ()
0 commit comments