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
6 changes: 4 additions & 2 deletions src/arkruntime/resources/environments/work.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ def heartbeat(
environment_id: str,
work_id: str,
*,
worker_id: str = "",
expected_last_heartbeat: str = "",
desired_ttl_seconds: int = 0,
extra_headers=None,
Expand All @@ -118,7 +119,7 @@ def heartbeat(
_path(environment_id, f"{quote(work_id, safe='')}/heartbeat"),
options=make_request_options(
query=query,
extra_headers=extra_headers,
extra_headers=_worker_headers(worker_id, extra_headers),
timeout=timeout,
),
cast_to=HeartbeatWorkResponse,
Expand Down Expand Up @@ -200,6 +201,7 @@ async def heartbeat(
environment_id: str,
work_id: str,
*,
worker_id: str = "",
expected_last_heartbeat: str = "",
desired_ttl_seconds: int = 0,
extra_headers=None,
Expand All @@ -216,7 +218,7 @@ async def heartbeat(
_path(environment_id, f"{quote(work_id, safe='')}/heartbeat"),
options=make_request_options(
query=query,
extra_headers=extra_headers,
extra_headers=_worker_headers(worker_id, extra_headers),
timeout=timeout,
),
cast_to=HeartbeatWorkResponse,
Expand Down
2 changes: 2 additions & 0 deletions src/arkruntime/selfhosted/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ def heartbeat_work(
environment_id: str,
work_id: str,
*,
worker_id: str = "",
expected_last_heartbeat: str,
desired_ttl_seconds: int = 30,
) -> HeartbeatResponse:
Expand All @@ -101,6 +102,7 @@ def heartbeat_work(
response = self.client.environments.work.heartbeat(
environment_id,
work_id,
worker_id=worker_id,
expected_last_heartbeat=(expected_last_heartbeat or EXPECTED_LAST_HEARTBEAT_NO_HEARTBEAT),
desired_ttl_seconds=desired_ttl_seconds,
timeout=max(1.0, min(float(desired_ttl_seconds or 30) / 2, 30.0)),
Expand Down
3 changes: 2 additions & 1 deletion src/arkruntime/selfhosted/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -332,7 +332,7 @@ def _handle_item(self, work: _ClaimedWork) -> None:
cause = heartbeat_cause["value"]
if _should_stop_item(cause):
reason = None if cause == "stop_requested" else stop_reason
# 外围关闭且 heartbeat 未给出其他原因时,需要由平台重新排队。
# 外围关闭且 heartbeat 未给出其他原因时,记录 worker 异常退出。
if cause == "" and self._stop.is_set():
reason = WorkStopReason.worker_abnormal
try:
Expand Down Expand Up @@ -360,6 +360,7 @@ def _heartbeat_loop(self, work: _ClaimedWork, stop: threading.Event, done: threa
resp = self.api.heartbeat_work(
work.environment_id,
work.id,
worker_id=self.options.worker_id,
expected_last_heartbeat=last,
desired_ttl_seconds=int(ttl),
)
Expand Down
7 changes: 0 additions & 7 deletions src/arkruntime/types/environment/env_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,13 +41,6 @@ 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 存储配置。
Expand Down
3 changes: 1 addition & 2 deletions src/arkruntime/types/environment/stop_work_body.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,5 @@ class StopWorkBody(BaseModel):
强制停止原因;仅允许与 `force=true` 一起提交,省略时按 `others` 记录。
Worker 正常完成时传 `completed`,异常退出时传 `worker_abnormal`;用户或管理员操作分别传
`user_cancelled` 或 `admin_stopped`,无法归类时传 `others`。`lease_expired` 仅由平台在租约
过期时写入,不应通过 StopWork 上报。该字段用于停止诊断,并在 Worker 上报 `completed`
且全部 tool_use 已闭合时作为清零恢复次数的信号;平台不会按停止原因排除积极恢复。
过期时写入,不应通过 StopWork 上报。该字段仅用于停止诊断;强停不会自动重新排队。
"""
10 changes: 2 additions & 8 deletions src/arkruntime/types/environment/work_item.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,14 +71,8 @@ class WorkItem(BaseModel):
"""
stop_reason: Optional[WorkStopReason] = None
"""
最近一次停止原因。首次停止前省略;state 为 stopping / stopped 时返回本次停止原因;
自动恢复到 queued / starting / active 后仍保留上一次停止原因,当前生命周期状态以
`state` 为准。
"""
recovery_count: Optional[int] = None
"""
自上次成功完成且全部 tool_use 已闭合后,平台已自动将本 Work 重新排队的次数;
首次执行或字段省略时按 0 处理。手动 Ensure 不会清零该计数。
最近一次停止原因。首次停止前省略;state 为 stopping / stopped 时返回本次停止原因。
原 owner 通过 Heartbeat 将 `stopped + lease_expired` 恢复为 active 时清除此字段。
"""
type: Literal["work"]
"""
Expand Down
6 changes: 3 additions & 3 deletions tests/selfhosted/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,6 @@ 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",
}

Expand All @@ -81,7 +80,6 @@ def handle(request: httpx.Request) -> httpx.Response:
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"]


Expand Down Expand Up @@ -113,7 +111,7 @@ def handle(request: httpx.Request) -> httpx.Response:
client.close()


def test_heartbeat_uses_generated_response_model() -> None:
def test_heartbeat_uses_generated_response_model_and_worker_identity() -> None:
heartbeat = {
"last_heartbeat": "2026-08-24T10:00:00Z",
"lease_extended": True,
Expand All @@ -123,13 +121,15 @@ def test_heartbeat_uses_generated_response_model() -> None:
}

def handle(request: httpx.Request) -> httpx.Response:
assert request.headers["Ark-Worker-ID"] == "worker-1"
return httpx.Response(httpx.codes.OK, json=heartbeat)

client = _ark_client(httpx.MockTransport(handle))
try:
response = ClientAPI(client).heartbeat_work(
"env-1",
"work-1",
worker_id="worker-1",
expected_last_heartbeat="NO_HEARTBEAT",
desired_ttl_seconds=30,
)
Expand Down
19 changes: 7 additions & 12 deletions tests/selfhosted/test_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,9 +76,7 @@ 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, "reason": WorkStopReason.others})
]
assert api.stops == [("env-1", "work-1", {"force": True, "reason": WorkStopReason.others})]


def test_poller_does_not_swallow_programming_errors() -> None:
Expand All @@ -104,8 +102,10 @@ def test_default_worker_id_is_unique_per_worker() -> None:
class _StoppingHeartbeatAPI:
def __init__(self) -> None:
self.stops = []
self.heartbeat_worker_ids = []

def heartbeat_work(self, environment_id, work_id, **kwargs):
self.heartbeat_worker_ids.append(kwargs.get("worker_id"))
return HeartbeatResponse(
last_heartbeat="2026-08-24T10:00:00Z",
state="stopping",
Expand Down Expand Up @@ -136,6 +136,7 @@ def test_heartbeat_stop_cancels_session_runner(tmp_path) -> None:

assert time.monotonic() - started < 1
assert api.stops == [("env-1", "work-1", {"force": True})]
assert api.heartbeat_worker_ids == [worker.options.worker_id]


class _ActiveHeartbeatAPI(_StoppingHeartbeatAPI):
Expand Down Expand Up @@ -164,9 +165,7 @@ def test_worker_failure_stops_work_as_abnormal(tmp_path) -> None:
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})
]
assert api.stops == [("env-1", "work-1", {"force": True, "reason": WorkStopReason.worker_abnormal})]


class _TerminatedWorkerAPI(_FailingWorkerAPI):
Expand All @@ -183,9 +182,7 @@ def test_session_termination_stops_work_as_completed(tmp_path) -> None:

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})
]
assert api.stops == [("env-1", "work-1", {"force": True, "reason": WorkStopReason.completed})]


def test_worker_close_stops_work_as_abnormal(tmp_path) -> None:
Expand All @@ -198,9 +195,7 @@ def test_worker_close_stops_work_as_abnormal(tmp_path) -> None:

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})
]
assert api.stops == [("env-1", "work-1", {"force": True, "reason": WorkStopReason.worker_abnormal})]


class _LeaseLostHeartbeatAPI(_StoppingHeartbeatAPI):
Expand Down
Loading