feat: 跨 worker 真 kill（kill_requests 意图表 + 收殓器）与自动化触发入闸
机制 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
兼容，未改用例语义）。
This commit is contained in:
host committed
1 parent
b306b56517
commit
3a06205d8b
8 files changed
+938
-12
No files matched your search
@@ -0,0 +1,112 @@
|
||||
"""自动化执行器入口 run_due_automation：抢单 + 配额门单测。"""
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
|
||||
from cloud.automations import (
|
||||
create_automation,
|
||||
ensure_automations_tables,
|
||||
list_automation_runs,
|
||||
run_due_automation,
|
||||
)
|
||||
from cloud.server import init_db
|
||||
|
||||
|
||||
def _tasks(db_path):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
return conn.execute("SELECT id, prompt, status FROM tasks").fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _audit_count(db_path, action):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
return conn.execute(
|
||||
"SELECT COUNT(*) FROM audit_log WHERE action = ?",
|
||||
(action,)).fetchone()[0]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _set_quota(db_path, user_id, micros):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS user_quotas("
|
||||
"user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT, "
|
||||
" updated_at DOUBLE PRECISION)")
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO user_quotas(user_id,"
|
||||
" quota_usd_micros, updated_at) VALUES(?, ?, ?)",
|
||||
(user_id, micros, time.time()))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class RunDueAutomationTestCase(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.conn = sqlite3.connect(self.db_path)
|
||||
self.conn.row_factory = sqlite3.Row
|
||||
ensure_automations_tables(self.conn)
|
||||
|
||||
def tearDown(self):
|
||||
self.conn.close()
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _auto(self, owner="owner1"):
|
||||
return create_automation(self.conn, "proj", "daily", "do thing",
|
||||
owner, interval_seconds=60)
|
||||
|
||||
def test_run_creates_task_within_quota(self):
|
||||
auto = self._auto()
|
||||
run = run_due_automation(self.conn, auto["id"])
|
||||
self.assertIsNotNone(run)
|
||||
self.assertEqual(run["status"], "created")
|
||||
self.assertIsNotNone(run["task_id"])
|
||||
tasks = _tasks(self.db_path)
|
||||
self.assertEqual(len(tasks), 1)
|
||||
self.assertEqual(tasks[0][1], "do thing")
|
||||
self.assertEqual(tasks[0][2], "queued")
|
||||
self.assertEqual(
|
||||
_audit_count(self.db_path, "automation_run_created"), 1)
|
||||
# 同一间隔内第二次跑：已抢过，返回 None，不重复建任务。
|
||||
self.assertIsNone(run_due_automation(self.conn, auto["id"]))
|
||||
self.assertEqual(len(_tasks(self.db_path)), 1)
|
||||
|
||||
def test_run_skips_when_over_quota(self):
|
||||
auto = self._auto()
|
||||
_set_quota(self.db_path, "owner1", 0)
|
||||
run = run_due_automation(self.conn, auto["id"])
|
||||
self.assertIsNotNone(run)
|
||||
self.assertEqual(run["status"], "skipped_quota")
|
||||
self.assertIsNone(run["task_id"])
|
||||
self.assertEqual(_tasks(self.db_path), [])
|
||||
self.assertEqual(
|
||||
_audit_count(self.db_path, "automation_run_quota_skipped"), 1)
|
||||
# 跳过仍记了产出（不是失败），同间隔内不再重跑。
|
||||
runs = list_automation_runs(self.conn, auto["id"])
|
||||
self.assertEqual(len(runs), 1)
|
||||
self.assertEqual(runs[0]["status"], "skipped_quota")
|
||||
|
||||
def test_run_not_due_returns_none(self):
|
||||
auto = self._auto()
|
||||
self.conn.execute(
|
||||
"UPDATE automations SET last_run_at = ? WHERE id = ?",
|
||||
(time.time(), auto["id"]))
|
||||
self.conn.commit()
|
||||
self.assertIsNone(run_due_automation(self.conn, auto["id"]))
|
||||
self.assertEqual(_tasks(self.db_path), [])
|
||||
self.assertEqual(
|
||||
list_automation_runs(self.conn, auto["id"]), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,130 @@
|
||||
"""协调员子任务创建时配额门单测。"""
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from cloud.server import init_db
|
||||
from cloud.worker import Worker
|
||||
|
||||
|
||||
def _insert_parent(db_path, task_id="coord", created_by="owner1"):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
parts = json.dumps([{"name": "后端", "prompt": "做后端"},
|
||||
{"name": "前端", "prompt": "做前端"}])
|
||||
conn.execute(
|
||||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||||
" created_by, created_at, updated_at, coordinator_parts)"
|
||||
" VALUES(?, 'proj', 'coord', 'running', ?, ?, ?, ?)",
|
||||
(task_id, created_by, 1000.0, 1000.0, parts))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _children(db_path, parent_id):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
return conn.execute(
|
||||
"SELECT id FROM tasks WHERE parent_task_id = ?",
|
||||
(parent_id,)).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _audit_count(db_path, action):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
try:
|
||||
return conn.execute(
|
||||
"SELECT COUNT(*) FROM audit_log WHERE action = ?",
|
||||
(action,)).fetchone()[0]
|
||||
except sqlite3.OperationalError:
|
||||
# 审计是 best-effort：写审计的连接与任务写连接在
|
||||
# SQLite 上撞写锁时审计会静默丢失（既有行为），
|
||||
# 此时表可能根本不存在，计 0 而不是让测试炸掉。
|
||||
return 0
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _set_quota(db_path, user_id, micros):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS user_quotas("
|
||||
"user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT, "
|
||||
" updated_at DOUBLE PRECISION)")
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO user_quotas(user_id,"
|
||||
" quota_usd_micros, updated_at) VALUES(?, ?, ?)",
|
||||
(user_id, micros, 1000.0))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class CoordQuotaGateTestCase(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.env = {"XCM_API_KEY": "test-placeholder"}
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _worker(self):
|
||||
return Worker(self.db_path, "/fake/host.py", env=self.env)
|
||||
|
||||
def _parts(self):
|
||||
return [{"name": "后端", "prompt": "做后端"},
|
||||
{"name": "前端", "prompt": "做前端"}]
|
||||
|
||||
def test_skip_when_over_quota(self):
|
||||
_insert_parent(self.db_path)
|
||||
_set_quota(self.db_path, "owner1", 0)
|
||||
worker = self._worker()
|
||||
task = worker.get_task("coord")
|
||||
children = worker._coord_sync_children(
|
||||
task, task["coordinator_parts"])
|
||||
self.assertEqual(children, [])
|
||||
self.assertEqual(_children(self.db_path, "coord"), [])
|
||||
# 两个分工各记一次跳过审计；第二轮重调不再重复记。
|
||||
self.assertEqual(
|
||||
_audit_count(self.db_path, "coordinator_child_quota_skipped"),
|
||||
2)
|
||||
worker._coord_sync_children(task, task["coordinator_parts"])
|
||||
self.assertEqual(
|
||||
_audit_count(self.db_path, "coordinator_child_quota_skipped"),
|
||||
2)
|
||||
self.assertEqual(_children(self.db_path, "coord"), [])
|
||||
|
||||
def test_create_when_within_quota(self):
|
||||
_insert_parent(self.db_path)
|
||||
_set_quota(self.db_path, "owner1", 100000000)
|
||||
worker = self._worker()
|
||||
task = worker.get_task("coord")
|
||||
children = worker._coord_sync_children(
|
||||
task, task["coordinator_parts"])
|
||||
self.assertEqual(len(children), 2)
|
||||
self.assertEqual(len(_children(self.db_path, "coord")), 2)
|
||||
self.assertEqual(
|
||||
_audit_count(self.db_path, "coordinator_child_quota_skipped"),
|
||||
0)
|
||||
|
||||
def test_create_with_default_quota_when_no_row(self):
|
||||
# 无配额行按默认 $10 计，未花过即放行（不是不限额，
|
||||
# 也不是零额度误拒）。
|
||||
_insert_parent(self.db_path)
|
||||
worker = self._worker()
|
||||
task = worker.get_task("coord")
|
||||
children = worker._coord_sync_children(
|
||||
task, task["coordinator_parts"])
|
||||
self.assertEqual(len(children), 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,241 @@
|
||||
"""跨 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()
|
||||
Reference in new issue
Block a user