P1 security batch: quota execution checks, cost accounting, goal/coordinator semantics, PR channel, orphan signature, control-plane and console hardening
- Docker quota watchdog + fail-closed quota checks, system tasks need real owner - max_cost_usd actually accumulated, cancel records cost, estimated rows excluded from unpriced cap, cache_write declared non-billed - goal-resume keeps original verify on breakpoint resume, goal retry falls back to resume, goal_max_rounds capped at 100 at create and execution - coordinator children inherit parent model/effort, orphans marked interrupted with audit - PR jobs on docker workers return explicit needs_local_worker instead of silent empty - worker code signature covers 6 files and is checked before claim, fail-closed - sub2api lazy provisioning with rotation, backup whitelist + 0600, push SSRF block, repos/workers desensitized, key/project audits and limits, console buttons/quota/status wording, coordinator/automation quota reserve
This commit is contained in:
1 parent
d254240d84
commit
1402dc67fd
20 files changed
+2616
-179
No files matched your search
@@ -33,6 +33,9 @@ class CapsServerTestCase(unittest.TestCase):
|
||||
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")
|
||||
# P1:完整 workers 清单只限全局管理员;owner 在此测试里即管理员。
|
||||
self._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
|
||||
os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
|
||||
self.server = create_server(self.db_path, "127.0.0.1", 0)
|
||||
self.port = self.server.server_address[1]
|
||||
self.thread = threading.Thread(
|
||||
@@ -48,6 +51,10 @@ class CapsServerTestCase(unittest.TestCase):
|
||||
except Exception:
|
||||
pass
|
||||
self.tmp.cleanup()
|
||||
if getattr(self, "_old_admins", None) is None:
|
||||
os.environ.pop("CODEX_GLOBAL_ADMINS", None)
|
||||
else:
|
||||
os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins
|
||||
|
||||
def _req(self, method, path, body=None, key=None):
|
||||
conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)
|
||||
@@ -82,11 +89,21 @@ class CapsServerTestCase(unittest.TestCase):
|
||||
"capacity": 4, "running": 2},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 200, body)
|
||||
# 读权限:任何已认证用户可读(viewer 也能看到基础设施行)。
|
||||
workers = self._workers(self.viewer_key)
|
||||
# P1:完整清单只限全局管理员(owner 在此测试里即管理员)。
|
||||
workers = self._workers(self.owner_key)
|
||||
self.assertEqual(workers["w1"]["capacity"], 4)
|
||||
self.assertEqual(workers["w1"]["running"], 2)
|
||||
|
||||
def test_workers_list_non_admin_gets_aggregate_only(self):
|
||||
self._req("POST", "/workers/heartbeat",
|
||||
{"worker_id": "w1", "executor": "local",
|
||||
"capacity": 4, "running": 2},
|
||||
self.owner_key)
|
||||
status, body = self._req("GET", "/workers", key=self.viewer_key)
|
||||
self.assertEqual(status, 200, body)
|
||||
self.assertNotIn("workers", body)
|
||||
self.assertIn("workers_total", body)
|
||||
|
||||
def test_heartbeat_without_running_preserves_old(self):
|
||||
self._req("POST", "/workers/heartbeat",
|
||||
{"worker_id": "w1", "capacity": 2, "running": 1},
|
||||
|
||||
@@ -184,7 +184,16 @@ class WorkerRegistryTestCase(_ServerCase):
|
||||
self.assertEqual(status, 200, body)
|
||||
self.assertTrue(body["alive"])
|
||||
self.assertEqual(body["executor"], "docker")
|
||||
status, body = self._req("GET", "/workers", None, self.owner_key)
|
||||
# P1:完整 workers 清单只限全局管理员,本测试里 owner 即管理员。
|
||||
old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
|
||||
os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
|
||||
try:
|
||||
status, body = self._req("GET", "/workers", None, self.owner_key)
|
||||
finally:
|
||||
if old_admins is None:
|
||||
os.environ.pop("CODEX_GLOBAL_ADMINS", None)
|
||||
else:
|
||||
os.environ["CODEX_GLOBAL_ADMINS"] = old_admins
|
||||
self.assertEqual(status, 200, body)
|
||||
self.assertEqual(len(body["workers"]), 1)
|
||||
self.assertTrue(body["workers"][0]["alive"])
|
||||
|
||||
@@ -95,13 +95,23 @@ class ServerGithubReposTestCase(unittest.TestCase):
|
||||
|
||||
def test_list_returns_seeded_installations(self):
|
||||
self._seed_installations()
|
||||
# P1:repos 只返回调用者成员项目对应的仓库;先把项目仓库
|
||||
# 绑到其中一个已安装仓库,再断言只返回该仓库,非成员为空。
|
||||
s, d = _req(self.port, "PATCH",
|
||||
"/projects/" + self.project_id,
|
||||
{"repo_url": "https://github.com/openqj/codex-app-server.git"},
|
||||
self.owner_key)
|
||||
self.assertEqual(s, 200, d)
|
||||
s, d = _req(self.port, "GET", "/github/repos", key=self.owner_key)
|
||||
self.assertEqual(s, 200, d)
|
||||
names = [r["repo_full_name"] for r in d["repos"]]
|
||||
self.assertEqual(names, ["openqj/codex-app-server", "openqj/demo"])
|
||||
self.assertEqual(names, ["openqj/codex-app-server"])
|
||||
self.assertEqual(
|
||||
d["repos"][0]["clone_url"],
|
||||
"https://github.com/openqj/codex-app-server.git")
|
||||
s, d = _req(self.port, "GET", "/github/repos", key=self.other_key)
|
||||
self.assertEqual(s, 200, d)
|
||||
self.assertEqual(d["repos"], [])
|
||||
|
||||
def test_patch_repo_url_owner(self):
|
||||
url = "https://github.com/openqj/demo.git"
|
||||
|
||||
@@ -0,0 +1,598 @@
|
||||
"""P1 控制面/控制台单测(server.py / console.html / sub2api.py / push.py)。
|
||||
|
||||
只覆盖本批 P1 控制面改动,不跑真实模型回合、不联网、不碰生产:
|
||||
server 侧走本地临时端口 + 临时 SQLite;网关侧用假 request /
|
||||
假 admin(key、token 全是测试假值);推送 DNS 用注入的
|
||||
resolver。绝不改 worker.py / coordinator.py / cost.py / host.py。
|
||||
"""
|
||||
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import stat
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
try:
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.push import dispatch_event, validate_push_target
|
||||
from cloud.push import create_subscription
|
||||
from cloud.server import create_server
|
||||
from cloud import server as server_mod
|
||||
from cloud.sub2api import Sub2apiError
|
||||
from cloud.sub2api import ensure_gateway_keys_table
|
||||
from cloud.sub2api import get_user_gateway_key
|
||||
from cloud.sub2api import rotate_user_gateway_key
|
||||
from cloud.sub2api import validate_or_rotate_gateway_key
|
||||
except ImportError:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.push import dispatch_event, validate_push_target
|
||||
from cloud.push import create_subscription
|
||||
from cloud.server import create_server
|
||||
from cloud import server as server_mod
|
||||
from cloud.sub2api import Sub2apiError
|
||||
from cloud.sub2api import ensure_gateway_keys_table
|
||||
from cloud.sub2api import get_user_gateway_key
|
||||
from cloud.sub2api import rotate_user_gateway_key
|
||||
from cloud.sub2api import validate_or_rotate_gateway_key
|
||||
|
||||
|
||||
def _req(port, method, path, body=None, key=None, raw_body=None):
|
||||
conn = http.client.HTTPConnection("127.0.0.1", port, timeout=10)
|
||||
hdrs = {}
|
||||
if raw_body is not None:
|
||||
payload = raw_body
|
||||
elif body is not None:
|
||||
payload = json.dumps(body).encode("utf-8")
|
||||
hdrs["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 = data.decode("utf-8", "replace")
|
||||
return resp.status, parsed
|
||||
|
||||
|
||||
class _FakeGatewayAdmin:
|
||||
"""假网关管理员:建号/授余额只记调用,不联网。"""
|
||||
|
||||
def __init__(self):
|
||||
self.created = []
|
||||
self.deleted = []
|
||||
self.balances = []
|
||||
|
||||
def create_user(self, email, password, username):
|
||||
sub_id = 1000 + len(self.created)
|
||||
self.created.append((email, username))
|
||||
return sub_id
|
||||
|
||||
def delete_user(self, sub_user_id):
|
||||
self.deleted.append(sub_user_id)
|
||||
return True
|
||||
|
||||
def set_balance(self, sub_user_id, amount):
|
||||
self.balances.append((sub_user_id, amount))
|
||||
|
||||
def resolve_group_id(self):
|
||||
return 2
|
||||
|
||||
|
||||
class _FakeGatewayRequest:
|
||||
"""假网关 HTTP:models 状态可调,其余走固定成功路径。"""
|
||||
|
||||
def __init__(self, models_status=200):
|
||||
self.models_status = models_status
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, method, url, body=None, token=None):
|
||||
self.calls.append((method, url))
|
||||
if url.endswith("/v1/models"):
|
||||
return self.models_status, {"data": []}
|
||||
if url.endswith("/api/v1/auth/login"):
|
||||
return 200, {"data": {"access_token": "fake-user-token",
|
||||
"refresh_token": "fake-refresh",
|
||||
"expires_in": 3600}}
|
||||
if url.endswith("/api/v1/keys"):
|
||||
return 201, {"data": {"id": "fake-key-id-1",
|
||||
"key": "fake-gateway-key-1"}}
|
||||
if "/balance" in url:
|
||||
return 200, {"data": {}}
|
||||
return 200, {"data": {}}
|
||||
|
||||
|
||||
class Sub2apiP1TestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.conn = sqlite3.connect(":memory:")
|
||||
ensure_gateway_keys_table(self.conn)
|
||||
self.admin = _FakeGatewayAdmin()
|
||||
# 假 origin:请求本身全是 fake,不产生真实网络调用。
|
||||
self._old_origin = os.environ.get("SUB2API_BASE_URL")
|
||||
os.environ["SUB2API_BASE_URL"] = "http://127.0.0.1:9"
|
||||
|
||||
def tearDown(self):
|
||||
self.conn.close()
|
||||
if self._old_origin is None:
|
||||
os.environ.pop("SUB2API_BASE_URL", None)
|
||||
else:
|
||||
os.environ["SUB2API_BASE_URL"] = self._old_origin
|
||||
|
||||
def _seed_row(self, key="fake-old-key"):
|
||||
self.conn.execute(
|
||||
"INSERT INTO user_gateway_keys(user_id, sub2api_user_id,"
|
||||
" api_key, refresh_token, created_at, updated_at)"
|
||||
" VALUES(?, ?, ?, ?, ?, ?)",
|
||||
("u1", 42, key, "fake-refresh", 1.0, 1.0))
|
||||
self.conn.commit()
|
||||
|
||||
def test_validate_no_row_provisions(self):
|
||||
req = _FakeGatewayRequest()
|
||||
status, key, sub_id, key_id = validate_or_rotate_gateway_key(
|
||||
self.conn, "u1", admin=self.admin, request=req)
|
||||
self.assertEqual(status, "provisioned")
|
||||
self.assertEqual(key, "fake-gateway-key-1")
|
||||
self.assertEqual(key_id, "fake-key-id-1")
|
||||
self.assertIsNotNone(get_user_gateway_key(self.conn, "u1"))
|
||||
|
||||
def test_validate_healthy_key_not_rotated(self):
|
||||
self._seed_row()
|
||||
req = _FakeGatewayRequest(models_status=200)
|
||||
status, key, _sub, _kid = validate_or_rotate_gateway_key(
|
||||
self.conn, "u1", admin=self.admin, request=req)
|
||||
self.assertEqual(status, "ok")
|
||||
self.assertEqual(key, "fake-old-key")
|
||||
self.assertEqual(self.admin.created, [])
|
||||
|
||||
def test_validate_dead_key_rotates_not_pending(self):
|
||||
self._seed_row()
|
||||
req = _FakeGatewayRequest(models_status=403)
|
||||
status, key, _sub, _kid = validate_or_rotate_gateway_key(
|
||||
self.conn, "u1", admin=self.admin, request=req)
|
||||
self.assertEqual(status, "rotated")
|
||||
self.assertEqual(key, "fake-gateway-key-1")
|
||||
row = get_user_gateway_key(self.conn, "u1")
|
||||
self.assertEqual(row["api_key"], "fake-gateway-key-1")
|
||||
self.assertNotEqual(row["api_key"], "fake-old-key")
|
||||
|
||||
def test_validate_zero_balance_403_rotates(self):
|
||||
# 余额为 0 时网关鉴权层回 403:与 401 同属死 key,必须轮换。
|
||||
self._seed_row()
|
||||
req = _FakeGatewayRequest(models_status=403)
|
||||
status, _key, _sub, _kid = validate_or_rotate_gateway_key(
|
||||
self.conn, "u1", admin=self.admin, request=req)
|
||||
self.assertEqual(status, "rotated")
|
||||
|
||||
def test_validate_gateway_down_keeps_old_row(self):
|
||||
self._seed_row()
|
||||
req = _FakeGatewayRequest(models_status=500)
|
||||
with self.assertRaises(Sub2apiError):
|
||||
validate_or_rotate_gateway_key(
|
||||
self.conn, "u1", admin=self.admin, request=req)
|
||||
row = get_user_gateway_key(self.conn, "u1")
|
||||
self.assertEqual(row["api_key"], "fake-old-key")
|
||||
|
||||
def test_rotate_replaces_row(self):
|
||||
self._seed_row()
|
||||
req = _FakeGatewayRequest()
|
||||
status, key, sub_id, key_id = rotate_user_gateway_key(
|
||||
self.conn, "u1", admin=self.admin, request=req)
|
||||
self.assertEqual(status, "rotated")
|
||||
self.assertEqual(key_id, "fake-key-id-1")
|
||||
self.assertEqual(
|
||||
get_user_gateway_key(self.conn, "u1")["api_key"], key)
|
||||
|
||||
|
||||
class PushP1TestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.conn = sqlite3.connect(":memory:")
|
||||
|
||||
def tearDown(self):
|
||||
self.conn.close()
|
||||
|
||||
def test_literal_private_and_loopback_blocked(self):
|
||||
for target in (
|
||||
"http://127.0.0.1:9/hook",
|
||||
"http://10.0.0.5/hook",
|
||||
"http://172.16.4.4/hook",
|
||||
"http://192.168.1.10/hook",
|
||||
"http://169.254.169.254/latest/meta-data/",
|
||||
"http://[::1]/hook"):
|
||||
self.assertIsNotNone(
|
||||
validate_push_target(target), target)
|
||||
|
||||
def test_metadata_hostname_blocked(self):
|
||||
self.assertIsNotNone(validate_push_target(
|
||||
"http://metadata.google.internal/hook"))
|
||||
|
||||
def test_hostname_resolved_to_private_blocked(self):
|
||||
err = validate_push_target(
|
||||
"https://hooks.example.com/hook",
|
||||
resolver=lambda host: ["10.1.2.3"])
|
||||
self.assertIsNotNone(err)
|
||||
|
||||
def test_hostname_public_resolves_ok(self):
|
||||
err = validate_push_target(
|
||||
"https://hooks.example.com/hook",
|
||||
resolver=lambda host: ["93.184.216.34"])
|
||||
self.assertIsNone(err)
|
||||
|
||||
def test_dns_failure_rejected(self):
|
||||
def _boom(host):
|
||||
raise OSError("no dns in test")
|
||||
self.assertIsNotNone(validate_push_target(
|
||||
"https://hooks.example.com/hook", resolver=_boom))
|
||||
|
||||
def test_dispatch_loopback_never_posts(self):
|
||||
# 真实发送路径(sender=None,生产服务端走的路径):
|
||||
# loopback 目标在预检即被挡,不产生任何 POST。
|
||||
create_subscription(self.conn, "proj-a", "user-1", "webhook",
|
||||
"http://127.0.0.1:9/hook", "instant")
|
||||
result = dispatch_event(self.conn, "failed", "proj-a",
|
||||
{"task_id": "t1"}, sender=None)
|
||||
self.assertEqual(result["sent"], [])
|
||||
self.assertEqual(len(result["errors"]), 1)
|
||||
self.assertIn("blocked", result["errors"][0]["reason"])
|
||||
|
||||
|
||||
class ServerP1TestCase(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.member_id, self.member_key = self.auth.create_user("member")
|
||||
self.outsider_id, self.outsider_key = self.auth.create_user(
|
||||
"outsider")
|
||||
self._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
|
||||
self._old_gateway = os.environ.get("CLOUD_GATEWAY")
|
||||
self._old_backup_root = os.environ.get("CLOUD_BACKUP_ROOT")
|
||||
self._old_backup_dir = os.environ.get("CLOUD_BACKUP_DIR")
|
||||
os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
|
||||
os.environ.pop("CLOUD_GATEWAY", None)
|
||||
os.environ.pop("CLOUD_BACKUP_ROOT", None)
|
||||
os.environ.pop("CLOUD_BACKUP_DIR", None)
|
||||
self._old_max_projects = server_mod.MAX_PROJECTS_PER_USER
|
||||
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": "owner/repo"}, self.owner_key)
|
||||
self.assertEqual(s, 201, proj)
|
||||
self.project_id = proj["id"]
|
||||
self.auth.set_member(
|
||||
self.project_id, self.owner_id, self.member_id, "maintainer")
|
||||
|
||||
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()
|
||||
server_mod.MAX_PROJECTS_PER_USER = self._old_max_projects
|
||||
if self._old_admins is None:
|
||||
os.environ.pop("CODEX_GLOBAL_ADMINS", None)
|
||||
else:
|
||||
os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins
|
||||
if self._old_gateway is None:
|
||||
os.environ.pop("CLOUD_GATEWAY", None)
|
||||
else:
|
||||
os.environ["CLOUD_GATEWAY"] = self._old_gateway
|
||||
if self._old_backup_root is None:
|
||||
os.environ.pop("CLOUD_BACKUP_ROOT", None)
|
||||
else:
|
||||
os.environ["CLOUD_BACKUP_ROOT"] = self._old_backup_root
|
||||
if self._old_backup_dir is None:
|
||||
os.environ.pop("CLOUD_BACKUP_DIR", None)
|
||||
else:
|
||||
os.environ["CLOUD_BACKUP_DIR"] = self._old_backup_dir
|
||||
# 恢复猴子补丁过的 server 模块函数(relay/rotate 相关测试用)。
|
||||
for name in ("_validate_sub2api_key", "_list_sub2api_models",
|
||||
"_rotate_sub2api_key"):
|
||||
try:
|
||||
mock.patch.object(server_mod, name).stop()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# -- helpers ---------------------------------------------------
|
||||
def _audit_actions(self, user_id=None):
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
rows = conn.execute(
|
||||
"SELECT action, detail FROM audit_log"
|
||||
+ (" WHERE user_id = ?" if user_id else ""),
|
||||
(user_id,) if user_id else ()).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
return rows
|
||||
|
||||
def _set_quota(self, target, quota):
|
||||
return _req(self.port, "POST", "/admin/quotas",
|
||||
{"user_id": target, "quota_usd_micros": quota},
|
||||
self.owner_key)
|
||||
|
||||
# -- 备份 -------------------------------------------------------
|
||||
def test_backup_whitelist_and_mode(self):
|
||||
root = os.path.join(self.tmp.name, "backup-root")
|
||||
os.environ["CLOUD_BACKUP_ROOT"] = root
|
||||
out_dir = os.path.join(root, "daily")
|
||||
s, body = _req(self.port, "POST", "/admin/backup",
|
||||
{"out_dir": out_dir}, self.owner_key)
|
||||
self.assertEqual(s, 200, body)
|
||||
self.assertTrue(os.path.exists(body["path"]))
|
||||
self.assertEqual(stat.S_IMODE(os.stat(body["path"]).st_mode),
|
||||
0o600)
|
||||
details = [d for _a, d in self._audit_actions(self.owner_id)]
|
||||
self.assertTrue(any("admin_backup" in a
|
||||
for a, _d in self._audit_actions(
|
||||
self.owner_id)))
|
||||
self.assertTrue(any(out_dir in (d or "")
|
||||
for d in details))
|
||||
# 白名单外:系统目录与 .. 逃逸一律 400。
|
||||
for bad in ("/etc", os.path.join(root, "..", "escape")):
|
||||
s, body = _req(self.port, "POST", "/admin/backup",
|
||||
{"out_dir": bad}, self.owner_key)
|
||||
self.assertEqual(s, 400, (bad, body))
|
||||
|
||||
# -- workers 聚合 ------------------------------------------------
|
||||
def test_workers_admin_full_member_aggregate_only(self):
|
||||
s, _b = _req(self.port, "POST", "/workers/heartbeat",
|
||||
{"worker_id": "w-p1", "executor": "docker",
|
||||
"capacity": 2, "running": 0}, self.owner_key)
|
||||
self.assertEqual(s, 200, _b)
|
||||
s, body = _req(self.port, "GET", "/workers", key=self.owner_key)
|
||||
self.assertEqual(s, 200)
|
||||
self.assertIn("workers", body)
|
||||
self.assertTrue(any(w.get("worker_id") == "w-p1"
|
||||
for w in body["workers"]))
|
||||
s, body = _req(self.port, "GET", "/workers", key=self.member_key)
|
||||
self.assertEqual(s, 200)
|
||||
self.assertNotIn("workers", body)
|
||||
self.assertNotIn("w-p1", json.dumps(body))
|
||||
self.assertEqual(body["workers_total"], 1)
|
||||
self.assertEqual(body["workers_online"], 1)
|
||||
self.assertEqual(body["workers_offline"], 0)
|
||||
|
||||
# -- github repos 过滤 -------------------------------------------
|
||||
def test_github_repos_filtered_by_member_projects(self):
|
||||
s, proj = _req(self.port, "POST", "/projects",
|
||||
{"name": "with-repo",
|
||||
"repo_url": "https://github.com/acme/alpha.git"},
|
||||
self.owner_key)
|
||||
self.assertEqual(s, 201, proj)
|
||||
self.auth.set_member(proj["id"], self.owner_id, self.member_id,
|
||||
"maintainer")
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS github_app_installations("
|
||||
"installation_id TEXT, account_login TEXT,"
|
||||
" repo_full_name TEXT, created_at REAL,"
|
||||
" UNIQUE(installation_id, repo_full_name))")
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO github_app_installations"
|
||||
" VALUES(?, ?, ?, ?)", ("7", "acme", "acme/alpha", 1.0))
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO github_app_installations"
|
||||
" VALUES(?, ?, ?, ?)", ("7", "acme", "acme/beta", 1.0))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
s, body = _req(self.port, "GET", "/github/repos",
|
||||
key=self.member_key)
|
||||
self.assertEqual(s, 200, body)
|
||||
names = [r["repo_full_name"] for r in body["repos"]]
|
||||
self.assertEqual(names, ["acme/alpha"])
|
||||
s, body = _req(self.port, "GET", "/github/repos",
|
||||
key=self.outsider_key)
|
||||
self.assertEqual(s, 200, body)
|
||||
self.assertEqual(body["repos"], [])
|
||||
|
||||
# -- key 审计 -----------------------------------------------------
|
||||
def test_api_key_audit_only_id_and_label(self):
|
||||
s, body = _req(self.port, "POST", "/users/me/keys",
|
||||
{"label": "p1-test"}, self.owner_key)
|
||||
self.assertEqual(s, 201, body)
|
||||
key_id = body["id"]
|
||||
plain = body["api_key"]
|
||||
s, _b = _req(self.port, "DELETE",
|
||||
"/users/me/keys/" + key_id, key=self.owner_key)
|
||||
self.assertEqual(s, 200, _b)
|
||||
s, body = _req(self.port, "POST",
|
||||
"/users/me/keys/revoke-others", {},
|
||||
self.owner_key)
|
||||
self.assertEqual(s, 200, body)
|
||||
actions = {a for a, _d in self._audit_actions(self.owner_id)}
|
||||
self.assertIn("api_key_created", actions)
|
||||
self.assertIn("api_key_revoked", actions)
|
||||
self.assertIn("api_key_revoke_others", actions)
|
||||
blob = json.dumps(self._audit_actions(self.owner_id))
|
||||
self.assertNotIn(plain, blob)
|
||||
|
||||
# -- 项目上限 ------------------------------------------------------
|
||||
def test_project_limits(self):
|
||||
s, body = _req(self.port, "POST", "/projects",
|
||||
{"name": "x" * 101}, self.owner_key)
|
||||
self.assertEqual(s, 400, body)
|
||||
self.assertIn("100", body["error"])
|
||||
s, body = _req(self.port, "POST", "/projects",
|
||||
{"name": "ok", "id": "y" * 65}, self.owner_key)
|
||||
self.assertEqual(s, 400, body)
|
||||
self.assertIn("64", body["error"])
|
||||
server_mod.MAX_PROJECTS_PER_USER = 1
|
||||
# setUp 已建 1 个 owner 项目,限额 1 时再建即超限。
|
||||
s, body = _req(self.port, "POST", "/projects",
|
||||
{"name": "second"}, self.owner_key)
|
||||
self.assertEqual(s, 400, body)
|
||||
self.assertIn("max 1", body["error"])
|
||||
|
||||
# -- 协调员预留 ----------------------------------------------------
|
||||
def test_coordinator_reserve_blocks_parent(self):
|
||||
s, _b = self._set_quota(self.owner_id, 1000000)
|
||||
self.assertEqual(s, 200, _b)
|
||||
parts = [{"name": "a", "prompt": "做 a"},
|
||||
{"name": "b", "prompt": "做 b"},
|
||||
{"name": "c", "prompt": "做 c"}]
|
||||
s, body = _req(self.port, "POST", "/tasks",
|
||||
{"project_id": self.project_id, "prompt": "协调",
|
||||
"coordinator_parts": parts}, self.owner_key)
|
||||
self.assertEqual(s, 402, body)
|
||||
self.assertEqual(body["error"],
|
||||
"coordinator_quota_reserve_exceeded")
|
||||
self.assertEqual(body["coordinator_parts"], 3)
|
||||
s, _b = self._set_quota(self.owner_id, 10000000)
|
||||
self.assertEqual(s, 200, _b)
|
||||
s, body = _req(self.port, "POST", "/tasks",
|
||||
{"project_id": self.project_id, "prompt": "协调",
|
||||
"coordinator_parts": parts}, self.owner_key)
|
||||
self.assertEqual(s, 201, body)
|
||||
self.assertIn("parent task creator",
|
||||
body.get("quota_attribution", ""))
|
||||
self.assertEqual(body.get("quota_reserved_usd_micros"), 3000000)
|
||||
|
||||
# -- 自动化配额门 ----------------------------------------------------
|
||||
def test_automation_create_quota_gate(self):
|
||||
s, _b = self._set_quota(self.owner_id, 0)
|
||||
self.assertEqual(s, 200, _b)
|
||||
s, body = _req(self.port, "POST",
|
||||
"/projects/%s/automations" % self.project_id,
|
||||
{"name": "nightly", "prompt": "巡检"},
|
||||
self.owner_key)
|
||||
self.assertEqual(s, 402, body)
|
||||
s, _b = self._set_quota(self.owner_id, 10000000)
|
||||
self.assertEqual(s, 200, _b)
|
||||
s, body = _req(self.port, "POST",
|
||||
"/projects/%s/automations" % self.project_id,
|
||||
{"name": "nightly", "prompt": "巡检"},
|
||||
self.owner_key)
|
||||
self.assertEqual(s, 201, body)
|
||||
|
||||
# -- sub2api 懒开通 / 轮换端点 --------------------------------------
|
||||
def test_relay_models_lazy_provision_not_just_pending(self):
|
||||
os.environ["CLOUD_GATEWAY"] = "sub2api"
|
||||
orig_validate = server_mod._validate_sub2api_key
|
||||
orig_list = server_mod._list_sub2api_models
|
||||
try:
|
||||
server_mod._validate_sub2api_key = (
|
||||
lambda conn, uid: ("provisioned", "fake-key", 7, 9))
|
||||
server_mod._list_sub2api_models = (
|
||||
lambda origin, key: ["model-a"])
|
||||
s, body = _req(self.port, "GET", "/relay/models",
|
||||
key=self.owner_key)
|
||||
self.assertEqual(s, 200, body)
|
||||
self.assertEqual(body["models"], ["model-a"])
|
||||
self.assertEqual(body["provisioning"], "provisioned")
|
||||
server_mod._validate_sub2api_key = (
|
||||
lambda conn, uid: (_ for _ in ()).throw(
|
||||
Sub2apiError("网关暂不可用")))
|
||||
s, body = _req(self.port, "GET", "/relay/models",
|
||||
key=self.owner_key)
|
||||
self.assertEqual(s, 502, body)
|
||||
finally:
|
||||
server_mod._validate_sub2api_key = orig_validate
|
||||
server_mod._list_sub2api_models = orig_list
|
||||
|
||||
def test_gateway_rotate_endpoint(self):
|
||||
s, body = _req(self.port, "POST", "/users/me/gateway/rotate",
|
||||
{}, self.owner_key)
|
||||
self.assertEqual(s, 400, body)
|
||||
os.environ["CLOUD_GATEWAY"] = "sub2api"
|
||||
orig_rotate = server_mod._rotate_sub2api_key
|
||||
try:
|
||||
server_mod._rotate_sub2api_key = (
|
||||
lambda conn, uid: ("rotated", "fake-key", 7, 9))
|
||||
s, body = _req(self.port, "POST", "/users/me/gateway/rotate",
|
||||
{}, self.owner_key)
|
||||
self.assertEqual(s, 200, body)
|
||||
self.assertTrue(body["rotated"])
|
||||
self.assertNotIn("api_key", body)
|
||||
actions = {a for a, _d in self._audit_actions(self.owner_id)}
|
||||
self.assertIn("gateway_key_rotated", actions)
|
||||
finally:
|
||||
server_mod._rotate_sub2api_key = orig_rotate
|
||||
|
||||
# -- 推送审计与分发 SSRF ----------------------------------------------
|
||||
def test_push_audit_and_dispatch_ssrf(self):
|
||||
s, sub = _req(self.port, "POST",
|
||||
"/projects/%s/push/subscriptions" % self.project_id,
|
||||
{"channel": "webhook",
|
||||
"target": "https://example.com/hook",
|
||||
"preference": "instant"}, self.owner_key)
|
||||
self.assertEqual(s, 201, sub)
|
||||
actions = {a for a, _d in self._audit_actions(self.owner_id)}
|
||||
self.assertIn("push_subscription_created", actions)
|
||||
s, _b = _req(self.port, "DELETE",
|
||||
"/projects/%s/push/subscriptions/%s"
|
||||
% (self.project_id, sub["id"]), key=self.owner_key)
|
||||
self.assertEqual(s, 200, _b)
|
||||
actions = {a for a, _d in self._audit_actions(self.owner_id)}
|
||||
self.assertIn("push_subscription_deleted", actions)
|
||||
# 内网目标:分发不 POST、sent 为空、错误含 delivery failed。
|
||||
s, sub = _req(self.port, "POST",
|
||||
"/projects/%s/push/subscriptions" % self.project_id,
|
||||
{"channel": "webhook",
|
||||
"target": "http://127.0.0.1:9/hook",
|
||||
"preference": "instant"}, self.owner_key)
|
||||
self.assertEqual(s, 201, sub)
|
||||
s, body = _req(self.port, "POST", "/admin/push/dispatch",
|
||||
{"event_type": "failed",
|
||||
"project_id": self.project_id,
|
||||
"payload": {"task_id": "t1"}}, self.owner_key)
|
||||
self.assertEqual(s, 200, body)
|
||||
self.assertEqual(body["sent"], [])
|
||||
self.assertEqual(len(body["errors"]), 1)
|
||||
self.assertIn("delivery failed", body["errors"][0]["reason"])
|
||||
actions = {a for a, _d in self._audit_actions(self.owner_id)}
|
||||
self.assertIn("admin_push_dispatch", actions)
|
||||
|
||||
# -- 配额双字段与控制台文案 --------------------------------------------
|
||||
def test_quota_payload_has_unpriced_fields(self):
|
||||
s, body = _req(self.port, "GET", "/users/me/quota",
|
||||
key=self.owner_key)
|
||||
self.assertEqual(s, 200, body)
|
||||
self.assertIn("unpriced_tokens_used", body)
|
||||
self.assertIn("unpriced_token_cap", body)
|
||||
|
||||
def test_console_p1_markers(self):
|
||||
conn = http.client.HTTPConnection(
|
||||
"127.0.0.1", self.port, timeout=10)
|
||||
try:
|
||||
conn.request("GET", "/console")
|
||||
resp = conn.getresponse()
|
||||
text = resp.read().decode("utf-8")
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(resp.status, 200)
|
||||
for marker in (
|
||||
"登记评审需求",
|
||||
"记录允许",
|
||||
"记录拒绝",
|
||||
"当前版本只留回执,不拦截实时行为",
|
||||
"模型回合结束、改动在 worktree 里尚未合并",
|
||||
"控制面已标记 merged,不代表 git 主分支已合入",
|
||||
"协调员与 goal/计划互斥",
|
||||
"仅 failed/interrupted/merge-failed 可重试",
|
||||
"仅 maintainer 可合并"):
|
||||
self.assertIn(marker, text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,534 @@
|
||||
"""P1 执行面针对性单测(只测本路改动,不跑全量三套)。
|
||||
|
||||
运行(在 ~/workspace/codex-app-server-p1security 下):
|
||||
python3 -m unittest discover -s cloud/tests -p "test_p1_security_worker.py" -v
|
||||
|
||||
说明:只用注入的假 runner 与测试假值,不跑真实模型回合、不联网、
|
||||
不登录任何服务;env 里的值只是占位用的假名,绝不是真实密钥。
|
||||
本文件只归本路(执行面)所有,不碰 test_p1_security_server.py。
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
try:
|
||||
import cloud.sandbox as sandbox_mod
|
||||
from cloud.cost import CostStore
|
||||
from cloud.pr_exec import enqueue_pr_job
|
||||
from cloud.sandbox import DockerExecutor, LocalExecutor
|
||||
from cloud.server import init_db
|
||||
from cloud.worker import Worker, _code_signature
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
import cloud.sandbox as sandbox_mod
|
||||
from cloud.cost import CostStore
|
||||
from cloud.pr_exec import enqueue_pr_job
|
||||
from cloud.sandbox import DockerExecutor, LocalExecutor
|
||||
from cloud.server import init_db
|
||||
from cloud.worker import Worker, _code_signature
|
||||
|
||||
FAKE_ENV = {"XCM_API_KEY": "test-placeholder"}
|
||||
|
||||
|
||||
def _completed_runner(cmd):
|
||||
if "dispatch" in cmd:
|
||||
name = cmd[cmd.index("--name") + 1]
|
||||
else:
|
||||
name = cmd[cmd.index("--task") + 1]
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout=json.dumps({"task": name, "status": "completed",
|
||||
"text": "done"}),
|
||||
stderr="")
|
||||
|
||||
|
||||
def _usage_runner(usage, calls):
|
||||
def fake(cmd):
|
||||
calls.append(list(cmd))
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout=json.dumps({"status": "completed", "text": "done",
|
||||
"usage": usage}),
|
||||
stderr="")
|
||||
return fake
|
||||
|
||||
|
||||
def _insert_task(db_path, task_id, created_by="tester", status="queued",
|
||||
project_id="p1", **cols):
|
||||
now = time.time()
|
||||
conn = sqlite3.connect(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(?, ?, ?, ?, NULL, '[]', 0, ?, ?, ?, NULL)",
|
||||
(task_id, project_id, "do work", status, created_by, now, now),
|
||||
)
|
||||
for key, value in cols.items():
|
||||
conn.execute(
|
||||
"UPDATE tasks SET %s = ? WHERE id = ?" % key,
|
||||
(value, task_id))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _insert_quota(db_path, user_id, micros):
|
||||
conn = sqlite3.connect(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 INTO user_quotas(user_id, quota_usd_micros, updated_at)"
|
||||
" VALUES(?, ?, ?)",
|
||||
(user_id, micros, time.time()))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _insert_cost(db_path, task_id, model, input_tokens, output_tokens,
|
||||
micros, source="actual"):
|
||||
_insert_task(db_path, task_id, status="completed")
|
||||
store = CostStore(db_path)
|
||||
try:
|
||||
store.record(task_id, "p1", model, input_tokens=input_tokens,
|
||||
output_tokens=output_tokens, usd_micros=micros,
|
||||
source=source)
|
||||
finally:
|
||||
try:
|
||||
store._conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _status(db_path, task_id):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT status FROM tasks WHERE id = ?", (task_id,)).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
return row[0] if row else None
|
||||
|
||||
|
||||
def _error_text(db_path, task_id):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT error_text FROM tasks WHERE id = ?",
|
||||
(task_id,)).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
return row[0] if row else None
|
||||
|
||||
|
||||
class P1WorkerTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
init_db(self.db_path)
|
||||
self.calls = []
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _worker(self, runner=None, **kwargs):
|
||||
kwargs.setdefault("env", FAKE_ENV)
|
||||
return Worker(self.db_path, "/fake/host.py",
|
||||
runner=runner or _completed_runner, **kwargs)
|
||||
|
||||
# 1. 无配额行按默认 $10 额度(不再是不限额)
|
||||
def test_missing_quota_row_uses_default_quota(self):
|
||||
_insert_cost(self.db_path, "donor-a", "gpt-6.1-sol",
|
||||
1, 0, 11000000)
|
||||
_insert_cost(self.db_path, "donor-b", "gpt-6.1-sol",
|
||||
1, 0, 9000000)
|
||||
worker = self._worker()
|
||||
over = worker._quota_exceeded(
|
||||
{"created_by": "tester", "project_id": "p1"}, None)
|
||||
# donor-a 归 tester:已花 $11 > 默认 $10,挡住。
|
||||
self.assertIsNotNone(over)
|
||||
self.assertIn("quota_exceeded", over)
|
||||
# donor-b 同样归 tester 会把总额推得更高,单独测另一用户:
|
||||
_insert_task(self.db_path, "donor-c", created_by="rich",
|
||||
status="completed")
|
||||
store = CostStore(self.db_path)
|
||||
try:
|
||||
store.record("donor-c", "p1", "gpt-6.1-sol", input_tokens=1,
|
||||
output_tokens=0, usd_micros=9000000,
|
||||
source="actual")
|
||||
finally:
|
||||
try:
|
||||
store._conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
under = worker._quota_exceeded(
|
||||
{"created_by": "rich", "project_id": "p1"}, None)
|
||||
self.assertIsNone(under)
|
||||
|
||||
# 2. 显式 NULL 配额行仍是不限额
|
||||
def test_explicit_null_quota_row_stays_unlimited(self):
|
||||
_insert_quota(self.db_path, "tester", None)
|
||||
_insert_cost(self.db_path, "donor", "gpt-6.1-sol",
|
||||
1, 0, 999000000)
|
||||
worker = self._worker()
|
||||
self.assertIsNone(worker._quota_exceeded(
|
||||
{"created_by": "tester", "project_id": "p1"}, None))
|
||||
|
||||
# 3. 属主为空拒绝新回合(不再返回 None 放行)
|
||||
def test_ownerless_task_refused(self):
|
||||
worker = self._worker()
|
||||
for owner in (None, ""):
|
||||
err = worker._quota_exceeded(
|
||||
{"created_by": owner, "project_id": "p1"}, None)
|
||||
self.assertIsNotNone(err)
|
||||
self.assertIn("quota_exceeded", err)
|
||||
|
||||
# 4. 协调员子任务属主跟父任务走
|
||||
def test_child_inherits_parent_owner(self):
|
||||
_insert_quota(self.db_path, "tester", 100)
|
||||
_insert_cost(self.db_path, "donor", "gpt-6.1-sol",
|
||||
1, 0, 200)
|
||||
_insert_task(self.db_path, "parent", created_by="tester")
|
||||
_insert_task(self.db_path, "parent2", created_by="fresh")
|
||||
worker = self._worker()
|
||||
blocked = worker._quota_exceeded(
|
||||
{"created_by": None, "parent_task_id": "parent",
|
||||
"project_id": "p1"}, None)
|
||||
self.assertIsNotNone(blocked)
|
||||
self.assertIn("quota_exceeded", blocked)
|
||||
# fresh 用户无配额行、未花钱:默认额度内放行。
|
||||
self.assertIsNone(worker._quota_exceeded(
|
||||
{"created_by": None, "parent_task_id": "parent2",
|
||||
"project_id": "p1"}, None))
|
||||
|
||||
# 5. 查库失败 fail-closed(新回合拒绝,不自动放行)
|
||||
def test_quota_db_failure_is_fail_closed(self):
|
||||
worker = self._worker()
|
||||
with mock.patch("cloud.worker._db_connect",
|
||||
side_effect=RuntimeError("test db down")):
|
||||
err = worker._quota_exceeded(
|
||||
{"created_by": "tester", "project_id": "p1"}, None)
|
||||
self.assertIsNotNone(err)
|
||||
self.assertTrue(err.startswith("quota_check_failed:"))
|
||||
|
||||
# 6. Docker 配额看门狗:命中杀容器、未命中不杀
|
||||
def test_docker_quota_watchdog(self):
|
||||
killed = []
|
||||
|
||||
def fake(cmd, timeout=None):
|
||||
killed.append(list(cmd))
|
||||
return SimpleNamespace(returncode=0, stdout="", stderr="")
|
||||
|
||||
executor = DockerExecutor(runner=fake)
|
||||
worker = Worker(self.db_path, "/fake/host.py",
|
||||
env=FAKE_ENV, executor=executor)
|
||||
# 未命中:0.2 秒后停表,kill 不应被调。
|
||||
stop = threading.Event()
|
||||
threading.Timer(0.2, stop.set).start()
|
||||
worker._docker_quota_watchdog(
|
||||
"T1", "codex-test", stop, lambda: False, interval=0.05)
|
||||
self.assertEqual(killed, [])
|
||||
# 命中:立即杀容器并记命中标记。
|
||||
stop2 = threading.Event()
|
||||
worker._docker_quota_watchdog(
|
||||
"T2", "codex-test", stop2, lambda: True, interval=0.05)
|
||||
self.assertTrue(any("kill" in c for c in killed))
|
||||
self.assertTrue(worker._take_quota_hit("T2"))
|
||||
# 命中已消费,第二次取不到(不污染下一轮)。
|
||||
self.assertFalse(worker._take_quota_hit("T2"))
|
||||
|
||||
# 7. usage_cost_usd 按实际成本累加,max_cost_usd 不再是死参数
|
||||
def test_usage_cost_accumulates_and_budget_trips(self):
|
||||
_insert_task(self.db_path, "T1")
|
||||
worker = self._worker(
|
||||
runner=_usage_runner(
|
||||
{"input_tokens": 10000, "cached_input_tokens": 0,
|
||||
"output_tokens": 1000, "total_tokens": 11000},
|
||||
self.calls),
|
||||
model="gpt-6.1-sol", max_cost_usd=0.01)
|
||||
self.assertFalse(worker._budget_exceeded())
|
||||
result = worker.run_one()
|
||||
self.assertEqual(result["status"], "completed")
|
||||
# 10000*$2 + 1000*$10 = $0.03,超过 $0.01 上限。
|
||||
self.assertAlmostEqual(worker.usage_cost_usd, 0.03, places=6)
|
||||
self.assertTrue(worker._budget_exceeded())
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT usd_micros, source FROM cost_records"
|
||||
" WHERE task_id = 'T1'").fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(row, (30000, "actual"))
|
||||
|
||||
# 8. 取消分支先记成本再返回(与配额中断收尾同款记账点)
|
||||
def test_cancel_branch_records_cost(self):
|
||||
worker = self._worker()
|
||||
task = {"id": "T1", "project_id": "p1", "prompt": "do work"}
|
||||
with worker._db_lock:
|
||||
cost_id = worker._record_usage_cost(
|
||||
task, "gpt-6.1-sol",
|
||||
{"input_tokens": 10000, "cached_input_tokens": 0,
|
||||
"output_tokens": 1000, "total_tokens": 11000},
|
||||
"done", 0)
|
||||
self.assertIsNotNone(cost_id)
|
||||
self.assertAlmostEqual(worker.usage_cost_usd, 0.03, places=6)
|
||||
# 无用量时回落估算,不累加 usage_cost_usd。
|
||||
with worker._db_lock:
|
||||
worker._record_usage_cost(
|
||||
{"id": "T2", "project_id": "p1", "prompt": "do work"},
|
||||
"gpt-6.1-sol", None, "done", 0)
|
||||
self.assertAlmostEqual(worker.usage_cost_usd, 0.03, places=6)
|
||||
|
||||
# 9. 估算行不计入无价 token 上限
|
||||
def test_estimated_rows_excluded_from_unpriced_cap(self):
|
||||
_insert_cost(self.db_path, "donor-est", "gpt-5.6-terra",
|
||||
100000, 0, None, source="estimated")
|
||||
worker = self._worker()
|
||||
old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP")
|
||||
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100"
|
||||
try:
|
||||
# 10 万估算 token 在上限 100 下仍放行(估算不是真花费)。
|
||||
self.assertIsNone(worker._quota_exceeded(
|
||||
{"created_by": "tester", "project_id": "p1"}, None))
|
||||
_insert_cost(self.db_path, "donor-act", "gpt-5.6-terra",
|
||||
75, 75, None, source="actual")
|
||||
err = worker._quota_exceeded(
|
||||
{"created_by": "tester", "project_id": "p1"}, None)
|
||||
self.assertIsNotNone(err)
|
||||
self.assertIn("无价模型 token", err)
|
||||
finally:
|
||||
if old is None:
|
||||
os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None)
|
||||
else:
|
||||
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old
|
||||
|
||||
# 10. cache_write 不计入无价上限(只 input+output 计)
|
||||
def test_cache_write_excluded_from_unpriced_cap(self):
|
||||
worker = self._worker()
|
||||
old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP")
|
||||
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100"
|
||||
try:
|
||||
usage = {"input_tokens": 10, "output_tokens": 5,
|
||||
"cache_write_input_tokens": 9999999,
|
||||
"total_tokens": 15}
|
||||
self.assertIsNone(worker._quota_exceeded(
|
||||
{"created_by": "tester", "project_id": "p1",
|
||||
"model": "gpt-5.6-terra"}, usage))
|
||||
usage2 = {"input_tokens": 200, "output_tokens": 0,
|
||||
"cache_write_input_tokens": 0,
|
||||
"total_tokens": 200}
|
||||
err = worker._quota_exceeded(
|
||||
{"created_by": "tester", "project_id": "p1",
|
||||
"model": "gpt-5.6-terra"}, usage2)
|
||||
self.assertIsNotNone(err)
|
||||
finally:
|
||||
if old is None:
|
||||
os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None)
|
||||
else:
|
||||
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old
|
||||
|
||||
# 11. 协调员子任务继承父任务模型/强度,分工显式值优先
|
||||
def test_coordinator_child_inherits_model_effort(self):
|
||||
worker = self._worker()
|
||||
task = {"id": "coord", "project_id": "p1", "prompt": "coord",
|
||||
"priority": 0, "created_by": "tester", "model": "gpt-6.1-sol",
|
||||
"effort": "high", "executor": None, "environment_id": None}
|
||||
children = worker._coord_sync_children(task, [
|
||||
{"name": "a", "prompt": "do a", "model": None,
|
||||
"effort": None, "executor": None},
|
||||
{"name": "b", "prompt": "do b", "model": "gpt-5.6-terra",
|
||||
"effort": "low", "executor": None},
|
||||
])
|
||||
self.assertEqual(len(children), 2)
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
rows = {r[0]: (r[1], r[2]) for r in conn.execute(
|
||||
"SELECT coordinator_part, model, effort FROM tasks"
|
||||
" WHERE parent_task_id = 'coord'").fetchall()}
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(rows["a"], ("gpt-6.1-sol", "high"))
|
||||
self.assertEqual(rows["b"], ("gpt-5.6-terra", "low"))
|
||||
|
||||
# 12. 断点续跑:Local/Docker/popen 三路都不传 --goal-text
|
||||
def test_goal_resume_breakpoint_omits_goal_text(self):
|
||||
calls = []
|
||||
|
||||
def fake(cmd):
|
||||
calls.append(list(cmd))
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout="goal met after rounds\n"
|
||||
"usage input=10 cached=0 cache_write=0 output=5 "
|
||||
"reasoning=0 total=15 completions=1\n",
|
||||
stderr="")
|
||||
|
||||
task = {"id": "T1", "project_id": "p1", "prompt": "orig goal",
|
||||
"verify_command": "true", "goal_max_rounds": 3}
|
||||
local = LocalExecutor("/fake/host.py", runner=fake)
|
||||
local.goal_resume_task(task, None)
|
||||
self.assertIn("goal-resume", calls[-1])
|
||||
self.assertNotIn("--goal-text", calls[-1])
|
||||
docker_calls = []
|
||||
|
||||
def dfake(cmd, timeout=None):
|
||||
docker_calls.append(list(cmd))
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout="goal met after rounds\n"
|
||||
"usage input=10 cached=0 cache_write=0 output=5 "
|
||||
"reasoning=0 total=15 completions=1\n",
|
||||
stderr="")
|
||||
|
||||
docker = DockerExecutor(runner=dfake, env=FAKE_ENV)
|
||||
with mock.patch("cloud.sandbox.shutil.which",
|
||||
return_value="/usr/bin/docker"):
|
||||
docker.goal_resume_task(task, None)
|
||||
self.assertIn("goal-resume", docker_calls[-1])
|
||||
self.assertNotIn("--goal-text", docker_calls[-1])
|
||||
with mock.patch.object(sandbox_mod.subprocess, "Popen") as popen:
|
||||
popen.return_value = SimpleNamespace(pid=9999)
|
||||
local.popen_task(
|
||||
{"id": "T1", "prompt": "resume msg",
|
||||
"_resume_kind": "breakpoint"})
|
||||
# popen_task 默认 kind=dispatch;显式测 goal-resume:
|
||||
local.popen_task(
|
||||
{"id": "T1", "prompt": "resume msg",
|
||||
"_resume_kind": "breakpoint"}, kind="goal-resume")
|
||||
cmd = popen.call_args[0][0]
|
||||
self.assertIn("goal-resume", cmd)
|
||||
self.assertNotIn("--goal-text", cmd)
|
||||
|
||||
# 13. 开新目标:goal-resume 带 --goal-text(旧行为保留)
|
||||
def test_goal_resume_new_goal_keeps_goal_text(self):
|
||||
calls = []
|
||||
|
||||
def fake(cmd):
|
||||
calls.append(list(cmd))
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout="goal met after rounds\n"
|
||||
"usage input=10 cached=0 cache_write=0 output=5 "
|
||||
"reasoning=0 total=15 completions=1\n",
|
||||
stderr="")
|
||||
|
||||
task = {"id": "T1", "project_id": "p1", "prompt": "new goal"}
|
||||
local = LocalExecutor("/fake/host.py", runner=fake)
|
||||
local.goal_resume_task(task, "new goal")
|
||||
self.assertIn("--goal-text", calls[-1])
|
||||
self.assertIn("new goal", calls[-1])
|
||||
|
||||
# 14. goal 轮次上限:sandbox loud fail + worker 认领后 failed
|
||||
def test_goal_max_rounds_cap(self):
|
||||
task = {"id": "T1", "project_id": "p1", "prompt": "g",
|
||||
"verify_command": "true", "goal_max_rounds": 101}
|
||||
local = LocalExecutor("/fake/host.py", runner=_completed_runner)
|
||||
with self.assertRaises(ValueError):
|
||||
local.goal_task(task)
|
||||
_insert_task(self.db_path, "T1", verify_command="true",
|
||||
goal_max_rounds=101)
|
||||
worker = self._worker()
|
||||
result = worker.run_one()
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertIn("上限", result["error"])
|
||||
self.assertEqual(_status(self.db_path, "T1"), "failed")
|
||||
self.assertEqual(self.calls, [])
|
||||
|
||||
# 15. PR 非 local:显式 needs_local_worker,任务与排队都不动
|
||||
def test_pr_non_local_needs_local_worker(self):
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
job_id = enqueue_pr_job(conn, "T1", "p1", "test-host")
|
||||
finally:
|
||||
conn.close()
|
||||
executor = DockerExecutor(runner=lambda cmd, timeout=None:
|
||||
SimpleNamespace(returncode=0, stdout="",
|
||||
stderr=""))
|
||||
worker = Worker(self.db_path, "/fake/host.py", env=FAKE_ENV,
|
||||
executor=executor, host_name="test-host")
|
||||
results = worker._run_pr_jobs()
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(results[0]["status"], "needs_local_worker")
|
||||
self.assertEqual(results[0]["pr_job_id"], job_id)
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT status FROM pr_jobs WHERE id = ?",
|
||||
(job_id,)).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(row[0], "queued")
|
||||
|
||||
# 16. 多文件签名 + 签名缺失/已变时认领 fail-closed
|
||||
def test_code_signature_multi_file_and_claim_refusal(self):
|
||||
sig = _code_signature("/fake/host.py")
|
||||
self.assertIsNotNone(sig)
|
||||
self.assertEqual(len(sig), 6)
|
||||
for entry in sig:
|
||||
self.assertEqual(len(entry), 2)
|
||||
_insert_task(self.db_path, "T1")
|
||||
worker = self._worker()
|
||||
worker._code_signature = ("stale-signature",)
|
||||
result = worker.run_one()
|
||||
self.assertFalse(result["dispatched"])
|
||||
self.assertEqual(result["reason"], "already_claimed")
|
||||
self.assertEqual(_status(self.db_path, "T1"), "queued")
|
||||
worker._code_signature = None
|
||||
result = worker.run_one()
|
||||
self.assertFalse(result["dispatched"])
|
||||
self.assertEqual(_status(self.db_path, "T1"), "queued")
|
||||
|
||||
# 17. goal 首轮撞已存在 → 断点续跑一次,不直接 failed
|
||||
def test_goal_already_exists_falls_back_to_breakpoint(self):
|
||||
calls = []
|
||||
|
||||
def fake(cmd):
|
||||
calls.append(list(cmd))
|
||||
if "goal-resume" in cmd:
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout="round 1: turn=completed verify_rc=0 "
|
||||
"verify_out=\ngoal met after rounds\n"
|
||||
"usage input=10 cached=0 cache_write=0 output=5 "
|
||||
"reasoning=0 total=15 completions=1\n",
|
||||
stderr="")
|
||||
if "goal" in cmd:
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout="host 任务已存在:T1 已存在",
|
||||
stderr="")
|
||||
return _completed_runner(cmd)
|
||||
|
||||
_insert_task(self.db_path, "T1", verify_command="true",
|
||||
goal_max_rounds=2)
|
||||
worker = self._worker(runner=fake)
|
||||
result = worker.run_one()
|
||||
self.assertEqual(result["status"], "completed")
|
||||
# 注意:cmd 是元素列表,"goal" in cmd 只匹配恰好等于
|
||||
# "goal" 的元素,goal-resume 命令里并没有 "goal" 元素,
|
||||
# 所以按“任一元素含 goal 子串”过滤。
|
||||
goal_cmds = [c for c in calls
|
||||
if any("goal" in e for e in c)]
|
||||
self.assertTrue(any("goal-resume" in c for c in goal_cmds))
|
||||
resume_cmd = [c for c in goal_cmds if "goal-resume" in c][0]
|
||||
self.assertNotIn("--goal-text", resume_cmd)
|
||||
self.assertEqual(_status(self.db_path, "T1"), "completed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -307,7 +307,12 @@ class WorkerPrJobsTestCase(unittest.TestCase):
|
||||
enqueue_pr_job(conn, "t1", "p1", worker.host_name)
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(worker._run_pr_jobs(), [])
|
||||
# P1:Docker worker 不再静默返回空,改为显式 needs_local_worker,
|
||||
# 任务保持 queued 等 local worker 认领。
|
||||
jobs = worker._run_pr_jobs()
|
||||
self.assertEqual(len(jobs), 1)
|
||||
self.assertEqual(jobs[0]["status"], "needs_local_worker")
|
||||
self.assertEqual(jobs[0]["task_id"], "t1")
|
||||
|
||||
def test_autoscaler_forces_host_name(self):
|
||||
scaler = Autoscaler(self.db_path, [])
|
||||
|
||||
@@ -41,6 +41,9 @@ class WebhookExecTestCase(unittest.TestCase):
|
||||
self.owner_id, self.owner_key = self.auth.create_user("owner")
|
||||
self._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
|
||||
os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
|
||||
# P1:备份 out_dir 限白名单根;本测试的 tmp 目录即白名单根。
|
||||
self._old_backup_root = os.environ.get("CLOUD_BACKUP_ROOT")
|
||||
os.environ["CLOUD_BACKUP_ROOT"] = self.tmp.name
|
||||
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]
|
||||
@@ -67,6 +70,10 @@ class WebhookExecTestCase(unittest.TestCase):
|
||||
os.environ.pop("CODEX_GLOBAL_ADMINS", None)
|
||||
else:
|
||||
os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins
|
||||
if getattr(self, "_old_backup_root", None) is None:
|
||||
os.environ.pop("CLOUD_BACKUP_ROOT", None)
|
||||
else:
|
||||
os.environ["CLOUD_BACKUP_ROOT"] = self._old_backup_root
|
||||
|
||||
def _req(self, method, path, body=None, key=None, headers=None,
|
||||
raw_body=None):
|
||||
|
||||
@@ -336,6 +336,11 @@ class TestRealSenders(PushTestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
# P1 SSRF 门默认挡 loopback;本地录制服务器测试显式放行,
|
||||
# 生产默认值仍关闭(不要把本开关当成生产默认)。
|
||||
self._old_push_allow = os.environ.get(
|
||||
"CODEX_PUSH_ALLOW_PRIVATE_TARGETS")
|
||||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = "1"
|
||||
_RecordingHandler.requests = []
|
||||
_RecordingHandler.status = 200
|
||||
self.server = http.server.ThreadingHTTPServer(
|
||||
@@ -349,6 +354,11 @@ class TestRealSenders(PushTestCase):
|
||||
self.server.shutdown()
|
||||
self.server.server_close()
|
||||
self.thread.join(timeout=5)
|
||||
if getattr(self, "_old_push_allow", None) is None:
|
||||
os.environ.pop("CODEX_PUSH_ALLOW_PRIVATE_TARGETS", None)
|
||||
else:
|
||||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = (
|
||||
self._old_push_allow)
|
||||
super().tearDown()
|
||||
|
||||
def _url(self, path="/hook"):
|
||||
|
||||
@@ -331,7 +331,8 @@ class TestGoalExecutors(unittest.TestCase):
|
||||
{"id": "t1", "prompt": "do it",
|
||||
"verify_command": "true"})
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertIn("请新建任务", result["goal_error"])
|
||||
# P1 新语义:goal 失败请用继续而非重试(旧“请新建任务”已废)。
|
||||
self.assertIn("请用继续而非重试", result["goal_error"])
|
||||
|
||||
def test_local_goal_task_without_verify_is_llm_decided(self):
|
||||
# 新语义:没有验收命令不再抛错——走模型自决分支:
|
||||
|
||||
@@ -114,6 +114,11 @@ class _Listener:
|
||||
|
||||
class GoalPrNotifyTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# P1 SSRF 门默认挡 loopback;本地监听器测试显式放行,
|
||||
# 生产默认值仍关闭。
|
||||
self._old_push_allow = os.environ.get(
|
||||
"CODEX_PUSH_ALLOW_PRIVATE_TARGETS")
|
||||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = "1"
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
self.auth = AuthStore(self.db_path)
|
||||
@@ -133,6 +138,11 @@ class GoalPrNotifyTestCase(unittest.TestCase):
|
||||
self.server.shutdown()
|
||||
self.server.server_close()
|
||||
self.thread.join(timeout=5)
|
||||
if getattr(self, "_old_push_allow", None) is None:
|
||||
os.environ.pop("CODEX_PUSH_ALLOW_PRIVATE_TARGETS", None)
|
||||
else:
|
||||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = (
|
||||
self._old_push_allow)
|
||||
try:
|
||||
self.auth.conn.close()
|
||||
except Exception:
|
||||
|
||||
@@ -74,6 +74,11 @@ def _completed_runner(cmd):
|
||||
|
||||
class WorkerTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# P1 SSRF 门默认挡 loopback;通知测试用本地监听器,显式放行,
|
||||
# 生产默认值仍关闭。
|
||||
self._old_push_allow = os.environ.get(
|
||||
"CODEX_PUSH_ALLOW_PRIVATE_TARGETS")
|
||||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = "1"
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
init_db(self.db_path)
|
||||
@@ -81,6 +86,11 @@ class WorkerTestCase(unittest.TestCase):
|
||||
self.calls = []
|
||||
|
||||
def tearDown(self):
|
||||
if getattr(self, "_old_push_allow", None) is None:
|
||||
os.environ.pop("CODEX_PUSH_ALLOW_PRIVATE_TARGETS", None)
|
||||
else:
|
||||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = (
|
||||
self._old_push_allow)
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _worker(self, runner=None, **kwargs):
|
||||
|
||||
Reference in new issue
Block a user