Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 71 additions & 4 deletions src/arkruntime/selfhosted/session_tool_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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"):
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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 = (
Expand All @@ -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()
Expand All @@ -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)

Expand Down Expand Up @@ -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()
Expand All @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down
10 changes: 10 additions & 0 deletions src/arkruntime/selfhosted/tool_result_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))

Expand Down
10 changes: 10 additions & 0 deletions src/arkruntime/selfhosted/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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(
Expand Down
130 changes: 130 additions & 0 deletions tests/selfhosted/test_session_tool_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"}
9 changes: 9 additions & 0 deletions tests/selfhosted/test_tool_result_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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() == ({}, {})
Loading