feat: 任务间互发消息（API + agent 标记块，不中断任务）
同项目任务互发消息、对方做完可回消息：API POST /tasks/<目标>/messages （from_task_id，目标项目 create_task + 发送项目 read 鉴权）与 agent 回合输出标记块（worker 回合后解析投递，不用联网、不用凭证，本地 与 Docker 都通）。投递不中断任务：对方运行中/排队中排成后续消息 （轮次边界执行，direct 优先）；对方终态则续跑复活（resume_message + queued），A→B→A 往返不受回合时序影响。拒绝项：自己、跨项目、goal 目标、已合并、复活位被占；每轮最多 5 块，仅 completed 轮次投递； 目标执行时配额门照常生效。控制台后续消息列显示发送方，新增给 其他任务发消息框（落已代理的 /tasks 下，不用改 nginx）。 测试：新增 test_peer_messages（15）、test_peer_messages_server（8）、 test_worker_peer（4）；全量 cloud 654、host 64、deploy 19 全过（cloud 4 跳过）。
This commit is contained in:
host committed
1 parent
4901a1090b
commit
ef30558f7c
11 files changed
+1171
-23
No files matched your search
@@ -64,7 +64,8 @@ class ConsoleTestCase(unittest.TestCase):
|
||||
self.assertIn("Codex 命令台", text)
|
||||
# 壳必须引用它实际调的接口，防 JS 与服务端路由漂移。
|
||||
for marker in ("/health", "/projects", "/tasks", "/workers",
|
||||
"X-API-Key", "agentBox", "/agents"):
|
||||
"X-API-Key", "agentBox", "/agents",
|
||||
"peerMessage", "/messages"):
|
||||
self.assertIn(marker, text)
|
||||
|
||||
def test_console_trailing_slash_same_page(self):
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
"""任务间互发消息单测：标记块提取 + 投递分支。
|
||||
|
||||
不跑真实模型、不联网：extract 是纯函数；deliver 用内存
|
||||
SQLite + 最小表结构（tasks 只含投递读到的列，followups
|
||||
表由 peer_messages 自己建）。
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from cloud.peer_messages import (
|
||||
deliver_peer_message,
|
||||
extract_peer_messages,
|
||||
peer_message_prefix,
|
||||
)
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.peer_messages import (
|
||||
deliver_peer_message,
|
||||
extract_peer_messages,
|
||||
peer_message_prefix,
|
||||
)
|
||||
|
||||
|
||||
def _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, verify_command TEXT,"
|
||||
" goal_mode INTEGER, resume_message TEXT, updated_at REAL)")
|
||||
return conn
|
||||
|
||||
|
||||
def _task(conn, task_id, project_id="p1", status="queued",
|
||||
created_by="tester", verify_command=None, goal_mode=None,
|
||||
resume_message=None):
|
||||
conn.execute(
|
||||
"INSERT INTO tasks(id, project_id, created_by, status,"
|
||||
" verify_command, goal_mode, resume_message, updated_at)"
|
||||
" VALUES(?, ?, ?, ?, ?, ?, ?, 0)",
|
||||
(task_id, project_id, created_by, status, verify_command,
|
||||
goal_mode, resume_message))
|
||||
conn.commit()
|
||||
|
||||
|
||||
def _followups(conn, task_id=None, sender=None):
|
||||
sql = "SELECT * FROM task_followups"
|
||||
conds = []
|
||||
args = []
|
||||
if task_id is not None:
|
||||
conds.append("task_id = ?")
|
||||
args.append(task_id)
|
||||
if sender is not None:
|
||||
conds.append("sender_task_id = ?")
|
||||
args.append(sender)
|
||||
if conds:
|
||||
sql += " WHERE " + " AND ".join(conds)
|
||||
try:
|
||||
return conn.execute(sql, args).fetchall()
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def _task_row(conn, task_id):
|
||||
return conn.execute(
|
||||
"SELECT * FROM tasks WHERE id = ?", (task_id,)).fetchone()
|
||||
|
||||
|
||||
class ExtractTestCase(unittest.TestCase):
|
||||
def test_single_block_default_queued(self):
|
||||
text = ("做完了。\n```task-message\n"
|
||||
"{\"to\": \"T2\", \"body\": \"请评审\"}\n```\n")
|
||||
msgs = extract_peer_messages(text)
|
||||
self.assertEqual(len(msgs), 1)
|
||||
self.assertEqual(msgs[0]["to_task_id"], "T2")
|
||||
self.assertEqual(msgs[0]["body"], "请评审")
|
||||
self.assertEqual(msgs[0]["mode"], "queued")
|
||||
|
||||
def test_direct_mode_and_alias_key(self):
|
||||
text = ("```task-message\n{\"to_task_id\": \"T2\","
|
||||
" \"body\": \"急\", \"mode\": \"direct\"}\n```")
|
||||
msgs = extract_peer_messages(text)
|
||||
self.assertEqual(msgs[0]["mode"], "direct")
|
||||
|
||||
def test_bad_mode_falls_back_to_queued(self):
|
||||
text = ("```task-message\n{\"to\": \"T2\", \"body\": \"x\","
|
||||
" \"mode\": \"smoke-signal\"}\n```")
|
||||
self.assertEqual(extract_peer_messages(text)[0]["mode"], "queued")
|
||||
|
||||
def test_invalid_blocks_skipped(self):
|
||||
text = (
|
||||
"```task-message\nnot json\n```\n"
|
||||
"```task-message\n{\"body\": \"no to\"}\n```\n"
|
||||
"```task-message\n{\"to\": \"T2\"}\n```\n"
|
||||
"```task-message\n{\"to\": \"T2\", \"body\": \" \"}\n```\n"
|
||||
"```task-message\n{\"to\": \"T2\", \"body\": \"ok\"}\n```\n")
|
||||
msgs = extract_peer_messages(text)
|
||||
self.assertEqual(len(msgs), 1)
|
||||
self.assertEqual(msgs[0]["body"], "ok")
|
||||
|
||||
def test_body_too_long_skipped(self):
|
||||
text = ("```task-message\n{\"to\": \"T2\", \"body\": \""
|
||||
+ "x" * 20001 + "\"}\n```")
|
||||
self.assertEqual(extract_peer_messages(text), [])
|
||||
|
||||
def test_max_five_per_turn(self):
|
||||
blocks = "".join(
|
||||
"```task-message\n{\"to\": \"T%d\", \"body\": \"m%d\"}\n```\n"
|
||||
% (i, i) for i in range(7))
|
||||
self.assertEqual(len(extract_peer_messages(blocks)), 5)
|
||||
|
||||
def test_no_fence_or_empty(self):
|
||||
self.assertEqual(extract_peer_messages("plain text"), [])
|
||||
self.assertEqual(extract_peer_messages(""), [])
|
||||
self.assertEqual(extract_peer_messages(None), [])
|
||||
|
||||
|
||||
class DeliverTestCase(unittest.TestCase):
|
||||
def test_running_target_becomes_followup(self):
|
||||
conn = _conn()
|
||||
_task(conn, "T1")
|
||||
_task(conn, "T2", status="running")
|
||||
status, detail = deliver_peer_message(conn, "T1", "T2", "请评审")
|
||||
self.assertEqual(status, "delivered")
|
||||
rows = _followups(conn, task_id="T2", sender="T1")
|
||||
self.assertEqual(len(rows), 1)
|
||||
self.assertEqual(rows[0]["body"], "请评审")
|
||||
self.assertEqual(rows[0]["mode"], "queued")
|
||||
self.assertEqual(rows[0]["status"], "pending")
|
||||
self.assertEqual(rows[0]["created_by"], "tester")
|
||||
self.assertEqual(rows[0]["id"], detail)
|
||||
|
||||
def test_queued_target_direct_mode(self):
|
||||
conn = _conn()
|
||||
_task(conn, "T1")
|
||||
_task(conn, "T2", status="queued")
|
||||
status, _ = deliver_peer_message(conn, "T1", "T2", "急", "direct")
|
||||
self.assertEqual(status, "delivered")
|
||||
self.assertEqual(_followups(conn, task_id="T2")[0]["mode"], "direct")
|
||||
|
||||
def test_completed_target_revives(self):
|
||||
conn = _conn()
|
||||
_task(conn, "T1")
|
||||
_task(conn, "T2", status="completed")
|
||||
status, detail = deliver_peer_message(conn, "T1", "T2", "回来一下")
|
||||
self.assertEqual(status, "revived")
|
||||
self.assertEqual(detail, "")
|
||||
row = _task_row(conn, "T2")
|
||||
self.assertEqual(row["status"], "queued")
|
||||
self.assertEqual(row["resume_message"],
|
||||
peer_message_prefix("T1") + "回来一下")
|
||||
self.assertEqual(_followups(conn, task_id="T2"), [])
|
||||
|
||||
def test_revive_blocked_when_resume_occupied(self):
|
||||
conn = _conn()
|
||||
_task(conn, "T1")
|
||||
_task(conn, "T2", status="completed", resume_message="旧消息")
|
||||
status, detail = deliver_peer_message(conn, "T1", "T2", "新消息")
|
||||
self.assertEqual(status, "rejected")
|
||||
self.assertIn("pending resume", detail)
|
||||
self.assertEqual(_task_row(conn, "T2")["resume_message"], "旧消息")
|
||||
|
||||
def test_failed_and_interrupted_revive(self):
|
||||
for st in ("failed", "interrupted"):
|
||||
conn = _conn()
|
||||
_task(conn, "T1")
|
||||
_task(conn, "T2", status=st)
|
||||
status, _ = deliver_peer_message(conn, "T1", "T2", "回来")
|
||||
self.assertEqual(status, "revived")
|
||||
self.assertEqual(_task_row(conn, "T2")["status"], "queued")
|
||||
|
||||
def test_goal_targets_rejected(self):
|
||||
conn = _conn()
|
||||
_task(conn, "T1")
|
||||
_task(conn, "T2", status="running", verify_command="true")
|
||||
status, _ = deliver_peer_message(conn, "T1", "T2", "x")
|
||||
self.assertEqual(status, "rejected")
|
||||
conn = _conn()
|
||||
_task(conn, "T1")
|
||||
_task(conn, "T2", status="running", goal_mode=1)
|
||||
status, _ = deliver_peer_message(conn, "T1", "T2", "x")
|
||||
self.assertEqual(status, "rejected")
|
||||
|
||||
def test_self_cross_project_merged_missing_rejected(self):
|
||||
conn = _conn()
|
||||
_task(conn, "T1")
|
||||
_task(conn, "T2", project_id="p2")
|
||||
_task(conn, "T3", status="merged")
|
||||
self.assertEqual(
|
||||
deliver_peer_message(conn, "T1", "T1", "x")[0], "rejected")
|
||||
self.assertEqual(
|
||||
deliver_peer_message(conn, "T1", "T2", "x")[0], "rejected")
|
||||
self.assertEqual(
|
||||
deliver_peer_message(conn, "T1", "T3", "x")[0], "rejected")
|
||||
self.assertEqual(
|
||||
deliver_peer_message(conn, "T1", "NOPE", "x")[0], "rejected")
|
||||
self.assertEqual(
|
||||
deliver_peer_message(conn, "NOPE", "T1", "x")[0], "rejected")
|
||||
# 拒绝不写任何东西、不改状态。
|
||||
self.assertEqual(_followups(conn), [])
|
||||
self.assertEqual(_task_row(conn, "T3")["status"], "merged")
|
||||
|
||||
def test_body_validation(self):
|
||||
conn = _conn()
|
||||
_task(conn, "T1")
|
||||
_task(conn, "T2", status="running")
|
||||
self.assertEqual(
|
||||
deliver_peer_message(conn, "T1", "T2", " ")[0], "rejected")
|
||||
self.assertEqual(
|
||||
deliver_peer_message(conn, "T1", "T2", "x" * 20001)[0],
|
||||
"rejected")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,248 @@
|
||||
"""任务间互发消息服务端单测：POST /tasks/<id>/messages 与
|
||||
GET /tasks/<id>/messages。
|
||||
|
||||
不跑真实模型、不联网：本地临时端口 + 临时 SQLite，与
|
||||
test_task_links.py 同款骨架。覆盖鉴权、权限、校验、三种
|
||||
投递结果（排队 / 复活 / 拒绝）与发出列表。
|
||||
"""
|
||||
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.server import create_server
|
||||
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.server import create_server
|
||||
|
||||
|
||||
class PeerMessageServerTestCase(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.outsider_id, self.outsider_key = self.auth.create_user("outsider")
|
||||
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()
|
||||
self.trace_dir = os.path.join(self.tmp.name, "host-traces")
|
||||
os.makedirs(self.trace_dir, exist_ok=True)
|
||||
self.server.trace_dir = self.trace_dir
|
||||
status, self.project = self._req(
|
||||
"POST", "/projects", {"name": "demo"}, self.owner_key)
|
||||
self.assertEqual(status, 201)
|
||||
status, self.project2 = self._req(
|
||||
"POST", "/projects", {"name": "other"}, 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):
|
||||
conn = http.client.HTTPConnection(
|
||||
"127.0.0.1", self.port, timeout=10)
|
||||
hdrs = {}
|
||||
payload = None
|
||||
if body is not None:
|
||||
payload = json.dumps(body).encode("utf-8")
|
||||
hdrs["Content-Type"] = "application/json"
|
||||
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 = None
|
||||
return resp.status, parsed
|
||||
|
||||
def _create_task(self, project_id, **extra):
|
||||
body = {"project_id": project_id, "prompt": "做事"}
|
||||
body.update(extra)
|
||||
status, task = self._req("POST", "/tasks", body, self.owner_key)
|
||||
self.assertEqual(status, 201, task)
|
||||
return task
|
||||
|
||||
def _set_status(self, task_id, status):
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
conn.execute("UPDATE tasks SET status = ? WHERE id = ?",
|
||||
(status, task_id))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _task_row(self, task_id):
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
return conn.execute(
|
||||
"SELECT status, resume_message FROM tasks WHERE id = ?",
|
||||
(task_id,)).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _followups(self, task_id):
|
||||
_status, body = self._req(
|
||||
"GET", "/tasks/%s/followups" % task_id, None, self.owner_key)
|
||||
return body["followups"]
|
||||
|
||||
def test_delivered_to_running_target(self):
|
||||
sender = self._create_task(self.project["id"])
|
||||
target = self._create_task(self.project["id"])
|
||||
self._set_status(target["id"], "running")
|
||||
status, body = self._req(
|
||||
"POST", "/tasks/%s/messages" % target["id"],
|
||||
{"from_task_id": sender["id"], "body": "请评审"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 201, body)
|
||||
self.assertEqual(body["delivery"], "delivered")
|
||||
self.assertIsNotNone(body["followup_id"])
|
||||
fus = self._followups(target["id"])
|
||||
self.assertEqual(len(fus), 1)
|
||||
self.assertEqual(fus[0]["body"], "请评审")
|
||||
self.assertEqual(fus[0]["sender_task_id"], sender["id"])
|
||||
self.assertEqual(fus[0]["status"], "pending")
|
||||
|
||||
def test_revived_completed_target(self):
|
||||
sender = self._create_task(self.project["id"])
|
||||
target = self._create_task(self.project["id"])
|
||||
self._set_status(target["id"], "completed")
|
||||
status, body = self._req(
|
||||
"POST", "/tasks/%s/messages" % target["id"],
|
||||
{"from_task_id": sender["id"], "body": "回来一下"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 201, body)
|
||||
self.assertEqual(body["delivery"], "revived")
|
||||
self.assertIsNone(body["followup_id"])
|
||||
row = self._task_row(target["id"])
|
||||
self.assertEqual(row[0], "queued")
|
||||
self.assertIn("回来一下", row[1])
|
||||
self.assertIn(sender["id"], row[1])
|
||||
|
||||
def test_direct_mode_writes_trace(self):
|
||||
sender = self._create_task(self.project["id"])
|
||||
target = self._create_task(self.project["id"])
|
||||
self._set_status(target["id"], "running")
|
||||
status, body = self._req(
|
||||
"POST", "/tasks/%s/messages" % target["id"],
|
||||
{"from_task_id": sender["id"], "body": "急事",
|
||||
"mode": "direct"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 201, body)
|
||||
self.assertEqual(body["delivery"], "delivered")
|
||||
names = os.listdir(self.trace_dir)
|
||||
self.assertEqual(len(names), 1)
|
||||
with open(os.path.join(self.trace_dir, names[0]),
|
||||
encoding="utf-8") as f:
|
||||
entry = json.loads(f.read().strip())
|
||||
self.assertEqual(entry["phase"], "followup-direct")
|
||||
self.assertIn("急事", entry["text"])
|
||||
self.assertIn(sender["id"], entry["text"])
|
||||
|
||||
def test_validation_and_not_found(self):
|
||||
sender = self._create_task(self.project["id"])
|
||||
target = self._create_task(self.project["id"])
|
||||
base = "/tasks/%s/messages" % target["id"]
|
||||
status, _ = self._req("POST", base, {"body": "x"}, self.owner_key)
|
||||
self.assertEqual(status, 400)
|
||||
status, _ = self._req(
|
||||
"POST", base, {"from_task_id": sender["id"]}, self.owner_key)
|
||||
self.assertEqual(status, 400)
|
||||
status, _ = self._req(
|
||||
"POST", base,
|
||||
{"from_task_id": sender["id"], "body": "x", "mode": "smoke"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 400)
|
||||
status, _ = self._req(
|
||||
"POST", base, {"from_task_id": "nope", "body": "x"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 404)
|
||||
status, _ = self._req(
|
||||
"POST", "/tasks/nope/messages",
|
||||
{"from_task_id": sender["id"], "body": "x"}, self.owner_key)
|
||||
self.assertEqual(status, 404)
|
||||
status, _ = self._req(
|
||||
"POST", base, {"from_task_id": sender["id"], "body": "x"},
|
||||
None)
|
||||
self.assertEqual(status, 401)
|
||||
|
||||
def test_forbidden_for_outsider(self):
|
||||
sender = self._create_task(self.project["id"])
|
||||
target = self._create_task(self.project["id"])
|
||||
status, _ = self._req(
|
||||
"POST", "/tasks/%s/messages" % target["id"],
|
||||
{"from_task_id": sender["id"], "body": "x"}, self.outsider_key)
|
||||
self.assertEqual(status, 403)
|
||||
|
||||
def test_cross_project_rejected(self):
|
||||
sender = self._create_task(self.project["id"])
|
||||
target = self._create_task(self.project2["id"])
|
||||
self._set_status(target["id"], "running")
|
||||
status, body = self._req(
|
||||
"POST", "/tasks/%s/messages" % target["id"],
|
||||
{"from_task_id": sender["id"], "body": "x"}, self.owner_key)
|
||||
self.assertEqual(status, 409, body)
|
||||
self.assertIn("cross-project", body["error"])
|
||||
|
||||
def test_self_and_goal_rejected(self):
|
||||
task = self._create_task(self.project["id"])
|
||||
self._set_status(task["id"], "running")
|
||||
status, body = self._req(
|
||||
"POST", "/tasks/%s/messages" % task["id"],
|
||||
{"from_task_id": task["id"], "body": "x"}, self.owner_key)
|
||||
self.assertEqual(status, 409)
|
||||
goal = self._create_task(self.project["id"],
|
||||
verify_command="true")
|
||||
self._set_status(goal["id"], "running")
|
||||
sender = self._create_task(self.project["id"])
|
||||
status, _ = self._req(
|
||||
"POST", "/tasks/%s/messages" % goal["id"],
|
||||
{"from_task_id": sender["id"], "body": "x"}, self.owner_key)
|
||||
self.assertEqual(status, 409)
|
||||
|
||||
def test_list_outgoing_messages(self):
|
||||
sender = self._create_task(self.project["id"])
|
||||
target = self._create_task(self.project["id"])
|
||||
self._set_status(target["id"], "running")
|
||||
status, _ = self._req(
|
||||
"POST", "/tasks/%s/messages" % target["id"],
|
||||
{"from_task_id": sender["id"], "body": "请评审"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 201)
|
||||
status, body = self._req(
|
||||
"GET", "/tasks/%s/messages" % sender["id"], None, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(len(body["messages"]), 1)
|
||||
self.assertEqual(body["messages"][0]["task_id"], target["id"])
|
||||
self.assertEqual(body["messages"][0]["sender_task_id"], sender["id"])
|
||||
status, body = self._req(
|
||||
"GET", "/tasks/%s/messages" % sender["id"], None,
|
||||
self.outsider_key)
|
||||
self.assertEqual(status, 403)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,192 @@
|
||||
"""cloud worker 任务间互发消息单测：agent 标记块 → 投递 →
|
||||
接收方排空前缀 / 终态复活。
|
||||
|
||||
不跑真实模型、不联网：注入假 runner（假 CLI），任务与
|
||||
目标都在临时 SQLite 里；env 里的值只是占位用的假名，绝
|
||||
不是真实密钥。
|
||||
"""
|
||||
|
||||
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.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.server import init_db
|
||||
from cloud.worker import Worker
|
||||
|
||||
MARKER_TO_T2 = ("结果如下。\n```task-message\n"
|
||||
"{\"to\": \"T2\", \"body\": \"请评审这份结论\"}\n```\n")
|
||||
PREFIX = "[来自任务 T1 的消息]\n"
|
||||
|
||||
|
||||
def _insert_task(db_path, task_id, created_by="tester", status="queued",
|
||||
project_id="p1"):
|
||||
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),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
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 _resume_message(db_path, task_id):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT resume_message FROM tasks WHERE id = ?",
|
||||
(task_id,)).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
return row[0] if row else None
|
||||
|
||||
|
||||
def _followups(db_path, task_id=None):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
if task_id is None:
|
||||
return conn.execute("SELECT * FROM task_followups").fetchall()
|
||||
return conn.execute(
|
||||
"SELECT * FROM task_followups WHERE task_id = ?",
|
||||
(task_id,)).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _prompt_values(calls):
|
||||
"""假 runner 收到的全部 prompt：dispatch 在 --task 后，
|
||||
run 在 --prompt 后，两个都扫。"""
|
||||
values = []
|
||||
for cmd in calls:
|
||||
for flag in ("--task", "--prompt"):
|
||||
if flag in cmd:
|
||||
values.append(cmd[cmd.index(flag) + 1])
|
||||
return values
|
||||
|
||||
|
||||
def _marker_runner(calls, marker_text):
|
||||
"""第一轮输出含标记块，之后一律 plain done。"""
|
||||
def fake(cmd):
|
||||
calls.append(list(cmd))
|
||||
text = marker_text if len(calls) == 1 else "done"
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout=json.dumps({"status": "completed", "text": text}),
|
||||
stderr="")
|
||||
return fake
|
||||
|
||||
|
||||
class WorkerPeerTestCase(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.env = {"XCM_API_KEY": "test-placeholder"}
|
||||
self.calls = []
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _worker(self, runner):
|
||||
return Worker(self.db_path, "/fake/host.py", runner=runner,
|
||||
env=self.env, model="gpt-6.1-sol",
|
||||
retry_backoff_seconds=0)
|
||||
|
||||
def test_marker_delivered_and_drain_prefixed(self):
|
||||
_insert_task(self.db_path, "T1")
|
||||
_insert_task(self.db_path, "T2")
|
||||
worker = self._worker(_marker_runner(self.calls, MARKER_TO_T2))
|
||||
result = worker.run_one(task=worker.get_task("T1"))
|
||||
self.assertEqual(result["status"], "completed")
|
||||
# T1 的标记块投递成 T2 的后续消息（发送方 id 在列里）。
|
||||
fus = _followups(self.db_path, "T2")
|
||||
self.assertEqual(len(fus), 1)
|
||||
row = fus[0]
|
||||
# 列序与 DDL 一致：id, task_id, project_id, body, mode,
|
||||
# status, created_by, sender_task_id, created_at, updated_at。
|
||||
self.assertEqual(row[3], "请评审这份结论")
|
||||
self.assertEqual(row[5], "pending")
|
||||
self.assertEqual(row[7], "T1")
|
||||
# T2 跑起来：首轮原 prompt，排空轮次 prompt 带发送方前缀。
|
||||
result = worker.run_one(task=worker.get_task("T2"))
|
||||
self.assertEqual(result["status"], "completed")
|
||||
prompts = _prompt_values(self.calls)
|
||||
self.assertTrue(
|
||||
any(PREFIX in p and "请评审这份结论" in p for p in prompts),
|
||||
prompts)
|
||||
self.assertEqual(_followups(self.db_path, "T2")[0][5], "done")
|
||||
|
||||
def test_marker_revives_completed_target(self):
|
||||
_insert_task(self.db_path, "T1")
|
||||
_insert_task(self.db_path, "T2", status="completed")
|
||||
worker = self._worker(_marker_runner(self.calls, MARKER_TO_T2))
|
||||
result = worker.run_one(task=worker.get_task("T1"))
|
||||
self.assertEqual(result["status"], "completed")
|
||||
# T2 被复活：状态回 queued，续跑消息带发送方前缀。
|
||||
self.assertEqual(_status(self.db_path, "T2"), "queued")
|
||||
self.assertEqual(_resume_message(self.db_path, "T2"),
|
||||
PREFIX + "请评审这份结论")
|
||||
self.assertEqual(_followups(self.db_path, "T2"), [])
|
||||
# 复活后的 T2 续跑：prompt 就是续跑消息，跑完清掉。
|
||||
result = worker.run_one(task=worker.get_task("T2"))
|
||||
self.assertEqual(result["status"], "completed")
|
||||
prompts = _prompt_values(self.calls)
|
||||
self.assertTrue(any(PREFIX in p for p in prompts), prompts)
|
||||
self.assertIsNone(_resume_message(self.db_path, "T2"))
|
||||
|
||||
def test_failed_round_markers_not_delivered(self):
|
||||
_insert_task(self.db_path, "T1")
|
||||
_insert_task(self.db_path, "T2")
|
||||
|
||||
def failing(cmd):
|
||||
self.calls.append(list(cmd))
|
||||
return SimpleNamespace(
|
||||
returncode=1,
|
||||
stdout=json.dumps({"status": "completed",
|
||||
"text": MARKER_TO_T2}),
|
||||
stderr="boom")
|
||||
|
||||
worker = self._worker(failing)
|
||||
result = worker.run_one(task=worker.get_task("T1"))
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertEqual(_followups(self.db_path), [])
|
||||
|
||||
def test_self_and_cross_project_markers_rejected(self):
|
||||
_insert_task(self.db_path, "T1")
|
||||
_insert_task(self.db_path, "T3", project_id="p2")
|
||||
text = ("```task-message\n{\"to\": \"T1\", \"body\": \"自己\"}\n```\n"
|
||||
"```task-message\n{\"to\": \"T3\", \"body\": \"别项目\"}\n```\n")
|
||||
worker = self._worker(_marker_runner(self.calls, text))
|
||||
result = worker.run_one(task=worker.get_task("T1"))
|
||||
self.assertEqual(result["status"], "completed")
|
||||
self.assertEqual(_followups(self.db_path), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user