Skip to content

Commit 5c0eaf6

Browse files
committed
feat(selfHosted): support interrupt and multimodal read
Hand-written change from the internal SDK repository. The original commit message is intentionally not mirrored; see the source repository for details. Sync-Source-Commit: d0250a00282194e6454b4494ddc1d78603bf3ec0 Hand-Written-Reason: No Ark-APIs provenance marker; treated as a hand-written source commit. Release-Version: 0.9.0
1 parent 42ec723 commit 5c0eaf6

5 files changed

Lines changed: 709 additions & 47 deletions

File tree

‎src/arkruntime/selfhosted/session_tool_runner.py‎

Lines changed: 161 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,9 @@
88
import random
99
import threading
1010
import time
11+
from collections import deque
1112
from 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

1415
from .tool_result_store import FileToolResultStore
1516
from .tools import Tool, ToolContext, ToolResult, ToolSet, error_result
@@ -27,6 +28,7 @@
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+
101122
class 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

Comments
 (0)