host
a2863b2b10
此前服务端只读共享链路目录，租户隔离的 Docker 任务链路文件落在 tenants/<项目>/work/host-traces/，/trace 与 /agents 会误报为空。改后按任务解析目录列表（租户优先、共享兜底），父链路与子链路文件多目录合并读；直接发送的续写跟到已有链路文件所在的目录，不把一条任务的链路拆成两处。
337 lines
14 KiB
Python
337 lines
14 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, _tenant_trace_dir
|
|
except ImportError:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.auth import AuthStore
|
|
from cloud.server import create_server, _tenant_trace_dir
|
|
|
|
|
|
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.tenants_root = os.path.join(self.tmp.name, "tenants")
|
|
os.makedirs(self.tenants_root, 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.server.tenants_root = self.tenants_root
|
|
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 _write_tenant_trace(self, lines, filename=None):
|
|
d = os.path.join(self.tenants_root, self.project["id"],
|
|
"work", "host-traces")
|
|
os.makedirs(d, exist_ok=True)
|
|
path = os.path.join(d, filename or (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"], [])
|
|
|
|
def test_tenant_trace_dir_mapping(self):
|
|
class S:
|
|
tenants_root = "/tenants"
|
|
self.assertEqual(
|
|
_tenant_trace_dir(S(), "abc-123"),
|
|
"/tenants/abc-123/work/host-traces")
|
|
self.assertIsNone(_tenant_trace_dir(S(), None))
|
|
self.assertIsNone(_tenant_trace_dir(S(), ""))
|
|
self.assertIsNone(_tenant_trace_dir(S(), ".."))
|
|
self.assertIsNone(_tenant_trace_dir(S(), "..."))
|
|
|
|
def test_trace_reads_tenant_dir_only(self):
|
|
# 租户 Docker 任务的链路文件只在租户目录里：旧版只读
|
|
# 共享目录会误报为空，修后 /trace 与 /agents 都看得见。
|
|
self._write_tenant_trace([
|
|
json.dumps({"kind": "exec", "id": "c1", "command": "make"}),
|
|
])
|
|
status, data = self._req(
|
|
"GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual([e["id"] for e in data["entries"]], ["c1"])
|
|
status, data = self._req(
|
|
"GET", "/tasks/%s/agents" % self.task_id, key=self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual([e["id"] for e in data["agents"][0]["calls"]],
|
|
["c1"])
|
|
|
|
def test_trace_merges_tenant_and_shared(self):
|
|
self._write_tenant_trace([
|
|
json.dumps({"kind": "exec", "id": "c1", "command": "tenant"}),
|
|
])
|
|
self._write_trace([
|
|
json.dumps({"kind": "exec", "id": "c2", "command": "shared"}),
|
|
])
|
|
status, data = self._req(
|
|
"GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual([e["id"] for e in data["entries"]], ["c1", "c2"])
|
|
# 同 id 冲突时后读的共享目录覆盖（顺序有文档、可复现）。
|
|
self._write_tenant_trace([
|
|
json.dumps({"kind": "exec", "id": "c1",
|
|
"command": "tenant-cmd"}),
|
|
])
|
|
self._write_trace([
|
|
json.dumps({"kind": "exec", "id": "c1",
|
|
"command": "shared-cmd"}),
|
|
])
|
|
status, data = self._req(
|
|
"GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
|
|
self.assertEqual(data["entries"][0]["command"], "shared-cmd")
|
|
|
|
def test_agents_tenant_child_file(self):
|
|
# 委派子任务的子链路文件也在租户目录里：父标记的
|
|
# trace 名要在租户目录里找得到，子任务调用才进链。
|
|
self._write_tenant_trace([
|
|
json.dumps({"kind": "agent-start", "id": "agent:rev",
|
|
"name": "rev", "prompt": "review",
|
|
"trace": "rev.jsonl", "ts": 1002}),
|
|
json.dumps({"kind": "agent-done", "id": "agent:rev:done",
|
|
"name": "rev", "status": "completed",
|
|
"trace": "rev.jsonl", "ts": 1004}),
|
|
])
|
|
self._write_tenant_trace([
|
|
json.dumps({"kind": "exec", "id": "c9", "command": "grep"}),
|
|
], filename="rev.jsonl")
|
|
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"], "completed")
|
|
self.assertEqual([e["id"] for e in delegate["calls"]], ["c9"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|