host
4428a721a6
Deploy restarts kill running tasks: the worker marks them interrupted (worker_shutdown on graceful exit, sweep on stale reclaim), and auto_resume_interrupted used to accept goal tasks only, so every non-goal task killed by a deploy stayed interrupted forever unless a human retried it (the 2026-10-30 demo task died this way). Now worker_shutdown/sweep interruptions are re-queued for non-goal tasks too: the platform killed the task, so the platform re-queues it. The re-queue is a fresh run, not a breakpoint resume: goal_resume_pending is forced to 0 and any stale resume_message is cleared, so the new run cannot be hijacked by the goal-resume path or an old continue message; files in the same worktree persist. All existing guards stay: user pause, coordinator orphans, quota (still goal-only), backoff, the 3-strike stall guard (non-goal tasks have no checkpoints, so each platform re-queue counts as no-progress), and the auto_resume_max cap. Other non-goal reasons (NULL, quota, user cancel) still do not auto-resume.
289 lines
12 KiB
Python
289 lines
12 KiB
Python
"""长任务自动恢复单测：只有用户暂停才暂停。
|
||
|
||
只测本路改动，不跑真实模型回合、不联网。
|
||
规则（见 worker.auto_resume_interrupted docstring）：
|
||
- 系统中断的 goal 任务 -> 回 queued + goal_resume_pending=1 + count+1
|
||
- 用户暂停/取消（paused_by_user=1）-> 不动
|
||
- 配额仍超限 -> 不动（不重排队）；配额已恢复 -> 自动续跑
|
||
- 协调员孤儿、非平台中断的非 goal 任务 -> 不动
|
||
- 平台中断（worker_shutdown/sweep）的非 goal 任务 -> 回 queued
|
||
全新跑（goal_resume_pending=0，resume_message 清掉）
|
||
- 自动续跑次数到上限 -> 不动
|
||
- no_time_limit 任务的 _run_attempt 不设 worker 级截止
|
||
"""
|
||
|
||
import http.client
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
|
||
try:
|
||
from cloud.auth import AuthStore
|
||
from cloud.server import create_server, init_db
|
||
from cloud.worker import Worker
|
||
except ImportError:
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.auth import AuthStore
|
||
from cloud.server import create_server, init_db
|
||
from cloud.worker import Worker
|
||
|
||
|
||
def _insert_task(db_path, task_id, status="interrupted", **cols):
|
||
now = time.time()
|
||
conn = sqlite3.connect(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(?, ?, ?, ?, NULL, '[]', 0, ?, ?, ?, NULL)",
|
||
(task_id, "p1", "do work", status, "tester", now, now),
|
||
)
|
||
for key, value in cols.items():
|
||
conn.execute(
|
||
"UPDATE tasks SET %s = ? WHERE id = ?" % key,
|
||
(value, task_id))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _row(db_path, task_id):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT status, goal_resume_pending, auto_resume_count,"
|
||
" paused_by_user, interrupted_reason, no_time_limit"
|
||
" FROM tasks WHERE id = ?", (task_id,)).fetchone()
|
||
return row
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _resume_msg(db_path, task_id):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT resume_message FROM tasks WHERE id = ?",
|
||
(task_id,)).fetchone()
|
||
return row[0] if row else None
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
class LongtaskAutoResumeTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "test.db")
|
||
init_db(self.db_path)
|
||
self.worker = Worker(
|
||
self.db_path, "/fake/host.py", env={"XCM_API_KEY": "x"},
|
||
task_timeout=1800)
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def test_system_interrupted_goal_resumes(self):
|
||
_insert_task(self.db_path, "T1", goal_mode=1, goal="g",
|
||
interrupted_reason="sweep", paused_by_user=0)
|
||
resumed = self.worker.auto_resume_interrupted()
|
||
self.assertEqual(resumed, ["T1"])
|
||
status, pending, count, _paused, reason, _ntl = _row(
|
||
self.db_path, "T1")
|
||
self.assertEqual(status, "queued")
|
||
self.assertEqual(pending, 1)
|
||
self.assertEqual(count, 1)
|
||
self.assertIsNone(reason)
|
||
|
||
def test_user_pause_and_cancel_do_not_resume(self):
|
||
_insert_task(self.db_path, "T1", goal_mode=1, goal="g",
|
||
interrupted_reason="user_pause", paused_by_user=1)
|
||
_insert_task(self.db_path, "T2", goal_mode=1, goal="g",
|
||
interrupted_reason="user_cancel", paused_by_user=1)
|
||
self.assertEqual(self.worker.auto_resume_interrupted(), [])
|
||
for tid in ("T1", "T2"):
|
||
status, pending, count, _p, _r, _n = _row(self.db_path, tid)
|
||
self.assertEqual(status, "interrupted")
|
||
self.assertEqual(pending, 0)
|
||
self.assertEqual(count, 0)
|
||
|
||
def test_quota_orphan_nongoal_do_not_resume(self):
|
||
# 配额仍超限时不重排队（mock 现查超限）。
|
||
self.worker._quota_exceeded = lambda task, usage, on_error="closed": (
|
||
"quota_exceeded: 本月配额已用完")
|
||
_insert_task(self.db_path, "T1", goal_mode=1, goal="g",
|
||
interrupted_reason="quota", paused_by_user=0)
|
||
_insert_task(self.db_path, "T2", goal_mode=1, goal="g",
|
||
interrupted_reason="coordinator_orphan",
|
||
paused_by_user=0, parent_task_id="P0")
|
||
_insert_task(self.db_path, "T3", goal_mode=0)
|
||
self.assertEqual(self.worker.auto_resume_interrupted(), [])
|
||
for tid in ("T1", "T2", "T3"):
|
||
status, _pending, _c, _p, _r, _n = _row(self.db_path, tid)
|
||
self.assertEqual(status, "interrupted")
|
||
|
||
def test_platform_interrupted_nongoal_resumes_fresh(self):
|
||
# 平台中断的非 goal 任务：全新跑，不置断点标记，
|
||
# 残留的续跑消息清掉，旧断点标记清掉。
|
||
_insert_task(self.db_path, "T1", goal_mode=0,
|
||
interrupted_reason="worker_shutdown",
|
||
paused_by_user=0, goal_resume_pending=1,
|
||
resume_message="stale-msg")
|
||
_insert_task(self.db_path, "T2", goal_mode=0,
|
||
interrupted_reason="sweep", paused_by_user=0)
|
||
resumed = self.worker.auto_resume_interrupted()
|
||
self.assertEqual(sorted(resumed), ["T1", "T2"])
|
||
for tid in ("T1", "T2"):
|
||
status, pending, count, _p, reason, _n = _row(
|
||
self.db_path, tid)
|
||
self.assertEqual(status, "queued")
|
||
self.assertEqual(pending, 0)
|
||
self.assertEqual(count, 1)
|
||
self.assertIsNone(reason)
|
||
self.assertIsNone(_resume_msg(self.db_path, "T1"))
|
||
|
||
def test_other_reason_nongoal_does_not_resume(self):
|
||
# 非平台中断的非 goal 任务仍不接：原因 NULL、配额
|
||
# （即使配额已恢复，配额口径只放 goal 任务）。
|
||
self.worker._quota_exceeded = lambda task, usage, on_error="closed": None
|
||
_insert_task(self.db_path, "T1", goal_mode=0)
|
||
_insert_task(self.db_path, "T2", goal_mode=0,
|
||
interrupted_reason="quota", paused_by_user=0)
|
||
self.assertEqual(self.worker.auto_resume_interrupted(), [])
|
||
for tid in ("T1", "T2"):
|
||
status, _pending, _c, _p, _r, _n = _row(self.db_path, tid)
|
||
self.assertEqual(status, "interrupted")
|
||
|
||
def test_quota_restored_resumes(self):
|
||
# 配额已恢复（现查为 None）时配额中断也自动续跑。
|
||
self.worker._quota_exceeded = lambda task, usage, on_error="closed": None
|
||
_insert_task(self.db_path, "T1", goal_mode=1, goal="g",
|
||
interrupted_reason="quota", paused_by_user=0)
|
||
self.assertEqual(self.worker.auto_resume_interrupted(), ["T1"])
|
||
status, pending, count, _p, reason, _n = _row(self.db_path, "T1")
|
||
self.assertEqual(status, "queued")
|
||
self.assertEqual(pending, 1)
|
||
self.assertEqual(count, 1)
|
||
self.assertIsNone(reason)
|
||
|
||
def test_auto_resume_max_cap(self):
|
||
_insert_task(self.db_path, "T1", goal_mode=1, goal="g",
|
||
interrupted_reason="sweep", paused_by_user=0,
|
||
auto_resume_count=50)
|
||
self.assertEqual(self.worker.auto_resume_interrupted(), [])
|
||
status, _p, count, _pa, _r, _n = _row(self.db_path, "T1")
|
||
self.assertEqual(status, "interrupted")
|
||
self.assertEqual(count, 50)
|
||
|
||
def test_no_time_limit_run_attempt_has_no_deadline(self):
|
||
captured = {}
|
||
|
||
def fake_execute(task, kind="dispatch", deadline_seconds=None,
|
||
seed_host_dir=None, quota_check=None):
|
||
captured["deadline"] = deadline_seconds
|
||
return {"status": "completed", "stdout": "", "stderr": ""}
|
||
|
||
self.worker._execute_via_process = fake_execute
|
||
self.worker._ensure_gateway = lambda task: (None, {})
|
||
self.worker.process_factory = object()
|
||
task = {"id": "T1", "project_id": "p1", "prompt": "x",
|
||
"goal_mode": 0, "no_time_limit": 1}
|
||
self.worker._run_attempt(task)
|
||
self.assertIsNone(captured["deadline"])
|
||
task2 = dict(task, no_time_limit=0)
|
||
self.worker._run_attempt(task2)
|
||
self.assertEqual(captured["deadline"], 1800)
|
||
|
||
|
||
class ServerPauseTestCase(unittest.TestCase):
|
||
"""POST /tasks/{id}/pause：用户暂停标记 + 自动恢复不接它。"""
|
||
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
self.auth = AuthStore(self.db_path)
|
||
self.owner_id, self.owner_key = self.auth.create_user("owner")
|
||
self._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
|
||
os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
|
||
self.server = create_server(self.db_path, "127.0.0.1", 0)
|
||
self.port = self.server.server_address[1]
|
||
self.thread = threading.Thread(
|
||
target=self.server.serve_forever, daemon=True)
|
||
self.thread.start()
|
||
status, project = self._req(
|
||
"POST", "/projects", {"name": "demo"}, self.owner_key)
|
||
self.assertEqual(status, 201, project)
|
||
self.project_id = project["id"]
|
||
_insert_task(self.db_path, "TP", status="running",
|
||
goal_mode=1, goal="g")
|
||
|
||
def tearDown(self):
|
||
self.server.shutdown()
|
||
self.server.server_close()
|
||
self.thread.join(timeout=5)
|
||
try:
|
||
self.auth.conn.close()
|
||
except Exception:
|
||
pass
|
||
self.tmp.cleanup()
|
||
if getattr(self, "_old_admins", None) is None:
|
||
os.environ.pop("CODEX_GLOBAL_ADMINS", None)
|
||
else:
|
||
os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins
|
||
|
||
def _req(self, method, path, body=None, key=None):
|
||
conn = http.client.HTTPConnection(
|
||
"127.0.0.1", self.port, timeout=10)
|
||
hdrs = {}
|
||
payload = None
|
||
if body is not None:
|
||
payload = json.dumps(body).encode("utf-8")
|
||
hdrs["Content-Type"] = "application/json"
|
||
if key is not None:
|
||
hdrs["X-API-Key"] = key
|
||
try:
|
||
conn.request(method, path, body=payload, headers=hdrs)
|
||
resp = conn.getresponse()
|
||
data = resp.read()
|
||
finally:
|
||
conn.close()
|
||
try:
|
||
parsed = json.loads(data.decode("utf-8")) if data else None
|
||
except ValueError:
|
||
parsed = None
|
||
return resp.status, parsed
|
||
|
||
def test_pause_marks_user_paused_and_auto_resume_skips(self):
|
||
# 注意：_insert_task 的 project_id 是 "p1"，与本用例项目
|
||
# 不同——pause 的权限校验按任务所属项目查，owner 不是
|
||
# p1 成员会 403；所以先把任务项目改成本用例项目。
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"UPDATE tasks SET project_id = ? WHERE id = ?",
|
||
(self.project_id, "TP"))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
status, body = self._req(
|
||
"POST", "/tasks/TP/pause", {}, self.owner_key)
|
||
self.assertEqual(status, 200, body)
|
||
self.assertEqual(body["status"], "interrupted")
|
||
self.assertEqual(body["paused_by_user"], 1)
|
||
self.assertEqual(body["interrupted_reason"], "user_pause")
|
||
worker = Worker(self.db_path, "/fake/host.py",
|
||
env={"XCM_API_KEY": "x"})
|
||
self.assertEqual(worker.auto_resume_interrupted(), [])
|
||
row = _row(self.db_path, "TP")
|
||
self.assertEqual(row[0], "interrupted")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|