host
c4d3ec9bff
host：链路条目带 thread 归属；delegate 向父链路文件写 agent-start/done 标记。服务端：GET /tasks/{id}/agents 聚合父链路 + 子链路文件成代理链（状态口径：委派以结束标记为准、无标记即进行中；模型子代理无结束信号时父任务运行中为进行中、否则未知）。命令台：新增子代理链路区，随 SSE 流防抖刷新；扁平链路也渲染 agent 标记。
252 lines
10 KiB
Python
252 lines
10 KiB
Python
"""任务实时链路（服务端）单测：GET /tasks/{id}/trace 与 /trace/stream。
|
|
|
|
覆盖：鉴权（401/403/404）、快照去重与坏行跳过、worktree 在别机
|
|
时的空链路 + 说明、已结束任务的流（snapshot 后立刻 done 关流）。
|
|
不跑真实模型，不启动 worker；链路目录经 server.trace_dir 注入。
|
|
"""
|
|
|
|
import http.client
|
|
import json
|
|
import os
|
|
import socket
|
|
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:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.auth import AuthStore
|
|
from cloud.server import create_server
|
|
|
|
|
|
class TaskTraceTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.db_path = os.path.join(self.tmp.name, "control.db")
|
|
self.trace_dir = os.path.join(self.tmp.name, "host-traces")
|
|
os.makedirs(self.trace_dir, exist_ok=True)
|
|
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.trace_dir = self.trace_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)
|
|
status, self.task = self._req(
|
|
"POST", "/tasks",
|
|
{"project_id": self.project["id"], "prompt": "hi"},
|
|
self.owner_key)
|
|
self.assertEqual(status, 201)
|
|
self.task_id = self.task["id"]
|
|
|
|
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)
|
|
headers = {"Content-Type": "application/json"}
|
|
if key:
|
|
headers["X-API-Key"] = key
|
|
conn.request(method, path,
|
|
json.dumps(body) if body is not None else None, headers)
|
|
resp = conn.getresponse()
|
|
raw = resp.read()
|
|
conn.close()
|
|
try:
|
|
data = json.loads(raw)
|
|
except ValueError:
|
|
data = None
|
|
return resp.status, data
|
|
|
|
def _set_task(self, **fields):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
for k, v in fields.items():
|
|
conn.execute(
|
|
"UPDATE tasks SET %s = ? WHERE id = ?" % k,
|
|
(v, self.task_id))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def _write_trace(self, lines):
|
|
path = os.path.join(self.trace_dir, self.task_id + ".jsonl")
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
for line in lines:
|
|
f.write(line + "\n")
|
|
|
|
def test_trace_auth_and_not_found(self):
|
|
status, _ = self._req(
|
|
"GET", "/tasks/%s/trace" % self.task_id)
|
|
self.assertEqual(status, 401)
|
|
status, _ = self._req(
|
|
"GET", "/tasks/%s/trace" % self.task_id, key=self.other_key)
|
|
self.assertEqual(status, 403)
|
|
status, _ = self._req(
|
|
"GET", "/tasks/nope/trace", key=self.owner_key)
|
|
self.assertEqual(status, 404)
|
|
|
|
def test_trace_snapshot_dedup_and_bad_lines(self):
|
|
self._write_trace([
|
|
json.dumps({"kind": "exec", "id": "c1", "command": "ls",
|
|
"exitCode": None}),
|
|
"this is not json",
|
|
json.dumps({"kind": "exec", "id": "c1", "command": "ls",
|
|
"exitCode": 0, "durationMs": 5}),
|
|
json.dumps({"kind": "file", "id": "f1:a.py", "path": "a.py"}),
|
|
])
|
|
status, data = self._req(
|
|
"GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["status"], "queued")
|
|
# 同 id 后到的完成态覆盖先到的开始态，保首次顺序。
|
|
self.assertEqual([e["id"] for e in data["entries"]],
|
|
["c1", "f1:a.py"])
|
|
self.assertEqual(data["entries"][0]["exitCode"], 0)
|
|
self.assertEqual(data["note"], "")
|
|
|
|
def test_trace_missing_file_is_empty(self):
|
|
status, data = self._req(
|
|
"GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["entries"], [])
|
|
|
|
def test_trace_remote_worktree_host_is_empty_with_note(self):
|
|
self._write_trace([
|
|
json.dumps({"kind": "exec", "id": "c1", "command": "ls"})])
|
|
self._set_task(worktree_host="some-other-host")
|
|
status, data = self._req(
|
|
"GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["entries"], [])
|
|
self.assertIn("some-other-host", data["note"])
|
|
|
|
def _read_sse(self, path, key):
|
|
conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)
|
|
headers = {}
|
|
if key:
|
|
headers["X-API-Key"] = key
|
|
conn.request("GET", path, headers=headers)
|
|
resp = conn.getresponse()
|
|
raw = resp.read()
|
|
conn.close()
|
|
events = []
|
|
for block in raw.decode("utf-8").split("\n\n"):
|
|
for line in block.split("\n"):
|
|
if line.startswith("data: "):
|
|
events.append(json.loads(line[len("data: "):]))
|
|
return resp.status, events
|
|
|
|
def test_stream_finished_task_snapshot_then_done(self):
|
|
self._write_trace([
|
|
json.dumps({"kind": "exec", "id": "c1", "command": "ls",
|
|
"exitCode": 0})])
|
|
self._set_task(status="completed")
|
|
status, events = self._read_sse(
|
|
"/tasks/%s/trace/stream" % self.task_id, self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual([e["type"] for e in events], ["snapshot", "done"])
|
|
self.assertEqual(events[0]["entries"][0]["id"], "c1")
|
|
self.assertEqual(events[1]["status"], "completed")
|
|
|
|
def test_stream_auth(self):
|
|
status, events = self._read_sse(
|
|
"/tasks/%s/trace/stream" % self.task_id, None)
|
|
self.assertEqual(status, 401)
|
|
self.assertEqual(events, [])
|
|
|
|
def test_agents_auth_and_not_found(self):
|
|
status, _ = self._req(
|
|
"GET", "/tasks/%s/agents" % self.task_id)
|
|
self.assertEqual(status, 401)
|
|
status, _ = self._req(
|
|
"GET", "/tasks/%s/agents" % self.task_id, key=self.other_key)
|
|
self.assertEqual(status, 403)
|
|
status, _ = self._req(
|
|
"GET", "/tasks/nope/agents", key=self.owner_key)
|
|
self.assertEqual(status, 404)
|
|
|
|
def test_agents_chain_groups_calls_by_agent(self):
|
|
self._set_task(status="running")
|
|
self._write_trace([
|
|
json.dumps({"kind": "exec", "id": "c1", "thread": "th-main",
|
|
"command": "make", "ts": 1000}),
|
|
json.dumps({"kind": "tool", "id": "t1", "thread": "th-sub",
|
|
"server": "s", "tool": "read", "ts": 1001}),
|
|
json.dumps({"kind": "agent-start", "id": "agent:rev",
|
|
"name": "rev", "prompt": "review code",
|
|
"trace": "rev.jsonl", "ts": 1002}),
|
|
json.dumps({"kind": "agent-done", "id": "agent:rev:done",
|
|
"name": "rev", "status": "completed",
|
|
"trace": "rev.jsonl", "ts": 1004}),
|
|
])
|
|
with open(os.path.join(self.trace_dir, "rev.jsonl"),
|
|
"w", encoding="utf-8") as f:
|
|
f.write(json.dumps({"kind": "exec", "id": "c9",
|
|
"command": "grep", "ts": 1003}) + "\n")
|
|
status, data = self._req(
|
|
"GET", "/tasks/%s/agents" % self.task_id, key=self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["status"], "running")
|
|
by_kind = {a["kind"]: a for a in data["agents"]}
|
|
main = by_kind["main"]
|
|
self.assertEqual(main["status"], "running")
|
|
self.assertEqual([e["id"] for e in main["calls"]], ["c1"])
|
|
model = by_kind["model"]
|
|
self.assertEqual(model["status"], "running")
|
|
self.assertEqual([e["id"] for e in model["calls"]], ["t1"])
|
|
delegate = [a for a in data["agents"]
|
|
if a["kind"] == "delegate"][0]
|
|
self.assertEqual(delegate["name"], "rev")
|
|
self.assertEqual(delegate["status"], "completed")
|
|
self.assertEqual(delegate["prompt"], "review code")
|
|
self.assertEqual([e["id"] for e in delegate["calls"]], ["c9"])
|
|
self.assertEqual(delegate["updated_ts"], 1004)
|
|
|
|
def test_agents_delegate_without_done_stays_running(self):
|
|
# 委派进程与父任务相互独立：父任务已结束、无结束标记时，
|
|
# 子任务仍算进行中（不能谎报它已中断）。
|
|
self._set_task(status="completed")
|
|
self._write_trace([
|
|
json.dumps({"kind": "agent-start", "id": "agent:rev",
|
|
"name": "rev", "prompt": "review",
|
|
"trace": "rev.jsonl", "ts": 1002}),
|
|
])
|
|
status, data = self._req(
|
|
"GET", "/tasks/%s/agents" % self.task_id, key=self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
delegate = [a for a in data["agents"]
|
|
if a["kind"] == "delegate"][0]
|
|
self.assertEqual(delegate["status"], "running")
|
|
self.assertEqual(data["agents"][0]["status"], "completed")
|
|
|
|
def test_agents_missing_file_is_main_only(self):
|
|
status, data = self._req(
|
|
"GET", "/tasks/%s/agents" % self.task_id, key=self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(len(data["agents"]), 1)
|
|
self.assertEqual(data["agents"][0]["kind"], "main")
|
|
self.assertEqual(data["agents"][0]["calls"], [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|