From dbfa026a4e6c7e29e753b05275a4519c643da78c Mon Sep 17 00:00:00 2001 From: host Date: Tue, 27 Oct 2026 17:18:18 +0800 Subject: [PATCH] =?UTF-8?q?sub2api=20=E5=BC=80=E9=80=9A=E6=8E=88=E5=88=9D?= =?UTF-8?q?=E5=A7=8B=E4=BD=99=E9=A2=9D=EF=BC=9A=E4=BD=99=E9=A2=9D=E4=B8=BA?= =?UTF-8?q?=200=20=E7=9A=84=20key=20=E5=9C=A8=E7=BD=91=E5=85=B3=E4=BE=A7?= =?UTF-8?q?=E7=9B=B4=E6=8E=A5=20403?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cloud/README.md | 4 ++++ cloud/sub2api.py | 40 +++++++++++++++++++++++++++++---- cloud/tests/test_sub2api.py | 45 +++++++++++++++++++++++++++++++++++++ 3 files changed, 85 insertions(+), 4 deletions(-) diff --git a/cloud/README.md b/cloud/README.md index 92aaca9..df27b73 100644 --- a/cloud/README.md +++ b/cloud/README.md @@ -1033,6 +1033,10 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳 任务直接 failed，绝不跑错网关。协调员子任务无 created_by 时 跟父任务属主走。key 明文存 PG，与 GitHub App 私钥同一信任 边界（库只在宿主 loopback 可达）；at-rest 加密是后续项。 + 开通同时授初始余额：`CLOUD_SUB2API_OPENING_BALANCE_USD` + （默认 10 美元，显式配 0 表示不授；非法/负值回落默认）—— + 网关鉴权层在余额 ≤ 0 时直接 403，不授余额的 key 立刻不可 + 用；授余额失败时 key 与账号一并清理再抛错，不留孤儿。 - 注入：worker 在三条执行路径（Local popen、Local 同步、Docker 同步）都注入任务级 env：`SUB2API_API_KEY`（属主的 key）、 `SUB2API_BASE_URL`，并把网关 host 追加进 NO_PROXY（否则 diff --git a/cloud/sub2api.py b/cloud/sub2api.py index 398ef2a..d6a51b6 100644 --- a/cloud/sub2api.py +++ b/cloud/sub2api.py @@ -73,6 +73,24 @@ def gateway_group_name(): or "codex-system" +def opening_balance(): + """开通时授的初始余额（美元）：默认 10，显式值优先。 + + 新用户余额为 0 时网关鉴权层直接 403（余额门在任何上游路由 + 之前），所以开通必须同时授余额，否则发出的 key 立刻不可用。 + 显式配 0 表示不授（调用方自理）；非法/负值回落默认 10。 + """ + raw = os.environ.get( + "CLOUD_SUB2API_OPENING_BALANCE_USD", "").strip() + if not raw: + return 10.0 + try: + value = float(raw) + except ValueError: + return 10.0 + return value if value >= 0 else 10.0 + + def ensure_gateway_keys_table(conn): conn.execute(GATEWAY_TABLE_SQL) @@ -181,6 +199,16 @@ class Sub2apiAdmin: "DELETE", "/api/v1/admin/users/%s" % sub_user_id) return st in (200, 204) + def set_balance(self, sub_user_id, amount): + st, data = self._call( + "POST", "/api/v1/admin/users/%s/balance" % sub_user_id, + {"balance": amount, "operation": "set", + "notes": "codex provisioning opening balance"}) + if st != 200: + raise Sub2apiError( + "sub2api 授余额失败: %s" + % ((data or {}).get("message") or st), status=st) + def list_groups(self): st, data = self._call("GET", "/api/v1/admin/groups") payload = (data or {}).get("data") @@ -277,11 +305,12 @@ def list_gateway_models(origin, api_key, request=None): def ensure_user_gateway_key(conn, user_id, admin=None, request=None): - """幂等开通：已有行直接返回 key；否则建号->登录->发 key->入库。 + """幂等开通：已有行直接返回 key；否则建号->登录->发 key->授余额->入库。 - 入库失败时尽力清理刚建的 key 与账号（避免孤儿 key），再抛原错。 - 调用方负责 conn 的提交语义：本函数在入库后 commit（开通是独立 - 事务，不跟调用方的业务事务绑在一起）。 + 入库失败时尽力清理刚建的 key 与账号（避免孤儿 key），再抛原错； + 授余额失败同样走这条清理（余额为 0 的 key 在网关侧直接 403， + 留着没有意义）。调用方负责 conn 的提交语义：本函数在入库后 + commit（开通是独立事务，不跟调用方的业务事务绑在一起）。 """ ensure_gateway_keys_table(conn) row = get_user_gateway_key(conn, user_id) @@ -304,6 +333,9 @@ def ensure_user_gateway_key(conn, user_id, admin=None, request=None): key_id, key = user_create_key(origin, user_token, "codex-task-key", group_id, request=req) + amount = opening_balance() + if amount > 0: + admin.set_balance(sub_id, amount) now = time.time() conn.execute( "INSERT INTO user_gateway_keys(user_id, sub2api_user_id, " diff --git a/cloud/tests/test_sub2api.py b/cloud/tests/test_sub2api.py index 6dfaab6..16b5341 100644 --- a/cloud/tests/test_sub2api.py +++ b/cloud/tests/test_sub2api.py @@ -36,6 +36,8 @@ class FakeApi: self.next_user_id = 10 self.next_key_id = 20 self.fail_refresh_once = False + self.balance_calls = [] + self.fail_balance = False def __call__(self, method, url, body=None, token=None): path = url.split("28084", 1)[-1] if "28084" in url else url @@ -52,6 +54,12 @@ class FakeApi: return 200, {"data": {"access_token": "new-tok", "refresh_token": "new-ref", "expires_in": 3600}} + if path.startswith("/api/v1/admin/users/") \ + and path.endswith("/balance") and method == "POST": + self.balance_calls.append(body) + if self.fail_balance: + return 400, {"message": "bad balance"} + return 200, {"data": {"balance": body.get("balance")}} if path == "/api/v1/admin/users" and method == "POST": self.next_user_id += 1 return 200, {"data": {"id": self.next_user_id}} @@ -135,6 +143,43 @@ class TestEnsureGatewayKey(unittest.TestCase): request=self.api) self.assertEqual(models, ["gpt-5", "claude-x"]) + def test_provision_grants_opening_balance(self): + conn = _conn() + key = ensure_user_gateway_key(conn, "user-bal", admin=self.admin, + request=self.api) + self.assertTrue(key.startswith("sk-probe-")) + self.assertEqual(len(self.api.balance_calls), 1) + self.assertEqual(self.api.balance_calls[0]["balance"], 10.0) + self.assertEqual(self.api.balance_calls[0]["operation"], "set") + + def test_provision_balance_failure_cleans_up(self): + self.api.fail_balance = True + conn = _conn() + with self.assertRaises(sub2api.Sub2apiError): + ensure_user_gateway_key(conn, "user-fail", admin=self.admin, + request=self.api) + paths = [(m, p) for (m, p, _b, _t) in self.api.calls] + self.assertIn(("DELETE", "/api/v1/keys/21"), paths) + self.assertIn(("DELETE", "/api/v1/admin/users/11"), paths) + row = conn.execute( + "SELECT COUNT(*) FROM user_gateway_keys" + " WHERE user_id = ?", ("user-fail",)).fetchone() + self.assertEqual(row[0], 0) + + def test_provision_zero_opening_balance_skips_grant(self): + old = os.environ.get("CLOUD_SUB2API_OPENING_BALANCE_USD") + os.environ["CLOUD_SUB2API_OPENING_BALANCE_USD"] = "0" + try: + conn = _conn() + ensure_user_gateway_key(conn, "user-zero", admin=self.admin, + request=self.api) + finally: + if old is None: + os.environ.pop("CLOUD_SUB2API_OPENING_BALANCE_USD", None) + else: + os.environ["CLOUD_SUB2API_OPENING_BALANCE_USD"] = old + self.assertEqual(self.api.balance_calls, []) + class TestProviderCommand(unittest.TestCase): def test_server_command_sub2api(self):