host
7c46ac8e5b
建任务时传 coordinator_parts（分工列表，1-20 项，分工名唯一）即为协调员任务：本身不跑模型回合，worker 认领后一次建齐子任务（普通任务，并行被各 worker 认领），等全部终态后合流成汇总报告写回。子任务带 parent_task_id 与 coordinator_part，控制台可用 parent_task_id 查到；协调员重跑按分工名复用已完成子任务、失败/中断的回 queued 重跑；取消与超时尽力取消未终态子任务。子任务的配额/key/执行器门照常生效，不构成绕过。协调员任务拒绝与 goal/plan/验收命令同用。控制台：任务列表 [协调]/[协调子任务] 标记、详情分工与子任务区、新建表单子任务 JSON 输入。新单测 26 个。
418 lines
17 KiB
Python
418 lines
17 KiB
Python
"""协调员单测：分工校验、汇总报告、worker 分发/等齐/合流、服务端接口。
|
||
|
||
不跑真实模型回合、不联网：worker 侧用临时 SQLite + 假 runner，
|
||
子任务终态由测试直接写库模拟（等齐逻辑只读子任务行，不跑模型）；
|
||
服务端侧用本地临时端口 + 临时 SQLite，与 test_peer_messages_server
|
||
同款骨架。
|
||
"""
|
||
|
||
import http.client
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
|
||
try:
|
||
from cloud.auth import AuthStore
|
||
from cloud.coordinator import build_report, parse_parts, validate_parts
|
||
from cloud.server import create_server, init_db
|
||
from cloud.worker import Worker
|
||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.auth import AuthStore
|
||
from cloud.coordinator import build_report, parse_parts, validate_parts
|
||
from cloud.server import create_server, init_db
|
||
from cloud.worker import Worker
|
||
|
||
|
||
def _completed_runner(cmd):
|
||
return SimpleNamespace(returncode=0, stdout="{}", stderr="")
|
||
|
||
|
||
def _insert_task(db_path, task_id, prompt="do work", status="queued",
|
||
project_id="p1", coordinator_parts=None,
|
||
parent_task_id=None, coordinator_part=None,
|
||
result_text=None, error_text=None):
|
||
now = time.time()
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||
" idempotency_key, depends_on, priority, created_by,"
|
||
" created_at, updated_at, last_heartbeat, coordinator_parts,"
|
||
" parent_task_id, coordinator_part, result_text, error_text)"
|
||
" VALUES(?, ?, ?, ?, NULL, '[]', 0, 'tester', ?, ?, NULL,"
|
||
" ?, ?, ?, ?, ?)",
|
||
(task_id, project_id, prompt, status, now, now,
|
||
json.dumps(coordinator_parts) if coordinator_parts else None,
|
||
parent_task_id, coordinator_part, result_text, error_text),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _row(db_path, task_id):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT id, status, result_text, error_text, parent_task_id,"
|
||
" coordinator_part, executor, model, effort FROM tasks"
|
||
" WHERE id = ?", (task_id,)).fetchone()
|
||
finally:
|
||
conn.close()
|
||
if row is None:
|
||
return None
|
||
return {
|
||
"id": row[0], "status": row[1], "result_text": row[2],
|
||
"error_text": row[3], "parent_task_id": row[4],
|
||
"coordinator_part": row[5], "executor": row[6],
|
||
"model": row[7], "effort": row[8],
|
||
}
|
||
|
||
|
||
def _children(db_path, parent_id):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
rows = conn.execute(
|
||
"SELECT id FROM tasks WHERE parent_task_id = ?"
|
||
" ORDER BY coordinator_part",
|
||
(parent_id,)).fetchall()
|
||
finally:
|
||
conn.close()
|
||
return [r[0] for r in rows]
|
||
|
||
|
||
class ValidatePartsTestCase(unittest.TestCase):
|
||
def test_valid_minimal(self):
|
||
parts, err = validate_parts([{"name": "a", "prompt": "做 a"}])
|
||
self.assertIsNone(err)
|
||
self.assertEqual(parts, [{"name": "a", "prompt": "做 a",
|
||
"model": None, "effort": None,
|
||
"executor": None}])
|
||
|
||
def test_valid_with_options(self):
|
||
parts, err = validate_parts([{
|
||
"name": "b", "prompt": "做 b", "model": "m1",
|
||
"effort": "high", "executor": "docker"}])
|
||
self.assertIsNone(err)
|
||
self.assertEqual(parts[0]["model"], "m1")
|
||
self.assertEqual(parts[0]["effort"], "high")
|
||
self.assertEqual(parts[0]["executor"], "docker")
|
||
|
||
def test_rejects_non_list(self):
|
||
_parts, err = validate_parts({"name": "a"})
|
||
self.assertIsNotNone(err)
|
||
|
||
def test_rejects_empty(self):
|
||
_parts, err = validate_parts([])
|
||
self.assertIsNotNone(err)
|
||
|
||
def test_rejects_too_many(self):
|
||
parts = [{"name": "p%d" % i, "prompt": "x"} for i in range(21)]
|
||
_parts, err = validate_parts(parts)
|
||
self.assertIsNotNone(err)
|
||
|
||
def test_rejects_duplicate_name(self):
|
||
_parts, err = validate_parts([
|
||
{"name": "a", "prompt": "x"}, {"name": "a", "prompt": "y"}])
|
||
self.assertIsNotNone(err)
|
||
|
||
def test_rejects_empty_name_or_prompt(self):
|
||
_parts, err = validate_parts([{"name": " ", "prompt": "x"}])
|
||
self.assertIsNotNone(err)
|
||
_parts, err = validate_parts([{"name": "a", "prompt": " "}])
|
||
self.assertIsNotNone(err)
|
||
|
||
def test_rejects_bad_effort_executor_model(self):
|
||
_parts, err = validate_parts(
|
||
[{"name": "a", "prompt": "x", "effort": "ultra"}])
|
||
self.assertIsNotNone(err)
|
||
_parts, err = validate_parts(
|
||
[{"name": "a", "prompt": "x", "executor": "k8s"}])
|
||
self.assertIsNotNone(err)
|
||
_parts, err = validate_parts(
|
||
[{"name": "a", "prompt": "x", "model": " "}])
|
||
self.assertIsNotNone(err)
|
||
|
||
def test_rejects_non_dict_part(self):
|
||
_parts, err = validate_parts(["a"])
|
||
self.assertIsNotNone(err)
|
||
|
||
|
||
class ParsePartsTestCase(unittest.TestCase):
|
||
def test_none_and_empty(self):
|
||
self.assertEqual(parse_parts(None), [])
|
||
self.assertEqual(parse_parts(""), [])
|
||
|
||
def test_bad_json(self):
|
||
self.assertEqual(parse_parts("{not json"), [])
|
||
self.assertEqual(parse_parts('{"name": "a"}'), [])
|
||
|
||
def test_valid_json(self):
|
||
parts = parse_parts('[{"name": "a", "prompt": "x"}]')
|
||
self.assertEqual(len(parts), 1)
|
||
self.assertEqual(parts[0]["name"], "a")
|
||
|
||
def test_list_input(self):
|
||
self.assertEqual(len(parse_parts([{"name": "a", "prompt": "x"}])), 1)
|
||
self.assertEqual(parse_parts([{"name": "a"}]), [])
|
||
|
||
|
||
class BuildReportTestCase(unittest.TestCase):
|
||
def test_all_completed(self):
|
||
report, overall = build_report([
|
||
{"name": "a", "id": "1", "status": "completed",
|
||
"result_text": "A 好了", "error_text": None},
|
||
{"name": "b", "id": "2", "status": "completed",
|
||
"result_text": "B 好了", "error_text": None}])
|
||
self.assertEqual(overall, "completed")
|
||
self.assertIn("2/2", report)
|
||
self.assertIn("A 好了", report)
|
||
|
||
def test_mixed_failed(self):
|
||
report, overall = build_report([
|
||
{"name": "a", "id": "1", "status": "completed",
|
||
"result_text": "ok", "error_text": None},
|
||
{"name": "b", "id": "2", "status": "failed",
|
||
"result_text": None, "error_text": "炸了"}])
|
||
self.assertEqual(overall, "failed")
|
||
self.assertIn("1/2", report)
|
||
self.assertIn("炸了", report)
|
||
|
||
def test_truncates_long_output(self):
|
||
report, _overall = build_report([
|
||
{"name": "a", "id": "1", "status": "completed",
|
||
"result_text": "x" * 600, "error_text": None}])
|
||
self.assertNotIn("x" * 600, report)
|
||
self.assertIn("x" * 500, report)
|
||
|
||
def test_empty_output_note(self):
|
||
report, _overall = build_report([
|
||
{"name": "a", "id": "1", "status": "completed",
|
||
"result_text": None, "error_text": None}])
|
||
self.assertIn("无输出", report)
|
||
|
||
|
||
class WorkerCoordinatorTestCase(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",
|
||
runner=_completed_runner, env=self.env)
|
||
|
||
def _parts(self):
|
||
return [{"name": "后端", "prompt": "做后端", "model": "m1",
|
||
"effort": "high"},
|
||
{"name": "前端", "prompt": "做前端"}]
|
||
|
||
def test_sync_creates_children_once(self):
|
||
_insert_task(self.db_path, "coord", coordinator_parts=self._parts())
|
||
worker = self._worker()
|
||
task = worker.get_task("coord")
|
||
first = worker._coord_sync_children(task, task["coordinator_parts"])
|
||
second = worker._coord_sync_children(task, task["coordinator_parts"])
|
||
self.assertEqual([c["id"] for c in first],
|
||
[c["id"] for c in second])
|
||
self.assertEqual(len(_children(self.db_path, "coord")), 2)
|
||
backend = _row(self.db_path, first[0]["id"])
|
||
self.assertEqual(backend["parent_task_id"], "coord")
|
||
self.assertEqual(backend["coordinator_part"], "后端")
|
||
self.assertEqual(backend["model"], "m1")
|
||
self.assertEqual(backend["effort"], "high")
|
||
self.assertEqual(backend["status"], "queued")
|
||
|
||
def test_sync_requeues_failed_keeps_completed(self):
|
||
_insert_task(self.db_path, "coord", coordinator_parts=self._parts())
|
||
_insert_task(self.db_path, "c-old", status="failed",
|
||
parent_task_id="coord", coordinator_part="后端",
|
||
result_text="旧输出", error_text="旧错")
|
||
_insert_task(self.db_path, "c-done", status="completed",
|
||
parent_task_id="coord", coordinator_part="前端",
|
||
result_text="已做完")
|
||
worker = self._worker()
|
||
task = worker.get_task("coord")
|
||
children = worker._coord_sync_children(task, task["coordinator_parts"])
|
||
self.assertEqual([c["id"] for c in children], ["c-old", "c-done"])
|
||
old = _row(self.db_path, "c-old")
|
||
self.assertEqual(old["status"], "queued")
|
||
self.assertIsNone(old["result_text"])
|
||
self.assertIsNone(old["error_text"])
|
||
done = _row(self.db_path, "c-done")
|
||
self.assertEqual(done["status"], "completed")
|
||
self.assertEqual(done["result_text"], "已做完")
|
||
|
||
def test_run_one_all_completed(self):
|
||
_insert_task(self.db_path, "coord", coordinator_parts=self._parts())
|
||
_insert_task(self.db_path, "c1", status="completed",
|
||
parent_task_id="coord", coordinator_part="后端",
|
||
result_text="后端好了")
|
||
_insert_task(self.db_path, "c2", status="completed",
|
||
parent_task_id="coord", coordinator_part="前端",
|
||
result_text="前端好了")
|
||
worker = self._worker()
|
||
result = worker.run_one(worker.get_task("coord"))
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertIn("2/2", result["stdout"])
|
||
self.assertIn("后端好了", result["stdout"])
|
||
self.assertEqual(_row(self.db_path, "coord")["status"], "completed")
|
||
|
||
def test_run_timeout_cancels_pending(self):
|
||
_insert_task(self.db_path, "coord", coordinator_parts=self._parts())
|
||
conn = sqlite3.connect(self.db_path)
|
||
conn.execute("UPDATE tasks SET status = 'running' WHERE id = 'coord'")
|
||
conn.commit()
|
||
conn.close()
|
||
worker = self._worker()
|
||
task = worker.get_task("coord")
|
||
result = worker._run_coordinator(
|
||
task, time.time(), poll_interval=0.01, timeout=0.05)
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertIn("timeout", result["error"])
|
||
self.assertEqual(_row(self.db_path, "coord")["status"], "failed")
|
||
for cid in _children(self.db_path, "coord"):
|
||
self.assertEqual(_row(self.db_path, cid)["status"], "interrupted")
|
||
|
||
def test_run_cancelled(self):
|
||
_insert_task(self.db_path, "coord", coordinator_parts=self._parts())
|
||
_insert_task(self.db_path, "c1", status="completed",
|
||
parent_task_id="coord", coordinator_part="后端",
|
||
result_text="ok")
|
||
_insert_task(self.db_path, "c2", status="completed",
|
||
parent_task_id="coord", coordinator_part="前端",
|
||
result_text="ok")
|
||
conn = sqlite3.connect(self.db_path)
|
||
conn.execute("UPDATE tasks SET status = 'running' WHERE id = 'coord'")
|
||
conn.commit()
|
||
conn.close()
|
||
worker = self._worker()
|
||
worker._cancel_requested.add("coord")
|
||
result = worker._run_coordinator(
|
||
worker.get_task("coord"), time.time(),
|
||
poll_interval=0.01, timeout=60)
|
||
self.assertEqual(result["status"], "interrupted")
|
||
self.assertEqual(result["error"], "cancelled")
|
||
self.assertEqual(_row(self.db_path, "coord")["status"], "interrupted")
|
||
|
||
|
||
class ServerCoordinatorTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
self.auth = AuthStore(self.db_path)
|
||
self.owner_id, self.owner_key = self.auth.create_user("owner")
|
||
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()
|
||
status, self.project = self._req(
|
||
"POST", "/projects", {"name": "demo"}, self.owner_key)
|
||
self.assertEqual(status, 201)
|
||
|
||
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)
|
||
hdrs = {}
|
||
payload = None
|
||
if body is not None:
|
||
payload = json.dumps(body).encode("utf-8")
|
||
hdrs["Content-Type"] = "application/json"
|
||
if key is not None:
|
||
hdrs["X-API-Key"] = key
|
||
try:
|
||
conn.request(method, path, body=payload, headers=hdrs)
|
||
resp = conn.getresponse()
|
||
data = resp.read()
|
||
finally:
|
||
conn.close()
|
||
try:
|
||
parsed = json.loads(data.decode("utf-8")) if data else None
|
||
except ValueError:
|
||
parsed = None
|
||
return resp.status, parsed
|
||
|
||
def _create_coord(self, parts=None, extra=None):
|
||
body = {"project_id": self.project["id"], "prompt": "协调干活",
|
||
"coordinator_parts": parts if parts is not None else [
|
||
{"name": "后端", "prompt": "做后端"},
|
||
{"name": "前端", "prompt": "做前端"}]}
|
||
if extra:
|
||
body.update(extra)
|
||
return self._req("POST", "/tasks", body, self.owner_key)
|
||
|
||
def test_create_coordinator_ok(self):
|
||
status, task = self._create_coord()
|
||
self.assertEqual(status, 201)
|
||
self.assertEqual(len(task["coordinator_parts"]), 2)
|
||
self.assertEqual(task["coordinator_parts"][0]["name"], "后端")
|
||
self.assertIsNone(task["parent_task_id"])
|
||
status, fetched = self._req(
|
||
"GET", "/tasks/" + task["id"], key=self.owner_key)
|
||
self.assertEqual(status, 200)
|
||
self.assertEqual(len(fetched["coordinator_parts"]), 2)
|
||
|
||
def test_create_coordinator_invalid(self):
|
||
status, body = self._create_coord(parts=[])
|
||
self.assertEqual(status, 400)
|
||
status, _body = self._create_coord(parts=[
|
||
{"name": "a", "prompt": "x"}, {"name": "a", "prompt": "y"}])
|
||
self.assertEqual(status, 400)
|
||
status, _body = self._create_coord(parts=[
|
||
{"name": "a", "prompt": "x", "effort": "ultra"}])
|
||
self.assertEqual(status, 400)
|
||
|
||
def test_create_coordinator_conflicts(self):
|
||
status, _body = self._create_coord(extra={"goal_mode": 1})
|
||
self.assertEqual(status, 400)
|
||
status, _body = self._create_coord(extra={"plan": 1})
|
||
self.assertEqual(status, 400)
|
||
status, _body = self._create_coord(
|
||
extra={"verify_command": "true"})
|
||
self.assertEqual(status, 400)
|
||
|
||
def test_list_children_filter(self):
|
||
status, coord = self._create_coord()
|
||
self.assertEqual(status, 201)
|
||
_insert_task(self.db_path, "kid1", parent_task_id=coord["id"],
|
||
coordinator_part="后端",
|
||
project_id=self.project["id"])
|
||
_insert_task(self.db_path, "other", project_id=self.project["id"])
|
||
status, data = self._req(
|
||
"GET", "/tasks?project_id=" + self.project["id"] +
|
||
"&parent_task_id=" + coord["id"], key=self.owner_key)
|
||
self.assertEqual(status, 200)
|
||
self.assertEqual([t["id"] for t in data["tasks"]], ["kid1"])
|
||
status, data = self._req(
|
||
"GET", "/tasks?project_id=" + self.project["id"],
|
||
key=self.owner_key)
|
||
self.assertEqual(status, 200)
|
||
self.assertEqual(len(data["tasks"]), 3)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|