diff --git a/src/arkruntime/selfhosted/session_tool_runner.py b/src/arkruntime/selfhosted/session_tool_runner.py index f69b4ea..aa64dc2 100644 --- a/src/arkruntime/selfhosted/session_tool_runner.py +++ b/src/arkruntime/selfhosted/session_tool_runner.py @@ -8,8 +8,9 @@ import random import threading import time +from collections import deque from dataclasses import dataclass, field, replace -from typing import Any, Dict, Iterable, List, Optional, Set +from typing import Any, Deque, Dict, Iterable, List, Optional, Set from .tool_result_store import FileToolResultStore from .tools import Tool, ToolContext, ToolResult, ToolSet, error_result @@ -27,6 +28,7 @@ EVENT_TYPE_SESSION_STATUS_RUNNING, EVENT_TYPE_SESSION_STATUS_TERMINATED, EVENT_TYPE_USER_CUSTOM_TOOL_RESULT, + EVENT_TYPE_USER_INTERRUPT, EVENT_TYPE_USER_TOOL_CONFIRMATION, EVENT_TYPE_USER_TOOL_RESULT, PERMISSION_ALLOW, @@ -98,6 +100,25 @@ class SessionToolRunnerOptions: on_tool_error: Any = None +@dataclass +class _PendingToolEvent: + event: Event + custom: bool + confirmation: str = "" + + +@dataclass +class _ActiveToolExecution: + pending: _PendingToolEvent + cancel_event: _ToolCancelEvent + + +@dataclass +class _ToolExecutionResult: + pending: _PendingToolEvent + result: ToolResult + + class SessionToolRunner: def __init__(self, api: Any, session_id: str, options: SessionToolRunnerOptions) -> None: if not session_id: @@ -128,26 +149,29 @@ def _is_stopped(self) -> bool: return self._stop.is_set() or bool(self.options.stop_event and self.options.stop_event.is_set()) 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"): - try: - self._consume_stream_loop() - return self._results - except EventStreamUnsupported: - pass - except (IdleTimeout, SessionTerminated): - raise - except Exception: - if self._is_stopped(): + try: + 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"): + try: + self._consume_stream_loop() return self._results - raise - self._consume_list() - return self._results + except EventStreamUnsupported: + pass + except (IdleTimeout, SessionTerminated): + raise + except Exception: + if self._is_stopped(): + return self._results + raise + self._consume_list() + return self._results + finally: + self._state.cancel_active_execution() def _consume_stream_loop(self) -> None: backoff = STREAM_BACKOFF_START @@ -158,6 +182,7 @@ def _consume_stream_loop(self) -> None: pump.start() self._state.reconcile() while not self._is_stopped() and (pump.is_alive() or not event_queue.empty()): + self._state.drain_execution_done() self._state.flush_results() self._raise_if_idle_expired() try: @@ -204,6 +229,7 @@ def _put_stream_item(self, event_queue: "queue.Queue[object]", item: object) -> def _consume_list(self) -> None: while not self._is_stopped(): + self._state.drain_execution_done() self._state.reconcile(reconcile=False) self._state.flush_results() self._raise_if_idle_expired() @@ -212,6 +238,7 @@ def _consume_list(self) -> None: def _sleep_or_idle(self, seconds: float) -> None: deadline = time.monotonic() + max(seconds, 0) while not self._is_stopped(): + self._state.drain_execution_done() self._raise_if_idle_expired() remaining = deadline - time.monotonic() if remaining <= 0: @@ -239,6 +266,11 @@ def __init__(self, runner: SessionToolRunner) -> None: self.pending_ask: Dict[str, Event] = {} self.confirmations: Dict[str, Event] = {} self.external_tools: Dict[str, Event] = {} + self.tool_use_events: Dict[str, Event] = {} + self.scheduled: Set[str] = set() + self.execution_queue: Deque[_PendingToolEvent] = deque() + self.active_execution: Optional[_ActiveToolExecution] = None + self.execution_done: "queue.Queue[_ToolExecutionResult]" = queue.Queue() self.session_tool_uses: Set[str] = set() self.tool_uses_since_status: Set[str] = set() self.blocking_event_ids: Set[str] = set() @@ -284,6 +316,7 @@ def _reconcile_once(self, *, reconcile: bool) -> None: def process_listed_events(self, events: Iterable[Event], reconcile: bool = False) -> None: pending: List[Event] = [] pending_ids: Dict[str, bool] = {} + replayed_tool_uses: Set[str] = set() touched_idle = False last_was_end_turn = False for event in events: @@ -299,10 +332,13 @@ def process_listed_events(self, events: Iterable[Event], reconcile: bool = False ) if event.type == EVENT_TYPE_USER_TOOL_CONFIRMATION: self.record_confirmation(event) + elif event.type == EVENT_TYPE_USER_INTERRUPT: + self.handle_interrupt(event, replayed_tool_uses if reconcile else None) elif event.type in (EVENT_TYPE_USER_TOOL_RESULT, EVENT_TYPE_USER_CUSTOM_TOOL_RESULT): self.mark_answered(tool_result_call_id(event)) elif event.type in (EVENT_TYPE_AGENT_TOOL_USE, EVENT_TYPE_AGENT_CUSTOM_TOOL_USE): call_id = tool_use_call_id(event) + replayed_tool_uses.add(call_id) if call_id and not pending_ids.get(call_id): pending.append(event) pending_ids[call_id] = True @@ -318,10 +354,7 @@ def process_listed_events(self, events: Iterable[Event], reconcile: bool = False self.handle_tool_use(event, event.type == EVENT_TYPE_AGENT_CUSTOM_TOOL_USE) self.release_confirmed_tool_uses() if touched_idle and last_was_end_turn: - if self.has_unblocked_outstanding_tool(pending): - self.disarm_idle() - else: - self.arm_idle() + self.arm_idle() def note_idle_event(self, event: Event) -> None: if event.type == EVENT_TYPE_USER_TOOL_CONFIRMATION: @@ -343,6 +376,8 @@ def handle_event(self, event: Event) -> None: if event.type == EVENT_TYPE_USER_TOOL_CONFIRMATION: self.record_confirmation(event) self.release_confirmed_tool_uses() + elif event.type == EVENT_TYPE_USER_INTERRUPT: + self.handle_interrupt(event) elif event.type in (EVENT_TYPE_USER_TOOL_RESULT, EVENT_TYPE_USER_CUSTOM_TOOL_RESULT): self.mark_answered(tool_result_call_id(event)) 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: def mark_event_seen(self, event: Event) -> bool: key = event.id or tool_use_call_id(event) + if not key and event.type == EVENT_TYPE_USER_INTERRUPT: + key = f"interrupt:{event.processed_at}:{event.session_thread_id}" if not key: return True if self.seen.get(key): @@ -368,6 +405,7 @@ def mark_answered(self, call_id: str) -> None: self.recovered_results.discard(call_id) self.pending_ask.pop(call_id, None) self.external_tools.pop(call_id, None) + self.tool_use_events.pop(call_id, None) self.maybe_arm_pending_idle() def is_answered(self, call_id: str) -> bool: @@ -384,19 +422,9 @@ def release_confirmed_tool_uses(self) -> None: self.pending_ask.pop(tool_use_call_id(event), None) self.handle_tool_use(event, event.type == EVENT_TYPE_AGENT_CUSTOM_TOOL_USE) - 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) 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 - return True - return False - def handle_tool_use(self, event: Event, custom: bool) -> None: 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 call_id in self.scheduled: return pending = self.pending_results.get(call_id) if pending is not None: @@ -424,8 +452,97 @@ def handle_tool_use(self, event: Event, custom: bool) -> None: self.pending_results[call_id] = decision.result self.send_result(call_id, event, custom, "", decision.result) return - result = self.execute_tool(event, custom) - self.post_result(event, custom, call_id, result, confirmation) + self.scheduled.add(call_id) + self.execution_queue.append(_PendingToolEvent(event, custom, confirmation)) + self.start_next_tool_execution() + + def start_next_tool_execution(self) -> None: + if self.active_execution is not None: + return + while self.execution_queue: + pending = self.execution_queue.popleft() + call_id = tool_use_call_id(pending.event) + if self.is_answered(call_id): + self.scheduled.discard(call_id) + continue + cancel_event = _ToolCancelEvent(self.runner.options.tool_context.cancel_event) + self.active_execution = _ActiveToolExecution(pending, cancel_event) + + def execute(item: _PendingToolEvent = pending, cancel: _ToolCancelEvent = cancel_event) -> None: + result = self.execute_tool(item.event, item.custom, cancel) + self.execution_done.put(_ToolExecutionResult(item, result)) + + threading.Thread(target=execute, name="ma-self-host-tool-dispatch", daemon=True).start() + return + + def drain_execution_done(self) -> None: + while True: + try: + completed = self.execution_done.get_nowait() + except queue.Empty: + return + self.finish_tool_execution(completed) + + def finish_tool_execution(self, completed: _ToolExecutionResult) -> None: + call_id = tool_use_call_id(completed.pending.event) + active = self.active_execution + if active is not None and tool_use_call_id(active.pending.event) == call_id: + active.cancel_event.set() + self.active_execution = None + self.scheduled.discard(call_id) + if not self.is_answered(call_id): + self.post_result( + completed.pending.event, + completed.pending.custom, + call_id, + completed.result, + completed.pending.confirmation, + ) + self.start_next_tool_execution() + + def cancel_active_execution(self) -> None: + if self.active_execution is not None: + self.active_execution.cancel_event.set() + + def handle_interrupt(self, event: Event, eligible: Optional[Set[str]] = None) -> None: + thread_id = event.session_thread_id + + def matches(call_id: str, tool_event: Event) -> bool: + if eligible is not None and call_id not in eligible: + return False + return not thread_id or tool_event.session_thread_id == thread_id + + active = self.active_execution + if active is not None: + call_id = tool_use_call_id(active.pending.event) + if matches(call_id, active.pending.event): + active.cancel_event.set() + self.settle_interrupted_tool_use(call_id) + retained: Deque[_PendingToolEvent] = deque() + for pending in self.execution_queue: + call_id = tool_use_call_id(pending.event) + if matches(call_id, pending.event): + self.settle_interrupted_tool_use(call_id) + else: + retained.append(pending) + self.execution_queue = retained + for call_id, tool_event in list(self.tool_use_events.items()): + if not self.is_answered(call_id) and matches(call_id, tool_event): + self.settle_interrupted_tool_use(call_id) + + def settle_interrupted_tool_use(self, call_id: str) -> None: + if not call_id or self.is_answered(call_id): + return + self.scheduled.discard(call_id) + self.mark_answered(call_id) + if self.runner.options.result_store is None: + return + try: + self.runner.options.result_store.discard(call_id) + except Exception as exc: # noqa: BLE001 - interruption must not stop event consumption. + self.runner.options.logger.warning( + "discard interrupted tool result failed tool_use_id=%s err=%s", call_id, exc + ) def owns_tool(self, event: Event, custom: bool) -> bool: if custom: @@ -453,13 +570,15 @@ def permission_allows(self, event: Event, custom: bool, call_id: str) -> tuple: self.pending_ask[call_id] = event return "", False - def execute_tool(self, event: Event, custom: bool) -> ToolResult: + def execute_tool( + self, event: Event, custom: bool, execution_cancel: Optional[_ToolCancelEvent] = None + ) -> ToolResult: context = replace(self.runner.options.tool_context) if self.runner.options.tool_timeout_seconds is not None and self.runner.options.tool_timeout_seconds > 0: context.tool_timeout_seconds = self.runner.options.tool_timeout_seconds if context.tool_timeout_seconds <= 0: context.tool_timeout_seconds = DEFAULT_TOOL_TIMEOUT_SECONDS - cancel_event = _ToolCancelEvent(context.cancel_event) + cancel_event = execution_cancel or _ToolCancelEvent(context.cancel_event) context.cancel_event = cancel_event results: "queue.Queue[ToolResult]" = queue.Queue(maxsize=1) @@ -579,6 +698,8 @@ def observe_session_state(self, event: Event) -> None: call_id = tool_use_call_id(event) if call_id: self.session_tool_uses.add(call_id) + if not self.is_answered(call_id): + self.tool_use_events[call_id] = event self.tool_uses_since_status.add(call_id) return if event.type == EVENT_TYPE_SESSION_STATUS_IDLE: @@ -642,7 +763,7 @@ def maybe_arm_pending_idle(self) -> None: self.idle_armed_at = time.monotonic() def has_idle_blockers(self) -> bool: - return bool(self.pending_ask or self.pending_results or self.external_tools) + return bool(self.pending_ask or self.pending_results or self.external_tools or self.scheduled) def idle_expired(self) -> bool: max_idle = self.max_idle_seconds() diff --git a/src/arkruntime/selfhosted/tools.py b/src/arkruntime/selfhosted/tools.py index d1d3121..cdb791b 100644 --- a/src/arkruntime/selfhosted/tools.py +++ b/src/arkruntime/selfhosted/tools.py @@ -3,8 +3,10 @@ from __future__ import annotations +import base64 import glob as globlib import json +import mimetypes import os import re import signal @@ -20,6 +22,9 @@ from .types import DEFAULT_TOOL_TIMEOUT_SECONDS, ContentBlock MAX_OUTPUT_BYTES = 100000 +DEFAULT_MAX_MEDIA_FILE_BYTES = 7 << 20 +DEFAULT_READ_LINES = 2000 +MAX_READ_LINE_CHARS = 2000 MAX_SEARCH_MATCHES = 1000 _SENSITIVE_ENV_PREFIXES = ( "AIME_", @@ -84,6 +89,9 @@ class ToolContext: unrestricted_paths: bool = False tool_timeout_seconds: float = DEFAULT_TOOL_TIMEOUT_SECONDS cancel_event: Any = None + max_input_file_bytes: int = MAX_OUTPUT_BYTES + # 0 follows max_input_file_bytes; a negative value disables the media limit. + max_media_file_bytes: int = DEFAULT_MAX_MEDIA_FILE_BYTES class FunctionTool(Tool): @@ -168,13 +176,152 @@ class ReadFileTool(Tool): def execute(self, tool_input: Any, context: ToolContext) -> ToolResult: args = _as_mapping(tool_input) - path = _safe_path(context, str(args.get("path") or args.get("file") or "")) - limit = min(max(int(args.get("limit") or 20000), 0), MAX_OUTPUT_BYTES) - offset = int(args.get("offset") or 0) - with open(path, "r", encoding="utf-8", errors="replace") as f: - if offset > 0: - f.seek(offset) - return text_result(f.read(limit)) + try: + path = _safe_path(context, _read_file_path(args)) + except ValueError as exc: + return error_result(str(exc)) + view_range = args.get("view_range") + if view_range in (None, []): + view_range = None + elif not isinstance(view_range, (list, tuple)) or len(view_range) != 2: + return error_result("view_range must be [start_line, end_line]") + try: + offset = _optional_integer(args, "offset") + limit = _optional_integer(args, "limit") + if view_range is not None: + start_line = _integer_value(view_range[0], "view_range") + end_line = _integer_value(view_range[1], "view_range") + except ValueError as exc: + return error_result(str(exc)) + if view_range is not None and (offset is not None or limit is not None): + return error_result("view_range cannot be combined with offset or limit") + media = _detect_read_media(path) + if media is not None: + if view_range is not None or offset is not None or limit is not None: + return error_result("view_range, offset, and limit are only supported for text files") + size = os.path.getsize(path) + media_limit = context.max_media_file_bytes + if media_limit == 0: + media_limit = context.max_input_file_bytes + if media_limit > 0 and size > media_limit: + return error_result(f"media file too large: {size} bytes") + block_type, media_type = media + with open(path, "rb") as file: + raw = file.read(media_limit + 1) if media_limit > 0 else file.read() + if media_limit > 0 and len(raw) > media_limit: + return error_result(f"media file too large: {max(size, len(raw))} bytes") + data = base64.b64encode(raw).decode("ascii") + return ToolResult( + [ContentBlock(type=block_type, source={"type": "base64", "media_type": media_type, "data": data})] + ) + text_limit = context.max_input_file_bytes if context.max_input_file_bytes > 0 else MAX_OUTPUT_BYTES + size = os.path.getsize(path) + if text_limit > 0 and size > text_limit: + return error_result(f"file too large: {size} bytes") + raw = Path(path).read_bytes() + try: + text = raw.decode("utf-8") + except UnicodeDecodeError: + return error_result("binary file cannot be read directly") + if view_range is not None: + lines = text.split("\n") + start = max(0, start_line - 1) + if start >= len(lines): + return text_result("") + end = len(lines) + if end_line > 0: + end = min(end, end_line) + if end < start: + return error_result(f"view_range end line {end_line} is before start line {start_line}") + return text_result("\n".join(lines[start:end])) + if offset is not None and offset < 1: + return error_result(f"offset is the 1-based start line and must be >= 1, got {offset}") + return text_result(_render_read_lines(text, offset, limit)) + + +def _read_file_path(args: Mapping[str, Any]) -> str: + path = "" + for name in ("file_path", "path", "file"): + value = args.get(name) + if value is not None and not isinstance(value, str): + raise ValueError(f"{name} must be a string") + candidate = value or "" + if not candidate: + continue + if path and path != candidate: + raise ValueError("file_path, path, and file must not conflict") + path = candidate + if not path: + raise ValueError("file_path is required") + return path + + +def _integer_value(value: Any, name: str) -> int: + if isinstance(value, bool): + raise ValueError(f"{name} must be an integer") + if isinstance(value, int): + return value + if isinstance(value, str) and re.fullmatch(r"[+-]?\d+", value): + return int(value) + raise ValueError(f"{name} must be an integer") + + +def _optional_integer(args: Mapping[str, Any], name: str) -> Optional[int]: + if name not in args or args[name] is None: + return None + return _integer_value(args[name], name) + + +def _render_read_lines(text: str, offset: Optional[int], limit: Optional[int]) -> str: + lines = text.split("\n") + if lines and lines[-1] == "" and text.endswith("\n"): + lines.pop() + start = (offset or 1) - 1 + start = min(start, len(lines)) + count = limit if limit is not None and limit > 0 else DEFAULT_READ_LINES + end = min(start + count, len(lines)) + output = [] + for index in range(start, end): + line = lines[index] + if len(line) > MAX_READ_LINE_CHARS: + line = f"{line[:MAX_READ_LINE_CHARS]} [line truncated]" + output.append(f"{index + 1:6d}\t{line}\n") + if end < len(lines): + output.append(f"\n[truncated: showing lines {start + 1}-{end} of {len(lines)}]\n") + return "".join(output) + + +def _detect_read_media(path: str) -> Optional[tuple[str, str]]: + with open(path, "rb") as file: + header = file.read(512) + if header.startswith(b"\xff\xd8\xff"): + return "image", "image/jpeg" + if header.startswith(b"\x89PNG\r\n\x1a\n"): + return "image", "image/png" + if header.startswith((b"GIF87a", b"GIF89a")): + return "image", "image/gif" + if len(header) >= 12 and header[:4] == b"RIFF" and header[8:12] == b"WEBP": + return "image", "image/webp" + if header.startswith(b"%PDF-"): + return "document", "application/pdf" + if not _looks_binary(header): + return None + media_type = mimetypes.guess_type(path)[0] or "" + if media_type in ("image/jpeg", "image/png", "image/gif", "image/webp"): + return "image", media_type + if media_type == "application/pdf": + return "document", media_type + return None + + +def _looks_binary(data: bytes) -> bool: + if b"\x00" in data: + return True + try: + data.decode("utf-8") + except UnicodeDecodeError: + return True + return False class WriteFileTool(Tool): diff --git a/src/arkruntime/selfhosted/types.py b/src/arkruntime/selfhosted/types.py index e5b026d..cc35804 100644 --- a/src/arkruntime/selfhosted/types.py +++ b/src/arkruntime/selfhosted/types.py @@ -23,6 +23,7 @@ EVENT_TYPE_AGENT_TOOL_USE = "agent.tool_use" EVENT_TYPE_AGENT_CUSTOM_TOOL_USE = "agent.custom_tool_use" EVENT_TYPE_USER_TOOL_CONFIRMATION = "user.tool_confirmation" +EVENT_TYPE_USER_INTERRUPT = "user.interrupt" 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" diff --git a/tests/selfhosted/test_session_tool_runner.py b/tests/selfhosted/test_session_tool_runner.py index 465ccfd..e62ea04 100644 --- a/tests/selfhosted/test_session_tool_runner.py +++ b/tests/selfhosted/test_session_tool_runner.py @@ -61,6 +61,11 @@ def send_event(self, session_id, event): self.sent_event.set() +def _finish_execution(runner: SessionToolRunner) -> None: + completed = runner._state.execution_done.get(timeout=1) + runner._state.finish_tool_execution(completed) + + def test_list_fallback_retries_full_history_and_converts_custom_tool_error(tmp_path) -> None: api = _ListAPI() @@ -126,6 +131,7 @@ def execute(_input, _context): runner._state.retry_send_event = lambda _event, _call_id: False runner._state.release_confirmed_tool_uses() + _finish_execution(runner) runner._state.release_confirmed_tool_uses() assert len(executions) == 1 @@ -297,6 +303,7 @@ def test_recovered_results_are_filtered_against_current_blockers(tmp_path) -> No ], reconcile=True, ) + _finish_execution(runner) assert len(executions) == 1 assert len(sent) == 1 @@ -306,6 +313,297 @@ def test_recovered_results_are_filtered_against_current_blockers(tmp_path) -> No assert not runner._state.recovered_results +def test_interrupt_cancels_active_tool_without_posting_result(tmp_path) -> None: + started = threading.Event() + canceled = threading.Event() + + def block(_input, context): + started.set() + context.cancel_event.wait(1) + canceled.set() + return text_result("late") + + store = _RecordingStore() + tool = FunctionTool("blocking", block) + runner = SessionToolRunner( + object(), + "session-1", + SessionToolRunnerOptions( + tools=ToolSet(), + tool_context=ToolContext(workdir=str(tmp_path)), + custom_tools={"blocking": tool}, + result_store=store, + ), + ) + sent = [] + runner._state.retry_send_event = lambda event, _call_id: sent.append(event) or True + tool_use = Event( + id="tool-1", + type="agent.custom_tool_use", + name="blocking", + custom_tool_use_id="call-1", + session_thread_id="thread-1", + input={}, + ) + + runner._state.handle_stream_event(tool_use) + assert started.wait(1) + runner._state.handle_stream_event(Event(id="interrupt-1", type="user.interrupt", session_thread_id="thread-1")) + assert canceled.wait(1) + _finish_execution(runner) + + assert not sent + assert runner._state.is_answered("call-1") + assert store.discarded == ["call-1"] + + +def test_interrupt_only_cancels_target_thread(tmp_path) -> None: + started = threading.Event() + release = threading.Event() + + def block(_input, context): + started.set() + while not release.is_set() and not context.cancel_event.is_set(): + time.sleep(0.005) + return text_result("done") + + tool = FunctionTool("blocking", block) + runner = SessionToolRunner( + object(), + "session-1", + SessionToolRunnerOptions( + tools=ToolSet(), + tool_context=ToolContext(workdir=str(tmp_path)), + custom_tools={"blocking": tool}, + ), + ) + runner._state.retry_send_event = lambda _event, _call_id: True + runner._state.handle_stream_event( + Event( + id="tool-1", + type="agent.custom_tool_use", + name="blocking", + custom_tool_use_id="call-1", + session_thread_id="thread-a", + input={}, + ) + ) + assert started.wait(1) + + runner._state.handle_stream_event(Event(id="interrupt-other", type="user.interrupt", session_thread_id="thread-b")) + assert not runner._state.is_answered("call-1") + assert not runner._state.active_execution.cancel_event.is_set() + + runner._state.handle_interrupt(Event(type="user.interrupt", session_thread_id="thread-a")) + _finish_execution(runner) + assert runner._state.is_answered("call-1") + assert runner._state.active_execution is None + + +def test_list_replay_of_anonymous_interrupt_does_not_redispatch_or_cancel_later_tool(tmp_path) -> None: + started = threading.Event() + release = threading.Event() + executions = [] + + def block(_input, context): + executions.append(True) + started.set() + while not release.is_set() and not context.cancel_event.is_set(): + time.sleep(0.005) + return text_result("done") + + tool = FunctionTool("blocking", block) + runner = SessionToolRunner( + object(), + "session-1", + SessionToolRunnerOptions( + tools=ToolSet(), + tool_context=ToolContext(workdir=str(tmp_path)), + custom_tools={"blocking": tool}, + ), + ) + runner._state.retry_send_event = lambda _event, _call_id: True + events = [ + Event( + id="old-call", + type="agent.custom_tool_use", + name="blocking", + custom_tool_use_id="old-call", + session_thread_id="thread-1", + input={}, + ), + Event(type="user.interrupt", session_thread_id="thread-1"), + Event( + id="new-call", + type="agent.custom_tool_use", + name="blocking", + custom_tool_use_id="new-call", + session_thread_id="thread-1", + input={}, + ), + ] + + runner._state.process_listed_events(events, reconcile=False) + assert started.wait(1) + runner._state.process_listed_events(events, reconcile=False) + runner._state.process_listed_events(events, reconcile=True) + + assert runner._state.is_answered("old-call") + assert not runner._state.is_answered("new-call") + assert len(executions) == 1 + assert not runner._state.active_execution.cancel_event.is_set() + assert "old-call" not in runner._state.tool_use_events + + runner._state.handle_interrupt(Event(type="user.interrupt")) + _finish_execution(runner) + + +def test_list_interrupt_cancels_tool_from_earlier_poll(tmp_path) -> None: + started = threading.Event() + canceled = threading.Event() + + def block(_input, context): + started.set() + while not context.cancel_event.is_set(): + time.sleep(0.005) + canceled.set() + return text_result("late") + + runner = SessionToolRunner( + object(), + "session-1", + SessionToolRunnerOptions( + tools=ToolSet(), + tool_context=ToolContext(workdir=str(tmp_path)), + custom_tools={"blocking": FunctionTool("blocking", block)}, + ), + ) + runner._state.retry_send_event = lambda _event, _call_id: True + tool_use = Event( + id="cross-poll-call", + type="agent.custom_tool_use", + name="blocking", + custom_tool_use_id="cross-poll-call", + session_thread_id="thread-1", + input={}, + ) + runner._state.process_listed_events([tool_use], reconcile=False) + assert started.wait(1) + + runner._state.process_listed_events( + [ + tool_use, + Event( + type="user.interrupt", + processed_at="2026-09-20T00:00:00Z", + session_thread_id="thread-1", + ), + ], + reconcile=False, + ) + + assert canceled.wait(1) + _finish_execution(runner) + assert runner._state.is_answered("cross-poll-call") + + +def test_serial_queue_interrupt_removes_matching_tool_only(tmp_path) -> None: + active_started = threading.Event() + active_canceled = threading.Event() + queued_started = threading.Event() + + def active_tool(_input, context): + active_started.set() + while not context.cancel_event.is_set(): + time.sleep(0.005) + active_canceled.set() + return text_result("late") + + def queued_tool(_input, _context): + queued_started.set() + return text_result("unexpected") + + runner = SessionToolRunner( + object(), + "session-1", + SessionToolRunnerOptions( + tools=ToolSet(), + tool_context=ToolContext(workdir=str(tmp_path)), + custom_tools={ + "active": FunctionTool("active", active_tool), + "queued": FunctionTool("queued", queued_tool), + }, + ), + ) + runner._state.retry_send_event = lambda _event, _call_id: True + runner._state.process_listed_events( + [ + Event( + id="active-call", + type="agent.custom_tool_use", + name="active", + custom_tool_use_id="active-call", + session_thread_id="thread-a", + input={}, + ), + Event( + id="queued-call", + type="agent.custom_tool_use", + name="queued", + custom_tool_use_id="queued-call", + session_thread_id="thread-b", + input={}, + ), + ], + reconcile=False, + ) + assert active_started.wait(1) + assert not queued_started.wait(0.05) + + runner._state.handle_interrupt(Event(type="user.interrupt", session_thread_id="thread-b")) + + assert runner._state.is_answered("queued-call") + assert not runner._state.is_answered("active-call") + assert not active_canceled.wait(0.05) + + runner._state.handle_interrupt(Event(type="user.interrupt")) + _finish_execution(runner) + assert not queued_started.is_set() + + +def test_list_end_turn_arms_idle_after_tool_completes(tmp_path) -> None: + runner = SessionToolRunner( + object(), + "session-1", + SessionToolRunnerOptions( + tools=ToolSet(), + tool_context=ToolContext(workdir=str(tmp_path)), + custom_tools={"custom": FunctionTool("custom", lambda _input, _context: text_result("ok"))}, + max_idle_seconds=60, + ), + ) + runner._state.retry_send_event = lambda _event, _call_id: True + tool_use = Event( + id="idle-call", + type="agent.custom_tool_use", + name="custom", + custom_tool_use_id="idle-call", + input={}, + ) + runner._state.process_listed_events([tool_use], reconcile=False) + runner._state.process_listed_events( + [tool_use, Event(id="idle-end-turn", type="session.status_idle", stop_reason={"type": "end_turn"})], + reconcile=False, + ) + assert runner._state.idle_arm_pending + assert not runner._state.idle_armed_at + + _finish_execution(runner) + + assert not runner._state.idle_arm_pending + assert runner._state.idle_armed_at + + def test_current_blocker_reuses_recovered_result_without_reexecution(tmp_path) -> None: executions = [] store = _RecordingStore() diff --git a/tests/selfhosted/test_tools.py b/tests/selfhosted/test_tools.py index de8e5fa..db1b4c2 100644 --- a/tests/selfhosted/test_tools.py +++ b/tests/selfhosted/test_tools.py @@ -3,6 +3,7 @@ from __future__ import annotations +import base64 import threading import time @@ -12,6 +13,7 @@ BashTool, EditFileTool, GrepTool, + ReadFileTool, ToolContext, WriteFileTool, _is_sensitive_env_key, @@ -70,6 +72,99 @@ def test_bash_honors_worker_cancellation(tmp_path) -> None: assert "canceled" in _text(result) +@pytest.mark.parametrize( + "name, data, block_type, media_type", + [ + ("photo.jpg", b"\xff\xd8\xff\xe0data", "image", "image/jpeg"), + ("image.png", b"\x89PNG\r\n\x1a\ndata", "image", "image/png"), + ("image.gif", b"GIF89adata", "image", "image/gif"), + ("image.webp", b"RIFF\x04\x00\x00\x00WEBPdata", "image", "image/webp"), + ("document.pdf", b"%PDF-1.7\ndata", "document", "application/pdf"), + ], +) +def test_read_returns_native_media_block(tmp_path, name, data, block_type, media_type) -> None: + (tmp_path / name).write_bytes(data) + + result = ReadFileTool().execute({"file_path": name}, ToolContext(workdir=str(tmp_path))) + + assert not result.is_error + block = list(result.content)[0] + assert block.type == block_type + assert block.source == { + "type": "base64", + "media_type": media_type, + "data": base64.b64encode(data).decode("ascii"), + } + + +def test_read_rejects_ranges_and_oversized_media(tmp_path) -> None: + data = b"\x89PNG\r\n\x1a\ndata" + (tmp_path / "image.png").write_bytes(data) + tool = ReadFileTool() + + ranged = tool.execute({"path": "image.png", "limit": 1}, ToolContext(workdir=str(tmp_path))) + invalid = tool.execute({"path": "image.png", "limit": "10KB"}, ToolContext(workdir=str(tmp_path))) + oversized = tool.execute( + {"path": "image.png"}, ToolContext(workdir=str(tmp_path), max_media_file_bytes=len(data) - 1) + ) + + assert ranged.is_error + assert _text(ranged) == "view_range, offset, and limit are only supported for text files" + assert invalid.is_error + assert _text(invalid) == "limit must be an integer" + assert oversized.is_error + assert _text(oversized) == f"media file too large: {len(data)} bytes" + + +def test_read_uses_managed_agent_line_range(tmp_path) -> None: + (tmp_path / "example.txt").write_text("a\nb\nc\nd\n") + tool = ReadFileTool() + context = ToolContext(workdir=str(tmp_path)) + + whole = tool.execute({"file_path": "example.txt"}, context) + ranged = tool.execute({"file_path": "example.txt", "offset": 2, "limit": 2}, context) + invalid = tool.execute({"file_path": "example.txt", "offset": 0}, context) + + assert not whole.is_error + assert _text(whole) == " 1\ta\n 2\tb\n 3\tc\n 4\td\n" + assert not ranged.is_error + assert _text(ranged) == " 2\tb\n 3\tc\n\n[truncated: showing lines 2-3 of 4]\n" + assert invalid.is_error + assert _text(invalid) == "offset is the 1-based start line and must be >= 1, got 0" + + +def test_read_marks_truncated_lines(tmp_path) -> None: + (tmp_path / "long.txt").write_text("a" * 2001) + + result = ReadFileTool().execute({"file_path": "long.txt"}, ToolContext(workdir=str(tmp_path))) + + assert not result.is_error + assert _text(result) == f" 1\t{'a' * 2000} [line truncated]\n" + + +def test_read_keeps_legacy_inputs_without_ambiguity(tmp_path) -> None: + (tmp_path / "example.txt").write_text("a\nb\nc\n") + tool = ReadFileTool() + context = ToolContext(workdir=str(tmp_path)) + + for tool_input in ( + {"path": "example.txt"}, + {"file": "example.txt"}, + {"file_path": "example.txt", "path": "example.txt"}, + ): + assert not tool.execute(tool_input, context).is_error + legacy_range = tool.execute({"file_path": "example.txt", "view_range": [2, 3]}, context) + path_conflict = tool.execute({"file_path": "example.txt", "path": "other.txt"}, context) + range_conflict = tool.execute({"file_path": "example.txt", "view_range": [1, 1], "offset": 1}, context) + + assert not legacy_range.is_error + assert _text(legacy_range) == "b\nc" + assert path_conflict.is_error + assert _text(path_conflict) == "file_path, path, and file must not conflict" + assert range_conflict.is_error + assert _text(range_conflict) == "view_range cannot be combined with offset or limit" + + def test_grep_skips_symlink_that_escapes_workdir(tmp_path) -> None: outside = tmp_path.parent / f"{tmp_path.name}-outside" outside.mkdir()