租户配额与计费：按用户月度 USD 上限 + 建任务/分支/重试准入 402 + 管理/自查端点 + 控制台展示

cloud/server.py

@@ -623,6 +623,19 @@ def _ensure_environments_table(conn):
         " WHERE status != 'archived'")
 
 
+def _ensure_quotas_table(conn):
+    """幂等建配额表：server init_db 与配额查询都可调它，老库
+    第一次跑到这里时自动补表，不丢数据。quota_usd_micros 为
+    NULL 表示不限额（与 0 区分：0 是一分钱也不能花）；金额
+    单位与 cost_records 一致，都是 USD 微单位（1 美元 =
+    1_000_000），避免浮点。"""
+    conn.execute(
+        "CREATE TABLE IF NOT EXISTS user_quotas("
+        "user_id TEXT PRIMARY KEY, "
+        "quota_usd_micros BIGINT, "
+        "updated_at DOUBLE PRECISION)")
+
+
 _ENV_BRANCH_RE = re.compile(r"^[A-Za-z0-9_.\-/]+$")
 
 
@@ -683,6 +696,7 @@ def init_db(db_path):
             migrate_projects_table(conn)
             _ensure_followups_table(conn)
             _ensure_environments_table(conn)
+            _ensure_quotas_table(conn)
             _ensure_uploads_tables(conn)
             conn.commit()
             return
@@ -714,6 +728,7 @@ def init_db(db_path):
         migrate_projects_table(conn)
         _ensure_followups_table(conn)
         _ensure_environments_table(conn)
+        _ensure_quotas_table(conn)
         _ensure_uploads_tables(conn)
         # worker 每轮 poll 都按 status 取任务（见 load_tasks 的
         # WHERE status IN (...)），无索引时是全表扫描；任务表随
@@ -1098,6 +1113,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
                 self._archive_environment(user_id, rest[: -len("/archive")])
             else:
                 self._send_json(404, {"error": "not found"})
+        elif path == "/users/me/quota" and method == "GET":
+            self._user_quota_self(user_id)
         elif path == "/tasks" and method == "POST":
             self._create_task(user_id)
         elif path == "/uploads" and method == "POST":
@@ -1108,6 +1125,10 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
             self._admin_sweep()
         elif path == "/admin/backup" and method == "POST":
             self._admin_backup(user_id)
+        elif path == "/admin/quotas" and method == "GET":
+            self._admin_quota_get(user_id, query)
+        elif path == "/admin/quotas" and method == "POST":
+            self._admin_quota_set(user_id)
         elif path == "/admin/push/dispatch" and method == "POST":
             self._push_dispatch_endpoint(user_id)
         elif path.startswith("/tasks/"):
@@ -2017,6 +2038,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
             if not store.check_permission(user_id, project_id, "create_task"):
                 self._send_json(403, {"error": "forbidden"})
                 return
+        if self._reject_if_over_quota(user_id):
+            return
         # 环境绑定：environment_id 非空时必须是本项目、已 ready
         # 的环境；creating/failed/archived 一律拒绝并把状态写进
         # 报错（creating 时用户就知道该等 worker 准备好再建任务，
@@ -2466,6 +2489,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
                     user_id, task["project_id"], "create_task"):
                 self._send_json(403, {"error": "forbidden"})
                 return
+        if self._reject_if_over_quota(user_id):
+            return
         message = body.get("message")
         if not message or not isinstance(message, str) \
                 or not message.strip():
@@ -2800,6 +2825,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
                     user_id, task["project_id"], "approve_merge"):
                 self._send_json(403, {"error": "forbidden"})
                 return
+        if self._reject_if_over_quota(user_id):
+            return
         if task["status"] not in ("failed", "interrupted", "merge-failed"):
             self._send_json(409, {
                 "error": "only failed/interrupted tasks can be retried",
@@ -2834,6 +2861,169 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
             store.audit(user_id, task["project_id"], "task_retried", task_id)
         self._send_json(200, task)
 
+    # -- 租户配额与计费 ----------------------------------------------
+    # 计费口径：只算 cost_records 里 usd_micros 非空的实际成本
+    # 行（estimated 行 usd 为 NULL，SUM 自动跳过，不算已花
+    # 的钱）；归属按任务创建者 created_by，不按项目——一个
+    # 用户在多个项目里花的钱都算进他自己的月度账单，项目维
+    # 度的拆分由 _project_cost 另行提供。窗口是 UTC 自然月，
+    # 与 cost_records.day（UTC 日期）同口径。
+    def _quota_month(self):
+        from datetime import datetime, timezone
+        return datetime.now(timezone.utc).strftime("%Y-%m")
+
+    def _quota_status(self, user_id):
+        """返回 (quota_usd_micros|None, used_usd_micros, month)。
+        quota 为 None 表示不限额：老用户与未设配额的用户默认
+        不限额，行为与过去一致，不会因为上了配额功能就突然
+        被挡住。"""
+        month = self._quota_month()
+        conn = self._db()
+        try:
+            _ensure_quotas_table(conn)
+            conn.commit()
+            # cost_records 由 CostStore 懒建：全新库上它还不存在，
+            # 直接 JOIN 会报 no such table——先实例化一次建表
+            # （与 _project_cost 同款做法），再跑配额查询。
+            _cost_store = CostStore(self.server.db_path)
+            try:
+                _cost_store._conn.close()
+            except Exception:
+                pass
+            row = conn.execute(
+                "SELECT quota_usd_micros FROM user_quotas"
+                " WHERE user_id = ?",
+                (user_id,)).fetchone()
+            quota = row[0] if row is not None else None
+            used = conn.execute(
+                "SELECT COALESCE(SUM(c.usd_micros), 0)"
+                " FROM cost_records c JOIN tasks t ON c.task_id = t.id"
+                " WHERE t.created_by = ? AND c.day LIKE ?",
+                (user_id, month + "%")).fetchone()[0]
+            return quota, int(used or 0), month
+        finally:
+            conn.close()
+
+    def _quota_payload(self, user_id, quota, used, month):
+        return {"user_id": user_id, "month": month,
+                "quota_usd_micros": quota, "used_usd_micros": used,
+                "remaining_usd_micros": (quota - used)
+                if quota is not None else None}
+
+    def _check_quota(self, user_id):
+        """配额准入：放行返回 None；超限返回 (quota, used,
+        month)，调用方一律回 402——这是钱用完了，不是请求
+        太快，所以不用 429。注意：只做建任务/分支/重跑时的
+        准入校验；跑起来的回合中途超限不会被掐断（worker
+        侧没有按金额中途停转的钩子，那是下一步的事）。"""
+        quota, used, month = self._quota_status(user_id)
+        if quota is not None and used >= quota:
+            return quota, used, month
+        return None
+
+    def _reject_if_over_quota(self, user_id):
+        blocked = self._check_quota(user_id)
+        if blocked is not None:
+            quota, used, month = blocked
+            self._send_json(402, {"error": "quota_exceeded",
+                                  "month": month,
+                                  "quota_usd_micros": quota,
+                                  "used_usd_micros": used})
+            return True
+        return False
+
+    def _is_global_maintainer(self, user_id):
+        """全局管理判定：沿用 _admin_backup 的口径——在任一
+        项目里是 maintainer/owner 即算（没有单设的全局管理
+        员表；没有任何项目角色的用户无权改任何人的配额）。"""
+        conn = self._db()
+        try:
+            row = conn.execute(
+                "SELECT COUNT(*) FROM project_members"
+                " WHERE user_id = ? AND role IN ('maintainer', 'owner')",
+                (user_id,)).fetchone()
+            return (row[0] if row is not None else 0) > 0
+        finally:
+            conn.close()
+
+    def _user_exists(self, target):
+        conn = self._db()
+        try:
+            return conn.execute(
+                "SELECT 1 FROM users WHERE id = ?",
+                (target,)).fetchone() is not None
+        finally:
+            conn.close()
+
+    def _user_quota_self(self, user_id):
+        quota, used, month = self._quota_status(user_id)
+        self._send_json(200,
+                        self._quota_payload(user_id, quota, used, month))
+
+    def _admin_quota_get(self, user_id, query):
+        if not self._is_global_maintainer(user_id):
+            self._send_json(403, {"error": "forbidden"})
+            return
+        target = (query.get("user_id") or [None])[0]
+        if not target or not isinstance(target, str):
+            self._send_json(400, {"error": "user_id required"})
+            return
+        if not self._user_exists(target):
+            self._send_json(404, {"error": "user not found"})
+            return
+        quota, used, month = self._quota_status(target)
+        self._send_json(200,
+                        self._quota_payload(target, quota, used, month))
+
+    def _admin_quota_set(self, user_id):
+        if not self._is_global_maintainer(user_id):
+            self._send_json(403, {"error": "forbidden"})
+            return
+        body, err = self._read_json()
+        if err:
+            self._send_json(400, {"error": err})
+            return
+        target = body.get("user_id")
+        if not target or not isinstance(target, str):
+            self._send_json(400, {"error": "user_id required"})
+            return
+        if not self._user_exists(target):
+            self._send_json(404, {"error": "user not found"})
+            return
+        quota = body.get("quota_usd_micros")
+        if quota is not None:
+            try:
+                quota = int(quota)
+            except (TypeError, ValueError):
+                self._send_json(400, {
+                    "error": "quota_usd_micros must be an integer or null"})
+                return
+            if quota < 0:
+                self._send_json(400, {
+                    "error": "quota_usd_micros must be >= 0"})
+                return
+        now = time.time()
+        conn = self._db()
+        try:
+            cur = conn.execute(
+                "UPDATE user_quotas SET quota_usd_micros = ?,"
+                " updated_at = ? WHERE user_id = ?",
+                (quota, now, target))
+            if cur.rowcount != 1:
+                conn.execute(
+                    "INSERT INTO user_quotas(user_id, quota_usd_micros,"
+                    " updated_at) VALUES(?, ?, ?)",
+                    (target, quota, now))
+            conn.commit()
+        finally:
+            conn.close()
+        with _auth_store(self.server.db_path) as store:
+            store.audit(user_id, None, "quota_set",
+                        "%s:%s" % (target, quota))
+        quota, used, month = self._quota_status(target)
+        self._send_json(200,
+                        self._quota_payload(target, quota, used, month))
+
     # -- worker 池 -----------------------------------------------------
     def _workers_list(self, user_id):
         """列出注册 worker：任何已认证用户可读（worker 是基础设施