host
3a06205d8b
机制 A 移植自旧 PR #101（分支 feat-residual-b，头 895e004），
不合旧 PR：旧基线 6e3eb7b 已 CONFLICTING，本批重做到现 main
6d4f87e（#151 合并后）上。
- kill_requests 意图表：cloud/db.py POSTGRES_SCHEMA + migrations
清单补表；server.ensure_kill_requests_table / request_kill
（同一任务只留一条 pending，写失败返回 False）；init_db 的
SQLite 分支补建表（PG 分支由 POSTGRES_SCHEMA 全量建表覆盖）。
- 取消/sweep 写请求 best-effort：_cancel_task 在 commit 后调
request_kill（失败不挡取消）；sweep_stale_tasks 在同一事务
内为被扫任务写 system 请求。
- worker 收殓器 _reap_kill_requests：owner 活着（心跳停滞未超
300 秒，COALESCE(last_heartbeat, updated_at)）绝不抢杀；只
有 owner 已死且任务仍 interrupted 时，按任务派生容器名
（executor.container_name，codex-<安全字符任务 id>）在本机
真杀，绝不按前缀扫、绝不杀自由传入的名字；本机无容器时
请求照常终结；任务行失踪或已不在 interrupted 时请求直接
终结；单条异常只记日志。run loop 每轮调一次，放在 sweep
块外（旧 PR 只在 sweep 块内，sweep 关掉时请求会积压）。
- 自动化入闸只补两条旁路：worker._coord_sync_children 建子任
务前过 _quota_exceeded（超限跳过建任务、记审计，同一 worker
内同一分工只记一次）；automations.run_due_automation
（claim 抢单 -> 配额门 -> 建任务 -> 记产出，超限记
skipped_quota 不进失败重试）。triage/webhook/续跑/后续/
分叉未动。
- 单测：旧 PR 三份单测原样移植（Worker/Server 签名与现 main
兼容，未改用例语义）。
242 lines
8.7 KiB
Python
242 lines
8.7 KiB
Python
"""跨 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()
|