"""成本/配额/文件标记单测：task-cost / quota-status / task-files 解析、sandbox 提取、worker 投递与 worktree 解析。 不跑真实模型、不联网：解析是纯函数；worker 投递用 init_db 建的临时 SQLite 库，直接调 _deliver_marker_tools，不走 run_one 全量流程；文件读取用临时目录冒充 worktree（测 试里直接给 Worker 换 tenants_root 与 host_py_path）。 """ import json import os import sqlite3 import sys import tempfile import time import unittest from pathlib import Path from types import SimpleNamespace try: from cloud.cost import CostStore from cloud.marker_tools import ( extract_quota_status, extract_task_costs, extract_task_files, format_cost, safe_worktree_path, ) from cloud.sandbox import LocalExecutor, _parse_host_result from cloud.server import init_db 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.cost import CostStore from cloud.marker_tools import ( extract_quota_status, extract_task_costs, extract_task_files, format_cost, safe_worktree_path, ) from cloud.sandbox import LocalExecutor, _parse_host_result from cloud.server import init_db from cloud.worker import Worker FAKE_ENV = {"XCM_API_KEY": "test-placeholder"} class ExtractTestCase(unittest.TestCase): def test_cost_and_files(self): self.assertEqual(extract_task_costs( "```task-cost\n{\"task_id\": \"t1\"}\n```"), [{"task_id": "t1"}]) self.assertEqual(extract_task_costs( "```task-cost\n{\"to\": \"t1\"}\n```"), [{"task_id": "t1"}]) self.assertEqual(extract_task_costs( "```task-cost\n{}\n```"), []) self.assertEqual(extract_task_files( "```task-files\n{\"task_id\": \"t1\"}\n```"), [{"task_id": "t1", "path": ""}]) self.assertEqual(extract_task_files( "```task-files\n{\"task_id\": \"t1\", \"path\": \"a.txt\"}\n```"), [{"task_id": "t1", "path": "a.txt"}]) self.assertEqual(extract_task_files( "```task-files\n{\"path\": \"a.txt\"}\n```"), []) def test_quota_only_object(self): self.assertEqual(extract_quota_status( "```quota-status\n{}\n```"), [{}]) self.assertEqual(extract_quota_status( "```quota-status\n{\"owner\": \"someone\"}\n```"), [{}]) self.assertEqual(extract_quota_status( "```quota-status\n[]\n```"), []) def test_safe_worktree_path(self): self.assertIsNone(safe_worktree_path("/tmp/w", "../etc/passwd")) self.assertIsNone(safe_worktree_path("/tmp/w", "/etc/passwd")) self.assertEqual(safe_worktree_path("/tmp/w", ""), "/tmp/w") self.assertEqual(safe_worktree_path("/tmp/w", "a/b.txt"), "/tmp/w/a/b.txt") def test_format_cost_unpriced_separate(self): line = format_cost("t1", {"usd_micros": 0, "total_tokens": 10, "actual_tokens": 10, "estimated_tokens": 0, "turns": 1, "unpriced_actual_tokens": 10}) self.assertIn("$0.0000", line) self.assertIn("无价 actual token 10 未计入美元", line) def test_sandbox_extracts_both_executors(self): text = ("```task-cost\n{\"task_id\": \"t1\"}\n```\n" "```quota-status\n{}\n```\n" "```task-files\n{\"task_id\": \"t1\"}\n```") fake = SimpleNamespace( returncode=0, stdout=json.dumps({"status": "completed", "text": text}), stderr="") local = LocalExecutor("/fake/host.py")._result_from_runner("t1", fake) self.assertEqual(len(local["task_costs"]), 1) self.assertEqual(len(local["quota_status"]), 1) self.assertEqual(len(local["task_files"]), 1) docker = _parse_host_result("t1", fake) self.assertEqual(len(docker["task_costs"]), 1) self.assertEqual(len(docker["quota_status"]), 1) self.assertEqual(len(docker["task_files"]), 1) class WorkerCostFilesTestCase(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._insert("PM", status="running", created_by="tester") self._member("tester", "maintainer") def tearDown(self): self.tmp.cleanup() def _insert(self, task_id, status="queued", created_by="tester", project_id="p1", prompt="do work", parent=None, **cols): now = time.time() 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(?, ?, ?, ?, NULL, '[]', 0, ?, ?, ?, NULL)", (task_id, project_id, prompt, status, created_by, now, now)) if parent is not None: conn.execute( "UPDATE tasks SET parent_task_id = ? WHERE id = ?", (parent, task_id)) for key, value in cols.items(): conn.execute( "UPDATE tasks SET %s = ? WHERE id = ?" % key, (value, task_id)) conn.commit() finally: conn.close() def _member(self, user_id, role): conn = sqlite3.connect(self.db_path) try: conn.execute( "CREATE TABLE IF NOT EXISTS users(id TEXT PRIMARY KEY," " username TEXT UNIQUE, api_key_hash TEXT," " created_at REAL)") conn.execute( "CREATE TABLE IF NOT EXISTS project_members(" "project_id TEXT, user_id TEXT, role TEXT," " PRIMARY KEY(project_id, user_id))") conn.execute( "INSERT OR REPLACE INTO project_members" "(project_id, user_id, role) VALUES('p1', ?, ?)", (user_id, role)) conn.execute( "INSERT OR IGNORE INTO users(id, username, api_key_hash," " created_at) VALUES(?, ?, '', 0)", (user_id, user_id)) conn.commit() finally: conn.close() def _worker(self): return Worker(self.db_path, "/fake/host.py", env=FAKE_ENV) def _followup_bodies(self, task_id="PM"): conn = sqlite3.connect(self.db_path) try: rows = conn.execute( "SELECT body FROM task_followups" " WHERE task_id = ?", (task_id,)).fetchall() return [row[0] for row in rows] finally: conn.close() def test_deliver_cost_and_quota(self): self._insert("E1", status="completed", parent="PM") store = CostStore(self.db_path) try: store.record("E1", "p1", "gpt-test", 100, 50, usd_micros=12345, source="actual") finally: try: store._conn.close() except Exception: pass worker = self._worker() worker._deliver_marker_tools(worker.get_task("PM"), { "status": "completed", "task_costs": [{"task_id": "E1"}], "quota_status": [{}], "task_files": []}) bodies = self._followup_bodies() self.assertEqual(len(bodies), 1) self.assertIn("$0.0123", bodies[0]) self.assertIn("配额", bodies[0]) def test_deliver_cost_other_creator_rejected(self): self._insert("E1", status="completed", created_by="someone-else") worker = self._worker() worker._deliver_marker_tools(worker.get_task("PM"), { "status": "completed", "task_costs": [{"task_id": "E1"}], "quota_status": [], "task_files": []}) bodies = self._followup_bodies() self.assertIn("不在派生链内且创建人不同", bodies[0]) def test_deliver_files_list_and_read(self): self._insert("E1", status="completed", parent="PM") worker = self._worker() # 冒充 worktree：租户路径下给 E1 建目录与文件。 worktree = os.path.join(self.tmp.name, "tenants", "p1", "work", "host-worktrees", "E1") os.makedirs(worktree) with open(os.path.join(worktree, "out.txt"), "w") as fh: fh.write("hello-files") with open(os.path.join(worktree, "bin.dat"), "wb") as fh: fh.write(b"\x00\x01\x02") worker.tenants_root = os.path.join(self.tmp.name, "tenants") worker._deliver_marker_tools(worker.get_task("PM"), { "status": "completed", "task_costs": [], "quota_status": [], "task_files": [{"task_id": "E1", "path": ""}, {"task_id": "E1", "path": "out.txt"}, {"task_id": "E1", "path": "bin.dat"}, {"task_id": "E1", "path": "../escape"}]}) bodies = self._followup_bodies() self.assertEqual(len(bodies), 1) self.assertIn("out.txt", bodies[0]) self.assertIn("hello-files", bodies[0]) self.assertIn("二进制文件", bodies[0]) self.assertIn("路径逃逸", bodies[0]) def test_deliver_files_cloud_and_no_worktree_rejected(self): self._insert("C1", status="completed", parent="PM", executor="cloud") self._insert("N1", status="completed", parent="PM") worker = self._worker() worker.tenants_root = os.path.join(self.tmp.name, "tenants") worker._deliver_marker_tools(worker.get_task("PM"), { "status": "completed", "task_costs": [], "quota_status": [], "task_files": [{"task_id": "C1", "path": ""}, {"task_id": "N1", "path": ""}]}) bodies = self._followup_bodies() self.assertIn("不在本机", bodies[0]) self.assertIn("没有本地 worktree", bodies[0]) def test_deliver_goal_sender_ignored(self): self._insert("G1", status="running", goal_mode=1) worker = self._worker() worker._deliver_marker_tools(worker.get_task("G1"), { "status": "completed", "task_costs": [{"task_id": "G1"}], "quota_status": [{}], "task_files": []}) self.assertEqual(self._followup_bodies("G1"), []) if __name__ == "__main__": unittest.main()