"""任务内标记工具单测：解析、门禁、快照格式化与 worker 投递。 不跑真实模型、不联网：解析与门禁是纯函数；worker 投递用 init_db 建的临时 SQLite 库 + 最小任务行，直接调 _deliver_marker_tools，不走 run_one 全量流程。 """ import json import os import sqlite3 import sys import tempfile import time import unittest from pathlib import Path from types import SimpleNamespace try: import cloud.marker_tools as mt from cloud.marker_tools import ( ancestry_relation, check_control_gate, extract_task_cancels, extract_task_creates, extract_task_polls, extract_task_retries, extract_wait_threads, format_snapshot, ) 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])) import cloud.marker_tools as mt from cloud.marker_tools import ( ancestry_relation, check_control_gate, extract_task_cancels, extract_task_creates, extract_task_polls, extract_task_retries, extract_wait_threads, format_snapshot, ) 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"} def _mem_conn(): conn = sqlite3.connect(":memory:") conn.row_factory = sqlite3.Row conn.execute( "CREATE TABLE tasks(id TEXT PRIMARY KEY, project_id TEXT," " created_by TEXT, status TEXT, parent_task_id TEXT," " coordinator_part TEXT, verify_command TEXT," " goal_mode INTEGER, result_text TEXT, error_text TEXT," " executor TEXT, model TEXT, actual_model TEXT," " turn_id TEXT, last_heartbeat REAL, worktree_host TEXT," " pr_url TEXT, pr_number INTEGER, environment_id TEXT)") return conn def _mem_task(conn, task_id, project_id="p1", status="queued", created_by="tester", parent=None, part=None, verify_command=None, goal_mode=None, result=None): conn.execute( "INSERT INTO tasks(id, project_id, created_by, status," " parent_task_id, coordinator_part, verify_command, goal_mode," " result_text, error_text, executor, model, actual_model," " turn_id, last_heartbeat, worktree_host, pr_url, pr_number," " environment_id) VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, NULL," " NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL)", (task_id, project_id, created_by, status, parent, part, verify_command, goal_mode, result)) conn.commit() class ExtractTestCase(unittest.TestCase): def test_create_basic_and_caps(self): blocks = "".join( "```task-create\n{\"name\": \"n%d\", \"prompt\": \"do %d\"}\n```\n" % (i, i) for i in range(7)) creates = extract_task_creates(blocks) self.assertEqual(len(creates), 5) self.assertEqual(creates[0]["name"], "n0") self.assertEqual(creates[0]["prompt"], "do 0") def test_create_invalid_blocks_skipped(self): text = ( "```task-create\nnot json\n```\n" "```task-create\n{\"prompt\": \" \"}\n```\n" "```task-create\n{\"prompt\": \"" + "x" * 20001 + "\"}\n```\n" "```task-create\n{\"prompt\": \"ok\", \"effort\": \"huge\"}\n```\n" "```task-create\n{\"prompt\": \"ok\", \"executor\": \"mars\"}\n```\n" "```task-create\n{\"prompt\": \"good\"}\n```\n") creates = extract_task_creates(text) self.assertEqual(len(creates), 1) self.assertEqual(creates[0]["prompt"], "good") def test_create_fence_at_line_start_truncates_to_invalid(self): # 行首的围栏会把块截断成非法 JSON 而被跳过（非贪婪匹配）； # 同行内的围栏不构成块边界，能正常解析（见下一断言的对照）。 text = ("```task-create\n{\"prompt\": \"a\n```\n b\"}\n```\n" "```task-create\n{\"prompt\": \"good\"}\n```\n") creates = extract_task_creates(text) self.assertEqual([c["prompt"] for c in creates], ["good"]) inline = extract_task_creates( "```task-create\n{\"prompt\": \"a ``` b\"}\n```\n") self.assertEqual([c["prompt"] for c in inline], ["a ``` b"]) def test_poll_clamps_and_read_alias(self): text = ( "```task-poll\n{\"task_id\": \"T1\", \"interval_seconds\": 999," " \"timeout_seconds\": 99999, \"chars\": 1}\n```\n" "```task-read\n{\"task_id\": \"T2\"}\n```\n" "```task-read\n{\"task_id\": \"T3\", \"wait_seconds\": 30}\n```\n") polls = extract_task_polls(text) self.assertEqual(len(polls), 3) self.assertEqual(polls[0]["interval_seconds"], 60) self.assertEqual(polls[0]["timeout_seconds"], 3600) self.assertEqual(polls[0]["chars"], 200) self.assertEqual(polls[1]["timeout_seconds"], 0) self.assertEqual(polls[2]["timeout_seconds"], 30) self.assertEqual(polls[2]["interval_seconds"], 30) def test_poll_max_five_and_missing_id_skipped(self): blocks = "".join( "```task-poll\n{\"task_id\": \"T%d\"}\n```\n" % i for i in range(6)) blocks += "```task-poll\n{\"interval_seconds\": 5}\n```\n" self.assertEqual(len(extract_task_polls(blocks)), 5) def test_wait_cap_eight_not_ten_and_max_two_blocks(self): ids = ["T%d" % i for i in range(12)] text = ( "```wait-threads\n%s\n```\n" "```wait_threads\n{\"task_ids\": [\"A\"]}\n```\n" "```wait-threads\n{\"task_ids\": [\"B\"]}\n```\n" % json.dumps({"task_ids": ids})) waits = extract_wait_threads(text) # 每块最多 8 个（与服务端 /wait 对齐，不是 spawn 原先的 10 个）， # 每轮最多 2 块。 self.assertEqual(len(waits), 2) self.assertEqual(len(waits[0]["task_ids"]), 8) self.assertEqual(waits[0]["task_ids"][0], "T0") self.assertEqual(waits[1]["task_ids"], ["A"]) def test_wait_invalid_blocks_skipped(self): text = ( "```wait-threads\nnot json\n```\n" "```wait-threads\n{\"task_ids\": []}\n```\n" "```wait-threads\n{\"task_ids\": \"not-a-list\"}\n```\n" "```wait-threads\n{\"task_ids\": [\"A\", \"A\", \" \", \"B\"]}\n```\n") waits = extract_wait_threads(text) self.assertEqual(len(waits), 1) self.assertEqual(waits[0]["task_ids"], ["A", "B"]) def test_control_extract_caps_and_aliases(self): cancels = "".join( "```task-cancel\n{\"task_id\": \"C%d\"}\n```\n" % i for i in range(6)) self.assertEqual(len(extract_task_cancels(cancels)), 5) retries = extract_task_retries( "```task_retry\n{\"task_id\": \"R1\"}\n```\n" "```task-retry\n{\"to_task_id\": \"R2\"}\n```\n" "```task-retry\n{\"no_id\": true}\n```\n") self.assertEqual([r["task_id"] for r in retries], ["R1", "R2"]) class GateTestCase(unittest.TestCase): def _sender(self, task_id="PM", project_id="p1", created_by="tester"): return {"id": task_id, "project_id": project_id, "created_by": created_by, "verify_command": None, "goal_mode": None} def _target(self, task_id="E1", project_id="p1", status="running", created_by="tester", verify_command=None, goal_mode=None): return {"id": task_id, "project_id": project_id, "status": status, "created_by": created_by, "verify_command": verify_command, "goal_mode": goal_mode} def test_cancel_terminal_rejected(self): for status in ("completed", "failed", "interrupted", "merged", "merge-failed"): ok, _ = check_control_gate( self._target(status=status), self._sender(), "descendant", "cancel") self.assertFalse(ok, status) def test_retry_only_three_states(self): for status in ("failed", "interrupted", "merge-failed"): ok, _ = check_control_gate( self._target(status=status), self._sender(), "descendant", "retry") self.assertTrue(ok, status) for status in ("completed", "merged", "queued", "running"): ok, _ = check_control_gate( self._target(status=status), self._sender(), "descendant", "retry") self.assertFalse(ok, status) def test_cross_project_self_goal_merged_rejected(self): sender = self._sender() ok, _ = check_control_gate( self._target(project_id="p2"), sender, "descendant", "cancel") self.assertFalse(ok) ok, _ = check_control_gate( self._target(task_id="PM"), sender, "descendant", "cancel") self.assertFalse(ok) ok, _ = check_control_gate( self._target(verify_command="make test"), sender, "descendant", "cancel") self.assertFalse(ok) ok, _ = check_control_gate( self._target(goal_mode=1), sender, "descendant", "retry") self.assertFalse(ok) ok, _ = check_control_gate( self._target(status="merged"), sender, "descendant", "cancel") self.assertFalse(ok) def test_relation_other_rejected_same_creator_passes(self): ok, _ = check_control_gate( self._target(created_by="other"), self._sender(), "other", "cancel") self.assertFalse(ok) ok, _ = check_control_gate( self._target(created_by="tester"), self._sender(), "same_creator", "cancel") self.assertTrue(ok) def test_ancestry_chain(self): conn = _mem_conn() _mem_task(conn, "PM") _mem_task(conn, "E1", parent="PM") _mem_task(conn, "E2", parent="E1") _mem_task(conn, "X1", created_by="other") sender = {"id": "PM", "project_id": "p1", "created_by": "tester"} target_e2 = mt.fetch_task_snapshot(conn, "E2") self.assertEqual(ancestry_relation(conn, sender, target_e2), "descendant") target_x = mt.fetch_task_snapshot(conn, "X1") self.assertEqual(ancestry_relation(conn, sender, target_x), "other") target_e1 = mt.fetch_task_snapshot(conn, "E1") self.assertEqual(ancestry_relation(conn, sender, target_e1), "descendant") class SnapshotTestCase(unittest.TestCase): def test_format_terminal_truncates(self): snap = {"id": "T1", "status": "completed", "result_text": "r" * 5000, "error_text": "", "executor": "docker", "model": "m", "actual_model": None, "turn_id": None, "last_heartbeat": None, "worktree_host": None} text = format_snapshot(snap, 200) self.assertIn("状态：completed", text) self.assertIn("r" * 200, text) self.assertNotIn("r" * 201, text) def test_format_nonterminal_timeout_cancelled(self): base = {"id": "T1", "status": "running", "result_text": "", "error_text": "", "executor": None, "model": None, "actual_model": None, "turn_id": None, "last_heartbeat": None, "worktree_host": None} self.assertIn("仍在运行或排队", format_snapshot(base, 200)) self.assertIn("等待超时", format_snapshot(base, 200, timed_out=True)) self.assertIn("父任务被取消", format_snapshot(base, 200, cancelled=True)) class SandboxExtractTestCase(unittest.TestCase): def _text(self, *blocks): return "\n".join(blocks) def test_local_and_docker_extract_all_four(self): text = self._text( "```task-create\n{\"prompt\": \"do x\"}\n```", "```task-poll\n{\"task_id\": \"T1\"}\n```", "```wait-threads\n{\"task_ids\": [\"T1\", \"T2\"]}\n```", "```task-cancel\n{\"task_id\": \"T2\"}\n```", "```task-retry\n{\"task_id\": \"T3\"}\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_creates"]), 1) self.assertEqual(len(local["task_polls"]), 1) self.assertEqual(len(local["wait_threads"]), 1) self.assertEqual(len(local["task_cancels"]), 1) self.assertEqual(len(local["task_retries"]), 1) docker = _parse_host_result("t1", fake) self.assertEqual(len(docker["task_creates"]), 1) self.assertEqual(len(docker["wait_threads"]), 1) self.assertEqual(len(docker["task_retries"]), 1) class WorkerDeliverTestCase(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", 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, "do work", 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 _sender(self, worker): return worker.get_task("PM") def _followup_bodies(self): conn = sqlite3.connect(self.db_path) try: rows = conn.execute( "SELECT body FROM task_followups" " WHERE task_id = 'PM'").fetchall() return [row[0] for row in rows] finally: conn.close() def _status(self, task_id): conn = sqlite3.connect(self.db_path) try: row = conn.execute( "SELECT status, result_text FROM tasks WHERE id = ?", (task_id,)).fetchone() return row[0], row[1] finally: conn.close() def test_deliver_create_and_poll(self): worker = self._worker() sender = self._sender(worker) worker._deliver_marker_tools(sender, { "status": "completed", "task_creates": [{"name": "E1", "prompt": "do e1", "model": None, "effort": None, "executor": None}], "task_polls": [], "wait_threads": [], "task_cancels": [], "task_retries": []}) conn = sqlite3.connect(self.db_path) try: row = conn.execute( "SELECT id, parent_task_id, status, created_by FROM tasks" " WHERE parent_task_id = 'PM'").fetchone() finally: conn.close() self.assertIsNotNone(row) self.assertEqual(row[1], "PM") self.assertEqual(row[2], "queued") self.assertEqual(row[3], "tester") def test_deliver_poll_reports_snapshot(self): self._insert("E1", status="completed", parent="PM", result_text="done-thing") worker = self._worker() worker._deliver_marker_tools(self._sender(worker), { "status": "completed", "task_creates": [], "task_polls": [{"task_id": "E1", "interval_seconds": 5, "timeout_seconds": 0, "chars": 2000}], "wait_threads": [], "task_cancels": [], "task_retries": []}) bodies = self._followup_bodies() self.assertEqual(len(bodies), 1) self.assertIn("done-thing", bodies[0]) def test_deliver_wait_reports_all_targets(self): self._insert("E1", status="completed", parent="PM", result_text="first-done") self._insert("E2", status="queued", parent="PM") worker = self._worker() worker._deliver_marker_tools(self._sender(worker), { "status": "completed", "task_creates": [], "task_polls": [], "wait_threads": [{"task_ids": ["E1", "E2"], "interval_seconds": 5, "timeout_seconds": 0, "chars": 2000}], "task_cancels": [], "task_retries": []}) bodies = self._followup_bodies() self.assertEqual(len(bodies), 1) self.assertIn("E1", bodies[0]) self.assertIn("E2", bodies[0]) self.assertIn("first-done", bodies[0]) def test_deliver_cancel_descendant(self): self._insert("E1", status="running", parent="PM") worker = self._worker() worker._deliver_marker_tools(self._sender(worker), { "status": "completed", "task_creates": [], "task_polls": [], "wait_threads": [], "task_cancels": [{"task_id": "E1"}], "task_retries": []}) self.assertEqual(self._status("E1")[0], "interrupted") def test_deliver_cancel_cross_project_and_terminal_rejected(self): self._insert("X1", status="running", project_id="p2") self._insert("E1", status="completed", parent="PM") worker = self._worker() worker._deliver_marker_tools(self._sender(worker), { "status": "completed", "task_creates": [], "task_polls": [], "wait_threads": [], "task_cancels": [{"task_id": "X1"}, {"task_id": "E1"}], "task_retries": []}) self.assertEqual(self._status("X1")[0], "running") self.assertEqual(self._status("E1")[0], "completed") bodies = self._followup_bodies() self.assertTrue(any("拒绝" in body for body in bodies)) def test_deliver_retry_clears_old_result(self): self._insert("E1", status="failed", parent="PM", result_text="old-result", error_text="old-error") conn = sqlite3.connect(self.db_path) try: conn.execute( "UPDATE tasks SET actual_model = 'm' WHERE id = 'E1'") conn.commit() finally: conn.close() worker = self._worker() worker._deliver_marker_tools(self._sender(worker), { "status": "completed", "task_creates": [], "task_polls": [], "wait_threads": [], "task_cancels": [], "task_retries": [{"task_id": "E1"}]}) status, result = self._status("E1") self.assertEqual(status, "queued") self.assertIsNone(result) def test_deliver_retry_denied_without_approve_merge(self): self._member("tester", "viewer") self._insert("E1", status="failed", parent="PM", result_text="old-result") worker = self._worker() worker._deliver_marker_tools(self._sender(worker), { "status": "completed", "task_creates": [], "task_polls": [], "wait_threads": [], "task_cancels": [], "task_retries": [{"task_id": "E1"}]}) self.assertEqual(self._status("E1")[0], "failed") bodies = self._followup_bodies() self.assertTrue(any("拒绝" in body for body in bodies)) def test_deliver_goal_and_coordinator_sender_ignored(self): self._insert("G1", status="running", verify_command="make test") worker = self._worker() goal_sender = worker.get_task("G1") worker._deliver_marker_tools(goal_sender, { "status": "completed", "task_creates": [{"name": "E1", "prompt": "do e1", "model": None, "effort": None, "executor": None}], "task_polls": [], "wait_threads": [], "task_cancels": [], "task_retries": []}) conn = sqlite3.connect(self.db_path) try: count = conn.execute( "SELECT COUNT(*) FROM tasks" " WHERE parent_task_id = 'G1'").fetchone()[0] finally: conn.close() self.assertEqual(count, 0) def test_deliver_create_quota_exceeded_skips(self): 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('tester', 0, 0)") conn.commit() finally: conn.close() worker = self._worker() worker._deliver_marker_tools(self._sender(worker), { "status": "completed", "task_creates": [{"name": "E1", "prompt": "do e1", "model": None, "effort": None, "executor": None}], "task_polls": [], "wait_threads": [], "task_cancels": [], "task_retries": []}) conn = sqlite3.connect(self.db_path) try: count = conn.execute( "SELECT COUNT(*) FROM tasks" " WHERE parent_task_id = 'PM'").fetchone()[0] finally: conn.close() self.assertEqual(count, 0) bodies = self._followup_bodies() self.assertTrue(any("配额" in body for body in bodies)) def test_deliver_create_direct_limit_rejected(self): for i in range(20): self._insert("OLD%d" % i, status="completed", parent="PM") worker = self._worker() worker._deliver_marker_tools(self._sender(worker), { "status": "completed", "task_creates": [{"name": "E1", "prompt": "do e1", "model": None, "effort": None, "executor": None}], "task_polls": [], "wait_threads": [], "task_cancels": [], "task_retries": []}) conn = sqlite3.connect(self.db_path) try: count = conn.execute( "SELECT COUNT(*) FROM tasks" " WHERE parent_task_id = 'PM'").fetchone()[0] finally: conn.close() self.assertEqual(count, 20) bodies = self._followup_bodies() self.assertTrue(any("上限" in body for body in bodies)) def test_deliver_create_depth_limit_rejected(self): self._insert("A", status="completed", parent="PM") self._insert("B", status="completed", parent="A") self._insert("C", status="running", parent="B") worker = self._worker() sender_c = worker.get_task("C") worker._deliver_marker_tools(sender_c, { "status": "completed", "task_creates": [{"name": "E1", "prompt": "do e1", "model": None, "effort": None, "executor": None}], "task_polls": [], "wait_threads": [], "task_cancels": [], "task_retries": []}) conn = sqlite3.connect(self.db_path) try: count = conn.execute( "SELECT COUNT(*) FROM tasks" " WHERE parent_task_id = 'C'").fetchone()[0] finally: conn.close() self.assertEqual(count, 0) def test_deliver_wait_timeout_reports_pending(self): self._insert("E1", status="queued", parent="PM") worker = self._worker() worker._deliver_marker_tools(self._sender(worker), { "status": "completed", "task_creates": [], "task_polls": [], "wait_threads": [{"task_ids": ["E1"], "interval_seconds": 1, "timeout_seconds": 1, "chars": 2000}], "task_cancels": [], "task_retries": []}) bodies = self._followup_bodies() self.assertEqual(len(bodies), 1) self.assertIn("E1", bodies[0]) self.assertIn("等待超时", bodies[0]) def test_deliver_control_auth_failure_falls_back_to_descendant(self): from unittest import mock self._insert("E1", status="running", parent="PM") self._insert("S1", status="running") worker = self._worker() with mock.patch("cloud.worker.AuthStore", side_effect=RuntimeError("auth down")): worker._deliver_marker_tools(self._sender(worker), { "status": "completed", "task_creates": [], "task_polls": [], "wait_threads": [], "task_cancels": [{"task_id": "E1"}, {"task_id": "S1"}], "task_retries": []}) # 权限查不到时：派生链内仍可取消，同创建人-only 被拒绝。 self.assertEqual(self._status("E1")[0], "interrupted") self.assertEqual(self._status("S1")[0], "running") def test_deliver_non_completed_round_ignored(self): worker = self._worker() worker._deliver_marker_tools(self._sender(worker), { "status": "failed", "task_creates": [{"name": "E1", "prompt": "do e1", "model": None, "effort": None, "executor": None}], "task_polls": [], "wait_threads": [], "task_cancels": [], "task_retries": []}) self.assertEqual(self._followup_bodies(), []) if __name__ == "__main__": unittest.main()