diff --git a/cloud/README.md b/cloud/README.md index 00b4c26..80edb28 100644 --- a/cloud/README.md +++ b/cloud/README.md @@ -88,7 +88,10 @@ client_secret（只走环境变量），不是无认证后门。 心跳即注册）、`POST /tasks/{id}/pr`（已 merged 任务的自动 PR， 默认关闭，缺条件时返回 pr_skipped 而不是失败；见下文“云端派发与 自动 PR”一节）。`POST /tasks` 新增可选 `executor` -（local/docker/cloud，未传表示任意 worker 可认领）； +（local/docker/cloud；未传且无环境继承时默认 docker——空 +executor（任意 worker 可认领）会让任务落到隔离较弱的 local +worker，是多用户下的隔离缺口；`CLOUD_DEFAULT_EXECUTOR` 可 +覆盖，空字符串恢复旧行为）； `POST /projects` 新增可选 `auto_pr`（0/1，默认 0）。 ### OAuth 真实浏览器登录（已实现，部署前置仍需用户动手） @@ -373,7 +376,7 @@ psycopg，写明“需安装 psycopg”），绝不静默回落到 SQLite， | `POST /tasks/{id}/cancel` | 该项目 `create_task` | `queued` / `running` 改为 `interrupted` 并记审计 `cancel_task`；已是终态（completed / failed / interrupted / merged / merge-failed）返回 `400`；不存在 `404` | | `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/sweep` | 任一项目 maintainer+（与备份管理同口径），否则 `403`；扫完记审计 `admin_sweep` | 调 `sweep_stale_tasks(db_path, threshold_seconds=300)`（body 可传 `threshold_seconds` 覆盖默认值，下限钳到 60 秒）；把 `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}`，绝不返回数据库密码或备份文件内容 | | `GET /tasks/{id}/cost` | 该任务所属项目 `read` | 成本聚合：`{task_id, project_id, total_tokens, usd_micros, turns, fallback_count, actual_tokens, estimated_tokens, cached_input_tokens, unpriced_actual_tokens}`；真数与估算分开，无价 actual 行进 `unpriced_actual_tokens`；不存在 `404`，越权 `403` | | `POST /projects/{id}/memory` | 该项目 `create_task` | 写线程摘要 body `{thread_id?, goal, conclusion, files_changed?, open_issues?}`，返回 `201 {id}`，并记审计 `create_memory` | @@ -956,11 +959,26 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳 - 准入：建任务、会话分支、重试在权限校验通过后、写库之前查 配额，本月用量 >= 配额时回 402 `quota_exceeded`（附 month、 quota_usd_micros、used_usd_micros），不建任务、不改状态。 - 老用户与未设配额的用户默认不限额，行为与过去一致。重试按 + 老用户与未设配额的用户默认不限额，行为与过去一致。但 + OAuth 新用户不再是“建号即不限额”：建号（created=True）时 + 写一行默认配额（`CLOUD_DEFAULT_QUOTA_USD_MICROS`，默认 + 10_000_000，即 $10/月；`INSERT OR IGNORE` 幂等，不覆盖管理 + 员后来改过的额度），并记 `quota_default` 审计。重试按 操作人查配额、费用仍记到原任务创建者名下（建任务与分支的 操作人就是创建者，没有这个错位）。系统任务（GitHub webhook 的 /fix-review，created_by 是 'system'）不记到任何用户名下、 不占配额；要把它算到项目所有者头上需另做归属。 +- 无价模型 token 上限：价目表目前只有两个 sol 模型有价，默认 + 模型 gpt-5.6-terra 等中转模型无价、成本行 usd 为 NULL，USD + 配额对它们永远不生效（按 0 计）。因此有配额行的用户另受一 + 道体积护栏：当月无价模型 token（input + output，已记账 + + 本轮，usd_micros IS NULL 的行）达到 + `CLOUD_QUOTA_UNPRICED_TOKEN_CAP`（默认 20_000_000/月）即 + 停——服务端建任务回 402 `quota_exceeded`（附 + unpriced_token_cap、unpriced_tokens_used），worker 侧认领/ + 重试/后续消息轮次同口径拦截。这道是防滥用的体积上限，不 + 是计价：成本报表里无价 token 仍进 unpriced_actual_tokens， + 绝不折算成假美元；无配额行的用户不受限。 - 执行时加固（worker 侧，`Worker._quota_exceeded`，口径与准入 一致：当月已记账 + 本轮已累计用量，按价目表换算）：认领后 开跑前再查一次（排队期间配额被花完/调低时不跑，任务标 diff --git a/cloud/console.html b/cloud/console.html index 41d80fc..5cf1265 100644 --- a/cloud/console.html +++ b/cloud/console.html @@ -151,9 +151,8 @@
diff --git a/cloud/cost.py b/cloud/cost.py index c9a68a1..0c119f8 100644 --- a/cloud/cost.py +++ b/cloud/cost.py @@ -111,6 +111,26 @@ def price_usd_micros(model, input_tokens, cached_input_tokens, return int(round(usd * 1000000)) +def unpriced_token_cap(): + """Monthly unpriced-token ceiling for users who have a quota row. + + Models with no price entry (today: the default gpt-5.6-terra and + the other relay models) record usd as unknown, so a USD quota + alone never binds them -- without this cap a quota user could + burn unlimited unpriced tokens. The cap is a volume guardrail, + NOT a price: cost reports still show those tokens as unpriced, + never converted to fake dollars. Only users with a quota row + are capped (no row still means unlimited, as before). Env + ``CLOUD_QUOTA_UNPRICED_TOKEN_CAP`` overrides the default + 20_000_000 tokens/month; unparsable values fall back to it. + """ + try: + return max(0, int(os.environ.get( + "CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "20000000"))) + except (TypeError, ValueError): + return 20000000 + + _USAGE_SUM_KEYS = ("input_tokens", "cached_input_tokens", "cache_write_input_tokens", "output_tokens", "reasoning_output_tokens", "total_tokens", "completions") diff --git a/cloud/server.py b/cloud/server.py index da15b2f..10b71c6 100644 --- a/cloud/server.py +++ b/cloud/server.py @@ -58,7 +58,7 @@ from urllib.parse import parse_qs, quote, urlparse from cloud.auth import AuthStore from cloud.coordinator import parse_parts as _coord_parse_parts from cloud.coordinator import validate_parts as _coord_validate_parts -from cloud.cost import CostStore +from cloud.cost import CostStore, unpriced_token_cap from cloud.db import connect as _db_connect from cloud.db import create_database as _create_database from cloud.db import init_postgres_schema as _init_postgres_schema @@ -936,6 +936,61 @@ def _ensure_quotas_table(conn): "updated_at DOUBLE PRECISION)") +def _default_executor(): + """新任务未指定执行器时的默认值：docker。 + + 空 executor（任意 worker 可认领）是多用户下的隔离缺口： + 任务可能落到隔离较弱的 local worker 上。默认 docker 后， + 只有显式 local/cloud 或系统创建（webhook/修复任务，executor + 为 NULL）才走别的路。环境变量 CLOUD_DEFAULT_EXECUTOR 可 + 覆盖：设为空字符串即恢复旧行为（None，任意 worker）， + 供单机运维在没有 docker worker 时兜底；非法值按 docker 计。 + """ + raw = os.environ.get("CLOUD_DEFAULT_EXECUTOR", "docker") + if raw is None: + return "docker" + raw = raw.strip() + if raw == "": + return None + if raw in ("local", "docker", "cloud"): + return raw + return "docker" + + +def _default_quota_usd_micros(): + """OAuth 新用户建号时写入的默认配额（微美元）。 + + 无配额行 = 不限额，所以 OAuth 自动建号过去等于无限额入驻。 + 新号现在建号即写一行默认配额（老用户无配额行不受影响， + 行为不变）。环境变量 CLOUD_DEFAULT_QUOTA_USD_MICROS 可 + 覆盖，默认 10_000_000（$10/月）；非法值按默认计。 + """ + try: + return max(0, int(os.environ.get( + "CLOUD_DEFAULT_QUOTA_USD_MICROS", "10000000"))) + except (TypeError, ValueError): + return 10000000 + + +def _apply_default_quota(conn, user_id): + """给新用户写默认配额行（已存在则不动，幂等）。 + + 只在 OAuth 建号（created=True）时调：INSERT OR IGNORE + 保证重复回调不覆盖管理员后来改过的额度。返回 True 表示 + 本次写了新行（调用方据此记 quota_default 审计）。 + """ + _ensure_quotas_table(conn) + cur = conn.execute( + "INSERT OR IGNORE INTO user_quotas(user_id, quota_usd_micros," + " updated_at) VALUES(?, ?, ?)", + (user_id, _default_quota_usd_micros(), time.time())) + try: + conn.commit() + except Exception: + pass + return cur.rowcount == 1 + + _ENV_BRANCH_RE = re.compile(r"^[A-Za-z0-9_.\-/]+$") @@ -1465,7 +1520,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): elif path == "/tasks" and method == "GET": self._list_tasks(user_id, query) elif path == "/admin/sweep" and method == "POST": - self._admin_sweep() + self._admin_sweep(user_id) elif path == "/admin/backup" and method == "POST": self._admin_backup(user_id) elif path == "/admin/quotas" and method == "GET": @@ -2582,6 +2637,11 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): return if executor is None and env["executor"]: executor = env["executor"] + if executor is None: + # 未指定且无环境继承时默认 docker：空 executor（任意 + # worker 可认领）会让任务落到隔离较弱的 local worker， + # 是多用户下的隔离缺口；显式 local/cloud 不受影响。 + executor = _default_executor() # 附件绑定（upload_id 由 POST /uploads 发出）：包必须存在、 # 归属当前用户、未使用、24 小时内、暂存目录还在；cloud # 执行器没有 worktree 可 seed，直接拒绝。校验在 INSERT @@ -3288,7 +3348,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): " VALUES(?, ?, ?, 'queued', NULL, '[]', ?, ?, ?, ?," " NULL, ?, NULL, NULL, ?, ?, ?, ?, ?, ?)", (new_id, task["project_id"], message, task["priority"], - user_id, now, now, task["executor"], + user_id, now, now, + task["executor"] or _default_executor(), task["goal_max_rounds"] or 3, task["model"], task["effort"], task["approval"], task["plan"] or 0, task_id), @@ -3809,6 +3870,56 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): "remaining_usd_micros": (quota - used) if quota is not None else None} + def _quota_unpriced_tokens(self, user_id, month): + """当月无价模型 token 数（usd_micros IS NULL 的成本行， + input + output 之和；cached 已含在 input 里，不重复计）。 + + 与 _quota_status 同款 CostStore 懒建表做法，全新库上 + cost_records 不存在时先建表再查。无配额行时调用方不 + 应以此挡人（见 _check_unpriced_token_cap）。 + """ + conn = self._db() + try: + _cost_store = CostStore(self.server.db_path) + try: + _cost_store._conn.close() + except Exception: + pass + row = conn.execute( + "SELECT COALESCE(SUM(c.input_tokens + c.output_tokens), 0)" + " FROM cost_records c JOIN tasks t ON c.task_id = t.id" + " WHERE t.created_by = ? AND c.day LIKE ?" + " AND c.usd_micros IS NULL", + (user_id, month + "%")).fetchone() + return int(row[0] if row is not None else 0) + finally: + conn.close() + + def _check_unpriced_token_cap(self, user_id): + """无价 token 上限准入：超限返回 (cap, used, month)，否则 None。 + + 只有有配额行的用户才受限（无配额行 = 不限额，与 USD + 口径一致）；USD 配额本身不计量无价模型，所以这道是 + 补 USD 口径的洞，不是第二套计费。""" + month = self._quota_month() + conn = self._db() + try: + _ensure_quotas_table(conn) + conn.commit() + row = conn.execute( + "SELECT quota_usd_micros FROM user_quotas" + " WHERE user_id = ?", + (user_id,)).fetchone() + if row is None: + return None + finally: + conn.close() + used = self._quota_unpriced_tokens(user_id, month) + cap = unpriced_token_cap() + if used >= cap: + return cap, used, month + return None + def _check_quota(self, user_id): """配额准入：放行返回 None；超限返回 (quota, used, month)，调用方一律回 402——这是钱用完了，不是请求 @@ -3829,6 +3940,14 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): "quota_usd_micros": quota, "used_usd_micros": used}) return True + capped = self._check_unpriced_token_cap(user_id) + if capped is not None: + cap, used, month = capped + self._send_json(402, {"error": "quota_exceeded", + "month": month, + "unpriced_token_cap": cap, + "unpriced_tokens_used": used}) + return True return False def _is_global_maintainer(self, user_id): @@ -4115,7 +4234,19 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): self._send_json(200, task) # -- 管理 ---------------------------------------------------------- - def _admin_sweep(self): + def _admin_sweep(self, user_id): + """POST /admin/sweep：把超阈值的陈旧 running 任务标 interrupted。 + + 权限：与 _admin_backup 同口径——调用方须在任一项目里 + 是 maintainer 以上，否则 403。过去这里没有任何权限检查， + 任意持 key 用户都能跨用户改任务状态。threshold_seconds + 下限钳到 60 秒：0 会把所有 running 任务一口气扫掉， + 只能是管理员的显式紧急操作才需要那么小的值，而 60 秒 + 以下的“陈旧”本就没有意义。扫完记 admin_sweep 审计。 + """ + 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}) @@ -4127,7 +4258,11 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): except (TypeError, ValueError): self._send_json(400, {"error": "threshold_seconds must be an integer"}) return + threshold = max(60, threshold) swept = sweep_stale_tasks(self.server.db_path, threshold) + with _auth_store(self.server.db_path) as store: + store.audit(user_id, "", "admin_sweep", + "threshold=%d swept=%d" % (threshold, len(swept))) self._send_json(200, {"swept": swept}) def _admin_backup(self, user_id): @@ -4366,6 +4501,10 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): auth, provider, identity) auth.audit(local_id, "", "oauth_login", "%s:%s" % (provider, identity.get("external_id"))) + if created and _apply_default_quota(auth.conn, local_id): + auth.audit(local_id, "", "quota_default", + "quota_usd_micros=%d" + % _default_quota_usd_micros()) if not created: auth.audit(local_id, "", "api_key_issued", provider) return local_id, created, api_key_plain, rotated, identity, False diff --git a/cloud/tests/test_multiuser_gaps.py b/cloud/tests/test_multiuser_gaps.py new file mode 100644 index 0000000..aecf232 --- /dev/null +++ b/cloud/tests/test_multiuser_gaps.py @@ -0,0 +1,279 @@ +"""多用户生产级缺口批单测：sweep 鉴权、默认执行器、新用户默认配额、 +无价模型 token 上限（服务端准入）。 + +只用本地临时端口 + 临时 SQLite，不跑真实模型，不联网。 +环境变量（默认执行器、默认配额额、token 上限）是服务端每次 +请求现读的，测试里用 try/finally 改回，不影响其他测试。 +""" + +import http.client +import json +import os +import sqlite3 +import sys +import tempfile +import threading +import time +import unittest +from datetime import datetime, timezone +from pathlib import Path + +try: + from cloud.auth import AuthStore + from cloud.cost import CostStore + from cloud.server import _apply_default_quota, _default_quota_usd_micros + from cloud.server import create_server +except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from cloud.auth import AuthStore + from cloud.cost import CostStore + from cloud.server import _apply_default_quota, _default_quota_usd_micros + from cloud.server import create_server + + +def _req(port, method, path, body=None, key=None): + conn = http.client.HTTPConnection("127.0.0.1", port, timeout=10) + hdrs = {} + payload = None + if body is not None: + payload = json.dumps(body).encode("utf-8") + hdrs["Content-Type"] = "application/json" + if key is not None: + hdrs["X-API-Key"] = key + try: + conn.request(method, path, body=payload, headers=hdrs) + resp = conn.getresponse() + data = resp.read() + finally: + conn.close() + try: + parsed = json.loads(data.decode("utf-8")) if data else None + except ValueError: + parsed = None + return resp.status, parsed + + +def _utc_today(): + return datetime.now(timezone.utc).strftime("%Y-%m-%d") + + +def _setenv(name, value): + old = os.environ.get(name) + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + return old + + +class MultiuserGapsTestCase(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.outsider_id, self.outsider_key = self.auth.create_user("outsider") + 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() + s, proj = _req(self.port, "POST", "/projects", + {"name": "gproj"}, self.owner_key) + self.assertEqual(s, 201, proj) + self.project_id = proj["id"] + 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() + + # -- helpers ------------------------------------------------------- + def _create_task(self, key=None, body_extra=None): + body = {"project_id": self.project_id, "prompt": "gap probe"} + if body_extra: + body.update(body_extra) + return _req(self.port, "POST", "/tasks", body, + key if key is not None else self.owner_key) + + def _insert_quota(self, user_id, micros): + conn = sqlite3.connect(self.db_path) + try: + conn.execute( + "CREATE TABLE IF NOT EXISTS user_quotas(" + "user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT," + " updated_at DOUBLE PRECISION)") + conn.execute( + "INSERT OR REPLACE INTO user_quotas" + "(user_id, quota_usd_micros, updated_at) VALUES(?, ?, ?)", + (user_id, micros, time.time())) + conn.commit() + finally: + conn.close() + + def _insert_unpriced_cost(self, task_id, input_tokens, output_tokens): + conn = sqlite3.connect(self.db_path) + try: + conn.execute( + "INSERT INTO tasks(id, project_id, prompt, status," + " idempotency_key, depends_on, priority, created_by," + " created_at, updated_at, last_heartbeat)" + " VALUES(?, ?, ?, 'completed', NULL, '[]', 0, ?, ?, ?, NULL)", + (task_id, self.project_id, "donor", self.owner_id, + time.time(), time.time())) + conn.commit() + finally: + conn.close() + # CostStore.record 建表并按 UTC 今天记 day，与服务端查询口径一致。 + store = CostStore(self.db_path) + try: + store.record(task_id, self.project_id, "gpt-5.6-terra", + input_tokens=input_tokens, + output_tokens=output_tokens, + usd_micros=None, source="actual") + finally: + try: + store._conn.close() + except Exception: + pass + + def _make_running_task(self): + s, task = self._create_task() + self.assertEqual(s, 201, task) + s, _ = _req(self.port, "POST", + "/tasks/%s/heartbeat" % task["id"], {}, + self.owner_key) + self.assertEqual(s, 200) + return task["id"] + + # -- sweep 鉴权 ----------------------------------------------------- + def test_sweep_viewer_and_outsider_forbidden(self): + for key in (self.viewer_key, self.outsider_key): + s, body = _req(self.port, "POST", "/admin/sweep", {}, + key) + self.assertEqual(s, 403, body) + + def test_sweep_owner_maintainer_ok_and_sweeps_stale(self): + tid = self._make_running_task() + conn = sqlite3.connect(self.db_path) + try: + conn.execute( + "UPDATE tasks SET last_heartbeat = ? WHERE id = ?", + (time.time() - 3600, tid)) + conn.commit() + finally: + conn.close() + s, body = _req(self.port, "POST", "/admin/sweep", + {"threshold_seconds": 300}, self.owner_key) + self.assertEqual(s, 200, body) + self.assertIn(tid, body["swept"]) + + def test_sweep_threshold_clamped_to_60(self): + tid = self._make_running_task() + s, body = _req(self.port, "POST", "/admin/sweep", + {"threshold_seconds": 0}, self.owner_key) + self.assertEqual(s, 200, body) + self.assertNotIn(tid, body["swept"]) + + # -- 默认执行器 ------------------------------------------------------ + def test_default_executor_is_docker(self): + s, task = self._create_task() + self.assertEqual(s, 201, task) + self.assertEqual(task["executor"], "docker") + + def test_explicit_executor_kept(self): + s, task = self._create_task(body_extra={"executor": "local"}) + self.assertEqual(s, 201, task) + self.assertEqual(task["executor"], "local") + + def test_default_executor_env_override(self): + old = _setenv("CLOUD_DEFAULT_EXECUTOR", "local") + try: + s, task = self._create_task() + self.assertEqual(s, 201, task) + self.assertEqual(task["executor"], "local") + finally: + _setenv("CLOUD_DEFAULT_EXECUTOR", old) + old = _setenv("CLOUD_DEFAULT_EXECUTOR", "") + try: + s, task = self._create_task() + self.assertEqual(s, 201, task) + self.assertIsNone(task["executor"]) + finally: + _setenv("CLOUD_DEFAULT_EXECUTOR", old) + + # -- 新用户默认配额 --------------------------------------------------- + def test_apply_default_quota_idempotent(self): + conn = sqlite3.connect(self.db_path) + try: + self.assertTrue(_apply_default_quota(conn, "newbie")) + row = conn.execute( + "SELECT quota_usd_micros FROM user_quotas" + " WHERE user_id = 'newbie'").fetchone() + self.assertEqual(row[0], _default_quota_usd_micros()) + # 第二次不再写，不覆盖管理员后来改过的额度。 + conn.execute( + "UPDATE user_quotas SET quota_usd_micros = 12345" + " WHERE user_id = 'newbie'") + conn.commit() + self.assertFalse(_apply_default_quota(conn, "newbie")) + row = conn.execute( + "SELECT quota_usd_micros FROM user_quotas" + " WHERE user_id = 'newbie'").fetchone() + self.assertEqual(row[0], 12345) + finally: + conn.close() + + def test_default_quota_env_override(self): + old = _setenv("CLOUD_DEFAULT_QUOTA_USD_MICROS", "5000") + try: + self.assertEqual(_default_quota_usd_micros(), 5000) + finally: + _setenv("CLOUD_DEFAULT_QUOTA_USD_MICROS", old) + + # -- 无价 token 上限（服务端准入）-------------------------------------- + def test_unpriced_token_cap_blocks_create(self): + self._insert_quota(self.owner_id, 10 ** 12) # USD 很大，不挡 + for i in range(7): # 7 * 150 = 1050 token + self._insert_unpriced_cost("donor-%d" % i, 100, 50) + old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000") + try: + s, body = self._create_task() + self.assertEqual(s, 402, body) + self.assertEqual(body["error"], "quota_exceeded") + self.assertEqual(body["unpriced_tokens_used"], 1050) + finally: + _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old) + + def test_unpriced_token_cap_under_cap_passes(self): + self._insert_quota(self.owner_id, 10 ** 12) + for i in range(6): # 6 * 150 = 900 < 1000 + self._insert_unpriced_cost("donor-%d" % i, 100, 50) + old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000") + try: + s, task = self._create_task() + self.assertEqual(s, 201, task) + finally: + _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old) + + def test_unpriced_tokens_unlimited_without_quota_row(self): + for i in range(7): + self._insert_unpriced_cost("donor-%d" % i, 100, 50) + old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000") + try: + s, task = self._create_task() + self.assertEqual(s, 201, task) + finally: + _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old) + + +if __name__ == "__main__": + unittest.main() diff --git a/cloud/tests/test_worker_quota.py b/cloud/tests/test_worker_quota.py index 1fb6017..556abe6 100644 --- a/cloud/tests/test_worker_quota.py +++ b/cloud/tests/test_worker_quota.py @@ -80,6 +80,26 @@ def _insert_donor_cost(db_path, micros): pass +def _insert_unpriced_cost(db_path, tokens): + """给 tester 记一笔无价成本：已完成任务 + usd_micros NULL 的行。 + + tokens 拆成 input/output 两半记（input + output = tokens， + 与服务端/ worker 的无价 token 口径一致）。 + """ + _insert_task(db_path, "donor-unpriced", status="completed") + store = CostStore(db_path) + try: + store.record("donor-unpriced", "p1", "gpt-5.6-terra", + input_tokens=tokens // 2, + output_tokens=tokens - tokens // 2, + usd_micros=None, source="actual") + finally: + try: + store._conn.close() + except Exception: + pass + + def _insert_followup(db_path, fu_id, task_id, body="follow up"): now = time.time() conn = sqlite3.connect(db_path) @@ -229,6 +249,73 @@ class WorkerQuotaTestCase(unittest.TestCase): result = worker.run_one() self.assertEqual(result["status"], "completed") + # -- 无价 token 上限 ------------------------------------------------- + def test_unpriced_token_cap_blocks_claim(self): + _insert_quota(self.db_path, "tester", 10 ** 12) # USD 很大，不挡 + _insert_unpriced_cost(self.db_path, 150) + _insert_task(self.db_path, "T1") + worker = self._worker(_usage_runner( + {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, + self.calls)) + old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP") + os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100" + try: + result = worker.run_one() + finally: + if old is None: + os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None) + else: + os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old + self.assertEqual(result["status"], "interrupted") + self.assertIn("无价模型 token", result["error"]) + self.assertEqual(self.calls, []) + + def test_unpriced_round_usage_counts_toward_cap(self): + _insert_quota(self.db_path, "tester", 10 ** 12) + _insert_unpriced_cost(self.db_path, 90) + _insert_task(self.db_path, "T1") + _insert_followup(self.db_path, "fu1", "T1") + worker = Worker(self.db_path, "/fake/host.py", runner=_usage_runner( + {"input_tokens": 50, "cached_input_tokens": 0, + "output_tokens": 0, "total_tokens": 50}, + self.calls), env=self.env, model="gpt-5.6-terra", + retry_backoff_seconds=0) + old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP") + os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100" + try: + result = worker.run_one() + finally: + if old is None: + os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None) + else: + os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old + # 首轮无价 50 token：90 + 50 = 140 >= 100，排空时停。 + self.assertEqual(result["status"], "interrupted") + self.assertIn("quota_exceeded", result["error"]) + self.assertEqual(_followup_status(self.db_path, "fu1"), "pending") + + def test_priced_usage_ignores_token_cap(self): + _insert_quota(self.db_path, "tester", 10 ** 12) + _insert_unpriced_cost(self.db_path, 90) + _insert_task(self.db_path, "T1") + _insert_followup(self.db_path, "fu1", "T1") + worker = self._worker(_usage_runner( + {"input_tokens": 100, "cached_input_tokens": 0, + "output_tokens": 10, "total_tokens": 110}, + self.calls)) + old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP") + os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100" + try: + result = worker.run_one() + finally: + if old is None: + os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None) + else: + os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old + # 有价模型用量不计入无价上限：90 < 100，正常跑完。 + self.assertEqual(result["status"], "completed") + self.assertEqual(_followup_status(self.db_path, "fu1"), "done") + # -- 轮询掐断 ------------------------------------------------------ def test_poll_loop_quota_kill(self): _insert_task(self.db_path, "T1") diff --git a/cloud/worker.py b/cloud/worker.py index 9353db9..ec9b12f 100644 --- a/cloud/worker.py +++ b/cloud/worker.py @@ -139,7 +139,8 @@ from cloud.coordinator import DEFAULT_POLL_INTERVAL as _COORD_POLL_INTERVAL from cloud.coordinator import DEFAULT_TIMEOUT_SECONDS as _COORD_TIMEOUT from cloud.coordinator import ensure_columns as _coord_ensure_columns from cloud.coordinator import parse_parts as _coord_parse_parts -from cloud.cost import CostStore, estimate_tokens, price_usd_micros, sum_usage +from cloud.cost import CostStore, estimate_tokens, price_usd_micros, \ + sum_usage, unpriced_token_cap from cloud.db import connect as _db_connect from cloud.logging_config import get_logger, log_event, setup_logging from cloud.models import ModelRouter, is_fallbackable, record_fallback @@ -832,9 +833,14 @@ class Worker: 口径与服务端准入（_quota_status）一致：按任务创建者算， 当月（UTC）已记账用量 + 本轮已累计用量（按价目表换算， 无价模型按 0 计，与服务端 SUM（usd NULL 不计）一致）。 - 无配额行 = 不限额。查询本身异常时按不限额处理（fail-open - 并记一条警告）：配额是第二道网，绝不能因查询抖动把全部 - 任务挡在门外；服务端建任务/分支/重试时的准入仍是第一道。 + 另有一道无价 token 上限（与服务端 + _check_unpriced_token_cap 同口径）：USD 配额不计量无价 + 模型，默认模型 terra 无价，若无此上限则配额对默认模型 + 永远不生效；有配额行的用户当月无价 token（input + + output，已记账 + 本轮）达到上限即停。无配额行 = 不限额。 + 查询本身异常时按不限额处理（fail-open 并记一条警告）： + 配额是第二道网，绝不能因查询抖动把全部任务挡在门外； + 服务端建任务/分支/重试时的准入仍是第一道。 """ owner = (task or {}).get("created_by") if not owner: @@ -873,6 +879,14 @@ class Worker: " WHERE t.created_by = ? AND c.day LIKE ?", (owner, month + "%")).fetchone() used = used_row[0] if used_row else 0 + unpriced_row = conn.execute( + "SELECT COALESCE(SUM(c.input_tokens + c.output_tokens), 0)" + " FROM cost_records c JOIN tasks t" + " ON c.task_id = t.id" + " WHERE t.created_by = ? AND c.day LIKE ?" + " AND c.usd_micros IS NULL", + (owner, month + "%")).fetchone() + unpriced_used = unpriced_row[0] if unpriced_row else 0 finally: conn.close() if usage_total: @@ -882,9 +896,20 @@ class Worker: usage_total.get("cached_input_tokens", 0), usage_total.get("output_tokens", 0)) used += micros or 0 + if micros is None: + # 无价模型：本轮 token 计入无价上限（USD 口 + # 径里它是 0，不会在上面那道里被拦住）。 + unpriced_used += ( + max(0, int(usage_total.get("input_tokens", 0))) + + max(0, int(usage_total.get("output_tokens", 0)))) if used >= quota: return ("quota_exceeded: 本月配额已用完（已用 %d / " "配额 %d 微美元，%s）" % (used, quota, month)) + cap = unpriced_token_cap() + if unpriced_used >= cap: + return ("quota_exceeded: 本月无价模型 token 已用完" + "（已用 %d / 上限 %d，%s）" + % (unpriced_used, cap, month)) return None except Exception as exc: try: