fix(selfHosted): align worker lifecycle behavior

src/arkruntime/selfhosted/session_tool_runner.py

@@ -17,6 +17,7 @@
     CONFIRMATION_ALLOW,
     CONFIRMATION_DENY,
     DEFAULT_MAX_IDLE_SECONDS,
+    DEFAULT_TOOL_TIMEOUT_SECONDS,
     EVENT_LIST_ORDER_ASC,
     EVENT_TYPE_AGENT_CUSTOM_TOOL_USE,
     EVENT_TYPE_AGENT_TOOL_USE,
@@ -50,6 +51,30 @@
 STREAM_QUEUE_SIZE = 256
 
 
+class _ToolCancelEvent:
+    def __init__(self, parent: Any = None) -> None:
+        self._parent = parent
+        self._local = threading.Event()
+
+    def set(self) -> None:
+        self._local.set()
+
+    def is_set(self) -> bool:
+        return self._local.is_set() or bool(self._parent and self._parent.is_set())
+
+    def wait(self, timeout: Optional[float] = None) -> bool:
+        deadline = None if timeout is None else time.monotonic() + max(timeout, 0)
+        while not self.is_set():
+            wait_for = 0.05
+            if deadline is not None:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    return self.is_set()
+                wait_for = min(wait_for, remaining)
+            self._local.wait(wait_for)
+        return True
+
+
 @dataclass
 class SessionToolRunnerOptions:
     work_id: str = ""
@@ -413,8 +438,37 @@ def permission_allows(self, event: Event, custom: bool, call_id: str) -> tuple:
 
     def execute_tool(self, event: Event, custom: bool) -> ToolResult:
         context = replace(self.runner.options.tool_context)
-        if self.runner.options.tool_timeout_seconds is not None:
+        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)
+        context.cancel_event = cancel_event
+        results: "queue.Queue[ToolResult]" = queue.Queue(maxsize=1)
+
+        def execute() -> None:
+            results.put(self._execute_tool(event, custom, context))
+
+        thread = threading.Thread(target=execute, name="ma-self-host-tool", daemon=True)
+        thread.start()
+        deadline = time.monotonic() + context.tool_timeout_seconds
+        while True:
+            if self.runner._is_stopped() or cancel_event.is_set():
+                cancel_event.set()
+                return error_result("tool execution canceled")
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                try:
+                    return results.get_nowait()
+                except queue.Empty:
+                    cancel_event.set()
+                    return error_result(f"tool execution timed out after {context.tool_timeout_seconds:g}s")
+            try:
+                return results.get(timeout=min(remaining, 0.05))
+            except queue.Empty:
+                continue
+
+    def _execute_tool(self, event: Event, custom: bool, context: ToolContext) -> ToolResult:
         if custom:
             tool = self.runner.options.custom_tools[event.name]
             try:

src/arkruntime/selfhosted/tools.py

@@ -69,6 +69,8 @@ class ToolResult:
 
 
 class Tool:
+    """Tool contract; the runner enforces timeouts and tools should release promptly on cancellation."""
+
     name: str
 
     def execute(self, tool_input: Any, context: "ToolContext") -> ToolResult:

src/arkruntime/selfhosted/worker.py

@@ -50,6 +50,12 @@ class WorkPollerOptions:
 
 
 class WorkPoller:
+    """Serial work poller with optional ownership cleanup.
+
+    ``auto_stop`` is intended for iterator-style serial processing. Callers
+    dispatching work concurrently must disable it and own heartbeat and stop.
+    """
+
     def __init__(self, api: Any, options: WorkPollerOptions) -> None:
         if api is None:
             raise ValueError("api is required")
@@ -196,6 +202,7 @@ class EnvironmentWorkerOptions:
     max_idle_seconds: Optional[float] = DEFAULT_MAX_IDLE_SECONDS
     custom_tools: Dict[str, Tool] = field(default_factory=dict)
     logger: logging.Logger = logging.getLogger("arkruntime.selfhosted.environment_worker")
+    tool_timeout_seconds: Optional[float] = None
 
 
 class EnvironmentWorker:
@@ -289,6 +296,7 @@ def _handle_item(self, work: _ClaimedWork, *, use_workdir_as_session: bool) -> N
                     custom_tools=self.options.custom_tools,
                     result_store=store,
                     max_idle_seconds=self.options.max_idle_seconds,
+                    tool_timeout_seconds=self.options.tool_timeout_seconds,
                     stop_event=work_stop,
                     logger=self.options.logger,
                 ),
@@ -381,11 +389,14 @@ def _heartbeat_loop(self, work: _ClaimedWork, stop: threading.Event, done: threa
     def _tool_context(self, workdir: str, cancel_event: Any) -> ToolContext:
         base = self.options.tool_context or ToolContext(workdir=workdir)
         env = None if base.env is None else dict(base.env)
+        tool_timeout_seconds = base.tool_timeout_seconds
+        if self.options.tool_timeout_seconds is not None and self.options.tool_timeout_seconds > 0:
+            tool_timeout_seconds = self.options.tool_timeout_seconds
         return ToolContext(
             workdir=workdir,
             env=env,
             unrestricted_paths=self.options.unrestricted_paths or base.unrestricted_paths,
-            tool_timeout_seconds=base.tool_timeout_seconds,
+            tool_timeout_seconds=tool_timeout_seconds,
             cancel_event=cancel_event,
         )
 

tests/selfhosted/test_client.py

@@ -26,6 +26,25 @@ def _ark_client(transport: httpx.MockTransport, *, max_retries: int = 0) -> Ark:
     )
 
 
+def test_selfhosted_event_stream_has_default_read_inactivity_timeout() -> None:
+    seen = []
+
+    class Events:
+        def stream(self, session_id, *, timeout):
+            seen.append((session_id, timeout))
+            return iter(())
+
+    class Sessions:
+        events = Events()
+
+    class Client:
+        sessions = Sessions()
+
+    list(ClientAPI(Client()).stream_events("session-1"))
+
+    assert seen == [("session-1", 30.0)]
+
+
 def test_poll_work_preserves_nested_session_data() -> None:
     work = {
         "id": "sesn-20260814050521-zb4l4",

tests/selfhosted/test_session_tool_runner.py

@@ -9,7 +9,7 @@
 import pytest
 
 from arkruntime.selfhosted import Event, ListEventsResponse, SessionToolRunner, SessionToolRunnerOptions
-from arkruntime.selfhosted.tools import FunctionTool, ToolContext, ToolSet
+from arkruntime.selfhosted.tools import FunctionTool, ToolContext, ToolSet, text_result
 
 
 class _ListAPI:
@@ -153,7 +153,8 @@ def test_reconcile_does_not_reset_idle_deadline_for_seen_history(tmp_path) -> No
     assert runner._state.idle_armed_at == armed_at
 
 
-def test_tool_execution_copies_context_and_preserves_configured_timeout(tmp_path) -> None:
+@pytest.mark.parametrize("override, expected", [(None, 7), (0, 7), (-1, 7), (3, 3)])
+def test_tool_execution_copies_context_and_preserves_configured_timeout(tmp_path, override, expected) -> None:
     contexts = []
 
     def capture(_input, context):
@@ -167,17 +168,53 @@ def capture(_input, context):
         SessionToolRunnerOptions(
             tools=ToolSet([FunctionTool("capture", capture)]),
             tool_context=original,
+            tool_timeout_seconds=override,
         ),
     )
     event = Event(id="tool-1", type="agent.tool_use", name="capture", tool_use_id="call-1", input={})
 
     runner._state.execute_tool(event, custom=False)
 
     assert contexts[0] is not original
-    assert contexts[0].tool_timeout_seconds == 7
+    assert contexts[0].tool_timeout_seconds == expected
     assert original.tool_timeout_seconds == 7
 
 
+@pytest.mark.parametrize("custom", [False, True])
+def test_tool_timeout_abandons_noncooperative_tool(tmp_path, custom) -> None:
+    release = threading.Event()
+    started = threading.Event()
+
+    def block(_input, _context):
+        started.set()
+        release.wait(2)
+        return text_result("late")
+
+    tool = FunctionTool("blocking", block)
+    runner = SessionToolRunner(
+        object(),
+        "session-1",
+        SessionToolRunnerOptions(
+            tools=ToolSet() if custom else ToolSet([tool]),
+            tool_context=ToolContext(workdir=str(tmp_path)),
+            custom_tools={"blocking": tool} if custom else {},
+            tool_timeout_seconds=0.02,
+        ),
+    )
+    event = Event(id="tool-1", type="agent.tool_use", name="blocking", tool_use_id="call-1", input={})
+
+    started_at = time.monotonic()
+    try:
+        result = runner._state.execute_tool(event, custom=custom)
+    finally:
+        release.set()
+
+    assert started.wait(1)
+    assert time.monotonic() - started_at < 0.5
+    assert result.is_error is True
+    assert result.content[0].text == "tool execution timed out after 0.02s"
+
+
 def test_successful_send_stays_answered_when_mark_sent_fails(tmp_path, caplog) -> None:
     class FailingStore:
         def mark_sent(self, _call_id):

tests/selfhosted/test_worker.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import logging
 import time
 
 import pytest
@@ -18,6 +19,7 @@
     WorkPoller,
     WorkPollerOptions,
 )
+from arkruntime.selfhosted.tools import ToolContext
 from arkruntime.selfhosted.types import WorkData, is_fatal_4xx
 
 
@@ -149,6 +151,49 @@ def test_poller_auto_stop_is_configurable(auto_stop, expected_stops) -> None:
     assert len(api.stops) == expected_stops
 
 
+def test_worker_tool_timeout_overrides_tool_context(tmp_path) -> None:
+    worker = EnvironmentWorker(
+        object(),
+        EnvironmentWorkerOptions(
+            workdir=str(tmp_path),
+            tool_timeout_seconds=0.02,
+        ),
+    )
+
+    context = worker._tool_context(str(tmp_path), None)
+
+    assert context.tool_timeout_seconds == 0.02
+
+
+@pytest.mark.parametrize("timeout", [None, 0, -1])
+def test_worker_nonpositive_tool_timeout_preserves_tool_context(tmp_path, timeout) -> None:
+    worker = EnvironmentWorker(
+        object(),
+        EnvironmentWorkerOptions(
+            workdir=str(tmp_path),
+            tool_context=ToolContext(workdir=str(tmp_path), tool_timeout_seconds=7),
+            tool_timeout_seconds=timeout,
+        ),
+    )
+
+    context = worker._tool_context(str(tmp_path), None)
+
+    assert context.tool_timeout_seconds == 7
+
+
+def test_worker_options_preserve_legacy_positional_order() -> None:
+    custom_tools = {"custom": object()}
+    logger = logging.getLogger("legacy-positional-worker")
+
+    options = EnvironmentWorkerOptions(
+        "env-1", "worker-1", ".", False, None, None, 60, custom_tools, logger
+    )
+
+    assert options.custom_tools is custom_tools
+    assert options.logger is logger
+    assert options.tool_timeout_seconds is None
+
+
 def test_session_id_cannot_escape_worker_root(tmp_path) -> None:
     worker = EnvironmentWorker(object(), EnvironmentWorkerOptions(workdir=str(tmp_path)))