From a97f95ec8fe6296e6855fc47caded0a58f2a4e20 Mon Sep 17 00:00:00 2001 From: "ark-hand[bot]" Date: Wed, 14 Oct 2026 02:48:01 +0000 Subject: [PATCH] fix(selfhosted): send worker identity on heartbeat Sync-Source-Commit: 3a50fd6cf984cc31bd147844160553bb3b0d539e Hand-Written-Reason: No Ark-APIs provenance marker; treated as a hand-written source commit. Release-Version: 0.11.0 --- src/arkruntime/resources/environments/work.py | 6 ++++-- src/arkruntime/selfhosted/client.py | 2 ++ src/arkruntime/selfhosted/worker.py | 3 ++- .../types/environment/env_config.py | 7 ------- .../types/environment/stop_work_body.py | 3 +-- src/arkruntime/types/environment/work_item.py | 10 ++-------- tests/selfhosted/test_client.py | 6 +++--- tests/selfhosted/test_worker.py | 19 +++++++------------ 8 files changed, 21 insertions(+), 35 deletions(-) diff --git a/src/arkruntime/resources/environments/work.py b/src/arkruntime/resources/environments/work.py index f311560..a711c46 100644 --- a/src/arkruntime/resources/environments/work.py +++ b/src/arkruntime/resources/environments/work.py @@ -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, @@ -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, @@ -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, @@ -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, diff --git a/src/arkruntime/selfhosted/client.py b/src/arkruntime/selfhosted/client.py index a55eabc..4893e39 100644 --- a/src/arkruntime/selfhosted/client.py +++ b/src/arkruntime/selfhosted/client.py @@ -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: @@ -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)), diff --git a/src/arkruntime/selfhosted/worker.py b/src/arkruntime/selfhosted/worker.py index 63b425c..827938c 100644 --- a/src/arkruntime/selfhosted/worker.py +++ b/src/arkruntime/selfhosted/worker.py @@ -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: @@ -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), ) diff --git a/src/arkruntime/types/environment/env_config.py b/src/arkruntime/types/environment/env_config.py index 7184ed5..a177f81 100644 --- a/src/arkruntime/types/environment/env_config.py +++ b/src/arkruntime/types/environment/env_config.py @@ -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 存储配置。 diff --git a/src/arkruntime/types/environment/stop_work_body.py b/src/arkruntime/types/environment/stop_work_body.py index 0079eb2..5f4078e 100644 --- a/src/arkruntime/types/environment/stop_work_body.py +++ b/src/arkruntime/types/environment/stop_work_body.py @@ -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 上报。该字段仅用于停止诊断；强停不会自动重新排队。 """ diff --git a/src/arkruntime/types/environment/work_item.py b/src/arkruntime/types/environment/work_item.py index 36c3875..4d42a20 100644 --- a/src/arkruntime/types/environment/work_item.py +++ b/src/arkruntime/types/environment/work_item.py @@ -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"] """ diff --git a/tests/selfhosted/test_client.py b/tests/selfhosted/test_client.py index 8b82b13..8db078c 100644 --- a/tests/selfhosted/test_client.py +++ b/tests/selfhosted/test_client.py @@ -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", } @@ -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"] @@ -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, @@ -123,6 +121,7 @@ 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)) @@ -130,6 +129,7 @@ def handle(request: httpx.Request) -> httpx.Response: response = ClientAPI(client).heartbeat_work( "env-1", "work-1", + worker_id="worker-1", expected_last_heartbeat="NO_HEARTBEAT", desired_ttl_seconds=30, ) diff --git a/tests/selfhosted/test_worker.py b/tests/selfhosted/test_worker.py index a033b56..6fcd3ef 100644 --- a/tests/selfhosted/test_worker.py +++ b/tests/selfhosted/test_worker.py @@ -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: @@ -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", @@ -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): @@ -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): @@ -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: @@ -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):