Merge pull request #5 from openqj/prod-grade-batch
Production-grade: webhook execution, Postgres backup, rate limit, worker timeout/retry
This commit is contained in:
9 files changed
+1366
-52
No files matched your search
+82
-4
@@ -273,6 +273,7 @@ psycopg,写明“需安装 psycopg”),绝不静默回落到 SQLite,
|
||||
| `POST /tasks/{id}/merge` | 该项目 `approve_merge`(maintainer 以上) | 只同步状态,不执行合并:仅 `completed` 可改为 `merged` 并记审计 `merge_task`,其它状态(含已 merged 的重复调用)返回 `409`;不存在 `404`,越权 `403`,无 key `401` |
|
||||
| `POST /tasks/{id}/heartbeat` | 该项目 `create_task` | 更新 `last_heartbeat = now`;状态为 `queued` 则改为 `running`;返回任务 |
|
||||
| `POST /admin/sweep` | 认证即可 | 调 `sweep_stale_tasks(db_path, threshold_seconds=300)`(body 可传 `threshold_seconds` 覆盖默认值);把 `running` 且 `last_heartbeat`(无心跳则用 `updated_at`)早于 `now - threshold` 的任务改为 `interrupted` 并记审计 `sweep_interrupted`;返回 `{"swept": [ids]}` |
|
||||
| `POST /admin/backup` | maintainer 以上 | Postgres 时调 `backup_postgres`、SQLite 时调 `backup_db`(body 可传 `out_dir`,否则用 `CLOUD_BACKUP_DIR` 或默认 `cloud-backups`);返回 `{backend, path, size_bytes, task_count}`,绝不返回数据库密码或备份文件内容 |
|
||||
| `POST /tasks/{id}/cost` | 该任务所属项目 `read` | 成本聚合:`{task_id, project_id, total_tokens, usd_micros, turns, fallback_count, actual_tokens, estimated_tokens}`;真数与估算分开;不存在 `404`,越权 `403` |
|
||||
| `POST /projects/{id}/memory` | 该项目 `create_task` | 写线程摘要 body `{thread_id?, goal, conclusion, files_changed?, open_issues?}`,返回 `201 {id}`,并记审计 `create_memory` |
|
||||
| `GET /projects/{id}/memory?top_k=5` | 该项目 `read` | 最近 active 摘要(默认 top_k=5)+ `prompt_context`(只含 top_k,不全量注入);不存在 `404`,越权 `403` |
|
||||
@@ -299,14 +300,34 @@ psycopg,写明“需安装 psycopg”),绝不静默回落到 SQLite,
|
||||
函数 `sweep_stale_tasks(db_path, threshold_seconds=300) -> list` 独立
|
||||
可单测,不依赖 HTTP;回收动作的审计 `user_id` 记为 `"system"`。
|
||||
|
||||
### webhook 端点(加分项,已实现;只解析,不执行任务)
|
||||
### webhook 端点(GitHub 已从“只解析”升级为“可执行且幂等”;Slack 仍只解析)
|
||||
|
||||
- `POST /webhooks/github`:用 `webhooks.verify_github_signature`
|
||||
验签(请求头 `X-Hub-Signature-256`),secret 从环境变量
|
||||
`GITHUB_WEBHOOK_SECRET` 读。未设 secret 返回 `503`;验签失败
|
||||
返回 `401`;成功返回 `parse_github_pr_comment` 的结果(PR 评论含
|
||||
`/fix-review` 时返回 fix_review 触发点,否则 `{"result": null}`)。
|
||||
真正执行 fix-review 由控制面后续调 host.py 完成,本端点不执行。
|
||||
返回 `401`(两者都在去重之前,避免未验签的重复送达污染去重表)。
|
||||
顺序:验签 -> delivery 去重 -> 解析 -> 执行建任务。
|
||||
- delivery 去重:请求头 `X-GitHub-Delivery` 做幂等,去重记录持久化
|
||||
到数据库新表 `webhook_deliveries`(delivery_id 主键、event、repo、
|
||||
created_at),不是只放内存;同一 delivery 第二次送达直接返回
|
||||
`200 {"duplicate": true, ...}`,不重复建任务、不重复记审计。
|
||||
并发/去重兜底:任务另有 `idempotency_key = webhook:<delivery_id>`,
|
||||
冲突时回读已建任务而不重复建。
|
||||
- PR 评论 `/fix-review` 执行:解析成功后在控制面数据库里建一个
|
||||
`queued` 任务,prompt 为 `[webhook /fix-review] repo=<repo> pr=<PR号>
|
||||
author=<作者>\n<评论原文>`,并记审计 `webhook_fix_review`,返回
|
||||
`task_id` 与 `project_id`;webhook 处理里不直接跑模型,模型执行
|
||||
仍是 worker 认领后做的事。任务归属按项目名 == repo full_name 匹配;
|
||||
没有匹配项目时不凭空建项目、不建任务,返回 `task_id: null` 并说明
|
||||
原因(worker 的 repo 仍是配置里的 repo,不是任意 GitHub 仓库的
|
||||
自动克隆)。
|
||||
- `check_run` / `check_suite` 解析:`parse_github_check_run` 只在
|
||||
conclusion 已定(success / failure)时返回结构化结果(repo、
|
||||
head_sha、conclusion、check_name),handler 直接返回该结果而不建
|
||||
任务、不改云端任务状态;结论进行中(None / queued / in_progress /
|
||||
neutral / cancelled 等)时返回 `{"result": null}`。CI 状态回填仍由
|
||||
`pr-checks` 主动拉取,webhook 只做通知解析,不要假装 webhook 已
|
||||
自动更新了云端任务状态。
|
||||
- `POST /webhooks/slack`:用 `verify_slack_signature` 验签(请求头
|
||||
`X-Slack-Signature` / `X-Slack-Request-Timestamp`,含 300 秒防重放),
|
||||
secret 从环境变量 `SLACK_SIGNING_SECRET` 读。未设返回 `503`;
|
||||
@@ -461,3 +482,60 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
|
||||
真实 docker build、真实 docker run、资源限额在真实负载下的
|
||||
表现、bwrap 在容器内的可用性、以及取消时容器确实消失。
|
||||
|
||||
## 生产级补齐(本轮新增:webhook 执行、备份、限流、超时/重试、清理兜底)
|
||||
|
||||
- Postgres 备份:`cloud/backup.py` 的 `backup_postgres(out_dir, keep=3,
|
||||
runner=None, pg_dump_cmd="pg_dump")`,数据库名/用户/密码只从
|
||||
`CLOUD_DB_NAME` / `CLOUD_DB_USER` / `CLOUD_DB_PASSWORD` 环境变量取,
|
||||
密码只经 `PGPASSWORD` 环境变量传给 pg_dump 子进程,绝不出现在
|
||||
命令参数、返回值或日志里;pg_dump 缺失或非零退出时明确抛错;
|
||||
成功返回备份文件路径(轮转只留最近 `keep` 份 `.sql`),server 的
|
||||
`POST /admin/backup`(需 maintainer 以上权限,body 可传 `out_dir`,
|
||||
否则用 `CLOUD_BACKUP_DIR` 或默认 `cloud-backups`)返回 backend、
|
||||
path、size_bytes、task_count,绝不返回数据库密码或备份文件内容;
|
||||
SQLite 分支(`backup_db`)原样不动。Postgres 备份仍需在真实环境
|
||||
验证一次真实 `pg_dump`(本轮只用 fake runner 单测)。
|
||||
- 限流:server 的 `RateLimiter`(内存滑动窗口,短锁只围住 deque
|
||||
剪枝/计数,不做全局锁、不串行正常请求),默认每分钟 120 次
|
||||
(环境变量 `CLOUD_RATE_LIMIT_PER_MINUTE` /
|
||||
`CLOUD_RATE_LIMIT_WINDOW_SECONDS` 可调,默认 120 / 60 秒);
|
||||
限流在健康/OAuth/webhook 之后、认证之前,超限返回 `429
|
||||
{"error": "rate_limited"}`,不记为认证失败;key 按 API key 的 sha256
|
||||
计(原始 key 不存),无 key 时按客户端 IP 计;测试可用
|
||||
`server.rate_limit_per_minute` 注入小阈值验证 429。
|
||||
- 健康与指标补齐:`GET /health`(不认证也通)除 `ok` 外返回
|
||||
`backend`(sqlite / postgres,取 `CLOUD_DB_BACKEND`,绝不返回
|
||||
连接串或密码)与 `db_ok`(一次轻量 `SELECT 1` 成功与否);
|
||||
`GET /metrics` 的 `by_status` 恒含七个状态键(queued / running /
|
||||
completed / failed / interrupted / merged / merge-failed),缺失
|
||||
状态为 0 而不是缺键。
|
||||
- 任务级超时:worker 的 `--task-timeout`(默认 1800 秒),与另外
|
||||
两个参数不要混淆:`--docker-timeout` 是执行器级的单次 docker run
|
||||
超时,`--max-seconds` 是 `run_loop` 的整体墙钟上限;三者互不覆盖,
|
||||
Docker 路径下 worker 取 `min(task_timeout, docker_timeout)`
|
||||
(docker_timeout 未设时取 task_timeout)作为本次 docker run 超时,
|
||||
Local popen 路径下 worker 在轮询循环里到截止时间就杀进程组;
|
||||
超时后真杀容器/进程组、任务置 `failed`(超时是失败,不是取消,
|
||||
绝不标 interrupted),并记 `worker_timeout` 审计与错误文本
|
||||
`task timeout after <秒> seconds`;Local 同步路径(假 runner / 无
|
||||
超时的 subprocess.run)无法中途抢占一个已阻塞的调用,超时只能靠
|
||||
执行器自身的 timeout,这是诚实边界,不要假装同步路径也能中途杀掉。
|
||||
- 失败重试与退避:worker 的 `--max-retries`(默认 1 次重试,0 表示
|
||||
不重试)与 `--retry-backoff-seconds`(默认 5 秒,测试可注入 0);
|
||||
取消不重试,已记 fallback 事件的失败不做同模型重试(fallback 本身
|
||||
就是模型级的重试机制),其它失败(含 worker 超时)重试前按固定
|
||||
退避等待;重试次数记在审计(`worker_retry`)与最终 `error_text`
|
||||
的 `[retries=N]` 前缀里,不另起新表;重试耗尽后任务才最终 failed,
|
||||
成本记录与 usage_turns 只按最终结果记一次,不按尝试次数重复记。
|
||||
- 容器清理兜底:`DockerExecutor.list_containers(prefix="codex-")`(
|
||||
`docker ps -a --format {{.Names}}`,失败返回 `[]`)与
|
||||
`cleanup_stale_containers()`(列出并逐个 kill+rm,单个失败只跳过,
|
||||
绝不抛异常),`Worker.cleanup_stale_containers()` 在 run_loop 启动
|
||||
与 CLI 启动时各调一次;取消/超时/失败路径上已有的容器清理保留
|
||||
(docker run `--rm`、超时 kill+rm、取消按名 kill)。
|
||||
- 仍未做(不要宣称已做):webhook 建的任务仍需 worker 认领执行,
|
||||
worker 的 repo 仍是配置里的 repo,不是任意 GitHub 仓库的自动克隆;
|
||||
`CloudExecutor` 仍是占位符;公网 https 入口仍需部署侧反向代理 +
|
||||
TLS;Docker 超时杀容器与真实 `pg_dump`、真实 GitHub webhook 送达
|
||||
仍需在真实环境验证(本轮只用 fake runner / 本地 HTTP 单测)。
|
||||
|
||||
+63
-2
@@ -15,7 +15,9 @@
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
from cloud.db import connect as _db_connect
|
||||
@@ -50,13 +52,72 @@ def backup_db(db_path, out_dir, keep=3, timestamp=None):
|
||||
return dest
|
||||
|
||||
|
||||
def backup_postgres(out_dir, keep=3, timestamp=None, runner=None,
|
||||
pg_dump_cmd="pg_dump"):
|
||||
"""Postgres 备份:调 pg_dump 写出 .sql 文件并轮转,返回文件路径。
|
||||
|
||||
连接参数只走 CLOUD_DB_* 环境变量(host/port/dbname/user/
|
||||
password),密码只以 PGPASSWORD 环境变量传给子进程,绝不
|
||||
出现在命令行参数、返回值或日志里。runner 可注入(单测用假
|
||||
runner):调用约定 runner(cmd, env) -> 具 returncode/stdout/
|
||||
stderr 的对象;未注入时用 subprocess.run。pg_dump 不存在时
|
||||
明确报错;pg_dump 非零退出时明确报错(stderr 截断,不含密码)。
|
||||
轮转规则与 SQLite 分支一致:只留最近 keep 份,keep < 1 报错。
|
||||
SQLite 分支原样不动,本函数不碰任何 .db 文件。
|
||||
"""
|
||||
if keep < 1:
|
||||
raise ValueError("keep 必须 >= 1,否则会删光所有备份")
|
||||
host = os.environ.get("CLOUD_DB_HOST") or "127.0.0.1"
|
||||
port = os.environ.get("CLOUD_DB_PORT") or "5432"
|
||||
dbname = os.environ.get("CLOUD_DB_NAME")
|
||||
user = os.environ.get("CLOUD_DB_USER")
|
||||
password = os.environ.get("CLOUD_DB_PASSWORD")
|
||||
if not dbname or not user:
|
||||
raise RuntimeError(
|
||||
"Postgres 备份需要 CLOUD_DB_NAME 与 CLOUD_DB_USER")
|
||||
if shutil.which(pg_dump_cmd) is None and runner is None:
|
||||
raise RuntimeError("pg_dump 不可用:本机找不到 %s" % pg_dump_cmd)
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
stamp = timestamp or time.strftime("%Y%m%d%H%M%S", time.localtime())
|
||||
dest = os.path.join(out_dir, "cloud-backup-%s.sql" % stamp)
|
||||
cmd = [pg_dump_cmd, "-h", host, "-p", str(port), "-U", user,
|
||||
"-d", dbname, "-f", dest]
|
||||
env = dict(os.environ)
|
||||
if password:
|
||||
env["PGPASSWORD"] = password
|
||||
if runner is not None:
|
||||
result = runner(cmd, env)
|
||||
else:
|
||||
result = subprocess.run(cmd, capture_output=True, text=True,
|
||||
env=env)
|
||||
if getattr(result, "returncode", 1) != 0:
|
||||
err = (getattr(result, "stderr", "") or "")[:500]
|
||||
raise RuntimeError("pg_dump 失败:%s" % err)
|
||||
if not os.path.exists(dest):
|
||||
raise RuntimeError("pg_dump 未生成备份文件:%s" % dest)
|
||||
backups = sorted(glob.glob(os.path.join(out_dir, "cloud-backup-*.sql")))
|
||||
for old in backups[:-keep] if len(backups) > keep else []:
|
||||
os.remove(old)
|
||||
return dest
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
parser = argparse.ArgumentParser(description="cloud control-plane backup")
|
||||
parser.add_argument("--db", required=True)
|
||||
parser.add_argument("--db", default=None,
|
||||
help="SQLite 库路径;postgres 模式可省略")
|
||||
parser.add_argument("--out", required=True)
|
||||
parser.add_argument("--keep", type=int, default=3)
|
||||
parser.add_argument("--backend", default="sqlite",
|
||||
choices=["sqlite", "postgres"],
|
||||
help="sqlite 时 --db 必填;postgres 时忽略 --db,"
|
||||
"连接参数走 CLOUD_DB_* 环境变量")
|
||||
args = parser.parse_args(argv)
|
||||
dest = backup_db(args.db, args.out, keep=args.keep)
|
||||
if args.backend == "postgres":
|
||||
dest = backup_postgres(args.out, keep=args.keep)
|
||||
else:
|
||||
if not args.db:
|
||||
parser.error("sqlite 模式需要 --db")
|
||||
dest = backup_db(args.db, args.out, keep=args.keep)
|
||||
print(dest)
|
||||
|
||||
|
||||
|
||||
@@ -405,6 +405,41 @@ class DockerExecutor(Executor):
|
||||
except TypeError:
|
||||
return self.runner(cmd)
|
||||
|
||||
def list_containers(self, prefix="codex-"):
|
||||
"""列出名字以 prefix 开头的容器(默认 codex-),失败返回 []。
|
||||
|
||||
用 `docker ps -a --format {{.Names}}`,输出按行取容器名;
|
||||
runner 抛异常或输出不可解析时返回空列表,调用方(worker
|
||||
启动清理)据此只记日志、不炸掉 worker。
|
||||
"""
|
||||
try:
|
||||
result = self._call_runner(
|
||||
["docker", "ps", "-a", "--format", "{{.Names}}"])
|
||||
except Exception:
|
||||
return []
|
||||
out = getattr(result, "stdout", "") or ""
|
||||
names = []
|
||||
for line in str(out).splitlines():
|
||||
name = line.strip()
|
||||
if name and name.startswith(prefix):
|
||||
names.append(name)
|
||||
return names
|
||||
|
||||
def cleanup_stale_containers(self, prefix="codex-"):
|
||||
"""清理上一轮残留的 codex- 容器:列出并逐个 kill+rm。
|
||||
|
||||
返回已清理的容器名列表;单个容器清理失败只跳过它(kill
|
||||
本身尽力而为),列出失败时返回 [],绝不抛异常。
|
||||
"""
|
||||
cleaned = []
|
||||
for name in self.list_containers(prefix=prefix):
|
||||
try:
|
||||
if self.kill_container(name):
|
||||
cleaned.append(name)
|
||||
except Exception:
|
||||
continue
|
||||
return cleaned
|
||||
|
||||
def kill_container(self, task_or_name):
|
||||
"""杀容器:docker kill + docker rm -f,尽力而为不抛异常。
|
||||
|
||||
|
||||
+290
-5
@@ -37,12 +37,15 @@ audit_log 三张表;本 server 用 projects/tasks 两张表,表名不冲突
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import html
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections import deque
|
||||
from contextlib import contextmanager
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
@@ -71,6 +74,8 @@ from cloud.push import (
|
||||
list_subscriptions as _push_list,
|
||||
)
|
||||
from cloud.webhooks import (
|
||||
WebhookDeliveryStore,
|
||||
parse_github_check_run,
|
||||
parse_github_pr_comment,
|
||||
parse_slack_command,
|
||||
verify_github_signature,
|
||||
@@ -81,6 +86,73 @@ TERMINAL_STATUSES = frozenset(
|
||||
{"completed", "failed", "interrupted", "merged", "merge-failed"}
|
||||
)
|
||||
DEFAULT_SWEEP_THRESHOLD_SECONDS = 300
|
||||
# 限流默认值:每 IP / 每 API key 每分钟 120 次;环境变量
|
||||
# CLOUD_RATE_LIMIT_PER_MINUTE / CLOUD_RATE_LIMIT_WINDOW_SECONDS 可调。
|
||||
# 限流是内存滑动窗口,不跨进程共享(多 worker 部署时每进程独立计数,
|
||||
# 这是诚实边界,不是全局限流)。
|
||||
DEFAULT_RATE_LIMIT_PER_MINUTE = 120
|
||||
DEFAULT_RATE_LIMIT_WINDOW_SECONDS = 60
|
||||
TASK_STATUS_KEYS = (
|
||||
"queued", "running", "completed", "failed",
|
||||
"interrupted", "merged", "merge-failed",
|
||||
)
|
||||
|
||||
|
||||
def _rate_limit_from_env(name, default):
|
||||
try:
|
||||
value = int(os.environ.get(name) or default)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
return value if value > 0 else default
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
"""内存滑动窗口限流:每 key 每窗口最多 per_minute 次。
|
||||
|
||||
只在 allow() 内部持一把短锁做 deque 修剪与计数,不持锁做
|
||||
任何 I/O,不会把正常请求串行化。key 由调用方给(本 server
|
||||
传“key:<API key 的 sha256>”或“ip:<客户端 IP>”,原始 API key
|
||||
绝不存进限流结构里)。窗口与阈值可在实例上改(测试注入小
|
||||
阈值验证 429 就靠改这两个属性)。
|
||||
"""
|
||||
|
||||
def __init__(self, per_minute=DEFAULT_RATE_LIMIT_PER_MINUTE,
|
||||
window_seconds=DEFAULT_RATE_LIMIT_WINDOW_SECONDS,
|
||||
time_fn=None):
|
||||
self.per_minute = per_minute
|
||||
self.window_seconds = window_seconds
|
||||
self._time_fn = time_fn or time.time
|
||||
self._lock = threading.Lock()
|
||||
self._hits = {}
|
||||
|
||||
def allow(self, key):
|
||||
now = self._time_fn()
|
||||
window = self.window_seconds
|
||||
limit = self.per_minute
|
||||
with self._lock:
|
||||
queue = self._hits.get(key)
|
||||
if queue is None:
|
||||
queue = deque()
|
||||
self._hits[key] = queue
|
||||
cutoff = now - window
|
||||
while queue and queue[0] <= cutoff:
|
||||
queue.popleft()
|
||||
if len(queue) >= limit:
|
||||
return False
|
||||
queue.append(now)
|
||||
# 键数兜底:长期运行的 server 不应让 _hits 无限增长;
|
||||
# 超过 10000 个 key 时清掉最久未用的空队列(不影响计数)。
|
||||
if len(self._hits) > 10000:
|
||||
for old_key in list(self._hits):
|
||||
if not self._hits[old_key]:
|
||||
del self._hits[old_key]
|
||||
break
|
||||
return True
|
||||
|
||||
|
||||
def _db_backend_name():
|
||||
"""当前数据库后端名:环境变量,未设则 sqlite;绝不返回连接串。"""
|
||||
return (os.environ.get("CLOUD_DB_BACKEND") or "sqlite").strip().lower()
|
||||
|
||||
|
||||
def _connect(db_path):
|
||||
@@ -343,7 +415,22 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
query = parse_qs(parsed.query)
|
||||
|
||||
if path == "/health" and method == "GET":
|
||||
self._send_json(200, {"ok": True, "service": "cloud-control-plane"})
|
||||
# 健康检查不走限流与认证;db_ok 只做一次轻量查询,
|
||||
# 绝不返回连接串、密码或任何凭据。
|
||||
db_ok = False
|
||||
try:
|
||||
conn = _connect(self.server.db_path)
|
||||
try:
|
||||
conn.execute("SELECT 1").fetchone()
|
||||
db_ok = True
|
||||
finally:
|
||||
conn.close()
|
||||
except Exception:
|
||||
db_ok = False
|
||||
self._send_json(200, {
|
||||
"ok": True, "service": "cloud-control-plane",
|
||||
"backend": _db_backend_name(), "db_ok": db_ok,
|
||||
})
|
||||
return
|
||||
if path.startswith("/oauth/"):
|
||||
# OAuth 登录/回调是 pre-auth 端点(用户还没本站 API key,
|
||||
@@ -373,6 +460,37 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
self._send_json(404, {"error": "not found"})
|
||||
return
|
||||
|
||||
# 限流在认证之前、健康/OAuth/webhook 之后:超限返回 429,
|
||||
# 不是 401(不要把“请求太频繁”记成“认证失败”)。
|
||||
# key:有 X-API-Key 时按 key 的 sha256 计(原始 key 不存),
|
||||
# 无 key 时按客户端 IP 计。
|
||||
limiter = getattr(self.server, "rate_limiter", None)
|
||||
if limiter is None:
|
||||
limiter = RateLimiter(
|
||||
per_minute=_rate_limit_from_env(
|
||||
"CLOUD_RATE_LIMIT_PER_MINUTE",
|
||||
DEFAULT_RATE_LIMIT_PER_MINUTE),
|
||||
window_seconds=_rate_limit_from_env(
|
||||
"CLOUD_RATE_LIMIT_WINDOW_SECONDS",
|
||||
DEFAULT_RATE_LIMIT_WINDOW_SECONDS))
|
||||
self.server.rate_limiter = limiter
|
||||
# 测试注入:server.rate_limit_per_minute / rate_limit_window_seconds
|
||||
# 改动后即时生效,不用重启 server。
|
||||
if getattr(self.server, "rate_limit_per_minute", None) is not None:
|
||||
limiter.per_minute = self.server.rate_limit_per_minute
|
||||
if getattr(self.server, "rate_limit_window_seconds", None) is not None:
|
||||
limiter.window_seconds = self.server.rate_limit_window_seconds
|
||||
api_key_header = self.headers.get("X-API-Key")
|
||||
if api_key_header:
|
||||
limit_key = "key:" + hashlib.sha256(
|
||||
api_key_header.encode("utf-8")).hexdigest()
|
||||
else:
|
||||
limit_key = "ip:" + (self.client_address[0]
|
||||
if self.client_address else "unknown")
|
||||
if not limiter.allow(limit_key):
|
||||
self._send_json(429, {"error": "rate_limited"})
|
||||
return
|
||||
|
||||
user_id = self._authenticate()
|
||||
if user_id is None:
|
||||
self._send_json(401, {"error": "unauthorized"})
|
||||
@@ -394,6 +512,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
self._list_tasks(user_id, query)
|
||||
elif path == "/admin/sweep" and method == "POST":
|
||||
self._admin_sweep()
|
||||
elif path == "/admin/backup" and method == "POST":
|
||||
self._admin_backup(user_id)
|
||||
elif path == "/admin/push/dispatch" and method == "POST":
|
||||
self._push_dispatch_endpoint(user_id)
|
||||
elif path.startswith("/tasks/"):
|
||||
@@ -1116,6 +1236,63 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
swept = sweep_stale_tasks(self.server.db_path, threshold)
|
||||
self._send_json(200, {"swept": swept})
|
||||
|
||||
def _admin_backup(self, user_id):
|
||||
"""POST /admin/backup:备份当前后端数据库,返回路径与规模信息。
|
||||
|
||||
权限:调用方须在至少一个项目里是 maintainer 以上(全局
|
||||
管理端点没有单个项目上下文,按“任一项目 maintainer+”
|
||||
判定;无此角色返回 403)。返回只含 backend、备份文件
|
||||
路径、文件大小与任务行数,绝不返回数据库密码、备份文件
|
||||
内容或任何凭据。out_dir 可由 body 指定,否则走环境变量
|
||||
CLOUD_BACKUP_DIR,未设则用 cloud-backups。
|
||||
"""
|
||||
body, err = self._read_json()
|
||||
if err:
|
||||
self._send_json(400, {"error": err})
|
||||
return
|
||||
conn = self._db()
|
||||
try:
|
||||
member_rows = conn.execute(
|
||||
"SELECT role FROM project_members WHERE user_id = ?",
|
||||
(user_id,),
|
||||
).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
role_levels = {"viewer": 0, "contributor": 1,
|
||||
"maintainer": 2, "owner": 3}
|
||||
is_maintainer = any(
|
||||
role_levels.get(row[0], 0) >= 2 for row in member_rows)
|
||||
if not is_maintainer:
|
||||
self._send_json(403, {"error": "forbidden"})
|
||||
return
|
||||
out_dir = body.get("out_dir") or os.environ.get(
|
||||
"CLOUD_BACKUP_DIR") or "cloud-backups"
|
||||
if not isinstance(out_dir, str) or not out_dir:
|
||||
self._send_json(400, {"error": "out_dir must be a string"})
|
||||
return
|
||||
backend = _db_backend_name()
|
||||
try:
|
||||
if backend == "postgres":
|
||||
from cloud.backup import backup_postgres as _backup_pg
|
||||
dest = _backup_pg(out_dir, keep=3)
|
||||
else:
|
||||
from cloud.backup import backup_db as _backup_sqlite
|
||||
dest = _backup_sqlite(self.server.db_path, out_dir, keep=3)
|
||||
size_bytes = os.path.getsize(dest)
|
||||
conn = self._db()
|
||||
try:
|
||||
task_count = conn.execute(
|
||||
"SELECT COUNT(*) FROM tasks").fetchone()[0]
|
||||
finally:
|
||||
conn.close()
|
||||
except Exception as exc:
|
||||
self._send_json(500, {"error": "backup failed: %s" % (exc,)})
|
||||
return
|
||||
self._send_json(200, {
|
||||
"backend": backend, "path": dest,
|
||||
"size_bytes": size_bytes, "task_count": task_count,
|
||||
})
|
||||
|
||||
def _metrics(self, user_id):
|
||||
"""可观测性:任务计数、成功率、平均耗时、fallback、成本合计。
|
||||
|
||||
@@ -1136,7 +1313,9 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
"SELECT project_id, status, started_at, finished_at"
|
||||
" FROM tasks"
|
||||
).fetchall()
|
||||
by_status = {}
|
||||
# 按状态分计数:七个状态键恒存在(缺失为 0),调用方
|
||||
# 不用再猜某个状态是“零”还是“没统计”。
|
||||
by_status = {key: 0 for key in TASK_STATUS_KEYS}
|
||||
durations = []
|
||||
for row in rows:
|
||||
if row["project_id"] not in member_ids:
|
||||
@@ -1413,6 +1592,21 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
|
||||
# -- webhooks(签名认证,不走 X-API-Key)---------------------------
|
||||
def _handle_github_webhook(self):
|
||||
"""GitHub webhook:验签 -> delivery 去重 -> 解析 -> 执行建任务。
|
||||
|
||||
顺序是故意的:secret 未配 503、签名无效 401(两者都在
|
||||
去重之前,避免未验签的重复送达污染去重表);同一
|
||||
X-GitHub-Delivery 第二次送达直接回 200 duplicate,不
|
||||
重复建任务、不重复记审计。/fix-review 解析成功后在
|
||||
控制面建一个 queued 任务(prompt 含 repo、PR 号、评论
|
||||
原文、作者),由 worker 认领执行——webhook 本身不跑模型。
|
||||
任务归属:按项目名 == repo full_name 匹配项目;没有匹配
|
||||
项目时不凭空建项目、不建任务,返回 task_id 为 None 并
|
||||
说明原因(worker 的 repo 仍是配置里的 repo,不是任意
|
||||
GitHub 仓库的自动克隆,见 README)。check_run /
|
||||
check_suite 只做通知解析,返回解析结果而不建任务、不改
|
||||
云端任务状态(CI 状态回填由 pr-checks 主动拉取)。
|
||||
"""
|
||||
raw = self._read_body()
|
||||
secret = os.environ.get("GITHUB_WEBHOOK_SECRET")
|
||||
if not secret:
|
||||
@@ -1430,12 +1624,93 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
if not isinstance(payload, dict):
|
||||
self._send_json(400, {"error": "body must be a JSON object"})
|
||||
return
|
||||
# 只解析不执行:/fix-review 的真正执行由控制面后续调 host.py 完成。
|
||||
event = self.headers.get("X-GitHub-Event")
|
||||
delivery_id = self.headers.get("X-GitHub-Delivery")
|
||||
repository = payload.get("repository") or {}
|
||||
repo_name = repository.get("full_name")
|
||||
store = WebhookDeliveryStore(self.server.db_path)
|
||||
try:
|
||||
if not store.claim(delivery_id, event=event, repo=repo_name):
|
||||
self._send_json(200, {
|
||||
"duplicate": True, "delivery_id": delivery_id,
|
||||
})
|
||||
return
|
||||
finally:
|
||||
try:
|
||||
store.close()
|
||||
except Exception:
|
||||
pass
|
||||
# check_run / check_suite:只解析,不建任务、不改状态。
|
||||
check = parse_github_check_run(payload)
|
||||
if check is not None:
|
||||
self._send_json(200, check)
|
||||
return
|
||||
# issue_comment 与 pull_request_review_comment:后者把
|
||||
# pull_request.number 归一化成 parser 要的 issue.number。
|
||||
result = parse_github_pr_comment(payload)
|
||||
if result is None and payload.get("pull_request") and \
|
||||
payload.get("comment"):
|
||||
normalized = dict(payload)
|
||||
normalized["issue"] = {
|
||||
"number": (payload["pull_request"] or {}).get("number")}
|
||||
result = parse_github_pr_comment(normalized)
|
||||
if result is None:
|
||||
self._send_json(200, {"result": None})
|
||||
else:
|
||||
self._send_json(200, result)
|
||||
return
|
||||
# /fix-review 执行:建 queued 任务(不跑模型,worker 认领)。
|
||||
conn = self._db()
|
||||
try:
|
||||
project_row = None
|
||||
if repo_name:
|
||||
project_row = conn.execute(
|
||||
"SELECT id FROM projects WHERE name = ?",
|
||||
(repo_name,),
|
||||
).fetchone()
|
||||
if project_row is None:
|
||||
self._send_json(200, dict(
|
||||
result, task_id=None, project_id=None,
|
||||
error="no matching project for repo"))
|
||||
return
|
||||
project_id = project_row[0]
|
||||
task_id = uuid.uuid4().hex
|
||||
now = time.time()
|
||||
prompt = (
|
||||
"[webhook /fix-review] repo=%s pr=%s author=%s\n%s"
|
||||
% (result.get("repo"), result.get("pr_number"),
|
||||
result.get("author"), result.get("comment")))
|
||||
try:
|
||||
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', ?, '[]', 0, 'system',"
|
||||
" ?, ?, NULL)",
|
||||
(task_id, project_id, prompt,
|
||||
"webhook:%s" % delivery_id if delivery_id else None,
|
||||
now, now),
|
||||
)
|
||||
conn.commit()
|
||||
except (sqlite3.IntegrityError, _PgIntegrityError):
|
||||
# 并发/去重兜底:同 delivery 已建过任务时回读它,
|
||||
# 不建重复任务、不重复记审计。
|
||||
row = conn.execute(
|
||||
"SELECT id FROM tasks WHERE project_id = ?"
|
||||
" AND idempotency_key = ?",
|
||||
(project_id,
|
||||
"webhook:%s" % delivery_id if delivery_id else None),
|
||||
).fetchone()
|
||||
if row:
|
||||
self._send_json(200, dict(
|
||||
result, task_id=row[0], project_id=project_id,
|
||||
duplicate=True))
|
||||
return
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
with _auth_store(self.server.db_path) as auth:
|
||||
auth.audit("system", project_id, "webhook_fix_review", task_id)
|
||||
self._send_json(200, dict(
|
||||
result, task_id=task_id, project_id=project_id))
|
||||
|
||||
def _handle_slack_webhook(self):
|
||||
raw = self._read_body()
|
||||
@@ -1488,6 +1763,16 @@ def create_server(db_path, host="127.0.0.1", port=8765, backend=None):
|
||||
server = ThreadingHTTPServer((host, port), CloudControlPlaneHandler)
|
||||
server.db_path = db_path
|
||||
server.daemon_threads = True
|
||||
# 限流器:内存滑动窗口,默认值与环境变量覆盖见 RateLimiter;
|
||||
# 测试可直接改 server.rate_limit_per_minute / window_seconds。
|
||||
server.rate_limiter = RateLimiter(
|
||||
per_minute=_rate_limit_from_env(
|
||||
"CLOUD_RATE_LIMIT_PER_MINUTE", DEFAULT_RATE_LIMIT_PER_MINUTE),
|
||||
window_seconds=_rate_limit_from_env(
|
||||
"CLOUD_RATE_LIMIT_WINDOW_SECONDS",
|
||||
DEFAULT_RATE_LIMIT_WINDOW_SECONDS))
|
||||
server.rate_limit_per_minute = None
|
||||
server.rate_limit_window_seconds = None
|
||||
# OAuth 回调注入点:生产为 None(走真实换 token + 真实 userinfo),
|
||||
# 测试可注入 fake http_client / userinfo_client;绝不注入真实 secret。
|
||||
server.oauth_http_client = None
|
||||
|
||||
@@ -0,0 +1,290 @@
|
||||
"""生产级 webhook / 备份 / 限流 / 指标单测。
|
||||
|
||||
只用标准库,不跑真实模型回合、不联网、不连真实 GitHub;
|
||||
webhook secret 均为单测内临时值,不是生产密钥。
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
try:
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.server import RateLimiter, create_server
|
||||
from cloud.webhooks import WebhookDeliveryStore, parse_github_check_run
|
||||
except ImportError:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.server import RateLimiter, create_server
|
||||
from cloud.webhooks import WebhookDeliveryStore, parse_github_check_run
|
||||
|
||||
|
||||
def _sig(secret, raw):
|
||||
return "sha256=" + hmac.new(
|
||||
secret.encode("utf-8"), raw, hashlib.sha256).hexdigest()
|
||||
|
||||
|
||||
class WebhookExecTestCase(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()
|
||||
# webhook 任务归属按项目名 == repo full_name 匹配。
|
||||
status, self.project = self._req(
|
||||
"POST", "/projects", {"name": "owner/repo"}, self.owner_key)
|
||||
self.assertEqual(status, 201, self.project)
|
||||
self.auth.set_member(
|
||||
self.project["id"], self.owner_id, self.viewer_id, "viewer")
|
||||
|
||||
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,
|
||||
raw_body=None):
|
||||
conn = http.client.HTTPConnection(
|
||||
"127.0.0.1", self.port, timeout=10)
|
||||
hdrs = dict(headers or {})
|
||||
if raw_body is not None:
|
||||
payload = raw_body
|
||||
elif body is not None:
|
||||
payload = json.dumps(body).encode("utf-8")
|
||||
hdrs.setdefault("Content-Type", "application/json")
|
||||
else:
|
||||
payload = None
|
||||
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 _webhook(self, payload, delivery="d-1", event="issue_comment",
|
||||
secret="test-secret"):
|
||||
raw = json.dumps(payload).encode("utf-8")
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"X-Hub-Signature-256": _sig(secret, raw),
|
||||
"X-GitHub-Event": event,
|
||||
"X-GitHub-Delivery": delivery,
|
||||
}
|
||||
with mock.patch.dict(
|
||||
os.environ, {"GITHUB_WEBHOOK_SECRET": secret}):
|
||||
return self._req(
|
||||
"POST", "/webhooks/github", raw_body=raw, headers=headers)
|
||||
|
||||
def _fix_payload(self):
|
||||
return {
|
||||
"action": "created",
|
||||
"comment": {"body": "请 /fix-review 看一下",
|
||||
"user": {"login": "alice"}},
|
||||
"repository": {"full_name": "owner/repo"},
|
||||
"issue": {"number": 42},
|
||||
}
|
||||
|
||||
def _tasks(self):
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
return conn.execute(
|
||||
"SELECT id, prompt, status FROM tasks").fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
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 test_fix_review_creates_queued_task_and_audit(self):
|
||||
status, body = self._webhook(self._fix_payload(), delivery="d-100")
|
||||
self.assertEqual(status, 200, body)
|
||||
self.assertEqual(body["type"], "fix_review")
|
||||
self.assertIsNotNone(body["task_id"])
|
||||
self.assertEqual(body["project_id"], self.project["id"])
|
||||
tasks = self._tasks()
|
||||
self.assertEqual(len(tasks), 1)
|
||||
task_id, prompt, task_status = tasks[0]
|
||||
self.assertEqual(task_id, body["task_id"])
|
||||
self.assertEqual(task_status, "queued")
|
||||
self.assertIn("owner/repo", prompt)
|
||||
self.assertIn("42", prompt)
|
||||
self.assertIn("请 /fix-review 看一下", prompt)
|
||||
self.assertIn("alice", prompt)
|
||||
self.assertIn("webhook_fix_review", self._audit_actions())
|
||||
|
||||
def test_duplicate_delivery_creates_no_second_task_or_audit(self):
|
||||
status, first = self._webhook(self._fix_payload(), delivery="d-dup")
|
||||
self.assertEqual(status, 200)
|
||||
self.assertIsNotNone(first["task_id"])
|
||||
status, second = self._webhook(self._fix_payload(), delivery="d-dup")
|
||||
self.assertEqual(status, 200)
|
||||
self.assertTrue(second["duplicate"])
|
||||
self.assertEqual(len(self._tasks()), 1)
|
||||
self.assertEqual(
|
||||
self._audit_actions().count("webhook_fix_review"), 1)
|
||||
|
||||
def test_fix_review_no_matching_project_creates_no_task(self):
|
||||
payload = self._fix_payload()
|
||||
payload["repository"] = {"full_name": "other/repo"}
|
||||
status, body = self._webhook(payload, delivery="d-nomatch")
|
||||
self.assertEqual(status, 200)
|
||||
self.assertIsNone(body["task_id"])
|
||||
self.assertEqual(self._tasks(), [])
|
||||
|
||||
def test_check_run_success_and_failure_parsed_no_task(self):
|
||||
for conclusion in ("success", "failure"):
|
||||
payload = {
|
||||
"repository": {"full_name": "owner/repo"},
|
||||
"check_run": {"name": "verify", "head_sha": "abc123",
|
||||
"conclusion": conclusion},
|
||||
}
|
||||
status, body = self._webhook(
|
||||
payload, delivery="d-check-%s" % conclusion,
|
||||
event="check_run")
|
||||
self.assertEqual(status, 200, body)
|
||||
self.assertEqual(body["type"], "check_run")
|
||||
self.assertEqual(body["repo"], "owner/repo")
|
||||
self.assertEqual(body["head_sha"], "abc123")
|
||||
self.assertEqual(body["conclusion"], conclusion)
|
||||
self.assertEqual(body["check_name"], "verify")
|
||||
self.assertEqual(self._tasks(), [])
|
||||
|
||||
def test_check_run_in_progress_returns_no_check_result(self):
|
||||
payload = {
|
||||
"repository": {"full_name": "owner/repo"},
|
||||
"check_run": {"name": "verify", "head_sha": "abc123",
|
||||
"conclusion": None},
|
||||
}
|
||||
status, body = self._webhook(
|
||||
payload, delivery="d-check-prog", event="check_run")
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(body, {"result": None})
|
||||
|
||||
def test_webhook_bad_signature_401_and_secret_missing_503(self):
|
||||
payload = self._fix_payload()
|
||||
raw = json.dumps(payload).encode("utf-8")
|
||||
with mock.patch.dict(
|
||||
os.environ, {"GITHUB_WEBHOOK_SECRET": "test-secret"}):
|
||||
status, _ = self._req(
|
||||
"POST", "/webhooks/github", raw_body=raw,
|
||||
headers={"Content-Type": "application/json",
|
||||
"X-Hub-Signature-256": "sha256=deadbeef",
|
||||
"X-GitHub-Event": "issue_comment",
|
||||
"X-GitHub-Delivery": "d-bad"})
|
||||
self.assertEqual(status, 401)
|
||||
with mock.patch.dict(os.environ, {}, clear=False):
|
||||
os.environ.pop("GITHUB_WEBHOOK_SECRET", None)
|
||||
status, _ = self._req(
|
||||
"POST", "/webhooks/github", raw_body=raw,
|
||||
headers={"Content-Type": "application/json",
|
||||
"X-Hub-Signature-256": _sig("test-secret", raw)})
|
||||
self.assertEqual(status, 503)
|
||||
|
||||
def test_admin_backup_sqlite_owner_ok_viewer_forbidden(self):
|
||||
out_dir = os.path.join(self.tmp.name, "backups")
|
||||
status, body = self._req(
|
||||
"POST", "/admin/backup", {"out_dir": out_dir}, self.owner_key)
|
||||
self.assertEqual(status, 200, body)
|
||||
self.assertEqual(body["backend"], "sqlite")
|
||||
self.assertTrue(os.path.exists(body["path"]))
|
||||
self.assertGreater(body["size_bytes"], 0)
|
||||
self.assertIn("task_count", body)
|
||||
self.assertNotIn("password", json.dumps(body))
|
||||
status, _ = self._req(
|
||||
"POST", "/admin/backup", {"out_dir": out_dir}, self.viewer_key)
|
||||
self.assertEqual(status, 403)
|
||||
status, _ = self._req(
|
||||
"POST", "/admin/backup", {"out_dir": out_dir}, key=None)
|
||||
self.assertEqual(status, 401)
|
||||
|
||||
def test_rate_limit_429_not_401(self):
|
||||
# 新限流器 + 小阈值注入:前 3 次通,第 4 次 429。
|
||||
self.server.rate_limiter = RateLimiter(per_minute=3)
|
||||
self.server.rate_limit_per_minute = 3
|
||||
for _ in range(3):
|
||||
status, _ = self._req(
|
||||
"GET", "/projects", key=self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
status, body = self._req("GET", "/projects", key=self.owner_key)
|
||||
self.assertEqual(status, 429)
|
||||
self.assertEqual(body, {"error": "rate_limited"})
|
||||
|
||||
def test_health_backend_and_db_ok_and_metrics_by_status(self):
|
||||
status, body = self._req("GET", "/health")
|
||||
self.assertEqual(status, 200)
|
||||
self.assertTrue(body["ok"])
|
||||
self.assertEqual(body["backend"], "sqlite")
|
||||
self.assertTrue(body["db_ok"])
|
||||
status, metrics = self._req(
|
||||
"GET", "/metrics", key=self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
for key in ("queued", "running", "completed", "failed",
|
||||
"interrupted", "merged"):
|
||||
self.assertIn(key, metrics["by_status"])
|
||||
|
||||
|
||||
class WebhookUnitTestCase(unittest.TestCase):
|
||||
def test_parse_check_run_only_success_failure(self):
|
||||
base = {"repository": {"full_name": "o/r"}}
|
||||
ok = dict(base, check_run={"name": "v", "head_sha": "h",
|
||||
"conclusion": "success"})
|
||||
self.assertEqual(parse_github_check_run(ok)["conclusion"], "success")
|
||||
bad = dict(base, check_run={"name": "v", "head_sha": "h",
|
||||
"conclusion": "neutral"})
|
||||
self.assertIsNone(parse_github_check_run(bad))
|
||||
self.assertIsNone(parse_github_check_run({}))
|
||||
suite = dict(base, check_suite={"head_sha": "h2",
|
||||
"conclusion": "failure"})
|
||||
parsed = parse_github_check_run(suite)
|
||||
self.assertEqual(parsed["conclusion"], "failure")
|
||||
self.assertEqual(parsed["head_sha"], "h2")
|
||||
|
||||
def test_delivery_store_claim_duplicate(self):
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
db_path = os.path.join(tmp.name, "c.db")
|
||||
store = WebhookDeliveryStore(db_path)
|
||||
try:
|
||||
self.assertTrue(store.claim("d1", event="issue_comment",
|
||||
repo="o/r"))
|
||||
self.assertFalse(store.claim("d1", event="issue_comment",
|
||||
repo="o/r"))
|
||||
self.assertTrue(store.claim("d2"))
|
||||
finally:
|
||||
store.close()
|
||||
tmp.cleanup()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,292 @@
|
||||
"""生产级 worker 超时 / 重试 / 容器清理兜底 + Postgres 备份单测。
|
||||
|
||||
只用标准库与本地 fake runner / 真实 sleep 进程,不跑模型、不联网;
|
||||
Postgres 备份用可注入的 fake pg_dump runner,不调真实 pg_dump,
|
||||
密码只经 PGPASSWORD 环境变量传给子进程(单测里不放真实密码)。
|
||||
"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
try:
|
||||
from cloud import backup as backup_mod
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.sandbox import DockerExecutor
|
||||
from cloud.worker import Worker
|
||||
except ImportError:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud import backup as backup_mod
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.sandbox import DockerExecutor
|
||||
from cloud.worker import Worker
|
||||
|
||||
|
||||
class _FakeRunner:
|
||||
"""同步 fake 执行器 runner:按 results 顺序返回结果。
|
||||
|
||||
LocalExecutor 的 runner 收的是命令列表、返回 SimpleNamespace
|
||||
(returncode/stdout/stderr),与真实 subprocess 同款接口;
|
||||
这里按顺序返回,不按任务 id 查表。
|
||||
"""
|
||||
|
||||
def __init__(self, results):
|
||||
self.results = list(results)
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, cmd):
|
||||
self.calls.append(list(cmd))
|
||||
spec = self.results[min(len(self.calls) - 1,
|
||||
len(self.results) - 1)]
|
||||
if isinstance(spec, Exception):
|
||||
raise spec
|
||||
return SimpleNamespace(
|
||||
returncode=0 if spec[0] == "completed" else 1,
|
||||
stdout=spec[1], stderr="")
|
||||
|
||||
|
||||
def _make_worker(tmp, runner=None, process_factory=None,
|
||||
task_timeout=1800, max_retries=1,
|
||||
retry_backoff_seconds=0):
|
||||
db_path = os.path.join(tmp, "control.db")
|
||||
# worker 本身不查项目成员关系:init_db + 直接插任务即可,
|
||||
# 不用走 AuthStore(其 create_project 只写成员表、不返项目)。
|
||||
from cloud.server import init_db
|
||||
init_db(db_path)
|
||||
project_id = "proj"
|
||||
task_id = "t1"
|
||||
conn = sqlite3.connect(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(?, ?, ?, 'queued', '[]', 0,"
|
||||
" 'system', 0, 0, NULL)",
|
||||
(task_id, project_id, "do it"))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
worker = Worker(db_path, "host.py", config_path=None, repo=None,
|
||||
runner=runner,
|
||||
env={"XCM_API_KEY": "test-placeholder"},
|
||||
env_key=None,
|
||||
process_factory=process_factory,
|
||||
task_timeout=task_timeout, max_retries=max_retries,
|
||||
retry_backoff_seconds=retry_backoff_seconds)
|
||||
task = {"id": task_id, "project_id": project_id,
|
||||
"prompt": "do it", "depends_on": "[]", "priority": 0}
|
||||
return worker, db_path, task
|
||||
|
||||
|
||||
def _task_row(db_path):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
return conn.execute(
|
||||
"SELECT status, error_text FROM tasks WHERE id='t1'").fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _audit_actions(db_path):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
return [r[0] for r in conn.execute(
|
||||
"SELECT action FROM audit_log").fetchall()]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class WorkerTimeoutTestCase(unittest.TestCase):
|
||||
def test_popen_timeout_kills_process_group_and_fails(self):
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
|
||||
def factory(task):
|
||||
return subprocess.Popen(
|
||||
["sleep", "30"], stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE, text=True, start_new_session=True)
|
||||
|
||||
worker, db_path, task = _make_worker(
|
||||
tmp.name, process_factory=factory, task_timeout=0.2,
|
||||
max_retries=0, retry_backoff_seconds=0)
|
||||
result = worker.run_one(task)
|
||||
self.assertEqual(result["status"], "failed")
|
||||
# 超时是失败,不是取消:库里必须是 failed,绝不能是 interrupted。
|
||||
status, error_text = _task_row(db_path)
|
||||
self.assertEqual(status, "failed")
|
||||
self.assertIn("timeout", error_text)
|
||||
self.assertIn("worker_timeout", _audit_actions(db_path))
|
||||
tmp.cleanup()
|
||||
|
||||
def test_docker_timeout_converts_to_failed_and_audit(self):
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
|
||||
def docker_runner(cmd, timeout=None):
|
||||
if cmd[1] in ("kill", "rm"):
|
||||
return SimpleNamespace(returncode=0, stdout="", stderr="")
|
||||
raise subprocess.TimeoutExpired(cmd, timeout)
|
||||
|
||||
worker, db_path, task = _make_worker(
|
||||
tmp.name, runner=_FakeRunner([]), max_retries=0,
|
||||
retry_backoff_seconds=0)
|
||||
# 测试占位 key(非真实密钥),只为通过执行器的“key 存在性”
|
||||
# 前置检查,让 runner 走到 TimeoutExpired 分支。
|
||||
worker.executor = DockerExecutor(
|
||||
runner=docker_runner, env={"XCM_API_KEY": "test"},
|
||||
timeout=5)
|
||||
# VM 上无 docker 二进制:只 mock 存在性检查(_check_docker),
|
||||
# 真正的 run 仍走 fake runner 的 TimeoutExpired 分支。
|
||||
with mock.patch("cloud.sandbox.shutil.which",
|
||||
return_value="/usr/bin/docker"):
|
||||
result = worker.run_one(task)
|
||||
self.assertEqual(result["status"], "failed")
|
||||
status, error_text = _task_row(db_path)
|
||||
self.assertEqual(status, "failed")
|
||||
self.assertIn("timeout", error_text)
|
||||
self.assertIn("worker_timeout", _audit_actions(db_path))
|
||||
tmp.cleanup()
|
||||
|
||||
|
||||
class WorkerRetryTestCase(unittest.TestCase):
|
||||
def test_generic_failure_retries_once_then_failed(self):
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
runner = _FakeRunner([("failed", "boom"), ("failed", "boom")])
|
||||
worker, db_path, task = _make_worker(
|
||||
tmp.name, runner=runner, max_retries=1,
|
||||
retry_backoff_seconds=0)
|
||||
result = worker.run_one(task)
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertEqual(len(runner.calls), 2)
|
||||
status, error_text = _task_row(db_path)
|
||||
self.assertEqual(status, "failed")
|
||||
self.assertIn("[retries=1]", error_text)
|
||||
self.assertIn("boom", error_text)
|
||||
self.assertIn("worker_retry", _audit_actions(db_path))
|
||||
tmp.cleanup()
|
||||
|
||||
def test_retry_success_recovers(self):
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
runner = _FakeRunner([("failed", "boom"), ("completed", "ok")])
|
||||
worker, db_path, task = _make_worker(
|
||||
tmp.name, runner=runner, max_retries=1,
|
||||
retry_backoff_seconds=0)
|
||||
result = worker.run_one(task)
|
||||
self.assertEqual(result["status"], "completed")
|
||||
self.assertEqual(len(runner.calls), 2)
|
||||
status, _ = _task_row(db_path)
|
||||
self.assertEqual(status, "completed")
|
||||
tmp.cleanup()
|
||||
|
||||
def test_max_retries_zero_means_no_retry(self):
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
runner = _FakeRunner([("failed", "boom")])
|
||||
worker, db_path, task = _make_worker(
|
||||
tmp.name, runner=runner, max_retries=0,
|
||||
retry_backoff_seconds=0)
|
||||
result = worker.run_one(task)
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertEqual(len(runner.calls), 1)
|
||||
_, error_text = _task_row(db_path)
|
||||
self.assertNotIn("[retries=", error_text)
|
||||
tmp.cleanup()
|
||||
|
||||
|
||||
class ContainerCleanupTestCase(unittest.TestCase):
|
||||
def _docker_executor(self, ps_stdout=None, fail_ps=False):
|
||||
calls = []
|
||||
|
||||
def runner(cmd, timeout=None):
|
||||
calls.append(list(cmd))
|
||||
if cmd[1] == "ps":
|
||||
if fail_ps:
|
||||
raise RuntimeError("docker down")
|
||||
return SimpleNamespace(
|
||||
returncode=0, stdout=ps_stdout or "", stderr="")
|
||||
return SimpleNamespace(returncode=0, stdout="", stderr="")
|
||||
|
||||
executor = DockerExecutor(runner=runner)
|
||||
return executor, calls
|
||||
|
||||
def test_cleanup_lists_by_prefix_and_removes(self):
|
||||
executor, calls = self._docker_executor(
|
||||
ps_stdout="codex-abc\nother\ncodex-def\n")
|
||||
cleaned = executor.cleanup_stale_containers()
|
||||
self.assertEqual(cleaned, ["codex-abc", "codex-def"])
|
||||
flat = [c for call in calls for c in call]
|
||||
self.assertNotIn("other", flat)
|
||||
# kill + rm 各调一次(每个容器)。
|
||||
self.assertEqual(
|
||||
sum(1 for call in calls if call[1] == "kill"), 2)
|
||||
self.assertEqual(
|
||||
sum(1 for call in calls if call[1] == "rm"), 2)
|
||||
|
||||
def test_cleanup_failure_does_not_crash_worker(self):
|
||||
executor, _ = self._docker_executor(fail_ps=True)
|
||||
self.assertEqual(executor.cleanup_stale_containers(), [])
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
worker, _, _ = _make_worker(tmp.name, runner=_FakeRunner([]))
|
||||
worker.executor = executor
|
||||
# worker 层兜底同样不抛异常。
|
||||
self.assertEqual(worker.cleanup_stale_containers(), [])
|
||||
tmp.cleanup()
|
||||
|
||||
def test_list_containers_unparsable_or_failed_returns_empty(self):
|
||||
executor, _ = self._docker_executor(ps_stdout="")
|
||||
self.assertEqual(executor.list_containers(), [])
|
||||
|
||||
|
||||
class PostgresBackupTestCase(unittest.TestCase):
|
||||
def test_backup_postgres_fake_runner_success(self):
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
created = {}
|
||||
|
||||
def fake_runner(cmd, env=None):
|
||||
self.assertNotIn("secret-pw", " ".join(cmd))
|
||||
self.assertEqual(env.get("PGPASSWORD"), "secret-pw")
|
||||
dest = cmd[cmd.index("-f") + 1]
|
||||
Path(dest).write_text("-- dummy dump\n")
|
||||
created["cmd"] = cmd
|
||||
return SimpleNamespace(returncode=0, stdout="", stderr="")
|
||||
|
||||
env = {"CLOUD_DB_NAME": "codexdb", "CLOUD_DB_USER": "codexsvc",
|
||||
"CLOUD_DB_PASSWORD": "secret-pw"}
|
||||
with mock.patch.dict(os.environ, env):
|
||||
dest = backup_mod.backup_postgres(
|
||||
tmp.name, runner=fake_runner)
|
||||
self.assertTrue(os.path.exists(dest))
|
||||
self.assertGreater(os.path.getsize(dest), 0)
|
||||
self.assertNotIn("secret-pw", dest)
|
||||
self.assertIn("pg_dump", created["cmd"][0])
|
||||
tmp.cleanup()
|
||||
|
||||
def test_backup_postgres_failure_raises_clear_error(self):
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
|
||||
def failing_runner(cmd, env=None):
|
||||
return SimpleNamespace(
|
||||
returncode=1, stdout="", stderr="pg_dump: boom")
|
||||
|
||||
env = {"CLOUD_DB_NAME": "codexdb", "CLOUD_DB_USER": "codexsvc"}
|
||||
with mock.patch.dict(os.environ, env):
|
||||
with self.assertRaises(RuntimeError) as ctx:
|
||||
backup_mod.backup_postgres(tmp.name, runner=failing_runner)
|
||||
self.assertIn("pg_dump", str(ctx.exception))
|
||||
tmp.cleanup()
|
||||
|
||||
def test_backup_postgres_missing_binary_raises(self):
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
env = {"CLOUD_DB_NAME": "codexdb"}
|
||||
with mock.patch.dict(os.environ, env):
|
||||
with self.assertRaises(RuntimeError):
|
||||
backup_mod.backup_postgres(
|
||||
tmp.name, pg_dump_cmd="nonexistent-pg-dump-xyz")
|
||||
tmp.cleanup()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -96,8 +96,12 @@ class ServerTestCase(unittest.TestCase):
|
||||
def test_health_no_auth(self):
|
||||
status, body = self._req("GET", "/health")
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(
|
||||
body, {"ok": True, "service": "cloud-control-plane"})
|
||||
# health 除 ok/service 外还返回 backend 与 db_ok(加法字段,
|
||||
# 不是改成功语义);旧的精确全等断言已随字段扩展而放宽。
|
||||
self.assertTrue(body["ok"])
|
||||
self.assertEqual(body["service"], "cloud-control-plane")
|
||||
self.assertEqual(body["backend"], "sqlite")
|
||||
self.assertTrue(body["db_ok"])
|
||||
|
||||
def test_no_key_401(self):
|
||||
status, body = self._req("GET", "/projects")
|
||||
|
||||
@@ -26,6 +26,8 @@ from cloud.db import connect as _db_connect
|
||||
__all__ = [
|
||||
"verify_github_signature",
|
||||
"parse_github_pr_comment",
|
||||
"parse_github_check_run",
|
||||
"WebhookDeliveryStore",
|
||||
"verify_slack_signature",
|
||||
"parse_slack_command",
|
||||
"SlackUserMap",
|
||||
@@ -94,6 +96,104 @@ def parse_github_pr_comment(payload):
|
||||
}
|
||||
|
||||
|
||||
def parse_github_check_run(payload):
|
||||
"""解析 GitHub check_run / check_suite 事件,返回结构化结果或 None。
|
||||
|
||||
只在结论已定的时候返回:check_run / check_suite 的 conclusion
|
||||
为 "success" 或 "failure" 时返回::
|
||||
|
||||
{"type": "check_run", "repo": ..., "head_sha": ...,
|
||||
"conclusion": ..., "check_name": ...}
|
||||
|
||||
其中 repo 取 repository full_name,head_sha 取 check_run /
|
||||
check_suite 的 head_sha,check_name 取 check_run 的 name
|
||||
(check_suite 没有单项名时为 None)。结论为 None / 进行中
|
||||
(如 queued / in_progress / neutral / cancelled 等)时返回 None,
|
||||
调用方不应据此建任务或改任务状态——CI 状态回填由 pr-checks
|
||||
主动拉取,webhook 只做通知解析。
|
||||
"""
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
for key in ("check_run", "check_suite"):
|
||||
run = payload.get(key)
|
||||
if not isinstance(run, dict):
|
||||
continue
|
||||
conclusion = run.get("conclusion")
|
||||
if conclusion not in ("success", "failure"):
|
||||
return None
|
||||
repository = payload.get("repository") or {}
|
||||
return {
|
||||
"type": "check_run",
|
||||
"repo": repository.get("full_name"),
|
||||
"head_sha": run.get("head_sha"),
|
||||
"conclusion": conclusion,
|
||||
"check_name": run.get("name"),
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
class WebhookDeliveryStore:
|
||||
"""GitHub webhook delivery 去重(持久化,不只放内存)。
|
||||
|
||||
表 webhook_deliveries:delivery_id(主键,X-GitHub-Delivery 头)、
|
||||
event、repo、created_at(epoch 秒)。claim() 第一次返回 True,
|
||||
同一 delivery_id 第二次返回 False,调用方应直接回 200 且标记
|
||||
duplicate,不重复建任务、不重复记审计。INSERT 冲突时先按
|
||||
IntegrityError 处理,再回查一次确认(PG 适配层抛的完整性错误
|
||||
类型与 SQLite 不同,这里按“插入失败后回查存在即重复”处理,
|
||||
不依赖具体异常类名)。
|
||||
"""
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = db_path
|
||||
self._conn = _db_connect(db_path)
|
||||
self._conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS webhook_deliveries("
|
||||
"delivery_id TEXT PRIMARY KEY, "
|
||||
"event TEXT, "
|
||||
"repo TEXT, "
|
||||
"created_at REAL)"
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def claim(self, delivery_id, event=None, repo=None):
|
||||
"""认领一次 delivery:首次 True,重复 False。空 id 不认领。"""
|
||||
if not delivery_id:
|
||||
return True
|
||||
row = self._conn.execute(
|
||||
"SELECT delivery_id FROM webhook_deliveries WHERE delivery_id = ?",
|
||||
(delivery_id,),
|
||||
).fetchone()
|
||||
if row is not None:
|
||||
return False
|
||||
try:
|
||||
self._conn.execute(
|
||||
"INSERT INTO webhook_deliveries"
|
||||
"(delivery_id, event, repo, created_at)"
|
||||
" VALUES(?, ?, ?, ?)",
|
||||
(delivery_id, event, repo, time.time()),
|
||||
)
|
||||
self._conn.commit()
|
||||
except Exception:
|
||||
try:
|
||||
self._conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
row = self._conn.execute(
|
||||
"SELECT delivery_id FROM webhook_deliveries"
|
||||
" WHERE delivery_id = ?",
|
||||
(delivery_id,),
|
||||
).fetchone()
|
||||
if row is not None:
|
||||
return False
|
||||
raise
|
||||
return True
|
||||
|
||||
def close(self):
|
||||
"""关闭数据库连接(可选,单测/退出时用)。"""
|
||||
self._conn.close()
|
||||
|
||||
|
||||
def verify_slack_signature(signing_secret, timestamp, body, signature_header):
|
||||
"""校验 Slack 请求签名(防重放)。
|
||||
|
||||
|
||||
+208
-39
@@ -46,6 +46,34 @@ queued/running 任务直接标 interrupted(已终态的任务不动),
|
||||
执行返回后会先查任务是否已被取消,已被取消的一律保持
|
||||
interrupted,不用执行结果覆盖它。
|
||||
|
||||
任务级超时与失败重试(生产级,见 Worker.__init__ 参数):
|
||||
- task_timeout(默认 1800 秒,CLI --task-timeout 可调)是
|
||||
worker 级的每任务截止时间,与下面两个参数不要混淆:
|
||||
--docker-timeout 是执行器级的单次 docker run 超时(Docker
|
||||
Executor 自身的 timeout),--max-seconds 是 run_loop 的
|
||||
整体墙钟上限(不是每任务)。三者互不覆盖:Docker 路径下
|
||||
worker 取 min(task_timeout, docker_timeout)(docker_timeout
|
||||
未设时取 task_timeout)作为本次 docker run 的超时,超时后
|
||||
真杀容器、任务置 failed(超时是失败,不是取消,绝不标
|
||||
interrupted);Local popen 路径下 worker 在轮询循环里到
|
||||
截止时间就杀进程组、同样置 failed;Local 同步路径(假
|
||||
runner / 无超时的 subprocess.run)无法在中途抢占一个已阻塞
|
||||
的调用,超时只能靠执行器自身的 timeout,这是诚实边界,
|
||||
不要假装同步路径也能中途杀掉。
|
||||
- max_retries(默认 1 次重试,CLI --max-retries 可调,0 表示
|
||||
不重试)是对模型任务失败的有限重试:取消不重试;已记
|
||||
fallback 事件的失败不做同模型重试(fallback 本身就是模型
|
||||
级的重试机制,再同模型重试会重复烧额度);其它失败(含
|
||||
worker 超时)重试前按 retry_backoff_seconds(默认 5 秒,
|
||||
测试可注入 0)等待,重试次数记在审计(worker_retry)与
|
||||
最终 error_text([retries=N] 前缀)里,不另起新表;重试
|
||||
耗尽后任务才最终 failed,成本记录与 usage_turns 只记一次
|
||||
(按最终结果记,不按尝试次数重复记)。
|
||||
- 容器清理兜底:Worker.cleanup_stale_containers() 列出
|
||||
codex- 前缀的残留容器并逐个 kill+rm,失败只记日志不炸
|
||||
worker;run_loop 启动时与 CLI 启动时各调一次(run_all /
|
||||
__init__ 里不调,避免单测构造 Worker 时多余 I/O)。
|
||||
|
||||
可注入:runner(假 CLI,见 sandbox.LocalExecutor)与 env 映射
|
||||
均可注入,单测不跑真实模型回合、不联网、不登录任何服务。
|
||||
|
||||
@@ -164,7 +192,9 @@ class Worker:
|
||||
runner=None, env=None, env_key=None,
|
||||
codex_config_path=None, model=None, router=None,
|
||||
max_turns=None, max_cost_usd=None, executor=None,
|
||||
cost_store=None, process_factory=None):
|
||||
cost_store=None, process_factory=None,
|
||||
task_timeout=1800, max_retries=1,
|
||||
retry_backoff_seconds=5.0):
|
||||
self.db_path = str(db_path)
|
||||
self.host_py_path = host_py_path
|
||||
self.config_path = config_path
|
||||
@@ -198,6 +228,10 @@ class Worker:
|
||||
self._running_container = None
|
||||
self._loop_stop = False
|
||||
self.logger = get_logger("cloud.worker")
|
||||
# 任务级超时 / 失败重试(见模块 docstring 的生产级说明)。
|
||||
self.task_timeout = task_timeout
|
||||
self.max_retries = max_retries
|
||||
self.retry_backoff_seconds = retry_backoff_seconds
|
||||
|
||||
# -- 成本存储(惰性建表,避免只认领不执行时也建表)--------------
|
||||
@property
|
||||
@@ -463,6 +497,28 @@ class Worker:
|
||||
project_id=task["project_id"], status="interrupted")
|
||||
return claimed
|
||||
|
||||
def cleanup_stale_containers(self):
|
||||
"""启动清理兜底:清掉上一轮残留的 codex- 容器。
|
||||
|
||||
只在执行器有 cleanup_stale_containers 时做(Docker);
|
||||
列出或删除失败只记日志,绝不抛异常炸掉 worker。
|
||||
返回已清理的容器名列表(非 Docker 执行器返回 [])。
|
||||
"""
|
||||
cleanup = getattr(self.executor, "cleanup_stale_containers", None)
|
||||
if cleanup is None:
|
||||
return []
|
||||
try:
|
||||
cleaned = cleanup()
|
||||
except Exception as exc:
|
||||
log_event(self.logger, "worker_cleanup_failed",
|
||||
message="cleanup stale containers failed: %s" % exc)
|
||||
return []
|
||||
if cleaned:
|
||||
log_event(self.logger, "worker_cleanup",
|
||||
message="cleaned stale containers: %s"
|
||||
% ",".join(cleaned))
|
||||
return cleaned
|
||||
|
||||
def _was_cancelled(self, task_id):
|
||||
if task_id in self._cancel_requested:
|
||||
return True
|
||||
@@ -493,12 +549,15 @@ class Worker:
|
||||
is _sandbox_default_runner
|
||||
and hasattr(self.executor, "popen_task"))
|
||||
|
||||
def _execute_via_process(self, task, kind="dispatch"):
|
||||
def _execute_via_process(self, task, kind="dispatch",
|
||||
deadline_seconds=None):
|
||||
"""进程模式执行一轮:启动、轮询、取消时杀进程组。
|
||||
|
||||
返回与同步路径同款 result dict(status/stdout/stderr/
|
||||
host_status/returncode)。取消时返回 status='cancelled'
|
||||
的标记,调用方据此保持 interrupted。
|
||||
的标记,调用方据此保持 interrupted;deadline_seconds
|
||||
到期时杀进程组并返回 status='timeout' 的标记,调用方
|
||||
据此置 failed(超时是失败,不是取消)。
|
||||
"""
|
||||
task_id = task["id"]
|
||||
if self.process_factory is not None:
|
||||
@@ -512,12 +571,19 @@ class Worker:
|
||||
except Exception:
|
||||
pass
|
||||
cancelled = False
|
||||
timed_out = False
|
||||
started = time.time()
|
||||
try:
|
||||
while proc.poll() is None:
|
||||
if task_id in self._cancel_requested or self._loop_stop:
|
||||
cancelled = True
|
||||
self._kill_proc_group(proc)
|
||||
break
|
||||
if deadline_seconds is not None and \
|
||||
time.time() - started >= deadline_seconds:
|
||||
timed_out = True
|
||||
self._kill_proc_group(proc)
|
||||
break
|
||||
time.sleep(0.05)
|
||||
try:
|
||||
out, err = proc.communicate(timeout=5)
|
||||
@@ -535,6 +601,11 @@ class Worker:
|
||||
"status": "cancelled", "stdout": out or "",
|
||||
"stderr": err or "", "returncode": proc.returncode,
|
||||
"host_status": None}
|
||||
if timed_out:
|
||||
return {"executor": "local", "task_id": task_id,
|
||||
"status": "timeout", "stdout": out or "",
|
||||
"stderr": err or "", "returncode": proc.returncode,
|
||||
"host_status": None}
|
||||
fake = SimpleNamespace(returncode=proc.returncode,
|
||||
stdout=out or "", stderr=err or "")
|
||||
if hasattr(self.executor, "_result_from_runner"):
|
||||
@@ -545,6 +616,85 @@ class Worker:
|
||||
"returncode": proc.returncode, "host_status": None}
|
||||
|
||||
# -- 执行 ----------------------------------------------------------
|
||||
def _run_attempt(self, task):
|
||||
"""单次执行尝试(不含认领、key 门、重试、成本与最终审计)。
|
||||
|
||||
返回 (result, error):result 是执行器结果 dict,error 是
|
||||
runner 异常时的文本(无异常时 None)。Docker 路径下把
|
||||
worker 级 task_timeout 与执行器级 docker_timeout 取较小
|
||||
值作为本次 docker run 超时(docker_timeout 未设时取
|
||||
task_timeout);Local popen 路径下把 task_timeout 作为
|
||||
轮询截止时间;Local 同步路径无法中途抢占,超时只能靠
|
||||
执行器自身 timeout(见模块 docstring 的诚实边界)。
|
||||
调用方负责:取消优先、超时转 failed、重试、成本与审计。
|
||||
"""
|
||||
task_id = task["id"]
|
||||
result = None
|
||||
error = None
|
||||
# Docker 同步路径:dispatch 阻塞在 docker run 里时,取消
|
||||
# 线程靠 _running_container 按名杀容器(见 cancel())。
|
||||
if hasattr(self.executor, "container_name"):
|
||||
try:
|
||||
self._running_container = \
|
||||
self.executor.container_name(task)
|
||||
self._running_task_id = task_id
|
||||
except Exception:
|
||||
self._running_container = None
|
||||
# Docker 超时取较小值:两个超时都不应静默延长另一个。
|
||||
saved_docker_timeout = None
|
||||
is_docker = isinstance(self.executor, DockerExecutor)
|
||||
if is_docker and self.task_timeout is not None:
|
||||
saved_docker_timeout = self.executor.timeout
|
||||
if saved_docker_timeout is None:
|
||||
self.executor.timeout = self.task_timeout
|
||||
else:
|
||||
try:
|
||||
self.executor.timeout = min(
|
||||
float(self.task_timeout),
|
||||
float(saved_docker_timeout))
|
||||
except (TypeError, ValueError):
|
||||
self.executor.timeout = self.task_timeout
|
||||
try:
|
||||
if self._use_popen():
|
||||
result = self._execute_via_process(
|
||||
task, deadline_seconds=self.task_timeout)
|
||||
combined = (result.get("stdout", "") or "") + \
|
||||
(result.get("stderr", "") or "")
|
||||
if result["status"] not in (
|
||||
"completed", "cancelled", "timeout") and \
|
||||
_ALREADY_EXISTS_MARKER in combined:
|
||||
result = self._execute_via_process(
|
||||
task, kind="run",
|
||||
deadline_seconds=self.task_timeout)
|
||||
else:
|
||||
result = self.executor.dispatch_task(task)
|
||||
combined = (result.get("stdout", "") or "") + \
|
||||
(result.get("stderr", "") or "")
|
||||
if result["status"] != "completed" and \
|
||||
_ALREADY_EXISTS_MARKER in combined:
|
||||
result = self.executor.run_task(task)
|
||||
except Exception as exc: # runner 异常也如实写回 failed
|
||||
result = {"executor": "local", "task_id": task_id,
|
||||
"status": "failed", "stdout": "", "returncode": None,
|
||||
"host_status": None, "stderr": ""}
|
||||
error = str(exc)
|
||||
finally:
|
||||
if is_docker:
|
||||
self.executor.timeout = saved_docker_timeout
|
||||
# 同步路径执行结束:清掉取消用的容器标记(取消路径里
|
||||
# 若已清过,这里再清一次无害)。
|
||||
self._running_container = None
|
||||
if self._running_task_id == task_id and \
|
||||
self._running_proc is None:
|
||||
self._running_task_id = None
|
||||
# Docker 超时标记:执行器已在超时时 kill+rm 容器并返回
|
||||
# failed(stdout 含 timeout),这里把状态统一成 timeout,
|
||||
# 调用方据此置 failed 并记 worker_timeout 审计。
|
||||
if result is not None and result.get("status") == "failed" and \
|
||||
"docker run timeout" in (result.get("stdout", "") or ""):
|
||||
result = dict(result, status="timeout")
|
||||
return result, error
|
||||
|
||||
def run_one(self, task=None):
|
||||
"""认领并执行一个任务,返回结果 dict(见模块 docstring)。
|
||||
|
||||
@@ -589,41 +739,40 @@ class Worker:
|
||||
fallback_count = 0
|
||||
result = None
|
||||
error = None
|
||||
# Docker 同步路径:dispatch 阻塞在 docker run 里时,取消
|
||||
# 线程靠 _running_container 按名杀容器(见 cancel())。
|
||||
if hasattr(self.executor, "container_name"):
|
||||
try:
|
||||
self._running_container = \
|
||||
self.executor.container_name(task)
|
||||
self._running_task_id = task_id
|
||||
except Exception:
|
||||
self._running_container = None
|
||||
try:
|
||||
if self._use_popen():
|
||||
result = self._execute_via_process(task)
|
||||
combined = (result.get("stdout", "") or "") + \
|
||||
(result.get("stderr", "") or "")
|
||||
if result["status"] not in ("completed", "cancelled") and \
|
||||
_ALREADY_EXISTS_MARKER in combined:
|
||||
result = self._execute_via_process(task, kind="run")
|
||||
else:
|
||||
result = self.executor.dispatch_task(task)
|
||||
combined = (result.get("stdout", "") or "") + \
|
||||
(result.get("stderr", "") or "")
|
||||
if result["status"] != "completed" and \
|
||||
_ALREADY_EXISTS_MARKER in combined:
|
||||
result = self.executor.run_task(task)
|
||||
except Exception as exc: # runner 异常也如实写回 failed
|
||||
result = {"executor": "local", "task_id": task_id,
|
||||
"status": "failed", "stdout": "", "returncode": None,
|
||||
"host_status": None, "stderr": ""}
|
||||
error = str(exc)
|
||||
# 同步路径执行结束:清掉取消用的容器标记(取消路径里
|
||||
# 若已清过,这里再清一次无害)。
|
||||
self._running_container = None
|
||||
if self._running_task_id == task_id and \
|
||||
self._running_proc is None:
|
||||
self._running_task_id = None
|
||||
retry_count = 0
|
||||
# 有限重试:取消不重试;已记 fallback 的失败不做同模型重试
|
||||
# (fallback 本身就是模型级重试机制);其它失败(含 worker
|
||||
# 超时)按 max_retries 重试,退避等待后重跑一次 _run_attempt。
|
||||
for attempt in range(self.max_retries + 1):
|
||||
result, error = self._run_attempt(task)
|
||||
if self._was_cancelled(task_id) or \
|
||||
result.get("status") == "cancelled":
|
||||
break
|
||||
if result.get("status") == "completed":
|
||||
break
|
||||
host_status_try = result.get("host_status")
|
||||
if self.router is not None and host_status_try is not None \
|
||||
and is_fallbackable(host_status_try) and \
|
||||
self.router.next_after(
|
||||
self.model, host_status_try) is not None:
|
||||
break
|
||||
if attempt < self.max_retries:
|
||||
retry_count += 1
|
||||
self._audit(
|
||||
task["project_id"], "worker_retry",
|
||||
"%s: attempt %d failed, retry %d/%d"
|
||||
% (task_id, attempt + 1, retry_count,
|
||||
self.max_retries))
|
||||
if self.retry_backoff_seconds:
|
||||
time.sleep(self.retry_backoff_seconds)
|
||||
continue
|
||||
break
|
||||
# 超时转 failed:超时是失败,不是取消;最终状态必须是
|
||||
# VALID_STATUSES 里的 failed,绝不能把 "timeout" 写进库。
|
||||
if result is not None and result.get("status") == "timeout":
|
||||
result = dict(result, status="failed")
|
||||
error = "task timeout after %s seconds" % self.task_timeout
|
||||
self._audit(task["project_id"], "worker_timeout", task_id)
|
||||
|
||||
# 取消优先:执行期间被 cancel 的任务保持 interrupted,
|
||||
# 绝不用执行结果覆盖它(见模块 docstring 的取消说明)。
|
||||
@@ -649,6 +798,8 @@ class Worker:
|
||||
if final != "completed":
|
||||
if error is None:
|
||||
error = stdout.strip() or "host execution failed"
|
||||
if retry_count:
|
||||
error = "[retries=%d] %s" % (retry_count, error)
|
||||
# fallback 事件:只记、不自动重试(见模块 docstring)。
|
||||
if self.router is not None and host_status is not None \
|
||||
and is_fallbackable(host_status):
|
||||
@@ -714,6 +865,9 @@ class Worker:
|
||||
_running_task_id 非空),先杀进程组、再标 interrupted,
|
||||
不让任务永远卡在 running。
|
||||
"""
|
||||
# 常驻启动清理兜底:清掉上一轮残留的 codex- 容器,
|
||||
# 失败只记日志(见 cleanup_stale_containers),不炸循环。
|
||||
self.cleanup_stale_containers()
|
||||
started = time.time()
|
||||
all_results = []
|
||||
try:
|
||||
@@ -807,6 +961,16 @@ def main(argv=None):
|
||||
"用于容器内 host.py 写 work/、host-state.json,"
|
||||
"以及挂载原仓库和 codex-home(同路径挂载,"
|
||||
"容器内配置无需改路径)")
|
||||
parser.add_argument("--task-timeout", type=float, default=1800,
|
||||
help="每任务截止时间(秒,默认 1800);超时后真杀"
|
||||
"容器/进程组、任务置 failed;与 --docker-timeout"
|
||||
"(单次 docker run 超时)取较小值,与"
|
||||
" --max-seconds(run_loop 整体墙钟上限)不同")
|
||||
parser.add_argument("--max-retries", type=int, default=1,
|
||||
help="模型任务失败重试次数(默认 1,0 表示不重试;"
|
||||
"取消与已记 fallback 的失败不重试)")
|
||||
parser.add_argument("--retry-backoff-seconds", type=float, default=5.0,
|
||||
help="重试前等待秒数(默认 5)")
|
||||
parser.add_argument("--loop", action="store_true",
|
||||
help="常驻循环,每 --interval 秒跑一轮")
|
||||
parser.add_argument("--interval", type=int, default=300)
|
||||
@@ -821,7 +985,12 @@ def main(argv=None):
|
||||
repo=args.repo, env_key=args.env_key,
|
||||
codex_config_path=args.codex_config, model=args.model,
|
||||
max_turns=args.max_turns,
|
||||
executor=_build_cli_executor(args, host_py))
|
||||
executor=_build_cli_executor(args, host_py),
|
||||
task_timeout=args.task_timeout,
|
||||
max_retries=args.max_retries,
|
||||
retry_backoff_seconds=args.retry_backoff_seconds)
|
||||
# CLI 启动清理兜底(常驻与单次都做):清掉上一轮残留容器。
|
||||
worker.cleanup_stale_containers()
|
||||
# CLI 起 worker 前先建表/迁移:全新库直接起循环不应因缺表而炸。
|
||||
try:
|
||||
from cloud.server import init_db as _init_db
|
||||
|
||||
Reference in new issue
Block a user