host
7c444a688b
server：POST /uploads（multipart，413 预检、路径/数量/大小校验、 24h 过期包、used 防一包多用）+ 建任务绑定（同事务写附件行 + 包置 used，cloud 执行器 400）+ 列表/详情带 attachments + 取消删暂存 + 分支带走附件清单；db.py PG 建表 + init_db 补表。 worker：_task_seed_dir（无附件/目录已删静默 None）+ dispatch/ goal 三路（同步本地、popen、Docker 只读挂载 /tmp/codex-seed） + cloud 执行器兜底失败 + 终局删暂存。 console：文件/文件夹/图片三路输入，先上传拿 upload_id 再建任务， 列表 [附件×n]、详情列附件路径。 附带：popen 的 goal 命令此前强制要求 verify_command，goal 新语义 任务到本地 worker（走 popen）会直接报错，改成与同步 goal_task 一致（无验收不传 --verify，轮次 0 原样传）。 测试：新增 test_task_attachments.py 15 个。
419 lines
17 KiB
Python
419 lines
17 KiB
Python
"""任务附件单测：上传包校验、建任务绑定、派发 seed、终局清理。
|
||
|
||
不跑真实模型回合、不联网：server 侧走本地临时端口 + 临时
|
||
SQLite；worker/sandbox 侧注入假 runner 录命令；host seed 侧用
|
||
临时目录。图片与普通文件不区分，一律是 worktree 里的文件。
|
||
"""
|
||
|
||
import http.client
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
|
||
try:
|
||
from cloud.auth import AuthStore
|
||
from cloud.sandbox import DockerExecutor, LocalExecutor
|
||
from cloud.server import create_server, init_db
|
||
from cloud.uploads import (
|
||
MAX_FILES,
|
||
parse_multipart,
|
||
parse_paths_field,
|
||
save_bundle,
|
||
validate_bundle,
|
||
validate_rel_path,
|
||
)
|
||
from cloud.worker import Worker
|
||
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.sandbox import DockerExecutor, LocalExecutor
|
||
from cloud.server import create_server, init_db
|
||
from cloud.uploads import (
|
||
MAX_FILES,
|
||
parse_multipart,
|
||
parse_paths_field,
|
||
save_bundle,
|
||
validate_bundle,
|
||
validate_rel_path,
|
||
)
|
||
from cloud.worker import Worker
|
||
|
||
|
||
def _multipart_body(boundary, fields, files):
|
||
"""拼一个 multipart body：fields {名: 文本}，files
|
||
[(part 名, 文件名, bytes)]。"""
|
||
out = b""
|
||
for name, value in fields.items():
|
||
out += ("--%s\r\nContent-Disposition: form-data; name=\"%s\"\r\n\r\n"
|
||
% (boundary, name)).encode("utf-8")
|
||
out += value.encode("utf-8") + b"\r\n"
|
||
for part_name, filename, data in files:
|
||
out += ("--%s\r\nContent-Disposition: form-data; name=\"%s\"; "
|
||
"filename=\"%s\"\r\nContent-Type: application/octet-stream"
|
||
"\r\n\r\n" % (boundary, part_name, filename)).encode("utf-8")
|
||
out += data + b"\r\n"
|
||
out += ("--%s--\r\n" % boundary).encode("utf-8")
|
||
return out
|
||
|
||
|
||
class UploadsModuleTestCase(unittest.TestCase):
|
||
def test_parse_roundtrip_nested_paths(self):
|
||
boundary = "TESTBOUNDARY"
|
||
body = _multipart_body(
|
||
boundary,
|
||
{"paths": json.dumps(["a.txt", "sub/b.png"])},
|
||
[("file0", "a.txt", b"hello"),
|
||
("file1", "b.png", b"\x89PNG")])
|
||
fields, files = parse_multipart(
|
||
body, "multipart/form-data; boundary=" + boundary)
|
||
self.assertEqual(fields["paths"],
|
||
json.dumps(["a.txt", "sub/b.png"]))
|
||
self.assertEqual([n for n, _d in files], ["a.txt", "b.png"])
|
||
paths = parse_paths_field(fields, files)
|
||
self.assertEqual(paths, ["a.txt", "sub/b.png"])
|
||
|
||
def test_parse_paths_count_mismatch(self):
|
||
with self.assertRaises(ValueError):
|
||
parse_paths_field({"paths": json.dumps(["a.txt"])},
|
||
[("file0", b"x"), ("file1", b"y")])
|
||
|
||
def test_validate_rel_path(self):
|
||
for good in ["a.txt", "sub/b/c.png", "a-b_c.d/e f.txt", ".hidden/x"]:
|
||
self.assertIsNone(validate_rel_path(good), good)
|
||
for bad in ["/abs", "../x", "a/../b", "a//b", "", "a\\b",
|
||
"C:/x", "a\x00b", ".", "..", "a/./b"]:
|
||
self.assertIsNotNone(validate_rel_path(bad), bad)
|
||
|
||
def test_validate_bundle_limits(self):
|
||
self.assertIsNotNone(validate_bundle(
|
||
["f%d.txt" % i for i in range(MAX_FILES + 1)],
|
||
[1] * (MAX_FILES + 1)))
|
||
self.assertIsNotNone(validate_bundle(
|
||
["big.bin"], [21 * 1024 * 1024]))
|
||
self.assertIsNotNone(validate_bundle(
|
||
["a.bin", "b.bin"], [60 * 1024 * 1024, 60 * 1024 * 1024]))
|
||
self.assertIsNone(validate_bundle(["a.txt"], [5]))
|
||
|
||
def test_save_bundle_roundtrip_and_traversal(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
bundle_id, saved = save_bundle(
|
||
tmp, ["a.txt", "sub/b.txt"], [b"hello", b"world"])
|
||
self.assertEqual([p for p, _s in saved],
|
||
["a.txt", "sub/b.txt"])
|
||
with open(os.path.join(tmp, bundle_id, "sub", "b.txt"),
|
||
"rb") as f:
|
||
self.assertEqual(f.read(), b"world")
|
||
with self.assertRaises(ValueError):
|
||
save_bundle(tmp, ["../escape.txt"], [b"x"])
|
||
|
||
|
||
class ServerAttachmentsTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
self.upload_dir = os.path.join(self.tmp.name, "uploads")
|
||
self.auth = AuthStore(self.db_path)
|
||
self.owner_id, self.owner_key = self.auth.create_user("owner")
|
||
self.other_id, self.other_key = self.auth.create_user("other")
|
||
self.server = create_server(self.db_path, "127.0.0.1", 0)
|
||
self.server.upload_dir = self.upload_dir
|
||
self.port = self.server.server_address[1]
|
||
self.thread = threading.Thread(
|
||
target=self.server.serve_forever, daemon=True)
|
||
self.thread.start()
|
||
status, self.project = self._req(
|
||
"POST", "/projects", {"name": "demo"}, self.owner_key)
|
||
self.assertEqual(status, 201)
|
||
|
||
def tearDown(self):
|
||
self.server.shutdown()
|
||
self.server.server_close()
|
||
self.thread.join(timeout=5)
|
||
try:
|
||
self.auth.conn.close()
|
||
except Exception:
|
||
pass
|
||
self.tmp.cleanup()
|
||
|
||
def _req(self, method, path, body=None, key=None,
|
||
content_type="application/json"):
|
||
conn = http.client.HTTPConnection(
|
||
"127.0.0.1", self.port, timeout=10)
|
||
hdrs = {}
|
||
payload = None
|
||
if key:
|
||
hdrs["X-API-Key"] = key
|
||
if body is not None:
|
||
if isinstance(body, bytes):
|
||
payload = body
|
||
else:
|
||
payload = json.dumps(body)
|
||
hdrs["Content-Type"] = content_type
|
||
conn.request(method, path, body=payload, headers=hdrs)
|
||
resp = conn.getresponse()
|
||
raw = resp.read()
|
||
conn.close()
|
||
try:
|
||
return resp.status, json.loads(raw.decode("utf-8"))
|
||
except (ValueError, UnicodeDecodeError):
|
||
return resp.status, {}
|
||
|
||
def _upload(self, paths, contents, key=None):
|
||
body = _multipart_body(
|
||
"TESTBOUNDARY", {"paths": json.dumps(paths)},
|
||
[("file%d" % i, os.path.basename(p) or "f", data)
|
||
for i, (p, data) in enumerate(zip(paths, contents))])
|
||
return self._req("POST", "/uploads", body,
|
||
key or self.owner_key,
|
||
content_type="multipart/form-data;"
|
||
" boundary=TESTBOUNDARY")
|
||
|
||
def _create_task(self, upload_id=None, executor=None):
|
||
body = {"project_id": self.project["id"], "prompt": "do work"}
|
||
if upload_id is not None:
|
||
body["upload_id"] = upload_id
|
||
if executor is not None:
|
||
body["executor"] = executor
|
||
return self._req("POST", "/tasks", body, self.owner_key)
|
||
|
||
def test_upload_and_bind(self):
|
||
status, up = self._upload(["a.txt", "sub/b.png"],
|
||
[b"hello", b"\x89PNG"])
|
||
self.assertEqual(status, 201, up)
|
||
self.assertEqual([f["path"] for f in up["files"]],
|
||
["a.txt", "sub/b.png"])
|
||
status, task = self._create_task(up["upload_id"])
|
||
self.assertEqual(status, 201, task)
|
||
self.assertEqual(task["attachments"], ["a.txt", "sub/b.png"])
|
||
status, got = self._req(
|
||
"GET", "/tasks/" + task["id"], key=self.owner_key)
|
||
self.assertEqual(status, 200)
|
||
self.assertEqual(got["attachments"], ["a.txt", "sub/b.png"])
|
||
status, listed = self._req(
|
||
"GET", "/tasks?project_id=" + self.project["id"],
|
||
key=self.owner_key)
|
||
self.assertEqual(status, 200)
|
||
self.assertEqual(listed["tasks"][0]["attachments"],
|
||
["a.txt", "sub/b.png"])
|
||
|
||
def test_upload_rejects_bad_path_and_empty(self):
|
||
status, body = self._upload(["../escape.txt"], [b"x"])
|
||
self.assertEqual(status, 400)
|
||
status, body = self._req("POST", "/uploads",
|
||
{"paths": "[]"}, self.owner_key)
|
||
self.assertEqual(status, 400)
|
||
|
||
def test_bind_rejects_used_unknown_other_cloud(self):
|
||
status, up = self._upload(["a.txt"], [b"hi"])
|
||
self.assertEqual(status, 201)
|
||
status, _task = self._create_task(up["upload_id"])
|
||
self.assertEqual(status, 201)
|
||
status, body = self._create_task(up["upload_id"])
|
||
self.assertEqual(status, 409)
|
||
status, body = self._create_task("no-such-bundle")
|
||
self.assertEqual(status, 400)
|
||
status, up2 = self._upload(["b.txt"], [b"hi"])
|
||
self.assertEqual(status, 201)
|
||
status, body = self._req(
|
||
"POST", "/tasks",
|
||
{"project_id": self.project["id"], "prompt": "x",
|
||
"upload_id": up2["upload_id"]}, self.other_key)
|
||
self.assertEqual(status, 403)
|
||
status, up3 = self._upload(["c.txt"], [b"hi"])
|
||
self.assertEqual(status, 201)
|
||
status, body = self._create_task(up3["upload_id"],
|
||
executor="cloud")
|
||
self.assertEqual(status, 400)
|
||
|
||
def test_cancel_deletes_bundle_dir(self):
|
||
status, up = self._upload(["a.txt"], [b"hi"])
|
||
self.assertEqual(status, 201)
|
||
bundle_dir = os.path.join(self.upload_dir, up["upload_id"])
|
||
self.assertTrue(os.path.isdir(bundle_dir))
|
||
status, task = self._create_task(up["upload_id"])
|
||
self.assertEqual(status, 201)
|
||
status, _body = self._req(
|
||
"POST", "/tasks/%s/cancel" % task["id"], {},
|
||
self.owner_key)
|
||
self.assertEqual(status, 200, _body)
|
||
self.assertFalse(os.path.exists(bundle_dir))
|
||
|
||
def test_fork_copies_attachment_rows(self):
|
||
status, up = self._upload(["a.txt"], [b"hi"])
|
||
self.assertEqual(status, 201)
|
||
status, task = self._create_task(up["upload_id"])
|
||
self.assertEqual(status, 201)
|
||
conn = sqlite3.connect(self.db_path)
|
||
conn.execute("UPDATE tasks SET status = 'completed' WHERE id = ?",
|
||
(task["id"],))
|
||
conn.commit()
|
||
conn.close()
|
||
status, forked = self._req(
|
||
"POST", "/tasks/%s/fork" % task["id"],
|
||
{"message": "continue"}, self.owner_key)
|
||
self.assertEqual(status, 201, forked)
|
||
self.assertEqual(forked["attachments"], ["a.txt"])
|
||
|
||
|
||
class SandboxSeedTestCase(unittest.TestCase):
|
||
def _recording(self, calls):
|
||
def fake(cmd):
|
||
calls.append(list(cmd))
|
||
return SimpleNamespace(returncode=0, stdout="ok", stderr="")
|
||
return fake
|
||
|
||
def test_local_dispatch_and_goal_seed_flags(self):
|
||
calls = []
|
||
ex = LocalExecutor("/fake/host.py", runner=self._recording(calls))
|
||
ex.dispatch_task({"id": "t1", "prompt": "do"},
|
||
seed_dir="/tmp/seed")
|
||
self.assertIn("--seed-dir", calls[0])
|
||
self.assertEqual(calls[0][calls[0].index("--seed-dir") + 1],
|
||
"/tmp/seed")
|
||
calls.clear()
|
||
ex.dispatch_task({"id": "t1", "prompt": "do"})
|
||
self.assertNotIn("--seed-dir", calls[0])
|
||
calls.clear()
|
||
ex.goal_task({"id": "t2", "prompt": "build it", "goal_mode": 1,
|
||
"goal": "build it", "goal_max_rounds": 0},
|
||
seed_dir="/tmp/seed")
|
||
self.assertIn("--seed-dir", calls[0])
|
||
self.assertNotIn("--verify", calls[0])
|
||
self.assertEqual(calls[0][calls[0].index("--max-rounds") + 1],
|
||
"0")
|
||
|
||
def test_popen_goal_new_semantics_and_seed(self):
|
||
import cloud.sandbox as sandbox_mod
|
||
calls = []
|
||
|
||
class FakePopen(object):
|
||
def __init__(self, cmd, **kwargs):
|
||
calls.append(list(cmd))
|
||
|
||
orig = sandbox_mod.subprocess.Popen
|
||
sandbox_mod.subprocess.Popen = FakePopen
|
||
try:
|
||
ex = LocalExecutor("/fake/host.py")
|
||
ex.popen_task({"id": "t2", "prompt": "build it",
|
||
"goal_mode": 1, "goal": "build it",
|
||
"goal_max_rounds": 0},
|
||
kind="goal", seed_dir="/tmp/seed")
|
||
ex.popen_task({"id": "t1", "prompt": "do"}, kind="run",
|
||
seed_dir="/tmp/seed")
|
||
finally:
|
||
sandbox_mod.subprocess.Popen = orig
|
||
goal_cmd, run_cmd = calls
|
||
self.assertNotIn("--verify", goal_cmd)
|
||
self.assertEqual(goal_cmd[goal_cmd.index("--max-rounds") + 1],
|
||
"0")
|
||
self.assertIn("--seed-dir", goal_cmd)
|
||
self.assertNotIn("--seed-dir", run_cmd)
|
||
|
||
def test_docker_seed_mount(self):
|
||
from unittest import mock
|
||
calls = []
|
||
ex = DockerExecutor(
|
||
"/fake/host.py", runner=self._recording(calls),
|
||
env={"XCM_API_KEY": "placeholder-not-a-real-key"})
|
||
with tempfile.TemporaryDirectory() as seed_dir, \
|
||
mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex.dispatch_task({"id": "t1", "prompt": "do"},
|
||
seed_host_dir=seed_dir)
|
||
flat = " ".join(calls[0])
|
||
self.assertIn("%s:/tmp/codex-seed:ro" % seed_dir, flat)
|
||
self.assertIn("/tmp/codex-seed", flat)
|
||
calls.clear()
|
||
ex.dispatch_task({"id": "t1", "prompt": "do"})
|
||
flat = " ".join(calls[0])
|
||
self.assertNotIn("codex-seed", flat)
|
||
|
||
|
||
class WorkerSeedTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
self.upload_dir = os.path.join(self.tmp.name, "uploads")
|
||
init_db(self.db_path)
|
||
self.env = {"XCM_API_KEY": "test-placeholder"}
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _insert_task(self, task_id, **extra):
|
||
now = time.time()
|
||
cols = ["id", "project_id", "prompt", "status",
|
||
"idempotency_key", "depends_on", "priority",
|
||
"created_by", "created_at", "updated_at",
|
||
"last_heartbeat"]
|
||
vals = [task_id, "p1", "do work", "queued", None, "[]", 0,
|
||
"tester", now, now, None]
|
||
for k, v in extra.items():
|
||
cols.append(k)
|
||
vals.append(v)
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO tasks(%s) VALUES(%s)"
|
||
% (", ".join(cols), ", ".join("?" * len(cols))), vals)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
def _insert_bundle(self, task_id, bundle_id):
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO task_attachments"
|
||
"(task_id, bundle_id, rel_path, size) VALUES(?, ?, ?, ?)",
|
||
(task_id, bundle_id, "a.txt", 2))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
def test_task_seed_dir_resolution(self):
|
||
worker = Worker(self.db_path, "/fake/host.py", env=self.env,
|
||
upload_dir=self.upload_dir)
|
||
self.assertIsNone(worker._task_seed_dir({"id": "nope"}))
|
||
self._insert_bundle("t1", "b1")
|
||
self.assertIsNone(worker._task_seed_dir({"id": "t1"}))
|
||
bundle_dir = os.path.join(self.upload_dir, "b1")
|
||
os.makedirs(bundle_dir)
|
||
self.assertEqual(worker._task_seed_dir({"id": "t1"}), bundle_dir)
|
||
|
||
def test_run_one_passes_seed_and_cleans_bundle(self):
|
||
self._insert_task("t1")
|
||
self._insert_bundle("t1", "b1")
|
||
bundle_dir = os.path.join(self.upload_dir, "b1")
|
||
os.makedirs(bundle_dir)
|
||
with open(os.path.join(bundle_dir, "a.txt"), "w") as f:
|
||
f.write("hi")
|
||
seen = []
|
||
|
||
def fake(cmd):
|
||
seen.append(list(cmd))
|
||
return SimpleNamespace(returncode=0, stdout="ok", stderr="")
|
||
|
||
worker = Worker(self.db_path, "/fake/host.py", runner=fake,
|
||
env=self.env, upload_dir=self.upload_dir)
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "completed")
|
||
dispatch_cmds = [c for c in seen if "dispatch" in c]
|
||
self.assertTrue(dispatch_cmds)
|
||
self.assertIn("--seed-dir", dispatch_cmds[0])
|
||
self.assertEqual(
|
||
dispatch_cmds[0][dispatch_cmds[0].index("--seed-dir") + 1],
|
||
bundle_dir)
|
||
self.assertFalse(os.path.exists(bundle_dir))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|