diff --git a/src/arkruntime/resources/environments/work.py b/src/arkruntime/resources/environments/work.py index 961b5e8..f311560 100644 --- a/src/arkruntime/resources/environments/work.py +++ b/src/arkruntime/resources/environments/work.py @@ -14,6 +14,7 @@ from ...types.environment.heartbeat_work_response import HeartbeatWorkResponse from ...types.environment.stop_work_body import StopWorkBody from ...types.environment.work_item import WorkItem +from ...types.environment.work_stop_reason import WorkStopReason __all__ = ["EnvironmentWork", "AsyncEnvironmentWork"] @@ -42,6 +43,11 @@ def _worker_headers(worker_id: str, extra_headers: Optional[dict]) -> dict: return headers +def _validate_stop_reason(force: bool, reason: Optional[WorkStopReason]) -> None: + if reason is not None and not force: + raise ValueError("reason requires force=True") + + class EnvironmentWork(SyncAPIResource): def poll( self, @@ -124,12 +130,14 @@ def stop( work_id: str, *, force: bool = False, + reason: Optional[WorkStopReason] = None, extra_headers=None, timeout: float | httpx.Timeout | None = None, ) -> WorkItem: if not work_id: raise ValueError("work_id is required") - body = StopWorkBody(force=True) if force else StopWorkBody() + _validate_stop_reason(force, reason) + body = StopWorkBody(force=True if force else None, reason=reason) return self._post_without_retry( _path(environment_id, f"{quote(work_id, safe='')}/stop"), body=dump_body(body.model_dump(exclude_none=True, by_alias=True)), @@ -220,12 +228,14 @@ async def stop( work_id: str, *, force: bool = False, + reason: Optional[WorkStopReason] = None, extra_headers=None, timeout: float | httpx.Timeout | None = None, ) -> WorkItem: if not work_id: raise ValueError("work_id is required") - body = StopWorkBody(force=True) if force else StopWorkBody() + _validate_stop_reason(force, reason) + body = StopWorkBody(force=True if force else None, reason=reason) return await self._post_without_retry( _path(environment_id, f"{quote(work_id, safe='')}/stop"), body=dump_body(body.model_dump(exclude_none=True, by_alias=True)), diff --git a/src/arkruntime/selfhosted/client.py b/src/arkruntime/selfhosted/client.py index f5a71b6..a55eabc 100644 --- a/src/arkruntime/selfhosted/client.py +++ b/src/arkruntime/selfhosted/client.py @@ -18,6 +18,7 @@ from arkruntime import Ark from arkruntime._constants import CLIENT_REQUEST_HEADER, SERVER_REQUEST_HEADER from arkruntime._types import NOT_GIVEN +from arkruntime.types.environment import WorkStopReason from .types import ( EXPECTED_LAST_HEARTBEAT_NO_HEARTBEAT, @@ -106,7 +107,14 @@ def heartbeat_work( ) return response - def stop_work(self, environment_id: str, work_id: str, *, force: bool = False) -> None: + def stop_work( + self, + environment_id: str, + work_id: str, + *, + force: bool = False, + reason: Optional[WorkStopReason] = None, + ) -> None: if not environment_id: raise ValueError("environment_id is required") if not work_id: @@ -115,6 +123,7 @@ def stop_work(self, environment_id: str, work_id: str, *, force: bool = False) - environment_id, work_id, force=force, + reason=reason, timeout=10.0, ) diff --git a/src/arkruntime/selfhosted/worker.py b/src/arkruntime/selfhosted/worker.py index 3c444c4..63b425c 100644 --- a/src/arkruntime/selfhosted/worker.py +++ b/src/arkruntime/selfhosted/worker.py @@ -17,6 +17,7 @@ import httpx from .._exceptions import ArkAPIError +from ..types.environment import WorkStopReason from .envinit import Initializer, InitializerOptions from .session_tool_runner import SessionToolRunner, SessionToolRunnerOptions from .tool_result_store import FileToolResultStore @@ -121,7 +122,7 @@ def next(self) -> Optional[WorkItem]: except _POLLER_TRANSIENT_ERRORS as exc: self.options.logger.warning("ack work failed work_id=%s err=%s", item.id, exc) if _is_poller_fatal_4xx(exc): - self._stop_item(item, force=True) + self._stop_item(item, force=True, reason=WorkStopReason.others) continue self._backoff_discard() continue @@ -150,12 +151,21 @@ def _discard_invalid_work(self, item: WorkItem) -> None: except Exception as exc: # noqa: BLE001 - invalid work still obeys ACK ownership. self.options.logger.warning("ack invalid work failed work_id=%s err=%s", item.id, exc) return - self._stop_item(item, force=True) + self._stop_item(item, force=True, reason=WorkStopReason.others) self._backoff_discard() - def _stop_item(self, item: WorkItem, *, force: bool) -> None: + def _stop_item( + self, + item: WorkItem, + *, + force: bool, + reason: Optional[WorkStopReason] = None, + ) -> None: try: - self.api.stop_work(item.environment_id, item.id, force=force) + if reason is None: + self.api.stop_work(item.environment_id, item.id, force=force) + else: + self.api.stop_work(item.environment_id, item.id, force=force, reason=reason) except Exception as exc: # noqa: BLE001 if not _is_resolved_status(exc): self.options.logger.warning("stop work failed work_id=%s err=%s", item.id, exc) @@ -262,6 +272,7 @@ def _handle_item(self, work: _ClaimedWork) -> None: heartbeat_cause = {"value": ""} heartbeat = None initializer = None + stop_reason = WorkStopReason.completed try: workdir = self._workdir() heartbeat_thread = threading.Thread( @@ -303,6 +314,12 @@ def _handle_item(self, work: _ClaimedWork) -> None: ), ) runner.run() + except (IdleTimeout, SessionTerminated): + raise + except BaseException: + # 进程级退出也需要触发恢复，不能被误记为正常完成。 + stop_reason = WorkStopReason.worker_abnormal + raise finally: if initializer is not None: try: @@ -314,8 +331,15 @@ def _handle_item(self, work: _ClaimedWork) -> None: heartbeat_done.wait(timeout=DEFAULT_HEARTBEAT_SECONDS + 1) cause = heartbeat_cause["value"] if _should_stop_item(cause): + reason = None if cause == "stop_requested" else stop_reason + # 外围关闭且 heartbeat 未给出其他原因时，需要由平台重新排队。 + if cause == "" and self._stop.is_set(): + reason = WorkStopReason.worker_abnormal try: - self.api.stop_work(work.environment_id, work.id, force=True) + if reason is None: + self.api.stop_work(work.environment_id, work.id, force=True) + else: + self.api.stop_work(work.environment_id, work.id, force=True, reason=reason) except Exception as exc: # noqa: BLE001 if not _is_resolved_status(exc): self.options.logger.warning("stop work failed: %s", exc) diff --git a/src/arkruntime/types/environment/__init__.py b/src/arkruntime/types/environment/__init__.py index c0c9def..60b6560 100644 --- a/src/arkruntime/types/environment/__init__.py +++ b/src/arkruntime/types/environment/__init__.py @@ -26,6 +26,7 @@ from .work_data import WorkData from .work_item import WorkItem from .work_state import WorkState +from .work_stop_reason import WorkStopReason __all__ = [ "CreateEnvironmentRequest", @@ -48,6 +49,7 @@ "WorkData", "WorkItem", "WorkState", + "WorkStopReason", ] # Hand-written extras (preserved across regen via Makefile rsync --exclude=*_shim.py). diff --git a/src/arkruntime/types/environment/env_config.py b/src/arkruntime/types/environment/env_config.py index a177f81..7184ed5 100644 --- a/src/arkruntime/types/environment/env_config.py +++ b/src/arkruntime/types/environment/env_config.py @@ -41,6 +41,13 @@ class EnvConfig(BaseModel): """ 沙箱启动阶段执行的初始化脚本。 """ + active_recovery: Optional[bool] = None + """ + 是否启用 Work 积极恢复；仅 self_hosted Environment 支持，省略时等同于 false。该配置 + 仅允许在 Environment 级设置，不允许通过 Session override/upgrade 覆盖。开启后，平台在 + Work 已 stopped、仍有未闭合 agent.tool_use 且 recovery_count 少于 5 次时将 Work 重新排队， + 不按 stop_reason 排除。只有 Worker 上报 completed 且全部 tool_use 已闭合后才会清零计数。 + """ tos: Optional[TosConfig] = None """ Environment outputs 的 TOS 存储配置。 diff --git a/src/arkruntime/types/environment/stop_work_body.py b/src/arkruntime/types/environment/stop_work_body.py index 8f8693b..0079eb2 100644 --- a/src/arkruntime/types/environment/stop_work_body.py +++ b/src/arkruntime/types/environment/stop_work_body.py @@ -10,6 +10,8 @@ from arkruntime._models import BaseModel +from .work_stop_reason import WorkStopReason + class StopWorkBody(BaseModel): """ @@ -20,3 +22,11 @@ class StopWorkBody(BaseModel): """ 是否强制停止。 """ + reason: Optional[WorkStopReason] = None + """ + 强制停止原因；仅允许与 `force=true` 一起提交，省略时按 `others` 记录。 + Worker 正常完成时传 `completed`，异常退出时传 `worker_abnormal`；用户或管理员操作分别传 + `user_cancelled` 或 `admin_stopped`，无法归类时传 `others`。`lease_expired` 仅由平台在租约 + 过期时写入，不应通过 StopWork 上报。该字段用于停止诊断，并在 Worker 上报 `completed` + 且全部 tool_use 已闭合时作为清零恢复次数的信号；平台不会按停止原因排除积极恢复。 + """ diff --git a/src/arkruntime/types/environment/work_item.py b/src/arkruntime/types/environment/work_item.py index 23690d8..36c3875 100644 --- a/src/arkruntime/types/environment/work_item.py +++ b/src/arkruntime/types/environment/work_item.py @@ -13,6 +13,7 @@ from .volc_tag import VolcTag from .work_data import WorkData from .work_state import WorkState +from .work_stop_reason import WorkStopReason class WorkItem(BaseModel): @@ -68,6 +69,17 @@ class WorkItem(BaseModel): """ Work 停止时间，RFC 3339。 """ + stop_reason: Optional[WorkStopReason] = None + """ + 最近一次停止原因。首次停止前省略；state 为 stopping / stopped 时返回本次停止原因； + 自动恢复到 queued / starting / active 后仍保留上一次停止原因，当前生命周期状态以 + `state` 为准。 + """ + recovery_count: Optional[int] = None + """ + 自上次成功完成且全部 tool_use 已闭合后，平台已自动将本 Work 重新排队的次数； + 首次执行或字段省略时按 0 处理。手动 Ensure 不会清零该计数。 + """ type: Literal["work"] """ 对象类型，固定为 `work`。 diff --git a/src/arkruntime/types/environment/work_stop_reason.py b/src/arkruntime/types/environment/work_stop_reason.py new file mode 100644 index 0000000..3c3b6e4 --- /dev/null +++ b/src/arkruntime/types/environment/work_stop_reason.py @@ -0,0 +1,22 @@ +# Generated by datamodel-code-generator from ark-apis typespec/openapi. +# DO NOT EDIT — regenerate with `make gen-py-` in ark-apis. +# +# Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from enum import Enum + + +class WorkStopReason(str, Enum): + """ + Work 停止原因。 + """ + + lease_expired = "lease_expired" + worker_abnormal = "worker_abnormal" + user_cancelled = "user_cancelled" + admin_stopped = "admin_stopped" + completed = "completed" + others = "others" diff --git a/tests/selfhosted/test_client.py b/tests/selfhosted/test_client.py index e63c1bd..8b82b13 100644 --- a/tests/selfhosted/test_client.py +++ b/tests/selfhosted/test_client.py @@ -3,13 +3,17 @@ from __future__ import annotations +import asyncio +import json import math import httpx +import pytest -from arkruntime import Ark +from arkruntime import Ark, AsyncArk from arkruntime.selfhosted import APIError, ClientAPI, SkillRef from arkruntime.selfhosted.types import work_session_id +from arkruntime.types.environment import WorkStopReason def _test_credential() -> str: @@ -57,6 +61,8 @@ def test_poll_work_preserves_nested_session_data() -> None: }, "environment_id": "env-20260813132936-wq8d4", "state": "queued", + "stop_reason": "worker_abnormal", + "recovery_count": 5, "type": "work", } @@ -74,6 +80,8 @@ def handle(request: httpx.Request) -> httpx.Response: assert item.id == work["id"] assert item.data.id == work["data"]["id"] assert item.data.type == "session" + assert item.stop_reason == WorkStopReason.worker_abnormal + assert item.recovery_count == 5 assert work_session_id(item) == work["data"]["id"] @@ -151,6 +159,7 @@ def test_atomic_environment_work_api_matches_openapi_contract() -> None: "type": "work_heartbeat", } calls = [] + stop_bodies = [] def handle(request: httpx.Request) -> httpx.Response: calls.append(request) @@ -171,7 +180,7 @@ def handle(request: httpx.Request) -> httpx.Response: return httpx.Response(httpx.codes.OK, json=heartbeat) assert request.url.path.endswith("/stop") assert request.method == "POST" - assert request.content == b"{}" + stop_bodies.append(request.content) return httpx.Response(httpx.codes.OK, json=work) client = _ark_client(httpx.MockTransport(handle)) @@ -197,6 +206,21 @@ def handle(request: httpx.Request) -> httpx.Response: is True ) assert resource.stop("env-1", "work-1").id == "work-1" + assert ( + resource.stop( + "env-1", + "work-1", + force=True, + reason=WorkStopReason.worker_abnormal, + ).id + == "work-1" + ) + ClientAPI(client).stop_work( + "env-1", + "work-1", + force=True, + reason=WorkStopReason.worker_abnormal, + ) finally: client.close() @@ -205,7 +229,57 @@ def handle(request: httpx.Request) -> httpx.Response: "/api/v3/environments/env-1/work/work-1/ack", "/api/v3/environments/env-1/work/work-1/heartbeat", "/api/v3/environments/env-1/work/work-1/stop", + "/api/v3/environments/env-1/work/work-1/stop", + "/api/v3/environments/env-1/work/work-1/stop", ] + assert [json.loads(body) for body in stop_bodies] == [ + {}, + {"force": True, "reason": "worker_abnormal"}, + {"force": True, "reason": "worker_abnormal"}, + ] + + +def test_stop_reason_requires_force_without_sending_request() -> None: + calls = 0 + + def handle(_: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + raise AssertionError("invalid stop request must not be sent") + + client = _ark_client(httpx.MockTransport(handle)) + try: + with pytest.raises(ValueError, match="reason requires force=True"): + client.environments.work.stop( + "env-1", + "work-1", + reason=WorkStopReason.worker_abnormal, + ) + with pytest.raises(ValueError, match="reason requires force=True"): + ClientAPI(client).stop_work( + "env-1", + "work-1", + reason=WorkStopReason.worker_abnormal, + ) + finally: + client.close() + + async def check_async_stop() -> None: + async with AsyncArk( + api_key=_test_credential(), + base_url="https://ark.example.com/api/v3", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handle)), + max_retries=0, + ) as async_client: + with pytest.raises(ValueError, match="reason requires force=True"): + await async_client.environments.work.stop( + "env-1", + "work-1", + reason=WorkStopReason.worker_abnormal, + ) + + asyncio.run(check_async_stop()) + assert calls == 0 def test_list_events_decodes_wire_data_into_tool_use_events() -> None: diff --git a/tests/selfhosted/test_worker.py b/tests/selfhosted/test_worker.py index 27b69a8..a033b56 100644 --- a/tests/selfhosted/test_worker.py +++ b/tests/selfhosted/test_worker.py @@ -15,6 +15,7 @@ HandleItemOptions, HeartbeatResponse, Session, + SessionTerminated, WorkItem, WorkPoller, WorkPollerOptions, @@ -22,6 +23,7 @@ from arkruntime.selfhosted.tools import ToolContext from arkruntime.selfhosted.types import WorkData, is_fatal_4xx from arkruntime.selfhosted.worker import default_worker_id +from arkruntime.types.environment import WorkStopReason def _work_item(environment_id: str) -> WorkItem: @@ -74,7 +76,9 @@ def test_fatal_ack_error_force_stops_work_and_continues() -> None: assert poller.next() is None assert poller.error is None - assert api.stops == [("env-1", "work-1", {"force": True})] + assert api.stops == [ + ("env-1", "work-1", {"force": True, "reason": WorkStopReason.others}) + ] def test_poller_does_not_swallow_programming_errors() -> None: @@ -134,6 +138,71 @@ def test_heartbeat_stop_cancels_session_runner(tmp_path) -> None: assert api.stops == [("env-1", "work-1", {"force": True})] +class _ActiveHeartbeatAPI(_StoppingHeartbeatAPI): + def heartbeat_work(self, environment_id, work_id, **kwargs): + return HeartbeatResponse( + last_heartbeat="2026-08-24T10:00:00Z", + state="active", + lease_extended=True, + ttl_seconds=30, + type="work_heartbeat", + ) + + +class _FailingWorkerAPI(_ActiveHeartbeatAPI): + def get_session(self, session_id): + raise ValueError("worker failed") + + +def test_worker_failure_stops_work_as_abnormal(tmp_path) -> None: + api = _FailingWorkerAPI() + worker = EnvironmentWorker( + api, + EnvironmentWorkerOptions(environment_id="env-1", workdir=str(tmp_path)), + ) + + with pytest.raises(ValueError, match="worker failed"): + worker.handle_item(HandleItemOptions(work_id="work-1", environment_id="env-1", session_id="session-1")) + + assert api.stops == [ + ("env-1", "work-1", {"force": True, "reason": WorkStopReason.worker_abnormal}) + ] + + +class _TerminatedWorkerAPI(_FailingWorkerAPI): + def get_session(self, session_id): + raise SessionTerminated("session terminated") + + +def test_session_termination_stops_work_as_completed(tmp_path) -> None: + api = _TerminatedWorkerAPI() + worker = EnvironmentWorker( + api, + EnvironmentWorkerOptions(environment_id="env-1", workdir=str(tmp_path)), + ) + + worker.handle_item(HandleItemOptions(work_id="work-1", environment_id="env-1", session_id="session-1")) + + assert api.stops == [ + ("env-1", "work-1", {"force": True, "reason": WorkStopReason.completed}) + ] + + +def test_worker_close_stops_work_as_abnormal(tmp_path) -> None: + api = _ActiveHeartbeatAPI() + worker = EnvironmentWorker( + api, + EnvironmentWorkerOptions(environment_id="env-1", workdir=str(tmp_path)), + ) + worker.close() + + worker.handle_item(HandleItemOptions(work_id="work-1", environment_id="env-1", session_id="session-1")) + + assert api.stops == [ + ("env-1", "work-1", {"force": True, "reason": WorkStopReason.worker_abnormal}) + ] + + class _LeaseLostHeartbeatAPI(_StoppingHeartbeatAPI): def heartbeat_work(self, environment_id, work_id, **kwargs): raise APIError(412, "lease lost")