"""跨 worker 真 kill：kill_requests 意图表 + 收殓器单测。""" import http.client import json import os import sqlite3 import tempfile import threading import time import unittest from cloud.auth import AuthStore from cloud.server import (create_server, init_db, request_kill, sweep_stale_tasks) from cloud.worker import Worker def _insert_task(db_path, task_id, status, heartbeat_age=1000): conn = sqlite3.connect(db_path) try: now = time.time() conn.execute( "INSERT INTO tasks(id, project_id, prompt, status," " created_by, created_at, updated_at, last_heartbeat)" " VALUES(?, 'proj', 'do', ?, 'u1', ?, ?, ?)", (task_id, status, now, now, now - heartbeat_age)) conn.commit() finally: conn.close() def _kill_rows(db_path, task_id=None): conn = sqlite3.connect(db_path) conn.row_factory = sqlite3.Row try: if task_id is None: rows = conn.execute( "SELECT * FROM kill_requests").fetchall() else: rows = conn.execute( "SELECT * FROM kill_requests WHERE task_id = ?", (task_id,)).fetchall() return [dict(r) for r in rows] finally: conn.close() def _audit_actions(db_path, action): conn = sqlite3.connect(db_path) try: rows = conn.execute( "SELECT * FROM audit_log WHERE action = ?", (action,)).fetchall() return rows finally: conn.close() class _FakeDockerExecutor: """假 docker 执行器：只记杀了哪个任务派生名，不碰真 docker。""" def __init__(self, succeed=True): self.killed = [] self.succeed = succeed def container_name(self, task): return "codex-" + str(task.get("id")) def kill_container(self, task_or_name): name = (self.container_name(task_or_name) if isinstance(task_or_name, dict) else str(task_or_name)) self.killed.append(name) return self.succeed class KillRequestTestCase(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) def tearDown(self): self.tmp.cleanup() def test_request_kill_idempotent(self): self.assertTrue(request_kill(self.db_path, "t1", "u1")) self.assertFalse(request_kill(self.db_path, "t1", "u1")) rows = _kill_rows(self.db_path, "t1") self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["status"], "pending") def test_sweep_writes_kill_request(self): _insert_task(self.db_path, "t-old", "running") swept = sweep_stale_tasks(self.db_path, 300) self.assertIn("t-old", swept) rows = _kill_rows(self.db_path, "t-old") self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["requested_by"], "system") # 第二次 sweep：任务已不是 running，不再重复写。 sweep_stale_tasks(self.db_path, 300) self.assertEqual(len(_kill_rows(self.db_path, "t-old")), 1) class CancelHttpTestCase(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) self.auth = AuthStore(self.db_path) self.user_a, self.key_a = self.auth.create_user("user-a") 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() 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() def _req(self, method, path, body=None, key=None): conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10) try: headers = {} payload = None if body is not None: payload = json.dumps(body) headers["Content-Type"] = "application/json" if key is not None: headers["X-API-Key"] = key conn.request(method, path, body=payload, headers=headers) resp = conn.getresponse() raw = resp.read().decode() try: parsed = json.loads(raw) except ValueError: parsed = raw return resp.status, parsed finally: conn.close() def test_cancel_writes_kill_request(self): status, proj = self._req("POST", "/projects", {"name": "p"}, key=self.key_a) self.assertEqual(status, 201) pid = proj["id"] if isinstance(proj, dict) else proj["project"]["id"] status, task = self._req("POST", "/tasks", {"project_id": pid, "prompt": "do"}, key=self.key_a) self.assertEqual(status, 201) tid = task["id"] if isinstance(task, dict) else task["task"]["id"] status, _ = self._req("POST", "/tasks/%s/cancel" % tid, {}, key=self.key_a) self.assertEqual(status, 200) rows = _kill_rows(self.db_path, tid) self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["status"], "pending") # 重复取消：任务已终态，400，不重复写请求。 status, _ = self._req("POST", "/tasks/%s/cancel" % tid, {}, key=self.key_a) self.assertEqual(status, 400) self.assertEqual(len(_kill_rows(self.db_path, tid)), 1) class ReaperTestCase(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) def tearDown(self): self.tmp.cleanup() def _worker(self, executor): return Worker(self.db_path, "/fake/host.py", env={"XCM_API_KEY": "test-placeholder"}, executor=executor) def _seed_request(self, task_id): conn = sqlite3.connect(self.db_path) try: conn.execute( "INSERT INTO kill_requests(id, task_id, requested_by," " status, requested_at, done_at)" " VALUES(?, ?, 'u1', 'pending', ?, NULL)", ("req-" + task_id, task_id, time.time())) conn.commit() finally: conn.close() def test_reap_kills_stale_owner_container(self): _insert_task(self.db_path, "t1", "interrupted", heartbeat_age=1000) self._seed_request("t1") fake = _FakeDockerExecutor() worker = self._worker(fake) done = worker._reap_kill_requests(300) self.assertEqual(done, ["req-t1"]) self.assertEqual(fake.killed, ["codex-t1"]) self.assertEqual(_kill_rows(self.db_path, "t1")[0]["status"], "done") self.assertTrue(_audit_actions(self.db_path, "kill_request_reaped")) def test_reap_skips_live_owner(self): _insert_task(self.db_path, "t1", "interrupted", heartbeat_age=10) self._seed_request("t1") fake = _FakeDockerExecutor() worker = self._worker(fake) done = worker._reap_kill_requests(300) self.assertEqual(done, []) self.assertEqual(fake.killed, []) self.assertEqual(_kill_rows(self.db_path, "t1")[0]["status"], "pending") def test_reap_stale_request_for_live_task(self): # 任务已重回 queued（同 id 复活）：旧请求直接终结，绝不杀新跑。 _insert_task(self.db_path, "t1", "queued", heartbeat_age=1000) self._seed_request("t1") fake = _FakeDockerExecutor() worker = self._worker(fake) done = worker._reap_kill_requests(300) self.assertEqual(done, ["req-t1"]) self.assertEqual(fake.killed, []) self.assertTrue(_audit_actions(self.db_path, "kill_request_stale")) def test_reap_missing_task_and_missing_container(self): self._seed_request("t-gone") _insert_task(self.db_path, "t2", "interrupted", heartbeat_age=1000) self._seed_request("t2") fake = _FakeDockerExecutor(succeed=False) worker = self._worker(fake) done = worker._reap_kill_requests(300) self.assertEqual(sorted(done), ["req-t-gone", "req-t2"]) # 本机无容器（kill 返回 False）也不报错，请求照常终结。 self.assertEqual(_kill_rows(self.db_path, "t2")[0]["status"], "done") if __name__ == "__main__": unittest.main()