host
306cdf0d0c
Worker pool: workers table + register/heartbeat (heartbeat is
registration, no 404), GET /workers with alive (120s default),
POST /workers/heartbeat (X-API-Key, missing worker_id 400).
Tasks gain executor (local/docker/cloud, None = any worker);
worker claim filters by executor in pick_next and in the CAS
WHERE clause, CAS semantics unchanged; CloudExecutor _name
now correctly returns cloud, not local.
CloudExecutor: generic HTTP runner interface (endpoint from
explicit arg or CLOUD_RUNNER_ENDPOINT, auth header name only,
token value only from env into the request header, never into
logs/body/errors); no endpoint configured = explicit failure,
never silent downgrade. Vendor adapters (ECS/Cloud Run) are
out of scope, documented as such.
Auto PR: projects.auto_pr (0/1, create-time only, no PATCH
route exists); POST /tasks/{id}/pr (maintainer+ via
approve_merge) only for merged tasks (else 409); skip reasons
in order: auto_pr_disabled, no_worktree, no_remote, no_gh,
gh_not_authenticated — all return 200 with pr_skipped, task
stays merged, no push attempted; push/pr-create failure
returns 502 with step, task stays merged; success writes
pr_url/pr_number and audits pr_created. No default auto-push
from the worker. Remote URLs in errors are scrubbed of
userinfo; gh credentials are managed by gh itself, never
stored by the control plane.
Tests: host 27 pass, cloud 323 pass (8 skipped are
environment-dependent real-Postgres tests).
463 lines
18 KiB
Python
463 lines
18 KiB
Python
"""多 worker 池派发 + 通用 CloudExecutor + 自动 PR 单测。
|
|
|
|
只用标准库，不跑真实模型回合、不联网、不登录任何服务；
|
|
git/gh 一律用假 runner 注入，不依赖本机是否装 gh、是否有
|
|
远端仓库；CloudExecutor 的 endpoint 一律用假 http_client。
|
|
"""
|
|
|
|
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
|
|
from unittest import mock
|
|
|
|
try:
|
|
from cloud.auth import AuthStore
|
|
from cloud.pr import PrSkipped, create_task_pr
|
|
from cloud.sandbox import CloudExecutor, DockerExecutor
|
|
from cloud.server import create_server, init_db
|
|
from cloud.worker import Worker
|
|
from cloud.workers import (
|
|
heartbeat_worker,
|
|
list_workers,
|
|
register_worker,
|
|
)
|
|
except ImportError:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.auth import AuthStore
|
|
from cloud.pr import PrSkipped, create_task_pr
|
|
from cloud.sandbox import CloudExecutor, DockerExecutor
|
|
from cloud.server import create_server, init_db
|
|
from cloud.worker import Worker
|
|
from cloud.workers import (
|
|
heartbeat_worker,
|
|
list_workers,
|
|
register_worker,
|
|
)
|
|
|
|
|
|
def _result(returncode=0, stdout="", stderr=""):
|
|
return SimpleNamespace(returncode=returncode, stdout=stdout,
|
|
stderr=stderr)
|
|
|
|
|
|
class _ServerCase(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.viewer_id, self.viewer_key = self.auth.create_user("viewer")
|
|
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": "p1", "auto_pr": 1},
|
|
self.owner_key)
|
|
self.assertEqual(status, 201, self.project)
|
|
self.auth.set_member(
|
|
self.project["id"], self.owner_id, self.viewer_id, "viewer")
|
|
self.worktrees = os.path.join(self.tmp.name, "worktrees")
|
|
os.makedirs(self.worktrees, exist_ok=True)
|
|
self.server.host_worktrees_dir = self.worktrees
|
|
|
|
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, headers=None):
|
|
conn = http.client.HTTPConnection(
|
|
"127.0.0.1", self.port, timeout=10)
|
|
hdrs = dict(headers or {})
|
|
payload = None
|
|
if body is not None:
|
|
payload = json.dumps(body).encode("utf-8")
|
|
hdrs.setdefault("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 _insert_task(self, task_id, status, project_id=None):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
"INSERT INTO tasks(id, project_id, prompt, status,"
|
|
" depends_on, priority, created_by, created_at,"
|
|
" updated_at, last_heartbeat) VALUES(?, ?, ?, ?, '[]',"
|
|
" 0, 'system', 0, 0, NULL)",
|
|
(task_id, project_id or self.project["id"],
|
|
"do it", status))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def _task_row(self, task_id):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
row = conn.execute(
|
|
"SELECT status, pr_url, pr_number FROM tasks"
|
|
" WHERE id = ?", (task_id,)).fetchone()
|
|
finally:
|
|
conn.close()
|
|
return row
|
|
|
|
def _audit_actions(self):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
return [r[0] for r in conn.execute(
|
|
"SELECT action FROM audit_log").fetchall()]
|
|
finally:
|
|
conn.close()
|
|
|
|
def _git_runner(self, push_fail=False, remote_fail=False):
|
|
def run(cmd, cwd):
|
|
if "remote" in cmd:
|
|
if remote_fail:
|
|
return _result(1, "", "no remote")
|
|
return _result(0, "https://example.com/repo.git\n")
|
|
if "push" in cmd:
|
|
if push_fail:
|
|
return _result(1, "", "push rejected")
|
|
return _result(0, "")
|
|
return _result(0, "")
|
|
return run
|
|
|
|
def _gh_runner(self, authed=True, pr_url=None):
|
|
def run(cmd, cwd):
|
|
if cmd[:2] == ["gh", "auth"]:
|
|
return _result(0 if authed else 1, "", "")
|
|
if cmd[:2] == ["gh", "pr"]:
|
|
return _result(0, (pr_url or
|
|
"https://example.com/repo/pull/42\n"))
|
|
return _result(0, "")
|
|
return run
|
|
|
|
def _worktree(self, task_id):
|
|
path = os.path.join(self.worktrees, task_id)
|
|
os.makedirs(path, exist_ok=True)
|
|
return path
|
|
|
|
|
|
class WorkerRegistryTestCase(_ServerCase):
|
|
def test_heartbeat_requires_worker_id(self):
|
|
status, body = self._req(
|
|
"POST", "/workers/heartbeat", {}, self.owner_key)
|
|
self.assertEqual(status, 400, body)
|
|
|
|
def test_heartbeat_registers_and_lists_alive(self):
|
|
status, body = self._req(
|
|
"POST", "/workers/heartbeat",
|
|
{"worker_id": "w-docker-1", "executor": "docker",
|
|
"capacity": 2},
|
|
self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertTrue(body["alive"])
|
|
self.assertEqual(body["executor"], "docker")
|
|
status, body = self._req("GET", "/workers", None, self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(len(body["workers"]), 1)
|
|
self.assertTrue(body["workers"][0]["alive"])
|
|
|
|
def test_heartbeat_preserves_executor_when_not_given(self):
|
|
self._req("POST", "/workers/heartbeat",
|
|
{"worker_id": "w1", "executor": "docker"}, self.owner_key)
|
|
status, body = self._req(
|
|
"POST", "/workers/heartbeat", {"worker_id": "w1"},
|
|
self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["executor"], "docker")
|
|
|
|
def test_stale_worker_lists_not_alive(self):
|
|
register_worker(self.db_path, "w-old", executor="local")
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
"UPDATE workers SET last_seen = ? WHERE worker_id = ?",
|
|
(time.time() - 3600, "w-old"))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
workers = list_workers(self.db_path)
|
|
self.assertEqual(len(workers), 1)
|
|
self.assertFalse(workers[0]["alive"])
|
|
|
|
def test_heartbeat_without_key_is_unauthorized(self):
|
|
status, body = self._req(
|
|
"POST", "/workers/heartbeat", {"worker_id": "w1"}, None)
|
|
self.assertEqual(status, 401, body)
|
|
|
|
|
|
class TaskExecutorFilterTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
init_db(os.path.join(self.tmp.name, "control.db"))
|
|
self.db_path = os.path.join(self.tmp.name, "control.db")
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
for tid, exc in (("t-local", "local"), ("t-none", None),
|
|
("t-docker", "docker")):
|
|
conn.execute(
|
|
"INSERT INTO tasks(id, project_id, prompt, status,"
|
|
" depends_on, priority, created_by, created_at,"
|
|
" updated_at, last_heartbeat, executor)"
|
|
" VALUES(?, 'proj', 'do it', 'queued', '[]', 0,"
|
|
" 'system', 0, 0, NULL, ?)", (tid, exc))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def _docker_worker(self):
|
|
return Worker(
|
|
self.db_path, "host.py", config_path=None, repo=None,
|
|
executor=DockerExecutor(runner=lambda cmd: _result()),
|
|
env={"XCM_API_KEY": "test-placeholder"}, env_key=None,
|
|
task_timeout=1800, max_retries=0, retry_backoff_seconds=0)
|
|
|
|
def test_claim_rejects_other_pool_task(self):
|
|
worker = self._docker_worker()
|
|
self.assertFalse(worker._claim("t-local", executor="docker"))
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
status = conn.execute(
|
|
"SELECT status FROM tasks WHERE id = 't-local'"
|
|
).fetchone()[0]
|
|
finally:
|
|
conn.close()
|
|
self.assertEqual(status, "queued")
|
|
|
|
def test_claim_accepts_unassigned_and_same_pool(self):
|
|
worker = self._docker_worker()
|
|
self.assertTrue(worker._claim("t-none", executor="docker"))
|
|
self.assertTrue(worker._claim("t-docker", executor="docker"))
|
|
|
|
def test_cas_race_only_one_worker_wins(self):
|
|
w1 = self._docker_worker()
|
|
w2 = self._docker_worker()
|
|
self.assertTrue(w1._claim("t-none", executor="docker"))
|
|
self.assertFalse(w2._claim("t-none", executor="docker"))
|
|
|
|
def test_pick_next_skips_other_pool(self):
|
|
# 只剩 local 池任务时，docker worker 无就绪任务。
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute("DELETE FROM tasks WHERE id != 't-local'")
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
worker = self._docker_worker()
|
|
self.assertIsNone(worker.pick_next())
|
|
|
|
def test_run_one_skips_other_pool_without_model(self):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute("DELETE FROM tasks WHERE id != 't-local'")
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
worker = self._docker_worker()
|
|
result = worker.run_one()
|
|
self.assertFalse(result["dispatched"])
|
|
self.assertEqual(result["reason"], "no_ready_task")
|
|
|
|
|
|
class CloudExecutorTestCase(unittest.TestCase):
|
|
def test_unconfigured_fails_without_fallback(self):
|
|
ex = CloudExecutor(endpoint=None, env={})
|
|
with self.assertRaises(RuntimeError) as ctx:
|
|
ex.execute({"id": "t1", "prompt": "do it"})
|
|
self.assertIn("not configured", str(ctx.exception))
|
|
|
|
def test_configured_posts_only_id_prompt_repo_no_secrets(self):
|
|
seen = {}
|
|
|
|
def fake_client(url, payload, headers):
|
|
seen["url"] = url
|
|
seen["payload"] = payload
|
|
seen["headers"] = headers
|
|
return {"status": "completed", "stdout": "done"}
|
|
|
|
ex = CloudExecutor(
|
|
endpoint="https://runner.example.invalid/run",
|
|
http_client=fake_client,
|
|
auth_header="X-Runner-Token",
|
|
env={"CLOUD_RUNNER_AUTH_TOKEN": "test-token-value"})
|
|
result = ex.execute({"id": "t9", "prompt": "do it",
|
|
"repo": "https://example.com/r.git"})
|
|
self.assertEqual(result["status"], "completed")
|
|
self.assertEqual(seen["url"], "https://runner.example.invalid/run")
|
|
self.assertEqual(set(seen["payload"].keys()),
|
|
{"id", "prompt", "repo_url"})
|
|
self.assertEqual(seen["payload"]["repo_url"],
|
|
"https://example.com/r.git")
|
|
self.assertNotIn("test-token-value",
|
|
json.dumps(seen["payload"]))
|
|
self.assertEqual(seen["headers"]["X-Runner-Token"],
|
|
"test-token-value")
|
|
# dispatch_task 是 worker 走的入口，必须与 execute 同结果。
|
|
result2 = ex.dispatch_task({"id": "t9", "prompt": "do it"})
|
|
self.assertEqual(result2["status"], "completed")
|
|
|
|
def test_failed_status_maps_to_failed(self):
|
|
ex = CloudExecutor(
|
|
endpoint="https://runner.example.invalid/run",
|
|
http_client=lambda u, p, h: {"status": "failed",
|
|
"stdout": "boom"})
|
|
result = ex.execute({"id": "t1", "prompt": "x"})
|
|
self.assertEqual(result["status"], "failed")
|
|
|
|
|
|
class AutoPrTestCase(_ServerCase):
|
|
def _pr(self, task_id, key):
|
|
return self._req("POST", "/tasks/%s/pr" % task_id, {},
|
|
key)
|
|
|
|
def test_auto_pr_disabled_skips_and_stays_merged(self):
|
|
status, proj = self._req(
|
|
"POST", "/projects", {"name": "p-off"}, self.owner_key)
|
|
self.assertEqual(status, 201, proj)
|
|
self._insert_task("toff", "merged", project_id=proj["id"])
|
|
status, body = self._pr("toff", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertTrue(body["pr_skipped"])
|
|
self.assertEqual(body["reason"], "auto_pr_disabled")
|
|
self.assertEqual(self._task_row("toff")[0], "merged")
|
|
|
|
def test_auto_pr_illegal_value_rejected(self):
|
|
status, body = self._req(
|
|
"POST", "/projects", {"name": "p-bad", "auto_pr": 2},
|
|
self.owner_key)
|
|
self.assertEqual(status, 400, body)
|
|
|
|
def test_non_merged_returns_409(self):
|
|
self._insert_task("tq", "queued")
|
|
status, body = self._pr("tq", self.owner_key)
|
|
self.assertEqual(status, 409, body)
|
|
|
|
def test_viewer_cannot_pr(self):
|
|
self._insert_task("tm", "merged")
|
|
status, body = self._pr("tm", self.viewer_key)
|
|
self.assertEqual(status, 403, body)
|
|
|
|
def test_no_worktree_skips(self):
|
|
self._insert_task("tnw", "merged")
|
|
self.server.pr_git_runner = self._git_runner()
|
|
self.server.pr_gh_runner = self._gh_runner()
|
|
self.server.pr_gh_which = "gh"
|
|
status, body = self._pr("tnw", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["reason"], "no_worktree")
|
|
self.assertEqual(self._task_row("tnw")[0], "merged")
|
|
|
|
def test_no_remote_skips(self):
|
|
self._insert_task("tnr", "merged")
|
|
self._worktree("tnr")
|
|
self.server.pr_git_runner = self._git_runner(remote_fail=True)
|
|
self.server.pr_gh_runner = self._gh_runner()
|
|
self.server.pr_gh_which = "gh"
|
|
status, body = self._pr("tnr", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["reason"], "no_remote")
|
|
|
|
def test_no_gh_skips(self):
|
|
self._insert_task("tngh", "merged")
|
|
self._worktree("tngh")
|
|
self.server.pr_git_runner = self._git_runner()
|
|
self.server.pr_gh_runner = self._gh_runner()
|
|
self.server.pr_gh_which = ""
|
|
status, body = self._pr("tngh", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["reason"], "no_gh")
|
|
|
|
def test_gh_not_authenticated_skips(self):
|
|
self._insert_task("tna", "merged")
|
|
self._worktree("tna")
|
|
self.server.pr_git_runner = self._git_runner()
|
|
self.server.pr_gh_runner = self._gh_runner(authed=False)
|
|
self.server.pr_gh_which = "gh"
|
|
status, body = self._pr("tna", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["reason"], "gh_not_authenticated")
|
|
|
|
def test_push_failure_returns_502_and_stays_merged(self):
|
|
self._insert_task("tpf", "merged")
|
|
self._worktree("tpf")
|
|
self.server.pr_git_runner = self._git_runner(push_fail=True)
|
|
self.server.pr_gh_runner = self._gh_runner()
|
|
self.server.pr_gh_which = "gh"
|
|
status, body = self._pr("tpf", self.owner_key)
|
|
self.assertEqual(status, 502, body)
|
|
self.assertEqual(body["step"], "push")
|
|
self.assertEqual(self._task_row("tpf")[0], "merged")
|
|
|
|
def test_success_writes_back_pr_and_audit(self):
|
|
self._insert_task("tok", "merged")
|
|
self._worktree("tok")
|
|
self.server.pr_git_runner = self._git_runner()
|
|
self.server.pr_gh_runner = self._gh_runner()
|
|
self.server.pr_gh_which = "gh"
|
|
status, body = self._pr("tok", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["pr_url"],
|
|
"https://example.com/repo/pull/42")
|
|
self.assertEqual(body["pr_number"], 42)
|
|
self.assertEqual(body["branch"], "host/tok")
|
|
row = self._task_row("tok")
|
|
self.assertEqual(row[0], "merged")
|
|
self.assertEqual(row[1], "https://example.com/repo/pull/42")
|
|
self.assertEqual(row[2], 42)
|
|
self.assertIn("pr_created", self._audit_actions())
|
|
|
|
def test_pr_module_auto_pr_gate(self):
|
|
with self.assertRaises(PrSkipped) as ctx:
|
|
create_task_pr("x", 0)
|
|
self.assertEqual(ctx.exception.reason, "auto_pr_disabled")
|
|
|
|
|
|
class ProjectTaskFieldTestCase(_ServerCase):
|
|
def test_create_task_executor_validation(self):
|
|
status, body = self._req(
|
|
"POST", "/tasks",
|
|
{"project_id": self.project["id"], "prompt": "x",
|
|
"executor": "bogus"},
|
|
self.owner_key)
|
|
self.assertEqual(status, 400, body)
|
|
status, body = self._req(
|
|
"POST", "/tasks",
|
|
{"project_id": self.project["id"], "prompt": "x",
|
|
"executor": "docker"},
|
|
self.owner_key)
|
|
self.assertEqual(status, 201, body)
|
|
self.assertEqual(body["executor"], "docker")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|