wehub-resource-sync
385 lines
15 KiB
Python
385 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
pytestmark = [pytest.mark.asyncio, pytest.mark.e2e, pytest.mark.slow]
|
|
|
|
RUN_TIMEOUT_SECONDS = int(os.getenv("E2E_RUN_TIMEOUT_SECONDS", "300"))
|
|
|
|
|
|
def _assert_ok(response: httpx.Response) -> None:
|
|
assert response.status_code < 400, response.text
|
|
|
|
|
|
async def _create_agent(
|
|
client: httpx.AsyncClient,
|
|
headers: dict[str, str],
|
|
payload: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
response = await client.post("/api/agent", json=payload, headers=headers)
|
|
_assert_ok(response)
|
|
agent = response.json().get("agent")
|
|
assert isinstance(agent, dict), response.text
|
|
return agent
|
|
|
|
|
|
async def _delete_agent(client: httpx.AsyncClient, headers: dict[str, str], slug: str) -> None:
|
|
response = await client.delete(f"/api/agent/{slug}", headers=headers)
|
|
assert response.status_code in {200, 404}, response.text
|
|
|
|
|
|
async def _cancel_run(client: httpx.AsyncClient, headers: dict[str, str], run_id: str | None) -> None:
|
|
if not run_id:
|
|
return
|
|
response = await client.post(f"/api/agent/runs/{run_id}/cancel", headers=headers)
|
|
assert response.status_code < 500, response.text
|
|
|
|
|
|
async def _create_thread(client: httpx.AsyncClient, headers: dict[str, str], agent_id: str, marker: str) -> str:
|
|
response = await client.post(
|
|
"/api/chat/thread",
|
|
json={"agent_id": agent_id, "title": f"subagent-stream-e2e-{marker}", "metadata": {"marker": marker}},
|
|
headers=headers,
|
|
)
|
|
_assert_ok(response)
|
|
payload = response.json()
|
|
thread_id = payload.get("thread_id") or payload.get("id")
|
|
assert thread_id, payload
|
|
return str(thread_id)
|
|
|
|
|
|
async def _create_run(
|
|
client: httpx.AsyncClient,
|
|
headers: dict[str, str],
|
|
*,
|
|
agent_slug: str,
|
|
thread_id: str,
|
|
query: str,
|
|
) -> str:
|
|
response = await client.post(
|
|
"/api/agent/runs",
|
|
json={
|
|
"query": query,
|
|
"agent_slug": agent_slug,
|
|
"thread_id": thread_id,
|
|
"meta": {"request_id": f"subagent-stream-e2e-{uuid.uuid4()}"},
|
|
},
|
|
headers=headers,
|
|
)
|
|
_assert_ok(response)
|
|
run_id = response.json().get("run_id")
|
|
assert run_id, response.text
|
|
return str(run_id)
|
|
|
|
|
|
async def _iter_sse(client: httpx.AsyncClient, headers: dict[str, str], run_id: str):
|
|
async with client.stream("GET", f"/api/agent/runs/{run_id}/events", headers=headers) as response:
|
|
_assert_ok(response)
|
|
event = "message"
|
|
event_id = None
|
|
data_lines: list[str] = []
|
|
async for line in response.aiter_lines():
|
|
if not line:
|
|
if data_lines:
|
|
data_text = "\n".join(data_lines)
|
|
yield event, json.loads(data_text), event_id
|
|
event = "message"
|
|
event_id = None
|
|
data_lines = []
|
|
continue
|
|
if line.startswith(":"):
|
|
continue
|
|
if line.startswith("event:"):
|
|
event = line[len("event:") :].strip() or "message"
|
|
elif line.startswith("id:"):
|
|
event_id = line[len("id:") :].strip()
|
|
elif line.startswith("data:"):
|
|
data_lines.append(line[len("data:") :].strip())
|
|
|
|
|
|
def _collect_message_chunks(payload: dict[str, Any]) -> list[dict[str, Any]]:
|
|
chunks = []
|
|
chunk = payload.get("chunk")
|
|
if isinstance(chunk, dict):
|
|
chunks.append(chunk)
|
|
items = payload.get("items")
|
|
if isinstance(items, list):
|
|
chunks.extend(item for item in items if isinstance(item, dict))
|
|
return chunks
|
|
|
|
|
|
async def _consume_run_stream(
|
|
client: httpx.AsyncClient,
|
|
headers: dict[str, str],
|
|
run_id: str,
|
|
) -> tuple[dict[str, int], dict[str, Any], list[dict[str, Any]]]:
|
|
event_counts: dict[str, int] = {}
|
|
latest_agent_state: dict[str, Any] = {}
|
|
message_chunks: list[dict[str, Any]] = []
|
|
terminal_status = ""
|
|
|
|
async def consume() -> None:
|
|
nonlocal latest_agent_state, terminal_status
|
|
async for event, payload, _event_id in _iter_sse(client, headers, run_id):
|
|
event_counts[event] = event_counts.get(event, 0) + 1
|
|
if event == "messages":
|
|
message_chunks.extend(_collect_message_chunks(payload))
|
|
if event == "custom" and payload.get("name") == "yuxi.agent_state":
|
|
agent_state = payload.get("agent_state")
|
|
if isinstance(agent_state, dict):
|
|
latest_agent_state = agent_state
|
|
assert event != "error", payload
|
|
if event == "end":
|
|
event_payload = payload.get("payload") if isinstance(payload.get("payload"), dict) else payload
|
|
terminal_status = str(event_payload.get("status") or "")
|
|
return
|
|
|
|
await asyncio.wait_for(consume(), timeout=RUN_TIMEOUT_SECONDS)
|
|
assert terminal_status == "completed", {"status": terminal_status, "event_counts": event_counts}
|
|
return event_counts, latest_agent_state, message_chunks
|
|
|
|
|
|
def _find_tool_call_ids(value: Any) -> set[str]:
|
|
ids: set[str] = set()
|
|
if isinstance(value, dict):
|
|
tool_calls = value.get("tool_calls")
|
|
if isinstance(tool_calls, list):
|
|
for tool_call in tool_calls:
|
|
if isinstance(tool_call, dict) and tool_call.get("id"):
|
|
ids.add(str(tool_call["id"]))
|
|
for child in value.values():
|
|
ids.update(_find_tool_call_ids(child))
|
|
elif isinstance(value, list):
|
|
for item in value:
|
|
ids.update(_find_tool_call_ids(item))
|
|
return ids
|
|
|
|
|
|
async def _read_thread_file(
|
|
client: httpx.AsyncClient,
|
|
headers: dict[str, str],
|
|
thread_id: str,
|
|
path: str,
|
|
) -> str:
|
|
response = await client.get(
|
|
f"/api/chat/thread/{thread_id}/files/content",
|
|
params={"path": path},
|
|
headers=headers,
|
|
)
|
|
_assert_ok(response)
|
|
content = response.json().get("content")
|
|
assert isinstance(content, list), response.text
|
|
return "\n".join(str(line) for line in content)
|
|
|
|
|
|
async def test_subagent_stream_records_run_and_shares_output_files(
|
|
e2e_client: httpx.AsyncClient,
|
|
e2e_headers: dict[str, str],
|
|
):
|
|
me_response = await e2e_client.get("/api/auth/me", headers=e2e_headers)
|
|
_assert_ok(me_response)
|
|
me = me_response.json()
|
|
if me.get("role") not in {"admin", "superadmin"}:
|
|
pytest.skip("Subagent E2E needs an admin user to create temporary agents.")
|
|
uid = str(me.get("uid") or "")
|
|
assert uid, me
|
|
|
|
suffix = uuid.uuid4().hex[:8]
|
|
marker = f"YUXI_SUBAGENT_STREAM_E2E_{suffix}"
|
|
sub_slug = f"e2e-subagent-{suffix}"
|
|
main_slug = f"e2e-main-{suffix}"
|
|
output_path = "/home/gem/user-data/outputs/subagents.txt"
|
|
expected_content = "由这个子智能体创建"
|
|
created_agents: list[str] = []
|
|
run_id: str | None = None
|
|
run_completed = False
|
|
|
|
default_response = await e2e_client.get("/api/agent/default", headers=e2e_headers)
|
|
_assert_ok(default_response)
|
|
default_context = ((default_response.json().get("agent") or {}).get("config_json") or {}).get("context") or {}
|
|
base_context: dict[str, Any] = {"tools": [], "knowledges": [], "mcps": [], "skills": []}
|
|
if default_context.get("model"):
|
|
base_context["model"] = default_context["model"]
|
|
|
|
share_config = {"access_level": "user", "department_ids": [], "user_uids": [uid]}
|
|
|
|
try:
|
|
sub_agent = await _create_agent(
|
|
e2e_client,
|
|
e2e_headers,
|
|
{
|
|
"name": f"E2E 子智能体 {suffix}",
|
|
"slug": sub_slug,
|
|
"backend_id": "SubAgentBackend",
|
|
"description": "真实流式 E2E 子智能体",
|
|
"config_json": {
|
|
"context": {
|
|
**base_context,
|
|
"system_prompt": (
|
|
"你是专门负责文件写入和文件校验的子智能体。收到任务后必须使用文件系统工具完成任务，"
|
|
"不要向用户提问。必须严格写入用户指定路径，文件内容必须完全符合用户要求，"
|
|
"不要自动追加句号、引号、说明或其他字符。完成后只回复写入的路径和文件内容。"
|
|
),
|
|
}
|
|
},
|
|
"share_config": share_config,
|
|
"is_subagent": True,
|
|
},
|
|
)
|
|
created_agents.append(sub_slug)
|
|
|
|
await _create_agent(
|
|
e2e_client,
|
|
e2e_headers,
|
|
{
|
|
"name": f"E2E 主智能体 {suffix}",
|
|
"slug": main_slug,
|
|
"backend_id": "ChatbotAgent",
|
|
"description": "真实流式 E2E 主智能体",
|
|
"config_json": {
|
|
"context": {
|
|
**base_context,
|
|
"subagents": [sub_slug],
|
|
"system_prompt": (
|
|
"你是主智能体。遇到用户要求创建、修改或验证文件的任务时，"
|
|
"必须调用 task 工具交给可用子智能体完成，不要自己写文件，"
|
|
"也不要通过 shell、curl 或 HTTP API 调用子智能体。子智能体完成后，简短汇总结果。"
|
|
),
|
|
}
|
|
},
|
|
"share_config": share_config,
|
|
"is_subagent": False,
|
|
},
|
|
)
|
|
created_agents.append(main_slug)
|
|
|
|
default_agents_response = await e2e_client.get("/api/agent", headers=e2e_headers)
|
|
_assert_ok(default_agents_response)
|
|
default_agent_slugs = {str(item.get("slug")) for item in default_agents_response.json().get("agents") or []}
|
|
assert sub_slug not in default_agent_slugs
|
|
|
|
management_agents_response = await e2e_client.get("/api/agent?include_subagents=true", headers=e2e_headers)
|
|
_assert_ok(management_agents_response)
|
|
management_agent_slugs = {
|
|
str(item.get("slug")) for item in management_agents_response.json().get("agents") or []
|
|
}
|
|
assert sub_slug in management_agent_slugs
|
|
|
|
thread_id = await _create_thread(e2e_client, e2e_headers, main_slug, marker)
|
|
query = (
|
|
f"请调用子智能体 {sub_slug} 在 outputs 目录创建文件 {output_path}。"
|
|
"必须通过 task 工具调用子智能体完成。"
|
|
f"文件内容必须完全等于下面一行：\n{expected_content}\n"
|
|
"不要添加句号、引号、说明或其他任何字符。完成后只需要回复文件路径。"
|
|
)
|
|
run_id = await _create_run(
|
|
e2e_client,
|
|
e2e_headers,
|
|
agent_slug=main_slug,
|
|
thread_id=thread_id,
|
|
query=query,
|
|
)
|
|
|
|
event_counts, stream_agent_state, message_chunks = await _consume_run_stream(
|
|
e2e_client,
|
|
e2e_headers,
|
|
run_id,
|
|
)
|
|
assert event_counts.get("messages", 0) > 0
|
|
|
|
run_response = await e2e_client.get(f"/api/agent/runs/{run_id}", headers=e2e_headers)
|
|
_assert_ok(run_response)
|
|
assert (run_response.json().get("run") or {}).get("status") == "completed"
|
|
|
|
state_response = await e2e_client.get(f"/api/chat/thread/{thread_id}/state", headers=e2e_headers)
|
|
_assert_ok(state_response)
|
|
final_agent_state = state_response.json().get("agent_state") or stream_agent_state
|
|
subagent_runs = final_agent_state.get("subagent_runs") or []
|
|
assert subagent_runs, final_agent_state
|
|
completed_run = next((item for item in subagent_runs if item.get("status") == "completed"), subagent_runs[0])
|
|
assert completed_run.get("subagent_slug") == sub_slug
|
|
assert completed_run.get("subagent_name") == sub_agent["name"]
|
|
assert completed_run.get("child_thread_id")
|
|
assert completed_run.get("id")
|
|
|
|
child_thread_id = str(completed_run["child_thread_id"])
|
|
child_state_response = await e2e_client.get(
|
|
f"/api/chat/thread/{child_thread_id}/state",
|
|
params={"include_messages": "true"},
|
|
headers=e2e_headers,
|
|
)
|
|
_assert_ok(child_state_response)
|
|
child_state_payload = child_state_response.json()
|
|
assert child_state_payload.get("parent_thread_id") == thread_id
|
|
child_subagent_run = child_state_payload.get("subagent_run") or {}
|
|
assert child_subagent_run.get("child_thread_id") == child_thread_id
|
|
assert child_subagent_run.get("run_id")
|
|
child_run_response = await e2e_client.get(
|
|
f"/api/agent/runs/{child_subagent_run['run_id']}",
|
|
headers=e2e_headers,
|
|
)
|
|
_assert_ok(child_run_response)
|
|
child_run = child_run_response.json().get("run") or {}
|
|
assert child_run.get("run_type") == "subagent"
|
|
assert child_run.get("conversation_thread_id") == child_thread_id
|
|
assert child_run.get("created_by_run_id") == run_id
|
|
assert child_run.get("status") == "completed"
|
|
assert child_state_payload.get("messages"), child_state_payload
|
|
|
|
leaked_child_chunks = [
|
|
chunk for chunk in message_chunks if child_thread_id in json.dumps(chunk, ensure_ascii=False, default=str)
|
|
]
|
|
assert leaked_child_chunks == []
|
|
|
|
history_response = await e2e_client.get(f"/api/chat/thread/{thread_id}/history", headers=e2e_headers)
|
|
_assert_ok(history_response)
|
|
history_payload = history_response.json()
|
|
tool_call_ids = _find_tool_call_ids(history_payload)
|
|
assert str(completed_run["id"]) in tool_call_ids
|
|
assert child_thread_id in json.dumps(history_payload, ensure_ascii=False)
|
|
|
|
files_response = await e2e_client.get(
|
|
f"/api/chat/thread/{thread_id}/files",
|
|
params={"path": "/home/gem/user-data/outputs", "recursive": "true"},
|
|
headers=e2e_headers,
|
|
)
|
|
_assert_ok(files_response)
|
|
files_payload = files_response.json()
|
|
file_paths = {str(item.get("path") or "") for item in files_payload.get("files") or []}
|
|
assert output_path in file_paths, {
|
|
"output_path": output_path,
|
|
"files": files_payload,
|
|
"subagent_run": completed_run,
|
|
}
|
|
assert (await _read_thread_file(e2e_client, e2e_headers, thread_id, output_path)).strip() == expected_content
|
|
|
|
tree_response = await e2e_client.get(
|
|
"/api/viewer/filesystem/tree",
|
|
params={"thread_id": thread_id, "path": "/home/gem/user-data/outputs"},
|
|
headers=e2e_headers,
|
|
)
|
|
_assert_ok(tree_response)
|
|
assert output_path in json.dumps(tree_response.json(), ensure_ascii=False)
|
|
|
|
viewer_file_response = await e2e_client.get(
|
|
"/api/viewer/filesystem/file",
|
|
params={"thread_id": thread_id, "path": output_path},
|
|
headers=e2e_headers,
|
|
)
|
|
_assert_ok(viewer_file_response)
|
|
assert expected_content in json.dumps(viewer_file_response.json(), ensure_ascii=False)
|
|
run_completed = True
|
|
|
|
finally:
|
|
if not run_completed:
|
|
await _cancel_run(e2e_client, e2e_headers, run_id)
|
|
for slug in reversed(created_agents):
|
|
await _delete_agent(e2e_client, e2e_headers, slug)
|