"""长任务自动恢复单测：只有用户暂停才暂停。 只测本路改动，不跑真实模型回合、不联网。 规则（见 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()