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
兼容，未改用例语义）。
113 lines
3.7 KiB
Python
113 lines
3.7 KiB
Python
"""自动化执行器入口 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()
|