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

Author: ark-hand[bot]

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,13 +23,16 @@
     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,
     EVENT_TYPE_USER_TOOL_RESULT,
     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

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))
 

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(

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"}

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() == ({}, {})