From 306cdf0d0cbb994607b64faf101ba84ec0a39a44 Mon Sep 17 00:00:00 2001 From: host Date: Mon, 26 Oct 2026 04:57:37 +0800 Subject: [PATCH] Cloud dispatch pool + auto PR (opt-in, skip clearly) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- cloud/README.md | 76 +++++ cloud/db.py | 15 +- cloud/pr.py | 158 +++++++++ cloud/sandbox.py | 98 +++++- cloud/server.py | 183 +++++++++- cloud/tests/test_cloud_dispatch_pr.py | 462 ++++++++++++++++++++++++++ cloud/worker.py | 132 ++++++-- cloud/workers.py | 166 +++++++++ 8 files changed, 1256 insertions(+), 34 deletions(-) create mode 100644 cloud/pr.py create mode 100644 cloud/tests/test_cloud_dispatch_pr.py create mode 100644 cloud/workers.py diff --git a/cloud/README.md b/cloud/README.md index 8476c5f..629b3ef 100644 --- a/cloud/README.md +++ b/cloud/README.md @@ -83,6 +83,14 @@ X-API-Key，因为外部服务无法持本站 API key，改靠各自的签名认 本站 API key——它的认证靠随机 state（10 分钟过期、一次性）加 client_secret（只走环境变量），不是无认证后门。 +本轮新增端点（同样要求 X-API-Key）：`GET /workers`（注册 worker +列表，含 alive）、`POST /workers/heartbeat`（body 含 worker_id， +心跳即注册）、`POST /tasks/{id}/pr`（已 merged 任务的自动 PR， +默认关闭，缺条件时返回 pr_skipped 而不是失败；见下文“云端派发与 +自动 PR”一节）。`POST /tasks` 新增可选 `executor` +（local/docker/cloud，未传表示任意 worker 可认领）； +`POST /projects` 新增可选 `auto_pr`（0/1，默认 0）。 + ### OAuth 真实浏览器登录（已实现，部署前置仍需用户动手） 流程（标准 OAuth2 授权码流程）： @@ -626,3 +634,71 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳 恢复演练、monitor、多用户验证四项的主机实跑还没做（本轮只在 VM 上 跑了单测、deploy 干跑、monitor 自测与 VM 实跑）。 +## 云端派发与自动 PR（本轮新增：多 worker 池、通用 runner、默认关闭的自动 PR） + +- 多 worker 池（`cloud/workers.py` + worker 注册/心跳）：worker 启动 + 时向控制面注册（workers 表：worker_id、executor、capacity、 + last_seen），`run_loop` 每轮结束写注册表心跳；`GET /workers` 列出 + 全部 worker 及其 alive（last_seen 在 120 秒内，阈值可经环境变量 + `CLOUD_WORKER_ALIVE_SECONDS` 调）；`POST /workers/heartbeat` + 要求 body 含 worker_id（缺了 400），心跳本身就是注册——worker_id + 不存在时自动注册，绝不 404。worker_id 可经 CLI `--worker-id` + 指定（生产部署应配稳定 id），未指定时默认“主机名-进程号” + （重启后进程号会变，旧行自然变 stale，不要用随机 uuid 让注册表 + 每次重启都多一行无法对应到同一台机器）；capacity 经 + `--worker-capacity` 指定（默认 1，注册表展示用，本轮 run_loop + 仍是串行一轮一任务）。两个端点都要 X-API-Key（与其它已认证 + 端点一致）；worker 直接写库注册时不需要 key，HTTP 心跳是给有 + key 的外部 worker 用的另一条路。 +- 任务按执行器池派发：tasks 表新增 `executor` 列（可空，老库幂等 + 补列）；`POST /tasks` 接受可选 executor（local/docker/cloud， + 非法值 400），未传时为 None 表示任意 worker 可认领；worker + 认领时只认领 executor 为空或等于自己执行器类型的任务—— + pick_next 先过滤，`_claim` 的 WHERE 里再加 + `(executor IS NULL OR executor = ?)` 兜底，CAS 语义不变（仅当 + queued 时标 running，两个同池 worker 抢同一任务时只有一个拿到 + rowcount == 1）；认领日志与审计里含 executor。注意 + `_executor_name` 的修正：CloudExecutor 以前会被算成 local， + 现在明确返回 cloud，否则 cloud worker 会认领 local 池的任务。 +- CloudExecutor 通用 runner 接口（厂商无关，`cloud/sandbox.py`）： + 本轮不假装接通某个云厂商（ECS、Cloud Run 等专属适配明确不在 + 范围）。CloudExecutor 接受一个 runner endpoint（显式参数或 + 环境变量 `CLOUD_RUNNER_ENDPOINT`）；未设时 execute 一律明确 + 失败“未配置云端 runner”，绝不静默降级到 Docker/Local；已设时 + 以 HTTP POST 把任务发给该 endpoint，请求体只有 id、prompt、 + repo_url（task["repo"]），不含任何密钥；endpoint 返回的 status + 为 completed 才算 completed，否则算 failed。认证只走 header + 名：`CLOUD_RUNNER_AUTH_HEADER` 给 header 名，token 值只从环境 + 变量取（`CLOUD_RUNNER_AUTH_TOKEN_ENV` 指定的变量名，默认 + `CLOUD_RUNNER_AUTH_TOKEN`），只放进请求头，绝不写进日志、 + 错误文本或返回结果。worker CLI 的 `--executor` 新增 cloud 选项， + 构造 CloudExecutor；worker 的 run_one 走 dispatch_task， + CloudExecutor 的 dispatch_task/run_task 直接委托给 execute。 +- 自动 PR（默认关闭，显式触发，`POST /tasks/{id}/pr`，需项目 + maintainer 以上权限 approve_merge）：仅当任务状态为 merged 时 + 可调（其它状态 409；任务不存在 404；权限不足 403）。逻辑按顺序 + 检查（见 `cloud/pr.py`）：项目 auto_pr 不为 1 时返回 200 且 + pr_skipped=true、reason=auto_pr_disabled，不失败、不推；本地 + 没有对应 host 任务 worktree（`<仓库根>/work/host-worktrees/ + `，host 任务名 = cloud 任务 id 经安全字符过滤， + 分支名 = host/，与 sandbox._host_name 同规则）时 + reason=no_worktree；worktree 里没有 git remote origin 时 + reason=no_remote；本机没有 gh 二进制时 reason=no_gh；gh 未登录 + （gh auth status 非零）时 reason=gh_not_authenticated；以上全过 + 才 push 分支 + gh pr create，成功后把 pr_url/pr_number 写回任务 + 记录并记审计 pr_created，返回 200 含 pr_url。push 或 pr create + 失败时返回 502 且错误文本写清是哪一步（push 还是 pr_create）， + 任务状态保持 merged 不变，绝不把任务标 failed。worker 不自动推： + 自动 PR 只由显式调本端点触发，默认自动推会把测试任务也推到远端。 + projects 表新增 `auto_pr` 列（默认 0，老库幂等补列）； + `POST /projects` 接受可选 auto_pr（只接受 0/1，非法值 400）； + 本轮没有现成的 PATCH /projects/{id} 路由，auto_pr 只能在创建时 + 指定，建后改它需直接改库或重建项目，不要新造一个只为改 auto_pr + 的路由。gh 的登录凭证由 gh 自己管，控制面不存、不打任何 token。 +- 仍未做（不要宣称已做）：厂商专属适配（ECS、Cloud Run 等，需 + 用户选厂商后另做）；默认自动推（本轮明确不做）；主机上 gh 的安装 + 与登录、远端写凭证（控制面不存，自动 PR 的真实成功路径仍需它们）； + 多 worker 的并发容量 enforcement（capacity 目前只是注册表展示， + run_loop 仍串行）；worker 经 HTTP 发心跳的凭证供给（worker 直接 + 写库时不需要，外部 worker 走 HTTP 时需自带 API key）。 + diff --git a/cloud/db.py b/cloud/db.py index 70e2cc5..5f62190 100644 --- a/cloud/db.py +++ b/cloud/db.py @@ -781,7 +781,8 @@ POSTGRES_SCHEMA = ( "name TEXT, " "owner_id TEXT, " "created_at DOUBLE PRECISION, " - "repo_url TEXT)", + "repo_url TEXT, " + "auto_pr INTEGER)", # tasks 含 server.migrate_tasks_table 补的那 6 列，建表时一次建齐； # 老库迁移仍走 migrate_tasks_table（PRAGMA 翻译后在 PG 上是 no-op）。 "CREATE TABLE IF NOT EXISTS tasks(" @@ -802,6 +803,9 @@ POSTGRES_SCHEMA = ( "started_at DOUBLE PRECISION, " "finished_at DOUBLE PRECISION, " "child_pid BIGINT, " + "executor TEXT, " + "pr_url TEXT, " + "pr_number INTEGER, " "UNIQUE(project_id, idempotency_key))", "CREATE TABLE IF NOT EXISTS summaries(" "id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, " @@ -875,6 +879,15 @@ def init_postgres_schema(conn): "PRAGMA table_info(projects)").fetchall()} if "repo_url" not in cols: conn.execute("ALTER TABLE projects ADD COLUMN repo_url TEXT") + if "auto_pr" not in cols: + conn.execute("ALTER TABLE projects ADD COLUMN auto_pr INTEGER") + tcols = {row[1] for row in conn.execute( + "PRAGMA table_info(tasks)").fetchall()} + for _name, _type in (("executor", "TEXT"), ("pr_url", "TEXT"), + ("pr_number", "INTEGER")): + if _name not in tcols: + conn.execute( + "ALTER TABLE tasks ADD COLUMN %s %s" % (_name, _type)) conn.commit() diff --git a/cloud/pr.py b/cloud/pr.py new file mode 100644 index 0000000..d5566d6 --- /dev/null +++ b/cloud/pr.py @@ -0,0 +1,158 @@ +"""云端自动 PR：默认关闭、显式触发、缺条件时跳过不失败。 + +零第三方依赖，只用 Python 标准库；不改 host.py 的任务逻辑， +只复用它的规则：host 任务名 = cloud 任务 id 经安全字符过滤 +（见 cloud/sandbox.py LocalExecutor._host_name），分支名是 +host/，worktree 在 <仓库根>/work/host-worktrees/ +。 + +触发规则（按顺序，任一步跳过即返回，不推、不改任务状态）： + 1. 任务所属项目的 auto_pr 不为 1 -> pr_skipped, auto_pr_disabled + 2. 本地没有对应 worktree -> pr_skipped, no_worktree + 3. worktree 里没有 git remote origin -> pr_skipped, no_remote + 4. 本机没有 gh 二进制 -> pr_skipped, no_gh + 5. gh 未登录（gh auth status 非零） -> pr_skipped, + gh_not_authenticated + 只有以上全过才 push 分支 + gh pr create；push 或 pr create + 失败时抛 PrStepError（调用方据此返回 502），任务状态保持 + merged 不变，绝不把任务标 failed。 + + worker 不自动推：自动 PR 只由显式调 POST /tasks/{id}/pr + 触发。默认自动推会把测试任务也推到远端，本模块不做。 + + 本模块不存、不打任何密钥：gh 的登录凭证由 gh 自己管 + （gh auth login），push/pr 的命令与错误文本里不出现任何 + token；错误文本只含 git/gh 的原文摘要，不含远端 URL 里的 + userinfo（若 remote URL 含 userinfo，摘要里把它抹掉）。 +""" + +import os +import re +import shutil +import subprocess + +_SKIP_AUTO_PR = "auto_pr_disabled" +_SKIP_NO_WORKTREE = "no_worktree" +_SKIP_NO_REMOTE = "no_remote" +_SKIP_NO_GH = "no_gh" +_SKIP_NOT_AUTHED = "gh_not_authenticated" + + +class PrSkipped(Exception): + """跳过：reason 见上列常量，调用方返回 200 + pr_skipped。""" + + def __init__(self, reason): + super().__init__(reason) + self.reason = reason + + +class PrStepError(Exception): + """push / pr create 哪一步失败：step 为 push 或 pr_create。""" + + def __init__(self, step, message): + super().__init__(message) + self.step = step + self.message = message + + +def host_name_for_task(task_id): + """cloud 任务 id -> host 任务名（与 sandbox._host_name 同规则）。""" + return re.sub(r"[^A-Za-z0-9_.-]", "-", str(task_id)) + + +def default_worktrees_root(): + """默认 worktree 根：<仓库根>/work/host-worktrees。""" + here = os.path.dirname(os.path.abspath(__file__)) + return os.path.join(os.path.dirname(here), "work", "host-worktrees") + + +def resolve_worktrees_root(override=None): + """worktree 根优先级：显式参数 > 环境变量 > 默认。""" + if override: + return override + return os.environ.get("CLOUD_HOST_WORKTREES_DIR") or \ + default_worktrees_root() + + +def _scrub_url(text): + """抹掉 URL 里的 userinfo（https://user:pass@host -> https://host）， + 防止错误文本把凭证带出去。""" + if not text: + return text + return re.sub(r"(https?://)[^/@\s]+@", r"\1", text) + + +def _run(cmd, cwd=None, runner=None): + """命令唯一入口：runner 供单测注入，返回带 returncode/stdout/ + stderr 的对象；真实路径走 subprocess。""" + if runner is not None: + return runner(cmd, cwd) + return subprocess.run(cmd, cwd=cwd, capture_output=True, text=True) + + +def _git_remote_url(worktree_dir, git_runner=None): + r = _run(["git", "-C", worktree_dir, "remote", "get-url", "origin"], + cwd=worktree_dir, runner=git_runner) + url = (getattr(r, "stdout", "") or "").strip() + if getattr(r, "returncode", 1) != 0 or not url: + return None + return url + + +def check_pr_ready(worktree_dir, git_runner=None, gh_which=None, + gh_runner=None): + """按顺序做跳过检查，全过返回 (worktree_dir, branch)；任一步 + 不过抛 PrSkipped。branch 由调用方按 host 任务名算好传进来 + 的是 worktree 检查——本函数只查 worktree 与 remote/gh/登录， + 分支名由 create_task_pr 组装。""" + if not os.path.isdir(worktree_dir): + raise PrSkipped(_SKIP_NO_WORKTREE) + if _git_remote_url(worktree_dir, git_runner) is None: + raise PrSkipped(_SKIP_NO_REMOTE) + gh_bin = gh_which if gh_which is not None else shutil.which("gh") + if not gh_bin: + raise PrSkipped(_SKIP_NO_GH) + r = _run(["gh", "auth", "status"], cwd=worktree_dir, runner=gh_runner) + if getattr(r, "returncode", 1) != 0: + raise PrSkipped(_SKIP_NOT_AUTHED) + return True + + +def create_task_pr(task_id, auto_pr, worktrees_root=None, + git_runner=None, gh_runner=None, gh_which=None): + """给一个已 merged 的云端任务建 PR。 + + auto_pr 不为 1 时直接抛 PrSkipped(auto_pr_disabled)，调用方 + 在查 worktree 之前就应先查项目 auto_pr——本函数也再查一次， + 防止直接调用者漏查。成功返回 {pr_url, pr_number, branch}。 + """ + if auto_pr != 1: + raise PrSkipped(_SKIP_AUTO_PR) + host_name = host_name_for_task(task_id) + worktree_dir = os.path.join( + resolve_worktrees_root(worktrees_root), host_name) + branch = "host/" + host_name + check_pr_ready(worktree_dir, git_runner=git_runner, + gh_which=gh_which, gh_runner=gh_runner) + push = _run(["git", "-C", worktree_dir, "push", "origin", branch], + cwd=worktree_dir, runner=git_runner) + if getattr(push, "returncode", 1) != 0: + raise PrStepError( + "push", "git push origin %s 失败：%s" % ( + branch, _scrub_url( + (getattr(push, "stderr", "") or "") + or (getattr(push, "stdout", "") or "")))) + title = "cloud: %s" % task_id + r = _run(["gh", "pr", "create", "--title", title, + "--body", "cloud 任务 %s 的自动 PR。" % task_id, + "--head", branch], + cwd=worktree_dir, runner=gh_runner) + out = (getattr(r, "stdout", "") or "").strip() + if getattr(r, "returncode", 1) != 0 or not out: + raise PrStepError( + "pr_create", "gh pr create 失败：%s" % _scrub_url( + (getattr(r, "stderr", "") or "") or out)) + pr_url = out.splitlines()[0].strip() + m = re.search(r"/(\d+)\s*$", pr_url) + pr_number = int(m.group(1)) if m else None + return {"pr_url": pr_url, "pr_number": pr_number, "branch": branch} diff --git a/cloud/sandbox.py b/cloud/sandbox.py index b7df802..efc3d97 100644 --- a/cloud/sandbox.py +++ b/cloud/sandbox.py @@ -593,14 +593,98 @@ def _docker_default_runner(cmd, timeout=None): class CloudExecutor(Executor): - """云端执行器（占位，未接通）。 + """云端执行器：可插拔的通用 HTTP runner 接口（厂商无关）。 - execute(task) 一律 raise RuntimeError，绝不假装已接通云端。 - 接通前提：真实云账号 + 预算上限配置，届时由控制面注入凭证与 - 限额后再替换本占位实现。 + 本轮不假装接通某个云厂商（ECS、Cloud Run 等专属适配明确 + 不在范围）。它只做一件事：把任务 POST 到一个配置好的通用 + runner endpoint，endpoint 返回任务状态，本类据此返回结果。 + + 配置（只走参数或环境变量，不写文件）： + - endpoint：显式参数 > CLOUD_RUNNER_ENDPOINT；未设时 execute + 一律 raise RuntimeError，绝不静默降级到 Docker/Local。 + - 认证：auth_header 参数或 CLOUD_RUNNER_AUTH_HEADER 给 header + 名；token 值只从环境变量取（CLOUD_RUNNER_AUTH_TOKEN_ENV + 指定的变量名，默认 CLOUD_RUNNER_AUTH_TOKEN），只放进请求头， + 绝不写进日志、错误文本或返回结果。 + - 请求体只有 id、prompt、repo_url（task["repo"]），不含任何密钥。 + - http_client 供单测注入：http_client(url, payload, headers) + 返回带 status 字段的 dict（或具 .status/.json 的对象）； + 真实路径走 urllib.request。 + - dispatch_task/run_task：worker 的 run_one 走 dispatch_task， + 这里直接委托给 execute，保证 cloud worker 池能真正派发。 """ + def __init__(self, endpoint=None, http_client=None, auth_header=None, + auth_token_env=None, env=None): + self._env = env if env is not None else os.environ + self.endpoint = endpoint or self._env.get("CLOUD_RUNNER_ENDPOINT") + self.http_client = http_client + self.auth_header = auth_header or self._env.get( + "CLOUD_RUNNER_AUTH_HEADER") + self.auth_token_env = auth_token_env or self._env.get( + "CLOUD_RUNNER_AUTH_TOKEN_ENV") or "CLOUD_RUNNER_AUTH_TOKEN" + + def _headers(self): + headers = {"Content-Type": "application/json"} + if self.auth_header: + try: + token = self._env.get(self.auth_token_env) + except Exception: + token = None + if token: + headers[self.auth_header] = token + return headers + + def _post(self, payload): + if self.http_client is not None: + return self.http_client(self.endpoint, payload, self._headers()) + import urllib.request + req = urllib.request.Request( + self.endpoint, + data=json.dumps(payload).encode("utf-8"), + headers=self._headers(), method="POST") + with urllib.request.urlopen(req, timeout=30) as resp: + return json.loads(resp.read().decode("utf-8")) + def execute(self, task) -> dict: - raise RuntimeError( - "cloud executor not connected: missing cloud account " - "and budget limit") + if not isinstance(task, dict) or task.get("id") is None: + raise ValueError("任务缺少 id") + if not self.endpoint: + raise RuntimeError( + "cloud runner not configured: set CLOUD_RUNNER_ENDPOINT " + "to a generic runner endpoint; no cloud account/runner " + "configured, refusing to silently fall back to " + "Docker/Local") + payload = { + "id": task["id"], + "prompt": task.get("prompt") or "", + "repo_url": task.get("repo"), + } + result = self._post(payload) + if isinstance(result, dict): + status = result.get("status") + stdout = result.get("stdout", "") + else: + status = getattr(result, "status", None) + try: + stdout = result.json().get("stdout", "") + except Exception: + stdout = "" + if status == "completed": + final = "completed" + else: + final = "failed" + return { + "executor": "cloud", + "task_id": task["id"], + "status": final, + "stdout": _truncate_stdout(stdout), + "returncode": 0 if final == "completed" else 1, + "host_status": None, + } + + def dispatch_task(self, task) -> dict: + return self.execute(task) + + def run_task(self, task) -> dict: + return self.execute(task) diff --git a/cloud/server.py b/cloud/server.py index e185964..d7b18fd 100644 --- a/cloud/server.py +++ b/cloud/server.py @@ -59,6 +59,11 @@ from cloud.db import PostgresIntegrityError as _PgIntegrityError from cloud.db import PostgresOperationalError as _PgOperationalError from cloud.kb import KBStore from cloud.memory import MemoryStore +from cloud.pr import PrSkipped, PrStepError, create_task_pr +from cloud.workers import ( + heartbeat_worker as _registry_heartbeat, + list_workers as _registry_list, +) from cloud.oauth import ( OAuthError, OAuthStateStore, @@ -227,6 +232,9 @@ def migrate_tasks_table(conn): "started_at": "REAL", "finished_at": "REAL", "child_pid": "INTEGER", + "executor": "TEXT", + "pr_url": "TEXT", + "pr_number": "INTEGER", } for name, coltype in new_cols.items(): if name not in cols: @@ -246,6 +254,8 @@ def migrate_projects_table(conn): "PRAGMA table_info(projects)").fetchall()} if "repo_url" not in cols: conn.execute("ALTER TABLE projects ADD COLUMN repo_url TEXT") + if "auto_pr" not in cols: + conn.execute("ALTER TABLE projects ADD COLUMN auto_pr INTEGER") def _is_postgres_backend(): @@ -275,7 +285,8 @@ def init_db(db_path): "name TEXT, " "owner_id TEXT, " "created_at REAL, " - "repo_url TEXT)" + "repo_url TEXT, " + "auto_pr INTEGER)" ) conn.execute( "CREATE TABLE IF NOT EXISTS tasks(" @@ -346,6 +357,9 @@ def _task_to_dict(row): "started_at": _row_get(row, "started_at"), "finished_at": _row_get(row, "finished_at"), "child_pid": _row_get(row, "child_pid"), + "executor": _row_get(row, "executor"), + "pr_url": _row_get(row, "pr_url"), + "pr_number": _row_get(row, "pr_number"), } @@ -358,6 +372,9 @@ def _project_to_dict(row): # 老库行缺 repo_url 列时 _row_get 给 None，不炸（列本身由 # migrate_projects_table 在 init_db 时补，老行读到也是 None）。 "repo_url": _row_get(row, "repo_url"), + # auto_pr 默认 0（关闭）：老行缺列时 None 按 0 计，绝不因 + # 迁移就把老项目默认改成自动推。 + "auto_pr": _row_get(row, "auto_pr") or 0, } @@ -569,6 +586,10 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): self._create_project(user_id) elif path == "/projects" and method == "GET": self._list_projects(user_id) + elif path == "/workers" and method == "GET": + self._workers_list(user_id) + elif path == "/workers/heartbeat" and method == "POST": + self._workers_heartbeat(user_id) elif path == "/metrics" and method == "GET": self._metrics(user_id) elif path == "/models" and method == "GET": @@ -595,6 +616,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): self._cancel_task(user_id, rest[: -len("/cancel")]) elif rest.endswith("/merge") and method == "POST": self._merge_task(user_id, rest[: -len("/merge")]) + elif rest.endswith("/pr") and method == "POST": + self._task_pr(user_id, rest[: -len("/pr")]) elif rest.endswith("/heartbeat") and method == "POST": self._heartbeat_task(user_id, rest[: -len("/heartbeat")]) else: @@ -626,6 +649,18 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): except ValueError as exc: self._send_json(400, {"error": str(exc)}) return + # auto_pr 可选，默认 0（关闭）：只接受 0/1（布尔 True/False + # 按 1/0 计），其它值 400 且不建项目。本轮没有现成的 + # PATCH /projects/{id} 路由，auto_pr 只能在创建时指定， + # 建后改它需直接改库或重建项目（见 README）。 + auto_pr_raw = body.get("auto_pr", 0) + if isinstance(auto_pr_raw, bool): + auto_pr = 1 if auto_pr_raw else 0 + elif isinstance(auto_pr_raw, int) and auto_pr_raw in (0, 1): + auto_pr = auto_pr_raw + else: + self._send_json(400, {"error": "auto_pr must be 0 or 1"}) + return now = time.time() conn = self._db() try: @@ -637,8 +672,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): return conn.execute( "INSERT INTO projects(id, name, owner_id, created_at," - " repo_url) VALUES(?, ?, ?, ?, ?)", - (project_id, name, user_id, now, repo_url), + " repo_url, auto_pr) VALUES(?, ?, ?, ?, ?, ?)", + (project_id, name, user_id, now, repo_url, auto_pr), ) conn.commit() project = { @@ -647,6 +682,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): "owner_id": user_id, "created_at": now, "repo_url": repo_url, + "auto_pr": auto_pr, } except (sqlite3.IntegrityError, _PgIntegrityError): self._send_json(409, {"error": "project exists"}) @@ -671,7 +707,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): ).fetchall() } rows = conn.execute( - "SELECT id, name, owner_id, created_at, repo_url" + "SELECT id, name, owner_id, created_at, repo_url, auto_pr" " FROM projects ORDER BY created_at, id" ).fetchall() finally: @@ -1095,6 +1131,14 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): except (TypeError, ValueError): self._send_json(400, {"error": "priority must be an integer"}) return + # executor 可选：None 表示任意 worker 可认领；指定时只能是 + # local/docker/cloud，非法值 400 且不建任务。 + executor = body.get("executor") + if executor is not None and executor not in ( + "local", "docker", "cloud"): + self._send_json(400, { + "error": "executor must be local, docker or cloud"}) + return idempotency_key = body.get("idempotency_key") if idempotency_key is not None and not isinstance(idempotency_key, str): self._send_json(400, {"error": "idempotency_key must be a string"}) @@ -1120,10 +1164,11 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): conn.execute( "INSERT INTO tasks(id, project_id, prompt, status," " idempotency_key, depends_on, priority, created_by," - " created_at, updated_at, last_heartbeat)" - " VALUES(?, ?, ?, 'queued', ?, ?, ?, ?, ?, ?, NULL)", + " created_at, updated_at, last_heartbeat, executor)" + " VALUES(?, ?, ?, 'queued', ?, ?, ?, ?, ?, ?, NULL, ?)", (task_id, project_id, prompt, idempotency_key, - json.dumps(depends_on), priority, user_id, now, now), + json.dumps(depends_on), priority, user_id, now, now, + executor), ) conn.commit() except (sqlite3.IntegrityError, _PgIntegrityError): @@ -1270,6 +1315,124 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): store.audit(user_id, task["project_id"], "merge_task", task_id) self._send_json(200, task) + # -- worker 池 ----------------------------------------------------- + def _workers_list(self, user_id): + """列出注册 worker：任何已认证用户可读（worker 是基础设施 + 信息，不属于某个项目的数据，不按项目成员过滤）。""" + workers = _registry_list(self.server.db_path) + self._send_json(200, {"workers": workers}) + + def _workers_heartbeat(self, user_id): + """worker 心跳：body 含 worker_id（必填，缺了 400）与可选 + executor/capacity。心跳本身就是注册：worker_id 不存在时 + 自动注册，绝不 404。""" + body, err = self._read_json() + if err: + self._send_json(400, {"error": err}) + return + worker_id = body.get("worker_id") + if not worker_id or not isinstance(worker_id, str): + self._send_json(400, {"error": "worker_id required"}) + return + executor = body.get("executor") + capacity = body.get("capacity") + if capacity is not None: + try: + capacity = int(capacity) + except (TypeError, ValueError): + self._send_json(400, {"error": "capacity must be an integer"}) + return + try: + worker = _registry_heartbeat( + self.server.db_path, worker_id, + executor=executor, capacity=capacity) + except ValueError as exc: + self._send_json(400, {"error": str(exc)}) + return + worker["alive"] = True + self._send_json(200, worker) + + # -- 自动 PR（默认关闭，显式触发）---------------------------------- + def _task_pr(self, user_id, task_id): + """给已 merged 的任务建 PR（见 cloud/pr.py 的跳过规则）。 + + 权限：项目 maintainer 以上（approve_merge，403）；任务 + 不存在 404；非 merged 状态 409；跳过时返回 200 + + pr_skipped 与 reason，任务状态保持 merged 不变；push/pr + 创建失败时返回 502 并写清哪一步失败，任务同样保持 merged。 + 成功时把 pr_url/pr_number 写回任务并记审计 pr_created。 + """ + row = self._fetch_task(task_id) + if row is None: + self._send_json(404, {"error": "task not found"}) + return + task = _task_to_dict(row) + with _auth_store(self.server.db_path) as store: + if not store.check_permission( + user_id, task["project_id"], "approve_merge"): + self._send_json(403, {"error": "forbidden"}) + return + if task["status"] != "merged": + self._send_json(409, { + "error": "only merged tasks can create a PR", + "status": task["status"], + }) + return + conn = self._db() + try: + prow = conn.execute( + "SELECT auto_pr FROM projects WHERE id = ?", + (task["project_id"],)).fetchone() + finally: + conn.close() + auto_pr = 0 + if prow is not None: + try: + auto_pr = prow["auto_pr"] or 0 + except (KeyError, IndexError, TypeError): + auto_pr = 0 + try: + result = create_task_pr( + task_id, auto_pr, + worktrees_root=getattr( + self.server, "host_worktrees_dir", None), + git_runner=getattr(self.server, "pr_git_runner", None), + gh_runner=getattr(self.server, "pr_gh_runner", None), + gh_which=getattr(self.server, "pr_gh_which", None)) + except PrSkipped as exc: + self._send_json(200, { + "task_id": task_id, "status": task["status"], + "pr_skipped": True, "reason": exc.reason, + }) + return + except PrStepError as exc: + self._send_json(502, { + "error": exc.message, "step": exc.step, + "task_id": task_id, "status": task["status"], + }) + return + now = time.time() + conn = self._db() + try: + conn.execute( + "UPDATE tasks SET pr_url = ?, pr_number = ?," + " updated_at = ? WHERE id = ?", + (result["pr_url"], result["pr_number"], now, task_id), + ) + conn.commit() + task = _task_to_dict( + conn.execute( + "SELECT * FROM tasks WHERE id = ?", + (task_id,)).fetchone()) + finally: + conn.close() + with _auth_store(self.server.db_path) as store: + store.audit(user_id, task["project_id"], "pr_created", + result["pr_url"]) + self._send_json(200, dict(task, pr_url=result["pr_url"], + pr_number=result["pr_number"], + branch=result["branch"])) + def _heartbeat_task(self, user_id, task_id): row = self._fetch_task(task_id) if row is None: @@ -1857,6 +2020,12 @@ def create_server(db_path, host="127.0.0.1", port=8765, backend=None): # 测试可注入 fake http_client / userinfo_client；绝不注入真实 secret。 server.oauth_http_client = None server.oauth_userinfo_client = None + # 自动 PR 注入点：生产为 None（走真实 git/gh），测试可注入 + # fake git/gh runner 与 worktree 根；绝不注入真实凭证。 + server.pr_git_runner = None + server.pr_gh_runner = None + server.pr_gh_which = None + server.host_worktrees_dir = None return server diff --git a/cloud/tests/test_cloud_dispatch_pr.py b/cloud/tests/test_cloud_dispatch_pr.py new file mode 100644 index 0000000..7427e34 --- /dev/null +++ b/cloud/tests/test_cloud_dispatch_pr.py @@ -0,0 +1,462 @@ +"""多 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() diff --git a/cloud/worker.py b/cloud/worker.py index 347dc41..e1bac16 100644 --- a/cloud/worker.py +++ b/cloud/worker.py @@ -101,6 +101,7 @@ import json import os import re import signal +import socket import subprocess import sys import time @@ -113,8 +114,9 @@ from cloud.logging_config import get_logger, log_event, setup_logging from cloud.models import ModelRouter, is_fallbackable, record_fallback from cloud.queue import budget_exceeded, ready_tasks from cloud.repos import RepoError, resolve_repo -from cloud.sandbox import DockerExecutor, LocalExecutor +from cloud.sandbox import CloudExecutor, DockerExecutor, LocalExecutor from cloud.sandbox import _default_runner as _sandbox_default_runner +from cloud.workers import heartbeat_worker, register_worker # 与 server.TERMINAL_STATUSES 对齐：这些状态的任务不再被认领或取消。 TERMINAL_STATUSES = frozenset( @@ -178,6 +180,11 @@ def _task_from_row(row): depends_on = [] if not isinstance(depends_on, list): depends_on = [depends_on] + def _col(name, default=None): + try: + return row[name] + except (KeyError, IndexError, TypeError): + return default return { "id": row["id"], "project_id": row["project_id"], @@ -190,6 +197,9 @@ def _task_from_row(row): "created_at": row["created_at"], "updated_at": row["updated_at"], "last_heartbeat": row["last_heartbeat"], + # 老库行缺 executor 列时给 None（表示任意 worker 可认领）， + # 不炸；列本身由 _ensure_result_columns 在写回时补。 + "executor": _col("executor"), } @@ -208,7 +218,8 @@ class Worker: cost_store=None, process_factory=None, task_timeout=1800, max_retries=1, retry_backoff_seconds=5.0, repos_dir="./repos", - repos_runner=None, heartbeat_file=None): + repos_runner=None, heartbeat_file=None, + worker_id=None, capacity=1): self.db_path = str(db_path) self.host_py_path = host_py_path self.config_path = config_path @@ -259,6 +270,39 @@ class Worker: self.heartbeat_file = heartbeat_file or \ os.environ.get("CLOUD_WORKER_HEARTBEAT_FILE") or \ "./worker-heartbeat.json" + # 多 worker 池身份：显式 --worker-id 优先（生产部署应配稳定 + # id）；未给时用“主机名-进程号”，重启后进程号会变，旧行 + # 自然变 stale，不要用随机 uuid 让注册表无法对应同一台机器。 + self.worker_id = worker_id or "%s-%d" % ( + socket.gethostname(), os.getpid()) + try: + self.capacity = int(capacity) + except (TypeError, ValueError): + self.capacity = 1 + + def register(self): + """向控制面注册本 worker（写 workers 表，失败只记日志）。 + + 注册与心跳共用同一张表：心跳就是带 last_seen 刷新的注册， + 未注册的 worker 发心跳时自动注册，不 404。 + """ + try: + register_worker( + self.db_path, self.worker_id, + executor=self._executor_name(), capacity=self.capacity) + except Exception as exc: + log_event(self.logger, "worker_register_failed", + message="worker register failed: %s" % exc) + + def _registry_heartbeat(self): + """向控制面写注册表心跳（与文件心跳并行，失败只记日志）。""" + try: + heartbeat_worker( + self.db_path, self.worker_id, + executor=self._executor_name(), capacity=self.capacity) + except Exception as exc: + log_event(self.logger, "worker_registry_heartbeat_failed", + message="worker registry heartbeat failed: %s" % exc) def _write_heartbeat(self): """写一次运维心跳：写失败只记日志，绝不炸 worker。 @@ -312,17 +356,34 @@ class Worker: conn.close() return _task_from_row(row) if row is not None else None - def _claim(self, task_id): - """CAS 认领：仅当任务仍为 queued 时标 running 并记心跳。""" + def _claim(self, task_id, executor=None): + """CAS 认领：仅当任务仍为 queued 时标 running 并记心跳。 + + executor 过滤（多 worker 池）：executor 给定时，WHERE 里 + 再加 (executor IS NULL OR executor = ?)，即只认领未指定 + 执行器、或指定为本 worker 执行器类型的任务；CAS 语义不变， + 仍是“仅当 queued 时标 running”，两个同池 worker 抢同一 + 任务时只有一个能拿到 rowcount == 1。 + """ now = time.time() conn = _db_connect(self.db_path) try: - cur = conn.execute( - "UPDATE tasks SET status = 'running'," - " last_heartbeat = ?, updated_at = ?" - " WHERE id = ? AND status = 'queued'", - (now, now, task_id), - ) + self._ensure_result_columns(conn) + if executor is None: + cur = conn.execute( + "UPDATE tasks SET status = 'running'," + " last_heartbeat = ?, updated_at = ?" + " WHERE id = ? AND status = 'queued'", + (now, now, task_id), + ) + else: + cur = conn.execute( + "UPDATE tasks SET status = 'running'," + " last_heartbeat = ?, updated_at = ?" + " WHERE id = ? AND status = 'queued'" + " AND (executor IS NULL OR executor = ?)", + (now, now, task_id, executor), + ) conn.commit() return cur.rowcount == 1 finally: @@ -355,7 +416,9 @@ class Worker: for name, coltype in { "result_text": "TEXT", "error_text": "TEXT", "actual_model": "TEXT", "started_at": "REAL", - "finished_at": "REAL", "child_pid": "INTEGER"}.items(): + "finished_at": "REAL", "child_pid": "INTEGER", + "executor": "TEXT", "pr_url": "TEXT", + "pr_number": "INTEGER"}.items(): if name not in cols: try: conn.execute( @@ -484,6 +547,12 @@ class Worker: ready = ready_tasks(tasks, completed_ids) except ValueError: return None + # 多 worker 池过滤：只取未指定执行器、或指定为本 worker + # 执行器类型的任务；_claim 里还有一道 WHERE 过滤兜底， + # pick_next 这里是第一道，避免 worker 老去抢别人的任务。 + mine = self._executor_name() + ready = [t for t in ready + if t.get("executor") in (None, "", mine)] return ready[0] if ready else None # -- 取消（server cancel 语义，见模块 docstring）------------------- @@ -572,9 +641,13 @@ class Worker: # -- 缺 key 门 ------------------------------------------------------ def _executor_name(self): - """执行器名：DockerExecutor -> docker，其它 -> local。""" + """执行器名：DockerExecutor -> docker，CloudExecutor -> cloud， + 其它 -> local。cloud 绝不能回落成 local，否则 cloud worker + 会认领 local 池的任务，池过滤就失效了。""" if isinstance(self.executor, DockerExecutor): return "docker" + if isinstance(self.executor, CloudExecutor): + return "cloud" return "local" def key_available(self): """只查存在性：名在环境里且值非空即 True；值本身绝不返回。""" @@ -800,11 +873,15 @@ class Worker: if task is None: return {"dispatched": False, "reason": "no_ready_task"} task_id = task["id"] - if not self._claim(task_id): + if not self._claim(task_id, executor=self._executor_name()): return {"dispatched": False, "reason": "already_claimed", "task_id": task_id} task = self.get_task(task_id) or task - self._audit(task["project_id"], "worker_claim", task_id) + self._audit(task["project_id"], "worker_claim", + "%s executor=%s" % (task_id, self._executor_name())) + # log_event 只接受 task_id/project_id/status/message 四个 + # 字段，executor 只进审计 detail，不进日志调用（否则 + # TypeError 会炸掉整个 run_one）。 log_event(self.logger, "worker_claim", task_id=task_id, project_id=task["project_id"], status="running") started_at = time.time() @@ -982,6 +1059,9 @@ class Worker: # 运维心跳：启动时先写一次，之后每轮结束再写，server 的 # /health 据此算 worker_ok（见 server.py 的健康检查）。 self._write_heartbeat() + # 多 worker 池注册：启动时注册一次，之后每轮结束写注册表 + # 心跳（见本轮结束处）；注册失败只记日志，不炸循环。 + self.register() started = time.time() all_results = [] try: @@ -1001,8 +1081,9 @@ class Worker: results = [] all_results.extend(results) # 每轮结束写心跳：哪怕本轮一个任务都没派，worker 活着 - # 这件事本身也要被 server 看到。 + # 这件事本身也要被 server 看到（文件心跳 + 注册表心跳）。 self._write_heartbeat() + self._registry_heartbeat() # 按 0.2 秒切片睡，期间可响应 stop_event/_loop_stop。 deadline = time.time() + interval while time.time() < deadline: @@ -1032,7 +1113,11 @@ class Worker: def _build_cli_executor(args, host_py): """CLI 的执行器构造：local 时返回 None（Worker 默认 LocalExecutor）， docker 时返回 DockerExecutor（env_key 名用 Worker 的解析规则， - 值本身绝不从这里传——DockerExecutor 只查存在性，-e 只传名）。""" + 值本身绝不从这里传——DockerExecutor 只查存在性，-e 只传名）， + cloud 时返回 CloudExecutor（通用 HTTP runner，未配 endpoint + 时执行即明确失败，绝不降级）。""" + if getattr(args, "executor", "local") == "cloud": + return CloudExecutor() if getattr(args, "executor", "local") != "docker": return None return DockerExecutor( @@ -1075,10 +1160,17 @@ def main(argv=None): parser.add_argument("--max-turns", type=int, default=None) parser.add_argument("--max-tasks", type=int, default=None) parser.add_argument("--executor", default="local", - choices=["local", "docker"], + choices=["local", "docker", "cloud"], help="执行器：local 调 host.py 子进程；docker" " 调 DockerExecutor（无 docker 时任务" - "明确失败，不静默降级）") + "明确失败，不静默降级）；cloud 调通用" + " HTTP runner（需 CLOUD_RUNNER_ENDPOINT）") + parser.add_argument("--worker-id", default=None, + help="多 worker 池身份（默认：主机名-进程号；" + "生产部署应配稳定 id，否则重启后旧行变 stale）") + parser.add_argument("--worker-capacity", type=int, default=1, + help="worker 并发容量声明（注册表里展示用，" + "本轮 run_loop 仍是串行一轮一任务）") parser.add_argument("--docker-image", default="codex-app-server:latest") parser.add_argument("--docker-cpus", default="2") parser.add_argument("--docker-memory", default="2g") @@ -1119,7 +1211,9 @@ def main(argv=None): max_retries=args.max_retries, retry_backoff_seconds=args.retry_backoff_seconds, repos_dir=args.repos_dir, - heartbeat_file=args.heartbeat_file) + heartbeat_file=args.heartbeat_file, + worker_id=args.worker_id, + capacity=args.worker_capacity) # CLI 启动清理兜底（常驻与单次都做）：清掉上一轮残留容器。 worker.cleanup_stale_containers() # CLI 起 worker 前先建表/迁移：全新库直接起循环不应因缺表而炸。 diff --git a/cloud/workers.py b/cloud/workers.py new file mode 100644 index 0000000..fb6a789 --- /dev/null +++ b/cloud/workers.py @@ -0,0 +1,166 @@ +"""Worker 注册表：多 worker 池的注册、心跳与存活判断。 + +零第三方依赖，只用 Python 标准库。 + +设计： +- 注册表存在控制面数据库的 workers 表里（worker 与 server 共用 + 同一个库，所以 worker 直接写库注册，不需要为此再发 HTTP； + HTTP 侧的 GET /workers 与 POST /workers/heartbeat 只是同一张 + 表的 API 面，给有 API key 的外部 worker 用）。 +- worker_id 默认由调用方给（生产 worker 应配稳定的 --worker-id）； + 未给时由 worker.py 用“主机名-进程号”生成——重启后进程号会变， + 旧行会自然变 stale（alive=false），不要用随机 uuid 让注册表 + 每次重启都多一行无法对应到同一台机器。 +- 心跳本身就是注册：heartbeat 时 worker_id 不存在就 INSERT， + 存在就 UPDATE，绝不因“未注册”返回 404 报错。 +- alive 定义：last_seen 在阈值内（默认 120 秒，可经环境变量 + CLOUD_WORKER_ALIVE_SECONDS 调，测试可传参覆盖）。 +- 本模块不存、不打任何密钥：注册信息只有 worker_id、executor、 + capacity、last_seen，没有任何凭证字段。 +""" + +import os +import time + +from cloud.db import connect as _db_connect + +DEFAULT_ALIVE_SECONDS = 120 + + +def alive_threshold_seconds(override=None): + """存活阈值：显式参数 > 环境变量 > 默认 120 秒。""" + if override is not None: + try: + return float(override) + except (TypeError, ValueError): + return float(DEFAULT_ALIVE_SECONDS) + try: + return float(os.environ.get( + "CLOUD_WORKER_ALIVE_SECONDS") or DEFAULT_ALIVE_SECONDS) + except (TypeError, ValueError): + return float(DEFAULT_ALIVE_SECONDS) + + +def ensure_workers_table(db_path): + """建 workers 表（不存在时建，已存在时不动）。""" + conn = _db_connect(str(db_path)) + try: + conn.execute( + "CREATE TABLE IF NOT EXISTS workers(" + "worker_id TEXT PRIMARY KEY, " + "executor TEXT, " + "capacity INTEGER, " + "last_seen REAL, " + "created_at REAL)" + ) + conn.commit() + finally: + conn.close() + + +def register_worker(db_path, worker_id, executor=None, capacity=1): + """注册（或刷新）一个 worker：不存在 INSERT，存在 UPDATE。 + + 返回注册行 dict。worker_id 为空时抛 ValueError，调用方 + （server 的 heartbeat 端点）据此返回 400。 + """ + if not worker_id or not isinstance(worker_id, str): + raise ValueError("worker_id required") + ensure_workers_table(db_path) + now = time.time() + conn = _db_connect(str(db_path)) + try: + cur = conn.execute( + "UPDATE workers SET executor = ?, capacity = ?," + " last_seen = ? WHERE worker_id = ?", + (executor, capacity, now, worker_id), + ) + if cur.rowcount == 0: + conn.execute( + "INSERT INTO workers(worker_id, executor, capacity," + " last_seen, created_at) VALUES(?, ?, ?, ?, ?)", + (worker_id, executor, capacity, now, now), + ) + conn.commit() + row = conn.execute( + "SELECT * FROM workers WHERE worker_id = ?", + (worker_id,)).fetchone() + finally: + conn.close() + return _worker_to_dict(row) + + +def heartbeat_worker(db_path, worker_id, executor=None, capacity=None): + """心跳：就是带 last_seen 刷新的注册。 + + executor/capacity 未给时保留旧值（已注册的 worker 只刷时间， + 不把 executor 洗成 None）；未注册时按给定值（或默认）注册。 + """ + if not worker_id or not isinstance(worker_id, str): + raise ValueError("worker_id required") + ensure_workers_table(db_path) + now = time.time() + conn = _db_connect(str(db_path)) + try: + row = conn.execute( + "SELECT * FROM workers WHERE worker_id = ?", + (worker_id,)).fetchone() + if row is None: + conn.execute( + "INSERT INTO workers(worker_id, executor, capacity," + " last_seen, created_at) VALUES(?, ?, ?, ?, ?)", + (worker_id, executor, capacity if capacity is not None + else 1, now, now), + ) + else: + old = _worker_to_dict(row) + conn.execute( + "UPDATE workers SET executor = ?, capacity = ?," + " last_seen = ? WHERE worker_id = ?", + (executor if executor is not None else old["executor"], + capacity if capacity is not None else old["capacity"], + now, worker_id), + ) + conn.commit() + row = conn.execute( + "SELECT * FROM workers WHERE worker_id = ?", + (worker_id,)).fetchone() + finally: + conn.close() + return _worker_to_dict(row) + + +def list_workers(db_path, alive_seconds=None): + """列出全部 worker，每行含 alive（last_seen 在阈值内）。""" + threshold = alive_threshold_seconds(alive_seconds) + ensure_workers_table(db_path) + now = time.time() + conn = _db_connect(str(db_path)) + try: + rows = conn.execute( + "SELECT * FROM workers ORDER BY worker_id").fetchall() + finally: + conn.close() + out = [] + for row in rows: + item = _worker_to_dict(row) + last = item.get("last_seen") + item["alive"] = bool( + last is not None and (now - last) <= threshold) + out.append(item) + return out + + +def _worker_to_dict(row): + def _get(name): + try: + return row[name] + except (KeyError, IndexError, TypeError): + return None + return { + "worker_id": _get("worker_id"), + "executor": _get("executor"), + "capacity": _get("capacity"), + "last_seen": _get("last_seen"), + "created_at": _get("created_at"), + }