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
兼容，未改用例语义）。
131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
"""协调员子任务创建时配额门单测。"""
|
|
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()
|