From ccb22c960db832e4131cb00380b389b9d218faf1 Mon Sep 17 00:00:00 2001 From: host Date: Thu, 29 Oct 2026 04:40:47 +0800 Subject: [PATCH] worker: heartbeat parent task during marker poll/wait loops Marker waits run after _run_attempt has cleared _running, so the run_loop heartbeat thread no longer refreshes the waiting parent. sweep_stale_tasks (default 300s) then kills a parent waiting up to the 3600s marker timeout, as seen in the marker live test (PM2 swept mid wait-threads). Poll and wait loops now heartbeat the parent each iteration (interval clamped to <=60s, well under the 300s threshold), best-effort, without touching _running, so no new capacity leak and cancel still works (empty _running entry has no proc/container to kill). Tests: 2 new sweep-survival tests (stale heartbeat at -299s, 2s wait, sweep must spare the parent); both fail 2/2 on pre-fix code and pass with the fix. --- cloud/tests/test_marker_wait_heartbeat.py | 118 ++++++++++++++++++++++ cloud/worker.py | 15 +++ 2 files changed, 133 insertions(+) create mode 100644 cloud/tests/test_marker_wait_heartbeat.py diff --git a/cloud/tests/test_marker_wait_heartbeat.py b/cloud/tests/test_marker_wait_heartbeat.py new file mode 100644 index 0000000..27be938 --- /dev/null +++ b/cloud/tests/test_marker_wait_heartbeat.py @@ -0,0 +1,118 @@ +"""标记等待中父任务心跳单测：sweep 误杀回归。 + +根因：标记等待在 _run_attempt 返回后才跑，那时 _running 已被清掉， +run_loop 的心跳线程刷不到等候中的父任务；sweep 默认阈值 300 秒， +而标记等待标称可到 3600 秒，父任务等超 300 秒就被当死任务回收。 +修法：轮询 / 等待循环每次迭代自己刷一次心跳（间隔至多 interval， +钳制后 <=60 秒，远小于 300 秒阈值），不依赖心跳线程。 + +不跑真实模型、不联网：用 init_db 建的临时 SQLite 库 + 最小任务行， +直接调 _deliver_marker_tools，不走 run_one 全量流程。 +""" + +import os +import sqlite3 +import sys +import tempfile +import time +import unittest +from pathlib import Path + +try: + from cloud.server import init_db, sweep_stale_tasks + from cloud.worker import Worker +except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from cloud.server import init_db, sweep_stale_tasks + from cloud.worker import Worker + +FAKE_ENV = {"XCM_API_KEY": "test-placeholder"} + + +class MarkerWaitHeartbeatTestCase(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.db_path = os.path.join(self.tmp.name, "control.db") + init_db(self.db_path) + # 父任务 PM：running，但心跳已陈旧到离 sweep 阈值只差 1 秒。 + # 不修时，等 2 秒后心跳就过期（301 > 300），sweep 会误杀； + # 修了时，等待循环每次迭代刷心跳，sweep 应放过它。 + stale = time.time() - 299 + self._insert("PM", status="running", last_heartbeat=stale, + updated_at=stale) + self._insert("E1", status="queued", parent="PM") + + def tearDown(self): + self.tmp.cleanup() + + def _insert(self, task_id, status="queued", parent=None, + last_heartbeat=None, updated_at=None): + now = time.time() + conn = sqlite3.connect(self.db_path) + try: + conn.execute( + "INSERT INTO tasks(id, project_id, prompt, status," + " idempotency_key, depends_on, priority, created_by," + " created_at, updated_at, last_heartbeat)" + " VALUES(?, 'p1', 'do work', ?, NULL, '[]', 0," + " 'tester', ?, ?, ?)", + (task_id, status, now, + updated_at if updated_at is not None else now, + last_heartbeat)) + if parent is not None: + conn.execute( + "UPDATE tasks SET parent_task_id = ? WHERE id = ?", + (parent, task_id)) + conn.commit() + finally: + conn.close() + + def _worker(self): + return Worker(self.db_path, "/fake/host.py", env=FAKE_ENV) + + def _heartbeat(self, task_id): + conn = sqlite3.connect(self.db_path) + try: + row = conn.execute( + "SELECT status, last_heartbeat FROM tasks WHERE id = ?", + (task_id,)).fetchone() + return row[0], row[1] + finally: + conn.close() + + def test_poll_wait_refreshes_heartbeat_and_survives_sweep(self): + worker = self._worker() + worker._deliver_marker_tools(worker.get_task("PM"), { + "status": "completed", "task_creates": [], + "task_polls": [{"task_id": "E1", "interval_seconds": 1, + "timeout_seconds": 2, "chars": 2000}], + "wait_threads": [], "task_cancels": [], "task_retries": []}) + status, heartbeat = self._heartbeat("PM") + self.assertEqual(status, "running") + # 等待循环刷过心跳：心跳应比陈旧的 -299 秒新得多。 + self.assertGreater(heartbeat, time.time() - 60) + swept = sweep_stale_tasks(self.db_path, 300) + self.assertNotIn("PM", swept) + status, _ = self._heartbeat("PM") + self.assertEqual(status, "running") + + def test_wait_threads_refreshes_heartbeat_and_survives_sweep(self): + worker = self._worker() + worker._deliver_marker_tools(worker.get_task("PM"), { + "status": "completed", "task_creates": [], + "task_polls": [], + "wait_threads": [{"task_ids": ["E1"], + "interval_seconds": 1, + "timeout_seconds": 2, "chars": 2000}], + "task_cancels": [], "task_retries": []}) + status, heartbeat = self._heartbeat("PM") + self.assertEqual(status, "running") + self.assertGreater(heartbeat, time.time() - 60) + swept = sweep_stale_tasks(self.db_path, 300) + self.assertNotIn("PM", swept) + status, _ = self._heartbeat("PM") + self.assertEqual(status, "running") + + +if __name__ == "__main__": + unittest.main() diff --git a/cloud/worker.py b/cloud/worker.py index 5f9d66a..08260d5 100644 --- a/cloud/worker.py +++ b/cloud/worker.py @@ -2402,6 +2402,15 @@ class Worker: if self._was_cancelled(task["id"]): cancelled = True break + # 等候中父任务自己刷心跳：标记等待在 _run_attempt 返回后 + # 才跑，那时 _running 已被清掉，run_loop 的心跳线程刷 + # 不到它；不自刷的话，sweep（默认 300 秒）会在父任务 + # 等到 3600 秒上限之前把它当死任务误杀。每次循环刷一次， + # 间隔至多 interval（钳制后 <=60 秒），远小于 300 秒阈值。 + try: + self.heartbeat(task["id"]) + except Exception: + pass try: conn = _db_connect(self.db_path) try: @@ -2450,6 +2459,12 @@ class Worker: if self._was_cancelled(task["id"]): cancelled = True break + # 与 _poll_task_snapshot 同款：等候中父任务自己刷心跳， + # 否则 sweep（默认 300 秒）会误杀最长可等 3600 秒的父任务。 + try: + self.heartbeat(task["id"]) + except Exception: + pass pending = [tid for tid in target_ids if tid not in resolved] if not pending: break