diff --git a/cloud/automations.py b/cloud/automations.py index 8c109c5..6082fa3 100644 --- a/cloud/automations.py +++ b/cloud/automations.py @@ -12,16 +12,15 @@ 数据库里(automations 表，带 project_id)，不是 host state 文件。两者存储不相交，同一个自动化绝不会同时出现在两边， 所以本批不存在双重执行。 -- 本批不给云端自动化加执行器、不动 host.py 的 run-due 区： - host 的 run-due 跑的是 host 任务模型(codex 协议线程)，云 - 端任务模型(POST /tasks 建、worker 认领)是另一条链路，把 - 云端定义硬塞进 host run-due 会造出第二个调度器。云端执行 - 器是后续一步，它必须走 claim_due_automation 先抢单(条件 - UPDATE，rowcount 为 1 才算抢到)，抢到后再建云端任务、再 - 记 automation_runs 一行；抢不到的执行器绝不能建任务。这 - 个 claim 门就是防重复执行的唯一关口，比 host run-due 的 - “跑完才写 last_run”更严(中途被杀也不会在当间隔内重跑， - 代价是崩掉的那次要等到下个间隔，见 claim 文档)。 +- 云端执行器入口是本模块的 run_due_automation（残留批次 B + 落地，host.py 的 run-due 区不动）：它先走 claim_due_automation + 抢单，抢到后再过配额门、建云端任务、记 automation_runs 一 + 行；抢不到的执行器绝不能建任务。这个 claim 门就是防重复 + 执行的唯一关口，比 host run-due 的“跑完才写 last_run”更严 + （中途被杀也不会在当间隔内重跑，代价是崩掉的那次要等到 + 下个间隔才补跑，见 claim 文档）。配额门与服务端建任务口 + 径一致（见 _automation_quota_exceeded）：超配额时不建任务， + 产出记 skipped_quota（不是失败，不进失败重试循环），并记审计。 - 分诊收件箱是只读视图：到期自动化、失败或中断的任务、待 处理的后续消息、最近的自动化产出记录。它不跑模型、不改 任务状态，只把待处理项聚到一处。 @@ -44,8 +43,10 @@ import；表迁移写法仿 peer_messages(CREATE TABLE IF NOT EXISTS + PRAGMA 查列补齐，SQLite 与 Postgres 适配层同款)。 """ +import os import time import uuid +from datetime import datetime, timezone MIN_INTERVAL_SECONDS = 60 MAX_INTERVAL_SECONDS = 30 * 86400 @@ -461,6 +462,202 @@ def record_automation_run(conn, automation_id, project_id, task_id=None, } +def _default_quota_usd_micros(): + """默认月配额（微美元）：与服务端 _default_quota_usd_micros + 同值同环境变量（CLOUD_DEFAULT_QUOTA_USD_MICROS，默认 $10）。""" + try: + return max(0, int(os.environ.get( + "CLOUD_DEFAULT_QUOTA_USD_MICROS", "10000000"))) + except (TypeError, ValueError): + return 10000000 + + +def _unpriced_token_cap(): + """无价模型月 token 上限：与 cost.unpriced_token_cap 同值 + 同环境变量（CLOUD_QUOTA_UNPRICED_TOKEN_CAP，默认 2000 万）。""" + try: + return max(0, int(os.environ.get( + "CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "20000000"))) + except (TypeError, ValueError): + return 20000000 + + +def _automation_quota_exceeded(conn, owner): + """自动化建任务前的配额门：超限返回错误文本，否则 None。 + + 口径与服务端 _quota_status / worker _quota_exceeded 一致： + 按创建人算当月（UTC）已记账 USD + 无价模型 actual token； + 无配额行按默认额度（不是不限额），显式 NULL 行才不限额； + 查库失败 fail-closed（返回 quota_check_failed，调用方按跳 + 过处理，绝不建任务）。cost_records 表不存在时视为零花费 + （CostStore 首次记账时才建表，不存在即从未记过账），其 + 它查库异常才 fail-closed——全新库上不能因此误拒全部运行。 + 本模块只依赖标准库，不 import cloud 其它模块，故此处就 + 地实现 SQL 版口径，注释里标明三处同口径点，改口径时三 + 处要一起改（server._quota_status、worker._quota_exceeded、 + 本函数）。 + """ + if not owner: + return ("quota_exceeded: 自动化缺少创建人，无法查配额，" + "拒绝建任务") + try: + conn.execute( + "CREATE TABLE IF NOT EXISTS user_quotas(" + "user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT, " + " updated_at DOUBLE PRECISION)") + row = conn.execute( + "SELECT quota_usd_micros FROM user_quotas" + " WHERE user_id = ?", + (owner,)).fetchone() + if row is None: + quota = _default_quota_usd_micros() + else: + quota = row_get(row, "quota_usd_micros") + if quota is None: + return None + month = datetime.now(timezone.utc).strftime("%Y-%m") + try: + conn.execute("SELECT COUNT(*) FROM cost_records").fetchone() + table_exists = True + except Exception: + # PG 上失败语句会毒化当前事务：先 rollback，后面的 + # 任务 INSERT / 产出记录才能继续写；SQLite 上无害。 + try: + conn.rollback() + except Exception: + pass + table_exists = False + if table_exists: + used_row = conn.execute( + "SELECT COALESCE(SUM(c.usd_micros), 0) AS total" + " FROM cost_records c JOIN tasks t" + " ON c.task_id = t.id" + " WHERE t.created_by = ? AND c.day LIKE ?", + (owner, month + "%")).fetchone() + used = row_get(used_row, "total") or 0 + unpriced_row = conn.execute( + "SELECT COALESCE(SUM(c.input_tokens + c.output_tokens), 0)" + " AS total" + " FROM cost_records c JOIN tasks t" + " ON c.task_id = t.id" + " WHERE t.created_by = ? AND c.day LIKE ?" + " AND c.usd_micros IS NULL AND c.source = 'actual'", + (owner, month + "%")).fetchone() + unpriced_used = row_get(unpriced_row, "total") or 0 + else: + used = 0 + unpriced_used = 0 + except Exception as exc: + # 同上：先 rollback 再回 fail-closed，调用方仍能记 + # skipped_quota 产出（不记成失败重试循环）。 + try: + conn.rollback() + except Exception: + pass + return "quota_check_failed: %s" % exc + if used >= quota: + return ("quota_exceeded: 本月配额已用完（已用 %d / " + "配额 %d 微美元，%s）" % (used, quota, month)) + if unpriced_used >= _unpriced_token_cap(): + return ("quota_exceeded: 本月无价模型 token 已用完" + "（已用 %d / 上限 %d，%s）" + % (unpriced_used, _unpriced_token_cap(), month)) + return None + + +def _audit_row(conn, project_id, action, detail, now): + """best-effort 写审计行：失败不影响自动化运行本身。""" + try: + conn.execute( + "CREATE TABLE IF NOT EXISTS audit_log(" + "id TEXT PRIMARY KEY, " + "user_id TEXT, " + "project_id TEXT, " + "action TEXT, " + "detail TEXT, " + "created_at DOUBLE PRECISION)") + conn.execute( + "INSERT INTO audit_log(id, user_id, project_id, action," + " detail, created_at) VALUES(?, 'system', ?, ?, ?, ?)", + (uuid.uuid4().hex, project_id, action, detail, now)) + conn.commit() + except Exception: + try: + conn.rollback() + except Exception: + pass + + +def run_due_automation(conn, automation_id, now=None): + """云端自动化执行器的唯一入口：抢单 -> 配额门 -> 建任务 -> 记产出。 + + 抢不到（不存在、已取消、未到期、已被别的执行器抢走）时 + 返回 None，不记产出、不建任务。抢到后先过配额门：创建人 + 超配额时不建任务，产出记 skipped_quota（不是失败，不进 + 失败重试循环）并记审计，返回该产出；配额内才 INSERT 云 + 端任务（queued，executor/model 留 NULL 走 worker 默认， + 与服务端建任务缺省行为一致），产出记 created 并记审计。 + 任务 INSERT 与产出记录由 record_automation_run 的 commit + 一并落盘。调用方（未来的定时执行器）只调本函数，绝不 + 自己拼 claim + INSERT，否则配额门可被绕过。 + """ + if now is None: + now = time.time() + ensure_automations_tables(conn) + if not claim_due_automation(conn, automation_id, now): + return None + auto = conn.execute( + "SELECT id, project_id, prompt, created_by FROM automations" + " WHERE id = ?", + (automation_id,)).fetchone() + if auto is None: + return None + project_id = row_get(auto, "project_id") + prompt = row_get(auto, "prompt") + owner = row_get(auto, "created_by") + gate = _automation_quota_exceeded(conn, owner) + if gate is not None: + run = record_automation_run( + conn, automation_id, project_id, task_id=None, + status="skipped_quota", detail=gate, now=now) + _audit_row(conn, project_id, "automation_run_quota_skipped", + "%s: %s" % (automation_id, gate), now) + return run + task_id = uuid.uuid4().hex + # 任务列补齐（与 ensure_automations_tables 的补列惯例同款）： + # 裸库/老库的 tasks 表可能没有 executor 等列，INSERT 前补齐， + # 否则建任务直接报错；PG 适配层会翻译 PRAGMA 查询。 + existing_tasks = _existing_columns(conn, "tasks") + for col, ddl in ( + ("executor", "TEXT"), + ("model", "TEXT"), + ("effort", "TEXT"), + ("environment_id", "TEXT"), + ("parent_task_id", "TEXT"), + ("coordinator_part", "TEXT")): + if col not in existing_tasks: + try: + conn.execute( + "ALTER TABLE tasks ADD COLUMN %s %s" % (col, ddl)) + except Exception: + pass + conn.execute( + "INSERT INTO tasks(id, project_id, prompt, status," + " idempotency_key, depends_on, priority, created_by," + " created_at, updated_at, last_heartbeat, executor," + " model, effort, environment_id, parent_task_id," + " coordinator_part)" + " VALUES(?, ?, ?, 'queued', NULL, '[]', 0, ?, ?, ?," + " NULL, NULL, NULL, NULL, NULL, NULL, NULL)", + (task_id, project_id, prompt, owner, now, now)) + run = record_automation_run( + conn, automation_id, project_id, task_id=task_id, + status="created", detail="", now=now) + _audit_row(conn, project_id, "automation_run_created", + "%s -> %s" % (automation_id, task_id), now) + return run + + def list_automation_runs(conn, automation_id, limit=TRIAGE_LIMIT): """列某自动化的最近产出记录，按时间倒序。""" ensure_automations_tables(conn) diff --git a/cloud/db.py b/cloud/db.py index 277f460..db6b565 100644 --- a/cloud/db.py +++ b/cloud/db.py @@ -1030,6 +1030,19 @@ POSTGRES_SCHEMA = ( "memory_id TEXT, " "created_at DOUBLE PRECISION, " "processed_at DOUBLE PRECISION)", + # kill_requests：跨 worker 真 kill 的意图表。取消/sweep 只改库 + # 状态杀不掉 owner 已死时遗留的容器；写一条 pending 请求后， + # 任何活着的 worker 在 run loop 的收殓器里按任务派生容器名 + # 在本机 docker 杀掉并标 done。同一任务只留一条 pending + # （request_kill 的 SELECT 守卫，见 server.py；并发重复写 + # 撞上时多出的行同样幂等终结，不会多杀——杀本身按名幂等）。 + "CREATE TABLE IF NOT EXISTS kill_requests(" + "id TEXT PRIMARY KEY, " + "task_id TEXT, " + "requested_by TEXT, " + "status TEXT, " + "requested_at DOUBLE PRECISION, " + "done_at DOUBLE PRECISION)", ) diff --git a/cloud/migrations.py b/cloud/migrations.py index ea1889b..952cd8d 100644 --- a/cloud/migrations.py +++ b/cloud/migrations.py @@ -57,17 +57,19 @@ MIGRATION_TABLES = ( # （本批后无行按默认额度处理，不再是无限额，但显式 # 管理员额度仍需原样迁过去，不能丢）。 "user_quotas", - # P2-ops 对齐：以下 4 张表在 POSTGRES_SCHEMA 里有、 + # P2-ops 对齐：以下 5 张表在 POSTGRES_SCHEMA 里有、 # 原清单里没有，补进后清单与 schema 表集合一致。 # uploads / task_attachments：附件包与任务附件； # task_followups：任务间互发消息的后续消息队列； # environments：项目环境（分支/executor）。 - # 四张表主键均为 TEXT 或复合主键，没有身份列， + # kill_requests：跨 worker 真 kill 的意图表（残留批次 B）。 + # 五张表主键均为 TEXT 或复合主键，没有身份列， # 不需要序列拨动（见 SEQUENCE_TABLES）。 "uploads", "task_attachments", "task_followups", "environments", + "kill_requests", # Cloud identity batch: TEXT primary keys, no sequence to sync. "user_profiles", "identity_memories", diff --git a/cloud/server.py b/cloud/server.py index 3e48b25..1bd0eef 100644 --- a/cloud/server.py +++ b/cloud/server.py @@ -1694,6 +1694,7 @@ def init_db(db_path): _ensure_uploads_tables(conn) _ensure_task_checkpoints_table(conn) _ensure_task_steps_table(conn) + ensure_kill_requests_table(conn) # worker 每轮 poll 都按 status 取任务（见 load_tasks 的 # WHERE status IN (...)），无索引时是全表扫描；任务表随 # 时间变大后索引是必须的。IF NOT EXISTS 保证老库幂等。 @@ -1803,6 +1804,56 @@ def _project_to_dict(row): } +def ensure_kill_requests_table(conn): + """建 kill_requests 表（幂等，SQLite/PG 通用写法）。""" + conn.execute( + "CREATE TABLE IF NOT EXISTS kill_requests(" + "id TEXT PRIMARY KEY, " + "task_id TEXT, " + "requested_by TEXT, " + "status TEXT, " + "requested_at DOUBLE PRECISION, " + "done_at DOUBLE PRECISION)") + + +def request_kill(db_path, task_id, requested_by): + """写一条 kill 请求：同一任务只留一条 pending，返回 True/False。 + + 已有 pending 时不重复写（返回 False）；否则写一行 pending + 并返回 True。调用方（取消/sweep）必须 best-effort 调用： + 写请求失败绝不能挡住取消本身——任务状态已是 interrupted， + owner 活着时的 1 秒看门狗仍会杀，收殓器只是补 owner 已死 + 的情形。表不存在的老库先建表（与 worker 侧同款守卫）。 + """ + try: + conn = _connect(db_path) + except Exception: + return False + try: + ensure_kill_requests_table(conn) + row = conn.execute( + "SELECT id FROM kill_requests" + " WHERE task_id = ? AND status = 'pending' LIMIT 1", + (task_id,)).fetchone() + if row is not None: + return False + conn.execute( + "INSERT INTO kill_requests(id, task_id, requested_by," + " status, requested_at, done_at)" + " VALUES(?, ?, ?, 'pending', ?, NULL)", + (uuid.uuid4().hex, task_id, requested_by, time.time())) + conn.commit() + return True + except Exception: + try: + conn.rollback() + except Exception: + pass + return False + finally: + conn.close() + + def sweep_stale_tasks(db_path, threshold_seconds=300): """把 running 且心跳过期的任务改为 interrupted，返回被扫任务 id 列表。 @@ -1834,6 +1885,23 @@ def sweep_stale_tasks(db_path, threshold_seconds=300): ) if cur.rowcount: swept.append(row["id"]) + # 被扫任务的 owner 可能已死：光改库状态杀不掉它遗留的 + # 容器，顺手写 kill 请求，活着的 worker 在收殓器里按 + # 任务派生容器名在本机真杀（见 worker._reap_kill_requests）。 + # 同一事务内写，状态与请求要么一起落盘、要么一起丢。 + if swept: + ensure_kill_requests_table(conn) + for task_id in swept: + exists = conn.execute( + "SELECT id FROM kill_requests" + " WHERE task_id = ? AND status = 'pending' LIMIT 1", + (task_id,)).fetchone() + if exists is None: + conn.execute( + "INSERT INTO kill_requests(id, task_id," + " requested_by, status, requested_at, done_at)" + " VALUES(?, ?, 'system', 'pending', ?, NULL)", + (uuid.uuid4().hex, task_id, now)) conn.commit() stale = [(row["id"], row["project_id"]) for row in rows if row["id"] in swept] @@ -6748,6 +6816,13 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): "SELECT * FROM tasks WHERE id = ?", (task_id,) ).fetchone() ) + # owner 活着时它自己的 1 秒看门狗会杀；owner 已死时 + # 靠这条 kill 请求让别的活着的 worker 在收殓器里杀 + # 遗留容器。写请求失败绝不挡取消本身（best-effort）。 + try: + request_kill(self.server.db_path, task_id, user_id) + except Exception: + pass # 取消后 worker 不会再跑它：暂存包没人 seed 了，顺手删掉 # （任务终局后 worker 侧也会删，两边删同一目录都容忍 # 不存在；删失败不挡取消本身）。 diff --git a/cloud/tests/test_automation_run_due.py b/cloud/tests/test_automation_run_due.py new file mode 100644 index 0000000..af2d00b --- /dev/null +++ b/cloud/tests/test_automation_run_due.py @@ -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() diff --git a/cloud/tests/test_coord_quota_gate.py b/cloud/tests/test_coord_quota_gate.py new file mode 100644 index 0000000..ae0787a --- /dev/null +++ b/cloud/tests/test_coord_quota_gate.py @@ -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() diff --git a/cloud/tests/test_kill_requests.py b/cloud/tests/test_kill_requests.py new file mode 100644 index 0000000..995641d --- /dev/null +++ b/cloud/tests/test_kill_requests.py @@ -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() diff --git a/cloud/worker.py b/cloud/worker.py index 9fc8fd5..a995183 100644 --- a/cloud/worker.py +++ b/cloud/worker.py @@ -614,6 +614,11 @@ class Worker: # 配额中断就续跑下去了。受 _quota_hit_lock 保护。 self._quota_hit_tasks = set() self._quota_hit_lock = threading.Lock() + # 协调员子任务因配额跳过建任务时已记过审计的分工 + # （(parent_id, part_name)）：协调员每轮都会重调 + # _coord_sync_children，超配额期间每轮都记审计会刷屏， + # 同一 worker 进程内同一分工只记一次。 + self._coord_quota_skipped = set() def _code_changed(self): """启动后关键代码是否被换过（部署换目录/手动改代码）： @@ -4016,6 +4021,24 @@ class Worker: "%s -> %s" % (name, prev["id"])) children.append({"name": name, "id": prev["id"]}) continue + # 自动化入闸（创建时门）：子任务认领时虽有配额门， + # 但不建任务才是扇出不烧钱的硬保证。父任务创建人 + # 超配额时跳过建这个子任务、记审计，父任务本身不 + # 失败（协调员按缺失子任务的既有语义等超时）；同 + # 一 worker 进程内同一分工只记一次审计，避免每轮 + # 重调都刷屏。查库失败时 _quota_exceeded 同样返回 + # 非空串（fail-closed），一并跳过。 + gate = self._quota_exceeded( + {"created_by": task["created_by"]}, None) + if gate is not None: + key = (task["id"], name) + if key not in self._coord_quota_skipped: + self._coord_quota_skipped.add(key) + self._audit( + task["project_id"], + "coordinator_child_quota_skipped", + "%s: %s" % (name, gate)) + continue child_id = uuid.uuid4().hex conn.execute( "INSERT INTO tasks(id, project_id, prompt, status," @@ -4124,6 +4147,125 @@ class Worker: conn.close() return marked + def _reap_kill_requests(self, sweep_threshold=300): + """跨 worker 真 kill 的收殓器（run loop 每轮调一次）。 + + 分工（与 owner 活着时的 1 秒看门狗互补，绝不抢杀）： + owner 活着时，它的 _docker_cancel_watchdog 在任务变 + interrupted 后 1 秒内按容器名真杀，本函数不碰；只有 + owner 已死（任务心跳停滞超过 sweep_threshold，与陈旧 + 回收同一阈值、同一口径 COALESCE(last_heartbeat, + updated_at)）时，本函数才动手——因为这时已没有任何 + 进程会来杀这个容器，不杀就是一直烧钱。 + + 判定链（任一不满足即跳过，请求留待下轮；任务行失踪 + 或任务已不在 interrupted 时请求直接终结，避免复活的 + 任务（同 id 重回 queued）被旧请求误杀）： + 1. 请求 status 仍为 pending； + 2. 任务存在且 status == 'interrupted'； + 3. 任务心跳停滞超过 sweep_threshold； + 4. 本机 executor 有 container_name/kill_container + （docker 执行器才有；local 执行器无容器名，跳过）。 + 杀本身按任务派生容器名（executor.container_name， + codex-<安全字符任务 id>），绝不按前缀扫、绝不杀自由 + 传入的名字；本机无此容器时 kill_container 返回 False + 也不报错，请求照常终结。终结时写 kill_request_reaped + 审计；单条请求异常只记日志、不炸整轮。 + 返回终结的请求 id 列表。 + """ + done = [] + try: + conn = _db_connect(self.db_path) + except Exception as exc: + self.logger.warning("kill reap: db connect failed: %s", exc) + return done + try: + try: + conn.execute( + "CREATE TABLE IF NOT EXISTS kill_requests(" + "id TEXT PRIMARY KEY, " + "task_id TEXT, " + "requested_by TEXT, " + "status TEXT, " + "requested_at DOUBLE PRECISION, " + "done_at DOUBLE PRECISION)") + conn.commit() + except Exception: + pass + try: + reqs = conn.execute( + "SELECT id, task_id FROM kill_requests" + " WHERE status = 'pending'").fetchall() + except Exception as exc: + self.logger.warning("kill reap: list failed: %s", exc) + return done + now = time.time() + kill = getattr(self.executor, "kill_container", None) + name_of = getattr(self.executor, "container_name", None) + for req in reqs: + req_id = req["id"] + task_id = req["task_id"] + try: + row = conn.execute( + "SELECT status, project_id, last_heartbeat," + " updated_at FROM tasks WHERE id = ?", + (task_id,)).fetchone() + if row is None: + # 任务行失踪：请求已无意义，直接终结。 + conn.execute( + "UPDATE kill_requests SET status = 'done'," + " done_at = ? WHERE id = ?", + (now, req_id)) + conn.commit() + done.append(req_id) + continue + if row["status"] != "interrupted": + # 任务已不在 interrupted（重回 queued 等）： + # 旧请求直接终结，绝不能杀新跑。 + conn.execute( + "UPDATE kill_requests SET status = 'done'," + " done_at = ? WHERE id = ?", + (now, req_id)) + conn.commit() + done.append(req_id) + self._audit(row["project_id"], + "kill_request_stale", + "%s not_interrupted" % task_id) + continue + hb = row["last_heartbeat"] + if hb is None: + hb = row["updated_at"] + if hb is not None and hb >= now - sweep_threshold: + # owner 还活着：留待下轮，绝不抢杀。 + continue + # owner 已死：按任务派生容器名在本机真杀。 + # 容器名只由任务 id 派生，不接受外部输入； + # 本机无此容器时 kill 返回 False 也不报错。 + if kill is not None and name_of is not None: + try: + kill({"id": task_id}) + except Exception: + pass + conn.execute( + "UPDATE kill_requests SET status = 'done'," + " done_at = ? WHERE id = ?", + (now, req_id)) + conn.commit() + done.append(req_id) + self._audit(row["project_id"], "kill_request_reaped", + task_id) + except Exception as exc: + self.logger.warning( + "kill reap: request %s failed: %s", req_id, exc) + try: + conn.rollback() + except Exception: + pass + continue + finally: + conn.close() + return done + def sweep_orphaned_coordinator_children(self): """常驻循环每轮的孤儿扫尾（P1）：父已终态或父行失踪、子 仍未终态时，把子标 interrupted 并记审计。 @@ -5495,6 +5637,20 @@ class Worker: except Exception as exc: log_event(self.logger, "worker_orphan_sweep_failed", message="orphan sweep failed: %s" % exc) + # 跨 worker 真 kill 的收殓器：放在 sweep 块外、每轮 + # 都跑（旧 PR 只在 sweep 块内跑，sweep 关掉时请 + # 求会积压；本函数自带心跳停滞判定，sweep 关掉 + # 也不抢杀活 owner）。陈旧回收刚把死 worker 的 + # running 任务标 interrupted 并写了 kill 请求， + # 这里按请求杀本机遗留容器；owner 活着的任务心 + # 跳新鲜，收殓器绝不抢杀。 + try: + self._reap_kill_requests( + sweep_threshold + if sweep_threshold is not None else 300) + except Exception as exc: + log_event(self.logger, "worker_reap_failed", + message="worker reap failed: %s" % exc) # 长任务自动恢复：接在 sweep 之后——刚被 sweep # 标 interrupted 的系统中断 goal 任务，本轮即可 # 回 queued、下一轮认领即断点续跑；用户暂停、