From 8d6686b5464fb7cfd428a6be293527f052c4419a Mon Sep 17 00:00:00 2001 From: "ark-hand[bot]" <315378070+ark-hand[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 03:42:06 +0000 Subject: [PATCH] fix(selfHosted): align worker lifecycle behavior MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 简述 补齐 Python self-hosted worker 的工具超时配置，并对齐 Anthropic 的长连接和工具执行生命周期语义。 ## 修改前 - EnvironmentWorkerOptions 无法直接配置单次工具执行超时。 - timeout 只进入 ToolContext，非协作 custom tool 如果不主动检查 cancel_event，仍会阻塞 SessionToolRunner。 - SSE timeout 处理不清晰，误设为 None 会让静默断流无法及时退出。 - 新字段若插入 dataclass 中间会改变既有位置参数含义。 ## 修改后 - EnvironmentWorkerOptions 增加 tool_timeout_seconds，并透传 SessionToolRunner 和 ToolContext。 - 正值覆盖 ToolContext；None、0、负值保留用户 ToolContext 或 120s 默认值。 - SessionToolRunner 在 daemon thread 中执行工具，并在外层强制 deadline；超时后立即返回 error result，不等待忽略取消的工具，对齐 Anthropic anyio.fail_after 与 abandon_on_cancel 行为。 - 超时或 worker 停止时设置组合 cancel_event，配合型工具可及时释放进程和资源。 - SSE 保留 30s 读空闲 timeout；httpx 不设置整次 stream 总时长限制，静默断流仍可触发重连。 - WorkPoller 文档明确 AutoStop 是串行 iterator 语义；并发 dispatch 必须关闭并自行维护 heartbeat 和 stop。 - 新字段放在 dataclass 尾部，保持原位置参数顺序。 ## 边界 - 非协作工具的后台线程可能继续运行到自身返回，但结果会被丢弃；这是 Anthropic abandon-on-cancel 的相同行为，工具仍应响应 cancel_event 释放资源。 - 本 MR 不修改 SIGTERM 时 force-stop 行为，不实现 work 迁移、release 或 requeue。 ## 验证 - Python 全量单测：52 passed - Ruff 全量检查通过 - wheel 和 sdist 构建成功 - test/run.sh --sdk：Go、Python、Java 真实 STG worker 3/3 通过 - test/run.sh --all：真实 STG 全量默认套件通过，Docker 和直接 work-contract 用例按独立模式开关跳过 See merge request: !92 Sync-Source-Commit: 71cc0eae2668d0f5dc0fa57d1f450ef63d40d910 Hand-Written-Reason: Hand-written self-hosted worker lifecycle alignment; not produced by ark-apis generation. Release-Version: 0.4.0 --- .../selfhosted/session_tool_runner.py | 56 ++++++++++++++++++- src/arkruntime/selfhosted/tools.py | 2 + src/arkruntime/selfhosted/worker.py | 13 ++++- tests/selfhosted/test_client.py | 19 +++++++ tests/selfhosted/test_session_tool_runner.py | 43 +++++++++++++- tests/selfhosted/test_worker.py | 45 +++++++++++++++ 6 files changed, 173 insertions(+), 5 deletions(-) diff --git a/src/arkruntime/selfhosted/session_tool_runner.py b/src/arkruntime/selfhosted/session_tool_runner.py index 948dec4..cbc25b0 100644 --- a/src/arkruntime/selfhosted/session_tool_runner.py +++ b/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: diff --git a/src/arkruntime/selfhosted/tools.py b/src/arkruntime/selfhosted/tools.py index ce64b1d..d1d3121 100644 --- a/src/arkruntime/selfhosted/tools.py +++ b/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: diff --git a/src/arkruntime/selfhosted/worker.py b/src/arkruntime/selfhosted/worker.py index d1ff5fb..a549846 100644 --- a/src/arkruntime/selfhosted/worker.py +++ b/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, ) diff --git a/tests/selfhosted/test_client.py b/tests/selfhosted/test_client.py index 3619f6d..1fcce49 100644 --- a/tests/selfhosted/test_client.py +++ b/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", diff --git a/tests/selfhosted/test_session_tool_runner.py b/tests/selfhosted/test_session_tool_runner.py index 94208ac..000fc0e 100644 --- a/tests/selfhosted/test_session_tool_runner.py +++ b/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,6 +168,7 @@ 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={}) @@ -174,10 +176,45 @@ def capture(_input, context): 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): diff --git a/tests/selfhosted/test_worker.py b/tests/selfhosted/test_worker.py index 0456a3c..81798bb 100644 --- a/tests/selfhosted/test_worker.py +++ b/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)))