From e49065ffa38f05a2a4bc764aa5730fb49f6bf7fa Mon Sep 17 00:00:00 2001 From: "ark-hand[bot]" Date: Fri, 11 Sep 2026 03:06:28 +0000 Subject: [PATCH] fix(selfHosted): filter recovered tool results by session blockers Sync-Source-Commit: 28753fc065eef422f30ff63813019b6430b3f005 Hand-Written-Reason: No Ark-APIs provenance marker; treated as a hand-written source commit. Release-Version: 0.7.0 --- .../selfhosted/session_tool_runner.py | 75 +++++++++- .../selfhosted/tool_result_store.py | 10 ++ src/arkruntime/selfhosted/types.py | 10 ++ tests/selfhosted/test_session_tool_runner.py | 130 ++++++++++++++++++ tests/selfhosted/test_tool_result_store.py | 9 ++ 5 files changed, 230 insertions(+), 4 deletions(-) diff --git a/src/arkruntime/selfhosted/session_tool_runner.py b/src/arkruntime/selfhosted/session_tool_runner.py index cbc25b0..f69b4ea 100644 --- a/src/arkruntime/selfhosted/session_tool_runner.py +++ b/src/arkruntime/selfhosted/session_tool_runner.py @@ -9,7 +9,7 @@ import threading import time from dataclasses import dataclass, field, replace -from typing import Any, Dict, Iterable, List, Optional +from typing import Any, Dict, Iterable, List, Optional, Set from .tool_result_store import FileToolResultStore from .tools import Tool, ToolContext, ToolResult, ToolSet, error_result @@ -23,6 +23,8 @@ EVENT_TYPE_AGENT_TOOL_USE, EVENT_TYPE_SESSION_DELETED, EVENT_TYPE_SESSION_STATUS_IDLE, + EVENT_TYPE_SESSION_STATUS_RESCHEDULED, + EVENT_TYPE_SESSION_STATUS_RUNNING, EVENT_TYPE_SESSION_STATUS_TERMINATED, EVENT_TYPE_USER_CUSTOM_TOOL_RESULT, EVENT_TYPE_USER_TOOL_CONFIRMATION, @@ -30,6 +32,7 @@ PERMISSION_ALLOW, PERMISSION_DENY, SESSION_STOP_REASON_END_TURN, + SESSION_STOP_REASON_REQUIRES_ACTION, ContentBlock, Event, EventStreamUnsupported, @@ -128,6 +131,7 @@ def run(self) -> List[ToolCallResult]: if self.options.result_store is not None: pending, processed = self.options.result_store.recover() self._state.pending_results.update(pending) + self._state.recovered_results.update(pending) self._state.processed.update(processed) self._state.answered.update(processed) if self.options.prefer_stream and hasattr(self.api, "stream_events"): @@ -200,8 +204,8 @@ def _put_stream_item(self, event_queue: "queue.Queue[object]", item: object) -> def _consume_list(self) -> None: while not self._is_stopped(): - self._state.flush_results() self._state.reconcile(reconcile=False) + self._state.flush_results() self._raise_if_idle_expired() self._sleep_or_idle(self.options.event_poll_interval_seconds) @@ -231,9 +235,14 @@ def __init__(self, runner: SessionToolRunner) -> None: self.seen: Dict[str, bool] = {} self.answered: Dict[str, bool] = {} self.pending_results: Dict[str, Event] = {} + self.recovered_results: Set[str] = set() self.pending_ask: Dict[str, Event] = {} self.confirmations: Dict[str, Event] = {} self.external_tools: Dict[str, Event] = {} + self.session_tool_uses: Set[str] = set() + self.tool_uses_since_status: Set[str] = set() + self.blocking_event_ids: Set[str] = set() + self.blocking_events_known = False self.idle_armed_at = 0.0 self.idle_arm_pending = False @@ -281,6 +290,7 @@ def process_listed_events(self, events: Iterable[Event], reconcile: bool = False seen_now = self.mark_event_seen(event) if not reconcile and not seen_now: continue + self.observe_session_state(event) if seen_now and event.type != EVENT_TYPE_USER_TOOL_CONFIRMATION: touched_idle = True last_was_end_turn = ( @@ -298,10 +308,12 @@ def process_listed_events(self, events: Iterable[Event], reconcile: bool = False pending_ids[call_id] = True elif event.type in (EVENT_TYPE_SESSION_STATUS_TERMINATED, EVENT_TYPE_SESSION_DELETED): raise SessionTerminated("session terminated") + self.reconcile_recovered_results() if touched_idle: self.disarm_idle() for event in pending: - if self.is_answered(tool_use_call_id(event)): + call_id = tool_use_call_id(event) + if self.is_answered(call_id) or not self.should_handle_tool_use(call_id): continue self.handle_tool_use(event, event.type == EVENT_TYPE_AGENT_CUSTOM_TOOL_USE) self.release_confirmed_tool_uses() @@ -322,6 +334,8 @@ def note_idle_event(self, event: Event) -> None: def handle_stream_event(self, event: Event) -> None: if not self.mark_event_seen(event): return + self.observe_session_state(event) + self.reconcile_recovered_results() self.note_idle_event(event) self.handle_event(event) @@ -351,6 +365,7 @@ def mark_answered(self, call_id: str) -> None: self.answered[call_id] = True self.processed[call_id] = True self.pending_results.pop(call_id, None) + self.recovered_results.discard(call_id) self.pending_ask.pop(call_id, None) self.external_tools.pop(call_id, None) self.maybe_arm_pending_idle() @@ -372,7 +387,7 @@ def release_confirmed_tool_uses(self) -> None: def has_unblocked_outstanding_tool(self, pending: Iterable[Event]) -> bool: for event in pending: call_id = tool_use_call_id(event) - if not call_id or self.is_answered(call_id): + if not call_id or self.is_answered(call_id) or not self.should_handle_tool_use(call_id): continue if call_id in self.pending_ask or call_id in self.pending_results: continue @@ -385,6 +400,8 @@ def handle_tool_use(self, event: Event, custom: bool) -> None: return pending = self.pending_results.get(call_id) if pending is not None: + if call_id in self.recovered_results: + return self.send_result(call_id, event, custom, "", pending) return if not self.owns_tool(event, custom): @@ -541,6 +558,8 @@ def retry_send_event(self, event: Event, call_id: str) -> bool: def flush_results(self) -> None: for call_id, event in list(self.pending_results.items()): + if call_id in self.recovered_results: + continue if self.retry_send_event(event, call_id): self.mark_answered(call_id) if self.runner.options.result_store is not None: @@ -555,6 +574,54 @@ def flush_results(self) -> None: ) self.maybe_arm_pending_idle() + def observe_session_state(self, event: Event) -> None: + if event.type in (EVENT_TYPE_AGENT_TOOL_USE, EVENT_TYPE_AGENT_CUSTOM_TOOL_USE): + call_id = tool_use_call_id(event) + if call_id: + self.session_tool_uses.add(call_id) + self.tool_uses_since_status.add(call_id) + return + if event.type == EVENT_TYPE_SESSION_STATUS_IDLE: + self.blocking_events_known = True + self.blocking_event_ids = set() + if event.stop_reason_type() == SESSION_STOP_REASON_REQUIRES_ACTION: + self.blocking_event_ids.update(event.stop_reason_event_ids()) + self.tool_uses_since_status.clear() + return + if event.type in (EVENT_TYPE_SESSION_STATUS_RUNNING, EVENT_TYPE_SESSION_STATUS_RESCHEDULED): + self.blocking_events_known = True + self.blocking_event_ids.clear() + self.tool_uses_since_status.clear() + + def should_handle_tool_use(self, call_id: str) -> bool: + if not self.blocking_events_known: + return True + return call_id in self.blocking_event_ids or call_id in self.tool_uses_since_status + + def reconcile_recovered_results(self) -> None: + if not self.blocking_events_known: + return + for call_id in list(self.recovered_results): + if call_id in self.blocking_event_ids and call_id in self.session_tool_uses: + self.recovered_results.discard(call_id) + continue + if call_id in self.tool_uses_since_status: + continue + self.recovered_results.discard(call_id) + self.pending_results.pop(call_id, None) + self.runner.options.logger.warning("discard stale recovered tool result tool_use_id=%s", call_id) + if self.runner.options.result_store is not None: + discard = getattr(self.runner.options.result_store, "discard", None) + if not callable(discard): + continue + try: + discard(call_id) + except Exception as exc: # noqa: BLE001 - optional custom store cleanup must not stop the runner. + self.runner.options.logger.warning( + "discard persisted tool result failed tool_use_id=%s err=%s", call_id, exc + ) + self.maybe_arm_pending_idle() + def arm_idle(self) -> None: if not self.max_idle_seconds(): return diff --git a/src/arkruntime/selfhosted/tool_result_store.py b/src/arkruntime/selfhosted/tool_result_store.py index 40a38bf..6df4fe5 100644 --- a/src/arkruntime/selfhosted/tool_result_store.py +++ b/src/arkruntime/selfhosted/tool_result_store.py @@ -99,6 +99,16 @@ def mark_sent(self, call_id: str) -> None: record["state"] = STATE_SENT self._write_record(record) + def discard(self, call_id: str) -> None: + """Remove a recovered result that is no longer blocked by this session.""" + if not call_id: + raise ValueError("call id must not be empty") + try: + self._path(call_id).unlink() + except FileNotFoundError: + return + _sync_directory(self.dir) + def _read(self, call_id: str) -> dict: return self._read_path(self._path(call_id)) diff --git a/src/arkruntime/selfhosted/types.py b/src/arkruntime/selfhosted/types.py index eb5a242..e5b026d 100644 --- a/src/arkruntime/selfhosted/types.py +++ b/src/arkruntime/selfhosted/types.py @@ -26,6 +26,8 @@ EVENT_TYPE_USER_TOOL_RESULT = "user.tool_result" EVENT_TYPE_USER_CUSTOM_TOOL_RESULT = "user.custom_tool_result" EVENT_TYPE_SESSION_STATUS_IDLE = "session.status_idle" +EVENT_TYPE_SESSION_STATUS_RUNNING = "session.status_running" +EVENT_TYPE_SESSION_STATUS_RESCHEDULED = "session.status_rescheduled" EVENT_TYPE_SESSION_STATUS_TERMINATED = "session.status_terminated" EVENT_TYPE_SESSION_DELETED = "session.deleted" @@ -311,6 +313,14 @@ def stop_reason_type(self) -> str: return self.stop_reason return "" + def stop_reason_event_ids(self) -> List[str]: + if not isinstance(self.stop_reason, Mapping): + return [] + event_ids = self.stop_reason.get("event_ids") + if not isinstance(event_ids, list): + return [] + return [str(event_id) for event_id in event_ids if event_id] + def to_dict(self) -> Dict[str, Any]: out = dict(self.extra) out.update( diff --git a/tests/selfhosted/test_session_tool_runner.py b/tests/selfhosted/test_session_tool_runner.py index 000fc0e..465ccfd 100644 --- a/tests/selfhosted/test_session_tool_runner.py +++ b/tests/selfhosted/test_session_tool_runner.py @@ -9,9 +9,28 @@ import pytest from arkruntime.selfhosted import Event, ListEventsResponse, SessionToolRunner, SessionToolRunnerOptions +from arkruntime.selfhosted.tool_result_store import ToolCallStoreDecision from arkruntime.selfhosted.tools import FunctionTool, ToolContext, ToolSet, text_result +class _RecordingStore: + def __init__(self) -> None: + self.discarded = [] + self.marked = [] + + def begin(self, _call_id, _event): + return ToolCallStoreDecision() + + def save_result(self, _call_id, _result) -> None: + pass + + def mark_sent(self, call_id) -> None: + self.marked.append(call_id) + + def discard(self, call_id) -> None: + self.discarded.append(call_id) + + class _ListAPI: def __init__(self) -> None: self.calls = 0 @@ -240,3 +259,114 @@ def mark_sent(self, _call_id): assert runner._state.answered["call-1"] is True assert "call-1" not in runner._state.pending_results assert "mark tool result sent failed" in caplog.text + + +def test_recovered_results_are_filtered_against_current_blockers(tmp_path) -> None: + executions = [] + store = _RecordingStore() + tool = FunctionTool("custom", lambda _input, _context: executions.append(True) or text_result("ok")) + runner = SessionToolRunner( + object(), + "session-1", + SessionToolRunnerOptions( + tools=ToolSet(), + tool_context=ToolContext(workdir=str(tmp_path)), + custom_tools={"custom": tool}, + result_store=store, + ), + ) + sent = [] + runner._state.retry_send_event = lambda event, _call_id: sent.append(event) or True + runner._state.pending_results.update( + { + "foreign-call": Event(type="user.custom_tool_result", custom_tool_use_id="foreign-call"), + "stale-call": Event(type="user.custom_tool_result", custom_tool_use_id="stale-call"), + } + ) + runner._state.recovered_results.update(("foreign-call", "stale-call")) + + runner._state.process_listed_events( + [ + Event(id="stale-call", type="agent.custom_tool_use", name="custom", input={}), + Event(id="current-call", type="agent.custom_tool_use", name="custom", input={}), + Event( + id="idle", + type="session.status_idle", + stop_reason={"type": "requires_action", "event_ids": ["current-call"]}, + ), + ], + reconcile=True, + ) + + assert len(executions) == 1 + assert len(sent) == 1 + assert sent[0].custom_tool_use_id == "current-call" + assert set(store.discarded) == {"foreign-call", "stale-call"} + assert not runner._state.pending_results + assert not runner._state.recovered_results + + +def test_current_blocker_reuses_recovered_result_without_reexecution(tmp_path) -> None: + executions = [] + store = _RecordingStore() + tool = FunctionTool("custom", lambda _input, _context: executions.append(True) or text_result("ok")) + runner = SessionToolRunner( + object(), + "session-1", + SessionToolRunnerOptions( + tools=ToolSet(), + tool_context=ToolContext(workdir=str(tmp_path)), + custom_tools={"custom": tool}, + result_store=store, + ), + ) + sent = [] + result = Event(type="user.custom_tool_result", custom_tool_use_id="current-call") + runner._state.pending_results["current-call"] = result + runner._state.recovered_results.add("current-call") + runner._state.retry_send_event = lambda event, _call_id: sent.append(event) or True + + runner._state.process_listed_events( + [ + Event(id="current-call", type="agent.custom_tool_use", name="custom", input={}), + Event( + id="idle", + type="session.status_idle", + stop_reason={"type": "requires_action", "event_ids": ["current-call"]}, + ), + ], + reconcile=True, + ) + + assert not executions + assert sent == [result] + assert store.marked == ["current-call"] + assert not store.discarded + + +def test_recovered_result_waits_for_authoritative_status(tmp_path) -> None: + executions = [] + tool = FunctionTool("custom", lambda _input, _context: executions.append(True) or text_result("ok")) + runner = SessionToolRunner( + object(), + "session-1", + SessionToolRunnerOptions( + tools=ToolSet(), + tool_context=ToolContext(workdir=str(tmp_path)), + custom_tools={"custom": tool}, + ), + ) + sent = [] + runner._state.pending_results["current-call"] = Event( + type="user.custom_tool_result", custom_tool_use_id="current-call" + ) + runner._state.recovered_results.add("current-call") + runner._state.retry_send_event = lambda event, _call_id: sent.append(event) or True + + runner._state.process_listed_events( + [Event(id="current-call", type="agent.custom_tool_use", name="custom", input={})], reconcile=True + ) + + assert not executions + assert not sent + assert runner._state.recovered_results == {"current-call"} diff --git a/tests/selfhosted/test_tool_result_store.py b/tests/selfhosted/test_tool_result_store.py index 7d5e91f..dc92a39 100644 --- a/tests/selfhosted/test_tool_result_store.py +++ b/tests/selfhosted/test_tool_result_store.py @@ -40,3 +40,12 @@ def test_session_store_sanitizes_session_id(tmp_path) -> None: assert store.dir.parent == base assert store.dir.name.startswith("session-") + + +def test_discard_removes_recovered_record(tmp_path) -> None: + store = FileToolResultStore(str(tmp_path), "session-a") + store.begin("call-1", Event(id="call-1", type="agent.tool_use", name="bash")) + + store.discard("call-1") + + assert store.recover() == ({}, {})